mlx-dfloat 0.1.0__py3-none-any.whl

mlx_dfloat/mflux/flux1/lifecycle.py

@@ -0,0 +1,159 @@
+"""The encoder / compressed-set lifecycle: the text encoders and the compressed transformer are never resident together.
+
+On a 32 GB Mac the T5 encoder (8.9 GiB) and the compressed FLUX.1 set (15.2 GiB) do not fit under the
+fit rule, so a prompt is encoded first, its embeddings evaluated, the encoders dropped, and only then
+the set loaded. A new prompt after the set is resident drops the set, reloads the encoders, encodes,
+drops them and reloads the set. The object here holds no arrays and no modules: the callbacks own
+them, so a drop releases them.
+"""
+
+import gc
+import logging
+import time
+import traceback
+from collections.abc import Callable
+from dataclasses import asdict, dataclass
+
+import mlx.core as mx
+
+from mlx_dfloat.errors import DFloatResourceError
+
+log = logging.getLogger("mlx_dfloat.mflux.flux1")
+_eval = mx.eval  # looked up on the module, so a test can record the evaluation
+
+
+@dataclass(slots=True)
+class LifecycleCounters:
+    """Reloads and their seconds, for the report."""
+
+    encoder_loads: int = 0
+    encoder_seconds: float = 0.0
+    set_loads: int = 0
+    set_seconds: float = 0.0
+    forced_set_drops: int = 0
+
+    def as_dict(self) -> dict[str, float | int]:
+        """The counters as a JSON-ready dict."""
+        return asdict(self)
+
+
+class Lifecycle:
+    """Drives the encoder and set callbacks so that the two are never resident at once."""
+
+    def __init__(
+        self,
+        *,
+        load_encoders: Callable[[], None],
+        unload_encoders: Callable[[], None],
+        encode: Callable[[str], tuple[mx.array, mx.array]],
+        load_set: Callable[[], None],
+        unload_set: Callable[[], None],
+        prompt_cache: dict[str, tuple[mx.array, mx.array]],
+        retained_bound: Callable[[], int],
+        encoders_loaded: bool = False,
+    ) -> None:
+        """Bind the callbacks, the prompt cache mflux reads, and the active-memory bound a drop must reach.
+
+        ``retained_bound`` is called at each check, so what legitimately stays resident (cached
+        embeddings, a decoded VAE) can grow between drops.
+        """
+        self._load_encoders = load_encoders
+        self._unload_encoders = unload_encoders
+        self._encode = encode
+        self._load_set = load_set
+        self._unload_set = unload_set
+        self._cache = prompt_cache
+        self.retained_bound = retained_bound
+        self.encoders_loaded = encoders_loaded
+        self.set_resident = False
+        self.counters = LifecycleCounters()
+
+    def ensure_embeddings(self, *prompts: str) -> None:
+        """Encode the prompts not yet cached: the set is dropped first, the encoders after.
+
+        Each pair is evaluated before it is cached: a lazy pair would keep every encoder weight
+        alive past the drop. A failure in the encode or the evaluation still drops the encoders
+        (instead of leaving them resident next to a set loaded later) and re-raises the original
+        error; a memory check that fails on that path is logged, never raised over it.
+        """
+        missing = [p for p in prompts if p not in self._cache]
+        if not missing:
+            return
+        if self.set_resident:
+            log.info(
+                "new prompt with the compressed set resident: dropping the set to load the text encoders"
+            )
+            self.counters.forced_set_drops += 1
+            self.drop_set()
+        if not self.encoders_loaded:
+            start = time.perf_counter()
+            self._load_encoders()
+            self.encoders_loaded = True
+            self.counters.encoder_loads += 1
+            self.counters.encoder_seconds += time.perf_counter() - start
+        pair: tuple[mx.array, mx.array] | None = None
+        try:
+            for prompt in missing:
+                pair = self._encode(prompt)
+                _eval(*pair)
+                self._cache[prompt] = pair
+        except BaseException as exc:
+            # The traceback keeps the failing frames alive, and with them a lazy pair whose graph
+            # holds the encoder weights; clear them before the memory check.
+            pair = None
+            traceback.clear_frames(exc.__traceback__)
+            self._unload_encoders()
+            self.encoders_loaded = False
+            try:
+                self._reclaim("text encoders")
+            except DFloatResourceError as reclaim_error:
+                log.warning("%s (while handling %s: %s)", reclaim_error, type(exc).__name__, exc)
+            raise
+        del pair
+        self.drop_encoders()
+
+    def ensure_set(self) -> None:
+        """Load and attach the compressed set unless it is resident.
+
+        Raises:
+            DFloatResourceError: MLX memory above the retained bound is still active before the
+                load (a previous set not released), so a second copy is never loaded next to it.
+        """
+        if self.set_resident:
+            return
+        gc.collect()
+        active = int(mx.get_active_memory())
+        bound = self.retained_bound()
+        if active > bound:
+            raise DFloatResourceError(
+                f"a previous set is still resident: {active / 1024**3:.2f} GiB of MLX memory is active "
+                f"before the load (bound {bound / 1024**3:.2f} GiB)"
+            )
+        start = time.perf_counter()
+        self._load_set()
+        self.set_resident = True
+        self.counters.set_loads += 1
+        self.counters.set_seconds += time.perf_counter() - start
+
+    def drop_set(self) -> None:
+        """Detach and release the compressed set; assert the memory came back."""
+        self._unload_set()
+        self.set_resident = False
+        self._reclaim("compressed set")
+
+    def drop_encoders(self) -> None:
+        """Release the text encoders; assert the memory came back."""
+        self._unload_encoders()
+        self.encoders_loaded = False
+        self._reclaim("text encoders")
+
+    def _reclaim(self, what: str) -> None:
+        gc.collect()
+        mx.clear_cache()
+        active = int(mx.get_active_memory())
+        bound = self.retained_bound()
+        if active > bound:
+            raise DFloatResourceError(
+                f"after dropping the {what}, {active / 1024**3:.2f} GiB of MLX memory is still active "
+                f"(bound {bound / 1024**3:.2f} GiB): something else holds it"
+            )

