mlx-taef 0.1.0__py3-none-any.whl

mlx_taef/__init__.py

@@ -0,0 +1,9 @@
+"""mlx-taef: Tiny AutoEncoder family ported to Apple MLX."""
+
+import logging
+
+from mlx_taef.api import TAEF1, TAEF2, TAESD, TAESDXL, Taef
+
+__all__ = ["TAEF1", "TAEF2", "TAESD", "TAESDXL", "Taef"]
+
+logging.getLogger("mlx_taef").addHandler(logging.NullHandler())

mlx_taef/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

mlx_taef/api.py

@@ -0,0 +1,204 @@
+"""User-facing Taef family API: load weights and decode/encode latents.
+
+Tensor layout: all public methods use NHWC mx.array.
+Value space: decode() outputs [0, 1] float; encode() expects [0, 1] float.
+"""
+
+import logging
+from pathlib import Path
+from typing import cast
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_taef.model import make_decoder, make_encoder
+from mlx_taef.variants import (
+    TAEF1_CONFIG,
+    TAEF2_CONFIG,
+    TAESD_CONFIG,
+    TAESDXL_CONFIG,
+    TaesdVariantConfig,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Taef(nn.Module):  # type: ignore[misc,name-defined]
+    """Base class for TAESD-family models.
+
+    Subclasses set the `_config` class variable to a `TaesdVariantConfig`.
+
+    Args:
+        None — instantiated via `from_pretrained_local(path)`.
+
+    Tensor layout:
+        All public methods use NHWC `mx.array`:
+        - latent shape: (B, H, W, latent_channels)
+        - image shape: (B, H*8, W*8, 3)
+
+    Value space:
+        - `decode()` returns float in [0, 1] (clipped on output).
+        - `decode_image()` returns uint8 in [0, 255].
+        - Latents are RAW (post-Clamp scale), values roughly in [-3, 3].
+        - `scale_latents` / `unscale_latents` convert between raw and [0, 1].
+    """
+
+    _config: TaesdVariantConfig
+
+    def __init__(self) -> None:
+        """Initialise decoder and encoder from the subclass `_config`."""
+        super().__init__()
+        self.decoder = make_decoder(self._config)
+        self.encoder = make_encoder(self._config)
+
+    @classmethod
+    def from_pretrained_local(
+        cls,
+        decoder_path: Path | str,
+        encoder_path: Path | str | None = None,
+        *,
+        dtype: mx.Dtype = mx.float32,
+    ) -> "Taef":
+        """Instantiate from already-converted MLX safetensors on disk.
+
+        Args:
+            decoder_path: path to converted MLX decoder safetensors.
+            encoder_path: optional path to converted MLX encoder safetensors (Phase 2).
+            dtype: target dtype for model weights. Default fp32 to keep parity tests sharp.
+
+        Returns:
+            A loaded `Taef` instance ready for `decode()`.
+        """
+        instance = cls()
+        d_weights = cast("dict[str, mx.array]", mx.load(str(decoder_path)))
+        weights_list: list[tuple[str, mx.array]] = [
+            (f"decoder.{k}", v) for k, v in d_weights.items()
+        ]
+        if encoder_path is not None and hasattr(instance, "encoder"):
+            e_weights = cast("dict[str, mx.array]", mx.load(str(encoder_path)))
+            weights_list.extend((f"encoder.{k}", v) for k, v in e_weights.items())
+        instance.load_weights(weights_list, strict=False)
+        if dtype is not mx.float32:
+            instance.set_dtype(dtype)
+        instance.eval()
+        return instance
+
+    @classmethod
+    def from_pretrained(  # pragma: no cover
+        cls,
+        repo_id: str | None = None,
+        *,
+        dtype: mx.Dtype = mx.float32,
+        include_encoder: bool = True,
+    ) -> "Taef":
+        """Auto-download weights from HF Hub, convert to MLX, and load.
+
+        On first call, downloads upstream weights from `cls._config.hf_repo` and
+        caches converted output at ~/.cache/mlx-taef/. Subsequent calls return
+        instantly from cache.
+
+        Args:
+            repo_id: optional HF repo override. If set, must match `cls._config.hf_repo`.
+            dtype: target dtype for weights. Default fp32 to keep parity strict.
+            include_encoder: whether to also load the encoder side (default True).
+
+        Returns:
+            A loaded `Taef` instance.
+        """
+        from mlx_taef.download import get_or_convert
+
+        config = cls._config
+        if repo_id is not None and repo_id != config.hf_repo:
+            raise ValueError(
+                f"repo_id mismatch: requested {repo_id!r} but variant {config.name!r} "
+                f"uses {config.hf_repo!r}"
+            )
+        decoder_path = get_or_convert(config, role="decoder")
+        encoder_path = get_or_convert(config, role="encoder") if include_encoder else None
+        return cls.from_pretrained_local(decoder_path, encoder_path=encoder_path, dtype=dtype)
+
+    def decode(self, latents: mx.array) -> mx.array:
+        """Decode raw latents (NHWC) to image (NHWC, [0, 1] float).
+
+        Args:
+            latents: NHWC array with shape (B, H, W, latent_channels).
+
+        Returns:
+            NHWC float array with shape (B, H*8, W*8, 3), values in [0, 1].
+        """
+        out = self.decoder(latents)
+        return mx.clip(out, 0.0, 1.0)
+
+    def decode_image(self, latents: mx.array) -> mx.array:
+        """Decode raw latents to a uint8 NHWC image suitable for PIL/PNG.
+
+        Args:
+            latents: NHWC array with shape (B, H, W, latent_channels).
+
+        Returns:
+            NHWC uint8 array with shape (B, H*8, W*8, 3), values in [0, 255].
+        """
+        img_float = self.decode(latents)
+        return (img_float * 255.0).astype(mx.uint8)
+
+    def encode(self, image: mx.array) -> mx.array:
+        """Encode an NHWC RGB image (B, H, W, 3) in [0, 1] to a latent (B, H/8, W/8, latent_channels).
+
+        Args:
+            image: NHWC float array with shape (B, H, W, 3), values in [0, 1].
+
+        Returns:
+            NHWC latent array with shape (B, H/8, W/8, latent_channels).
+        """
+        return cast("mx.array", self.encoder(image))
+
+    def scale_latents(self, raw: mx.array) -> mx.array:
+        """Map raw latents to [0, 1] using config.latent_magnitude/shift.
+
+        Args:
+            raw: raw latent array, values roughly in [-3, 3].
+
+        Returns:
+            Scaled latent array clipped to [0, 1].
+        """
+        c = self._config
+        return mx.clip(raw / (2.0 * c.latent_magnitude) + c.latent_shift, 0.0, 1.0)
+
+    def unscale_latents(self, scaled: mx.array) -> mx.array:
+        """Inverse of scale_latents: [0, 1] -> raw.
+
+        Args:
+            scaled: latent array in [0, 1].
+
+        Returns:
+            Raw latent array.
+        """
+        c = self._config
+        return (scaled - c.latent_shift) * (2.0 * c.latent_magnitude)
+
+
+class TAESD(Taef):
+    """TAESD for Stable Diffusion 1.x."""
+
+    _config = TAESD_CONFIG
+
+
+class TAESDXL(Taef):
+    """TAESD for SDXL."""
+
+    _config = TAESDXL_CONFIG
+
+
+class TAEF1(Taef):
+    """TAEF1 for FLUX.1."""
+
+    _config = TAEF1_CONFIG
+
+
+class TAEF2(Taef):
+    """TAEF2 for FLUX.2 Klein."""
+
+    _config = TAEF2_CONFIG
+
+
+__all__ = ["TAEF1", "TAEF2", "TAESD", "TAESDXL", "Taef"]

mlx_taef/cli.py

@@ -0,0 +1,92 @@
+"""mlx-taef command-line interface: convert / info / bench."""
+
+import argparse
+import logging
+import time
+from pathlib import Path
+
+import mlx.core as mx
+
+from mlx_taef.variants import ALL_VARIANTS
+
+logger = logging.getLogger(__name__)
+
+
+def main(argv: list[str] | None = None) -> int:
+    """CLI entry point. Returns process exit code."""
+    parser = argparse.ArgumentParser(
+        prog="mlx-taef",
+        description="Tiny AutoEncoder family for diffusion on Apple MLX.",
+    )
+    sub = parser.add_subparsers(dest="cmd", required=True)
+
+    variant_names = [v.name for v in ALL_VARIANTS]
+
+    p_convert = sub.add_parser("convert", help="Download upstream weights and convert to MLX")
+    p_convert.add_argument("--variant", required=True, choices=variant_names)
+    p_convert.add_argument("--role", default="decoder", choices=["decoder", "encoder"])
+    p_convert.add_argument("--dst", required=True, type=Path, help="Output .safetensors path")
+
+    p_info = sub.add_parser("info", help="Print info about a converted MLX safetensors file")
+    p_info.add_argument("path", type=Path)
+
+    p_bench = sub.add_parser("bench", help="Decode benchmark on the current Mac")
+    p_bench.add_argument("--variant", default="taef2", choices=variant_names)
+
+    args = parser.parse_args(argv)
+    if args.cmd == "convert":  # pragma: no cover
+        return _cmd_convert(args)
+    if args.cmd == "info":
+        return _cmd_info(args)
+    if args.cmd == "bench":  # pragma: no cover
+        return _cmd_bench(args)
+    return 1  # pragma: no cover
+
+
+def _cmd_convert(args: argparse.Namespace) -> int:  # pragma: no cover
+    from mlx_taef.convert import convert_hf_decoder_to_mlx, convert_hf_encoder_to_mlx
+
+    config = next(v for v in ALL_VARIANTS if v.name == args.variant)
+    if args.role == "encoder":
+        convert_hf_encoder_to_mlx(out_path=args.dst, config=config)
+    else:
+        convert_hf_decoder_to_mlx(out_path=args.dst, config=config)
+    print(f"Wrote {args.dst}")
+    return 0
+
+
+def _cmd_info(args: argparse.Namespace) -> int:
+    weights = mx.load(str(args.path))
+    print(f"File: {args.path}")
+    print(f"Total tensors: {len(weights)}")
+    total_params = sum(int(w.size) for w in weights.values())  # type: ignore[misc,union-attr]
+    print(f"Total params: {total_params:,}")
+    return 0
+
+
+def _cmd_bench(args: argparse.Namespace) -> int:  # pragma: no cover
+    from mlx_taef.api import TAEF1, TAEF2, TAESD, TAESDXL
+
+    cls_by_name = {"taesd": TAESD, "taesdxl": TAESDXL, "taef1": TAEF1, "taef2": TAEF2}
+    cls = cls_by_name[args.variant]
+    config = next(v for v in ALL_VARIANTS if v.name == args.variant)
+
+    model = cls.from_pretrained(include_encoder=False)
+    # 1024x1024 image with 8x downsample = 128x128 latent
+    latent = mx.random.normal((1, 128, 128, config.latent_channels)).astype(mx.float16)
+    mx.eval(latent)
+
+    # Warm-up
+    mx.eval(model.decode(latent))
+
+    times: list[float] = []
+    for _ in range(5):
+        start = time.perf_counter()
+        mx.eval(model.decode(latent))
+        times.append(time.perf_counter() - start)
+    median_ms = sorted(times)[len(times) // 2] * 1000
+    print(f"{args.variant} decode median: {median_ms:.1f} ms over {len(times)} runs")
+    return 0
+
+
+__all__ = ["main"]

mlx_taef/convert.py

@@ -0,0 +1,201 @@
+"""HF safetensors -> MLX safetensors conversion.
+
+Zero PyTorch dependency: reads source files with `safetensors.numpy.load_file`
+and writes MLX safetensors directly. Runtime users never need torch.
+"""
+
+import logging
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+import numpy as np
+from huggingface_hub import hf_hub_download
+from safetensors.numpy import load_file as safetensors_load_numpy
+
+from mlx_taef.model import make_decoder
+from mlx_taef.variants import TaesdVariantConfig
+
+logger = logging.getLogger(__name__)
+
+
+def convert_diffusers_to_sequential(
+    sd: dict[str, Any],
+    *,
+    role: str,
+) -> dict[str, np.ndarray]:
+    """Map Diffusers-VAE keys to upstream Sequential-key format.
+
+    Per the TAEF2 model card, the decoder gets a +1 index shift because the
+    Diffusers VAE prepends one layer that the upstream Sequential decoder
+    doesn't have. Encoder keys have no offset.
+
+    Args:
+        sd: source state dict with Diffusers keys like 'decoder.layers.0.weight'.
+        role: 'decoder' (apply +1 offset) or 'encoder' (no offset).
+
+    Returns:
+        State dict with upstream-Sequential keys like '0.weight', '1.weight'.
+        Keys not matching role-prefix are filtered out.
+    """
+    out: dict[str, np.ndarray] = {}
+    prefix = f"{role}."
+    for k, v in sd.items():
+        if not k.startswith(prefix):
+            continue
+        suffix = k[len(prefix) :]
+        if suffix.startswith("layers."):
+            parts = suffix.split(".")
+            idx = int(parts[1])
+            if role == "decoder":
+                idx += 1
+            new_key = f"{idx}." + ".".join(parts[2:])
+        else:  # pragma: no cover
+            new_key = suffix
+        out[new_key] = v
+    return out
+
+
+def _load_role_state_dict(  # pragma: no cover
+    config: TaesdVariantConfig,
+    role: str,
+) -> dict[str, np.ndarray]:
+    """Download and load weights for (variant, role) into a Sequential-keyed dict."""
+    if config.key_format == "diffusers":
+        if config.hf_filename is None:
+            raise ValueError(f"Diffusers variant {config.name!r} has no hf_filename")
+        path = hf_hub_download(repo_id=config.hf_repo, filename=config.hf_filename)
+        full_sd = safetensors_load_numpy(path)
+        return convert_diffusers_to_sequential(full_sd, role=role)
+    if config.key_format == "upstream":
+        filename = config.hf_decoder_filename if role == "decoder" else config.hf_encoder_filename
+        if filename is None:
+            raise ValueError(f"Upstream variant {config.name!r} has no {role} filename")
+        path = hf_hub_download(repo_id=config.hf_repo, filename=filename)
+        return safetensors_load_numpy(path)
+    raise ValueError(f"Unknown key_format: {config.key_format!r}")
+
+
+def convert_hf_decoder_to_mlx(  # pragma: no cover
+    *,
+    out_path: Path | str,
+    config: TaesdVariantConfig,
+) -> None:
+    """Download upstream decoder weights, convert to MLX safetensors at `out_path`.
+
+    Handles both upstream-Sequential and Diffusers key formats. Transposes
+    Conv2d weights from NCHW to NHWC. Writes the result with MLX-flat keys
+    like 'layers.0.weight', 'layers.1.weight', ...
+
+    Args:
+        out_path: where to write the MLX safetensors file.
+        config: variant configuration.
+    """
+    sd = _load_role_state_dict(config, role="decoder")
+    decoder = make_decoder(config)
+    expected = _flatten_module_param_shapes(decoder)
+    converted = _build_mlx_state_dict(sd, expected_shapes=expected)
+    mx.save_safetensors(str(out_path), converted)
+
+
+def _sequential_key_to_mlx(src_key: str) -> str:
+    """Convert an upstream-Sequential key to an MLX-flat dotted key.
+
+    MLX's `nn.Sequential` stores its children under `.layers`, so every
+    integer path segment (after the first top-level layer index) must be
+    wrapped as `layers.<N>` rather than a bare `<N>`.
+
+    Examples::
+
+        "1.weight"           -> "layers.1.weight"
+        "3.conv.0.weight"    -> "layers.3.conv.layers.0.weight"
+        "3.pool.1.bias"      -> "layers.3.pool.layers.1.bias"
+
+    Args:
+        src_key: upstream-Sequential key like '3.conv.0.weight'.
+
+    Returns:
+        MLX-flat dotted key like 'layers.3.conv.layers.0.weight'.
+    """
+    parts = src_key.split(".")
+    out = ["layers", parts[0]]
+    for part in parts[1:]:
+        if part.isdigit():
+            out.extend(["layers", part])
+        else:
+            out.append(part)
+    return ".".join(out)
+
+
+def _build_mlx_state_dict(
+    sd: dict[str, np.ndarray],
+    *,
+    expected_shapes: dict[str, tuple[int, ...]],
+) -> dict[str, mx.array]:
+    """Apply NCHW->NHWC transpose for Conv weights and prefix keys with 'layers.'."""
+    converted: dict[str, mx.array] = {}
+    for src_key, arr in sd.items():
+        dst_key = _sequential_key_to_mlx(src_key)
+        if dst_key not in expected_shapes:
+            # Skip keys that don't map to the MLX module structure
+            # (e.g., extra Diffusers-specific keys we don't need)
+            continue
+        # Conv2d weight transpose NCHW (out, in, kH, kW) -> NHWC (out, kH, kW, in)
+        # Detected when source is 4D and expected MLX shape matches the transposed shape.
+        if arr.ndim == 4 and expected_shapes[dst_key] == (
+            arr.shape[0],
+            arr.shape[2],
+            arr.shape[3],
+            arr.shape[1],
+        ):
+            arr = np.transpose(arr, (0, 2, 3, 1)).copy()
+        converted[dst_key] = mx.array(arr)
+    return converted
+
+
+def _flatten_module_param_shapes(module: Any, prefix: str = "") -> dict[str, tuple[int, ...]]:
+    """Walk module.parameters() and return a flat dict of dotted-key -> shape."""
+    out: dict[str, tuple[int, ...]] = {}
+
+    def _walk(obj: Any, p: str) -> None:
+        if isinstance(obj, dict):
+            for k, v in obj.items():
+                _walk(v, f"{p}.{k}" if p else k)
+        elif isinstance(obj, list):
+            for i, item in enumerate(obj):
+                _walk(item, f"{p}.{i}")
+        elif hasattr(obj, "shape"):
+            out[p] = tuple(obj.shape)
+
+    _walk(module.parameters(), prefix)
+    return out
+
+
+def convert_hf_encoder_to_mlx(  # pragma: no cover
+    *,
+    out_path: Path | str,
+    config: TaesdVariantConfig,
+) -> None:
+    """Download upstream encoder weights, convert to MLX safetensors at `out_path`.
+
+    Mirrors `convert_hf_decoder_to_mlx` but introspects via `make_encoder` so
+    Conv weights are transposed against the correct shapes.
+
+    Args:
+        out_path: where to write the MLX safetensors file.
+        config: variant configuration.
+    """
+    from mlx_taef.model import make_encoder
+
+    sd = _load_role_state_dict(config, role="encoder")
+    encoder = make_encoder(config)
+    expected = _flatten_module_param_shapes(encoder)
+    converted = _build_mlx_state_dict(sd, expected_shapes=expected)
+    mx.save_safetensors(str(out_path), converted)
+
+
+__all__ = [
+    "convert_diffusers_to_sequential",
+    "convert_hf_decoder_to_mlx",
+    "convert_hf_encoder_to_mlx",
+]

mlx_taef/download.py

@@ -0,0 +1,45 @@
+"""HF Hub auto-download + cache. Zero PyTorch dependency at runtime."""
+
+import logging
+from pathlib import Path
+
+from mlx_taef.convert import convert_hf_decoder_to_mlx, convert_hf_encoder_to_mlx
+from mlx_taef.variants import TaesdVariantConfig
+
+logger = logging.getLogger(__name__)
+
+CACHE_ROOT = Path.home() / ".cache" / "mlx-taef"
+
+
+def get_or_convert(config: TaesdVariantConfig, *, role: str = "decoder") -> Path:
+    """Return the local path to converted MLX weights for (variant, role).
+
+    On cache miss, triggers the full conversion pipeline (HF download + key
+    remap + NHWC transpose + safetensors write). Subsequent calls return the
+    cached path without any network access.
+
+    Args:
+        config: variant configuration (selects HF repo + filename).
+        role: 'decoder' (default) or 'encoder'.
+
+    Returns:
+        Local filesystem path to the MLX safetensors file.
+    """
+    cache_dir = CACHE_ROOT / config.name
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    out_path = cache_dir / f"{config.name}_{role}.safetensors"
+    if out_path.exists():
+        logger.debug("Using cached weights at %s", out_path)
+        return out_path
+
+    logger.info("Downloading + converting %s %s weights from %s", config.name, role, config.hf_repo)
+    if role == "decoder":
+        convert_hf_decoder_to_mlx(out_path=out_path, config=config)
+    elif role == "encoder":
+        convert_hf_encoder_to_mlx(out_path=out_path, config=config)
+    else:
+        raise ValueError(f"role must be 'decoder' or 'encoder', got {role!r}")
+    return out_path
+
+
+__all__ = ["CACHE_ROOT", "get_or_convert"]

mlx_taef/integrations/__init__.py

@@ -0,0 +1 @@
+"""Optional integrations for mlx-taef (e.g. mflux)."""

mlx_taef/integrations/mflux.py

@@ -0,0 +1,168 @@
+"""mflux integration for mlx-taef.
+
+Provides:
+- `unpack_flux2_latent`: convert mflux's packed FLUX.2 latents to TAEF2-compatible NHWC.
+- `LivePreviewCallback`: drop-in mflux callback that writes preview PNGs every N steps.
+
+Install with: `pip install "mlx-taef[mflux]"`.
+"""
+
+import logging
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+
+try:
+    from mflux.callbacks.callback import InLoopCallback  # type: ignore[import-untyped]
+except ImportError as e:  # pragma: no cover
+    raise ImportError(
+        "mflux is required for mlx_taef.integrations.mflux. "
+        "Install with: pip install 'mlx-taef[mflux]'"
+    ) from e
+
+from mlx_taef.api import TAEF1, TAEF2, Taef
+
+logger = logging.getLogger(__name__)
+
+
+def unpack_flux2_latent(
+    packed: mx.array,
+    *,
+    latent_height: int,
+    latent_width: int,
+    bn_mean: mx.array | None = None,
+    bn_var: mx.array | None = None,
+    bn_eps: float = 1e-4,
+) -> mx.array:
+    """Unpack mflux's packed FLUX.2 latent into NHWC `(B, lH*2, lW*2, 32)` for TAEF2.
+
+    Mirrors mflux's `Flux2VAE.decode_packed_latents` preprocessing:
+    1. BN denormalize: latents = packed * sqrt(var + eps) + mean  (128-channel stats)
+    2. Unpatchify: (B, 128, lH, lW) -> (B, 32, lH*2, lW*2)
+    3. NCHW -> NHWC transpose.
+
+    Without `bn_mean`/`bn_var`, the helper assumes identity BN (correct shape,
+    slight value offset). Preview structure will still be visible but colors
+    may shift. For exact preview, pass the BN stats from the loaded Flux2VAE.
+
+    Args:
+        packed: in-loop latent from mflux, shape (B, lH*lW, 128).
+        latent_height: latent spatial height (image_height // 16).
+        latent_width: latent spatial width.
+        bn_mean: optional BN running_mean for exact value recovery (128 elements).
+        bn_var: optional BN running_var for exact value recovery (128 elements).
+        bn_eps: BN epsilon. Matches mflux's Flux2BatchNormStats default of 1e-4,
+            verified at mflux 0.17.5.
+
+    Returns:
+        NHWC tensor of shape (B, latent_height*2, latent_width*2, 32) — ready
+        for `TAEF2.decode()`.
+
+    See `notes/mflux-latent-layout.md` for the analysis behind this transform.
+    """
+    b, _, c = packed.shape
+    if c != 128:
+        raise ValueError(f"Expected 128-channel packed latent, got {c}")
+
+    # Step 1: reshape and transpose to NCHW: (B, lH*lW, 128) -> (B, lH, lW, 128) -> (B, 128, lH, lW)
+    latents = packed.reshape(b, latent_height, latent_width, c).transpose(0, 3, 1, 2)
+
+    # Step 2: BN denormalize
+    if bn_mean is not None and bn_var is not None:
+        mean = bn_mean.reshape(1, -1, 1, 1)
+        std = mx.sqrt(bn_var.reshape(1, -1, 1, 1) + bn_eps)
+        latents = latents * std + mean
+    # else: identity BN (mean=0, var=1)
+
+    # Step 3: Unpatchify (B, 128, lH, lW) -> (B, 32, lH*2, lW*2)
+    batch, _, h, w = latents.shape
+    latents = latents.reshape(batch, 32, 2, 2, h, w)
+    latents = latents.transpose(0, 1, 4, 2, 5, 3)
+    latents = latents.reshape(batch, 32, h * 2, w * 2)
+
+    # Step 4: NCHW -> NHWC for TAEF2
+    return latents.transpose(0, 2, 3, 1)
+
+
+class LivePreviewCallback(InLoopCallback):  # type: ignore[misc]
+    """mflux callback that writes a low-quality preview PNG every N iterations.
+
+    Drops into `mflux.Flux2Klein.generate_image(callbacks=[...])`. On the
+    iteration indices that match `every`, unpacks the in-flight latent, runs
+    TAEF2 decode (via `TAEF1` for FLUX.1), and writes a PIL PNG to disk.
+
+    Args:
+        variant: 'taef1' (for FLUX.1 latents) or 'taef2' (for FLUX.2 Klein).
+        every: emit a preview every Nth iteration. Default 5.
+        save_to: filesystem path to write previews. Overwritten each emission.
+        latent_height: latent spatial height; for 512x512 image with FLUX.2 Klein, this is 32.
+        latent_width: latent spatial width.
+        bn_mean: optional BN running_mean for TAEF2 (see `unpack_flux2_latent`).
+        bn_var: optional BN running_var for TAEF2.
+    """
+
+    def __init__(
+        self,
+        *,
+        variant: str = "taef2",
+        every: int = 5,
+        save_to: str | Path = "preview.png",
+        latent_height: int = 32,
+        latent_width: int = 32,
+        bn_mean: mx.array | None = None,
+        bn_var: mx.array | None = None,
+    ) -> None:
+        """Initialise LivePreviewCallback. See class docstring for argument descriptions."""
+        if variant == "taef1":  # pragma: no cover
+            self.model: Taef = TAEF1.from_pretrained(include_encoder=False)
+        elif variant == "taef2":
+            self.model = TAEF2.from_pretrained(include_encoder=False)
+        else:  # pragma: no cover
+            raise ValueError(f"variant must be 'taef1' or 'taef2', got {variant!r}")
+        self.every = every
+        self.save_to = Path(save_to)
+        self.latent_height = latent_height
+        self.latent_width = latent_width
+        self.bn_mean = bn_mean
+        self.bn_var = bn_var
+        self._iter = 0
+
+    def call_in_loop(
+        self,
+        t: object,
+        seed: object,
+        prompt: object,
+        latents: mx.array,
+        config: object,
+        time_steps: object,
+    ) -> None:
+        """Decode latent + save PNG every Nth iteration (counted by our own iter, not `t`)."""
+        idx = self._iter
+        self._iter += 1
+        if idx % self.every != 0:
+            return  # pragma: no cover
+
+        if self.model._config.name == "taef2":
+            unpacked = unpack_flux2_latent(
+                latents,
+                latent_height=self.latent_height,
+                latent_width=self.latent_width,
+                bn_mean=self.bn_mean,
+                bn_var=self.bn_var,
+            )
+        else:  # pragma: no cover
+            # TAEF1 / FLUX.1: latents come in NCHW (B, 16, H, W). Transpose to NHWC.
+            unpacked = mx.transpose(latents, (0, 2, 3, 1)) if latents.shape[1] == 16 else latents
+
+        img = self.model.decode_image(unpacked)
+        self._save_png(img[0])
+
+    def _save_png(self, img_nhwc_uint8: mx.array) -> None:
+        from PIL import Image
+
+        arr = np.array(img_nhwc_uint8)
+        Image.fromarray(arr).save(self.save_to)
+
+
+__all__ = ["LivePreviewCallback", "unpack_flux2_latent"]

mlx_taef/model.py

@@ -0,0 +1,197 @@
+"""MLX layer modules for the TAESD-family of tiny autoencoders.
+
+Layers are direct ports of the upstream PyTorch implementations at
+https://github.com/madebyollin/taesd (MIT licensed).
+
+All public layers accept and return NHWC `mx.array` tensors.
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_taef.variants import TaesdVariantConfig
+
+
+class Clamp(nn.Module):  # type: ignore[misc,name-defined]
+    """Soft clamp via `tanh(x / 3) * 3` — bounds output to [-3, 3].
+
+    Direct port of upstream `taesd.Clamp`. Used at the start of the decoder
+    to bound latent magnitudes before the first conv.
+    """
+
+    def __call__(self, x: mx.array) -> mx.array:
+        """Apply the soft clamp.
+
+        Args:
+            x: input tensor of any shape.
+
+        Returns:
+            Element-wise `tanh(x / 3) * 3` of the same shape.
+        """
+        return mx.tanh(x / 3.0) * 3.0
+
+
+def make_conv(n_in: int, n_out: int, *, stride: int = 1, bias: bool = True) -> "nn.Conv2d":  # type: ignore[name-defined]
+    """3x3 Conv2d with padding=1 — matches upstream `taesd.conv(n_in, n_out, **kwargs)`.
+
+    Args:
+        n_in: input channel count.
+        n_out: output channel count.
+        stride: spatial stride for the convolution (default 1).
+        bias: whether to include a learned bias term (default True).
+
+    Returns:
+        Configured `mlx.nn.Conv2d` with kernel_size=3, padding=1.
+    """
+    return nn.Conv2d(n_in, n_out, kernel_size=3, stride=stride, padding=1, bias=bias)  # type: ignore[attr-defined]
+
+
+class _Identity(nn.Module):  # type: ignore[misc,name-defined]
+    """No-op layer used as the skip connection when in/out channels match."""
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return x
+
+
+class Block(nn.Module):  # type: ignore[misc,name-defined]
+    """Residual block: `ReLU(Conv -> ReLU -> Conv -> ReLU -> Conv + skip)`.
+
+    Direct port of upstream `taesd.Block`. Optionally adds a parallel
+    `midblock_gn` pool branch (used by the flux_2 TAEF2 variant).
+
+    Args:
+        n_in: input channel count.
+        n_out: output channel count.
+        use_midblock_gn: when True, prepend `x = x + pool(x)` where `pool` is
+            a 1x1 conv expand -> GroupNorm -> ReLU -> 1x1 conv contract chain.
+            Used by TAEF2's flux_2 arch_variant. Default False.
+    """
+
+    def __init__(self, n_in: int, n_out: int, *, use_midblock_gn: bool = False) -> None:
+        """Initialise Block layers.
+
+        Args:
+            n_in: input channel count.
+            n_out: output channel count.
+            use_midblock_gn: when True, prepend the midblock GroupNorm pool
+                branch before the main conv path. Default False.
+        """
+        super().__init__()
+        self.conv = nn.Sequential(  # type: ignore[attr-defined]
+            make_conv(n_in, n_out),
+            nn.ReLU(),  # type: ignore[attr-defined]
+            make_conv(n_out, n_out),
+            nn.ReLU(),  # type: ignore[attr-defined]
+            make_conv(n_out, n_out),
+        )
+        self.skip = (
+            nn.Conv2d(n_in, n_out, kernel_size=1, bias=False) if n_in != n_out else _Identity()  # type: ignore[attr-defined]
+        )
+        self.fuse = nn.ReLU()  # type: ignore[attr-defined]
+        self.pool = None
+        if use_midblock_gn:
+            n_gn = n_in * 4
+            # CRITICAL: pytorch_compatible=True is required for parity. MLX
+            # defaults to a different group-channel ordering than PyTorch's
+            # nn.GroupNorm.
+            self.pool = nn.Sequential(  # type: ignore[attr-defined]
+                nn.Conv2d(n_in, n_gn, kernel_size=1, bias=False),  # type: ignore[attr-defined]
+                nn.GroupNorm(num_groups=4, dims=n_gn, pytorch_compatible=True),  # type: ignore[attr-defined]
+                nn.ReLU(),  # type: ignore[attr-defined]
+                nn.Conv2d(n_gn, n_in, kernel_size=1, bias=False),  # type: ignore[attr-defined]
+            )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        """Apply the residual block.
+
+        Args:
+            x: NHWC input tensor with `n_in` channels.
+
+        Returns:
+            NHWC output tensor with `n_out` channels.
+        """
+        if self.pool is not None:
+            x = x + self.pool(x)
+        return self.fuse(self.conv(x) + self.skip(x))  # type: ignore[no-any-return]
+
+
+def make_decoder(config: TaesdVariantConfig) -> nn.Sequential:  # type: ignore[name-defined]
+    """Build the decoder for a TAESD-family variant.
+
+    Mirrors upstream `taesd.Decoder(latent_channels, use_midblock_gn)`:
+    Clamp -> Conv -> ReLU -> Block x3 -> Up -> Conv -> Block x3 -> Up -> Conv ->
+    Block x3 -> Up -> Conv -> Block -> Conv.
+
+    For arch_variant="flux_2" (TAEF2), the first three Blocks have
+    `use_midblock_gn=True`; the rest are vanilla.
+
+    Args:
+        config: variant configuration that selects `latent_channels` and
+            whether the first three Blocks use midblock_gn.
+
+    Returns:
+        `nn.Sequential` decoder ready for weight loading.
+    """
+    mb = config.use_midblock_gn
+    layers: list[nn.Module] = [  # type: ignore[name-defined]
+        Clamp(),
+        make_conv(config.latent_channels, 64),
+        nn.ReLU(),  # type: ignore[attr-defined]
+        Block(64, 64, use_midblock_gn=mb),
+        Block(64, 64, use_midblock_gn=mb),
+        Block(64, 64, use_midblock_gn=mb),
+        nn.Upsample(scale_factor=2, mode="nearest"),  # type: ignore[attr-defined]
+        make_conv(64, 64, bias=False),
+        Block(64, 64),
+        Block(64, 64),
+        Block(64, 64),
+        nn.Upsample(scale_factor=2, mode="nearest"),  # type: ignore[attr-defined]
+        make_conv(64, 64, bias=False),
+        Block(64, 64),
+        Block(64, 64),
+        Block(64, 64),
+        nn.Upsample(scale_factor=2, mode="nearest"),  # type: ignore[attr-defined]
+        make_conv(64, 64, bias=False),
+        Block(64, 64),
+        make_conv(64, 3),
+    ]
+    return nn.Sequential(*layers)  # type: ignore[attr-defined]
+
+
+def make_encoder(config: TaesdVariantConfig) -> nn.Sequential:  # type: ignore[name-defined]
+    """Build the encoder for a TAESD-family variant.
+
+    Mirrors upstream `taesd.Encoder(latent_channels, use_midblock_gn)`:
+    Conv -> Block -> (StridedConv -> Block x3) x3 -> Conv.
+    Three strided convs each halve spatial dims (8x downsample total).
+
+    For arch_variant="flux_2" (TAEF2), the LAST three Blocks use `use_midblock_gn=True`.
+
+    Args:
+        config: variant configuration that selects `latent_channels` and `arch_variant`.
+
+    Returns:
+        `nn.Sequential` encoder ready for weight loading.
+    """
+    mb = config.use_midblock_gn
+    layers: list[nn.Module] = [  # type: ignore[name-defined]
+        make_conv(3, 64),
+        Block(64, 64),
+        make_conv(64, 64, stride=2, bias=False),
+        Block(64, 64),
+        Block(64, 64),
+        Block(64, 64),
+        make_conv(64, 64, stride=2, bias=False),
+        Block(64, 64),
+        Block(64, 64),
+        Block(64, 64),
+        make_conv(64, 64, stride=2, bias=False),
+        Block(64, 64, use_midblock_gn=mb),
+        Block(64, 64, use_midblock_gn=mb),
+        Block(64, 64, use_midblock_gn=mb),
+        make_conv(64, config.latent_channels),
+    ]
+    return nn.Sequential(*layers)  # type: ignore[attr-defined]
+
+
+__all__ = ["Block", "Clamp", "make_conv", "make_decoder", "make_encoder"]

mlx_taef/variants.py

@@ -0,0 +1,115 @@
+"""Variant configurations for the TAESD-family of tiny autoencoders.
+
+Field reference:
+- `key_format`: "upstream" means the HF safetensors uses upstream Sequential keys
+  like "0.weight", "1.weight"; "diffusers" means the HF safetensors uses Diffusers
+  VAE keys like "encoder.layers.0.weight" and requires a +1 decoder index offset.
+- `arch_variant`: None | "flux_2" | "f32". "flux_2" enables midblock_gn pool
+  branches in three blocks of encoder and decoder. "f32" is reserved for
+  TAESANA (future).
+- `latent_magnitude`, `latent_shift`: from upstream TAESD.scale_latents /
+  unscale_latents.
+- `hf_filename` (Diffusers single-file format) vs `hf_decoder_filename` /
+  `hf_encoder_filename` (upstream two-file format) — depends on the actual
+  repo layout for each variant.
+"""
+
+import logging
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class TaesdVariantConfig:
+    """Configuration for a single TAESD-family variant.
+
+    Attributes:
+        name: short identifier like "taef2".
+        latent_channels: latent channel count (4 for SD1/SDXL, 16 for FLUX.1, 32 for FLUX.2 Klein).
+        arch_variant: None for the standard arch, "flux_2" for TAEF2's midblock_gn, "f32" for TAESANA.
+        key_format: "upstream" (Sequential keys) or "diffusers" (requires remap).
+        hf_repo: e.g. "madebyollin/taef2".
+        hf_filename: filename within hf_repo when key_format == "diffusers" (single file).
+        hf_decoder_filename: decoder filename when key_format == "upstream".
+        hf_encoder_filename: encoder filename when key_format == "upstream".
+        latent_magnitude: scale factor for scale_latents/unscale_latents.
+        latent_shift: shift factor for scale_latents/unscale_latents.
+    """
+
+    name: str
+    latent_channels: int
+    arch_variant: str | None
+    key_format: str
+    hf_repo: str
+    hf_filename: str | None
+    hf_decoder_filename: str | None
+    hf_encoder_filename: str | None
+    latent_magnitude: float = 3.0
+    latent_shift: float = 0.5
+
+    @property
+    def use_midblock_gn(self) -> bool:
+        """Whether the variant's Block layers use the midblock GroupNorm pool branch."""
+        return self.arch_variant == "flux_2"
+
+
+TAESD_CONFIG = TaesdVariantConfig(
+    name="taesd",
+    latent_channels=4,
+    arch_variant=None,
+    key_format="upstream",
+    hf_repo="madebyollin/taesd",
+    hf_filename=None,
+    hf_decoder_filename="taesd_decoder.safetensors",
+    hf_encoder_filename="taesd_encoder.safetensors",
+)
+
+TAESDXL_CONFIG = TaesdVariantConfig(
+    name="taesdxl",
+    latent_channels=4,
+    arch_variant=None,
+    key_format="upstream",
+    hf_repo="madebyollin/taesdxl",
+    hf_filename=None,
+    hf_decoder_filename="taesdxl_decoder.safetensors",
+    hf_encoder_filename="taesdxl_encoder.safetensors",
+)
+
+TAEF1_CONFIG = TaesdVariantConfig(
+    name="taef1",
+    latent_channels=16,
+    arch_variant=None,
+    key_format="diffusers",
+    hf_repo="madebyollin/taef1",
+    hf_filename="diffusion_pytorch_model.safetensors",
+    hf_decoder_filename=None,
+    hf_encoder_filename=None,
+)
+
+TAEF2_CONFIG = TaesdVariantConfig(
+    name="taef2",
+    latent_channels=32,
+    arch_variant="flux_2",
+    key_format="diffusers",
+    hf_repo="madebyollin/taef2",
+    hf_filename="taef2.safetensors",
+    hf_decoder_filename=None,
+    hf_encoder_filename=None,
+)
+
+ALL_VARIANTS: tuple[TaesdVariantConfig, ...] = (
+    TAESD_CONFIG,
+    TAESDXL_CONFIG,
+    TAEF1_CONFIG,
+    TAEF2_CONFIG,
+)
+
+__all__ = [
+    "ALL_VARIANTS",
+    "TAEF1_CONFIG",
+    "TAEF2_CONFIG",
+    "TAESDXL_CONFIG",
+    "TAESD_CONFIG",
+    "TaesdVariantConfig",
+]

mlx_taef-0.1.0.dist-info/METADATA

@@ -0,0 +1,125 @@
+Metadata-Version: 2.4
+Name: mlx-taef
+Version: 0.1.0
+Summary: Tiny AutoEncoders for diffusion (TAESD family) on Apple MLX.
+Project-URL: Homepage, https://github.com/ionden/mlx-taef
+Project-URL: Documentation, https://ionden.github.io/mlx-taef
+Project-URL: Issues, https://github.com/ionden/mlx-taef/issues
+Project-URL: Changelog, https://github.com/ionden/mlx-taef/blob/main/CHANGELOG.md
+Author-email: Denis Ineshin <denis.ineshin@gmail.com>
+License-Expression: MIT
+Keywords: apple-silicon,autoencoder,flux,mlx,stable-diffusion,taesd
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: huggingface-hub>=0.24
+Requires-Dist: mlx>=0.20
+Requires-Dist: numpy>=1.26
+Requires-Dist: safetensors>=0.4
+Provides-Extra: image
+Requires-Dist: pillow>=10.0; extra == 'image'
+Provides-Extra: mflux
+Requires-Dist: mflux>=0.17; extra == 'mflux'
+Description-Content-Type: text/markdown
+
+# mlx-taef
+
+Tiny AutoEncoders for diffusion latents on Apple Silicon, in pure MLX.
+
+`mlx-taef` is the first MLX port of the TAESD family — TAESD (SD1.x), TAESDXL (SDXL), TAEF1 (FLUX.1), TAEF2 (FLUX.2 Klein) — distilled mini-autoencoders that decode diffusion latents to RGB in milliseconds using a few-MB model instead of multi-GB full VAEs.
+
+Use it for:
+- **Live previews** during long generations on Mac — see each step refresh in <100 ms instead of waiting 30 s for the full VAE.
+- **Low-memory fallbacks** when the full VAE OOMs on 16 GB Macs (TAEF2 peaks at ~1 GB for 1024×1024 vs ~9.6 GB for the full Flux VAE).
+- **Quick latent inspection** in notebooks and ML research.
+
+```python
+import mlx.core as mx
+from mlx_taef import TAEF2
+
+taef = TAEF2.from_pretrained()              # downloads + converts on first call
+img = taef.decode(latents)                  # NHWC float in [0, 1]
+img_uint8 = taef.decode_image(latents)      # uint8 NHWC ready for PIL
+```
+
+## Install
+
+```bash
+pip install mlx-taef
+# With mflux preview callback:
+pip install "mlx-taef[mflux]"
+```
+
+Requires Python ≥ 3.11 and Apple Silicon. Runtime install has **zero PyTorch dependency**.
+
+## Variants
+
+| Variant | latent_channels | For | HF source |
+|---|---|---|---|
+| `TAESD` | 4 | Stable Diffusion 1.x | [madebyollin/taesd](https://huggingface.co/madebyollin/taesd) |
+| `TAESDXL` | 4 | Stable Diffusion XL | [madebyollin/taesdxl](https://huggingface.co/madebyollin/taesdxl) |
+| `TAEF1` | 16 | FLUX.1 | [madebyollin/taef1](https://huggingface.co/madebyollin/taef1) |
+| `TAEF2` | 32 | FLUX.2 Klein | [madebyollin/taef2](https://huggingface.co/madebyollin/taef2) |
+
+All four share one API.
+
+## Benchmarks (M1 Max, fp16)
+
+| Metric | TAEF2 (this library) | Full Flux VAE (reference) | Win |
+|---|---|---|---|
+| Decode latency 1024×1024 | **~100 ms** | seconds | 50–100× |
+| Peak unified memory 1024×1024 | **~1 GB** | ~9.6 GB | **9.4×** |
+| Output cosine sim vs PyTorch reference | > 0.999 | — | (parity verified) |
+
+Numbers from `tests/test_perf.py` on M1 Max 32 GB. See `notes/phase1-benchmarks.md` for details.
+
+## mflux live previews
+
+```python
+from mflux.models.flux2 import Flux2Klein
+from mlx_taef.integrations.mflux import LivePreviewCallback
+
+model = Flux2Klein.from_pretrained("4bit")
+preview = LivePreviewCallback(
+    variant="taef2",
+    every=5,
+    save_to="preview.png",
+    latent_height=32,  # 512 / 16
+    latent_width=32,
+)
+model.callbacks.register(preview)
+model.generate_image(
+    prompt="a red apple on a wooden table",
+    num_inference_steps=25,
+    width=512,
+    height=512,
+    seed=42,
+)
+```
+
+For exact value-space recovery, also pass `bn_mean=flux2_vae.bn.running_mean, bn_var=flux2_vae.bn.running_var` to the callback. Without them, previews show correct structure but colors may shift.
+
+See `docs/manual-verification.md` for the full verification recipe.
+
+## Status
+
+- v0.1.0 — initial public release. All four variants, encoder + decoder, mflux integration, CI, 100% test coverage.
+
+## License
+
+MIT. Mirrors upstream [madebyollin/taesd](https://github.com/madebyollin/taesd) license. Pretrained weights belong to their respective authors (madebyollin).
+
+## Acknowledgements
+
+- [madebyollin](https://github.com/madebyollin) for the upstream TAESD-family models and weights.
+- [Apple ML Explore](https://github.com/ml-explore/mlx) for MLX.
+- [filipstrand/mflux](https://github.com/filipstrand/mflux) for the MLX-native FLUX runner this library integrates with.

mlx_taef-0.1.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+mlx_taef/__init__.py,sha256=hn12r3kNI6DnIKnzNpx6EjU6zfYXpHT034qtk3OQ8ic,261
+mlx_taef/_version.py,sha256=n_5vdJsPNu7wZ57LGuRL585uvll-hiuvZUBWzdG0RQU,520
+mlx_taef/api.py,sha256=_Si8Mi0uVkjbr5n9Ni_kxES-CVhyW3M0nNGedwmtaQo,6502
+mlx_taef/cli.py,sha256=-8nvCxga9tf4YWU6CJADuv0ch7qi2fa9MkwneKJrg0Q,3253
+mlx_taef/convert.py,sha256=5mLO-nfat6HuUpUoQFbIVu2Vi6wbEW7qlGetybL7syo,7002
+mlx_taef/download.py,sha256=QDrOz54aRtByHj_x4idkcAV3k1X-LjgNsFwl3f-OSQw,1608
+mlx_taef/model.py,sha256=FtI8PeDUQi_n01YGvOz_hlJoNkirprQPawPZ2lnSBQM,7340
+mlx_taef/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mlx_taef/variants.py,sha256=jAsiU0C2jG38fwazE8XfKf-6gVkh1A0AaBz6n8jnoNI,3655
+mlx_taef/integrations/__init__.py,sha256=XS1-Ml6Mj8iPVF9E3CvFbK18mg-o6dBDiYvg4-MIo08,55
+mlx_taef/integrations/mflux.py,sha256=PyDtgZp1Bgsb3ZYzKYQWg3DvJz_MUxnSsSXaWv2hs2k,6214
+mlx_taef-0.1.0.dist-info/METADATA,sha256=cNI3F-seg2WdGCcdvp5VVRkrhH1AyStHSN3RppctkPA,4883
+mlx_taef-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+mlx_taef-0.1.0.dist-info/entry_points.txt,sha256=x8owUzlm7qAoK-CsSc0ScgAnTzu0Ipkn9evR_kTso1I,47
+mlx_taef-0.1.0.dist-info/RECORD,,

mlx_taef-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

mlx_taef-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+mlx-taef = mlx_taef.cli:main