mlx_dfloat/mflux/flux1/memory.py

@@ -0,0 +1,201 @@
+"""FLUX.1 memory rules: the cache limit and the fit phases. Measured at schnell 1024² only; elsewhere predicted."""
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from mlx_dfloat.format import DF11Checkpoint
+from mlx_dfloat.integrate.memory import FitEstimate, fit_estimate
+from mlx_dfloat.mflux.flux1.names import DOUBLE_PREFIX, SINGLE_PREFIX
+
+ALLOWANCE_AT_REFERENCE = (
+    1_500_000_000  # recycled activation volume measured at 1024², 256 text tokens
+)
+REFERENCE_TOKENS = 4096 + 256
+ALLOWANCE_FLOOR = 500_000_000
+MEASURED_LIMIT_AT_1024 = 2_500_000_000
+
+
+def text_tokens(model_config: Any) -> int:
+    """The model's text sequence length (256 for schnell, 512 for dev and Krea-dev)."""
+    return int(model_config.max_sequence_length)
+
+
+def activation_allowance(*, height: int, width: int, text_tokens: int) -> int:
+    """Cache room for the activation buffers a block frees, scaled linearly with the token count."""
+    tokens = height * width // 256 + text_tokens
+    return max(ALLOWANCE_FLOOR, int(ALLOWANCE_AT_REFERENCE * tokens / REFERENCE_TOKENS))
+
+
+def cache_limit_for(
+    largest: Mapping[str, int],
+    *,
+    policy: str,
+    height: int,
+    width: int,
+    text_tokens: int,
+    override: int | None = None,
+) -> int:
+    """The MLX buffer-cache limit for one generate call: room for a decoded buffer of each kind next to the activations."""
+    if override is not None:
+        return override
+    limit = (
+        largest[DOUBLE_PREFIX]
+        + largest[SINGLE_PREFIX]
+        + activation_allowance(height=height, width=width, text_tokens=text_tokens)
+    )
+    if policy == "depth2":
+        limit += max(largest.values())
+    if height * width >= 1024 * 1024:
+        limit = max(limit, MEASURED_LIMIT_AT_1024)
+    return limit
+
+
+def flux_phases(
+    *,
+    compressed_bytes: int,
+    extras_bytes: int,
+    largest: Mapping[str, int],
+    policy: str,
+    cache_limit: int,
+    allowance: int,
+    encoders_bytes: int,
+    vae_bytes: int,
+    vae_transient_bytes: int,
+    overhead_bytes: int,
+    denoise_activation_bytes: int = 0,
+) -> dict[str, dict[str, int]]:
+    """The three phases of a generation, each term once."""
+    in_flight = max(largest.values()) * (2 if policy == "depth2" else 1)
+    return {
+        "encode": {
+            "encoders": encoders_bytes,
+            "activations": allowance,
+            "overhead": overhead_bytes,
+        },
+        "denoise": {
+            "compressed": compressed_bytes,
+            "extras": extras_bytes,
+            "decoded": in_flight,
+            "cache": cache_limit,
+            "activations": denoise_activation_bytes,
+            "overhead": overhead_bytes,
+        },
+        "vae": {
+            "compressed": compressed_bytes,
+            "extras": extras_bytes,
+            "vae": vae_bytes,
+            "transient": vae_transient_bytes,
+            "overhead": overhead_bytes,
+        },
+    }
+
+
+# Measured on 2026-09-28 (schnell 1024², one process: encode, drop, set, one step, VAE decode; M1 Max 32 GB,
+# macOS 27.0, mlx 0.32.2, mflux 0.20.0, git 4601a8b). OVERHEAD is the footprint that MLX's counters do not see
+# after construction (Python, torch and transformers imports, the runtime); VAE_TRANSIENT is what the float32
+# decode adds at 1024² on top of everything resident after the set load (the VAE phase peaked at 23.29 GiB with
+# the set resident, over the 23.0 GiB fit rule, which is why a call may drop the set before decoding).
+# Re-measure before changing either.
+OVERHEAD_BYTES = 247_712_510  # 0.23 GiB: build-phase footprint minus MLX active and cache
+VAE_TRANSIENT_BYTES = (
+    8_392_982_528  # 7.82 GiB: VAE-phase footprint peak minus the footprint after the set load
+)
+# Measured on 2026-09-28 (the `mlx-dfloat generate` runs, schnell 1024², 4 steps, per-block, 2.5 GB cache limit;
+# M1 Max 32 GB, mlx 0.32.2, mflux 0.20.0, git 2fe3b2c): the process footprint peaked at 19.94 GiB against a
+# denoise estimate of 18.38 GiB without this term. The 1.56 GiB gap is the activation volume the step holds
+# beyond the cache limit, at 4096 image + 256 text tokens; dev (512 text tokens) measured 20.10 GiB.
+DENOISE_ACTIVATION_AT_REFERENCE = 1_675_000_000
+MAX_MEASURED_PIXELS = 1024 * 1024  # no run above 1024² on this path
+
+
+def vae_transient_bytes(*, height: int, width: int) -> int:
+    """What the float32 VAE decode adds on top of what is resident: the measured 1024² value as a floor.
+
+    At or below 1024² this is the measured value (smaller images were not measured lower, so the
+    floor stays). Above 1024² it grows with the pixel count: a prediction, not a measurement.
+    """
+    return int(VAE_TRANSIENT_BYTES * max(1.0, (height * width) / (1024 * 1024)))
+
+
+def denoise_activation_bytes(*, height: int, width: int, text_tokens: int) -> int:
+    """The activation volume a denoise step holds beyond the cache limit, linear in the token count.
+
+    Calibrated on the measured schnell 1024² peak (4096 image + 256 text tokens); every other size
+    is a prediction.
+    """
+    tokens = height * width // 256 + text_tokens
+    return int(DENOISE_ACTIVATION_AT_REFERENCE * tokens / REFERENCE_TOKENS)
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class FluxSizes:
+    """Resident bytes of the components, from file sizes (exact) and the extras' tensor sizes."""
+
+    compressed: int
+    extras: int
+    encoders: int
+    vae: int
+
+
+def safetensors_bytes(root: Path, *subdirs: str) -> int:
+    """The size of every ``*.safetensors`` file directly under each ``root/subdir`` (missing subdirs count zero)."""
+    return sum(
+        p.stat().st_size
+        for sub in subdirs
+        for p in (root / sub).glob("*.safetensors")
+        if p.is_file()
+    )
+
+
+def sizes_for(ckpt: DF11Checkpoint, base_root: Path) -> FluxSizes:
+    """Sizes of a DF11 checkpoint (compressed set apart from its extras) and of a base's encoders and VAE."""
+    total = sum(p.stat().st_size for p in ckpt.root.glob("*.safetensors") if p.is_file())
+    extras = sum(info.nbytes for _path, info in ckpt.extras.values())
+    return FluxSizes(
+        compressed=total - extras,
+        extras=extras,
+        encoders=safetensors_bytes(base_root, "text_encoder", "text_encoder_2"),
+        vae=safetensors_bytes(base_root, "vae"),
+    )
+
+
+def fit_for(
+    *,
+    sizes: FluxSizes,
+    largest: Mapping[str, int],
+    policy: str,
+    cache_limit: int,
+    allowance: int,
+    budget: int,
+    height: int,
+    width: int,
+    text_tokens: int,
+    vae_with_set: bool = True,
+) -> FitEstimate:
+    """The phase estimate for one generate call against ``budget`` (a prediction, labelled as such by the caller).
+
+    ``height`` and ``width`` scale the VAE transient (above 1024²) and, with ``text_tokens``, the
+    denoise activation term. ``vae_with_set=False`` plans the VAE phase after the compressed set has
+    been dropped (what a call does when the resident variant would not fit).
+    """
+    phases = flux_phases(
+        compressed_bytes=sizes.compressed,
+        extras_bytes=sizes.extras,
+        largest=largest,
+        policy=policy,
+        cache_limit=cache_limit,
+        allowance=allowance,
+        encoders_bytes=sizes.encoders,
+        vae_bytes=sizes.vae,
+        vae_transient_bytes=vae_transient_bytes(height=height, width=width),
+        overhead_bytes=OVERHEAD_BYTES,
+        denoise_activation_bytes=denoise_activation_bytes(
+            height=height, width=width, text_tokens=text_tokens
+        ),
+    )
+    if not vae_with_set:
+        phases["vae"].pop("compressed")
+        phases["vae"].pop("extras")
+    return fit_estimate(phases, budget_bytes=budget)