From 5f98825d72920d2416cf861e8dd6665f8cea1876 Mon Sep 17 00:00:00 2001 From: Denis Ineshin Date: Wed, 28 Oct 2026 19:31:45 +0100 Subject: [PATCH 1/4] Run on mflux 0.23 and pin each family's base repository Raise the mflux extra and the bench group to mflux>=0.23,<0.24 with pillow<13 (mflux 0.23.0 fails at import with Pillow 13). Load FLUX.1's norm_out.linear.bias, which mflux now builds and loads. Run Qwen-Image 2.1 through mflux's packed projections one block at a time, keep its compiled step functions out of the generation path, and refuse a step function compiled between calls. Export the empty COMPUTED_PARAMS set from the Qwen-Image 2.1 transformer, as the Krea 2 and ERNIE transformers do. Set the generate-time attributes mflux 0.23 reads for Z-Image and FLUX.2 Klein and refuse its precision options. Name the missing tensors when a base checkpoint lacks required weights. Pin the base repository revision for FLUX.1, Z-Image and FLUX.2 Klein. Record the macOS version in generate reports. --- pyproject.toml | 6 +- src/mlx_dfloat/mflux/_compile.py | 2 +- src/mlx_dfloat/mflux/_hub.py | 33 +- src/mlx_dfloat/mflux/_pipeline.py | 18 ++ src/mlx_dfloat/mflux/ernie/model.py | 2 +- src/mlx_dfloat/mflux/ernie/names.py | 4 +- src/mlx_dfloat/mflux/ernie/transformer.py | 4 +- src/mlx_dfloat/mflux/families.py | 36 ++- src/mlx_dfloat/mflux/flux1/model.py | 7 +- src/mlx_dfloat/mflux/flux1/names.py | 5 +- src/mlx_dfloat/mflux/flux1/transformer.py | 2 - src/mlx_dfloat/mflux/flux2/init.py | 2 +- src/mlx_dfloat/mflux/flux2/model.py | 20 +- src/mlx_dfloat/mflux/flux2/names.py | 2 +- src/mlx_dfloat/mflux/flux2/transformer.py | 2 +- src/mlx_dfloat/mflux/generate.py | 24 +- src/mlx_dfloat/mflux/krea2/init.py | 2 +- src/mlx_dfloat/mflux/krea2/model.py | 2 +- src/mlx_dfloat/mflux/krea2/names.py | 4 +- src/mlx_dfloat/mflux/krea2/transformer.py | 5 +- src/mlx_dfloat/mflux/qwen21/init.py | 4 +- src/mlx_dfloat/mflux/qwen21/memory.py | 24 +- src/mlx_dfloat/mflux/qwen21/model.py | 27 +- src/mlx_dfloat/mflux/qwen21/transformer.py | 332 ++++++++++++++++++--- src/mlx_dfloat/mflux/zimage/memory.py | 25 +- src/mlx_dfloat/mflux/zimage/model.py | 21 +- src/mlx_dfloat/mflux/zimage/names.py | 2 +- tests/_qwen21_tiny.py | 70 ++++- tests/test_cli.py | 19 ++ tests/test_flux1_model.py | 39 +++ tests/test_flux1_names.py | 22 +- tests/test_flux1_transformer.py | 45 ++- tests/test_flux2_model.py | 44 +++ tests/test_krea2_transformer.py | 2 +- tests/test_mflux_families.py | 101 ++++++- tests/test_mflux_generate.py | 87 ++++++ tests/test_mflux_hub.py | 146 +++++++++ tests/test_mflux_pins.py | 34 +++ tests/test_mflux_pipeline.py | 21 ++ tests/test_qwen21_build.py | 26 +- tests/test_qwen21_init.py | 15 +- tests/test_qwen21_memory.py | 64 ++-- tests/test_qwen21_mflux023.py | 248 +++++++++++++++ tests/test_qwen21_model.py | 50 +++- tests/test_qwen21_packing.py | 216 ++++++++++++++ tests/test_qwen21_tiny.py | 5 +- tests/test_qwen21_tiny_compile.py | 57 ++++ tests/test_qwen21_transformer.py | 61 ++-- tests/test_verify_image.py | 98 ++++++ tests/test_zimage_memory.py | 117 +++++--- tests/test_zimage_model.py | 58 +++- uv.lock | 68 ++++- 52 files changed, 2049 insertions(+), 281 deletions(-) create mode 100644 tests/test_mflux_generate.py create mode 100644 tests/test_mflux_hub.py create mode 100644 tests/test_mflux_pins.py create mode 100644 tests/test_qwen21_mflux023.py create mode 100644 tests/test_qwen21_packing.py create mode 100644 tests/test_qwen21_tiny_compile.py diff --git a/pyproject.toml b/pyproject.toml index 7c1f51a..fb042b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,8 @@ mlx-dfloat = "mlx_dfloat.cli:main" [project.optional-dependencies] mflux = [ - "mflux>=0.20,<0.21; sys_platform == 'darwin'", + "mflux>=0.23,<0.24; sys_platform == 'darwin'", + "pillow<13", "psutil>=6", ] @@ -63,7 +64,8 @@ test = [ typecheck = ["mypy>=1.13"] lint = ["ruff>=0.7"] bench = [ - "mflux>=0.20,<0.21; sys_platform == 'darwin'", + "mflux>=0.23,<0.24; sys_platform == 'darwin'", + "pillow<13", "pyobjc-framework-Metal>=10; sys_platform == 'darwin'", ] dev = [ diff --git a/src/mlx_dfloat/mflux/_compile.py b/src/mlx_dfloat/mflux/_compile.py index ab6e20f..d6bed33 100644 --- a/src/mlx_dfloat/mflux/_compile.py +++ b/src/mlx_dfloat/mflux/_compile.py @@ -1,6 +1,6 @@ """mflux's compiled ``predict``, turned off for the duration of one factory call. -mflux 0.20's ``_predict`` factories (Z-Image, FLUX.2 Klein, ERNIE-Image, Krea 2) return ``mx.compile(predict)`` +mflux 0.23's ``_predict`` factories (Z-Image, FLUX.2 Klein, ERNIE-Image, Krea 2) return ``mx.compile(predict)`` unless ``AppleSiliconUtil.is_m1_or_m2()``, which they read at call time. The block seam evaluates at block boundaries, which cannot happen inside ``mx.compile``, so the adapters take the plain function on every chip. The override is process-wide but lasts only while the factory runs (it builds a closure; nothing else executes). A diff --git a/src/mlx_dfloat/mflux/_hub.py b/src/mlx_dfloat/mflux/_hub.py index f41f287..91529a4 100644 --- a/src/mlx_dfloat/mflux/_hub.py +++ b/src/mlx_dfloat/mflux/_hub.py @@ -4,6 +4,7 @@ adapters load the other components through mflux's own loader and applier with weight definitions that name only them. """ +import re import sys from dataclasses import dataclass from pathlib import Path @@ -12,6 +13,12 @@ from mlx_dfloat.errors import DFloatAccessError, DFloatFormatError, DFloatIntegrationError _SHA_LENGTH = 40 +# mflux 0.23's required-weights refusal (WeightLoader._describe_missing_weights, weight_loader.py:500-505): the +# component, the total count and at most three of the missing checkpoint names. +_MISSING_WEIGHTS = re.compile( + r"The (?P\S+) weights in .*? do not fit this model: (?P\d+) required " + r"weights? (?:has|have) no match \(expected names like (?P[^;]+); found " +) @dataclass(frozen=True, slots=True, kw_only=True) @@ -141,15 +148,31 @@ def subset_definition(source: type, names: tuple[str, ...], patterns: tuple[str, ) +def missing_weights(message: str) -> str | None: + """``" [ and N more]"`` from mflux's required-weights refusal; None for any other message.""" + match = _MISSING_WEIGHTS.search(message) + if match is None: + return None + shown = match["names"].split(", ") + more = int(match["count"]) - len(shown) + return f"{match['component']} {', '.join(shown)}" + (f" and {more} more" if more > 0 else "") + + def load_into(root: Path, definition: Any, models: dict[str, Any], *, what: str) -> None: """Mflux's loader and applier over ``definition`` into ``models``; the weights stay lazy until evaluated. The ``LoadedWeights`` object is not kept: it references every lazily loaded array, and the encoders must be droppable. + ``root`` is a local directory (``resolve`` already fetched it at its pin), so mflux takes it as given and its own + pinned revisions never apply. ``definition`` needs only ``get_components``, ``get_download_patterns`` and + ``quantization_predicate`` (``get_tokenizers`` for the tokenizer loader); mflux 0.23's loader and applier read + nothing else. + Raises: - DFloatFormatError: The directory lacks the components (a checkpoint or model repository without them), or - holds an mflux-saved quantized model. + DFloatFormatError: The directory lacks the components (a checkpoint or model repository without them; mflux + reports a component missing its required weights as ``ModelConfigError``, a ``ValueError``; that + refusal is reported with the missing tensors' names first), or holds an mflux-saved quantized model. DFloatIntegrationError: The applier quantized the components unasked. """ from mflux.models.common.weights.loading.weight_applier import WeightApplier @@ -163,6 +186,12 @@ def load_into(root: Path, definition: Any, models: dict[str, Any], *, what: str) download_patterns=definition.get_download_patterns(), ) except (FileNotFoundError, ValueError) as exc: + missing = missing_weights(str(exc)) + if missing is not None: + raise DFloatFormatError( + f"{root}: not a {what} with {names}: {missing} not in the base (a missing shard, a renamed " + f"tensor or another model's checkpoint?): {exc}" + ) from exc raise DFloatFormatError( f"{root}: not a {what} with {names} (a DFloat11 checkpoint or model repository?): {exc}" ) from exc diff --git a/src/mlx_dfloat/mflux/_pipeline.py b/src/mlx_dfloat/mflux/_pipeline.py index 2116625..2beefcb 100644 --- a/src/mlx_dfloat/mflux/_pipeline.py +++ b/src/mlx_dfloat/mflux/_pipeline.py @@ -76,6 +76,24 @@ def refuse_construction_args( refuse("bake_lora", "LoRA of any kind") +def refuse_precision_args(*, float32: bool, compute_precision: Any) -> None: + """Refuse mflux's precision options (Z-Image's ``float32`` stream, ``compute_precision``), accepted only to refuse. + + ``compute_precision`` casts the attention and feed-forward weights in place and ``float32`` changes the + transformer's stream dtype; neither run was checked or measured on this path. + + Raises: + DFloatUnsupportedError: ``compute_precision`` set or ``float32=True``. + """ + if compute_precision is not None: + refuse( + "compute_precision", + "a lower compute precision changes the output the format exists to keep", + ) + if float32: + refuse("float32", "the float32 transformer stream") + + _wired_refusal_warned = False diff --git a/src/mlx_dfloat/mflux/ernie/model.py b/src/mlx_dfloat/mflux/ernie/model.py index b113aee..7c58807 100644 --- a/src/mlx_dfloat/mflux/ernie/model.py +++ b/src/mlx_dfloat/mflux/ernie/model.py @@ -360,7 +360,7 @@ def cfg_prompts( ) -> tuple[str, ...]: """The prompts a call encodes, in mflux's batch order: the prompt alone, or the negative then the prompt. - mflux 0.20's rule: guidance is read as ``Config`` stores it (None is 0.0); above 1.0 the call runs + mflux 0.23's rule: guidance is read as ``Config`` stores it (None is 0.0); above 1.0 the call runs classifier-free guidance with the negative prompt, a blank or missing one replaced by ``" "``. """ g = 0.0 if guidance is None else float(guidance) diff --git a/src/mlx_dfloat/mflux/ernie/names.py b/src/mlx_dfloat/mflux/ernie/names.py index 1cc115b..ba401f6 100644 --- a/src/mlx_dfloat/mflux/ernie/names.py +++ b/src/mlx_dfloat/mflux/ernie/names.py @@ -1,6 +1,6 @@ """ERNIE-Image naming: the block kind, the decoded sub-paths, the map derived from mflux, the group check. -The block kind and its run order: mflux 0.20.0 ``ErnieTransformer`` runs one list, ``layers``. The decoded sub-paths +The block kind and its run order: mflux 0.23.0 ``ErnieTransformer`` runs one list, ``layers``. The decoded sub-paths and the three non-block groups are the ``pattern_dict`` of both mingyi456 DF11 checkpoints (ERNIE-Image and ERNIE-Image-Turbo, identical). mflux's ``ErnieWeightMapping`` already templates its block targets with ``{layer}``, so the shared derivation reads it as it is. @@ -31,7 +31,7 @@ "adaLN_modulation.1": ("adaLN_modulation.1.weight",), "final_norm.linear": ("final_norm.linear.weight",), } -DEFAULT_LAYERS = 36 # mflux 0.20.0 ``ErnieTransformer.__init__(num_layers=36)`` +DEFAULT_LAYERS = 36 # mflux 0.23.0 ``ErnieTransformer.__init__(num_layers=36)`` def ernie_name_map() -> StaticNameMap: diff --git a/src/mlx_dfloat/mflux/ernie/transformer.py b/src/mlx_dfloat/mflux/ernie/transformer.py index 7d86e2f..cc92c96 100644 --- a/src/mlx_dfloat/mflux/ernie/transformer.py +++ b/src/mlx_dfloat/mflux/ernie/transformer.py @@ -1,6 +1,6 @@ """mflux's ERNIE-Image transformer with the class-swap seam; the build from a DF11 checkpoint. -mflux 0.20.0 ``ErnieTransformer.__call__`` runs its one block list, ``layers``, inline. The pipeline's ``_predict`` +mflux 0.23.0 ``ErnieTransformer.__call__`` runs its one block list, ``layers``, inline. The pipeline's ``_predict`` factory compiles the step function off base and Pro M1/M2 chips; the model runs it through ``mlx_dfloat.mflux._compile.uncompiled`` on every chip, because the block seam evaluates at block boundaries, which cannot happen inside ``mx.compile``. @@ -40,7 +40,7 @@ MAX_BUILD_ACTIVE_BYTES = 2 * 1024**3 -# mflux 0.20.0's ErnieTransformer builds no parameter of its own: the RoPE embedder holds no arrays, the timestep +# mflux 0.23.0's ErnieTransformer builds no parameter of its own: the RoPE embedder holds no arrays, the timestep # frequencies are computed per call, and ``_pos_cache`` (an underscore attribute) is not a parameter. Kept as a named # constant so the build reads like the other families' and a future mflux table has one place to go. COMPUTED_PARAMS: frozenset[str] = frozenset() diff --git a/src/mlx_dfloat/mflux/families.py b/src/mlx_dfloat/mflux/families.py index 73f36c4..677179e 100644 --- a/src/mlx_dfloat/mflux/families.py +++ b/src/mlx_dfloat/mflux/families.py @@ -169,7 +169,9 @@ def _krea2_names() -> NameMap: } -def _flux(name: str, label: str, steps: int, df11: str, base: str) -> ModelEntry: +def _flux( + name: str, label: str, steps: int, df11: str, base: str, *, base_revision: str +) -> ModelEntry: return ModelEntry( name=name, family="flux1", @@ -182,11 +184,14 @@ def _flux(name: str, label: str, steps: int, df11: str, base: str) -> ModelEntry default_scheduler="linear", cfg_two_calls=False, uses_negative_prompt=False, + base_revision=base_revision, ) -def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> ModelEntry: - """A FLUX.2 Klein entry, with mflux 0.20.0's defaults for it. +def _klein( + name: str, label: str, steps: int, revision: str, *, base: bool, base_revision: str +) -> ModelEntry: + """A FLUX.2 Klein entry, with mflux's defaults for it. The default steps are 50 for a base model and 4 for a distilled one; guidance defaults to 1.0, and the scheduler is mflux's fixed one. A base model calls the transformer twice per step above guidance 1.0 (the negative is @@ -208,26 +213,37 @@ def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> M cfg_two_calls=base, uses_negative_prompt=False, fixed_guidance=None if base else 1.0, + base_revision=base_revision, ) MODELS: dict[str, ModelEntry] = { e.name: e for e in ( + # FLUX.1 bases: the snapshots of the identity and MEASURED tier runs. _flux( "schnell", "FLUX.1-schnell", 4, "DFloat11/FLUX.1-schnell-DF11", "black-forest-labs/FLUX.1-schnell", + base_revision="741f7c3ce8b383c54771c7003378a50191e9efe9", + ), + _flux( + "dev", + "FLUX.1-dev", + 25, + "DFloat11/FLUX.1-dev-DF11", + "black-forest-labs/FLUX.1-dev", + base_revision="3de623fc3c33e44ffbe2bad470d0f45bccf2eb21", ), - _flux("dev", "FLUX.1-dev", 25, "DFloat11/FLUX.1-dev-DF11", "black-forest-labs/FLUX.1-dev"), _flux( "krea-dev", "FLUX.1-Krea-dev", 25, "DFloat11/FLUX.1-Krea-dev-DF11", "black-forest-labs/FLUX.1-Krea-dev", + base_revision="8162a9c7b05a641be098422bf2fcf335615c2f28", ), ModelEntry( name="z-image", @@ -241,6 +257,8 @@ def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> M default_scheduler=None, cfg_two_calls=True, uses_negative_prompt=True, + # The snapshot the identity check and the MEASURED row ran on. + base_revision="04cc4abb7c5069926f75c9bfde9ef43d49423021", ), ModelEntry( name="z-image-turbo", @@ -254,13 +272,18 @@ def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> M default_scheduler=None, cfg_two_calls=False, uses_negative_prompt=False, + # The snapshot the sampled parity run compared against. + base_revision="f332072aa78be7aecdf3ee76d5c247082da564a6", ), + # FLUX.2 Klein bases: the identity and MEASURED tier snapshots (base models) and the sampled parity snapshots + # (distilled models). _klein( "flux2-klein-base-4b", "FLUX.2-klein-base-4B", 50, "b887c73c5cbc3f4d50887a04c41509fb25a0a0f0", base=True, + base_revision="a3b4f4849157f664bdbc776fd7453c2783562f4d", ), _klein( "flux2-klein-4b", @@ -268,6 +291,7 @@ def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> M 4, "d29a2c249ff0afeb678da101e707bd134596c96f", base=False, + base_revision="e7b7dc27f91deacad38e78976d1f2b499d76a294", ), _klein( "flux2-klein-base-9b", @@ -275,6 +299,7 @@ def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> M 50, "50dc8e7cba7a41eeaf9e9c4fcec557d1a6888ded", base=True, + base_revision="32773329fbe7e81a90ef971740e8ba4b0364ecf3", ), _klein( "flux2-klein-9b", @@ -282,8 +307,9 @@ def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> M 4, "45a202a0bd19ec01ce8db2d5236890586cbac6a6", base=False, + base_revision="92196c8e11f7b6cf2b7493e037d8c5345c559216", ), - # mflux 0.20.0: 40 steps (cli/defaults/defaults.py:52), guidance 1.0 (models/qwen21/cli/qwen21_generate.py:60), + # mflux's defaults, cited at 0.20.0: 40 steps (cli/defaults/defaults.py:52), guidance 1.0 (qwen21_generate.py:60), # scheduler "linear" (cli/parser/parsers.py:180). A second transformer call per step runs only above guidance # 1.0 with a non-empty negative prompt (models/qwen21/variants/txt2img/qwen_image_21.py:89). ModelEntry( diff --git a/src/mlx_dfloat/mflux/flux1/model.py b/src/mlx_dfloat/mflux/flux1/model.py index b96762a..800ba88 100644 --- a/src/mlx_dfloat/mflux/flux1/model.py +++ b/src/mlx_dfloat/mflux/flux1/model.py @@ -30,6 +30,7 @@ from mlx_dfloat.integrate.names import NameMap, Shapes from mlx_dfloat.integrate.providers import DF11Provider from mlx_dfloat.mflux import _pipeline, require_mflux +from mlx_dfloat.mflux import families as registry from mlx_dfloat.mflux.flux1 import init as base_init from mlx_dfloat.mflux.flux1.init import BaseComponents, ResolvedRepo from mlx_dfloat.mflux.flux1.lifecycle import Lifecycle @@ -170,7 +171,11 @@ def __init__( df11 = base_init.resolve(df11_path or df11_default, patterns=base_init.DF11_PATTERNS) ckpt = open_checkpoint(df11.root) check_flux_groups(ckpt) - base = base_init.resolve(base_path or base_default, patterns=base_init.BASE_PATTERNS) + # The default base is pinned to the snapshot the recorded runs used; a user's own base is taken as given. + base_pin = None if base_path else registry.MODELS[model].base_revision + base = base_init.resolve( + base_path or base_default, patterns=base_init.BASE_PATTERNS, revision=base_pin + ) components = base_init.load_base(base.root, model_config) transformer, shapes = build_transformer(model_config, ckpt) self._assemble( diff --git a/src/mlx_dfloat/mflux/flux1/names.py b/src/mlx_dfloat/mflux/flux1/names.py index 3bdbe68..17fd3e9 100644 --- a/src/mlx_dfloat/mflux/flux1/names.py +++ b/src/mlx_dfloat/mflux/flux1/names.py @@ -9,9 +9,8 @@ DOUBLE_PREFIX = "transformer_blocks" SINGLE_PREFIX = "single_transformer_blocks" MATRICES_PER_KIND = {DOUBLE_PREFIX: 14, SINGLE_PREFIX: 6} -# mflux 0.20 builds `norm_out.linear` with bias=False and drops this bias through update(strict=False); -# the seam follows mflux so its output equals mflux's. -DROPPED_EXTRAS: frozenset[str] = frozenset({"norm_out.linear.bias"}) +# Every extra the checkpoint carries is loaded. mflux 0.23 builds `norm_out.linear` with a bias and loads it, so the +# checkpoint's `norm_out.linear.bias` is an extra like any other and the coverage check refuses a build without it. def flux_name_map() -> StaticNameMap: diff --git a/src/mlx_dfloat/mflux/flux1/transformer.py b/src/mlx_dfloat/mflux/flux1/transformer.py index f79474c..f1fdde4 100644 --- a/src/mlx_dfloat/mflux/flux1/transformer.py +++ b/src/mlx_dfloat/mflux/flux1/transformer.py @@ -20,7 +20,6 @@ from mlx_dfloat.mflux import require_mflux from mlx_dfloat.mflux.flux1.names import ( DOUBLE_PREFIX, - DROPPED_EXTRAS, SINGLE_PREFIX, check_flux_groups, flux_name_map, @@ -217,7 +216,6 @@ def build_transformer( ckpt, names, counts={DOUBLE_PREFIX: n_double, SINGLE_PREFIX: n_single}, - dropped=DROPPED_EXTRAS, extras=extras, ) check_extras_cover(params, (name for name, _path, _info in plan), matrix_paths) diff --git a/src/mlx_dfloat/mflux/flux2/init.py b/src/mlx_dfloat/mflux/flux2/init.py index ae7f2c4..c123b89 100644 --- a/src/mlx_dfloat/mflux/flux2/init.py +++ b/src/mlx_dfloat/mflux/flux2/init.py @@ -23,7 +23,7 @@ BASE_PATTERNS: tuple[str, ...] = ENCODER_PATTERNS + VAE_PATTERNS + TOKENIZER_PATTERNS _WHAT = "FLUX.2 Klein base" -# Parameters no checkpoint holds: built at construction, so a base never replaces them. mflux 0.20.0: +# Parameters no checkpoint holds: built at construction, so a base never replaces them. mflux 0.23.0: # ``Qwen3TextEncoder.rotary_emb`` (qwen3_text_encoder.py:47-52) computes ``inv_freq`` in # ``Qwen3TextRotaryEmbedding.__init__`` (qwen3_text_rotary_embedding.py:18). The VAE has none: its batch-norm # statistics ``bn.running_mean`` / ``bn.running_var`` come from the file (flux2_weight_mapping.py:202-209). diff --git a/src/mlx_dfloat/mflux/flux2/model.py b/src/mlx_dfloat/mflux/flux2/model.py index bdda41b..c0e2605 100644 --- a/src/mlx_dfloat/mflux/flux2/model.py +++ b/src/mlx_dfloat/mflux/flux2/model.py @@ -57,7 +57,7 @@ # call's allowance, at least 500_000_000 B, + this slack) keeps those runs under it at every size, and an encode that # runs all 36 layers over it at 1024² (4B: 9 more layers of 201_861_632 B). ENCODE_SLACK_BYTES = 256 * 1024**2 -# mflux 0.20.0 ``Flux2Klein.generate_image``: the negative prompt it encodes, its guidance and scheduler defaults. +# mflux 0.23.0 ``Flux2Klein.generate_image``: the negative prompt it encodes, its guidance and scheduler defaults. NEGATIVE_PROMPT = " " DEFAULT_GUIDANCE = 1.0 DEFAULT_SCHEDULER = "flow_match_euler_discrete" @@ -101,6 +101,7 @@ def __init__( lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, bake_lora: bool = True, + compute_precision: Any = None, ) -> None: """Resolve the checkpoint and the base repository and build the model. @@ -109,8 +110,8 @@ def __init__( estimate and VAE decision. Raises: - DFloatUnsupportedError: ``quantize``, ``lora_paths``, ``lora_scales`` or ``bake_lora=False`` - (accepted only to refuse), an eval policy other than ``"per-block"`` / ``"depth2"``, the KV-cache + DFloatUnsupportedError: ``quantize``, ``lora_paths``, ``lora_scales``, ``bake_lora=False`` or a + ``compute_precision`` (accepted only to refuse), an eval policy other than ``"per-block"`` / ``"depth2"``, the KV-cache model, or an unknown ``model``. DFloatBackendError: The GPU decoder fails its check on the two packaged self-check groups (raised before anything is resolved or loaded), or Metal is unavailable. @@ -124,6 +125,7 @@ def __init__( _pipeline.refuse_construction_args( quantize=quantize, lora_paths=lora_paths, lora_scales=lora_scales, bake_lora=bake_lora ) + _pipeline.refuse_precision_args(float32=False, compute_precision=compute_precision) _pipeline.check_eval_policy(eval_policy, POLICIES) if "kv" in model.lower(): raise DFloatUnsupportedError( @@ -143,7 +145,10 @@ def __init__( pin = None if df11_path else registry.MODELS[model].df11_revision df11 = finit.resolve(df11_path or df11_default, patterns=finit.DF11_PATTERNS, revision=pin) ckpt = open_checkpoint(df11.root) - base = finit.resolve(base_path or base_default, patterns=finit.BASE_PATTERNS) + base_pin = None if base_path else registry.MODELS[model].base_revision + base = finit.resolve( + base_path or base_default, patterns=finit.BASE_PATTERNS, revision=base_pin + ) build = build_transformer(ckpt, transformer_overrides=model_config.transformer_overrides) components = finit.load_base(base.root, model_config) self._assemble( @@ -186,6 +191,7 @@ def _assemble( budget_bytes: int | None = None, decode: Callable[[MxGroup], DecodeResult] | None = None, ) -> None: + from mflux.models.common.compute_precision import ComputePrecision from mflux.models.flux2.flux2_initializer import Flux2Initializer _pipeline.check_eval_policy(eval_policy, POLICIES) @@ -193,6 +199,8 @@ def _assemble( Flux2Initializer._init_config( self, model_config ) # prompt_cache, model_config, callbacks, tiling_config + # mflux's default (Flux2Initializer.init sets it): no compute precision; generate_image records it. + self.compute_precision = ComputePrecision() self.vae = components.vae self.text_encoder = components.text_encoder self.tokenizers = components.tokenizers @@ -258,7 +266,7 @@ def _unload_encoders(self) -> None: def _encode(self, prompt: str) -> tuple[mx.array, ...]: """The prompt's embeddings and text ids, as ``Flux2Klein._encode_prompt_pair`` computes them. - mflux 0.20.0 hard-codes 512 tokens and hidden states (9, 18, 27) there; the model config's token count and + mflux 0.23.0 hard-codes 512 tokens and hidden states (9, 18, 27) there; the model config's token count and ``memory.TEXT_ENCODER_OUT_LAYERS`` carry the same values (a test pins them to mflux's source). """ from mflux.models.flux2.model.flux2_text_encoder.prompt_encoder import Flux2PromptEncoder @@ -372,7 +380,7 @@ def plan_call(self, *, height: int, width: int) -> CallPlan: def cfg_prompts(self, prompt: str, *, guidance: float | None) -> tuple[str, ...]: """The prompts a call encodes: the prompt, and mflux's blank negative ``" "`` when guidance is above 1. - mflux 0.20's rule (``Flux2Klein.generate_image`` and ``_encode_prompt_pair``): the negative branch runs, and + mflux 0.23's rule (``Flux2Klein.generate_image`` and ``_encode_prompt_pair``): the negative branch runs, and the transformer is called twice per step, only for a guidance above 1.0. """ if guidance is None or guidance <= 1.0: diff --git a/src/mlx_dfloat/mflux/flux2/names.py b/src/mlx_dfloat/mflux/flux2/names.py index 2c04bd1..84c9b6e 100644 --- a/src/mlx_dfloat/mflux/flux2/names.py +++ b/src/mlx_dfloat/mflux/flux2/names.py @@ -1,6 +1,6 @@ """FLUX.2 Klein naming: the block kinds in run order, the compressed sub-paths, the map derived from mflux, the check. -Block kinds and run order: mflux 0.20.0 ``Flux2Transformer.__call__`` (double blocks, then single blocks). Compressed +Block kinds and run order: mflux 0.23.0 ``Flux2Transformer.__call__`` (double blocks, then single blocks). Compressed sub-paths: the ``pattern_dict`` of the mingyi456 FLUX.2 Klein DF11 checkpoints (base and distilled, 4B and 9B). An empty pattern list there means the group is the module itself, so its one matrix is ``.weight``. """ diff --git a/src/mlx_dfloat/mflux/flux2/transformer.py b/src/mlx_dfloat/mflux/flux2/transformer.py index 2a27718..f1c0d41 100644 --- a/src/mlx_dfloat/mflux/flux2/transformer.py +++ b/src/mlx_dfloat/mflux/flux2/transformer.py @@ -28,7 +28,7 @@ klein_name_map, ) -# mflux 0.20.0 ``Flux2Transformer.__init__`` defaults (the 4B depth), used when the overrides leave them out. +# mflux 0.23.0 ``Flux2Transformer.__init__`` defaults (the 4B depth), used when the overrides leave them out. DEFAULT_DOUBLE, DEFAULT_SINGLE = 5, 20 # ``ModelConfig.transformer_overrides["num_attention_heads"]`` per Klein size (24 heads x 128 = 3072 wide; 32 = 4096). KLEIN_SIZES: dict[int, str] = {24: "4b", 32: "9b"} diff --git a/src/mlx_dfloat/mflux/generate.py b/src/mlx_dfloat/mflux/generate.py index 88ed090..497785f 100644 --- a/src/mlx_dfloat/mflux/generate.py +++ b/src/mlx_dfloat/mflux/generate.py @@ -6,10 +6,12 @@ import argparse import json +import platform import sys import time import traceback from collections.abc import Callable, Mapping +from importlib import metadata from pathlib import Path from typing import Any @@ -355,13 +357,23 @@ def _scheduler(args: argparse.Namespace) -> str | None: return args.scheduler if args.scheduler is not None else entry(args.model).default_scheduler +def _versions() -> dict[str, str | None]: + """The mlx and mflux versions, the generate report's ``versions`` (mflux None when it is not installed).""" + try: + mflux: str | None = metadata.version("mflux") + except metadata.PackageNotFoundError: + mflux = None + return {"mlx": mx.__version__, "mflux": mflux} # type: ignore[attr-defined] + + def _run_context( args: argparse.Namespace, *, tier_gb: int, label: str, limits: Mapping[str, int] ) -> dict[str, Any]: """What the watchdog's abort artifact records about the run it may stop. ``limits`` is the MLX ``memory`` / ``cache`` / ``wired`` limits read back after the install, - so a stopped run still says which limits it ran under. + so a stopped run still says which limits it ran under; ``versions`` and ``macos`` have the generate report's + shape, so its tier-table row still names the mflux it ran on. """ return { "model": args.model, @@ -372,6 +384,8 @@ def _run_context( "tier_gb": tier_gb, "label": label, "limits": {name: int(limits[name]) for name in ("memory", "cache", "wired")}, + "versions": _versions(), + "macos": platform.mac_ver()[0], # as the generate report records it } @@ -510,6 +524,7 @@ def run( "label": None, "watchdog_ceiling_bytes": None, "limits": None, + "macos": platform.mac_ver()[0], # as the bench provenance records it } try: host_ram, host_recommended = host_facts() @@ -575,9 +590,12 @@ def run( call["negative_prompt"] = args.negative_prompt image = model.generate_image(**call) final = Path(resolve(output)) - image.save(str(final), export_json_metadata=args.metadata, overwrite=True) + try: + image.save(str(final), export_json_metadata=args.metadata, overwrite=True) + except OSError as exc: # mflux 0.23 raises when the image file itself cannot be written + raise DFloatError(f"{final}: the image was not written: {exc}") from exc if not final.is_file() or final.stat().st_size == 0: - # mflux's save logs a write failure and returns normally + # an empty or missing file after a save that returned normally is still no image raise DFloatError(f"{final}: the image was not written") report["output"] = str(final) report.update(exit_code=EXIT_OK, **model.report()) diff --git a/src/mlx_dfloat/mflux/krea2/init.py b/src/mlx_dfloat/mflux/krea2/init.py index 46d5df7..ed8ef44 100644 --- a/src/mlx_dfloat/mflux/krea2/init.py +++ b/src/mlx_dfloat/mflux/krea2/init.py @@ -35,7 +35,7 @@ BASE_PATTERNS: tuple[str, ...] = ENCODER_PATTERNS + VAE_PATTERNS + TOKENIZER_PATTERNS _WHAT = "Krea 2 base" -# Parameters no checkpoint holds: built at construction, so a base never replaces them. mflux 0.20.0: the text +# Parameters no checkpoint holds: built at construction, so a base never replaces them. mflux 0.23.0: the text # encoder's ``rotary_emb.inv_freq`` is computed in ``Qwen3TextRotaryEmbedding.__init__`` (the base file holds only # ``language_model.*`` and ``visual.*`` tensors). The VAE has none: every one of its parameters is a target of mflux's # VAE mapping (Qwen-Image 1's). diff --git a/src/mlx_dfloat/mflux/krea2/model.py b/src/mlx_dfloat/mflux/krea2/model.py index e9253c0..6dd7dc4 100644 --- a/src/mlx_dfloat/mflux/krea2/model.py +++ b/src/mlx_dfloat/mflux/krea2/model.py @@ -407,7 +407,7 @@ def cfg_prompts( ) -> tuple[str, ...]: """The prompts a call encodes: the prompt, then the negative one when classifier-free guidance runs. - mflux 0.20's rule for Krea 2: the negative branch runs for any guidance other than 1.0 (below 1.0 too); a + mflux 0.23's rule for Krea 2: the negative branch runs for any guidance other than 1.0 (below 1.0 too); a blank or missing negative prompt is encoded as ``" "``. """ if guidance == 1.0: diff --git a/src/mlx_dfloat/mflux/krea2/names.py b/src/mlx_dfloat/mflux/krea2/names.py index a0773e1..aa1adcd 100644 --- a/src/mlx_dfloat/mflux/krea2/names.py +++ b/src/mlx_dfloat/mflux/krea2/names.py @@ -1,6 +1,6 @@ """Krea 2 naming: the block kind, the decoded sub-paths, the map derived from mflux, the group and variant checks. -The block kind and its run order: mflux 0.20.0 ``Krea2Transformer`` runs one list, ``blocks``, inline. The decoded +The block kind and its run order: mflux 0.23.0 ``Krea2Transformer`` runs one list, ``blocks``, inline. The decoded sub-paths and the seven non-block groups (the four text-fusion blocks, ``tmlp``, ``tproj``, ``txtmlp``) are the ``pattern_dict`` of the two ComfyUI single-file DF11 exports (Krea 2 Raw and Krea 2 Turbo, identical); their stored order was checked against the originals' bytes. mflux's ``Krea2WeightMapping`` templates its block targets with @@ -46,7 +46,7 @@ "txtmlp": ("txtmlp.1.weight", "txtmlp.3.weight"), **{g: tuple(f"{g}.{s}.weight" for s in TEXT_SUBS) for g in TEXT_FUSION_GROUPS}, } -DEFAULT_LAYERS = 28 # mflux 0.20.0 ``Krea2Transformer.__init__(layers=28)`` +DEFAULT_LAYERS = 28 # mflux 0.23.0 ``Krea2Transformer.__init__(layers=28)`` LAYOUT_FOR_MODEL: dict[str, str] = { "krea-2-raw": "krea-2-raw-comfyui", "krea-2": "krea-2-turbo-comfyui", diff --git a/src/mlx_dfloat/mflux/krea2/transformer.py b/src/mlx_dfloat/mflux/krea2/transformer.py index 3e3b9da..9ec155d 100644 --- a/src/mlx_dfloat/mflux/krea2/transformer.py +++ b/src/mlx_dfloat/mflux/krea2/transformer.py @@ -1,6 +1,6 @@ """mflux's Krea 2 transformer with the class-swap seam. -mflux 0.20.0 ``Krea2Transformer.__call__`` runs its one block list, ``blocks``, inline, after the text fusion, the +mflux 0.23.0 ``Krea2Transformer.__call__`` runs its one block list, ``blocks``, inline, after the text fusion, the timestep MLP and the text MLP have run once. The pipeline's ``_predict`` factory compiles the step function off base and Pro M1/M2 chips (an M1 Max or Ultra, and every M3 or later, compile); the model runs it through ``mlx_dfloat.mflux._compile.uncompiled`` on every chip, because the block seam evaluates at block boundaries, which @@ -64,7 +64,8 @@ MAX_BUILD_ACTIVE_BYTES = 2 * 1024**3 -# mflux 0.20.0's Krea2Transformer builds no parameter of its own: the RoPE embedder holds no arrays, the timestep +# mflux 0.23.0's Krea2Transformer builds no parameter of its own (the depth ControlNet's ``control_first`` is added +# only by its own initializer; the txt2img path never passes ``control``, and the seam forwards it unchanged): the RoPE embedder holds no arrays, the timestep # frequencies are computed per call, and the modulation tables and norm scales come from the checkpoint. Kept as a # named constant so the build reads like the other families' and a future mflux table has one place to go. COMPUTED_PARAMS: frozenset[str] = frozenset() diff --git a/src/mlx_dfloat/mflux/qwen21/init.py b/src/mlx_dfloat/mflux/qwen21/init.py index dbe15ec..5bf4e8f 100644 --- a/src/mlx_dfloat/mflux/qwen21/init.py +++ b/src/mlx_dfloat/mflux/qwen21/init.py @@ -35,10 +35,10 @@ BASE_PATTERNS: tuple[str, ...] = ENCODER_PATTERNS + VAE_PATTERNS + TOKENIZER_PATTERNS _WHAT = "Qwen-Image 2.1 base" -# Parameters no checkpoint holds: built at construction, so a base never replaces them. mflux 0.20.0: the text +# Parameters no checkpoint holds: built at construction, so a base never replaces them. mflux 0.23.0: the text # encoder's ``rotary_emb.inv_freq`` is computed in ``Qwen3VLRotaryEmbedding.__init__`` (qwen3_vl_rope.py:21). The # VAE has none: every one of its parameters is a target of mflux's VAE mapping, and its latent mean and standard -# deviation are NumPy class attributes, not parameters (qwen21_vae.py:12-32). +# deviation are NumPy class attributes, not parameters (qwen21_vae.py:9-10, from qwen21_vae_config.py:6-17). COMPUTED: dict[str, frozenset[str]] = { "text_encoder": frozenset({"rotary_emb.inv_freq"}), "vae": frozenset(), diff --git a/src/mlx_dfloat/mflux/qwen21/memory.py b/src/mlx_dfloat/mflux/qwen21/memory.py index 9e9051d..ebf99a1 100644 --- a/src/mlx_dfloat/mflux/qwen21/memory.py +++ b/src/mlx_dfloat/mflux/qwen21/memory.py @@ -7,8 +7,8 @@ The transformer has one block kind, so the cache limit holds one decoded block next to the activations (per-block evaluation frees one block-sized buffer at a time), and the family carries its own activation allowance. The VAE -transient is mflux's VAE decode, which runs in float32 (mflux 0.20.0 ``qwen21_vae.py:46-48``: the float32 latent -mean and standard deviation promote the latents): re-measure it when an mflux upgrade changes that dtype. No Qwen-Image 2.1 per-block overhead scenario is benched. +transient is mflux's VAE decode, which runs in float32 (mflux 0.23.0 ``qwen21_vae.py:17-18``, and ``28-29`` on the +encode side: the float32 latent mean and standard deviation promote the latents): re-measure it when an mflux upgrade changes that dtype. No Qwen-Image 2.1 per-block overhead scenario is benched. """ from dataclasses import replace @@ -42,14 +42,18 @@ OVERHEAD = 459_445_098 # 0.428 GiB # VAE_TRANSIENT, by the overall-peak rule over every 1024² sample: the highest measured peak of the VAE decode on the # resident set minus the VAE phase's other terms (compressed + extras + decoded non-block + VAE file + OVERHEAD): -# 22_788_580_400 - (9_587_103_379 + 137_388_032 + 134_217_728 + 1_350_989_512 + 459_445_098). Eleven samples on -# 2026-10-08 (the calibration run above, four calibration reruns of one step at two cache limits, three 40-step MEASURED -# runs, three 4-step identity runs with CFG; the highest is the second one-step rerun at the old cache limit) -# peaked at 21.76 to 22.79 GB in the VAE phase, with MLX's own VAE peak the same in every calibration run: about 1 GiB of -# run-to-run spread outside MLX's counters, as Z-Image's VAE phase shows too. A fit check must not under-predict, so -# the term covers the highest sample; the first run alone (22_139_643_024) would put the VAE phase 0.60 GiB under it. -# About 1.32x the FLUX.1 VAE's 7.82 GiB. -VAE_TRANSIENT = 11_119_436_651 # 10.356 GiB +# 22_831_539_272 - (9_587_103_379 + 137_388_032 + 134_217_728 + 1_350_989_512 + 459_445_098). Thirteen samples, all +# peaking in the VAE phase at 21.76 to 22.83 GB. Eleven on mflux 0.20 (2026-10-08: the calibration run above, four +# calibration reruns of one step at two cache limits, three 40-step MEASURED runs, three 4-step identity runs with CFG; +# bench/results/calibration/qwen-image-2.1-1024.json), with MLX's own VAE peak the same in every calibration run: about +# 1 GiB of run-to-run spread outside MLX's counters, as Z-Image's VAE phase shows too. Two on mflux 0.23 (2026-10-24, +# git cbed93b: the 40-step MEASURED run, 21_877_154_888, and the identity check's df11 side, 22_831_539_272; +# bench/results/calibration/qwen-image-2.1-1024-mflux023.json names the record fields). The 0.23 identity run binds; +# the highest 0.20 sample (the second one-step rerun at the old cache limit, 22_788_580_400, the 0.20 value's source) +# is 42_958_872 B under it and the 0.23 MEASURED run 954_384_384 B. A fit check must not under-predict, so the term +# covers the highest sample; the first run alone (22_139_643_024) would put the VAE phase 0.64 GiB under it. +# About 1.33x the FLUX.1 VAE's 7.82 GiB. +VAE_TRANSIENT = 11_162_395_523 # 10.396 GiB # DENOISE_ACTIVATION: denoise footprint peak minus the denoise phase's other terms at the planner's cache limit # (compressed + extras + non-block + in-flight + cache limit + OVERHEAD), at REFERENCE_TOKENS: # 13_939_106_072 - (9_587_103_379 + 137_388_032 + 134_217_728 + 436_207_616 + 2_731_761_634 + 459_445_098). Measured diff --git a/src/mlx_dfloat/mflux/qwen21/model.py b/src/mlx_dfloat/mflux/qwen21/model.py index db6a68b..5bb90ab 100644 --- a/src/mlx_dfloat/mflux/qwen21/model.py +++ b/src/mlx_dfloat/mflux/qwen21/model.py @@ -54,7 +54,7 @@ # accepted, and the report labels it unmeasured. MEASURED_POLICIES: frozenset[str] = frozenset({"per-block"}) RETAINED_SLACK_BYTES = 2 * 1024**3 -# mflux 0.20.0 ``QwenImage21.generate_image``: its guidance and scheduler defaults. +# mflux 0.23.0 ``QwenImage21.generate_image``: its guidance and scheduler defaults (unchanged since 0.20.0). DEFAULT_GUIDANCE = 1.0 DEFAULT_SCHEDULER = "linear" @@ -107,6 +107,7 @@ def __init__( lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, bake_lora: bool = True, + compute_precision: Any = None, ) -> None: """Resolve the checkpoint and the base repository and build the model. @@ -115,7 +116,7 @@ def __init__( estimate and VAE decision. Raises: - DFloatUnsupportedError: ``quantize``, ``lora_paths``, ``lora_scales`` or ``bake_lora=False`` + DFloatUnsupportedError: ``quantize``, ``lora_paths``, ``lora_scales``, ``bake_lora=False`` or a ``compute_precision`` (accepted only to refuse), an eval policy other than ``"per-block"`` / ``"depth2"``, or an unknown ``model``. DFloatBackendError: The GPU decoder fails its check on the two packaged self-check groups (raised before @@ -131,6 +132,7 @@ def __init__( _pipeline.refuse_construction_args( quantize=quantize, lora_paths=lora_paths, lora_scales=lora_scales, bake_lora=bake_lora ) + _pipeline.refuse_precision_args(float32=False, compute_precision=compute_precision) _pipeline.check_eval_policy(eval_policy, POLICIES) if model not in MODELS: raise DFloatUnsupportedError(f"model {model!r}: this path runs {sorted(MODELS)}") @@ -194,11 +196,12 @@ def _assemble( budget_bytes: int | None = None, decode: Callable[[MxGroup], DecodeResult] | None = None, ) -> None: + from mflux.models.common.compute_precision import ComputePrecision from mflux.models.qwen21.qwen21_initializer import Qwen21Initializer _pipeline.check_eval_policy(eval_policy, POLICIES) nn.Module.__init__(self) # type: ignore[attr-defined] # not QwenImage21.__init__: mflux's initializer loads the whole base - Qwen21Initializer._init_config( + Qwen21Initializer.init_config( self, model_config ) # prompt_cache, model_config, callbacks, tiling_config self.vae = components.vae @@ -206,6 +209,10 @@ def _assemble( self.tokenizers = components.tokenizers self.transformer = build.transformer self.bits = None + # mflux's defaults (Qwen21Initializer.init sets them): no compute precision, no LoRA; generate_image records them. + self.compute_precision = ComputePrecision() + self.lora_paths: list[str] = [] + self.lora_scales: list[float] = [] self._model = model self._constants = qmem.CONSTANTS self._ckpt, self._df11, self._base, self._sizes = ckpt, df11, base, sizes @@ -373,7 +380,7 @@ def cfg_prompts( ) -> tuple[str, ...]: """The prompts a call encodes: the prompt, and the negative prompt when classifier-free guidance runs. - mflux 0.20's rule (``QwenImage21.generate_image``): the negative branch runs, and the transformer is called + mflux 0.23's rule (``QwenImage21.generate_image``, unchanged since 0.20): the negative branch runs, and the transformer is called twice per step, only for a guidance above 1.0 together with a non-empty negative prompt. """ if guidance is None or guidance <= 1.0 or not negative_prompt: @@ -385,7 +392,7 @@ def generate_image( self, seed: int, prompt: str, - num_inference_steps: int = 40, + num_inference_steps: int | None = None, height: int = 1024, width: int = 1024, guidance: float | None = DEFAULT_GUIDANCE, @@ -393,19 +400,25 @@ def generate_image( image_strength: float | None = None, scheduler: str | None = DEFAULT_SCHEDULER, negative_prompt: str | None = None, + step_cache_ratio: float | None = None, + teacache_ratio: float | None = None, ) -> Any: """Mflux's ``generate_image`` behind a prelude: refusals, cache limit, fit check, prompt encoding, set load. - ``guidance`` or ``scheduler`` given as None take mflux's defaults (1.0, ``"linear"``). As in mflux, + ``guidance`` or ``scheduler`` given as None take mflux's defaults (1.0, ``"linear"``); ``num_inference_steps`` + None takes mflux's (``Qwen21TurboScheduler.default_steps``). As in mflux, classifier-free guidance runs only for a guidance above 1.0 with a non-empty ``negative_prompt``. Raises: - DFloatUnsupportedError: ``image_path`` / ``image_strength`` (img2img). + DFloatUnsupportedError: ``image_path`` / ``image_strength`` (img2img), ``step_cache_ratio`` / + ``teacache_ratio`` (step reuse: it skips transformer calls, a different image). DFloatResourceError: The fit check refuses the call (the estimate, or a size above 1024²), or memory stayed active after a drop. DFloatFormatError: A block's decode reported an error; the set is dropped for a clean retry. DFloatIntegrationError: mflux compiled the transformer's forward pass. """ + if step_cache_ratio is not None or teacache_ratio is not None: + _pipeline.refuse("step_cache_ratio/teacache_ratio", "step reuse") if image_path is not None or image_strength is not None: _pipeline.refuse("image_path/image_strength", "img2img") guidance = DEFAULT_GUIDANCE if guidance is None else guidance diff --git a/src/mlx_dfloat/mflux/qwen21/transformer.py b/src/mlx_dfloat/mflux/qwen21/transformer.py index cc67bf0..37181bb 100644 --- a/src/mlx_dfloat/mflux/qwen21/transformer.py +++ b/src/mlx_dfloat/mflux/qwen21/transformer.py @@ -1,8 +1,17 @@ -"""mflux's Qwen-Image 2.1 transformer with the class-swap seam and its forward pass kept uncompiled. +"""mflux's Qwen-Image 2.1 transformer with the block seam, its forward passes uncompiled and its projections packed per block. -mflux 0.20.0 ``Qwen21Transformer.__call__`` compiles its ``_forward`` on the first call, on every chip, and keeps it -in ``_step_fn``. The block seam evaluates at block boundaries, which cannot happen inside ``mx.compile``, so -``EagerForward`` sets ``_step_fn`` to the plain bound ``_forward`` first and refuses any other value. +mflux 0.23.0 ``Qwen21Transformer.__call__`` runs, for a prompt without padding, the text-prefix cache: each block's +``text_forward`` once per prompt, then each block's ``image_forward`` every step, compiled into ``_image_step_fn``; +padded prompts run the joint ``_forward`` compiled into ``_step_fn``. The block seam evaluates at block boundaries, +which cannot happen inside ``mx.compile``, so ``EagerForward`` sets both step functions to the plain bound methods +first and refuses any other value, and every block entry point (``__call__``, ``text_forward``, ``image_forward``) +runs through the seam. + +mflux 0.23.0 ``inference_projections()`` packs every block's Q/K/V and gate/up into one matrix each, once per +generation, and the blocks then run one matmul against it; on this machine that changes output bits against the +separate matmuls at the text-prefix shapes. Our blocks hold placeholders until their decode, so ``PackPerBlock`` +turns packing on for the window and each block packs its decoded weights with mflux's own ``_pack_linears`` for its +call, then drops the packed matrices. The build installs placeholders on the seven block matrices and on ``modulation.1`` (a non-block group, decoded once when the set loads), loads every other parameter from the checkpoint's extras (or a BF16 base's) and leaves the @@ -10,7 +19,8 @@ """ import types -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager from dataclasses import dataclass from functools import cache from pathlib import Path @@ -22,7 +32,8 @@ from mlx_dfloat._safetensors import TensorInfo from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError from mlx_dfloat.format import DF11Checkpoint -from mlx_dfloat.integrate.blockseam import StepSeam, seam_blocks +from mlx_dfloat.integrate import seam +from mlx_dfloat.integrate.blockseam import SeamCell, StepSeam from mlx_dfloat.integrate.coverage import check_extras_cover, extras_plan, load_extras from mlx_dfloat.integrate.names import NameMap, Shapes from mlx_dfloat.integrate.placeholders import PLACEHOLDER, install_placeholders @@ -40,65 +51,280 @@ MAX_BUILD_ACTIVE_BYTES = 2 * 1024**3 -# Parameters mflux computes when it builds the transformer; no checkpoint holds them, so the extras coverage leaves -# them out: the RoPE tables (mflux 0.20.0 qwen21_rope.py:21-26) and the timestep frequencies -# (qwen21_time_text_embed.py:16). -COMPUTED_PARAMS: frozenset[str] = frozenset( - { - "pos_embed.cos_tables.0", - "pos_embed.cos_tables.1", - "pos_embed.cos_tables.2", - "pos_embed.sin_tables.0", - "pos_embed.sin_tables.1", - "pos_embed.sin_tables.2", - "time_text_embed.time_proj.freqs", - } + +_BINDING = "_dfloat_qwen21_block" +_SWITCH = "_dfloat_pack_switch" +# The block methods mflux 0.23's transformer calls: ``__call__`` (the joint pass, padded prompts or the text cache +# off), ``text_forward`` (the text-prefix K/V cache, built once per prompt) and ``image_forward`` (every step after). +ENTRY_POINTS: tuple[str, ...] = ("__call__", "text_forward", "image_forward") + +Packer = Callable[[tuple[Any, ...]], Any] + + +class PackSwitch: + """Whether the blocks pack their projections: on inside ``inference_projections()`` (shared, never a parameter).""" + + __slots__ = ("on",) + + def __init__(self) -> None: + """Start off: outside ``inference_projections()`` mflux runs the separate projections.""" + self.on = False + + +@dataclass(frozen=True, slots=True) +class Qwen21Binding: + """A seamed block's name, the seam cell it reads its state from, the packing switch and the packer.""" + + name: str + cell: SeamCell + switch: PackSwitch + pack: Packer + + +def _pack(block: Any, pack: Packer) -> None: + """Pack the block's decoded weights as mflux 0.23 ``inference_projections()`` packs them, one block at a time.""" + attention, mlp = block.attn, block.img_mlp + attention._qkv_weight = pack((attention.to_q, attention.to_k, attention.to_v)) + mlp._gate_proj_weight = pack((mlp.gate_layer, mlp.proj)) + + +def _unpack(block: Any) -> None: + """Drop the packed matrices (mflux 0.23 ``_clear_inference_projections`` per block).""" + block.attn._qkv_weight = None + block.img_mlp._gate_proj_weight = None + + +class Qwen21BlockSeam: + """Placed in front of mflux's block class: every entry point runs through the seam, packed when switched on.""" + + __slots__ = () + + def _dfloat_run(self, entry: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + binding: Qwen21Binding = vars(self)[_BINDING] + state = binding.cell.state + if state is None: + raise DFloatIntegrationError( + f"{binding.name}: call attach(provider, shapes) before running a step" + ) + own = getattr(super(), entry) + + def run() -> Any: + # The weights are assigned by now; the packed matrices hold copies of them and go with the block's call. + try: + if binding.switch.on: + _pack(self, binding.pack) + return own(*args, **kwargs) + finally: + _unpack(self) + + return seam.run_block(state, binding.name, self, run, inputs=(args, kwargs)) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """The joint pass through the seam.""" + return self._dfloat_run("__call__", args, kwargs) + + def text_forward(self, *args: Any, **kwargs: Any) -> Any: + """The text-prefix pass through the seam.""" + return self._dfloat_run("text_forward", args, kwargs) + + def image_forward(self, *args: Any, **kwargs: Any) -> Any: + """The image-token pass through the seam.""" + return self._dfloat_run("image_forward", args, kwargs) + + +@cache +def seamed_block_class(cls: type) -> type: + """``Qwen21BlockSeam`` composed in front of ``cls`` (one class per block class).""" + return type( + f"Seam{cls.__name__}", (Qwen21BlockSeam, cls), {"__slots__": (), "__module__": __name__} + ) + + +def seam_qwen21_blocks( + block_lists: Sequence[tuple[str, Sequence[Any]]], + cell: SeamCell, + switch: PackSwitch, + *, + pack: Packer, +) -> tuple[str, ...]: + """Swap every block's class for its seamed class and bind it as ``.``; the names in run order. + + ``pack`` packs a tuple of ``nn.Linear`` layers into one matrix or returns None (mflux 0.23's + ``Qwen21Transformer._pack_linears``); a block packs only while ``switch.on``. + + Raises: + DFloatIntegrationError: A block is already seamed. + """ + names: list[str] = [] + for kind, blocks in block_lists: + for idx, block in enumerate(blocks): + name = f"{kind}.{idx}" + if isinstance(block, Qwen21BlockSeam): + raise DFloatIntegrationError(f"{name}: already seamed") + block_cls: type = type(block) + block.__class__ = seamed_block_class(block_cls) + vars(block)[_BINDING] = Qwen21Binding(name=name, cell=cell, switch=switch, pack=pack) + names.append(name) + return tuple(names) + + +# The tables mflux computes when it builds the transformer, which no checkpoint holds: the RoPE tables, one per axis +# (mflux 0.23.0 qwen21_rope.py:21-26, private lists), and the timestep frequencies (qwen21_time_text_embed.py:16, a +# private array). MLX leaves underscore attributes out of ``parameters()``, so the extras coverage never sees them +# and the build leaves them as mflux built them; one that turned into a parameter again would be refused as +# uncovered. +COMPUTED_TABLES: tuple[str, ...] = ( + "pos_embed._cos_tables", + "pos_embed._sin_tables", + "time_text_embed.time_proj._freqs", ) +# The parameters mflux computes at construction, as krea2 and ernie export them: none here, the computed tables above +# are not parameters. +COMPUTED_PARAMS: frozenset[str] = frozenset() -def is_eager(transformer: Any) -> bool: - """Whether the transformer's ``_step_fn`` is unset or its own plain, uncompiled ``_forward`` (imports mflux). - ``mx.compile`` returns ``mlx.gc_func``, a ``FunctionType`` subclass, so only an exact type check tells a compiled - function from a bound method. +def computed_tables(transformer: Any) -> dict[str, mx.array]: + """The computed tables by path, a list's entries as ``.``. + + Raises: + DFloatIntegrationError: A table is missing or is neither an array nor a list of arrays (mflux moved it). + """ + tables: dict[str, mx.array] = {} + for path in COMPUTED_TABLES: + owner_path, _, attr = path.rpartition(".") + owner = transformer + for part in owner_path.split("."): + owner = getattr(owner, part, None) + value = getattr( + owner, attr, None + ) # nn.Module keeps arrays and lists as its dict items, not in __dict__ + items = value if isinstance(value, list) else [value] + if value is None or not items or not all(isinstance(t, mx.array) for t in items): + raise DFloatIntegrationError(f"{path}: not a computed table on this mflux transformer") + if isinstance(value, list): + tables.update({f"{path}.{i}": t for i, t in enumerate(value)}) + else: + tables[path] = value + return tables + + +# The compiled step functions mflux 0.23 ``Qwen21Transformer.__call__`` builds on first use, and the plain methods +# they compile: ``_step_fn`` the joint pass (padded prompts, or ``use_text_cache`` off), ``_image_step_fn`` the +# image-token pass behind the text-prefix cache (the default). +STEP_FNS: tuple[tuple[str, str], ...] = ( + ("_step_fn", "_forward"), + ("_image_step_fn", "_image_forward"), +) + + +def _plain(transformer: Any, fn: Any, method: str) -> bool: + """Whether ``fn`` is unset or ``transformer``'s own plain ``method``. + + ``mx.compile`` returns ``mlx.gc_func``, a + ``FunctionType`` subclass, so only an exact type check tells a compiled function from a bound method. """ - require_mflux() from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer - fn = transformer._step_fn return fn is None or ( type(fn) is types.MethodType - and fn.__func__ is Qwen21Transformer._forward + and fn.__func__ is getattr(Qwen21Transformer, method) and fn.__self__ is transformer ) -class EagerForward: - """Compose in front of ``Qwen21Transformer``: its forward pass runs as the plain method, never compiled.""" +def is_eager(transformer: Any) -> bool: + """Whether both step functions are unset or the transformer's own plain, uncompiled methods (imports mflux).""" + require_mflux() + return all(_plain(transformer, getattr(transformer, attr), method) for attr, method in STEP_FNS) - _step_fn: Any - _forward: Any - def __call__(self, *args: Any, **kwargs: Any) -> Any: - """Set ``_step_fn`` to the plain ``_forward`` when unset, then run mflux's own call. +def refuse_compiled(transformer: Any) -> None: + """Refuse a step function that is set and is not the transformer's own plain method (imports mflux). - Raises: - DFloatIntegrationError: ``_step_fn`` is set to anything else (a compiled forward pass). - """ - if self._step_fn is None: - self._step_fn = self._forward - elif not is_eager(self): - fn = self._step_fn + Raises: + DFloatIntegrationError: A step function is set to anything else (a compiled forward pass). + """ + require_mflux() + for attr, method in STEP_FNS: + fn = getattr(transformer, attr) + if not _plain(transformer, fn, method): raise DFloatIntegrationError( - f"Qwen21Transformer._step_fn is {type(fn).__module__}.{type(fn).__qualname__}: mflux compiled the " + f"Qwen21Transformer.{attr} is {type(fn).__module__}.{type(fn).__qualname__}: mflux compiled the " "forward pass, and the block seam cannot evaluate inside mx.compile" ) + + +class EagerForward: + """Compose in front of ``Qwen21Transformer``: both forward passes run as the plain methods, never compiled.""" + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Refuse a compiled step function, set each unset one to its plain method, then run mflux's own call. + + Raises: + DFloatIntegrationError: A step function is set to anything else (a compiled forward pass). + """ + refuse_compiled(self) + for attr, method in STEP_FNS: + if getattr(self, attr) is None: + setattr(self, attr, getattr(self, method)) return super().__call__(*args, **kwargs) # type: ignore[misc] +class PackPerBlock: + """Compose in front of ``Qwen21Transformer``: ``inference_projections()`` packs per block, at each block's decode. + + mflux 0.23 packs every block's Q/K/V and gate/up once on entry; here the blocks hold placeholders, so entry only + turns the blocks' packing on (when ``use_packed_projections`` is set, as mflux reads it on entry) and each seamed + block packs its decoded weights with mflux's own ``_pack_linears`` for its call. Entry and exit clear what mflux + clears (packed matrices, step functions, text caches). + + On mflux 0.23 ``_pack_linears`` evaluates the packed matrices synchronously inside each block, so the ``depth2`` + policy's look-ahead decode does not overlap the block before it: ``depth2`` runs as ``per-block`` here (not + measured). + """ + + use_packed_projections: bool + _clear_inference_projections: Callable[[], None] + + @property + def pack_switch(self) -> PackSwitch: + """The switch this transformer's blocks read (created on first use).""" + switch = vars(self).get(_SWITCH) + if switch is None: + switch = PackSwitch() + vars(self)[_SWITCH] = switch + return switch + + @contextmanager + def inference_projections(self) -> Iterator[None]: + """The inference-only window of mflux 0.23, with the packing deferred to each block's decode. + + Raises: + DFloatIntegrationError: A step function is compiled on entry, before mflux's clear would hide it. + """ + refuse_compiled(self) + self._clear_inference_projections() + self.pack_switch.on = bool(self.use_packed_projections) + try: + yield + finally: + self.pack_switch.on = False + self._clear_inference_projections() + + +def mflux_packer() -> Packer: + """The packer of mflux 0.23 itself, ``Qwen21Transformer._pack_linears`` (imports mflux).""" + require_mflux() + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + pack: Packer = Qwen21Transformer._pack_linears + return pack + + @cache def seam_transformer_class() -> type: - """``StepSeam`` and ``EagerForward`` composed in front of mflux's ``Qwen21Transformer`` (imports mflux). + """``StepSeam``, ``EagerForward`` and ``PackPerBlock`` in front of mflux's ``Qwen21Transformer`` (imports mflux). Raises: DFloatDependencyError: The optional ``mflux`` extra is not installed. @@ -106,7 +332,9 @@ def seam_transformer_class() -> type: require_mflux() from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer - return type("SeamQwen21Transformer", (StepSeam, EagerForward, Qwen21Transformer), {}) + return type( + "SeamQwen21Transformer", (StepSeam, EagerForward, PackPerBlock, Qwen21Transformer), {} + ) def block_lists(transformer: Any) -> list[tuple[str, Sequence[Any]]]: @@ -142,8 +370,8 @@ def build_transformer( (a BF16 base's tensors, via ``base_extras``); with ``nonblock_from_extras`` ``modulation.1`` loads from them as a plain weight instead of a placeholder (the BF16 side of an identity check). A checkpoint that stores ``modulation.1`` as an extra loads it as a plain weight either way. The parameters mflux computes at construction - (``COMPUTED_PARAMS``) are left out of the coverage check and keep their values. The active memory the build added - is asserted. + (``COMPUTED_TABLES``) are not parameters, so the build never loads them. The active memory the build added is + asserted. Raises: DFloatFormatError: The groups are not Qwen-Image 2.1's; the checkpoint's block count differs from the @@ -166,7 +394,7 @@ def build_transformer( tf = make(**{**kwargs, "num_layers": n}) lists = block_lists(tf) shapes = install_placeholders(lists, names) - seam_blocks(lists, tf.seam_cell) + seam_qwen21_blocks(lists, tf.seam_cell, tf.pack_switch, pack=mflux_packer()) nonblock_matrices = ( [] if nonblock_from_extras @@ -176,7 +404,7 @@ def build_transformer( matrix_paths = {f"{b}.{a}.weight" for b, per in shapes.items() for a in per} | set(nonblock) built = {KIND: n} plan = extras_plan(ckpt, names, counts=built, extras=extras) - params = set(dict(tree_flatten(tf.parameters()))) - COMPUTED_PARAMS + params = set(dict(tree_flatten(tf.parameters()))) check_extras_cover(params, (p for p, _f, _i in plan), matrix_paths) load_extras(tf, plan, names) mx.eval(PLACEHOLDER) @@ -202,13 +430,25 @@ def base_extras( __all__ = [ "COMPUTED_PARAMS", + "COMPUTED_TABLES", + "ENTRY_POINTS", "MAX_BUILD_ACTIVE_BYTES", + "STEP_FNS", "EagerForward", + "PackPerBlock", + "PackSwitch", + "Qwen21Binding", + "Qwen21BlockSeam", "Qwen21Build", "base_extras", "base_transformer_files_index", "block_lists", "build_transformer", + "computed_tables", "is_eager", + "mflux_packer", + "refuse_compiled", + "seam_qwen21_blocks", "seam_transformer_class", + "seamed_block_class", ] diff --git a/src/mlx_dfloat/mflux/zimage/memory.py b/src/mlx_dfloat/mflux/zimage/memory.py index 7ed1a8e..5cf169a 100644 --- a/src/mlx_dfloat/mflux/zimage/memory.py +++ b/src/mlx_dfloat/mflux/zimage/memory.py @@ -23,15 +23,22 @@ # VAE_TRANSIENT (see below). Re-measure before changing any. # OVERHEAD: the footprint MLX's counters do not see after the build phase (turbo 510_030_936, base 470_479_960). OVERHEAD_BYTES = 510_030_936 # 0.475 GiB -# VAE_TRANSIENT: calibrated on the two measured 1024² generate runs (2026-10-08, same versions, Z-Image base and -# Turbo; footprint peaks 14_115_020_240 and 14_074_273_664, both in the VAE phase) so the VAE phase reproduces the -# larger one: base peak minus the phase's other terms (8_369_046_600 + 6_088_832 + 19_660_800 + 167_666_902 + -# OVERHEAD_BYTES) = 5_042_526_170 (Turbo needs 4_996_540_342). It absorbs the float32 decoder's own activations -# and the footprint the OS has not yet taken back from the denoise phase when the decode starts: MLX's cache is -# cleared by then, yet at the VAE start the footprint sits about 2.2 GiB above MLX's active memory plus OVERHEAD -# (turbo 11_254_325_096 - 8_392_187_384 - 510_030_936; base 11_312_586_144 - 8_389_484_140 - 510_030_936). The -# one-step value it replaces (VAE-phase peak minus the footprint after denoise, 3_982_295_040) left that part out. -VAE_TRANSIENT_BYTES = 5_042_526_170 # 4.696 GiB +# VAE_TRANSIENT, by the overall-peak rule over every committed 1024² sample with the set resident through the VAE +# decode (each peak reached in the VAE phase): the highest peak minus the VAE phase's other terms of the same run. +# The binding sample is the MEASURED Z-Image run on mflux 0.23 (2026-10-24, git cbed93b, mlx 0.32.2, 50 steps, +# guidance 4; bench/results/tiers/z-image-1024.json watched_peak_bytes): 15_462_030_800 - (8_369_046_600 + +# 6_088_832 + 19_660_800 + 167_666_902 + OVERHEAD_BYTES) = 6_389_536_730. Margins over the other samples: the 0.23 +# identity df11 side 376_356_864 (15_085_673_936); on mflux 0.20, Z-Image's MEASURED run 1_347_010_560 +# (14_115_020_240) and identity 1_395_834_640 (14_066_196_160), and on Turbo's 5_239_252 B larger set its MEASURED +# run 1_392_996_388 (14_074_273_664) and five repeats 158_937_556 to 756_412_836 (15_308_332_496 to +# 14_710_857_216). The 0.20 value, 5_042_526_170 from the 0.20 MEASURED run alone, is under both 0.23 runs and four +# of the Turbo repeats. Z-Image-Turbo shares the term (its 0.23 peak is not measured); a larger term is conservative +# for it. The term absorbs the float32 decoder's own activations and the footprint the OS has not yet taken back +# from the denoise phase when the decode starts: MLX's cache is cleared by then, yet at the VAE start the footprint +# sits about 2.2 GiB above MLX's active memory plus OVERHEAD (0.20 turbo 11_254_325_096 - 8_392_187_384 - +# 510_030_936; base 11_312_586_144 - 8_389_484_140 - 510_030_936), and the run-to-run spread of that footprint is +# about 1 GiB (Turbo's repeats). +VAE_TRANSIENT_BYTES = 6_389_536_730 # 5.951 GiB # DENOISE_ACTIVATION: denoise footprint peak minus (compressed + extras + non-block + in-flight + cache limit + # overhead) of the same unit: turbo 12_781_477_784 - 11_583_577_874 = 1_197_899_910; base 13_061_987_888 - # 11_538_787_646 = 1_523_200_242, at 4096 image + 512 text tokens. diff --git a/src/mlx_dfloat/mflux/zimage/model.py b/src/mlx_dfloat/mflux/zimage/model.py index b0820d4..9f0a062 100644 --- a/src/mlx_dfloat/mflux/zimage/model.py +++ b/src/mlx_dfloat/mflux/zimage/model.py @@ -117,6 +117,8 @@ def __init__( lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, bake_lora: bool = True, + float32: bool = False, + compute_precision: Any = None, ) -> None: """Resolve the checkpoint and the base repository and build the model. @@ -125,9 +127,9 @@ def __init__( estimate and VAE decision. Raises: - DFloatUnsupportedError: ``quantize``, ``lora_paths``, ``lora_scales`` or ``bake_lora=False`` - (accepted only to refuse), an eval policy other than ``"per-block"`` / ``"depth2"``, a - ControlNet model, or an unknown ``model``. + DFloatUnsupportedError: ``quantize``, ``lora_paths``, ``lora_scales``, ``bake_lora=False``, + ``float32=True`` or a ``compute_precision`` (accepted only to refuse), an eval policy other than + ``"per-block"`` / ``"depth2"``, a ControlNet model, or an unknown ``model``. DFloatBackendError: The GPU decoder fails its check on the two packaged self-check groups (raised before anything is resolved or loaded), or Metal is unavailable. DFloatFormatError: The checkpoint is not a Z-Image DF11 checkpoint, or the base lacks a component. @@ -147,6 +149,7 @@ def __init__( _refuse("lora_scales", "LoRA of any kind") if not bake_lora: _refuse("bake_lora", "LoRA of any kind") + _pipeline.refuse_precision_args(float32=float32, compute_precision=compute_precision) _check_eval_policy(eval_policy) if "controlnet" in model: raise DFloatUnsupportedError( @@ -167,7 +170,10 @@ def __init__( df11 = zinit.resolve(df11_path or df11_default, patterns=zinit.DF11_PATTERNS, revision=pin) ckpt = open_checkpoint(df11.root) check_zimage_groups(ckpt) - base = zinit.resolve(base_path or base_default, patterns=zinit.BASE_PATTERNS) + base_pin = None if base_path else registry.MODELS[model].base_revision + base = zinit.resolve( + base_path or base_default, patterns=zinit.BASE_PATTERNS, revision=base_pin + ) components = zinit.load_base(base.root) build = build_transformer(ckpt) self._assemble( @@ -211,11 +217,16 @@ def _assemble( budget_bytes: int | None = None, decode: Callable[[MxGroup], DecodeResult] | None = None, ) -> None: + from mflux.models.common.compute_precision import ComputePrecision from mflux.models.z_image.z_image_initializer import ZImageInitializer _check_eval_policy(eval_policy) nn.Module.__init__(self) # type: ignore[attr-defined] # not ZImage.__init__: mflux's initializer runs over the whole base ZImageInitializer._init_config(self, model_config) # model_config, callbacks, tiling_config + # mflux's defaults (ZImageInitializer.init sets them): the bfloat16 stream, no compute precision; generate_image + # records both. The transformer's own float32 flag starts False in ZImageTransformer.__init__. + self.float32 = False + self.compute_precision = ComputePrecision() self.vae = components.vae self.text_encoder = components.text_encoder self.tokenizers = components.tokenizers @@ -437,7 +448,7 @@ def cfg_prompts( ) -> tuple[str, ...]: """The prompts a call encodes: one, or the prompt and its negative when classifier-free guidance runs. - mflux 0.20's rules (``ZImage.generate_image`` and ``_encode_prompts``): a model without guidance support runs at guidance 0, guidance + mflux 0.23's rules (``ZImage.generate_image`` and ``_encode_prompts``): a model without guidance support runs at guidance 0, guidance at or below 1 encodes the prompt alone, and a blank negative prompt is a single space. """ effective = ( diff --git a/src/mlx_dfloat/mflux/zimage/names.py b/src/mlx_dfloat/mflux/zimage/names.py index a4fbe04..3286903 100644 --- a/src/mlx_dfloat/mflux/zimage/names.py +++ b/src/mlx_dfloat/mflux/zimage/names.py @@ -1,6 +1,6 @@ """Z-Image naming: the block kinds in run order, the compressed sub-paths, the map derived from mflux, the group check. -Block kinds and run order: mflux 0.20.0 ``ZImageTransformer.__call__``. Compressed +Block kinds and run order: mflux 0.23.0 ``ZImageTransformer.__call__``. Compressed sub-paths: the ``pattern_dict`` of both mingyi456 DF11 checkpoints (Z-Image and Z-Image-Turbo, identical). """ diff --git a/tests/_qwen21_tiny.py b/tests/_qwen21_tiny.py index 2ed7cbd..da47803 100644 --- a/tests/_qwen21_tiny.py +++ b/tests/_qwen21_tiny.py @@ -1,6 +1,9 @@ """Tiny Qwen-Image 2.1 transformer helpers for the mflux lane. Every mflux import happens inside a function.""" +import os import zlib +from collections.abc import Iterator +from contextlib import contextmanager from types import SimpleNamespace import mlx.core as mx @@ -15,10 +18,12 @@ "TINY", "TINY_ROW_SPLITS", "TINY_SEED", + "WIDE", "StubTextEncoder", "StubTokenizer", "StubVAE", "fake_model", + "mflux_uncompiled", "stub_components", "tiny_config", "tiny_groups", @@ -27,6 +32,34 @@ "write_tiny_checkpoint", ] + +# How many mflux_uncompiled blocks are open: MLX has no getter for its compile state, so an inner block must know an +# outer one switched compile off. +_UNCOMPILED_DEPTH = [0] + + +@contextmanager +def mflux_uncompiled() -> Iterator[None]: + """``mx.compile`` off for the block: stock mflux's own graph run uncompiled, the parity reference. + + mflux 0.23 compiles ``_step_fn`` / ``_image_step_fn`` on every call path (qwen21_transformer.py:96-98, 111-113), + and the compiled bits differ from the same graph run plain; our seam runs the plain graph. + + On exit the previous state comes back: compile stays off when it was off on entry. MLX exposes only + ``mx.enable_compile`` / ``mx.disable_compile``, no getter, so "off on entry" means an enclosing block of this + helper, or ``MLX_DISABLE_COMPILE`` set (MLX's documented switch). + """ + was_off = _UNCOMPILED_DEPTH[0] > 0 or "MLX_DISABLE_COMPILE" in os.environ + mx.disable_compile() + _UNCOMPILED_DEPTH[0] += 1 + try: + yield + finally: + _UNCOMPILED_DEPTH[0] -= 1 + if not was_off: + mx.enable_compile() + + # Head dim 16 = the sum of the three RoPE axes (4 + 6 + 6, each even: mflux 0.20.0 qwen21_rope.py:23-26 builds one # table per axis at half its width); inner dim 2 x 16 = 32. in_channels stays the default 64: Qwen21LatentCreator # always makes 64-channel latents (latent_creator/qwen21_latent_creator.py:13-14), and the model-level tests run @@ -71,16 +104,21 @@ def tiny_inputs(seed=0): } -def tiny_seamed(n_layers=2): - """A seamed mflux ``Qwen21Transformer`` at ``TINY`` size with placeholders installed; returns (tf, shapes).""" - from mlx_dfloat.integrate.blockseam import seam_blocks +def tiny_seamed(n_layers=2, dims=None): + """A seamed mflux ``Qwen21Transformer`` at ``dims`` (``TINY`` by default) with placeholders; returns (tf, shapes).""" + from mlx_dfloat.integrate.placeholders import install_placeholders from mlx_dfloat.mflux.qwen21.names import qwen21_name_map - from mlx_dfloat.mflux.qwen21.transformer import block_lists, seam_transformer_class + from mlx_dfloat.mflux.qwen21.transformer import ( + block_lists, + mflux_packer, + seam_qwen21_blocks, + seam_transformer_class, + ) - tf = seam_transformer_class()(**{**TINY, "num_layers": n_layers}) + tf = seam_transformer_class()(**{**(TINY if dims is None else dims), "num_layers": n_layers}) shapes = install_placeholders(block_lists(tf), qwen21_name_map()) - seam_blocks(block_lists(tf), tf.seam_cell) + seam_qwen21_blocks(block_lists(tf), tf.seam_cell, tf.pack_switch, pack=mflux_packer()) return tf, shapes @@ -128,7 +166,7 @@ def write_tiny_checkpoint( from mlx_dfloat.integrate.placeholders import install_placeholders from mlx_dfloat.mflux.qwen21.names import qwen21_name_map - from mlx_dfloat.mflux.qwen21.transformer import COMPUTED_PARAMS, block_lists + from mlx_dfloat.mflux.qwen21.transformer import block_lists dims = {**TINY, "num_layers": n_layers} if inner_override is not None: @@ -143,7 +181,9 @@ def write_tiny_checkpoint( gate_up = np.concatenate([mats["img_mlp.gate_layer"], mats["img_mlp.proj"]]) stored = {**mats, "img_mlp.gate_up": gate_up} groups[block] = [stored[sub] for sub in STORED_SUBS] - taken = {f"{b}.{a}.weight" for b, per in shapes.items() for a in per} | set(COMPUTED_PARAMS) + # mflux 0.23 keeps its computed tables outside the parameters (``COMPUTED_TABLES``), so every parameter that is + # not a block matrix is an extra the checkpoint stores. + taken = {f"{b}.{a}.weight" for b, per in shapes.items() for a in per} patterns = {r"transformer_blocks\.\d+": STORED_SUBS} if modulation_as_group: matrices["modulation.1"] = {"weight": random_bf16(rng, modulation_shape)} @@ -282,3 +322,17 @@ def fake_model(tmp_path, monkeypatch, *, tokenizer=None, **overrides): } parts.update(overrides) return model_module.DFloatQwenImage21._from_parts(**parts) + + +# Wide enough for the two mflux 0.23 paths TINY cannot reach: head dim 64 turns the fused QK-norm/RoPE kernel on +# (qwen21_fused_kernels.py:107-114: a multiple of 64, at most 512; the axes 16 + 24 + 24 = 64), and inner dim +# 16 x 64 = 1024 with a few text tokens is a shape where the packed Q/K/V matmul's bits differ from the separate ones +# (tests/test_qwen21_packing.py). mlp_ratio 1 keeps a block at 7 x 1024^2 weights. +WIDE = { + "num_layers": 1, + "attention_head_dim": 64, + "num_attention_heads": 16, + "context_in_dim": 32, + "mlp_ratio": 1, + "axes_dims_rope": (16, 24, 24), +} diff --git a/tests/test_cli.py b/tests/test_cli.py index cb398ef..be07545 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2,11 +2,14 @@ import importlib.util import json import os +import platform import subprocess import sys +from importlib import metadata from pathlib import Path from types import SimpleNamespace +import mlx.core as mx import pytest from mlx_dfloat import DFloatDependencyError, DFloatResourceError, DFloatUnsupportedError @@ -239,6 +242,7 @@ def test_generate_builds_the_model_from_the_flags_writes_the_image_and_the_repor "tier_gb": 32, "label": "MEASURED", "limits": {"memory": 111, "cache": 222, "wired": 333}, # what read_limits found + **_HOST_VERSIONS, # as the generate report records them } assert report["footprint_peak_bytes"] == 10**15 assert ( @@ -258,6 +262,7 @@ def test_the_watchdog_context_carries_the_models_default_steps(tmp_path): "tier_gb": 32, "label": "MEASURED", "limits": {"memory": 111, "cache": 222, "wired": 333}, # what read_limits found + **_HOST_VERSIONS, # as the generate report records them } @@ -1568,3 +1573,17 @@ def test_the_help_names_the_krea_defaults_and_the_cards_recipe(): "classifier-free guidance, two transformer calls per step, for any value other than 1.0" ) in text assert "Krea 2: er_sde (default), euler" in text + + +def _installed_mflux() -> str | None: + try: + return metadata.version("mflux") + except metadata.PackageNotFoundError: + return None + + +# The versions and macOS the watchdog context records, read here from the installed packages and the OS. +_HOST_VERSIONS = { + "versions": {"mlx": mx.__version__, "mflux": _installed_mflux()}, + "macos": platform.mac_ver()[0], +} diff --git a/tests/test_flux1_model.py b/tests/test_flux1_model.py index 546e40d..36658bf 100644 --- a/tests/test_flux1_model.py +++ b/tests/test_flux1_model.py @@ -172,6 +172,45 @@ def faulty(group, **kwargs): DFloatFlux1("schnell") +def test_the_default_bases_resolve_at_their_pins_and_a_user_base_unpinned(tmp_path, monkeypatch): + # Bug caught: the base pin not reaching the resolver (T5, CLIP, the VAE and the tokenizers following whatever lands + # on main instead of the snapshot the recorded runs used), a pin taken from another model's entry, or a pin applied + # to a user's own --base. + from mlx_dfloat import _metal_decode + from mlx_dfloat.mflux.flux1 import model as model_module + from mlx_dfloat.mflux.flux1.model import DFloatFlux1 + + class StopError(Exception): + pass + + base_calls = [] + + def stop_at_base(spec, *, patterns, revision=None): + if patterns == base_init.DF11_PATTERNS: + return ResolvedRepo(root=tmp_path, repo_id=spec, revision=revision) + base_calls.append((spec, revision)) + raise StopError + + monkeypatch.setattr(_metal_decode, "ensure_canary", lambda **kwargs: None) + monkeypatch.setattr(model_module, "open_checkpoint", lambda root: None) + monkeypatch.setattr(model_module, "check_flux_groups", lambda ckpt: None) + monkeypatch.setattr(base_init, "resolve", stop_at_base) + for kwargs in ( + {"model": "schnell"}, + {"model": "dev"}, + {"model": "krea-dev"}, + {"model": "dev", "base_path": "me/base"}, + ): + with pytest.raises(StopError): + DFloatFlux1(**kwargs) + assert base_calls == [ + ("black-forest-labs/FLUX.1-schnell", "741f7c3ce8b383c54771c7003378a50191e9efe9"), + ("black-forest-labs/FLUX.1-dev", "3de623fc3c33e44ffbe2bad470d0f45bccf2eb21"), + ("black-forest-labs/FLUX.1-Krea-dev", "8162a9c7b05a641be098422bf2fcf335615c2f28"), + ("me/base", None), + ] + + def test_an_unknown_model_and_the_none_policy_are_refused(monkeypatch, tmp_path): # Bug caught: mflux's from_name accepting "dev-fill" (another class upstream) and this path # building a FLUX.1 model for it; or policy "none" (every decode resident) accepted, whether diff --git a/tests/test_flux1_names.py b/tests/test_flux1_names.py index 5297348..0a99415 100644 --- a/tests/test_flux1_names.py +++ b/tests/test_flux1_names.py @@ -18,7 +18,7 @@ from mlx_dfloat.errors import DFloatFormatError from mlx_dfloat.format import open_checkpoint from mlx_dfloat.integrate.placeholders import install_placeholders -from mlx_dfloat.mflux.flux1.names import DROPPED_EXTRAS, check_flux_groups +from mlx_dfloat.mflux.flux1.names import check_flux_groups def _group(name: str) -> SimpleNamespace: @@ -91,10 +91,20 @@ def test_flux_groups_refuses_a_group_with_the_wrong_matrix_count(): check_flux_groups(bad) -def test_dropped_extras_is_the_one_bias_mflux_has_no_parameter_for(): - # Bug caught: DROPPED_EXTRAS growing to include a bias mflux actually loads, silently dropping - # a real parameter from the build. - assert frozenset({"norm_out.linear.bias"}) == DROPPED_EXTRAS +@pytest.mark.mflux +def test_mflux_builds_the_output_norm_bias_the_checkpoint_carries(): + # Bug caught: an mflux release that builds FLUX.1's norm_out.linear without a bias again (as 0.16 to 0.21 + # did): the checkpoint's bias would then have no parameter and every build would refuse it. mflux 0.23.0: + # AdaLayerNormContinuous(..., bias=True) at models/flux/model/flux_transformer/ada_layer_norm_continuous.py:8. + from mflux.models.common.config.model_config import ModelConfig + from mflux.models.flux.model.flux_transformer.transformer import Transformer + from mlx.utils import tree_flatten + + tf = Transformer( + ModelConfig.schnell(), num_transformer_blocks=0, num_single_transformer_blocks=0 + ) + params = dict(tree_flatten(tf.parameters())) + assert tuple(params["norm_out.linear.bias"].shape) == (2 * 3072,) @pytest.mark.parametrize("pattern", [r"transformer_blocks.\d+", r"transformer_blocks\.\d+"]) @@ -127,7 +137,7 @@ def test_flux_name_map_derived_from_mflux_equals_the_measured_table(): # Bug caught: the lane silently running a different mflux than the one the table was measured on # (mflux has no __version__; the distribution metadata is the source). - assert importlib.metadata.version("mflux").startswith("0.20.") + assert importlib.metadata.version("mflux").startswith("0.23.") derived = flux_name_map() assert derived.kinds == FLUX_TABLE.kinds for kind in FLUX_TABLE.kinds: diff --git a/tests/test_flux1_transformer.py b/tests/test_flux1_transformer.py index 999a59b..2c250a2 100644 --- a/tests/test_flux1_transformer.py +++ b/tests/test_flux1_transformer.py @@ -5,6 +5,7 @@ from types import SimpleNamespace import mlx.core as mx +import mlx.nn as nn import numpy as np import pytest from mlx.utils import tree_flatten @@ -92,8 +93,8 @@ def build(model_config, *, num_transformer_blocks, num_single_transformer_blocks return build -def _write_flux(tmp_path, *, n_double=2, n_single=1): - """A DF11 checkpoint for the fake transformer: one group per block and every non-matrix extra.""" +def _write_flux(tmp_path, *, n_double=2, n_single=1, more=None): + """A DF11 checkpoint for the fake transformer: one group per block and every non-matrix extra (plus ``more``).""" rng = np.random.default_rng(51) fake = FakeTransformer(Recorder(), n_double=n_double, n_single=n_single) shapes = install_placeholders(block_lists(fake), FLUX_TABLE) @@ -104,6 +105,7 @@ def _write_flux(tmp_path, *, n_double=2, n_single=1): for name, p in tree_flatten(fake.parameters()) if p.size > 0 } + extras.update(more or {}) write_flux_checkpoint(tmp_path / "ckpt", shapes, rng, extras=extras) return open_checkpoint(tmp_path / "ckpt"), extras @@ -130,6 +132,45 @@ def test_build_transformer_loads_every_extra_and_leaves_every_matrix_a_placehold ) +class _OutputNorm(nn.Module): + """mflux 0.23's FLUX.1 ``norm_out`` shape: ``AdaLayerNormContinuous`` with ``linear`` built with a bias.""" + + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(4, 8, bias=True) + + +def _fake_seam_class_with_output_norm(): + plain = _fake_seam_class() + + def build(model_config, *, num_transformer_blocks, num_single_transformer_blocks): + tf = plain( + model_config, + num_transformer_blocks=num_transformer_blocks, + num_single_transformer_blocks=num_single_transformer_blocks, + ) + tf.norm_out = _OutputNorm() + return tf + + return build + + +def test_build_transformer_loads_the_output_norm_bias_from_the_checkpoint(tmp_path, monkeypatch): + # Bug caught: norm_out.linear.bias still left out of the extras plan (what mflux 0.20, which built that layer + # without a bias, needed) while mflux 0.23 builds and loads it: the build refuses the uncovered parameter, or, + # with the coverage check loosened, the output norm runs on mflux's zero bias instead of the checkpoint's. + monkeypatch.setattr(transformer, "seam_transformer_class", _fake_seam_class_with_output_norm) + rng = np.random.default_rng(839) + bias = np.asarray(rng.integers(0x3C00, 0x3F80, size=(8,)), dtype=np.uint16) + weight = np.asarray(rng.integers(0x3C00, 0x3F80, size=(8, 4)), dtype=np.uint16) + ckpt, _extras = _write_flux( + tmp_path, more={"norm_out.linear.weight": weight, "norm_out.linear.bias": bias} + ) + tf, _shapes = build_transformer(None, ckpt, name_map=FLUX_TABLE) + assert np.array_equal(np.array(tf.norm_out.linear.bias.view(mx.uint16)), bias) + assert np.array_equal(np.array(tf.norm_out.linear.weight.view(mx.uint16)), weight) + + def test_build_transformer_names_the_active_memory_the_extras_added(tmp_path, monkeypatch): # Bug caught: the build's active-memory guard dropped (a resident block matrix or a float32 # extras copy would go unnoticed), or its message blaming a block matrix, which a lazy init diff --git a/tests/test_flux2_model.py b/tests/test_flux2_model.py index c146c41..4499d56 100644 --- a/tests/test_flux2_model.py +++ b/tests/test_flux2_model.py @@ -49,6 +49,8 @@ def test_every_attribute_upstreams_generate_image_reads_exists_on_the_model(tmp_ assert [n for n in sorted(names) if not hasattr(model, n)] == [] assert isinstance(model, Flux2Klein) assert (model.bits, model.lora_paths, model.lora_scales) == (None, [], []) + # mflux 0.23 records the compute precision in the image metadata; ours is mflux's default (none). + assert model.compute_precision.dtype is None @pytest.mark.parametrize( @@ -58,6 +60,7 @@ def test_every_attribute_upstreams_generate_image_reads_exists_on_the_model(tmp_ ({"lora_paths": ["a"]}, "lora_paths"), ({"lora_scales": [1.0]}, "lora_scales"), ({"bake_lora": False}, "bake_lora"), + ({"compute_precision": mx.float16}, "compute_precision"), ({"eval_policy": "none"}, "none"), ({"model": "flux2-klein-9b-kv"}, "KV"), ], @@ -121,6 +124,47 @@ def fake_resolve(spec, *, patterns, revision=None): ] +def test_the_default_bases_resolve_at_their_pins_and_a_user_base_unpinned(tmp_path, monkeypatch): + # Bug caught: the base pin not reaching the resolver (the Qwen3 encoder, VAE and tokenizer following whatever lands + # on main instead of the snapshot the recorded runs used), a pin taken from another model's entry, or a pin applied + # to a user's own --base. + from mlx_dfloat import _metal_decode + from mlx_dfloat.mflux._hub import ResolvedRepo + from mlx_dfloat.mflux.flux2 import model as model_module + from mlx_dfloat.mflux.flux2.model import DFloatFlux2Klein + + class StopError(Exception): + pass + + base_calls = [] + + def stop_at_base(spec, *, patterns, revision=None): + if patterns == finit.DF11_PATTERNS: + return ResolvedRepo(root=tmp_path, repo_id=spec, revision=revision) + base_calls.append((spec, revision)) + raise StopError + + monkeypatch.setattr(_metal_decode, "ensure_canary", lambda **kwargs: None) + monkeypatch.setattr(model_module, "open_checkpoint", lambda root: None) + monkeypatch.setattr(finit, "resolve", stop_at_base) + for kwargs in ( + {"model": "flux2-klein-base-4b"}, + {"model": "flux2-klein-4b"}, + {"model": "flux2-klein-base-9b"}, + {"model": "flux2-klein-9b"}, + {"model": "flux2-klein-9b", "base_path": "me/base"}, + ): + with pytest.raises(StopError): + DFloatFlux2Klein(**kwargs) + assert base_calls == [ + ("black-forest-labs/FLUX.2-klein-base-4B", "a3b4f4849157f664bdbc776fd7453c2783562f4d"), + ("black-forest-labs/FLUX.2-klein-4B", "e7b7dc27f91deacad38e78976d1f2b499d76a294"), + ("black-forest-labs/FLUX.2-klein-base-9B", "32773329fbe7e81a90ef971740e8ba4b0364ecf3"), + ("black-forest-labs/FLUX.2-klein-9B", "92196c8e11f7b6cf2b7493e037d8c5345c559216"), + ("me/base", None), + ] + + def test_the_constructor_builds_from_local_dirs_with_the_models_overrides_and_the_used_encoder_bytes( tmp_path, monkeypatch ): diff --git a/tests/test_krea2_transformer.py b/tests/test_krea2_transformer.py index 7a37326..f9a32d6 100644 --- a/tests/test_krea2_transformer.py +++ b/tests/test_krea2_transformer.py @@ -172,7 +172,7 @@ def test_the_swap_keeps_parameter_paths_and_nothing_is_computed(): else: targets.add(t.to_pattern) assert frozenset(params - targets) == COMPUTED_PARAMS - assert len(COMPUTED_PARAMS) == 0 # mflux 0.20.0: every parameter is a mapping target + assert len(COMPUTED_PARAMS) == 0 # mflux 0.23.0: every parameter is a mapping target @pytest.mark.mflux diff --git a/tests/test_mflux_families.py b/tests/test_mflux_families.py index cc95c4d..5212c1b 100644 --- a/tests/test_mflux_families.py +++ b/tests/test_mflux_families.py @@ -240,21 +240,94 @@ def test_the_qwen_entry_names_its_published_repos_and_pinned_checkpoint(): ) -def test_pinned_bases_are_the_snapshots_the_recorded_runs_used_and_no_other_base_is(): - # Bug caught: a base (text encoder, VAE, tokenizer) following whatever lands on main instead of the snapshot its - # recorded runs used (Qwen-Image 2.1: d26bb61, bench/results/tiers/qwen-image-2.1-1024.json; ERNIE-Image and - # ERNIE-Image-Turbo: the snapshots their parity and calibration runs read, 2026-10-08), or a pin set on a family whose - # base no recorded run names. - assert { - n: e.base_revision for n, e in families.MODELS.items() if e.base_revision is not None - } == { - "qwen-image-2.1": "d26bb61231c349cf6b7896fa83353113880e1ba3", - "ernie-image": "5346b31d68c9c23758ba56ef8be5e9dc174c7f99", - "ernie-image-turbo": "bc68c81e2a1730a394d5fc9fae70713dee940140", +@pytest.mark.parametrize( + ("family", "pins"), + [ + # FLUX.1: the base snapshots of the identity and MEASURED tier records (bench/results/identity/ + # flux1-schnell-1024, bench/results/tiers/{schnell,dev,krea-dev}-1024.json). + ( + "flux1", + { + "schnell": ( + "black-forest-labs/FLUX.1-schnell", + "741f7c3ce8b383c54771c7003378a50191e9efe9", + ), + "dev": ("black-forest-labs/FLUX.1-dev", "3de623fc3c33e44ffbe2bad470d0f45bccf2eb21"), + "krea-dev": ( + "black-forest-labs/FLUX.1-Krea-dev", + "8162a9c7b05a641be098422bf2fcf335615c2f28", + ), + }, + ), + # Z-Image: the identity and tier records (bench/results/identity/z-image-1024, tiers/z-image-1024.json); + # Z-Image-Turbo: the sampled parity record (bench/results/parity/z-image-turbo), the one record of that base. + ( + "zimage", + { + "z-image": ("Tongyi-MAI/Z-Image", "04cc4abb7c5069926f75c9bfde9ef43d49423021"), + "z-image-turbo": ( + "Tongyi-MAI/Z-Image-Turbo", + "f332072aa78be7aecdf3ee76d5c247082da564a6", + ), + }, + ), + # FLUX.2 Klein: the base models' identity and tier records (tiers/flux2-klein-base-{4b,9b}-1024.json); the + # distilled models' sampled parity records (bench/results/parity/flux2-klein-{4b,9b}). + ( + "flux2", + { + "flux2-klein-base-4b": ( + "black-forest-labs/FLUX.2-klein-base-4B", + "a3b4f4849157f664bdbc776fd7453c2783562f4d", + ), + "flux2-klein-4b": ( + "black-forest-labs/FLUX.2-klein-4B", + "e7b7dc27f91deacad38e78976d1f2b499d76a294", + ), + "flux2-klein-base-9b": ( + "black-forest-labs/FLUX.2-klein-base-9B", + "32773329fbe7e81a90ef971740e8ba4b0364ecf3", + ), + "flux2-klein-9b": ( + "black-forest-labs/FLUX.2-klein-9B", + "92196c8e11f7b6cf2b7493e037d8c5345c559216", + ), + }, + ), + # Qwen-Image 2.1: d26bb61, bench/results/tiers/qwen-image-2.1-1024.json. + ( + "qwen21", + {"qwen-image-2.1": ("Qwen/Qwen-Image-2.1", "d26bb61231c349cf6b7896fa83353113880e1ba3")}, + ), + # ERNIE-Image and ERNIE-Image-Turbo: the snapshots their parity and calibration runs read, 2026-10-08. + ( + "ernie", + { + "ernie-image": ("baidu/ERNIE-Image", "5346b31d68c9c23758ba56ef8be5e9dc174c7f99"), + "ernie-image-turbo": ( + "baidu/ERNIE-Image-Turbo", + "bc68c81e2a1730a394d5fc9fae70713dee940140", + ), + }, + ), # Krea 2: the gated base snapshots the range-read parity reads and every Krea run uses (2026-10-08). - "krea-2": "98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7", - "krea-2-raw": "6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603", - } + ( + "krea2", + { + "krea-2": ("krea/Krea-2-Turbo", "98e0fe118d17c9e3547fbb2e25acdbae2cadf7c7"), + "krea-2-raw": ("krea/Krea-2-Raw", "6b0ece7fffb640c5e3bcbe0a7f10f66b8e60a603"), + }, + ), + ], +) +def test_each_familys_default_bases_are_pinned_to_the_snapshots_the_recorded_runs_used( + family, pins +): + # Bug caught: a base (text encoders, VAE, tokenizer) following whatever lands on main instead of the snapshot its + # recorded runs used, so a re-upload changes encoder bytes with no error; or a model left unpinned. + assert { + n: (e.base_repo, e.base_revision) for n, e in families.MODELS.items() if e.family == family + } == pins def test_family_of_returns_the_spec_of_the_models_family(): diff --git a/tests/test_mflux_generate.py b/tests/test_mflux_generate.py new file mode 100644 index 0000000..57f172e --- /dev/null +++ b/tests/test_mflux_generate.py @@ -0,0 +1,87 @@ +"""`mlx-dfloat generate` against mflux 0.23's image save, which raises when the image file cannot be written. + +Offline: the command's own fakes (``tests.test_cli``) stand in for the model, the watchdog and the host. +""" + +import json +import platform + +import pytest +from tests.test_cli import _Image, _Model, _run + + +class _RaisingImage(_Image): + """mflux 0.23's ``GeneratedImage.save``: a failed image write propagates (``ImageUtil.save_image``).""" + + def save(self, path, export_json_metadata=False, overwrite=False): + self.log.append(("save", str(path), export_json_metadata, overwrite)) + raise PermissionError(13, "Permission denied", str(path)) + + +class _RaisingModel(_Model): + def generate_image(self, **kwargs): + self.log.append(("generate", kwargs)) + return _RaisingImage(self.log) + + +def test_a_save_that_raises_is_reported_as_the_image_not_written(tmp_path, capsys): + # Bug caught: the write failure mflux 0.23 now raises falling through to the generic handler, which prints a + # traceback instead of naming the file that was not written (exit 2 either way). + log = [] + out = tmp_path / "o.png" + code, _ = _run( + ["--prompt", "p", "--output", str(out), "--report", str(tmp_path / "r.json")], + log, + tmp_path, + factory=lambda **kw: _RaisingModel(log, **kw), + ) + err = capsys.readouterr().err + assert code == 2 + assert f"{out}: the image was not written" in err + assert "Traceback" not in err + report = json.loads((tmp_path / "r.json").read_text()) + assert report["exit_code"] == 2 + assert "Permission denied" in report["error"] + + +@pytest.mark.parametrize("model_class", [_Model, _RaisingModel]) +def test_the_report_records_the_macos_version_of_the_run(tmp_path, monkeypatch, model_class): + # Bug caught: a generate report with no OS field (the tier records could not say which macOS a MEASURED row ran + # on), or the field written only when the run succeeds. The value is the bench provenance's: platform.mac_ver()[0]. + monkeypatch.setattr(platform, "mac_ver", lambda: ("26.9.1", ("", "", ""), "arm64")) + log = [] + _run( + [ + "--prompt", + "p", + "--output", + str(tmp_path / "o.png"), + "--report", + str(tmp_path / "r.json"), + ], + log, + tmp_path, + factory=lambda **kw: model_class(log, **kw), + ) + assert json.loads((tmp_path / "r.json").read_text())["macos"] == "26.9.1" + + +def test_the_watchdogs_run_context_records_the_versions_and_macos_as_the_report_does(): + # Red when: the context generate hands its watchdog (all a stopped run's abort artifact keeps) lacks the mlx or + # mflux version or the macOS version, so the tier table cannot say what a stopped run ran on. + from importlib import metadata + + import mlx.core as mx + + from mlx_dfloat.mflux import generate as gen + + try: + mflux = metadata.version("mflux") + except metadata.PackageNotFoundError: + mflux = None + args = gen.build_parser().parse_args(["--prompt", "p"]) + context = gen._run_context( + args, tier_gb=16, label="CAPPED", limits={"memory": 1, "cache": 2, "wired": 3} + ) + assert context["versions"] == {"mlx": mx.__version__, "mflux": mflux} + assert context["macos"] == platform.mac_ver()[0] diff --git a/tests/test_mflux_hub.py b/tests/test_mflux_hub.py new file mode 100644 index 0000000..96faa57 --- /dev/null +++ b/tests/test_mflux_hub.py @@ -0,0 +1,146 @@ +"""``_hub.load_into`` against mflux 0.23's loader, which refuses a component whose required weights are missing. + +Offline: the two mflux modules ``load_into`` imports are replaced by fakes in ``sys.modules``. +""" + +import sys +import types + +import pytest + +from mlx_dfloat.errors import DFloatFormatError +from mlx_dfloat.mflux import _hub + + +class ModelConfigError(ValueError): + """mflux 0.23.0's class, as declared in ``mflux/utils/exceptions.py:44`` (a ``ValueError``).""" + + +MESSAGE = ( + "The vae weights in /base/vae do not fit this model: 2 required weights have no match " + "(expected names like encoder.conv_in.weight; found none under the names this component reads)." +) + + +def _install_loader(monkeypatch, load): + loader = types.ModuleType("mflux.models.common.weights.loading.weight_loader") + loader.WeightLoader = types.SimpleNamespace(load=load) # type: ignore[attr-defined] + applier = types.ModuleType("mflux.models.common.weights.loading.weight_applier") + applier.WeightApplier = types.SimpleNamespace( # type: ignore[attr-defined] + apply_and_quantize=lambda **kw: pytest.fail("applied after a refused load") + ) + monkeypatch.setitem(sys.modules, loader.__name__, loader) + monkeypatch.setitem(sys.modules, applier.__name__, applier) + + +def _definition(): + return types.SimpleNamespace( + get_components=lambda: [types.SimpleNamespace(name="vae")], + get_download_patterns=lambda: ["vae/*.safetensors"], + ) + + +def test_mfluxs_missing_weights_refusal_becomes_a_format_error_naming_the_base( + monkeypatch, tmp_path +): + # Bug caught: the required-weights refusal (mflux 0.23's ModelConfigError) escaping untranslated, so the + # command prints a traceback for a wrong base directory instead of naming it. + def load(**kwargs): + raise ModelConfigError(MESSAGE) + + _install_loader(monkeypatch, load) + with pytest.raises(DFloatFormatError) as info: + _hub.load_into(tmp_path, _definition(), {"vae": object()}, what="Z-Image base") + text = str(info.value) + assert text.startswith(f"{tmp_path}: not a Z-Image base with ['vae']") + assert MESSAGE in text + assert isinstance(info.value.__cause__, ModelConfigError) + + +# mflux 0.23.0's wording (weight_loader.py:500-505, WeightLoader._describe_missing_weights): at most three names, the +# total count, then what the directory holds instead. +ONE_MISSING = ( + "The vae weights in /base/vae do not fit this model: 1 required weight has no match " + "(expected names like bn.running_mean; found bn.num_batches_tracked). The checkpoint was probably converted " + "for another program or model. Use one in the original layout, or one written by mflux-save." +) +MANY_MISSING = ( + "The text_encoder weights in /base/text_encoder do not fit this model: 374 required weights have no match " + "(expected names like model.language_model.layers.2.input_layernorm.weight, " + "model.language_model.layers.2.post_attention_layernorm.weight, " + "model.language_model.layers.2.self_attn.q_proj.weight; found none under the names this component reads). " + "The checkpoint was probably converted for another program or model." +) + + +def _refused_with(monkeypatch, tmp_path, message, error=ModelConfigError): + def load(**kwargs): + raise error(message) + + _install_loader(monkeypatch, load) + with pytest.raises(DFloatFormatError) as info: + _hub.load_into(tmp_path, _definition(), {"vae": object()}, what="FLUX.2 Klein base") + return str(info.value) + + +def test_a_missing_required_tensor_is_refused_by_its_name_before_mfluxs_wording( + monkeypatch, tmp_path +): + # Bug caught: mflux 0.23's required-weights refusal wrapped as a generic "not a … base", so the message never says + # which tensor the base lacks (the name was only buried in mflux's own sentence, after ours). + text = _refused_with(monkeypatch, tmp_path, ONE_MISSING) + assert text.startswith( + f"{tmp_path}: not a FLUX.2 Klein base with ['vae']: vae bn.running_mean not in the base" + ) + assert ONE_MISSING in text + + +def test_more_missing_tensors_than_mflux_names_are_counted(monkeypatch, tmp_path): + # Bug caught: the count dropped (a base missing 374 tensors read as missing three), or computed from the names + # mflux quotes instead of the total it reports (off by the three shown). + text = _refused_with(monkeypatch, tmp_path, MANY_MISSING) + assert ( + "text_encoder model.language_model.layers.2.input_layernorm.weight, " + "model.language_model.layers.2.post_attention_layernorm.weight, " + "model.language_model.layers.2.self_attn.q_proj.weight and 371 more not in the base" + ) in text + + +def test_another_value_error_keeps_the_generic_refusal(monkeypatch, tmp_path): + # Bug caught: any ValueError from the loader (a bad shard header, an unreadable index) reported as missing tensors + # it never named. + text = _refused_with(monkeypatch, tmp_path, "Unsupported safetensors dtype F8_E4M3", ValueError) + assert text.startswith( + f"{tmp_path}: not a FLUX.2 Klein base with ['vae'] (a DFloat11 checkpoint" + ) + assert "not in the base" not in text + + +# The contract: the message is built by mflux 0.23's own WeightLoader._describe_missing_weights, not copied. Its inputs +# are minimal: a component namespace with the three fields it reads (name, num_blocks, num_layers; the counts given, so +# mflux does not detect them from the weights), an empty mapping (every raw name is then unmapped) and raw weights of +# one tiny array. +@pytest.mark.mflux +@pytest.mark.parametrize( + ("missing", "expected"), + [ + pytest.param(["bn.running_mean"], "vae bn.running_mean", id="one"), + pytest.param( + [f"decoder.up.{i}.weight" for i in range(5)], + "vae decoder.up.0.weight, decoder.up.1.weight, decoder.up.2.weight and 2 more", + id="five", + ), + ], +) +def test_the_parser_reads_mflux_023s_own_missing_weights_message(missing, expected): + # Red when: mflux's wording drifts from the pattern (the singular "weight has", the count, the "found" clause), or + # the count is dropped or taken from the names mflux shows. + import mlx.core as mx + + loader = pytest.importorskip("mflux.models.common.weights.loading.weight_loader") + component = types.SimpleNamespace(name="vae", num_blocks=0, num_layers=0) + raw = {"bn.num_batches_tracked": mx.zeros((1,))} + message = loader.WeightLoader._describe_missing_weights( + component, "/base/vae", [], missing, raw + ) + assert _hub.missing_weights(message) == expected diff --git a/tests/test_mflux_pins.py b/tests/test_mflux_pins.py new file mode 100644 index 0000000..9d1e855 --- /dev/null +++ b/tests/test_mflux_pins.py @@ -0,0 +1,34 @@ +"""The mflux pin: the ``[mflux]`` extra and the ``bench`` group name the same checked range and cap Pillow with it.""" + +import tomllib +from pathlib import Path + +import pytest + +PYPROJECT = Path(__file__).resolve().parents[1] / "pyproject.toml" +MFLUX_PIN = "mflux>=0.23,<0.24; sys_platform == 'darwin'" +PILLOW_CAP = "pillow<13" + + +def _requirement_lists() -> dict[str, list[str]]: + data = tomllib.loads(PYPROJECT.read_text()) + return { + "the [mflux] extra": data["project"]["optional-dependencies"]["mflux"], + "the bench group": data["dependency-groups"]["bench"], + } + + +@pytest.mark.parametrize("where", ["the [mflux] extra", "the bench group"]) +def test_each_mflux_requirement_list_pins_the_checked_range(where): + # Bug caught: one list left on the old range (the lock then resolves an mflux the adapters were not + # checked against), or the range widened past the release the seams were read against. + requirements = _requirement_lists()[where] + assert [r for r in requirements if r.startswith("mflux")] == [MFLUX_PIN] + + +@pytest.mark.parametrize("where", ["the [mflux] extra", "the bench group"]) +def test_each_mflux_requirement_list_caps_pillow_below_13(where): + # Bug caught: the cap dropped from one list; mflux 0.23.0 imports a private Pillow module that + # Pillow 13 removed, so a fresh install fails at `import mflux`. + requirements = _requirement_lists()[where] + assert [r for r in requirements if r.lower().startswith("pillow")] == [PILLOW_CAP] diff --git a/tests/test_mflux_pipeline.py b/tests/test_mflux_pipeline.py index 9dc3078..78ba247 100644 --- a/tests/test_mflux_pipeline.py +++ b/tests/test_mflux_pipeline.py @@ -75,6 +75,27 @@ def test_the_default_construction_arguments_pass(): _pipeline.refuse_construction_args(quantize=None, lora_paths=[], lora_scales=[], bake_lora=True) +@pytest.mark.parametrize( + ("kwargs", "name"), + [ + ({"compute_precision": mx.float16}, "compute_precision"), + ({"float32": True}, "float32"), + ], +) +def test_each_precision_argument_this_path_cannot_honour_is_refused_by_name(kwargs, name): + # Bug caught: mflux 0.23's compute_precision (casts the attention and feed-forward weights in place) or + # Z-Image's float32 stream silently accepted: the output would no longer follow the run the checkpoint was + # checked against, and the memory constants were not measured for it. + args = {"float32": False, "compute_precision": None, **kwargs} + with pytest.raises(DFloatUnsupportedError, match=rf"^{name}: .*not on the DFloat11 path"): + _pipeline.refuse_precision_args(**args) + + +def test_the_default_precision_arguments_pass(): + # Bug caught: mflux's defaults (no compute precision, the model's own stream dtype) refused. + _pipeline.refuse_precision_args(float32=False, compute_precision=None) + + def test_an_eval_policy_outside_the_familys_is_refused(): # Bug caught: policy "none" (every decoded block kept resident) accepted. _pipeline.check_eval_policy("depth2", ("per-block", "depth2")) diff --git a/tests/test_qwen21_build.py b/tests/test_qwen21_build.py index f5db582..0749995 100644 --- a/tests/test_qwen21_build.py +++ b/tests/test_qwen21_build.py @@ -105,22 +105,28 @@ def test_the_build_installs_placeholders_and_loads_every_extra_by_name(tmp_path) @pytest.mark.mflux def test_the_computed_tables_are_left_as_mflux_built_them(tmp_path): - # Bug caught: the coverage exemption replaced by loading something into the computed tables (a zero or stale - # RoPE table would wreck the positions while every check passes). + # Bug caught: the build loading something into the computed tables (a zero or stale RoPE table would wreck the + # positions while every check passes). mflux 0.23 keeps them outside the parameters; the oracle is a fresh module. from mflux.models.qwen21.model.qwen21_transformer.qwen21_rope import Qwen21Rope from mflux.models.qwen21.model.qwen21_transformer.qwen21_time_text_embed import ( Qwen21Timesteps, ) + from mlx_dfloat.mflux.qwen21.transformer import computed_tables + ckpt, _m, _c = _open(tmp_path) - params = _params(build_transformer(ckpt, transformer_kwargs=TINY)) - rope = Qwen21Rope(axes_dim=[4, 6, 6]) - for i in range(3): - for table, fresh in (("cos", rope.cos_tables[i]), ("sin", rope.sin_tables[i])): - got = params[f"pos_embed.{table}_tables.{i}"] - assert mx.array_equal(got.view(mx.uint32), fresh.view(mx.uint32)).item(), (table, i) - freqs = params["time_text_embed.time_proj.freqs"] - assert mx.array_equal(freqs.view(mx.uint32), Qwen21Timesteps().freqs.view(mx.uint32)).item() + tables = computed_tables(build_transformer(ckpt, transformer_kwargs=TINY).transformer) + rope = Qwen21Rope(axes_dim=(4, 6, 6)) + fresh = { + **{f"pos_embed._cos_tables.{i}": t for i, t in enumerate(rope._cos_tables)}, + **{f"pos_embed._sin_tables.{i}": t for i, t in enumerate(rope._sin_tables)}, + "time_text_embed.time_proj._freqs": Qwen21Timesteps()._freqs, + } + assert set(tables) == set(fresh) + for name, want in fresh.items(): + got = tables[name] + assert got.dtype == want.dtype, name + assert mx.array_equal(got.view(mx.uint32), want.view(mx.uint32)).item(), name @pytest.mark.mflux diff --git a/tests/test_qwen21_init.py b/tests/test_qwen21_init.py index f1b658b..a267c11 100644 --- a/tests/test_qwen21_init.py +++ b/tests/test_qwen21_init.py @@ -105,7 +105,7 @@ def test_df11_patterns_fetch_the_one_comfyui_file(): @pytest.mark.mflux def test_the_definitions_name_only_their_component_and_the_processor_tokenizer(): # Bug caught: a subset keeping the transformer component (mflux's loader would read transformer/), or the - # tokenizer looked up anywhere but processor/ (mflux 0.20.0 qwen21_weight_definition.py:38-50). + # tokenizer looked up anywhere but processor/ (mflux 0.23.0 qwen21_weight_definition.py:38-50). from mlx_dfloat.mflux.qwen21.init import encoder_definition, vae_definition assert [c.name for c in encoder_definition().get_components()] == ["text_encoder"] @@ -121,7 +121,7 @@ def test_the_definitions_name_only_their_component_and_the_processor_tokenizer() @pytest.mark.mflux def test_the_computed_exemptions_are_exactly_what_mfluxs_mapping_does_not_load(): # Bug caught: a computed buffer missing from COMPUTED (every real base refused), or a mapped weight listed there - # (a missing tensor never noticed). mflux 0.20.0: the VAE maps every parameter (qwen21_weight_mapping.py:91 on), + # (a missing tensor never noticed). mflux 0.23.0: the VAE maps every parameter (qwen21_weight_mapping.py:91 on), # the text encoder all but rotary_emb.inv_freq (qwen3_vl_rope.py:21). A two-layer encoder: the mapping names 36. from mflux.models.qwen21.model.qwen21_text_encoder.qwen21_text_encoder import Qwen21TextEncoder from mflux.models.qwen21.model.qwen21_vae.qwen21_vae import Qwen21VAE @@ -156,11 +156,13 @@ def wrong_width(root, definition, models, *, what): load_text_encoder(tmp_path) -# A two-layer Qwen3-VL text stack (tiny widths; mflux 0.20.0 Qwen21TextEncoder's constructor arguments). +# A 36-layer Qwen3-VL text stack at tiny widths (mflux 0.23.0 Qwen21TextEncoder's constructor arguments). Every +# layer is there: mflux's mapping lists layers 0-35 as required (Qwen21WeightMapping.NUM_TEXT_LAYERS = 36), and +# 0.23's loader refuses a checkpoint missing any of them before it applies a weight (weight_loader.py:193-203). TINY_ENCODER = { "vocab_size": 100, "hidden_size": 64, - "num_hidden_layers": 2, + "num_hidden_layers": 36, "num_attention_heads": 2, "num_key_value_heads": 1, "intermediate_size": 128, @@ -232,7 +234,8 @@ def test_a_vae_tensor_missing_from_the_base_is_refused_by_name(monkeypatch, tmp_ # keeps its random init and the decode returns noise. from mlx_dfloat.mflux.qwen21.init import load_vae - missing = "decoder.conv_out.conv.weight" + # mflux 0.23: decoder.conv_out is a plain nn.Conv2d, no wrapped .conv (vae.py:60). + missing = "decoder.conv_out.weight" monkeypatch.setattr(_hub, "load_into", _replace_all_but("vae", missing)) with pytest.raises(DFloatFormatError, match=re.escape(missing) + ".*not in the base"): load_vae(tmp_path) @@ -246,7 +249,7 @@ def test_a_complete_vae_loads(monkeypatch, tmp_path): monkeypatch.setattr(_hub, "load_into", _replace_all_but("vae", None)) vae = load_vae(tmp_path) - assert vae.decoder.conv_out.conv.weight.dtype == mx.bfloat16 + assert vae.decoder.conv_out.weight.dtype == mx.bfloat16 @pytest.mark.mflux diff --git a/tests/test_qwen21_memory.py b/tests/test_qwen21_memory.py index b33de37..e49f1f4 100644 --- a/tests/test_qwen21_memory.py +++ b/tests/test_qwen21_memory.py @@ -15,12 +15,14 @@ The denoise term is re-derived at the planner's limit (2026-10-08, the first one-step calibration run at that limit, in the cache-limit A/B; same model, prompt, size, seed and guidance): denoise footprint peak 13_939_106_072 at the -2_731_761_634 limit. The VAE term covers the highest of eleven 1024² VAE-phase footprint peaks (21.76-22.79 GB; MLX's -VAE peak was the same in every calibration run): 22_788_580_400, from the second one-step run of the A/B at the old limit. - -The calibration runs' phase peaks are committed in ``bench/results/calibration/qwen-image-2.1-1024.json`` (each run's -label, date, cache limit and footprint peaks; ``constants_from`` names the run each term was derived from), and the -tests below read them from there. +2_731_761_634 limit. The VAE term covers the highest of thirteen 1024² VAE-phase footprint peaks (21.76-22.83 GB): eleven +on mflux 0.20 (MLX's VAE peak was the same in every calibration run) and two on mflux 0.23 (the 40-step MEASURED run and +the identity check's df11 side, 2026-10-24); the highest is the 0.23 identity run's 22_831_539_272. + +The calibration runs' phase peaks are committed in ``bench/results/calibration/qwen-image-2.1-1024.json`` (mflux 0.20) +and ``qwen-image-2.1-1024-mflux023.json`` (mflux 0.23, copied from the records its ``source`` keys name): each run's +label, date, cache limit and footprint peaks; ``constants_from`` names the run each term was derived from (the 0.23 +record's overrides the 0.20 record's for the VAE term). The tests below read them from there. """ import json @@ -59,13 +61,16 @@ TEXT_TOKENS = 33 # the calibration prompt's (the negative " " is 9) REPO = Path(__file__).resolve().parents[1] CALIBRATION = json.loads((REPO / "bench/results/calibration/qwen-image-2.1-1024.json").read_text()) -CALIBRATION_RUNS = {run["label"]: run for run in CALIBRATION["runs"]} +CALIBRATION_023 = json.loads( + (REPO / "bench/results/calibration/qwen-image-2.1-1024-mflux023.json").read_text() +) +CALIBRATION_RUNS = {run["label"]: run for run in CALIBRATION["runs"] + CALIBRATION_023["runs"]} # The phase peaks the constants were derived from: the denoise peak is the first one-step A/B run's at CACHE_LIMIT; -# the VAE peak the highest of the eleven 1024² samples (the second A/B run at the old limit); the encode peak the first +# the VAE peak the highest of the thirteen 1024² samples (the mflux 0.23 identity run); the encode peak the first # calibration run's. MEASURED = { phase: CALIBRATION_RUNS[label]["footprint_peaks"][phase] - for phase, label in CALIBRATION["constants_from"].items() + for phase, label in (CALIBRATION["constants_from"] | CALIBRATION_023["constants_from"]).items() } VAE_SAMPLES = {label: run["footprint_peaks"]["vae"] for label, run in CALIBRATION_RUNS.items()} # This Mac (M1 Max 32 GB): mx.device_info()["max_recommended_working_set_size"] and the RAM, from the calibration record. @@ -250,23 +255,32 @@ def _plan(budget, *, fit_check=True, cache_limit_override=None): def test_the_calibration_record_holds_the_runs_the_constants_name(): # Bug caught: the committed record and the constants drifting apart (a re-derived term whose run is not recorded, # or a record edit that moves a peak): the three source peaks are the literals the constants were derived from. - assert CALIBRATION["constants_from"] == { - "encode": "derisk-1", - "denoise": "ab-new-limit-1", - "vae": "ab-old-limit-2", - } - assert MEASURED == {"encode": 15_857_938_288, "denoise": 13_939_106_072, "vae": 22_788_580_400} + assert CALIBRATION_023["constants_from"] == {"vae": "identity-4-steps-cfg-mflux023"} + assert MEASURED == {"encode": 15_857_938_288, "denoise": 13_939_106_072, "vae": 22_831_539_272} assert CALIBRATION_RUNS["ab-new-limit-1"]["cache_limit"] == CACHE_LIMIT +def test_the_mflux_023_samples_are_the_committed_records_peaks(): + # Bug caught: the derived calibration record drifting from the run records it was copied from (a peak typed in by + # hand, or a record re-run without the calibration following it). + for run in CALIBRATION_023["runs"]: + source = json.loads((REPO / run["source"]["record"]).read_text()) + assert run["footprint_peaks"]["vae"] == source[run["source"]["field"]] + assert source["exit_code"] == 0 + assert {run["footprint_peaks"]["vae"] for run in CALIBRATION_023["runs"]} == { + 21_877_154_888, + 22_831_539_272, + } + + def test_the_vae_term_covers_every_1024_sample_and_equals_the_highest(): - # Bug caught: the VAE transient derived from one run instead of the highest of the eleven samples (a fit check that - # under-predicts the runs it was calibrated on: the first calibration run alone puts the phase 648_937_376 B low), or a - # sample dropped from the record. - assert len(VAE_SAMPLES) == 11 + # Bug caught: the VAE transient derived from one run instead of the highest of the thirteen samples (a fit check + # that under-predicts the runs it was calibrated on: the mflux 0.20 constant puts the VAE phase 42_958_872 B under + # the 0.23 identity run, the first calibration run alone 691_896_248 B), or a sample dropped from either record. + assert len(VAE_SAMPLES) == 13 fit = _fit() - assert fit.phases["vae"] == max(VAE_SAMPLES.values()) assert all(peak <= fit.phases["vae"] for peak in VAE_SAMPLES.values()) + assert fit.phases["vae"] == max(VAE_SAMPLES.values()) == 22_831_539_272 def test_the_estimate_at_1024_equals_the_calibration_run_phase_for_phase(): @@ -308,9 +322,9 @@ def test_a_32_gb_mac_keeps_the_set_through_the_vae_decode(monkeypatch): # The shared rule (_pipeline.plan_call_for) drops the set before the VAE decode iff the VAE phase with the set is # over the budget, with no margin beyond the budget's own: the recommended working set minus a 2 GiB reserve # (integrate.memory.budget_bytes), 26_800_603_136 - 2_147_483_648 = 24_653_119_488 B here. The measured VAE peak, - # 22_788_580_400 B (the estimate's by calibration), is 1_864_539_088 B (1.74 GiB) under it: the set stays resident, + # 22_831_539_272 B (the estimate's by calibration), is 1_821_580_216 B (1.70 GiB) under it: the set stays resident, # and a 32 GB Mac pays no set reload per call. Bug caught: the decision inverted, a margin added on top of the - # reserve that flips it (more than 1.74 GiB), or the budget read without the reserve. + # reserve that flips it (more than 1.70 GiB), or the budget read without the reserve. monkeypatch.setattr( imem.mx, "device_info", lambda: {"max_recommended_working_set_size": RECOMMENDED_32} ) @@ -328,9 +342,9 @@ def test_a_32_gb_mac_keeps_the_set_through_the_vae_decode(monkeypatch): @pytest.mark.parametrize("tier", [16, 24]) def test_a_16_or_24_gb_tier_is_refused_in_the_encode_phase(tier): # The prediction for the CAPPED tiers at 1024² (the calibration prompt, 33 tokens): the VAE phase with the set - # (22_788_580_400 B) is over either tier's budget, so the set is dropped before the decode: the VAE phase is then + # (22_831_539_272 B) is over either tier's budget, so the set is dropped before the decode: the VAE phase is then # the extras (they stay on the transformer) + the VAE + its transient + overhead, 137_388_032 + 1_350_989_512 + - # 11_119_436_651 + 459_445_098 = 13_067_259_293 B; the encode phase (the language model 15_136_811_008 + the encode term 66_567_800 + overhead 459_445_098 = + # 11_162_395_523 + 459_445_098 = 13_110_218_165 B; the encode phase (the language model 15_136_811_008 + the encode term 66_567_800 + overhead 459_445_098 = # 15_662_823_906 B, 14.59 GiB) is then the peak and over both fit budgets. The fit budget is what a real Mac of # that size applies (generate passes limits.fit_budget_bytes): the tier's recommended working set (2/3 of its RAM) # minus the fit rule's 2 GiB reserve. 24 GB: 17_179_869_184 - 2_147_483_648 = 15_032_385_536 B (14.00 GiB), over by @@ -348,7 +362,7 @@ def test_a_16_or_24_gb_tier_is_refused_in_the_encode_phase(tier): assert forced.drop_set_before_vae is True assert forced.estimate.peak_phase == "encode" assert forced.estimate.peak_bytes == 15_662_823_906 - assert forced.estimate.phases["vae"] == 13_067_259_293 + assert forced.estimate.phases["vae"] == 13_110_218_165 # What the prompt encode may hold before the model warns (model._check_encode_peak): MLX's encode peak minus what was diff --git a/tests/test_qwen21_mflux023.py b/tests/test_qwen21_mflux023.py new file mode 100644 index 0000000..529e3e8 --- /dev/null +++ b/tests/test_qwen21_mflux023.py @@ -0,0 +1,248 @@ +"""Our DF11 Qwen-Image 2.1 transformer against mflux 0.23's own default generation path, bit for bit. `mflux` lane. + +mflux 0.23's default for a padding-free prompt inside ``inference_projections()``: packed Q/K/V and gate/up, the +fused QK-norm/RoPE kernel, the text-prefix K/V cache (``text_forward`` once per prompt) and the image-token step +(``image_forward``). Our side decodes each block's DF11 group just in time, packs it per block and runs both passes +uncompiled. The reference is stock mflux under ``mx.disable_compile()``: mflux compiles both step functions, and the +compiled bits differ from the same graph run plain (up to about 80 % of the elements by an ulp or more), so the +identity bar is the plain graph, as for every other family. Both sides of a comparison run under +``mflux_uncompiled()``: MLX's own activations are compiled too (``nn.silu`` is ``mx.compile(shapeless=True)``), and +for some weights the compiled bits differ from the plain body's. That a seamed transformer keeps mflux's step +functions plain is checked on its own below. Same weights on both sides; the comparison is ``view(uint16)``. +""" + +import types + +import mlx.core as mx +import numpy as np +import pytest +from mlx.utils import tree_flatten, tree_unflatten +from tests._qwen21_tiny import WIDE, mflux_uncompiled, tiny_config, tiny_groups, tiny_seamed + +from mlx_dfloat.decode import decode_group +from mlx_dfloat.errors import DFloatIntegrationError +from mlx_dfloat.integrate.providers import DF11Provider + +pytestmark = pytest.mark.mflux + +COMPILED = type(mx.compile(lambda x: x)) # mlx.gc_func +STEP_METHODS = {"_step_fn": "_forward", "_image_step_fn": "_image_forward"} + + +def _inputs(*, text=6, padded=False, t=0, seed=0): + """``Qwen21Transformer.__call__`` inputs at WIDE: 16 image tokens of 64 channels, ``text`` prompt tokens.""" + keys = mx.random.split(mx.random.key(seed), 2) + mask = mx.ones((1, text), dtype=mx.int32) + if padded: + mask = mx.concatenate([mask[:, :-1], mx.zeros((1, 1), dtype=mx.int32)], axis=1) + return { + "t": t, + "config": tiny_config(steps=2), + "hidden_states": mx.random.normal((1, 16, 64), key=keys[0]).astype(mx.bfloat16), + "encoder_hidden_states": mx.random.normal((1, text, 32), key=keys[1]).astype(mx.bfloat16), + "encoder_hidden_states_mask": mask, + } + + +def _pair(decoded=None): + """(ours, native): our seamed WIDE transformer on a DF11 provider, and stock mflux holding the same weights. + + The non-block parameters are cast to bfloat16, as a checkpoint's extras load (the computed RoPE tables and + timestep frequencies are not parameters in mflux 0.23 and keep mflux's own values on both sides); the block + matrices are the DF11 groups' source bits. + + The non-block parameters come from MLX's global generator, seeded here so a test does not depend on the tests run + before it. Seed 57 (mlx 0.32.2) gives weights where MLX's compiled ``nn.silu`` and its plain body differ in the + last bit, so a comparison that runs one side compiled and the other plain fails on it. + """ + from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer + + from mlx_dfloat.mflux.qwen21.names import qwen21_name_map + + mx.random.seed(57) + ours, shapes = tiny_seamed(n_layers=1, dims=WIDE) + groups, names, sources = tiny_groups(shapes, np.random.default_rng(7)) + native = Qwen21Transformer(**WIDE) + matrices = {m for per in names.values() for m in per} + shared = [ + (k, v.astype(mx.bfloat16)) + for k, v in tree_flatten(native.parameters()) + if k not in matrices + ] + native.update(tree_unflatten(shared)) + ours.update(tree_unflatten(shared)) + native.update( + tree_unflatten( + [ + (m, mx.array(bits).view(mx.bfloat16)) + for per in sources.values() + for m, bits in per.items() + ] + ) + ) + + def decode(group): + if decoded is not None: + decoded.append(group.name) + return decode_group(group, backend="reference") + + ours.attach( + DF11Provider(groups, names, qwen21_name_map(), decode=decode), shapes, verify_in_call=True + ) + return ours, native + + +def _bits(x): + mx.eval(x) + return x.view(mx.uint16).tolist() + + +def test_two_steps_match_mflux_s_default_packed_cached_path_uncompiled_bit_for_bit(): + # Bugs caught: the seam leaving text_forward or image_forward unseamed (a placeholder matmul raises), running the + # separate projections (test_packing_changes_the_bits_at_this_shape shows they differ here), packing outside + # the window or keeping a packed matrix across blocks, or a step reusing the wrong text cache. The precondition + # lines catch a WIDE that no longer reaches the fused kernel. + ours, native = _pair() + assert native.transformer_blocks[0].attn.use_fused_prologue + assert ours.transformer_blocks[0].attn.use_fused_prologue + # One prompt array for both steps: mflux keys its text cache on the array's identity, as its loop passes it. + steps = [_inputs(t=t) for t in (0, 1)] + for step in steps[1:]: + step["encoder_hidden_states"] = steps[0]["encoder_hidden_states"] + with mflux_uncompiled(), native.inference_projections(): + want = [_bits(native(**step)) for step in steps] + with mflux_uncompiled(), ours.inference_projections(): + got = [_bits(ours(**step)) for step in steps] + assert got == want + + +def test_packing_changes_the_bits_at_this_shape(): + # Bug caught (the control): a WIDE where mflux's packed and separate projections agree, so the test above could + # not tell a seam that skips the packing. Stock mflux only, uncompiled as the reference above, with + # use_packed_projections off on one side. + _ours, native = _pair() + inputs = _inputs() + with mflux_uncompiled(): + with native.inference_projections(): + packed = _bits(native(**inputs)) + native.use_packed_projections = False + with native.inference_projections(): + separate = _bits(native(**inputs)) + assert packed != separate + + +def test_mflux_compiles_both_step_functions_and_ours_runs_them_plain(): + # Bug caught: the parity reference above made uncompiled while mflux itself also stopped compiling (the reference + # would then be the only mflux path), or the seamed transformer letting mflux compile either pass. mflux 0.23 + # compiles on every call path (qwen21_transformer.py:96-98 joint, 111-113 image step), with no chip check. + for padded, attr in ((True, "_step_fn"), (False, "_image_step_fn")): + ours, native = ( + _pair() + ) # one pair per path: mflux caches the geometry (and its mask) per shape + inputs = _inputs(padded=padded) + with native.inference_projections(): + mx.eval(native(**inputs)) + assert type(getattr(native, attr)) is COMPILED, attr + with ours.inference_projections(): + mx.eval(ours(**inputs)) + fn = getattr(ours, attr) + assert type(fn) is types.MethodType, attr + assert fn.__func__ is getattr(type(native), STEP_METHODS[attr]), attr + + +def test_a_padded_prompt_takes_the_joint_pass_and_matches_mflux(): + # Bug caught: the block's own __call__ (the joint pass mflux runs for a padded prompt, qwen21_transformer.py:92-99) + # unseamed or unpacked. _step_fn left compiled is test_mflux_compiles_both_step_functions_and_ours_runs_them_plain's. + ours, native = _pair() + inputs = _inputs(padded=True) + with mflux_uncompiled(), native.inference_projections(): + want = _bits(native(**inputs)) + with mflux_uncompiled(), ours.inference_projections(): + got = _bits(ours(**inputs)) + assert got == want + + +def test_the_text_prefix_decodes_each_block_once_per_prompt(): + # Bug caught: the text pass rebuilt every step (a second decode per block per step) or never built (mflux's + # cache keyed off another array). Two steps on one prompt: text pass + 2 image passes = 3 decodes of the block. + decoded = [] + ours, _native = _pair(decoded) + inputs = _inputs() + with ours.inference_projections(): + for t in (0, 1): + mx.eval(ours(**{**inputs, "t": t})) + assert decoded == ["transformer_blocks.0"] * 3 + + +@pytest.mark.parametrize("attr", ["_step_fn", "_image_step_fn"]) +def test_a_compiled_step_function_is_refused_before_any_decode(attr): + # Bug caught: the guard checking only _step_fn (0.20's one compiled function): a compiled _image_step_fn runs the + # seam inside mx.compile, where the per-block eval raises. + from mlx_dfloat.mflux.qwen21.transformer import is_eager + + decoded = [] + ours, _native = _pair(decoded) + method = "_forward" if attr == "_step_fn" else "_image_forward" + setattr(ours, attr, mx.compile(getattr(ours, method))) + assert not is_eager(ours) + with pytest.raises(DFloatIntegrationError, match=rf"{attr} is .*compiled the forward pass"): + ours(**_inputs()) + assert decoded == [] + + +@pytest.mark.parametrize("attr", ["_step_fn", "_image_step_fn"]) +def test_a_compiled_step_function_set_before_the_window_is_refused_on_entry(attr): + # Bug caught: inference_projections() clearing a compiled step function before the guard looks (mflux 0.23's + # _clear_inference_projections sets both to None), so a compile planted outside the window goes unreported by + # the call while report() says "compiled"; the refusal must also come before any decode. + from mlx_dfloat.mflux.qwen21.transformer import is_eager + + decoded = [] + ours, _native = _pair(decoded) + compiled = mx.compile(getattr(ours, STEP_METHODS[attr])) + setattr(ours, attr, compiled) + with ( + pytest.raises(DFloatIntegrationError, match=rf"{attr} is .*compiled the forward pass"), + ours.inference_projections(), + ): + ours(**_inputs()) + assert decoded == [] + assert getattr(ours, attr) is compiled + assert not is_eager(ours) + assert ours.pack_switch.on is False + + +def test_the_seamed_transformer_runs_both_passes_plain(): + # Bug caught: mflux's compile reaching through the seamed class on the default (text-cache) path. + from mlx_dfloat.mflux.qwen21.transformer import is_eager + + ours, _native = _pair() + with ours.inference_projections(): + mx.eval(ours(**_inputs())) + assert type(ours._image_step_fn) is types.MethodType + assert is_eager(ours) + + +def test_the_generate_call_refuses_step_reuse(tmp_path, monkeypatch): + # Bug caught: step_cache_ratio / teacache_ratio passed through (#779): mflux would skip transformer calls, an + # image no identity check here has covered, or the refusal coming after the prompts were encoded. + from tests._qwen21_tiny import fake_model + + from mlx_dfloat.errors import DFloatUnsupportedError + + model = fake_model(tmp_path, monkeypatch) + for kwargs in ({"step_cache_ratio": 0.25}, {"teacache_ratio": 0.25}): + with pytest.raises(DFloatUnsupportedError, match="step_cache_ratio/teacache_ratio"): + model.generate_image( + seed=1, prompt="a", num_inference_steps=2, height=64, width=64, **kwargs + ) + assert model.prompt_cache == {} + + +def test_the_constructor_refuses_a_compute_precision(): + # Bug caught: compute_precision (#820) accepted: mflux casts the attention and FFN modules, another image. + from mlx_dfloat.errors import DFloatUnsupportedError + from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 + + with pytest.raises(DFloatUnsupportedError, match="compute_precision"): + DFloatQwenImage21(compute_precision=mx.float16) diff --git a/tests/test_qwen21_model.py b/tests/test_qwen21_model.py index 9519193..cb41b32 100644 --- a/tests/test_qwen21_model.py +++ b/tests/test_qwen21_model.py @@ -335,8 +335,10 @@ def test_the_generate_defaults_are_mfluxs(tmp_path, monkeypatch): for n, p in inspect.signature(DFloatQwenImage21.generate_image).parameters.items() if p.default is not inspect.Parameter.empty } + # mflux 0.23.0 qwen_image_21.py:52-66: the steps default is None (Qwen21TurboScheduler.default_steps resolves + # it), and step_cache_ratio / teacache_ratio are new (refused here: test_the_generate_call_refuses_step_reuse). assert defaults == { - "num_inference_steps": 40, + "num_inference_steps": None, "height": 1024, "width": 1024, "guidance": 1.0, @@ -344,7 +346,17 @@ def test_the_generate_defaults_are_mfluxs(tmp_path, monkeypatch): "image_strength": None, "scheduler": "linear", "negative_prompt": None, + "step_cache_ratio": None, + "teacache_ratio": None, } + from mflux.models.qwen21.variants.txt2img.qwen_image_21 import QwenImage21 + + upstream = { + n: p.default + for n, p in inspect.signature(QwenImage21.generate_image).parameters.items() + if p.default is not inspect.Parameter.empty + } + assert defaults == upstream def test_a_missing_guidance_and_scheduler_become_mflux_defaults(tmp_path, monkeypatch): @@ -387,8 +399,12 @@ def test_an_empty_prompt_and_a_negative_equal_to_the_prompt_run_from_the_cache( ) assert (encodes, model._cfg_calls) == (["", "x"], 2) assert set(model.prompt_cache) == {"", "x"} - # 1 step x 1 call x 2 blocks, then 1 step x 2 calls x 2 blocks: the second call ran the negative branch. - assert model.report()["decode_launches"] == 2 + 4 + # mflux 0.23 runs the text-prefix pass (each block's text_forward) once per prompt array per generation, then + # each block's image_forward per transformer call (qwen21_transformer.py:101-113); its text cache is keyed on + # the array's identity, and a negative equal to the prompt is the same cached array (qwen21_prompt_encoder.py:22). + # Call 1: 1 text pass x 2 blocks + 1 step x 1 call x 2 blocks. Call 2: 1 text pass x 2 blocks (one array for + # both branches) + 1 step x 2 calls x 2 blocks: the second call ran the negative branch. + assert model.report()["decode_launches"] == (2 + 2) + (2 + 4) def test_text_tokens_follow_the_longest_prompt_of_the_call(tmp_path, monkeypatch): @@ -484,7 +500,7 @@ def test_plan_call_reads_the_familys_constants_sizes_and_groups(tmp_path, monkey def test_the_vae_guard_follows_the_plan_and_zeroes_the_cache_limit(tmp_path, monkeypatch): # Bug caught: the guard not reading this call's plan (the set kept through a decode the budget cannot hold, or # dropped on a roomy Mac: a reload per call), or a transformer-sized cache limit left for the VAE decode. - # Tight: VAE phase 12e9 + 11_119_436_651 + 459_445_098 = 23_578_881_749 > 16e9; denoise ~15.2e9 fits. + # Tight: VAE phase 12e9 + 11_162_395_523 + 459_445_098 = 23_621_840_621 > 16e9; denoise ~15.2e9 fits. from mlx_dfloat.mflux._phases import FamilySizes big = FamilySizes(compressed=12_000_000_000, extras=0, nonblock=0, encoders=0, vae=0) @@ -768,7 +784,9 @@ def test_a_compiled_forward_is_refused_in_the_loop_and_the_report_says_compiled( # report() still printing "uncompiled" after mflux compiled the forward pass. model = fake_model(tmp_path, monkeypatch) assert model.report()["predict"] == "uncompiled" - model.transformer._step_fn = mx.compile(model.transformer._forward) + # mflux 0.23's default path compiles _image_step_fn, and inference_projections() clears it on entry, so the + # call must refuse it there, as report() says. + model.transformer._image_step_fn = mx.compile(model.transformer._image_forward) assert model.report()["predict"] == "compiled" with pytest.raises(DFloatIntegrationError, match="compiled"): model.generate_image(seed=1, prompt="a", num_inference_steps=1, height=64, width=64) @@ -949,7 +967,8 @@ def __call__(self, input_ids, attention_mask=None): def _stock_qwen(tokenizer, vae, matrices, extras): """Stock mflux ``QwenImage21`` over a plain tiny ``Qwen21Transformer`` holding the BF16 source of a tiny checkpoint - (``write_tiny_checkpoint``'s matrices and extras), uncompiled like our side.""" + (``write_tiny_checkpoint``'s matrices and extras); the caller runs it under ``mflux_uncompiled()``.""" + from mflux.models.common.compute_precision.compute_precision import ComputePrecision from mflux.models.common.config.model_config import ModelConfig from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer from mflux.models.qwen21.qwen21_initializer import Qwen21Initializer @@ -974,17 +993,18 @@ def _stock_qwen(tokenizer, vae, matrices, extras): weights += [(name, mx.array(bits).view(mx.bfloat16)) for name, bits in extras.items()] assert all(tuple(shapes[name]) == bits.shape for name, bits in extras.items()) transformer.update(tree_unflatten(weights)) - transformer._step_fn = ( - transformer._forward - ) # mflux compiles it otherwise; our side runs it plain stock = QwenImage21.__new__(QwenImage21) nn.Module.__init__(stock) - Qwen21Initializer._init_config(stock, ModelConfig.qwen_image_21()) + Qwen21Initializer.init_config(stock, ModelConfig.qwen_image_21()) stock.vae = vae stock.text_encoder = _VaryingTextEncoder() stock.tokenizers = {"qwen21": tokenizer} stock.transformer = transformer stock.bits = None + stock.lora_paths, stock.lora_scales = None, None # read only for the image metadata + stock.compute_precision = ComputePrecision( + None + ) # mflux 0.23 sets it after the components (qwen21_initializer.py:44) return stock @@ -999,7 +1019,7 @@ def test_a_tiny_generation_gives_stock_mfluxs_latents_bit_for_bit( # mask (the negative prompt is 12 ids here, the prompt 8, so a swap changes the latents or fails). The reference is # stock mflux over the same BF16 weights, uncompiled as ours; equality is bit for bit (an integer view of the # latents). The extras are random: constant ones make the forward pass blind to its inputs. - from tests._qwen21_tiny import TINY, StubTokenizer, stub_components + from tests._qwen21_tiny import TINY, StubTokenizer, mflux_uncompiled, stub_components from mlx_dfloat.format import open_checkpoint from mlx_dfloat.mflux.qwen21.transformer import build_transformer @@ -1028,8 +1048,12 @@ def test_a_tiny_generation_gives_stock_mfluxs_latents_bit_for_bit( "guidance": guidance, "negative_prompt": negative, } - model.generate_image(**kwargs) - stock.generate_image(**kwargs, scheduler="linear") + # Both sides run plain: mflux 0.23 compiles both step functions inside its window, and MLX's own activations are + # compiled too (nn.silu), whose bits can differ from the plain graph's for some weights. + with mflux_uncompiled(): + model.generate_image(**kwargs) + with mflux_uncompiled(): + stock.generate_image(**kwargs, scheduler="linear") (ours,), (theirs,) = ours_vae.seen, stock_vae.seen assert model.report()["cfg_calls_per_step"] == (2 if negative else 1) assert bool(mx.all(mx.isfinite(theirs))) diff --git a/tests/test_qwen21_packing.py b/tests/test_qwen21_packing.py new file mode 100644 index 0000000..a127ea3 --- /dev/null +++ b/tests/test_qwen21_packing.py @@ -0,0 +1,216 @@ +"""The Qwen-Image 2.1 block seam packs Q/K/V and gate/up per block, as mflux 0.23 packs them per generation. + +mflux 0.23 ``Qwen21Transformer.inference_projections()`` (qwen21_transformer.py:126-140) concatenates each block's +``to_q``/``to_k``/``to_v`` and ``gate_layer``/``proj`` weights once per generation; the blocks then run one matmul +against the packed matrix (qwen21_attention.py:208-212, qwen21_feed_forward.py:17-22). Our blocks hold zero-size +placeholders until their decode, so the seam packs each block's decoded weights itself. No mflux import here: the +block below has the 0.23 block's shape (the same attribute names and the same packed/unpacked branches), and the +packer is handed in, as ``build_transformer`` hands in mflux's own ``Qwen21Transformer._pack_linears``. +""" + +import mlx.core as mx +import mlx.nn as nn +import pytest + +from mlx_dfloat.errors import DFloatIntegrationError +from mlx_dfloat.integrate.blockseam import SeamCell +from mlx_dfloat.integrate.names import StaticNameMap +from mlx_dfloat.integrate.placeholders import PLACEHOLDER, install_placeholders +from mlx_dfloat.integrate.providers import ResidentProvider +from mlx_dfloat.integrate.seam import attach_state +from mlx_dfloat.mflux.qwen21.transformer import PackSwitch, seam_qwen21_blocks + +# Small M and K = 1024: the shape where a packed matmul's bits differ from the separate ones on this machine +# (measured 2026-10-23, mlx 0.32.2, M1 Max: M=16, K=1024 differs in 6 / 3 / 0 of 16384 Q/K/V outputs; at the real +# width K=4096 with 24 text tokens 34 / 32 / 36 of 98304). test_the_parity_shape_tells_packed_from_separate pins it. +TOKENS, DIM = 16, 1024 +SUBS = ( + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "img_mlp.gate_layer", + "img_mlp.proj", + "img_mlp.out", +) +NAME_MAP = StaticNameMap({"transformer_blocks": {sub: sub for sub in SUBS}}) + + +class Attn(nn.Module): + def __init__(self): + super().__init__() + self.to_q = nn.Linear(DIM, DIM, bias=False) + self.to_k = nn.Linear(DIM, DIM, bias=False) + self.to_v = nn.Linear(DIM, DIM, bias=False) + self.to_out = [nn.Linear(DIM, DIM, bias=False)] + self._qkv_weight = None + + def qkv(self, x): + # mflux 0.23 Qwen21Attention._project_qkv + if self._qkv_weight is not None: + return mx.split(x @ self._qkv_weight.T, 3, axis=-1) + return self.to_q(x), self.to_k(x), self.to_v(x) + + +class Mlp(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(DIM, DIM, bias=False) + self.out = nn.Linear(DIM, DIM, bias=False) + self.gate_layer = nn.Linear(DIM, DIM, bias=False) + self._gate_proj_weight = None + + def __call__(self, x): + # mflux 0.23 Qwen21SwiGLUFeedForward.__call__ + if self._gate_proj_weight is None: + gate, proj = self.gate_layer(x), self.proj(x) + else: + gate, proj = mx.split(x @ self._gate_proj_weight.T, 2, axis=-1) + return self.out(nn.silu(gate) * proj) + + +class Block(nn.Module): + """The three entry points mflux 0.23's transformer calls on a block, with the 0.23 projections inside.""" + + def __init__(self): + super().__init__() + self.attn = Attn() + self.img_mlp = Mlp() + + def _qkv_out(self, x): + q, k, v = self.attn.qkv(x) + return q, k, v, self.attn.to_out[0](q * k + v) + + def __call__(self, x): + *_qkv, attn = self._qkv_out(x) + return x + attn + self.img_mlp(x) + + def text_forward(self, x): + _q, k, v, attn = self._qkv_out(x) + return x + attn + self.img_mlp(x), k, v + + def image_forward(self, x): + *_qkv, attn = self._qkv_out(x) + return x + attn + self.img_mlp(x) + + +def pack_linears(layers): + """mflux 0.23 ``Qwen21Transformer._pack_linears`` (qwen21_transformer.py:159-175), for nn.Linear layers.""" + packed = mx.concatenate([layer.weight for layer in layers], axis=0) + mx.eval(packed) + for layer, weight in zip(layers, mx.split(packed, len(layers), axis=0), strict=True): + layer.weight = weight + return packed + + +def _weights(seed): + keys = mx.random.split(mx.random.key(seed), len(SUBS)) + return { + sub: (mx.random.normal((DIM, DIM), key=k) * 0.03).astype(mx.bfloat16) + for sub, k in zip(SUBS, keys, strict=True) + } + + +def _x(seed=11): + return mx.random.normal((1, TOKENS, DIM), key=mx.random.key(seed)).astype(mx.bfloat16) + + +def _native(weights, *, packed): + """mflux's own block with the BF16 weights, packed as ``inference_projections()`` leaves it (or not).""" + block = Block() + for sub, w in weights.items(): + node = block + for part in sub.split("."): + node = node[int(part)] if part.isdigit() else getattr(node, part) + node.weight = w + if packed: + block.attn._qkv_weight = pack_linears((block.attn.to_q, block.attn.to_k, block.attn.to_v)) + block.img_mlp._gate_proj_weight = pack_linears( + (block.img_mlp.gate_layer, block.img_mlp.proj) + ) + return block + + +def _seamed(weights, *, on): + blocks = [Block()] + lists = [("transformer_blocks", blocks)] + shapes = install_placeholders(lists, NAME_MAP) + cell, switch = SeamCell(), PackSwitch() + seam_qwen21_blocks(lists, cell, switch, pack=pack_linears) + cell.state = attach_state(ResidentProvider({"transformer_blocks.0": weights}), shapes) + switch.on = on + return blocks[0] + + +def _bits(out): + arrays = out if isinstance(out, tuple) else (out,) + mx.eval(*arrays) + return [a.view(mx.uint16).tolist() for a in arrays] + + +def test_the_parity_shape_tells_packed_from_separate(): + # Bug caught: a parity shape where packing changes no bit, so the seam tests below would pass with the packing + # left out. Our oracle is mflux's own native block, packed vs unpacked, at TOKENS x DIM. + weights, x = _weights(3), _x() + assert _bits(_native(weights, packed=True)(x)) != _bits(_native(weights, packed=False)(x)) + + +@pytest.mark.parametrize("entry", ["__call__", "text_forward", "image_forward"]) +def test_each_entry_point_computes_the_native_packed_block_bit_for_bit(entry): + # Bugs caught: the seam leaving an entry point unseamed (mflux 0.23's text-prefix cache calls text_forward, its + # image step image_forward: a placeholder matmul raises), or running it on the separate matmuls (the bits of the + # unpacked block, which the test above shows differ). + weights, x = _weights(3), _x() + want = _bits(getattr(_native(weights, packed=True), entry)(x)) + assert _bits(getattr(_seamed(weights, on=True), entry)(x)) == want + + +def test_outside_inference_projections_the_block_runs_unpacked(): + # Bug caught: packing with the switch off (mflux packs only inside inference_projections(), and only with + # use_packed_projections set): the bits would be the packed block's, not the unpacked one mflux runs then. + weights, x = _weights(3), _x() + assert _bits(_seamed(weights, on=False)(x)) == _bits(_native(weights, packed=False)(x)) + + +def test_the_packed_matrices_and_the_decoded_weights_are_released_after_each_block(): + # Bug caught: a packed matrix left on the block after its call (it holds 3 + 2 decoded matrices' worth of memory, + # 288 MiB a block at the real width, until the next decode of the block), or a weight left in place of the + # placeholder. + block = _seamed(_weights(3), on=True) + mx.eval(block.text_forward(_x())) + assert block.attn._qkv_weight is None + assert block.img_mlp._gate_proj_weight is None + for sub in SUBS: + node = block + for part in sub.split("."): + node = node[int(part)] if part.isdigit() else getattr(node, part) + assert node.weight is PLACEHOLDER + + +def test_a_block_whose_call_raises_still_releases_the_packed_matrices(): + # Bug caught: clearing the packed matrices only on success, so a failed step leaves them pinned. + block = _seamed(_weights(3), on=True) + with pytest.raises(ValueError, match="matmul"): + block(mx.zeros((1, TOKENS, DIM + 1), dtype=mx.bfloat16)) + assert block.attn._qkv_weight is None + assert block.img_mlp._gate_proj_weight is None + + +def test_a_detached_block_refuses_every_entry_point(): + # Bug caught: an entry point that skips the attach check runs on the zero-size placeholders. + blocks = [Block()] + lists = [("transformer_blocks", blocks)] + install_placeholders(lists, NAME_MAP) + seam_qwen21_blocks(lists, SeamCell(), PackSwitch(), pack=pack_linears) + for entry in ("__call__", "text_forward", "image_forward"): + with pytest.raises(DFloatIntegrationError, match=r"transformer_blocks\.0: call attach"): + getattr(blocks[0], entry)(_x()) + + +def test_seaming_a_block_twice_is_refused(): + # Bug caught: a second swap stacking two seams, which would decode and pack the block twice per call. + blocks = [Block()] + lists = [("transformer_blocks", blocks)] + seam_qwen21_blocks(lists, SeamCell(), PackSwitch(), pack=pack_linears) + with pytest.raises(DFloatIntegrationError, match="already seamed"): + seam_qwen21_blocks(lists, SeamCell(), PackSwitch(), pack=pack_linears) diff --git a/tests/test_qwen21_tiny.py b/tests/test_qwen21_tiny.py index 0c97856..8609cca 100644 --- a/tests/test_qwen21_tiny.py +++ b/tests/test_qwen21_tiny.py @@ -26,7 +26,8 @@ def test_one_tiny_generation_decodes_once_per_block_per_transformer_call( # Bug caught: anything between our prelude and mflux's loop not composing (the eager forward pass, the pairs in # mflux's prompt cache, the VAE guard), or the negative branch run at guidance 1.0 or with an empty negative # (mflux's rule: guidance > 1.0 and bool(negative_prompt), qwen_image_21.py:89). Tiny transformer: 2 blocks; 2 - # steps: launches == 2 steps x calls x 2 blocks == 8 with CFG, 4 without. + # steps: launches == (2 steps + 1 text pass per prompt, mflux 0.23's text-prefix cache) x calls x 2 blocks == 12 + # with CFG, 6 without. from mflux.utils.apple_silicon import AppleSiliconUtil # mflux compiles this transformer's forward on every chip; a Max chip shows no chip check is relied on. @@ -57,7 +58,7 @@ def call_after_loop(self, seed, prompt, latents, config): assert seen["latents"].shape == (1, 16, 64) assert np.isfinite(seen["latents"]).all() report = model.report() - assert report["decode_launches"] == 2 * calls * 2 + assert report["decode_launches"] == (2 + 1) * calls * 2 assert report["cfg_calls_per_step"] == calls assert report["predict"] == "uncompiled" # The tiny block group (4 x 32 x 32 + 2 x 96 x 32 + 32 x 96 = 13_312 elements, 26_624 B) next to the allowance's diff --git a/tests/test_qwen21_tiny_compile.py b/tests/test_qwen21_tiny_compile.py new file mode 100644 index 0000000..4aaf801 --- /dev/null +++ b/tests/test_qwen21_tiny_compile.py @@ -0,0 +1,57 @@ +"""``mflux_uncompiled`` gives back the compile state it found: off stays off, on comes back on.""" + +import os + +import mlx.core as mx +import pytest +from tests._qwen21_tiny import mflux_uncompiled + + +def _compile_is_on() -> bool: + # A compiled function traces its Python body once per input shape; uncompiled, it runs the body on every call. + calls = [] + + def body(x): + calls.append(1) + return x + 1 + + fn = mx.compile(body) + x = mx.zeros((1,)) + mx.eval(fn(x)) + mx.eval(fn(x)) + return len(calls) == 1 + + +@pytest.fixture +def restore_compile(): + yield + if "MLX_DISABLE_COMPILE" not in os.environ: + mx.enable_compile() + + +def test_off_on_entry_from_the_env_var_stays_off_on_exit(restore_compile, monkeypatch): + # Bug caught: the block switches compile back on at exit under MLX_DISABLE_COMPILE, so the rest of the session + # runs compiled after the user turned compile off. restore_compile comes first so it tears down last, after + # monkeypatch has removed the variable, and switches compile back on for the next test. + monkeypatch.setenv("MLX_DISABLE_COMPILE", "1") + mx.disable_compile() # what MLX does at start when the variable is set + with mflux_uncompiled(): + assert not _compile_is_on() + assert not _compile_is_on() + + +def test_a_nested_block_leaves_the_outer_block_uncompiled(restore_compile): + # Bug caught: an inner block's exit switches compile on while the outer block still needs the plain graph. + with mflux_uncompiled(): + with mflux_uncompiled(): + pass + assert not _compile_is_on() + assert _compile_is_on() + + +def test_on_on_entry_is_on_again_on_exit(restore_compile): + # Bug caught: the block never switches compile off, or leaves it off for every later test. + assert _compile_is_on() + with mflux_uncompiled(): + assert not _compile_is_on() + assert _compile_is_on() diff --git a/tests/test_qwen21_transformer.py b/tests/test_qwen21_transformer.py index 13aad3b..9993ec1 100644 --- a/tests/test_qwen21_transformer.py +++ b/tests/test_qwen21_transformer.py @@ -85,14 +85,18 @@ def count(group): @pytest.mark.mflux def test_the_stock_transformer_compiles_its_forward_on_the_first_call(): - # Bug caught (the control): mflux no longer compiling _forward in __call__ (qwen21_transformer.py:69-72), which + # Bug caught (the control): mflux 0.23 no longer compiling its default (text-cache) image step in __call__ + # (qwen21_transformer.py:111-113), which # would leave the guard and its note unjustified. mflux compiles there on every chip, no chip check. from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer import Qwen21Transformer tf = Qwen21Transformer(**TINY) assert tf._step_fn is None mx.eval(tf(**tiny_inputs())) - assert type(tf._step_fn) is COMPILED + assert type(tf._image_step_fn) is COMPILED + assert ( + tf._step_fn is None + ) # the joint pass runs only for padded prompts or with the text cache off @pytest.mark.mflux @@ -105,9 +109,9 @@ def test_the_seamed_transformer_runs_the_plain_forward(): tf, shapes = tiny_seamed() tf.attach(ResidentProvider(_constant_weights(shapes)), shapes) mx.eval(tf(**tiny_inputs())) - assert type(tf._step_fn) is types.MethodType - assert tf._step_fn.__func__ is Qwen21Transformer._forward - assert tf._step_fn.__self__ is tf + assert type(tf._image_step_fn) is types.MethodType + assert tf._image_step_fn.__func__ is Qwen21Transformer._image_forward + assert tf._image_step_fn.__self__ is tf assert is_eager(tf) @@ -163,35 +167,44 @@ def test_the_compiled_forward_cannot_run_the_seam(): @pytest.mark.mflux def test_the_swap_keeps_parameter_paths_and_the_computed_tables(): # Bug caught: the class swap breaking on Qwen21TransformerBlock, parameters moving under another path - # (placeholders, load_weights and coverage would miss them), or a computed table renamed upstream (the extras - # coverage exemption would go stale silently). + # (placeholders, load_weights and coverage would miss them), a computed table renamed or moved upstream + # (computed_tables raises), a computed table turned back into a parameter (the extras coverage would then refuse + # every checkpoint: none stores it), or the swap replacing a table mflux built. from mflux.models.qwen21.model.qwen21_transformer.qwen21_transformer_block import ( Qwen21TransformerBlock, ) - from mlx_dfloat.integrate.blockseam import seam_blocks - from mlx_dfloat.mflux.qwen21.transformer import COMPUTED_PARAMS + from mlx_dfloat.mflux.qwen21.transformer import ( + computed_tables, + mflux_packer, + seam_qwen21_blocks, + ) tf = seam_transformer_class()(**TINY) before = sorted(k for k, _v in tree_flatten(tf.parameters())) - assert seam_blocks(block_lists(tf), tf.seam_cell) == ( + tables = computed_tables(tf) + assert seam_qwen21_blocks( + block_lists(tf), tf.seam_cell, tf.pack_switch, pack=mflux_packer() + ) == ( "transformer_blocks.0", "transformer_blocks.1", ) after = sorted(k for k, _v in tree_flatten(tf.parameters())) assert after == before assert isinstance(tf.transformer_blocks[1], Qwen21TransformerBlock) - # The seven names mflux 0.20.0 computes at construction (qwen21_rope.py:21-26, qwen21_time_text_embed.py:16). - assert set(COMPUTED_PARAMS) == { - "pos_embed.cos_tables.0", - "pos_embed.cos_tables.1", - "pos_embed.cos_tables.2", - "pos_embed.sin_tables.0", - "pos_embed.sin_tables.1", - "pos_embed.sin_tables.2", - "time_text_embed.time_proj.freqs", + # The seven tables mflux 0.23.0 computes at construction, as private lists and a private array + # (qwen21_rope.py:21-26, qwen21_time_text_embed.py:16): MLX leaves underscore attributes out of the parameters. + assert set(tables) == { + "pos_embed._cos_tables.0", + "pos_embed._cos_tables.1", + "pos_embed._cos_tables.2", + "pos_embed._sin_tables.0", + "pos_embed._sin_tables.1", + "pos_embed._sin_tables.2", + "time_text_embed.time_proj._freqs", } - assert set(COMPUTED_PARAMS) <= set(after) + assert not [k for k in after if k.startswith("pos_embed.") or "time_proj" in k] + assert all(computed_tables(tf)[k] is v for k, v in tables.items()) def _eager_reference(tf, shapes, weights): @@ -210,6 +223,7 @@ def _eager_reference(tf, shapes, weights): idx = int(block.partition(".")[2]) for attr, w in per.items(): get_attr_path(plain.transformer_blocks[idx], attr).weight = w + plain.use_text_cache = False # the joint pass, as the seamed side runs it here plain._step_fn = plain._forward return plain @@ -221,6 +235,7 @@ def test_a_forward_runs_each_block_with_its_own_weights_and_ends_on_placeholders # stock mflux holding the same weights; a block output not evaluated at its boundary; or the placeholders not # restored afterwards. tf, shapes = tiny_seamed() + tf.use_text_cache = False # the joint pass: one array per block (the text-cache path is in test_qwen21_mflux023.py) weights = _random_weights(shapes) evals = [] real_eval = seam._eval @@ -257,6 +272,7 @@ def test_depth2_drains_the_last_block(monkeypatch): # look-ahead evaluating in the wrong order. Expected (seam._seam_eval + end_step): async r0, async r1, eval r0, # eval r1. tf, shapes = tiny_seamed() + tf.use_text_cache = False # the joint pass: one output per block events = [] real_eval, real_async = seam._eval, seam._async_eval @@ -285,7 +301,8 @@ def recording_async(*args): @pytest.mark.mflux def test_two_calls_per_step_decode_once_per_block_per_call(): # Bug caught: the second (negative-prompt) call of a CFG step (qwen_image_21.py:104-118) refused for the first - # call's unchecked status words, or a block decoded twice in one call. 2 calls x 2 blocks = 4 decodes. + # call's unchecked status words, or a block decoded more than once per pass. mflux 0.23's default path runs a + # text pass (a new prompt's K/V cache) then the image pass in each call: 2 calls x 2 passes x 2 blocks = 8. tf, shapes = tiny_seamed() calls = [] tf.attach(_counting_provider(shapes, calls), shapes, verify_in_call=True) @@ -294,6 +311,6 @@ def test_two_calls_per_step_decode_once_per_block_per_call(): pos = tf(**positive) neg = tf(**negative) mx.eval(pos, neg) - assert calls == BLOCKS * 2 + assert calls == BLOCKS * 2 * 2 assert pos.shape == neg.shape == (1, 16, 64) assert not mx.array_equal(pos, neg).item() diff --git a/tests/test_verify_image.py b/tests/test_verify_image.py index d802fec..66dad9f 100644 --- a/tests/test_verify_image.py +++ b/tests/test_verify_image.py @@ -1,5 +1,6 @@ import json import sys +from contextlib import contextmanager import mlx.core as mx import pytest @@ -752,6 +753,103 @@ def test_qwen_embeds_hold_the_positive_pair_and_the_negative_one_when_cfg_ran(): assert two["prompt_embeds"] is pos[0] +class _QwenScheduler: + """mflux's scheduler surface for the loop body: the scale is the identity, a step adds the noise.""" + + def scale_model_input(self, latents, t): + return latents + + def step(self, *, noise, timestep, latents): + return latents + noise + + +class _QwenConfig: + def __init__(self, steps, guidance=1.0): + self.time_steps = list(range(steps)) + self.scheduler = _QwenScheduler() + self.guidance = guidance + + +class _UnpackedQwenTransformer: + """A transformer without mflux 0.23's packing window (mflux 0.20, or any family without it). + + Its noise is the prompt embedding plus 1, so the positive and negative calls give different noise. + """ + + def __call__( + self, *, t, config, hidden_states, encoder_hidden_states, encoder_hidden_states_mask + ): + return mx.zeros_like(hidden_states) + encoder_hidden_states + self.offset() + + def offset(self): + return 1.0 + + +class _PackingQwenTransformer(_UnpackedQwenTransformer): + """mflux 0.23's surface: ``inference_projections()`` switches the packed projections on for its window. + + Inside the window the noise is 10 higher, so a loop run outside it gives other latents. + """ + + def __init__(self): + self.packed = False + self.windows = 0 + + def offset(self): + return 11.0 if self.packed else 1.0 + + @contextmanager + def inference_projections(self): + self.windows += 1 + self.packed = True + try: + yield + finally: + self.packed = False + + +def _qwen_denoise(transformer, config, *, negative=None): + latents = mx.zeros((1, 2, 3)) + prompt = (mx.zeros((1, 1, 3)), mx.ones((1, 1), dtype=mx.int32)) + neg = (None, None) if negative is None else (mx.full((1, 1, 3), negative), prompt[1]) + return vi.qwen21_denoise( + transformer, + config, + latents, + prompt_embeds=prompt[0], + prompt_mask=prompt[1], + negative_prompt_embeds=neg[0], + negative_prompt_mask=neg[1], + ) + + +def test_the_qwen_bf16_loop_runs_inside_the_transformers_packing_window(): + # Bug caught: the bf16 side's loop body run without `with transformer.inference_projections():`. On mflux 0.23 + # its blocks would then run unpacked while the df11 side (mflux's own generate_image) runs packed, and the + # identity check would compare two different computations. Two steps of noise 11 (packed) give 22; outside the + # window they give 2. + transformer = _PackingQwenTransformer() + latents = _qwen_denoise(transformer, _QwenConfig(2)) + assert latents.tolist() == [[[22.0, 22.0, 22.0], [22.0, 22.0, 22.0]]] + assert transformer.windows == 1 + assert transformer.packed is False # the window is closed before the VAE decode + + +def test_the_qwen_bf16_loop_runs_a_transformer_without_a_packing_window_unchanged(): + # Bug caught: the window required (an AttributeError on mflux 0.20's transformer), or entered on some other + # object. Two steps of noise 1 give 2. + latents = _qwen_denoise(_UnpackedQwenTransformer(), _QwenConfig(2)) + assert latents.tolist() == [[[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]] + + +def test_the_qwen_bf16_loop_combines_the_negative_call_with_mfluxs_guidance(): + # Bug caught: CFG lost or its combine flipped in the loop body. mflux qwen_image_21.py: neg + g * (pos - neg). + # Positive noise 11 (prompt 0, packed), negative 15 (prompt 4, packed), guidance 3: 15 + 3 * (11 - 15) = 3 per + # step, 6 after two steps. + latents = _qwen_denoise(_PackingQwenTransformer(), _QwenConfig(2, guidance=3.0), negative=4.0) + assert latents.tolist() == [[[6.0, 6.0, 6.0], [6.0, 6.0, 6.0]]] + + # --- ERNIE-Image ---------------------------------------------------------------------------------------------- ERNIE = "ernie-image" diff --git a/tests/test_zimage_memory.py b/tests/test_zimage_memory.py index c836d1b..662107d 100644 --- a/tests/test_zimage_memory.py +++ b/tests/test_zimage_memory.py @@ -3,10 +3,12 @@ Constants come from ``derisk-{turbo,base}/derisk.json`` (2026-10-08, git 3da18bb, mlx 0.32.2, mflux 0.20.0): overhead max(510_030_936, 470_479_960); denoise activation max(12_781_477_784 - 11_583_577_874, 13_061_987_888 - 11_538_787_646) = 1_523_200_242. The VAE transient comes from the measured 1024² generate run of the base model -(bench/results/tiers/z-image-1024.json): footprint peak 14_115_020_240 minus (8_369_046_600 + 6_088_832 + -19_660_800 + 167_666_902 + 510_030_936) = 5_042_526_170. +(bench/results/tiers/z-image-1024.json, mflux 0.23): footprint peak 15_462_030_800 minus (8_369_046_600 + 6_088_832 + +19_660_800 + 167_666_902 + 510_030_936) = 6_389_536_730, the highest of every committed 1024² sample. """ +import json +from pathlib import Path from types import SimpleNamespace import numpy as np @@ -37,7 +39,7 @@ def test_the_constants_are_the_derisk_measurements(): # Bug caught: a constant taken from one unit instead of the larger of the two, or a FLUX.1 value left in. assert OVERHEAD_BYTES == 510_030_936 - assert VAE_TRANSIENT_BYTES == 5_042_526_170 + assert VAE_TRANSIENT_BYTES == 6_389_536_730 assert DENOISE_ACTIVATION_AT_REFERENCE == 1_523_200_242 @@ -78,9 +80,9 @@ def test_the_allowance_scales_with_tokens_and_has_a_floor(): def test_the_vae_transient_is_floored_at_the_measurement_and_linear_above_1024_squared(): # Bug caught: a smaller image predicting less than the measured 1024² value, or no growth above it. - assert vae_transient_bytes(height=512, width=512) == 5_042_526_170 - assert vae_transient_bytes(height=1024, width=1024) == 5_042_526_170 - assert vae_transient_bytes(height=2048, width=1024) == 10_085_052_340 # 2 x 5_042_526_170 + assert vae_transient_bytes(height=512, width=512) == 6_389_536_730 + assert vae_transient_bytes(height=1024, width=1024) == 6_389_536_730 + assert vae_transient_bytes(height=2048, width=1024) == 12_779_073_460 # 2 x 6_389_536_730 def test_the_denoise_activation_is_linear_in_the_token_count(): @@ -157,7 +159,7 @@ def test_fit_for_sums_the_phases_and_drops_the_set_from_the_vae_phase_when_asked # Bug caught: vae_with_set=False keeping the compressed set / non-block bytes in the VAE phase or dropping the # extras (they stay on the transformer when the set goes), or the constants not the measured ones. Hand sums (denoise at 1024², 512 tokens): # 8_000_000_000 + 6_000_000 + 20_000_000 + 360_000_000 + 2_300_000_000 + 1_523_200_242 + 510_030_936 - # = 12_719_231_178 ; vae with set: 8_026_000_000 + 160_000_000 + 5_042_526_170 + 510_030_936 = 13_738_557_106. + # = 12_719_231_178 ; vae with set: 8_026_000_000 + 160_000_000 + 6_389_536_730 + 510_030_936 = 15_085_567_666. sizes = ZImageSizes( compressed=8_000_000_000, extras=6_000_000, @@ -178,52 +180,63 @@ def test_fit_for_sums_the_phases_and_drops_the_set_from_the_vae_phase_when_asked } kept = fit_for(**common) assert kept.phases["denoise"] == 12_719_231_178 - assert kept.phases["vae"] == 13_738_557_106 + assert kept.phases["vae"] == 15_085_567_666 assert kept.phases["encode"] == 8_000_000_000 + 1_500_000_000 + 510_030_936 assert kept.peak_phase == "vae" dropped = fit_for(vae_with_set=False, **common) - # 6_000_000 + 160_000_000 + 5_042_526_170 + 510_030_936 - assert dropped.phases["vae"] == 5_718_557_106 + # 6_000_000 + 160_000_000 + 6_389_536_730 + 510_030_936 + assert dropped.phases["vae"] == 7_065_567_666 assert dropped.peak_phase == "denoise" assert dropped.phases["denoise"] == kept.phases["denoise"] -# The two measured 1024² generate runs (bench/results/tiers/z-image{,-turbo}-1024.json, 2026-10-08): the report's -# own sizes and cache limit, its footprint peak (reached in the VAE phase), and the in-flight decoded group the -# report's denoise estimate implies: 13_101_538_864 - (8_369_046_600 + 6_088_832 + 19_660_800 + 2_311_752_734 -# + 1_523_200_242 + 510_030_936) = 361_759_720 (the turbo report gives the same). -_MEASURED_1024 = { - "z-image": ( - ZImageSizes( - compressed=8_369_046_600, - extras=6_088_832, - nonblock=19_660_800, - encoders=8_044_982_000, - vae=167_666_902, - ), - 14_115_020_240, +REPO = Path(__file__).resolve().parents[1] +_BASE_SIZES = ZImageSizes( + compressed=8_369_046_600, + extras=6_088_832, + nonblock=19_660_800, + encoders=8_044_982_000, + vae=167_666_902, +) + + +def _record_sample(path): + """A committed 1024² record's sizes and its watched peak (verify_image's df11 side: the report and its footprint).""" + record = json.loads((REPO / path).read_text()) + if "report" in record: + return ZImageSizes(**record["report"]["sizes"]), record["footprint_peak_bytes"] + return ZImageSizes(**record["sizes"]), record["watched_peak_bytes"] + + +# Every committed 1024² sample with the set resident through the VAE decode, each peak reached in the VAE phase: +# the records in the tree (mflux 0.23: Z-Image's MEASURED run and identity df11 side; mflux 0.20: Z-Image-Turbo's +# MEASURED run and its five repeats) and the two mflux 0.20 Z-Image records they replaced (git ed2b0ca, the 0.2.0 tag: +# bench/results/tiers/z-image-1024.json watched_peak_bytes and identity/z-image-1024/df11-result.json +# footprint_peak_bytes, the same sizes as the 0.23 records). +_SAMPLES_1024 = { + "z-image-measured-mflux023": _record_sample("bench/results/tiers/z-image-1024.json"), + "z-image-identity-mflux023": _record_sample( + "bench/results/identity/z-image-1024/df11-result.json" ), - "z-image-turbo": ( - ZImageSizes( - compressed=8_374_285_852, - extras=6_088_832, - nonblock=19_660_800, - encoders=8_044_982_000, - vae=167_666_902, - ), - 14_074_273_664, + "z-image-turbo-measured-mflux020": _record_sample( + "bench/results/tiers/z-image-turbo-1024.json" ), + **{ + f"z-image-turbo-repeat-{n}-mflux020": _record_sample( + f"bench/results/repeats/z-image-turbo-1024/turbo-run-{n}.json" + ) + for n in range(1, 6) + }, + "z-image-measured-mflux020": (_BASE_SIZES, 14_115_020_240), + "z-image-identity-mflux020": (_BASE_SIZES, 14_066_196_160), } -@pytest.mark.parametrize("model", sorted(_MEASURED_1024)) -def test_the_fit_model_reproduces_its_measured_1024_peaks_in_the_vae_phase(model): - # Bug caught: a fit model that predicts less than its own calibration runs measured (12.16 GiB VAE phase - # against 13.15 GiB measured, and the peak placed in the denoise phase), so a tier between the two would be - # told it fits. The VAE phase must reach the larger measured peak (base: exactly; turbo, 5.2 MB larger set and - # a 40.7 MB smaller peak: at most 64 MiB above it) and be the peak phase. - sizes, measured = _MEASURED_1024[model] - est = fit_for( +def _fit_1024(sizes): + # The 1024² working point of the measured runs: their cache limit and the in-flight decoded group the reports' + # denoise estimate implies, 13_101_538_864 - (8_369_046_600 + 6_088_832 + 19_660_800 + 2_311_752_734 + # + 1_523_200_242 + 510_030_936) = 361_759_720 (the turbo report gives the same). + return fit_for( sizes=sizes, largest={"layers": 361_759_720}, policy="per-block", @@ -234,12 +247,28 @@ def test_the_fit_model_reproduces_its_measured_1024_peaks_in_the_vae_phase(model width=1024, text_tokens=512, ) - assert measured <= est.phases["vae"] <= measured + 64 * 1024**2 - if model == "z-image": - assert est.phases["vae"] == 14_115_020_240 + + +@pytest.mark.parametrize("sample", sorted(_SAMPLES_1024)) +def test_the_vae_phase_never_under_predicts_a_committed_1024_sample(sample): + # Bug caught: a VAE transient calibrated on a subset of the samples, so a tier whose budget falls between the + # estimate and a measured peak is told the set can stay through the decode. On mflux 0.23 the 0.20 constant + # puts Z-Image's VAE phase 1_347_010_560 B under its MEASURED run and 970_653_696 B under its identity run, and + # it was already 1_188_073_004 B under the 0.20 Turbo repeat 4. + sizes, measured = _SAMPLES_1024[sample] + est = _fit_1024(sizes) + assert measured <= est.phases["vae"] assert est.peak_phase == "vae" +def test_the_vae_phase_equals_the_binding_sample(): + # Bug caught: a transient padded beyond the overall-peak rule (more set drops than the samples justify), or taken + # from another sample than the highest: the mflux 0.23 MEASURED Z-Image run binds, exactly. + sizes, measured = _SAMPLES_1024["z-image-measured-mflux023"] + assert measured == 15_462_030_800 + assert _fit_1024(sizes).phases["vae"] == 15_462_030_800 + + def _checkpoint(tmp_path): probe = FakeZImageFull(None, n_refiner_layers=1, n_layers=2) shapes = install_placeholders(block_lists(probe), ZIMAGE_FULL_TABLE) diff --git a/tests/test_zimage_model.py b/tests/test_zimage_model.py index 1fb61f8..18b7c9e 100644 --- a/tests/test_zimage_model.py +++ b/tests/test_zimage_model.py @@ -32,10 +32,22 @@ def test_every_attribute_upstreams_generate_image_reads_exists_on_the_model(tmp_ assert [n for n in sorted(names) if not hasattr(model, n)] == [] assert isinstance(model, ZImage) assert (model.bits, model.lora_paths, model.lora_scales) == (None, [], []) + # mflux 0.23 records both in the image metadata; ours are mflux's defaults (bfloat16 stream, no compute precision). + assert model.float32 is False + assert model.transformer.float32 is False + assert model.compute_precision.dtype is None @pytest.mark.parametrize( - "kwargs", [{"quantize": 8}, {"lora_paths": ["a"]}, {"lora_scales": [1.0]}, {"bake_lora": False}] + "kwargs", + [ + {"quantize": 8}, + {"lora_paths": ["a"]}, + {"lora_scales": [1.0]}, + {"bake_lora": False}, + {"float32": True}, + {"compute_precision": mx.float16}, + ], ) def test_constructor_refusals_come_before_any_resolution(monkeypatch, kwargs): # Bug caught: resolving (downloading 8 GB) before refusing an option this path cannot honour. @@ -102,6 +114,44 @@ def fake_resolve(spec, *, patterns, revision=None): ] +def test_the_default_base_resolves_at_its_pin_and_a_user_base_unpinned(tmp_path, monkeypatch): + # Bug caught: the base pin not reaching the resolver (the text encoder, VAE and tokenizer following whatever lands + # on main instead of the snapshot the recorded runs used), a pin taken from the other model's entry, or a pin + # applied to a user's own --base. + from mlx_dfloat import _metal_decode + from mlx_dfloat.mflux._hub import ResolvedRepo + from mlx_dfloat.mflux.zimage import model as model_module + from mlx_dfloat.mflux.zimage.model import DFloatZImage + + class StopError(Exception): + pass + + base_calls = [] + + def stop_at_base(spec, *, patterns, revision=None): + if patterns == zinit.DF11_PATTERNS: + return ResolvedRepo(root=tmp_path, repo_id=spec, revision=revision) + base_calls.append((spec, revision)) + raise StopError + + monkeypatch.setattr(_metal_decode, "ensure_canary", lambda **kwargs: None) + monkeypatch.setattr(model_module, "open_checkpoint", lambda root: None) + monkeypatch.setattr(model_module, "check_zimage_groups", lambda ckpt: None) + monkeypatch.setattr(zinit, "resolve", stop_at_base) + for kwargs in ( + {"model": "z-image"}, + {"model": "z-image-turbo"}, + {"model": "z-image-turbo", "base_path": "me/base"}, + ): + with pytest.raises(StopError): + DFloatZImage(**kwargs) + assert base_calls == [ + ("Tongyi-MAI/Z-Image", "04cc4abb7c5069926f75c9bfde9ef43d49423021"), + ("Tongyi-MAI/Z-Image-Turbo", "f332072aa78be7aecdf3ee76d5c247082da564a6"), + ("me/base", None), + ] + + def test_controlnet_unknown_models_and_the_none_policy_are_refused(monkeypatch, tmp_path): # Bug caught: "z-image-turbo-controlnet-union-2.1" (another class upstream) built as plain Turbo; policy "none" # (every decode resident) accepted, whether the check runs early (__init__) or only late (_assemble). @@ -408,8 +458,8 @@ def unforced_largest(root): def test_the_vae_decision_follows_the_budget_and_an_over_budget_call_is_refused( tmp_path, monkeypatch, caplog ): - # Bug caught: the set kept through a decode that does not fit (VAE phase 12e9 + 5_042_526_170 + 510_030_936 = - # 17_552_557_106 against a 16e9 budget), the drop made unconditional (a reload per call on a big Mac), the + # Bug caught: the set kept through a decode that does not fit (VAE phase 12e9 + 6_389_536_730 + 510_030_936 = + # 18_899_567_666 against a 16e9 budget), the drop made unconditional (a reload per call on a big Mac), the # constructor's budget ignored, or the fit rule not enforced. from mlx_dfloat.errors import DFloatResourceError from mlx_dfloat.mflux.zimage.memory import ZImageSizes @@ -420,7 +470,7 @@ def test_the_vae_decision_follows_the_budget_and_an_over_budget_call_is_refused( assert plan.drop_set_before_vae assert plan.estimate.fits assert plan.estimate.peak_phase == "denoise" - assert plan.estimate.phases["vae"] == 5_042_526_170 + 510_030_936 # the set left the phase + assert plan.estimate.phases["vae"] == 6_389_536_730 + 510_030_936 # the set left the phase roomy = fake_model(tmp_path / "r", monkeypatch, sizes=big, budget_bytes=40 * GIB) assert not roomy.plan_call(height=1024, width=1024).drop_set_before_vae refused = fake_model(tmp_path / "x", monkeypatch, sizes=big, budget_bytes=10 * GIB) diff --git a/uv.lock b/uv.lock index e31daf4..ab3e090 100644 --- a/uv.lock +++ b/uv.lock @@ -837,11 +837,13 @@ wheels = [ [[package]] name = "mflux" -version = "0.20.0" +version = "0.23.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "anyio", marker = "sys_platform == 'darwin'" }, { name = "filelock", marker = "sys_platform == 'darwin'" }, { name = "fonttools", marker = "sys_platform == 'darwin'" }, + { name = "fsspec", marker = "sys_platform == 'darwin'" }, { name = "hf-transfer", marker = "sys_platform == 'darwin'" }, { name = "huggingface-hub", marker = "sys_platform == 'darwin'" }, { name = "matplotlib", marker = "sys_platform == 'darwin'" }, @@ -865,9 +867,9 @@ dependencies = [ { name = "transformers", marker = "sys_platform == 'darwin'" }, { name = "urllib3", marker = "sys_platform == 'darwin'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7f/08/33ff2acce0f84e51f6c86b6a8d33b4e6703ef485a5d9c40a9e1ebe9af797/mflux-0.20.0.tar.gz", hash = "sha256:8fec3407bbc821cfba9cebd8ae55e35450b9429ec9b5445c9557aaeb78087de0", size = 998374, upload-time = "2026-09-21T17:32:05.081Z" } +sdist = { url = "https://files.pythonhosted.org/packages/84/a1/fd8fee0d4d64dd2de448b28fdc9d26feada981c3ac4eee2ac8bfb54b0c8c/mflux-0.23.0.tar.gz", hash = "sha256:cefddbcb7d98fab4f4ddc1e2b706d8f616547f32c9829fdc065c3edd3db86b9a", size = 1140421, upload-time = "2026-10-16T12:20:23.667Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c4/a8/c281e949682c317e6c551931bd536c495e2ac88d7b7eae2c306f9de45821/mflux-0.20.0-py3-none-any.whl", hash = "sha256:5d60a278d67289be0da2e4ba0e8dcf938dd3c676a01c9064f64e3fb10e30767d", size = 1351276, upload-time = "2026-09-21T17:32:03.425Z" }, + { url = "https://files.pythonhosted.org/packages/e9/ab/0a6be01917e696e4865959ffade1cbab8663c21c51910163655133388949/mflux-0.23.0-py3-none-any.whl", hash = "sha256:059dbcac2f921cd29df148fb8e765b55ab3dcbe94428b782c36862578942c0b0", size = 1542580, upload-time = "2026-10-16T12:20:21.397Z" }, ] [[package]] @@ -935,12 +937,14 @@ dependencies = [ [package.optional-dependencies] mflux = [ { name = "mflux", marker = "sys_platform == 'darwin'" }, + { name = "pillow" }, { name = "psutil" }, ] [package.dev-dependencies] bench = [ { name = "mflux", marker = "sys_platform == 'darwin'" }, + { name = "pillow" }, { name = "pyobjc-framework-metal", marker = "sys_platform == 'darwin'" }, ] dev = [ @@ -971,16 +975,18 @@ typecheck = [ [package.metadata] requires-dist = [ { name = "huggingface-hub", specifier = ">=0.24" }, - { name = "mflux", marker = "sys_platform == 'darwin' and extra == 'mflux'", specifier = ">=0.20,<0.21" }, + { name = "mflux", marker = "sys_platform == 'darwin' and extra == 'mflux'", specifier = ">=0.23,<0.24" }, { name = "mlx", specifier = ">=0.31" }, { name = "numpy", specifier = ">=1.26" }, + { name = "pillow", marker = "extra == 'mflux'", specifier = "<13" }, { name = "psutil", marker = "extra == 'mflux'", specifier = ">=6" }, ] provides-extras = ["mflux"] [package.metadata.requires-dev] bench = [ - { name = "mflux", marker = "sys_platform == 'darwin'", specifier = ">=0.20,<0.21" }, + { name = "mflux", marker = "sys_platform == 'darwin'", specifier = ">=0.23,<0.24" }, + { name = "pillow", specifier = "<13" }, { name = "pyobjc-framework-metal", marker = "sys_platform == 'darwin'", specifier = ">=10" }, ] dev = [ @@ -1335,29 +1341,81 @@ sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7 wheels = [ { url = "https://files.pythonhosted.org/packages/fb/c8/0a78b0e02d7ac54bc03e5321c9220da52f0c2ea83b21f7c40e7f3169c502/pillow-12.3.0-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:00808c5e14ef63ac5161091d242999076604ff74b883423a11e5d7bbb38bf756", size = 5392415, upload-time = "2026-07-01T11:53:47.162Z" }, { url = "https://files.pythonhosted.org/packages/b2/5b/a02d30018abd97ced9f5a6c63d28597694a00d066516b9c1c6de45859fc9/pillow-12.3.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:37d6d0a00072fd2948eb22bce7e1475f34569d90c87c59f7a2ec59541b77f7a6", size = 4785266, upload-time = "2026-07-01T11:53:49.079Z" }, + { url = "https://files.pythonhosted.org/packages/c8/98/766667a4be768150a202836acd9fad19c06824ca86c4286d3cf6b274964e/pillow-12.3.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bcb46e2f9feff8d06323983bd83ed00c201fdcab3d74973e7072a889b3979fcd", size = 6263814, upload-time = "2026-07-01T11:53:51.32Z" }, + { url = "https://files.pythonhosted.org/packages/3b/2d/ede717bc1144f63886c21fd349bb95860b0d1a21149ff16f2bb362b612b6/pillow-12.3.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:23d27a3e0307ec2244cc51e7287b919aa68d097504ebe19df4e76a98a3eea5bd", size = 6934408, upload-time = "2026-07-01T11:53:53.487Z" }, + { url = "https://files.pythonhosted.org/packages/a3/48/9c58b685e69d49c31af6c8eb9012055fab7e665785165c84796e2c73ce72/pillow-12.3.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:4f883547d4b7f0495ebe7056b0cc2aea76094e7a4abc8e933540f3271df27d9c", size = 6337160, upload-time = "2026-07-01T11:53:55.457Z" }, + { url = "https://files.pythonhosted.org/packages/ff/fa/dc2a5c0ba6df93f67c31d34b808b7ce440b40cdbf96f0b81cde1d1e6fa93/pillow-12.3.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:236ff70b9312fb68943c703aa842ca6a758abfa45ac187a5e7c1452e96ef72b5", size = 7045172, upload-time = "2026-07-01T11:53:57.736Z" }, + { url = "https://files.pythonhosted.org/packages/86/a5/444817a4d4c4c2417df00513086ca196f388d8f9ef40c2e4ccd1ad1af54b/pillow-12.3.0-cp311-cp311-win32.whl", hash = "sha256:10e41f0fbf1eec8cfd234b8fe17a4caac7c9d0db4c204d3c173a8f9f6ef3232b", size = 6472232, upload-time = "2026-07-01T11:53:59.767Z" }, + { url = "https://files.pythonhosted.org/packages/63/c6/4bad1b18d132a50b27e1365e1ab163616f7a5bb56d330f66f9d1d9d4f9d4/pillow-12.3.0-cp311-cp311-win_amd64.whl", hash = "sha256:8e95e1385e4998ae9694eeaa4730ba5457ff61185b3a55e2e7bea0880aef452a", size = 7233653, upload-time = "2026-07-01T11:54:02.066Z" }, + { url = "https://files.pythonhosted.org/packages/fd/16/00f91ab7760dc842f5aad55217e80fc4a7067a0604535249bc8a2d6d9870/pillow-12.3.0-cp311-cp311-win_arm64.whl", hash = "sha256:ebaea975e03d3141d9d3a507df75c9b3ec90fa9d2ffd07567b3a978d9d790b26", size = 2568195, upload-time = "2026-07-01T11:54:04.622Z" }, { url = "https://files.pythonhosted.org/packages/37/bf/fb3ebff8ddcb76aac5a01389251bbbb9519922a9b520d8247c1ca864a25d/pillow-12.3.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ba09209fbe443b4acccebe845d8a138b89a8f4fbaeedd44953490b5315d5e965", size = 5345969, upload-time = "2026-07-01T11:54:06.397Z" }, { url = "https://files.pythonhosted.org/packages/d8/66/9a386a92561f402389a4fc70c18838bf6d35eb5eb5c6850b4b2dc64f5048/pillow-12.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ffd0c5368496f41b0944be820fcb7a838aa6e623d250b01acf2643939c3f99d7", size = 4780323, upload-time = "2026-07-01T11:54:09.351Z" }, + { url = "https://files.pythonhosted.org/packages/25/27/ac8f99618ffd3dde21db0f4d4b1d2ab00c0880595bfd17df103f7f39fd0c/pillow-12.3.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d9c7f76c0673154f044e9d78c8655fb4213f6ca31a836df48b40fe5d187717b9", size = 6266838, upload-time = "2026-07-01T11:54:11.71Z" }, + { url = "https://files.pythonhosted.org/packages/84/21/a35af28dcc61f37ed850a2d64c65c701321dfbf25085e469d5559360cbbf/pillow-12.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78cb2c6865a35ab8ff8b75fd122f6033b92a62c82801110e48ddd6c936a45d91", size = 6940830, upload-time = "2026-07-01T11:54:13.732Z" }, + { url = "https://files.pythonhosted.org/packages/eb/51/8b08617af3ad95e33ce6d7dd2c99ed6c8298f7fb131636303956be022e25/pillow-12.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e491916b378fba47242221bb9ead245211b70d504f495d105d17b14a24b4907c", size = 6344383, upload-time = "2026-07-01T11:54:15.756Z" }, + { url = "https://files.pythonhosted.org/packages/1d/72/cf78ac9780bb93c28328f408973845a309d4d145041665f734572ced1b52/pillow-12.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0dd2064cbc55aaec028ef5fbb60fa47bb6c3e7918e07ff17935284b227a9d2df", size = 7052934, upload-time = "2026-07-01T11:54:17.721Z" }, + { url = "https://files.pythonhosted.org/packages/20/20/25e0f4dc178a6bc0696793720055519a0de89e7661dae886992decbd2f81/pillow-12.3.0-cp312-cp312-win32.whl", hash = "sha256:dbce0b29841537a2fa4a214c2bbf14de3587c9680caa9b4e217568472490b28f", size = 6472684, upload-time = "2026-07-01T11:54:19.839Z" }, + { url = "https://files.pythonhosted.org/packages/45/89/da2f7971a317f83d807fdd4065c0af40208e59e692cc43d315a71a0e96d1/pillow-12.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:a2b55dd6b2a4c4b7d87ffa56bdb33fdc5fdb9a462173861a7bc097f17d91cb09", size = 7227137, upload-time = "2026-07-01T11:54:22.025Z" }, + { url = "https://files.pythonhosted.org/packages/de/47/4845a0a6c0dbf1db8456bd9fc791f13c5ced7ced20606d08a0aacfd25b49/pillow-12.3.0-cp312-cp312-win_arm64.whl", hash = "sha256:331b624368d4f1d069149002f25f44bc61c8919ce8ddb3c45bdad8f6e2d89510", size = 2568267, upload-time = "2026-07-01T11:54:24.051Z" }, { url = "https://files.pythonhosted.org/packages/9d/ac/31fb64e1e7efb5a4b50cd3d92049ba89ac6e4d8d3bb6a74e15048ca3353e/pillow-12.3.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:21900ce7ba264168cd50defae43cd75d25c833ad4ad6e73ffc5596d12e25ac89", size = 4161684, upload-time = "2026-07-01T11:54:25.934Z" }, { url = "https://files.pythonhosted.org/packages/87/b4/9805e23d2b4d77842b468513841fda254ee42f0289d25088340e4ff46e2d/pillow-12.3.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:4e8c2a84d977f50b9daed6eeaf3baef67d00d5d74d932288f02cb94518ee3ace", size = 4255487, upload-time = "2026-07-01T11:54:27.935Z" }, { url = "https://files.pythonhosted.org/packages/df/39/ecf519435a200c693fe053a6ee4d835b41cf963a4dfc2551c4e637cb2a71/pillow-12.3.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:ae26d61dfa7a47befdc7572b521024e8745f3d809bd95ca9505a7bba9ef849ec", size = 3696433, upload-time = "2026-07-01T11:54:29.813Z" }, { url = "https://files.pythonhosted.org/packages/42/92/2fc3ffad878ae8dd5469ec1bc8eb83b71f48e13efdf68f02709003982a32/pillow-12.3.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:7a743ff716f746fc19a9557f60dab1600d4613255f8a7aeb3cdde4db7eb15a66", size = 5345889, upload-time = "2026-07-01T11:54:31.97Z" }, { url = "https://files.pythonhosted.org/packages/10/76/8803c13605b763d33d156c4678fc77f8443389c0c51c8aef707bb02015f4/pillow-12.3.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d69141514cc30b774ceea5e3ed3a6635c8d8a96edf664689b890f4089111fb35", size = 4780109, upload-time = "2026-07-01T11:54:34.026Z" }, + { url = "https://files.pythonhosted.org/packages/1f/01/e18aff37cb0b4aac47ac90f016d347a49aca667ef97f190b06ac2aabc928/pillow-12.3.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f7401aebd7f581d7f83a439d87d474999317ee099218e5ad25d125290990ba65", size = 6263736, upload-time = "2026-07-01T11:54:36.131Z" }, + { url = "https://files.pythonhosted.org/packages/f7/62/de5bdd77d935331f4f802edc11e4d82950f642caad6cb2f949837b8560e2/pillow-12.3.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0847a763afefb695bc912d7c131e7e0632d4edc1d8698f58ddabec8e46b8b6d3", size = 6937129, upload-time = "2026-07-01T11:54:38.216Z" }, + { url = "https://files.pythonhosted.org/packages/70/4d/105627a13300c5e0df1d174230b32fd1273062c96f7745fd552b945d1e1d/pillow-12.3.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:571b9fcb07b97ef3a492028fb3d2dc0993ca23a06138b0315286566d29ef718a", size = 6339562, upload-time = "2026-07-01T11:54:40.354Z" }, + { url = "https://files.pythonhosted.org/packages/6b/1d/f13de01a553988ab895ba1c722e06cf3144d4f57656fd5b81b6d881f1179/pillow-12.3.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:756c768d0c9c2955feb7a56c37ea24aea2e369f8d36a88da270b6a9f19e62b5e", size = 7049439, upload-time = "2026-07-01T11:54:42.489Z" }, + { url = "https://files.pythonhosted.org/packages/c9/f9/066794cca041b969964f779ee5fa66a9498bbf34248ac39c5d7954e4198f/pillow-12.3.0-cp313-cp313-win32.whl", hash = "sha256:a876864214e136f0eb367788dbd7df045f4806801518e2cfe9e13229cfe06d8f", size = 6473287, upload-time = "2026-07-01T11:54:44.9Z" }, + { url = "https://files.pythonhosted.org/packages/a6/9b/7a58e61d62be561da3a356fe2384d4059a6345fc130e23ef1c36a5b81d24/pillow-12.3.0-cp313-cp313-win_amd64.whl", hash = "sha256:1cca606cd25738df4ed873d5ad46bbdb3d83b5cbca291f6b4ff13a4df6b0bbe8", size = 7239691, upload-time = "2026-07-01T11:54:47.141Z" }, + { url = "https://files.pythonhosted.org/packages/aa/b0/c4ed4f0ef8f8fa5ee8351537db6650bb8189f7e118842978dd6589065692/pillow-12.3.0-cp313-cp313-win_arm64.whl", hash = "sha256:b629de27fda84b42cde7edef0d85f13b958b47f6e9bbcbba9b673c562a89bd8b", size = 2568185, upload-time = "2026-07-01T11:54:49.137Z" }, { url = "https://files.pythonhosted.org/packages/dc/01/001f65b68192f0228cc1dbbc8d2530ab5d58b61037ba0587f946fea607cd/pillow-12.3.0-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:9cf95fe4d0f84c82d282745d9bb08ad9f926efa00be4697e767b814ce40d4330", size = 4161736, upload-time = "2026-07-01T11:54:51.156Z" }, { url = "https://files.pythonhosted.org/packages/1a/d2/0219746d0fd16fc8a84498e79452375be3797d3ce4044596ce565164b84f/pillow-12.3.0-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:8728f216dcdb6e6d555cf971cb34076139ad74b31fc2c14da4fafc741c5f6217", size = 4255435, upload-time = "2026-07-01T11:54:53.414Z" }, { url = "https://files.pythonhosted.org/packages/c8/02/8d0bc62ef0302318c46ff2a512822d2610e81c7aa46c9b3abe6cbaca5ad0/pillow-12.3.0-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:a45650e8ce7fafffd731db8550230db6b0d306d181a90b67d3e6bca2f1990930", size = 3696262, upload-time = "2026-07-01T11:54:55.739Z" }, { url = "https://files.pythonhosted.org/packages/85/e2/73c77d218410b14f5f2d565e8a998d5317b7b9c75368d29985139f7a46f0/pillow-12.3.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:ba54cfebe86920a559a7c4d6b9050791c20513650a1952ebe3368c7dc70306f8", size = 5350344, upload-time = "2026-07-01T11:54:57.657Z" }, { url = "https://files.pythonhosted.org/packages/c7/da/32c752228ae345f489e3a42499d817b6c3996da7e8a3bc7a04fc806b243b/pillow-12.3.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:e158cb00350dc278f3b91551101aa7d12415a66ebf2c91d8d5ac14e56ddd3ad0", size = 4780131, upload-time = "2026-07-01T11:54:59.713Z" }, + { url = "https://files.pythonhosted.org/packages/b1/9d/8b2c807dbef61a5197c047afe99823787eb66f63daf9fb2432f91d6f0462/pillow-12.3.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e9aeb04d6aef139de265b29683e119b638208f88cf73cdd1658aa07221165321", size = 6263757, upload-time = "2026-07-01T11:55:01.778Z" }, + { url = "https://files.pythonhosted.org/packages/5c/44/c85361f65dbe00eea8576ee467c768d25129989efb76e94f205e9ca9bb46/pillow-12.3.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:251bf95b67017e27b13d82f5b326234ca62d70f9cf4c2b9032de2358a3b12c7b", size = 6936962, upload-time = "2026-07-01T11:55:03.93Z" }, + { url = "https://files.pythonhosted.org/packages/18/7e/e483414b35800b86b6f08dbbc7803fb5cd52c4d6f897f47d53ea2c7e6f65/pillow-12.3.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:fe3cca2e4e8a592be0f269a1ca4835c25199d9f3ce815c8491048f785b0a0198", size = 6339171, upload-time = "2026-07-01T11:55:05.989Z" }, + { url = "https://files.pythonhosted.org/packages/f0/f4/68c491844841ede6bed70189546b3ee9731cf9f2cbad396faff5e1ccba45/pillow-12.3.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:23aceaa007d6172b02c277f0cd359c79492bbb14f7072b4ede9fbcaf20648130", size = 7048116, upload-time = "2026-07-01T11:55:08.131Z" }, + { url = "https://files.pythonhosted.org/packages/a3/34/77f3f793fed8efc7d243f21b33c5a3f0d1c97ee70346d3db855587e155ff/pillow-12.3.0-cp314-cp314-win32.whl", hash = "sha256:af8d94b0db561cf68b88a267c5c44b49e134f525d0dc2cb7ed413a66bc23559a", size = 6467209, upload-time = "2026-07-01T11:55:10.408Z" }, + { url = "https://files.pythonhosted.org/packages/f1/e0/492879f69d94f91f60fc8cd05ba03650e9520afebb2fb7aa12777d7c7f38/pillow-12.3.0-cp314-cp314-win_amd64.whl", hash = "sha256:fdafc9cce40277e0f7a0feabce0ee50dd2fa1800f3b38015e51296b5e814048d", size = 7237707, upload-time = "2026-07-01T11:55:12.745Z" }, + { url = "https://files.pythonhosted.org/packages/c9/ac/6b11f2875f1c2ac040d84e1bbf9cf22a88038f901ca1037898b280b38365/pillow-12.3.0-cp314-cp314-win_arm64.whl", hash = "sha256:e91206ee562682b51b98ef4b26a6ef48fd84e15fd4c4bc5ec768eb641d206838", size = 2565995, upload-time = "2026-07-01T11:55:14.736Z" }, { url = "https://files.pythonhosted.org/packages/52/69/c2208e56af9bfc1913afb24020297a691eb1d4ef688474c8a04913f65e04/pillow-12.3.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:164b31cd1a0490ab6efae01aa5df49da7061be0af1b30e035b6e9a1bfe34ee6e", size = 5352503, upload-time = "2026-07-01T11:55:17.076Z" }, { url = "https://files.pythonhosted.org/packages/07/70/e5686d753e898a45d778ff1718dba8516ead6ab6b95d85fc8c4b70650cf2/pillow-12.3.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5afb51d599ea772b8365ae807ae557f18bccfe46ab261fd1c2a9ed700fc6eb17", size = 4782956, upload-time = "2026-07-01T11:55:19.448Z" }, + { url = "https://files.pythonhosted.org/packages/d5/37/25c6692f06927ee973ff18c8d9ee98ad0b4d84ee67a09610c2dd1447958e/pillow-12.3.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3edce1d53195db527e0191f84b71d02022de0540bf43a16ed734ed7537b07385", size = 6322855, upload-time = "2026-07-01T11:55:21.613Z" }, + { url = "https://files.pythonhosted.org/packages/cc/91/420637fcb8f1bc11029e403b4538e6694744428d8246118e45719f944556/pillow-12.3.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bf16ba1b4d0b6b7c8e534936632270cf70eb00dbe09005bc345b2677b726855c", size = 6989642, upload-time = "2026-07-01T11:55:24.006Z" }, + { url = "https://files.pythonhosted.org/packages/10/08/b94d7811281ccf0d143a1cf768d1c49e1e54af63e7b708ab2ee3eb87face/pillow-12.3.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:24870b09b224f7ae3c39ed07d10e819d06f8720bc551847b1d623832b5b0e28d", size = 6391281, upload-time = "2026-07-01T11:55:26.252Z" }, + { url = "https://files.pythonhosted.org/packages/d2/87/24233f785f55474dc02ce3e739c5528a77e3a862e9333d1dd7a25cc31f70/pillow-12.3.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:30f2aa603c41533cc25c05acd0da21636e84a315768feb631c937177db558931", size = 7096716, upload-time = "2026-07-01T11:55:28.318Z" }, + { url = "https://files.pythonhosted.org/packages/23/26/fcb2f6e37175b04f53570b59937867e2b80ee1685e744023153028fc14f9/pillow-12.3.0-cp314-cp314t-win32.whl", hash = "sha256:4b0a7fe987b14c31ebda6083f74f22b561fd3739bc0ac51e019622e3d72668c7", size = 6474125, upload-time = "2026-07-01T11:55:30.956Z" }, + { url = "https://files.pythonhosted.org/packages/90/de/3634abee5f1c9e13c56787b7d5517b0ba8d6de51700b95578cf338349c9f/pillow-12.3.0-cp314-cp314t-win_amd64.whl", hash = "sha256:962864dc93511324d51ddbb5b9f8731bf71675b93ca612a07441896f4688fb8c", size = 7242939, upload-time = "2026-07-01T11:55:34.044Z" }, + { url = "https://files.pythonhosted.org/packages/ce/2a/fd13f8eb24de5714a6eb444a3d67e2842c6c576e159a43793adf23051351/pillow-12.3.0-cp314-cp314t-win_arm64.whl", hash = "sha256:0740a512dc522224c77d9aa5a8d70d8b7d73fb91f2c21125d8d025d3b8990e45", size = 2567506, upload-time = "2026-07-01T11:55:35.988Z" }, { url = "https://files.pythonhosted.org/packages/5d/dc/8fdce34ec725a33c81c6ba122b904d6b9024e50ea9ac7bede62fab54506c/pillow-12.3.0-cp315-cp315-ios_13_0_arm64_iphoneos.whl", hash = "sha256:0feb2e9d6ad6c9e3c06effe9d00f3f1e618a6643273576b016f591e9315a7139", size = 4162063, upload-time = "2026-07-01T11:55:37.941Z" }, { url = "https://files.pythonhosted.org/packages/76/66/2044b9a63d3b84ff048228dfcb7cd9bf0df983e8470971bf7d4c57b693de/pillow-12.3.0-cp315-cp315-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:9e881fca225083806662a5c43d627d215f258ff43c890f831966c7d7ba9c7402", size = 4255549, upload-time = "2026-07-01T11:55:40.022Z" }, { url = "https://files.pythonhosted.org/packages/52/7e/1f67e6f4ece6b582ee4b539decbcc9f848dc245a93ed8cd7338bafef72f1/pillow-12.3.0-cp315-cp315-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:4998562bf62a445225f22e07c896bb04b35b1b1f2eb6d760584c9c51d7a5f78c", size = 3696331, upload-time = "2026-07-01T11:55:41.98Z" }, { url = "https://files.pythonhosted.org/packages/12/40/d306fc2c8e4d45d7f175c77edca7063be7b86fe7fe6e68f4353bf71d808c/pillow-12.3.0-cp315-cp315-macosx_10_15_x86_64.whl", hash = "sha256:dc624f6bc473dacdf7ef7eb8678d0d08edf15cd94fad6ae5c7d6cc67a4e4902f", size = 5350370, upload-time = "2026-07-01T11:55:44.028Z" }, { url = "https://files.pythonhosted.org/packages/dd/44/668fb1437e8ce420f62d6106eb66e44a5971602a4d794615bdf79315d82d/pillow-12.3.0-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:71d6097b330eea8fd15097780c8e89cb1a8ce7838669f48c5bacd6f663dd4701", size = 4780147, upload-time = "2026-07-01T11:55:46.073Z" }, + { url = "https://files.pythonhosted.org/packages/0c/08/93fa2e70e30a2d81547e481b6ee2bb9522117221fb1e0ce4b5df70967677/pillow-12.3.0-cp315-cp315-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:28ce87c5ab450a9dd970b52e5aca5fe63ed432d18a2eaddd1979a00a1ba24ace", size = 6273659, upload-time = "2026-07-01T11:55:48.264Z" }, + { url = "https://files.pythonhosted.org/packages/f8/6d/043e96ff814fc31a33077e4cba86082167db520c93632afdf2042febbb0c/pillow-12.3.0-cp315-cp315-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6b02afb9b97f65fbca5f31db6a2a3ba21aa93030225f150fa3f249717e938fb4", size = 6947439, upload-time = "2026-07-01T11:55:50.503Z" }, + { url = "https://files.pythonhosted.org/packages/af/92/ba71d2ee2ac0edf3fa33bd9d5ee9ee080da70b1766f3ca3934f9938ddac9/pillow-12.3.0-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:1182d52bc2d5e5d7d0949503aa7e36d12f42205dc287e4883f407b1988820d39", size = 6353577, upload-time = "2026-07-01T11:55:52.697Z" }, + { url = "https://files.pythonhosted.org/packages/0f/ce/e63064e2122923ff687c8ad792d0d736a7b3920a56a46982e81a7fdd25d6/pillow-12.3.0-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:e795b7eb908249c4e43c7c99fac7c2c75dab0c43566e37db472a355f63693d71", size = 7060394, upload-time = "2026-07-01T11:55:55.149Z" }, + { url = "https://files.pythonhosted.org/packages/54/76/a09cc3ccc8d773a7283d34c38bec1708f9e3cc932093cbc4c5e71ac4060b/pillow-12.3.0-cp315-cp315-win32.whl", hash = "sha256:57b3d78c95ba9059768b10e28b813002261d3f3dfc55cc48b0c988f625175827", size = 6467375, upload-time = "2026-07-01T11:55:57.769Z" }, + { url = "https://files.pythonhosted.org/packages/3e/03/1846c49ba3b1d5550392a4bbd06d6fb4578e1cd91a803198b5c90f5f7d53/pillow-12.3.0-cp315-cp315-win_amd64.whl", hash = "sha256:fa4ecea169a355be7a3ade2c783e2ed12f0e40d2c5621cda8b3297faf7fbb9f5", size = 7237048, upload-time = "2026-07-01T11:55:59.975Z" }, + { url = "https://files.pythonhosted.org/packages/fb/bb/89f35dcc79610423f9f195504d7def7f0d1416a711541b42867e25fe3412/pillow-12.3.0-cp315-cp315-win_arm64.whl", hash = "sha256:877c3f311ff35410f690861c4409e7ccbf0cd2f878e50628a28e5a0bb689e658", size = 2566006, upload-time = "2026-07-01T11:56:02.143Z" }, { url = "https://files.pythonhosted.org/packages/30/88/707027ba09942dfa2c28759b5c222d769290a41c6d20ea60ec250801941f/pillow-12.3.0-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:e9871b1ffbfa9656b60aeee92ed5136a5742696006fa322b29ea3d8da0ecc9cf", size = 5352509, upload-time = "2026-07-01T11:56:04.2Z" }, { url = "https://files.pythonhosted.org/packages/b0/6d/00352fa25332c2569cd387851f568cc5a4b75a9adbfb37ac4fbce4c02eec/pillow-12.3.0-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:53aa02d20d10c3d814d536aa4e5ac9b84ca0ff5a88377963b085ad6822f93e64", size = 4783167, upload-time = "2026-07-01T11:56:06.631Z" }, + { url = "https://files.pythonhosted.org/packages/13/4f/9e049dfa21af7c22427275720e2490267ba8138120add5c4c574deb69782/pillow-12.3.0-cp315-cp315t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:446c34dcc4324b084a53b705127dc15717b22c5e140ae0a3c38349d4efec071e", size = 6329237, upload-time = "2026-07-01T11:56:08.868Z" }, + { url = "https://files.pythonhosted.org/packages/36/16/cf6eeaae8d0fce8dd390a33437cf68c5d5bd73834a2bc6e2f14efda0ab45/pillow-12.3.0-cp315-cp315t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cf1845d02ad822a369a49f2bb9345b1614744267682e7a03527dc3bf6eea1777", size = 6997047, upload-time = "2026-07-01T11:56:11.379Z" }, + { url = "https://files.pythonhosted.org/packages/1e/69/dbf769bdd55f48bf5733cac28edc6364ffaa072ec9ba336266e4fe66be55/pillow-12.3.0-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:186941b6aef820ad110fb01fb06eb925374dc3a21b17e37ec9a53b250c6fe2d1", size = 6400440, upload-time = "2026-07-01T11:56:13.908Z" }, + { url = "https://files.pythonhosted.org/packages/a0/e1/ffc9cfc2eea0d178da8018e18e959301ad9d6bc9f3edb7181e748a474b97/pillow-12.3.0-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:f13c32a3abd6079a66d9526e18dad9b6d280384d49d7c54040cd57b6424041d9", size = 7105895, upload-time = "2026-07-01T11:56:16.575Z" }, + { url = "https://files.pythonhosted.org/packages/18/f0/a5595c1e8c3ae44b9828cb2f0fa8155e5095ef04d6327b8f61cf44a3df85/pillow-12.3.0-cp315-cp315t-win32.whl", hash = "sha256:1657923d2d45afb66526e5b933e5b3052e6bdea196c90d3abb2424e18c77dae8", size = 6474384, upload-time = "2026-07-01T11:56:18.855Z" }, + { url = "https://files.pythonhosted.org/packages/e4/04/62bcd9f844984c5938d3b05264a61d797a29d3e0812341a8204af70bbdee/pillow-12.3.0-cp315-cp315t-win_amd64.whl", hash = "sha256:8cd2f7bdda092d99c9fc2fb7391354f306d01443d22785d0cbfafa2e2c8bb418", size = 7243537, upload-time = "2026-07-01T11:56:21.214Z" }, + { url = "https://files.pythonhosted.org/packages/3d/68/1f3066acedf37673694a7141381d8f811ae97f30d34413d236abe7d489f1/pillow-12.3.0-cp315-cp315t-win_arm64.whl", hash = "sha256:06ff022112bc9cbf83b60f8e028d94ad87b60621706487e65f673de61610ab59", size = 2567491, upload-time = "2026-07-01T11:56:23.506Z" }, { url = "https://files.pythonhosted.org/packages/75/18/2e8b40223153ccbc60df07f9e8928dc0c76202aa4e55ae9f53962b6510d6/pillow-12.3.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:b3c777e849237620b022f7f297dd67705f9f5cf1685f09f02e46f93e92725468", size = 5302510, upload-time = "2026-07-01T11:56:25.736Z" }, { url = "https://files.pythonhosted.org/packages/46/3e/51fabf59d5ab801ceab709453d3ab6b180083496579549de4c45ced6528a/pillow-12.3.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:b343699e8308bdc51978310e1c959c584e7869cc8c40780058c87da7781a1e94", size = 4736058, upload-time = "2026-07-01T11:56:28.041Z" }, + { url = "https://files.pythonhosted.org/packages/bf/20/22fe9384b7949e25fb1293bcfc84fb82590ff4ea6b37c95b24d26d793d86/pillow-12.3.0-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fbd139c8447d25dd750ab79ee274cc5e1fe80fc56340ab10b18a195e1b6eca3e", size = 5237776, upload-time = "2026-07-01T11:56:30.263Z" }, + { url = "https://files.pythonhosted.org/packages/08/14/f6ba68107680ffa74b39985f3f30884e41318fbc4250caa423c79b4788bb/pillow-12.3.0-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e7e480451b9fa137494bccd3a7d69adbe8ac65a87d97be61e11f1b1050a5bac3", size = 5860358, upload-time = "2026-07-01T11:56:32.68Z" }, + { url = "https://files.pythonhosted.org/packages/36/54/0169bc772ec491108b62f644f8ecf1fe5d8ae5ebafde2ee2142210166903/pillow-12.3.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:04f01d28a6aaff387bf842a13be313df23ba0597a44f1a976c9feb3c6ff4711a", size = 7231786, upload-time = "2026-07-01T11:56:35.046Z" }, ] [[package]] From 2ae8dd21852158e91eb2a381a338dc0dccdadb25 Mon Sep 17 00:00:00 2001 From: Denis Ineshin Date: Wed, 28 Oct 2026 19:31:53 +0100 Subject: [PATCH 2/4] Label each tier row with its mflux version and run the CI lane on 0.23 Lock the integration lane to mflux 0.23.0 and fail it when the lock installs another release, so the lane's name cannot outlive the pin. Add an mflux column to the README tier table: each row names the version it ran on, rows not re-run on 0.23 say so, and records without a version read "not recorded". --- .github/workflows/ci.yml | 12 ++- scripts/_flux_rig.py | 3 +- scripts/verify_image.py | 84 +++++++++++----- src/mlx_dfloat/bench/table.py | 60 ++++++++++-- tests/test_bench_table.py | 19 +++- tests/test_bench_table_script.py | 5 +- tests/test_bench_table_versions.py | 151 +++++++++++++++++++++++++++++ tests/test_ci_config.py | 13 ++- tests/test_readme_bench_block.py | 120 +++++++++++++++++++++++ tests/test_verify_image.py | 9 +- 10 files changed, 425 insertions(+), 51 deletions(-) create mode 100644 tests/test_bench_table_versions.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 39e0e79..46d1d4c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -43,17 +43,19 @@ jobs: runs-on: macos-15 strategy: matrix: - # The lock pins mflux 0.20.0. A second row (another mflux release) would need the unlocked - # form instead of `uv sync --extra mflux`: `uv pip install "mflux=="` after the - # sync, then `uv run --no-sync` for the steps below. - mflux: ["0.20.0"] + # The lock pins mflux 0.23.0; the sanity step below fails when the lock installs any other + # release, so the row's name (the required check on main) cannot outlive the pin. A second + # row (another mflux release) would need the unlocked form instead of `uv sync --extra + # mflux`: `uv pip install "mflux=="` after the sync, then `uv run --no-sync` for the + # steps below. + mflux: ["0.23.0"] steps: - uses: actions/checkout@v4 with: fetch-depth: 0 - uses: astral-sh/setup-uv@v5 - run: uv sync --locked --group dev --extra mflux - - run: uv run --extra mflux --group dev python -c "import mflux, mlx.core, importlib.metadata as m; print('mflux', m.version('mflux'))" + - run: uv run --extra mflux --group dev python -c "import sys, mflux, mlx.core, importlib.metadata as m; v = m.version('mflux'); print('mflux', v); sys.exit(0 if v == '${{ matrix.mflux }}' else 'the lock installs mflux ' + v + ', the lane is named for ${{ matrix.mflux }}')" # No combined threshold: the steps after this one gate each model.py, and each init.py and transformer.py the required job omits, on its own. - name: Integration tests (mflux lane) env: diff --git a/scripts/_flux_rig.py b/scripts/_flux_rig.py index 84fdfab..dbf899e 100644 --- a/scripts/_flux_rig.py +++ b/scripts/_flux_rig.py @@ -42,7 +42,7 @@ ) from mlx_dfloat.integrate.seam import MAX_NONE_POLICY_LAUNCHING_BLOCKS, _async_eval, _eval from mlx_dfloat.mflux import require_mflux -from mlx_dfloat.mflux.flux1.names import DOUBLE_PREFIX, DROPPED_EXTRAS, SINGLE_PREFIX, flux_name_map +from mlx_dfloat.mflux.flux1.names import DOUBLE_PREFIX, SINGLE_PREFIX, flux_name_map from mlx_dfloat.mflux.flux1.transformer import ( MAX_BUILD_ACTIVE_BYTES, SeamMixin, @@ -311,7 +311,6 @@ def build_transformer( __all__ = [ "DOUBLE_PREFIX", - "DROPPED_EXTRAS", "EVAL_POLICIES", "FLUX_CACHE_LIMIT", "MAX_BUILD_ACTIVE_BYTES", diff --git a/scripts/verify_image.py b/scripts/verify_image.py index 98b02c0..9cf650f 100644 --- a/scripts/verify_image.py +++ b/scripts/verify_image.py @@ -996,11 +996,58 @@ def run_df11_qwen21(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, A } +def qwen21_denoise( + transformer: Any, + config: Any, + latents: mx.array, + *, + prompt_embeds: mx.array, + prompt_mask: mx.array, + negative_prompt_embeds: mx.array | None = None, + negative_prompt_mask: mx.array | None = None, +) -> mx.array: + """The Qwen-Image 2.1 loop body of mflux over ``config.time_steps``; returns the final latents (evaluated per step). + + Classifier-free guidance runs when ``negative_prompt_embeds`` is given: a second call on the negative pair per + step, combined as ``negative + guidance * (positive - negative)``. The whole loop runs inside the transformer's + ``inference_projections()`` window (packed Q/K/V and gate/up projections), as mflux's ``generate_image`` runs it + on the df11 side. + """ + with transformer.inference_projections(): + for t in config.time_steps: + # mflux qwen_image_21.py's step body: scale, positive call, negative call when CFG ran, combine, step, + # eval. + latents = config.scheduler.scale_model_input(latents, t) + noise = transformer( + t=t, + config=config, + hidden_states=latents, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_mask, + ) + if negative_prompt_embeds is not None: + noise_negative = transformer( + t=t, + config=config, + hidden_states=latents, + encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states_mask=negative_prompt_mask, + ) + noise = noise_negative + config.guidance * (noise - noise_negative) + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) + mx.eval(latents) + # No verify_step(): attach(..., verify_in_call=True) checks inside each call, and this provider defers + # nothing. + return latents + + def run_bf16_qwen21(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, Any]: """Stream the base's BF16 transformer shards through the seam and mflux's Qwen-Image 2.1 loop body, then the VAE. ``modulation.1`` loads from the base as a plain weight; the blocks stream as they run, the forward pass - uncompiled as on the df11 side (``EagerForward`` sets mflux's ``_step_fn`` to the plain method). + uncompiled as on the df11 side (``EagerForward`` sets mflux's ``_step_fn`` to the plain method and refuses a + compiled one on every call). The loop body (``qwen21_denoise``) runs inside the transformer's + ``inference_projections()`` window. Raises: VerifyImageError: The df11 side's saved embeddings are not present yet. @@ -1053,31 +1100,18 @@ def run_bf16_qwen21(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, A ) mx.eval(latents, prompt_embeds, prompt_mask) - for t in config.time_steps: - # mflux qwen_image_21.py's step body: scale, positive call, negative call when CFG ran, combine, step, eval. - latents = config.scheduler.scale_model_input(latents, t) - noise = transformer( - t=t, - config=config, - hidden_states=latents, - encoder_hidden_states=prompt_embeds, - encoder_hidden_states_mask=prompt_mask, - ) - if do_true_cfg: - noise_negative = transformer( - t=t, - config=config, - hidden_states=latents, - encoder_hidden_states=negative_prompt_embeds, - encoder_hidden_states_mask=negative_prompt_mask, - ) - noise = noise_negative + config.guidance * (noise - noise_negative) - latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - mx.eval(latents) - # No verify_step(): attach(..., verify_in_call=True) checks inside each call, and this provider defers nothing. + latents = qwen21_denoise( + transformer, + config, + latents, + prompt_embeds=prompt_embeds, + prompt_mask=prompt_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_prompt_mask=negative_prompt_mask, + ) - if not qtf.is_eager(transformer): - raise VerifyImageError("the bf16 side's forward pass was compiled; the df11 side's is not") + # No post-window is_eager check: the window clears both step functions on exit, so it would always pass. The + # check runs inside the window instead: the seam's EagerForward refuses a compiled step function on every call. mx.save_safetensors(str(out_dir / "latents.safetensors"), {"latents": latents}) vae = qinit.load_vae(base_root) diff --git a/src/mlx_dfloat/bench/table.py b/src/mlx_dfloat/bench/table.py index 4524c72..69018b9 100644 --- a/src/mlx_dfloat/bench/table.py +++ b/src/mlx_dfloat/bench/table.py @@ -17,6 +17,8 @@ _NOT_MEASURED = "not measured" _NOT_RECORDED = "not recorded" _STOPPED = "stopped by the watchdog" +# The bench command's flag; a report records that the run used it. +_SKIP_PREFLIGHT = "--skip-preflight" # A CAPPED row ran under the wired and memory caps mlx-dfloat installs on a Mac of that size. _TIER_CAPS_NOTE = "mlx-dfloat caps for the tier" # A CAPPED row ran under MLX's own defaults for a Mac of that size (no wired limit). @@ -50,6 +52,7 @@ class TierRow: limits_note: str source: str stopped_after_s: float | None = None # a run the watchdog stopped: its peaks are lower bounds + mflux: str | None = None # the record's versions.mflux; None when the record carries none @dataclass(frozen=True, slots=True, kw_only=True) @@ -118,9 +121,18 @@ def tier_row_from_generate_report(report: Mapping[str, Any], *, source: str) -> status="target" if watched <= ceiling else "over", limits_note=_applied_note(_need(limits, "applied", "limits")), source=source, + mflux=_mflux_version(report), ) +def _mflux_version(report: Mapping[str, Any]) -> str | None: + """The report's ``versions.mflux``; None when the report records no versions.""" + versions = report.get("versions") + if not isinstance(versions, Mapping) or versions.get("mflux") is None: + return None + return str(versions["mflux"]) + + def _applied_note(applied: Any) -> str: """The limits note for a report's ``limits.applied`` path. @@ -158,7 +170,8 @@ def tier_row_from_abort_artifact(artifact: Mapping[str, Any], *, source: str) -> watchdog); the DF11 size was never recorded, and the status says the run was stopped. On a CAPPED row the limits note comes from the context's ``limits`` (the MLX limits read back after the install): a wired limit above zero is mlx-dfloat's caps for the tier, zero is MLX's - defaults, and an older artifact without ``limits`` reads "not recorded". + defaults, and an older artifact without ``limits`` reads "not recorded". The mflux version is the context's + ``versions.mflux``; an older artifact without ``versions`` reads "not recorded" too. Raises: DFloatFormatError: The artifact has no run context, the context lacks ``model``, @@ -184,6 +197,7 @@ def tier_row_from_abort_artifact(artifact: Mapping[str, Any], *, source: str) -> limits_note=_abort_limits_note(label, context), source=source, stopped_after_s=float(_need(artifact, "elapsed", "abort artifact")), + mflux=_mflux_version(context), ) @@ -259,10 +273,32 @@ def _pct(x: float | None) -> str: _TIER_HEADER = ( "| Mac | Working set − reserve | Model | DF11 size | Peak (watched) | Peak footprint " # noqa: RUF001 - "| Peak MLX (sampled active + cache, or exact phase peak) | Label | Status | Limits | Result |" + "| Peak MLX (sampled active + cache, or exact phase peak) | Label | Status | Limits | Result | mflux |" ) +def _version_key(version: str) -> tuple[int, ...]: + """``"0.23.0"`` as ``(0, 23, 0)``, so versions compare as numbers, not strings. + + Only the first three numbers count: ``"0.23.0rc1"`` is ``(0, 23, 0)``, never above ``"0.23.0"``. + """ + return tuple(int(part) for part in re.findall(r"\d+", version)[:3]) + + +def _minor(version: str) -> str: + """``"0.20.0"`` as ``"0.20"``: the major and minor parts.""" + return ".".join(re.findall(r"\d+", version)[:2]) or version + + +def _mflux_cell(version: str | None, newest: str | None) -> str: + """The row's mflux version; a row older than the table's newest says it was not re-run on that one.""" + if version is None: + return _NOT_RECORDED + if newest is None or _version_key(version) >= _version_key(newest): + return _minor(version) + return f"{_minor(version)}, not re-run on {_minor(newest)}" + + def _peaks(r: TierRow) -> tuple[str, str, str]: """The three peak cells; a stopped run's are lower bounds, the watched one with the time it ran.""" watched, footprint = _gib(r.watched_peak_bytes), _gib(r.footprint_peak_bytes) @@ -278,14 +314,21 @@ def _peaks(r: TierRow) -> tuple[str, str, str]: def render_tier_table(rows: Sequence[TierRow]) -> str: - """Render the tier table: GiB with two decimals, one row per measured tier.""" - lines = [_TIER_HEADER, "|" + "---|" * 11] + """Render the tier table: GiB with two decimals, one row per measured tier. + + The last column is the mflux version from each row's record; a row on an older mflux than the newest row's + says it was not re-run on the newest, and a record without a version reads "not recorded". + """ + versions = [r.mflux for r in rows if r.mflux is not None] + newest = max(versions, key=_version_key) if versions else None + lines = [_TIER_HEADER, "|" + "---|" * 12] for r in rows: watched, footprint, mlx = _peaks(r) lines.append( f"| {r.mac_gb} GB | {_gib(r.ceiling_bytes)} | {_model_label(r.model)} " f"| {_gib_or_unrecorded(r.df11_bytes)} | {watched} " - f"| {footprint} | {mlx} | {r.label} | {r.status} | {r.limits_note} | `{r.source}` |" + f"| {footprint} | {mlx} | {r.label} | {r.status} | {r.limits_note} | `{r.source}` " + f"| {_mflux_cell(r.mflux, newest)} |" ) return "\n".join(lines) + "\n" @@ -309,7 +352,7 @@ def render_overhead_block( """Render each scenario's overhead line and recorded command, then the caption and cache note. ``preflight_skipped`` maps a scenario whose run skipped the launch check to the gates that - failed; its command line says so. + failed; its command carries ``--skip-preflight`` (the flag the run used) and says which gates failed. """ skipped = preflight_skipped or {} lines = [] @@ -324,7 +367,10 @@ def render_overhead_block( f"(one eval per step, same cache limit): {q8}" ) lines.append("") - command = f"Command: `{reproducers.get(key, '')}`" + reproducer = reproducers.get(key, "") + if key in skipped and _SKIP_PREFLIGHT not in reproducer.split(): + reproducer = f"{reproducer} {_SKIP_PREFLIGHT}".strip() + command = f"Command: `{reproducer}`" if key in skipped: command += f" (preflight skipped: {', '.join(skipped[key]) or 'no gate failed'})" lines += [command, ""] diff --git a/tests/test_bench_table.py b/tests/test_bench_table.py index 0131ab8..d68a0b0 100644 --- a/tests/test_bench_table.py +++ b/tests/test_bench_table.py @@ -413,9 +413,26 @@ def test_overhead_block_gives_each_scenario_its_command_and_a_skipped_preflight( lines = text.splitlines() dev = lines.index(next(x for x in lines if x.startswith("FLUX.1-dev"))) schnell = lines.index(next(x for x in lines if x.startswith("FLUX.1-schnell"))) - assert lines[dev + 2] == "Command: `uv run dev` (preflight skipped: not_charging)" + assert ( + lines[dev + 2] == "Command: `uv run dev --skip-preflight` (preflight skipped: not_charging)" + ) assert lines[schnell + 2] == "Command: `uv run schnell`" assert text.count("preflight skipped") == 1 + assert text.count("--skip-preflight") == 1 + + +def test_a_skipped_preflight_reproducer_that_already_carries_the_flag_prints_it_once(): + # Bug caught: the flag appended unconditionally, so a report whose recorded command already ends in + # --skip-preflight renders a command that names it twice. + text = render_overhead_block( + {"flux1-dev-1024": _summary()}, + caption="c", + reproducers={"flux1-dev-1024": "uv run dev --skip-preflight"}, + preflight_skipped={"flux1-dev-1024": ("not_charging",)}, + cache_limit_note="n", + ) + assert "Command: `uv run dev --skip-preflight` (preflight skipped: not_charging)" in text + assert text.count("--skip-preflight") == 1 def test_caption_keeps_the_dirty_suffix(): diff --git a/tests/test_bench_table_script.py b/tests/test_bench_table_script.py index 8d4eaa6..d33a52b 100644 --- a/tests/test_bench_table_script.py +++ b/tests/test_bench_table_script.py @@ -357,8 +357,9 @@ def test_the_committed_reports_render_each_command_and_the_dev_runs_skipped_pref assert collected.preflight_skipped == {"flux1-dev-1024": ("not_charging",)} block = bt.render_readme(README, collected, date=collected.date) for name in ("flux1-dev-1024", "flux1-schnell-1024"): - assert f"bench/scenarios/{name}.toml`" in block - assert "flux1-dev-1024.toml` (preflight skipped: not_charging)" in block + assert f"bench/scenarios/{name}.toml" in block + assert "flux1-dev-1024.toml --skip-preflight` (preflight skipped: not_charging)" in block + assert "flux1-schnell-1024.toml`\n" in block # the schnell run passed its launch check def test_collect_reads_the_watchdog_stops_next_to_the_tier_reports(tmp_path): diff --git a/tests/test_bench_table_versions.py b/tests/test_bench_table_versions.py new file mode 100644 index 0000000..41e70c7 --- /dev/null +++ b/tests/test_bench_table_versions.py @@ -0,0 +1,151 @@ +"""The tier table names the mflux version each row was measured on, from the row's own record.""" + +import json +from importlib import metadata +from pathlib import Path + +from mlx_dfloat.bench.table import ( + render_tier_table, + tier_row_from_abort_artifact, + tier_row_from_generate_report, +) + +_FIXTURE = Path(__file__).parent / "fixtures" / "generate_report_schnell_1024.json" + + +def _report(mflux: str | None) -> dict: + """A report ``mlx-dfloat generate`` wrote, its ``versions.mflux`` set (None: no ``versions`` key).""" + report = json.loads(_FIXTURE.read_text()) + if mflux is None: + del report["versions"] + else: + report["versions"] = {**report["versions"], "mflux": mflux} + return report + + +def _abort() -> dict: + # An abort artifact records no versions (the watchdog writes it, not generate). + return { + "reason": "memory", + "elapsed": 5.6, + "ceiling": 9_842_633_386, + "peak_watched": 9_900_000_000, + "peak_footprint": 9_800_000_000, + "peak_mlx": 9_700_000_000, + "context": {"model": "z-image-turbo", "tier_gb": 16, "label": "CAPPED"}, + } + + +def _last_cell(line: str) -> str: + return line.rstrip(" |").rsplit("|", 1)[1].strip() + + +def _cells(rows) -> list[str]: + lines = render_tier_table(rows).splitlines() + return [_last_cell(line) for line in lines[2:]] + + +def test_the_header_ends_with_an_mflux_column(): + # Red when: the table has no column saying which mflux the row ran on. + header = render_tier_table([tier_row_from_generate_report(_report("0.23.0"), source="a")]) + assert header.splitlines()[0].endswith("| Result | mflux |") + + +def test_a_generate_row_carries_the_reports_mflux_version(): + # Red when: the row ignores versions.mflux, or takes it from somewhere other than the report. + assert tier_row_from_generate_report(_report("0.23.0"), source="a").mflux == "0.23.0" + assert tier_row_from_generate_report(_report("0.20.0"), source="a").mflux == "0.20.0" + + +def test_a_report_without_versions_has_no_mflux_version(): + # Red when: a report written before versions existed is refused, or given a made-up version. + assert tier_row_from_generate_report(_report(None), source="a").mflux is None + + +def test_a_row_older_than_the_newest_says_it_was_not_re_run(): + # Red when: an older row reads like a current one (the label hard-coded, or taken from the newest row), or the + # newest row is also marked as not re-run. + rows = [ + tier_row_from_generate_report(_report("0.23.0"), source="new.json"), + tier_row_from_generate_report(_report("0.20.0"), source="old.json"), + ] + assert _cells(rows) == ["0.23", "0.20, not re-run on 0.23"] + + +def test_when_every_row_ran_on_one_version_no_row_says_not_re_run(): + # Red when: the "newest" version is a constant (0.23) rather than the newest among the rows. + rows = [tier_row_from_generate_report(_report("0.20.0"), source=s) for s in ("a", "b")] + assert _cells(rows) == ["0.20", "0.20"] + + +def test_the_newest_version_compares_numerically(): + # Red when: versions compare as strings ("0.9.1" > "0.23.0"), so the 0.23 row is called the old one. + rows = [ + tier_row_from_generate_report(_report("0.9.1"), source="a"), + tier_row_from_generate_report(_report("0.23.0"), source="b"), + ] + assert _cells(rows) == ["0.9, not re-run on 0.23", "0.23"] + + +def test_a_row_without_a_recorded_version_says_so(): + # Red when: an abort row (no versions in the artifact) is labelled with a guessed version, or a report without + # versions breaks the render. + rows = [ + tier_row_from_generate_report(_report("0.23.0"), source="a"), + tier_row_from_abort_artifact(_abort(), source="b"), + tier_row_from_generate_report(_report(None), source="c"), + ] + assert _cells(rows) == ["0.23", "not recorded", "not recorded"] + + +def _abort_on(versions: dict | None) -> dict: + """The abort artifact, its run context recording ``versions`` (None: an older artifact without them).""" + artifact = _abort() + if versions is not None: + artifact["context"] = {**artifact["context"], "versions": versions, "macos": "27.0.1"} + return artifact + + +def test_an_abort_row_carries_the_mflux_version_its_run_context_records(): + # Red when: an abort row ignores context.versions, so every stopped run reads "not recorded" even when the + # artifact says which mflux it ran on, or the row takes the newest version instead of its own. + rows = [ + tier_row_from_generate_report(_report("0.23.0"), source="a"), + tier_row_from_abort_artifact(_abort_on({"mlx": "0.32.2", "mflux": "0.20.0"}), source="b"), + tier_row_from_abort_artifact(_abort_on(None), source="c"), + ] + assert rows[1].mflux == "0.20.0" + assert _cells(rows) == ["0.23", "0.20, not re-run on 0.23", "not recorded"] + + +def test_a_stopped_runs_row_names_the_mflux_generate_ran_on(): + # Red when: generate's watchdog context (the abort artifact's) records no versions, so the row of a run the + # watchdog stopped cannot name its mflux. + from mlx_dfloat.mflux import generate as gen + + args = gen.build_parser().parse_args(["--prompt", "p"]) + limits = {"memory": 1, "cache": 2, "wired": 3} + context = gen._run_context(args, tier_gb=16, label="CAPPED", limits=limits) + row = tier_row_from_abort_artifact({**_abort(), "context": context}, source="b") + assert row.mflux == _installed("mflux") + + +def _installed(name: str) -> str | None: + try: + return metadata.version(name) + except metadata.PackageNotFoundError: + return None + + +def test_a_pre_release_does_not_sort_above_its_release(): + # Red when: the version key reads the 1 of "rc1" as a fourth number, so 0.23.0rc1 counts as newer than 0.23.0 and + # the 0.23.0 row is marked "not re-run on 0.23". Both orders: the newest is a max over the rows. + for order in (("0.23.0", "0.23.0rc1"), ("0.23.0rc1", "0.23.0")): + rows = [tier_row_from_generate_report(_report(v), source=v) for v in order] + assert _cells(rows) == ["0.23", "0.23"], order + + +def test_a_later_release_still_sorts_above_an_earlier_pre_release(): + # Red when: the fix drops the numbers that do count (a constant or a one-number key ties 0.24.1 with 0.23.0rc1). + rows = [tier_row_from_generate_report(_report(v), source=v) for v in ("0.23.0rc1", "0.24.1")] + assert _cells(rows) == ["0.23, not re-run on 0.24", "0.24"] diff --git a/tests/test_ci_config.py b/tests/test_ci_config.py index d4b8bb2..7db41d0 100644 --- a/tests/test_ci_config.py +++ b/tests/test_ci_config.py @@ -78,7 +78,7 @@ def test_the_mflux_lane_installs_the_locked_extra_and_proves_it_ran_tests(): # job dies before a single test), dropping `-m mflux` / `--run-network`, or a lane that goes # green having collected zero mflux tests, or a `uv run` step that could drop the extra. job = _integration_job() - assert job["strategy"]["matrix"]["mflux"] == ["0.20.0"] + assert job["strategy"]["matrix"]["mflux"] == ["0.23.0"] runs = [step.get("run", "") for step in job["steps"]] assert "uv sync --locked --group dev --extra mflux" in runs assert not any("uv pip install" in run or "--no-sync" in run for run in runs) @@ -113,6 +113,17 @@ def test_the_mflux_lane_installs_the_locked_extra_and_proves_it_ran_tests(): assert "-ge 2" in guard[0] +def test_the_mflux_lane_fails_when_the_lock_installs_another_mflux_than_its_row_names(): + # Bug caught: the lane named for one mflux release (the matrix row, which is the required check's name on main) + # while `uv sync --locked` installs whatever the lock pins, so a stale lock runs the old mflux under the new name + # and the lane stays green. The sanity step must compare the installed version with the row and exit non-zero. + runs = [step.get("run", "") for step in _integration_job()["steps"]] + sanity = [run for run in runs if "version('mflux')" in run] + assert len(sanity) == 1 + assert "== '${{ matrix.mflux }}'" in sanity[0] + assert "sys.exit(" in sanity[0] + + def test_the_mflux_lane_gates_the_model_module_on_its_own(): # Bug caught: model.py's lines diluted by init.py's in the lane's combined 80 % gate, so the # module with no offline test at all could drop far below 80 % while the job stays green; or diff --git a/tests/test_readme_bench_block.py b/tests/test_readme_bench_block.py index ccc9d79..6ffd1e9 100644 --- a/tests/test_readme_bench_block.py +++ b/tests/test_readme_bench_block.py @@ -1,7 +1,9 @@ """The README's generated numbers equal what the generator renders from the committed results.""" +import re from pathlib import Path +import pytest from scripts import bench_table as bt REPO = Path(__file__).resolve().parents[1] @@ -27,3 +29,121 @@ def test_hand_edit_is_caught(tmp_path): path.write_text(edited) assert bt.main(["--check", "--readme", str(path)]) == 1 assert path.read_text() == edited # --check never writes + + +# The "Does this help me?" table summarises the measured numbers in hand-written cells. They are read here and checked +# against the committed results, so a re-run that moves a number turns this file red until the cell is updated. +_GIB = 1024**3 +README_PATH = REPO / "README.md" + + +def _help_row(starts_with: str) -> str: + rows = [line for line in README_PATH.read_text().splitlines() if line.startswith(starts_with)] + assert len(rows) == 1, ( + f"expected one README table row starting {starts_with!r}, found {len(rows)}" + ) + return rows[0] + + +def _flux1_row() -> str: + return _help_row("| A 32 GB Mac, FLUX.1-schnell, FLUX.1-dev or FLUX.1-Krea-dev,") + + +def _q8_row() -> str: + return _help_row("| Less memory or a faster step,") + + +def _number(pattern: str, row: str) -> float: + match = re.search(pattern, row) + assert match is not None, f"{pattern!r} not found in the README row: {row}" + return float(match.group(1)) + + +def _summaries() -> dict: + return bt.collect(REPO / "bench" / "results").summaries + + +def test_the_flux1_rows_step_overheads_are_the_committed_per_block_medians(): + # Bug caught: a scenario re-run moves the dev or schnell per-block overhead (now +3.5 % and +4.0 %) and the + # hand-written "slower" cell keeps the old figure, or the two models' figures are swapped in the cell. + row = _flux1_row() + summaries = _summaries() + for model in ("dev", "schnell"): + stated = _number(rf"(\d+\.\d) % slower on {model}\b", row) + measured = summaries[f"flux1-{model}-1024"].overhead["per-block"] + assert measured > 0, f"{model}: the step is faster, not slower" + assert f"{stated:.1f}" == f"{measured * 100:.1f}", model + + +def test_the_flux1_rows_schnell_peak_is_its_tier_files_watched_peak(): + # Bug caught: a schnell tier re-run moves its watched peak (now 19.77 GiB, mflux 0.23.0) and the hand-written cell + # keeps the old figure, or the cell's number is taken from dev or Krea-dev's file. + stated = _number(r"Measured peak at 1024²: (\d+\.\d\d) GiB for schnell \(mflux", _flux1_row()) + measured = _tier_row("schnell-1024.json").watched_peak_bytes / _GIB + assert f"{stated:.2f}" == f"{measured:.2f}" + + +def test_the_flux1_rows_dev_and_krea_dev_peak_is_each_tier_rows_watched_peak_rounded_to_a_gib(): + # Bug caught: a dev or Krea-dev tier re-run peaks at, say, 21 GiB while the cell still says "about 20 GiB for dev + # and Krea-dev", or one of the two drops out of the tier files the claim covers. + stated = _number(r"about (\d+) GiB for dev and Krea-dev", _flux1_row()) + for file in ("dev-1024.json", "krea-dev-1024.json"): + assert round(_tier_row(file).watched_peak_bytes / _GIB) == stated, file + + +def test_the_q8_rows_ratio_lies_within_the_measured_q8_ratios(): + # Bug caught: a re-run moves the DF11-over-q8 ratios (now 1.26 on dev, 1.35 on schnell) away from the cell's + # "roughly 1.3 times faster"; "roughly" covers the measured span, so the stated figure must lie inside it. + stated = _number(r"roughly (\d+\.\d) times faster", _q8_row()) + ratios = [s.q8_ratio for s in _summaries().values()] + assert ratios + assert None not in ratios + assert min(ratios) <= stated <= max(ratios), ratios + + +def test_the_q8_rows_memory_saving_is_the_step_peak_gap_rounded_to_a_gib(): + # Bug caught: the q8 or DF11 step peak moves on a re-run (the gap is now 5.16 GiB on dev, 4.88 on schnell) and the + # cell's "about 5 GiB lower" stays, or the gap is taken against the wrong condition. + stated = _number(r"peaked about (\d+) GiB lower", _q8_row()) + for key, s in _summaries().items(): + gap = s.conditions["df11"]["step_watched_peak"] - s.conditions["q8"]["step_watched_peak"] + assert round(gap / _GIB) == stated, key + + +def _tier_row(file: str): + rows = [ + r for r in bt.collect(REPO / "bench" / "results").tier_rows if Path(r.source).name == file + ] + assert len(rows) == 1, f"expected one tier row from {file}, found {len(rows)}" + return rows[0] + + +@pytest.mark.parametrize( + ("row_start", "pattern", "tier_file"), + [ + pytest.param( + "| A 32 GB Mac, Z-Image-Turbo or Z-Image,", + r"(\d+\.\d\d) GiB for Z-Image \(mflux 0\.23\)", + "z-image-1024.json", + id="z-image", + ), + pytest.param( + "| A 32 GB Mac, FLUX.2 Klein (4B or 9B,", + r"Peak at 1024²: (\d+\.\d\d) to \d+\.\d\d GiB for the 4B models", + "flux2-klein-4b-1024.json", + id="klein-4b-distilled", + ), + pytest.param( + "| A 32 GB Mac, Qwen-Image-2.1,", + r"Peak at 1024²: (\d+\.\d\d) GiB \(mflux 0\.23\)", + "qwen-image-2.1-1024.json", + id="qwen-image-2.1", + ), + ], +) +def test_a_hand_written_peak_cell_is_its_tier_files_watched_peak(row_start, pattern, tier_file): + # Bug caught: a re-run of Z-Image (14.40 GiB), the distilled Klein 4B (11.85, the low end of the 4B span) or + # Qwen-Image-2.1 (20.37) moves its watched peak in bench/results/tiers and the hand-written cell keeps the old one. + stated = _number(pattern, _help_row(row_start)) + measured = _tier_row(tier_file).watched_peak_bytes / _GIB + assert f"{stated:.2f}" == f"{measured:.2f}", tier_file diff --git a/tests/test_verify_image.py b/tests/test_verify_image.py index 66dad9f..22ffcae 100644 --- a/tests/test_verify_image.py +++ b/tests/test_verify_image.py @@ -771,7 +771,7 @@ def __init__(self, steps, guidance=1.0): class _UnpackedQwenTransformer: - """A transformer without mflux 0.23's packing window (mflux 0.20, or any family without it). + """The packing transformer's base: its call outside any window. Its noise is the prompt embedding plus 1, so the positive and negative calls give different noise. """ @@ -835,13 +835,6 @@ def test_the_qwen_bf16_loop_runs_inside_the_transformers_packing_window(): assert transformer.packed is False # the window is closed before the VAE decode -def test_the_qwen_bf16_loop_runs_a_transformer_without_a_packing_window_unchanged(): - # Bug caught: the window required (an AttributeError on mflux 0.20's transformer), or entered on some other - # object. Two steps of noise 1 give 2. - latents = _qwen_denoise(_UnpackedQwenTransformer(), _QwenConfig(2)) - assert latents.tolist() == [[[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]] - - def test_the_qwen_bf16_loop_combines_the_negative_call_with_mfluxs_guidance(): # Bug caught: CFG lost or its combine flipped in the loop body. mflux qwen_image_21.py: neg + g * (pos - neg). # Positive noise 11 (prompt 0, packed), negative 15 (prompt 4, packed), guidance 3: 15 + 3 * (11 - 15) = 3 per From 5b64507368dd1822bae597d84f55456a38a7912f Mon Sep 17 00:00:00 2001 From: Denis Ineshin Date: Wed, 28 Oct 2026 19:32:01 +0100 Subject: [PATCH 3/4] Record the mflux 0.23 identity and tier runs Identity and tier records on mflux 0.23 for FLUX.1-schnell, Z-Image, FLUX.2 Klein 4B and Qwen-Image 2.1. Klein 4B gets its first identity record. --- .../qwen-image-2.1-1024-mflux023.json | 50 +++++++ .../calibration/qwen-image-2.1-1024.json | 2 +- .../flux1-schnell-1024/bf16-result.json | 16 +- .../identity/flux1-schnell-1024/compare.json | 8 +- .../flux1-schnell-1024/df11-result.json | 56 +++---- .../flux2-klein-4b-1024/bf16-result.json | 55 +++++++ .../identity/flux2-klein-4b-1024/compare.json | 34 +++++ .../flux2-klein-4b-1024/df11-result.json | 140 ++++++++++++++++++ .../qwen-image-2.1-1024/bf16-result.json | 16 +- .../identity/qwen-image-2.1-1024/compare.json | 8 +- .../qwen-image-2.1-1024/df11-result.json | 38 ++--- .../identity/z-image-1024/bf16-result.json | 16 +- .../identity/z-image-1024/compare.json | 8 +- .../identity/z-image-1024/df11-result.json | 48 +++--- bench/results/tiers/flux2-klein-4b-1024.json | 44 +++--- bench/results/tiers/qwen-image-2.1-1024.json | 39 ++--- bench/results/tiers/schnell-1024.json | 59 ++++---- bench/results/tiers/z-image-1024.json | 53 +++---- 18 files changed, 489 insertions(+), 201 deletions(-) create mode 100644 bench/results/calibration/qwen-image-2.1-1024-mflux023.json create mode 100644 bench/results/identity/flux2-klein-4b-1024/bf16-result.json create mode 100644 bench/results/identity/flux2-klein-4b-1024/compare.json create mode 100644 bench/results/identity/flux2-klein-4b-1024/df11-result.json diff --git a/bench/results/calibration/qwen-image-2.1-1024-mflux023.json b/bench/results/calibration/qwen-image-2.1-1024-mflux023.json new file mode 100644 index 0000000..37f26fb --- /dev/null +++ b/bench/results/calibration/qwen-image-2.1-1024-mflux023.json @@ -0,0 +1,50 @@ +{ + "model": "qwen-image-2.1", + "purpose": "the 1024² VAE-phase samples on mflux 0.23 the Qwen-Image 2.1 VAE transient is calibrated on together with qwen-image-2.1-1024.json (mlx_dfloat.mflux.qwen21.memory); copied from the committed records named in each run's source, no new run", + "df11": "mingyi456/Qwen-Image-2.1-DF11-ComfyUI@1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", + "base": "Qwen/Qwen-Image-2.1@d26bb61231c349cf6b7896fa83353113880e1ba3", + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "size": 1024, + "seed": 42, + "text_tokens": 33, + "device": "Apple M1 Max, 32 GB", + "macos": "27.0.1", + "versions": { + "mlx": "0.32.2", + "mflux": "0.23.0" + }, + "footprint_peaks_label": "OS phys_footprint in bytes: the run's watched peak, reached in the VAE phase", + "constants_from": { + "vae": "identity-4-steps-cfg-mflux023" + }, + "runs": [ + { + "label": "measured-40-steps-mflux023", + "date": "2026-10-24/25", + "run": "the MEASURED tier-32 generate run: 40 steps, the VAE decode on the resident set", + "steps": 40, + "cache_limit": 2731761634, + "source": { + "record": "bench/results/tiers/qwen-image-2.1-1024.json", + "field": "watched_peak_bytes" + }, + "footprint_peaks": { + "vae": 21877154888 + } + }, + { + "label": "identity-4-steps-cfg-mflux023", + "date": "2026-10-24", + "run": "the identity check's df11 side: 4 steps, guidance 4 with the negative prompt \" \"", + "steps": 4, + "cache_limit": 2731761634, + "source": { + "record": "bench/results/identity/qwen-image-2.1-1024/df11-result.json", + "field": "footprint_peak_bytes" + }, + "footprint_peaks": { + "vae": 22831539272 + } + } + ] +} diff --git a/bench/results/calibration/qwen-image-2.1-1024.json b/bench/results/calibration/qwen-image-2.1-1024.json index 1309719..2c5d9bd 100644 --- a/bench/results/calibration/qwen-image-2.1-1024.json +++ b/bench/results/calibration/qwen-image-2.1-1024.json @@ -17,7 +17,7 @@ "constants_from": { "encode": "derisk-1", "denoise": "ab-new-limit-1", - "vae": "ab-old-limit-2" + "vae": "identity-4-steps-cfg-mflux023" }, "runs": [ { diff --git a/bench/results/identity/flux1-schnell-1024/bf16-result.json b/bench/results/identity/flux1-schnell-1024/bf16-result.json index dbb29f1..cad946d 100644 --- a/bench/results/identity/flux1-schnell-1024/bf16-result.json +++ b/bench/results/identity/flux1-schnell-1024/bf16-result.json @@ -7,8 +7,8 @@ 64 ], "output_dtype": "mlx.core.float32", - "footprint_peak_bytes": 10350445792, - "mlx_peak_bytes_process": 9988562688, + "footprint_peak_bytes": 9552593976, + "mlx_peak_bytes_process": 9451114240, "key": { "model": "schnell", "seed": 42, @@ -19,7 +19,7 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", - "source": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2" }, "memory_caps_gb": [ @@ -27,10 +27,10 @@ 22 ], "provenance": { - "git": "89ff46a6614e9fd4611cfea28c809c84243cca11", - "source_hash": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "git": "cbed93b2b10483834f2b3eee4f8b5c453204cffa", + "source_hash": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2", - "mflux": "0.20.0", + "mflux": "0.23.0", "macos": "27.0.1", "device_info": { "device_name": "Apple M1 Max", @@ -48,8 +48,8 @@ "cache_memory": 0, "power": { "ac": true, - "battery_pct": 94 + "battery_pct": 100 }, - "date": "2026-10-08" + "date": "2026-10-24" } } diff --git a/bench/results/identity/flux1-schnell-1024/compare.json b/bench/results/identity/flux1-schnell-1024/compare.json index 41701ec..8a2de56 100644 --- a/bench/results/identity/flux1-schnell-1024/compare.json +++ b/bench/results/identity/flux1-schnell-1024/compare.json @@ -12,7 +12,7 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", - "source": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2" }, "bf16_key": { @@ -25,10 +25,10 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", - "source": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2" }, "pixels_identical": true, - "df11_png_sha256": "75b9af5d084de91d69e11a5a5e1a444aeba3a498cadc3a106ac6e2492db7e2c4", - "bf16_png_sha256": "e7c07257317efb656b88f18d65a45be0e914d44603888ca3312f0fc290e767c1" + "df11_png_sha256": "a60b5a88e1c8b06f46781aad345e154f57155b34b7f6da9d5b51cd0b7259ad35", + "bf16_png_sha256": "f348fd558d9881122e75368e061f82db075001677f3782e5ae7bf20f8b2fd59b" } diff --git a/bench/results/identity/flux1-schnell-1024/df11-result.json b/bench/results/identity/flux1-schnell-1024/df11-result.json index ed98c35..3ec629d 100644 --- a/bench/results/identity/flux1-schnell-1024/df11-result.json +++ b/bench/results/identity/flux1-schnell-1024/df11-result.json @@ -21,7 +21,7 @@ "phases": { "encode": 11518505446, "denoise": 21411230501, - "vae": 8808361940 + "vae": 8922261588 }, "peak_phase": "denoise", "peak_bytes": 21411230501, @@ -37,28 +37,28 @@ "peaks": { "label": "sampled at phase boundaries", "encode": { - "footprint_start": 363021464, - "mlx_peak": 10976952538, - "footprint_end": 9915833616, - "active_at_start": 115538184 + "footprint_start": 363136104, + "mlx_peak": 10283117198, + "footprint_end": 9807518992, + "active_at_start": 115550472 }, "set_load": { - "footprint_start": 9915833616, - "mlx_peak": 16315890952, - "footprint_end": 16608785104, - "active_at_start": 119734024 + "footprint_start": 9807518992, + "mlx_peak": 16315903240, + "footprint_end": 16609849944, + "active_at_start": 119746312 }, "denoise": { - "footprint_start": 16608785104, - "mlx_peak": 18725842260, - "footprint_end": 19131576992, - "active_at_start": 16315890952 + "footprint_start": 16609849944, + "mlx_peak": 18725854548, + "footprint_end": 19095040600, + "active_at_start": 16315903240 }, "vae": { - "footprint_start": 9832683784, - "mlx_peak": 10257014912, - "footprint_end": 583550464, - "active_at_start": 121831228 + "footprint_start": 10246560248, + "mlx_peak": 10257027200, + "footprint_end": 566068640, + "active_at_start": 121843516 } }, "decode_launches": 228, @@ -66,12 +66,12 @@ "encoder_loads": 0, "encoder_seconds": 0.0, "set_loads": 1, - "set_seconds": 26.188013666018378, + "set_seconds": 26.534756542067043, "forced_set_drops": 0 }, "versions": { "mlx": "0.32.2", - "mflux": "0.20.0" + "mflux": "0.23.0" } }, "output_shape": [ @@ -80,8 +80,8 @@ 64 ], "output_dtype": "mlx.core.float32", - "footprint_peak_bytes": 21237592736, - "mlx_peak_bytes_max_over_phases": 18725842260, + "footprint_peak_bytes": 21303128664, + "mlx_peak_bytes_max_over_phases": 18725854548, "key": { "model": "schnell", "seed": 42, @@ -92,7 +92,7 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", - "source": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2" }, "memory_caps_gb": [ @@ -100,10 +100,10 @@ 22 ], "provenance": { - "git": "89ff46a6614e9fd4611cfea28c809c84243cca11", - "source_hash": "d5992851091bf4a00217cc012c6382ffe6887d62655fc4bc306d4731d40cd801", + "git": "cbed93b2b10483834f2b3eee4f8b5c453204cffa", + "source_hash": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2", - "mflux": "0.20.0", + "mflux": "0.23.0", "macos": "27.0.1", "device_info": { "device_name": "Apple M1 Max", @@ -118,11 +118,11 @@ 22 ], "cache_limit": 32641751449, - "cache_memory": 1836212, + "cache_memory": 2884788, "power": { "ac": true, - "battery_pct": 94 + "battery_pct": 100 }, - "date": "2026-10-08" + "date": "2026-10-24" } } diff --git a/bench/results/identity/flux2-klein-4b-1024/bf16-result.json b/bench/results/identity/flux2-klein-4b-1024/bf16-result.json new file mode 100644 index 0000000..325a728 --- /dev/null +++ b/bench/results/identity/flux2-klein-4b-1024/bf16-result.json @@ -0,0 +1,55 @@ +{ + "exit_code": 0, + "reads": 400, + "output_shape": [ + 1, + 4096, + 128 + ], + "output_dtype": "mlx.core.bfloat16", + "footprint_peak_bytes": 7409108240, + "mlx_peak_bytes_process": 7796048558, + "key": { + "model": "flux2-klein-4b", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 1.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--FLUX.2-klein-4B-DF11/snapshots/d29a2c249ff0afeb678da101e707bd134596c96f", + "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots/e7b7dc27f91deacad38e78976d1f2b499d76a294", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", + "mlx": "0.32.2" + }, + "memory_caps_gb": [ + 20, + 22 + ], + "provenance": { + "git": "cbed93b2b10483834f2b3eee4f8b5c453204cffa", + "source_hash": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", + "mlx": "0.32.2", + "mflux": "0.23.0", + "macos": "27.0.1", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 0, + "cache_memory": 0, + "power": { + "ac": true, + "battery_pct": 100 + }, + "date": "2026-10-24" + } +} diff --git a/bench/results/identity/flux2-klein-4b-1024/compare.json b/bench/results/identity/flux2-klein-4b-1024/compare.json new file mode 100644 index 0000000..5c44650 --- /dev/null +++ b/bench/results/identity/flux2-klein-4b-1024/compare.json @@ -0,0 +1,34 @@ +{ + "exit_code": 0, + "equal": true, + "degenerate": [], + "df11_key": { + "model": "flux2-klein-4b", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 1.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--FLUX.2-klein-4B-DF11/snapshots/d29a2c249ff0afeb678da101e707bd134596c96f", + "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots/e7b7dc27f91deacad38e78976d1f2b499d76a294", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", + "mlx": "0.32.2" + }, + "bf16_key": { + "model": "flux2-klein-4b", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 1.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--FLUX.2-klein-4B-DF11/snapshots/d29a2c249ff0afeb678da101e707bd134596c96f", + "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots/e7b7dc27f91deacad38e78976d1f2b499d76a294", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", + "mlx": "0.32.2" + }, + "pixels_identical": true, + "df11_png_sha256": "cd02ee2e94edb36fc7f7411bd0d3e56ecc59eee612ae458e450c26fe8bb66015", + "bf16_png_sha256": "311460faf5ea56ff8761fe813a2420d9cfdb8432e679c520761673d01c0f9f7c" +} diff --git a/bench/results/identity/flux2-klein-4b-1024/df11-result.json b/bench/results/identity/flux2-klein-4b-1024/df11-result.json new file mode 100644 index 0000000..4d7fa8e --- /dev/null +++ b/bench/results/identity/flux2-klein-4b-1024/df11-result.json @@ -0,0 +1,140 @@ +{ + "exit_code": 0, + "degenerate": [], + "report": { + "family": "flux2", + "model": "flux2-klein-4b", + "size": "4b", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--FLUX.2-klein-4B-DF11/snapshots/d29a2c249ff0afeb678da101e707bd134596c96f", + "repo_id": null, + "revision": "d29a2c249ff0afeb678da101e707bd134596c96f" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots/e7b7dc27f91deacad38e78976d1f2b499d76a294", + "repo_id": null, + "revision": "e7b7dc27f91deacad38e78976d1f2b499d76a294" + }, + "eval_policy": "per-block", + "cfg_calls_per_step": 1, + "predict": "uncompiled", + "nonblock_groups": [ + "context_embedder", + "double_stream_modulation_img.linear", + "double_stream_modulation_txt.linear", + "norm_out.linear", + "single_stream_modulation.linear" + ], + "cache_limit_in_force": 2324335646, + "drop_set_before_vae": false, + "fit": { + "label": "predicted", + "phases": { + "encode": 8197581550, + "denoise": 9168776559, + "vae": 13621456239 + }, + "peak_phase": "vae", + "peak_bytes": 13621456239, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 5239066183, + "extras": 22035456, + "nonblock": 368050176, + "encoders": 6228181504, + "vae": 168120878 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 408257688, + "mlx_peak": 6757263340, + "footprint_end": 6387177328, + "active_at_start": 22035604 + }, + "set_load": { + "footprint_start": 6387177328, + "mlx_peak": 5637753040, + "footprint_end": 6032053432, + "active_at_start": 29910152 + }, + "denoise": { + "footprint_start": 6032053432, + "mlx_peak": 6977744204, + "footprint_end": 8136332472, + "active_at_start": 5384674952 + }, + "vae": { + "footprint_start": 8136332472, + "mlx_peak": 11709020462, + "footprint_end": 6719952008, + "active_at_start": 5386837676 + } + }, + "decode_launches": 100, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 8.064038374926895, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.23.0" + } + }, + "output_shape": [ + 1, + 4096, + 128 + ], + "output_dtype": "mlx.core.bfloat16", + "footprint_peak_bytes": 13186225288, + "mlx_peak_bytes_max_over_phases": 11709020462, + "key": { + "model": "flux2-klein-4b", + "seed": 42, + "steps": 4, + "size": 1024, + "prompt": "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water", + "guidance": 1.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--FLUX.2-klein-4B-DF11/snapshots/d29a2c249ff0afeb678da101e707bd134596c96f", + "base": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots/e7b7dc27f91deacad38e78976d1f2b499d76a294", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", + "mlx": "0.32.2" + }, + "memory_caps_gb": [ + 20, + 22 + ], + "provenance": { + "git": "cbed93b2b10483834f2b3eee4f8b5c453204cffa", + "source_hash": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", + "mlx": "0.32.2", + "mflux": "0.23.0", + "macos": "27.0.1", + "device_info": { + "device_name": "Apple M1 Max", + "max_recommended_working_set_size": 26800603136, + "memory_size": 34359738368, + "architecture": "applegpu_g13s", + "max_buffer_length": 20100448256, + "resource_limit": 499000 + }, + "memory_caps_gb": [ + 20, + 22 + ], + "cache_limit": 32641751449, + "cache_memory": 8914104, + "power": { + "ac": true, + "battery_pct": 100 + }, + "date": "2026-10-24" + } +} diff --git a/bench/results/identity/qwen-image-2.1-1024/bf16-result.json b/bench/results/identity/qwen-image-2.1-1024/bf16-result.json index 2e48dc0..ce2e07a 100644 --- a/bench/results/identity/qwen-image-2.1-1024/bf16-result.json +++ b/bench/results/identity/qwen-image-2.1-1024/bf16-result.json @@ -1,6 +1,6 @@ { "exit_code": 0, - "reads": 1792, + "reads": 2240, "cfg_calls_per_step": 2, "output_shape": [ 1, @@ -8,8 +8,8 @@ 64 ], "output_dtype": "mlx.core.bfloat16", - "footprint_peak_bytes": 13110216840, - "mlx_peak_bytes_process": 12887490430, + "footprint_peak_bytes": 14194198544, + "mlx_peak_bytes_process": 14098580478, "key": { "model": "qwen-image-2.1", "seed": 42, @@ -20,7 +20,7 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", "base": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", - "source": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2", "negative_prompt": " " }, @@ -29,10 +29,10 @@ 22 ], "provenance": { - "git": "341c865e4b5ab1f6f611d489c9c40e1e640dcd3c", - "source_hash": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "git": "cbed93b2b10483834f2b3eee4f8b5c453204cffa", + "source_hash": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2", - "mflux": "0.20.0", + "mflux": "0.23.0", "macos": "27.0.1", "device_info": { "device_name": "Apple M1 Max", @@ -52,6 +52,6 @@ "ac": true, "battery_pct": 100 }, - "date": "2026-10-08" + "date": "2026-10-24" } } diff --git a/bench/results/identity/qwen-image-2.1-1024/compare.json b/bench/results/identity/qwen-image-2.1-1024/compare.json index be738da..d075d91 100644 --- a/bench/results/identity/qwen-image-2.1-1024/compare.json +++ b/bench/results/identity/qwen-image-2.1-1024/compare.json @@ -12,7 +12,7 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", "base": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", - "source": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2", "negative_prompt": " " }, @@ -26,11 +26,11 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", "base": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", - "source": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2", "negative_prompt": " " }, "pixels_identical": true, - "df11_png_sha256": "172f2c760124b619ba1656b939585899bec3050808cfcf8d8ae06d322206773e", - "bf16_png_sha256": "7522a96aa16f2687f0df52a53f7f417b09b55cd88ef068580d821ba03d505829" + "df11_png_sha256": "f12bf107af82839a2d30b7b9b3c85e2a2412e28d85c556ca9bd29339356404e4", + "bf16_png_sha256": "a9593471778d2125445f4c3143066887150f32c1de8b8f0cddde1c1b34cea1bb" } diff --git a/bench/results/identity/qwen-image-2.1-1024/df11-result.json b/bench/results/identity/qwen-image-2.1-1024/df11-result.json index d030316..4a6fba3 100644 --- a/bench/results/identity/qwen-image-2.1-1024/df11-result.json +++ b/bench/results/identity/qwen-image-2.1-1024/df11-result.json @@ -47,41 +47,41 @@ "peaks": { "label": "sampled at phase boundaries", "encode": { - "footprint_start": 610518288, - "mlx_peak": 15353362024, - "footprint_end": 13757443256, + "footprint_start": 613729552, + "mlx_peak": 15358179296, + "footprint_end": 12818885456, "active_at_start": 146825876 }, "set_load": { - "footprint_start": 13757443256, + "footprint_start": 12818885456, "mlx_peak": 9869647680, - "footprint_end": 10344735096, + "footprint_end": 10346701200, "active_at_start": 147416384 }, "denoise": { - "footprint_start": 10344735096, - "mlx_peak": 11532343934, - "footprint_end": 13060727112, + "footprint_start": 10346701200, + "mlx_peak": 11453782656, + "footprint_end": 12553675152, "active_at_start": 9776175680 }, "vae": { - "footprint_start": 12856713544, + "footprint_start": 12536553872, "mlx_peak": 21174302250, - "footprint_end": 15545099360, + "footprint_end": 15547491424, "active_at_start": 9782073950 } }, - "decode_launches": 256, + "decode_launches": 320, "lifecycle": { "encoder_loads": 0, "encoder_seconds": 0.0, "set_loads": 1, - "set_seconds": 14.079782124958001, + "set_seconds": 14.527245334000327, "forced_set_drops": 0 }, "versions": { "mlx": "0.32.2", - "mflux": "0.20.0" + "mflux": "0.23.0" } }, "output_shape": [ @@ -90,7 +90,7 @@ 64 ], "output_dtype": "mlx.core.bfloat16", - "footprint_peak_bytes": 21935465568, + "footprint_peak_bytes": 22831539272, "mlx_peak_bytes_max_over_phases": 21174302250, "key": { "model": "qwen-image-2.1", @@ -102,7 +102,7 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--mingyi456--Qwen-Image-2.1-DF11-ComfyUI/snapshots/1b22a3a1f96293f3b328d03abe22ab2e51cbd9cc", "base": "~/.cache/huggingface/hub/models--Qwen--Qwen-Image-2.1/snapshots/d26bb61231c349cf6b7896fa83353113880e1ba3", - "source": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2", "negative_prompt": " " }, @@ -111,10 +111,10 @@ 22 ], "provenance": { - "git": "341c865e4b5ab1f6f611d489c9c40e1e640dcd3c", - "source_hash": "ff655ce6cbf4d5cc134863c936c49711d16a252c0009f67c4d05a3182ed9933a", + "git": "cbed93b2b10483834f2b3eee4f8b5c453204cffa", + "source_hash": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2", - "mflux": "0.20.0", + "mflux": "0.23.0", "macos": "27.0.1", "device_info": { "device_name": "Apple M1 Max", @@ -134,6 +134,6 @@ "ac": true, "battery_pct": 100 }, - "date": "2026-10-08" + "date": "2026-10-24" } } diff --git a/bench/results/identity/z-image-1024/bf16-result.json b/bench/results/identity/z-image-1024/bf16-result.json index e4265f3..6924578 100644 --- a/bench/results/identity/z-image-1024/bf16-result.json +++ b/bench/results/identity/z-image-1024/bf16-result.json @@ -8,8 +8,8 @@ 128 ], "output_dtype": "mlx.core.bfloat16", - "footprint_peak_bytes": 5583488512, - "mlx_peak_bytes_process": 5217124834, + "footprint_peak_bytes": 6542656472, + "mlx_peak_bytes_process": 5970935270, "key": { "model": "z-image", "seed": 42, @@ -20,7 +20,7 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--mingyi456--Z-Image-DF11/snapshots/36d582560630bfe2f5b78183c4164c66d6919e5b", "base": "~/.cache/huggingface/hub/models--Tongyi-MAI--Z-Image/snapshots/04cc4abb7c5069926f75c9bfde9ef43d49423021", - "source": "bdc540a383c402237ff4dd69c5c49a6d833f9822f86326332076ddd9008fdc4b", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2" }, "memory_caps_gb": [ @@ -28,10 +28,10 @@ 22 ], "provenance": { - "git": "1038f5ea0ce6a8ad9db48f982eb3c05ab6163165", - "source_hash": "bdc540a383c402237ff4dd69c5c49a6d833f9822f86326332076ddd9008fdc4b", + "git": "cbed93b2b10483834f2b3eee4f8b5c453204cffa", + "source_hash": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2", - "mflux": "0.20.0", + "mflux": "0.23.0", "macos": "27.0.1", "device_info": { "device_name": "Apple M1 Max", @@ -49,8 +49,8 @@ "cache_memory": 0, "power": { "ac": true, - "battery_pct": 79 + "battery_pct": 100 }, - "date": "2026-10-08" + "date": "2026-10-24" } } diff --git a/bench/results/identity/z-image-1024/compare.json b/bench/results/identity/z-image-1024/compare.json index d0eab57..38fe204 100644 --- a/bench/results/identity/z-image-1024/compare.json +++ b/bench/results/identity/z-image-1024/compare.json @@ -12,7 +12,7 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--mingyi456--Z-Image-DF11/snapshots/36d582560630bfe2f5b78183c4164c66d6919e5b", "base": "~/.cache/huggingface/hub/models--Tongyi-MAI--Z-Image/snapshots/04cc4abb7c5069926f75c9bfde9ef43d49423021", - "source": "bdc540a383c402237ff4dd69c5c49a6d833f9822f86326332076ddd9008fdc4b", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2" }, "bf16_key": { @@ -25,10 +25,10 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--mingyi456--Z-Image-DF11/snapshots/36d582560630bfe2f5b78183c4164c66d6919e5b", "base": "~/.cache/huggingface/hub/models--Tongyi-MAI--Z-Image/snapshots/04cc4abb7c5069926f75c9bfde9ef43d49423021", - "source": "bdc540a383c402237ff4dd69c5c49a6d833f9822f86326332076ddd9008fdc4b", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2" }, "pixels_identical": true, - "df11_png_sha256": "c319d85dd834266aa9eed5bdcf9dd0b5ad7bca43fa3e844d2881ed52cb5d0d6c", - "bf16_png_sha256": "46e6724149ac21e1ff3276cf240d9c10e5d25934ec26f82a0c00ba88f3c765d3" + "df11_png_sha256": "9e3f5d3266d0f29787a5dc508c8f07bac77a08710cea4b2bced6435cd2d78c09", + "bf16_png_sha256": "99269a4d125cf0519b6ce4821e56e8c1f932cda075bd8d260eba8e67af9f08a4" } diff --git a/bench/results/identity/z-image-1024/df11-result.json b/bench/results/identity/z-image-1024/df11-result.json index 8333313..9e78268 100644 --- a/bench/results/identity/z-image-1024/df11-result.json +++ b/bench/results/identity/z-image-1024/df11-result.json @@ -27,10 +27,10 @@ "phases": { "encode": 10143248230, "denoise": 13101538864, - "vae": 13054789110 + "vae": 14115020240 }, - "peak_phase": "denoise", - "peak_bytes": 13101538864, + "peak_phase": "vae", + "peak_bytes": 14115020240, "budget_bytes": 24653119488, "fits": true }, @@ -44,28 +44,28 @@ "peaks": { "label": "sampled at phase boundaries", "encode": { - "footprint_start": 480249176, + "footprint_start": 472679600, "mlx_peak": 9177190600, - "footprint_end": 8849448088, + "footprint_end": 8992119912, "active_at_start": 6580692 }, "set_load": { - "footprint_start": 8849448088, + "footprint_start": 8992119912, "mlx_peak": 8401205524, - "footprint_end": 8883248784, + "footprint_end": 8876400080, "active_at_start": 11823560 }, "denoise": { - "footprint_start": 8883248784, - "mlx_peak": 10637496148, - "footprint_end": 11316797072, + "footprint_start": 8876400080, + "mlx_peak": 9719201160, + "footprint_end": 11234663864, "active_at_start": 8387910856 }, "vae": { - "footprint_start": 11315814032, - "mlx_peak": 13100904742, - "footprint_end": 9278840512, - "active_at_start": 8389483772 + "footprint_start": 11190132152, + "mlx_peak": 14332359206, + "footprint_end": 9808682448, + "active_at_start": 8388959484 } }, "decode_launches": 272, @@ -73,12 +73,12 @@ "encoder_loads": 0, "encoder_seconds": 0.0, "set_loads": 1, - "set_seconds": 15.174747042008676, + "set_seconds": 13.422165415948257, "forced_set_drops": 0 }, "versions": { "mlx": "0.32.2", - "mflux": "0.20.0" + "mflux": "0.23.0" } }, "output_shape": [ @@ -88,8 +88,8 @@ 128 ], "output_dtype": "mlx.core.bfloat16", - "footprint_peak_bytes": 14066196160, - "mlx_peak_bytes_max_over_phases": 13100904742, + "footprint_peak_bytes": 15085673936, + "mlx_peak_bytes_max_over_phases": 14332359206, "key": { "model": "z-image", "seed": 42, @@ -100,7 +100,7 @@ "eval_policy": "per-block", "df11": "~/.cache/huggingface/hub/models--mingyi456--Z-Image-DF11/snapshots/36d582560630bfe2f5b78183c4164c66d6919e5b", "base": "~/.cache/huggingface/hub/models--Tongyi-MAI--Z-Image/snapshots/04cc4abb7c5069926f75c9bfde9ef43d49423021", - "source": "bdc540a383c402237ff4dd69c5c49a6d833f9822f86326332076ddd9008fdc4b", + "source": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2" }, "memory_caps_gb": [ @@ -108,10 +108,10 @@ 22 ], "provenance": { - "git": "1038f5ea0ce6a8ad9db48f982eb3c05ab6163165", - "source_hash": "bdc540a383c402237ff4dd69c5c49a6d833f9822f86326332076ddd9008fdc4b", + "git": "cbed93b2b10483834f2b3eee4f8b5c453204cffa", + "source_hash": "2cb25229c0305c6c99ceb949c28e00c4e9e967de940da5c7858f50e8e8089241", "mlx": "0.32.2", - "mflux": "0.20.0", + "mflux": "0.23.0", "macos": "27.0.1", "device_info": { "device_name": "Apple M1 Max", @@ -129,8 +129,8 @@ "cache_memory": 4687032, "power": { "ac": true, - "battery_pct": 80 + "battery_pct": 100 }, - "date": "2026-10-08" + "date": "2026-10-24" } } diff --git a/bench/results/tiers/flux2-klein-4b-1024.json b/bench/results/tiers/flux2-klein-4b-1024.json index 29d0ec6..957c0e2 100644 --- a/bench/results/tiers/flux2-klein-4b-1024.json +++ b/bench/results/tiers/flux2-klein-4b-1024.json @@ -1,6 +1,6 @@ { "exit_code": 0, - "output": "klein-4b.png", + "output": "flux2-klein-4b.png", "memory_caps_gb": [ 20, 22 @@ -8,7 +8,7 @@ "model_kwargs": { "model": "flux2-klein-4b", "df11_path": "~/.cache/huggingface/hub/models--mingyi456--FLUX.2-klein-4B-DF11/snapshots/d29a2c249ff0afeb678da101e707bd134596c96f", - "base_path": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-base-4B/snapshots/a3b4f4849157f664bdbc776fd7453c2783562f4d", + "base_path": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots/e7b7dc27f91deacad38e78976d1f2b499d76a294", "eval_policy": "per-block", "cache_limit": null, "fit_check": true @@ -30,6 +30,7 @@ "reserve_bytes": 2147483648, "budget_bytes": 26800603136, "ceiling_bytes": 24653119488, + "fit_budget_bytes": 24653119488, "budget_source": "device", "is_host": true, "label": "MEASURED" @@ -41,6 +42,7 @@ }, "applied": "host-caps" }, + "macos": "27.0.1", "family": "flux2", "model": "flux2-klein-4b", "size": "4b", @@ -50,9 +52,9 @@ "revision": "d29a2c249ff0afeb678da101e707bd134596c96f" }, "base": { - "root": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-base-4B/snapshots/a3b4f4849157f664bdbc776fd7453c2783562f4d", + "root": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.2-klein-4B/snapshots/e7b7dc27f91deacad38e78976d1f2b499d76a294", "repo_id": null, - "revision": "a3b4f4849157f664bdbc776fd7453c2783562f4d" + "revision": "e7b7dc27f91deacad38e78976d1f2b499d76a294" }, "eval_policy": "per-block", "cfg_calls_per_step": 1, @@ -88,27 +90,27 @@ "peaks": { "label": "sampled at phase boundaries", "encode": { - "footprint_start": 405734528, + "footprint_start": 406652056, "mlx_peak": 6759884780, - "footprint_end": 6757144336, + "footprint_end": 6704011072, "active_at_start": 22035604 }, "set_load": { - "footprint_start": 6757144336, - "mlx_peak": 5637753040, - "footprint_end": 6029169632, + "footprint_start": 6704011072, + "mlx_peak": 5637759836, + "footprint_end": 6030660696, "active_at_start": 29910152 }, "denoise": { - "footprint_start": 6029169632, - "mlx_peak": 6977744208, - "footprint_end": 8133268424, + "footprint_start": 6030660696, + "mlx_peak": 6977744204, + "footprint_end": 8135119912, "active_at_start": 5384674952 }, "vae": { - "footprint_start": 8133268424, - "mlx_peak": 12782765358, - "footprint_end": 7253398424, + "footprint_start": 8135119912, + "mlx_peak": 11709020462, + "footprint_end": 6181590056, "active_at_start": 5386837676 } }, @@ -117,16 +119,16 @@ "encoder_loads": 0, "encoder_seconds": 0.0, "set_loads": 1, - "set_seconds": 8.198075499967672, + "set_seconds": 7.515274209086783, "forced_set_drops": 0 }, "versions": { "mlx": "0.32.2", - "mflux": "0.20.0" + "mflux": "0.23.0" }, - "footprint_peak_bytes": 12752884632, + "footprint_peak_bytes": 12718953512, "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", - "watched_peak_bytes": 12752884632, - "mlx_peak_bytes": 11709023534, - "elapsed_seconds": 35.039 + "watched_peak_bytes": 12718953512, + "mlx_peak_bytes": 10635278638, + "elapsed_seconds": 34.476 } diff --git a/bench/results/tiers/qwen-image-2.1-1024.json b/bench/results/tiers/qwen-image-2.1-1024.json index 3e161d7..c964dd7 100644 --- a/bench/results/tiers/qwen-image-2.1-1024.json +++ b/bench/results/tiers/qwen-image-2.1-1024.json @@ -42,6 +42,7 @@ }, "applied": "host-caps" }, + "macos": "27.0.1", "family": "qwen21", "model": "qwen-image-2.1", "df11": { @@ -87,45 +88,45 @@ "peaks": { "label": "sampled at phase boundaries", "encode": { - "footprint_start": 611714344, - "mlx_peak": 15353361928, - "footprint_end": 13554035800, + "footprint_start": 611026168, + "mlx_peak": 15367452640, + "footprint_end": 13621226656, "active_at_start": 146825876 }, "set_load": { - "footprint_start": 13554035800, + "footprint_start": 13621226656, "mlx_peak": 9869450888, - "footprint_end": 10343883224, + "footprint_end": 10343539208, "active_at_start": 147219592 }, "denoise": { - "footprint_start": 10343883224, - "mlx_peak": 11503639112, - "footprint_end": 12186051032, + "footprint_start": 10343539208, + "mlx_peak": 11413002830, + "footprint_end": 12439347720, "active_at_start": 9775978888 }, "vae": { - "footprint_start": 12172665304, - "mlx_peak": 21171435010, - "footprint_end": 15541412888, - "active_at_start": 9779206710 + "footprint_start": 12410053128, + "mlx_peak": 21171418626, + "footprint_end": 15541167176, + "active_at_start": 9779190326 } }, - "decode_launches": 1280, + "decode_launches": 1312, "lifecycle": { "encoder_loads": 0, "encoder_seconds": 0.0, "set_loads": 1, - "set_seconds": 13.601246499980334, + "set_seconds": 14.332625124952756, "forced_set_drops": 0 }, "versions": { "mlx": "0.32.2", - "mflux": "0.20.0" + "mflux": "0.23.0" }, - "footprint_peak_bytes": 21929337856, + "footprint_peak_bytes": 21877154888, "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", - "watched_peak_bytes": 21929337856, - "mlx_peak_bytes": 19963475458, - "elapsed_seconds": 413.966 + "watched_peak_bytes": 21877154888, + "mlx_peak_bytes": 19963459074, + "elapsed_seconds": 397.941 } diff --git a/bench/results/tiers/schnell-1024.json b/bench/results/tiers/schnell-1024.json index 98bf445..c70c724 100644 --- a/bench/results/tiers/schnell-1024.json +++ b/bench/results/tiers/schnell-1024.json @@ -1,14 +1,14 @@ { "exit_code": 0, - "output": "/tmp/tier-schnell.png", + "output": "schnell.png", "memory_caps_gb": [ 20, 22 ], "model_kwargs": { "model": "schnell", - "df11_path": null, - "base_path": null, + "df11_path": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", + "base_path": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", "eval_policy": "per-block", "cache_limit": null, "fit_check": true @@ -30,6 +30,7 @@ "reserve_bytes": 2147483648, "budget_bytes": 26800603136, "ceiling_bytes": 24653119488, + "fit_budget_bytes": 24653119488, "budget_source": "device", "is_host": true, "label": "MEASURED" @@ -41,15 +42,16 @@ }, "applied": "host-caps" }, + "macos": "27.0.1", "model": "schnell", "df11": { "root": "~/.cache/huggingface/hub/models--DFloat11--FLUX.1-schnell-DF11/snapshots/51a428b928197e0531cb93d6e438941e2d0b247e", - "repo_id": "DFloat11/FLUX.1-schnell-DF11", + "repo_id": null, "revision": "51a428b928197e0531cb93d6e438941e2d0b247e" }, "base": { "root": "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9", - "repo_id": "black-forest-labs/FLUX.1-schnell", + "repo_id": null, "revision": "741f7c3ce8b383c54771c7003378a50191e9efe9" }, "eval_policy": "per-block", @@ -60,7 +62,7 @@ "phases": { "encode": 11518505446, "denoise": 21411230501, - "vae": 8808361940 + "vae": 8922261588 }, "peak_phase": "denoise", "peak_bytes": 21411230501, @@ -76,28 +78,28 @@ "peaks": { "label": "sampled at phase boundaries", "encode": { - "footprint_start": 358565016, - "mlx_peak": 10829791578, - "footprint_end": 10080869816, - "active_at_start": 115521672 + "footprint_start": 362808448, + "mlx_peak": 10976964826, + "footprint_end": 10218462816, + "active_at_start": 115550472 }, "set_load": { - "footprint_start": 10080869816, - "mlx_peak": 16315874440, - "footprint_end": 16579015424, - "active_at_start": 119717512 + "footprint_start": 10218462816, + "mlx_peak": 16315903240, + "footprint_end": 16609899264, + "active_at_start": 119746312 }, "denoise": { - "footprint_start": 16579015424, - "mlx_peak": 18725825760, - "footprint_end": 19133838008, - "active_at_start": 16315874440 + "footprint_start": 16609899264, + "mlx_peak": 18725854548, + "footprint_end": 19132789528, + "active_at_start": 16315903240 }, "vae": { - "footprint_start": 9721010320, - "mlx_peak": 10255949836, - "footprint_end": 853100104, - "active_at_start": 121814728 + "footprint_start": 8916276792, + "mlx_peak": 10255978624, + "footprint_end": 851920408, + "active_at_start": 121843516 } }, "decode_launches": 228, @@ -105,15 +107,16 @@ "encoder_loads": 0, "encoder_seconds": 0.0, "set_loads": 1, - "set_seconds": 27.269080250000115, + "set_seconds": 25.723129791091196, "forced_set_drops": 0 }, "versions": { "mlx": "0.32.2", - "mflux": "0.20.0" + "mflux": "0.23.0" }, - "footprint_peak_bytes": 21233332920, + "footprint_peak_bytes": 21232349976, "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", - "watched_peak_bytes": 21233332920, - "mlx_peak_bytes": 20438596158 -} \ No newline at end of file + "watched_peak_bytes": 21232349976, + "mlx_peak_bytes": 20489009842, + "elapsed_seconds": 106.051 +} diff --git a/bench/results/tiers/z-image-1024.json b/bench/results/tiers/z-image-1024.json index 84993c2..b727460 100644 --- a/bench/results/tiers/z-image-1024.json +++ b/bench/results/tiers/z-image-1024.json @@ -1,14 +1,14 @@ { "exit_code": 0, - "output": "/tmp/s2/z-image.png", + "output": "z-image.png", "memory_caps_gb": [ 20, 22 ], "model_kwargs": { "model": "z-image", - "df11_path": null, - "base_path": null, + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--Z-Image-DF11/snapshots/36d582560630bfe2f5b78183c4164c66d6919e5b", + "base_path": "~/.cache/huggingface/hub/models--Tongyi-MAI--Z-Image/snapshots/04cc4abb7c5069926f75c9bfde9ef43d49423021", "eval_policy": "per-block", "cache_limit": null, "fit_check": true @@ -30,6 +30,7 @@ "reserve_bytes": 2147483648, "budget_bytes": 26800603136, "ceiling_bytes": 24653119488, + "fit_budget_bytes": 24653119488, "budget_source": "device", "is_host": true, "label": "MEASURED" @@ -41,16 +42,17 @@ }, "applied": "host-caps" }, + "macos": "27.0.1", "family": "zimage", "model": "z-image", "df11": { "root": "~/.cache/huggingface/hub/models--mingyi456--Z-Image-DF11/snapshots/36d582560630bfe2f5b78183c4164c66d6919e5b", - "repo_id": "mingyi456/Z-Image-DF11", + "repo_id": null, "revision": "36d582560630bfe2f5b78183c4164c66d6919e5b" }, "base": { "root": "~/.cache/huggingface/hub/models--Tongyi-MAI--Z-Image/snapshots/04cc4abb7c5069926f75c9bfde9ef43d49423021", - "repo_id": "Tongyi-MAI/Z-Image", + "repo_id": null, "revision": "04cc4abb7c5069926f75c9bfde9ef43d49423021" }, "eval_policy": "per-block", @@ -66,10 +68,10 @@ "phases": { "encode": 10143248230, "denoise": 13101538864, - "vae": 13054789110 + "vae": 14115020240 }, - "peak_phase": "denoise", - "peak_bytes": 13101538864, + "peak_phase": "vae", + "peak_bytes": 14115020240, "budget_bytes": 24653119488, "fits": true }, @@ -83,28 +85,28 @@ "peaks": { "label": "sampled at phase boundaries", "encode": { - "footprint_start": 480806256, + "footprint_start": 472302768, "mlx_peak": 9177190600, - "footprint_end": 8999754928, + "footprint_end": 9054362704, "active_at_start": 6580692 }, "set_load": { - "footprint_start": 8999754928, + "footprint_start": 9054362704, "mlx_peak": 8401205524, - "footprint_end": 8883658144, + "footprint_end": 8875662776, "active_at_start": 11823560 }, "denoise": { - "footprint_start": 8883658144, - "mlx_peak": 10637496516, - "footprint_end": 11312586144, + "footprint_start": 8875662776, + "mlx_peak": 9719225080, + "footprint_end": 11233615240, "active_at_start": 8387910856 }, "vae": { - "footprint_start": 11312586144, - "mlx_peak": 13578549138, - "footprint_end": 9279348200, - "active_at_start": 8389484140 + "footprint_start": 11189083528, + "mlx_peak": 14811444118, + "footprint_end": 9806896592, + "active_at_start": 8388959852 } }, "decode_launches": 3400, @@ -112,15 +114,16 @@ "encoder_loads": 0, "encoder_seconds": 0.0, "set_loads": 1, - "set_seconds": 16.110344999993686, + "set_seconds": 12.381818374968134, "forced_set_drops": 0 }, "versions": { "mlx": "0.32.2", - "mflux": "0.20.0" + "mflux": "0.23.0" }, - "footprint_peak_bytes": 14115020240, + "footprint_peak_bytes": 15462030800, "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", - "watched_peak_bytes": 14115020240, - "mlx_peak_bytes": 13099464594 -} \ No newline at end of file + "watched_peak_bytes": 15462030800, + "mlx_peak_bytes": 14270362518, + "elapsed_seconds": 836.284 +} From b5c176c08d4753eb20bd51ffeb5aa7078d69f6ac Mon Sep 17 00:00:00 2001 From: Denis Ineshin Date: Wed, 28 Oct 2026 19:32:11 +0100 Subject: [PATCH 4/4] Document 0.2.1 --- CHANGELOG.md | 54 ++++++++++++- README.md | 221 ++++++++++++++++++++++++++++++--------------------- 2 files changed, 183 insertions(+), 92 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 479983f..d551842 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,57 @@ All notable changes to this project are documented here. The format follows ## [Unreleased] +## [0.2.1] - 2026-10-28 + +0.2.1 runs on mflux 0.23 and requires it; the `mflux` extra's pin excludes 0.20. FLUX.1-schnell, Z-Image, distilled +FLUX.2-klein-4B and Qwen-Image-2.1 were run again on mflux 0.23.0 on a 32 GB M1 Max (macOS 27.0.1, mlx 0.32.2). For +each, the final latents are equal to those of the same transformer streaming its BF16 weights (4 steps at 1024²), the +images are identical pixel for pixel, and the peaks stayed under the machine's 22.96 GiB budget +(`bench/results/identity/` and `bench/results/tiers/`). The distilled FLUX.2-klein-4B now has its first such record. ERNIE-Image and +ERNIE-Image-Turbo, Krea 2 Turbo and Raw, FLUX.2-klein-base-4B, base-9B and 9B, Z-Image-Turbo, FLUX.1-dev and +FLUX.1-Krea-dev, and the stopped 16 GB and 24 GB rows were not run again. Their numbers keep the mflux 0.20.0 they were +measured on, and the README's tier table says so in its new `mflux` column. + +Peaks at 1024² on that Mac, on mflux 0.23.0: FLUX.1-schnell 19.77 GiB (was 19.78), Z-Image 14.40 GiB (was 13.15 on mflux 0.20.0), +FLUX.2-klein-4B 11.85 GiB (was 11.88), Qwen-Image-2.1 20.37 GiB (was 20.42). Each is the file under +`bench/results/tiers/`; the 0.2.0 values were measured on mflux 0.20.0. + +### Changed + +- The `mflux` extra and the `bench` group require `mflux>=0.23,<0.24` and `pillow<13`, replacing `mflux>=0.20,<0.21`. + mflux 0.23.0 fails at import with Pillow 13. +- Qwen-Image-2.1 runs mflux's packed projections one block at a time. Its compiled step functions stay out of the + generation path, and a step function compiled between calls is refused. Both sides of its identity check run + without mflux's step compilation: one decodes the DF11 weights, the other streams the BF16 weights one block at a time + through a copy of mflux's own loop body. Stock mflux's compiled step can differ in the last bits from the same graph + run without that compilation, so an image from stock mflux's default path can differ in the last bits from this + package's, which matches the BF16 path run without it bit for bit. +- Z-Image and FLUX.2 Klein set the attributes mflux 0.23 reads when it generates, and refuse mflux's precision options. +- Each model's BF16 base repository (text encoders, VAE, tokenizer) is pinned to a fixed revision, the snapshot the + recorded runs used. A base you pass with `--base` is used as given. +- `mlx-dfloat generate` reports record the macOS version. +- The README's tier table has a last column, `mflux`, with the mflux version each row ran on. A row on the newest + version shows that version, an older row shows `0.20, not re-run on 0.23`, and a stopped run, whose file records no + version, shows `not recorded`. +- A tier run that the watchdog stops now records the mflux version in its abort file. The stopped rows already in the + table come from 0.2.0 files and still show `not recorded`. + +### Fixed + +- The memory estimate for Z-Image was too low on mflux 0.23. In the committed Z-Image run the VAE + step peaked at 14.40 GiB, above the old estimate of 13.15 GiB (the 0.20.0 peak it was calibrated on, not the 0.23 + one). The estimate is raised and now covers every committed Z-Image run. The raised estimate also applies to + Z-Image-Turbo, which shares it and was not measured on mflux 0.23. The estimate for Qwen-Image-2.1 was raised + too, by 0.04 GiB. One consequence for users: on a 24 GB Mac (not measured), a Z-Image or Z-Image-Turbo call is now + planned to drop the compressed transformer before the VAE decode, so the next call has to load it again. +- FLUX.1 loads `norm_out.linear.bias`, which mflux 0.23 builds and loads. +- When a base checkpoint lacks required weights, the error names the missing tensors. + +### Added + +- Identity records for the distilled FLUX.2-klein-4B (`bench/results/identity/flux2-klein-4b-1024/`): 4 steps at + guidance 1.0, latents equal and images identical pixel for pixel. + ## [0.2.0] - 2026-10-09 0.2.0 adds five image-model families to the FLUX.1 support of 0.1.0. Each runs from its DFloat11 checkpoint through @@ -338,6 +389,7 @@ The first release. - CI's mflux lane gets the Hub token for its test step only, so the three tokenizer tests against the gated FLUX.1 bases run there instead of skipping. -[Unreleased]: https://github.com/IonDen/mlx-dfloat/compare/v0.2.0...HEAD +[Unreleased]: https://github.com/IonDen/mlx-dfloat/compare/v0.2.1...HEAD +[0.2.1]: https://github.com/IonDen/mlx-dfloat/compare/v0.2.0...v0.2.1 [0.2.0]: https://github.com/IonDen/mlx-dfloat/compare/v0.1.0...v0.2.0 [0.1.0]: https://github.com/IonDen/mlx-dfloat/releases/tag/v0.1.0 diff --git a/README.md b/README.md index 69e55aa..10b5504 100644 --- a/README.md +++ b/README.md @@ -31,22 +31,22 @@ before it runs." width="720"> On a Mac, unified memory is the limit. The BF16 FLUX.1-dev transformer is about 24 GB,[^flux] more than the GPU on a 32 GB machine can comfortably hold. At 70% it comes in under that line, and unlike 4-bit or 8-bit -quantisation it changes nothing in the weights. On a 32 GB M1 Max a 1024² generation peaked at about 20 GiB for each -of the three FLUX.1 models, and between 11.88 GiB (FLUX.2-klein-4B) and 21.17 GiB (Krea 2 Turbo) for the others; no -other Mac has been measured. The numbers are under "Measured numbers" below. +quantisation it changes nothing in the weights. On a 32 GB M1 Max a 1024² generation peaked at 19.77 GiB for +FLUX.1-schnell (mflux 0.23.0) and at about 20 GiB for FLUX.1-dev and FLUX.1-Krea-dev (mflux 0.20.0), and between +11.85 GiB (FLUX.2-klein-4B) and 21.17 GiB (Krea 2 Turbo) for the others; no other Mac has been measured. The numbers are under "Measured numbers" below. ## Does this help me? | Your situation | What to use | |---|---| -| A 32 GB Mac, FLUX.1-schnell, FLUX.1-dev or FLUX.1-Krea-dev, and you want the BF16 weights unchanged | mlx-dfloat. Measured peak about 20 GiB at 1024² for all three. The final latents were checked bit for bit against BF16 on FLUX.1-schnell, not yet on the other two. Unpacking makes a denoising step 3.5 % slower on dev and 4.0 % slower on schnell than a control that runs the same graph on weights decoded in advance (Krea-dev's step was not timed). | -| Less memory or a faster step, and a different image is fine | mflux's own quantised transformer. Its 8-bit mode, the only one measured here, ran a step roughly 1.3 times faster and peaked about 5 GiB lower, under the same 2.5 GB cache limit; mflux's own generate sets no cache limit, so this is not quite its default setup (see "Measured numbers" for why the ratio is rough). It changes the weights and the output; by how much was not measured. | -| A 32 GB Mac, Z-Image-Turbo or Z-Image, with the weights unchanged | mlx-dfloat. Peak at 1024²: 13.1 to 14.3 GiB for Z-Image-Turbo over six runs, 13.15 GiB for Z-Image. See "Generate a Z-Image image". | -| A 32 GB Mac, FLUX.2 Klein (4B or 9B, base or distilled), with the weights unchanged | mlx-dfloat. Peak at 1024²: 11.88 to 12.78 GiB for the 4B models, 17.60 to 18.40 GiB for the 9B models. See "Generate a FLUX.2 Klein image". | +| A 32 GB Mac, FLUX.1-schnell, FLUX.1-dev or FLUX.1-Krea-dev, and you want the BF16 weights unchanged | mlx-dfloat. Measured peak at 1024²: 19.77 GiB for schnell (mflux 0.23.0), about 20 GiB for dev and Krea-dev (mflux 0.20.0, not re-run on 0.23). The final latents were checked bit for bit against BF16 on FLUX.1-schnell, not yet on the other two. Unpacking makes a denoising step 3.5 % slower on dev and 4.0 % slower on schnell than a control that runs the same graph on weights decoded in advance (mflux 0.20.0; the dev run skipped the charging gate; Krea-dev's step was not timed). | +| Less memory or a faster step, and a different image is fine | mflux's own quantised transformer. Its 8-bit mode, the only one measured here, ran a step roughly 1.3 times faster and peaked about 5 GiB lower (mflux 0.20.0; the dev run skipped the charging gate), under the same 2.5 GB cache limit; mflux's own generate sets no cache limit, so this is not quite its default setup (see "Measured numbers" for why the ratio is rough). It changes the weights and the output; by how much was not measured. | +| A 32 GB Mac, Z-Image-Turbo or Z-Image, with the weights unchanged | mlx-dfloat. Peak at 1024²: 13.1 to 14.3 GiB for Z-Image-Turbo over six runs (mflux 0.20, not re-run on 0.23), 14.40 GiB for Z-Image (mflux 0.23). See "Generate a Z-Image image". | +| A 32 GB Mac, FLUX.2 Klein (4B or 9B, base or distilled), with the weights unchanged | mlx-dfloat. Peak at 1024²: 11.85 to 12.78 GiB for the 4B models, 17.60 to 18.40 GiB for the 9B models. Only the distilled 4B was re-run on mflux 0.23; the other three are mflux 0.20 numbers. See "Generate a FLUX.2 Klein image". | | A 24 GB Mac, FLUX.2 Klein | CAPPED: FLUX.2-klein-base-4B passed, peak 12.29 GiB; FLUX.2-klein-base-9B is at risk (the fit check refuses it, and a forced run was stopped by the watchdog). See "Generate a FLUX.2 Klein image". | | A 16 GB Mac, FLUX.2-klein-base-4B | CAPPED: passed, peak 8.15 GiB. See "Generate a FLUX.2 Klein image". | | A 16 GB Mac, Z-Image-Turbo | CAPPED: at risk. The fit check refuses it, and a forced run was stopped by the watchdog at 9.26 GiB before the transformer ran. See "Generate a Z-Image image". | -| A 32 GB Mac, Qwen-Image-2.1, with the weights unchanged | mlx-dfloat. Peak at 1024²: 20.42 GiB. See "Generate a Qwen-Image-2.1 image". | +| A 32 GB Mac, Qwen-Image-2.1, with the weights unchanged | mlx-dfloat. Peak at 1024²: 20.37 GiB (mflux 0.23). See "Generate a Qwen-Image-2.1 image". | | A 16 GB or 24 GB Mac, Qwen-Image-2.1 | CAPPED: at risk. The fit check refuses it, and forced runs were stopped by the watchdog while the prompt was being encoded, before the transformer ran. See "Generate a Qwen-Image-2.1 image". | | A 32 GB Mac, ERNIE-Image or ERNIE-Image-Turbo, with the weights unchanged | mlx-dfloat. Peak at 1024²: 16.79 GiB for ERNIE-Image (50 steps), 17.34 GiB for ERNIE-Image-Turbo (8 steps). See "Generate an ERNIE-Image image". | | A 24 GB Mac, ERNIE-Image-Turbo | CAPPED: at risk. The fit check refuses it, 0.39 GiB over the budget; a forced run passed with a 14.39 GiB peak, 0.11 GiB under the watchdog ceiling. See "Generate an ERNIE-Image image". | @@ -60,15 +60,15 @@ other Mac has been measured. The numbers are under "Measured numbers" below. ## Status -Alpha: version 0.2.0 is on PyPI. +Alpha: version 0.2.1 is on PyPI. It requires mflux 0.23; the `mflux` extra's pin excludes 0.20. Released in 0.1.0: - All four published versions of the DFloat11 checkpoint format decode to exactly the BF16 originals: every compressed tensor of Qwen3-4B, and sampled blocks of FLUX.1-schnell, FLUX.1-Krea-dev, Qwen-Image-Edit and Qwen-Image-Edit-2509. -- A Metal kernel decodes to the same bits as the CPU reference decoder, at about 50 GB/s on an M1 Max on its default - path, timed on Qwen3-4B and FLUX.1-schnell groups with `uv run python -m scripts.bench_decode_kernel --df11 +- A Metal kernel decodes to the same bits as the CPU reference decoder, at about 50 GB/s of BF16 output per second on + an M1 Max on its default path, timed on Qwen3-4B and FLUX.1-schnell groups with `uv run python -m scripts.bench_decode_kernel --df11 --groups --out decode.json`. The records, `bench/results/decode/qwen3-4b-layer0.json` and `bench/results/decode/flux1-schnell.json`, give 49.9 to 54.2 GB/s, each the median of five timed repetitions (M1 Max, macOS 27.0, mlx 0.32.2). @@ -111,7 +111,8 @@ Released in 0.2.0: equals the BF16 original in `Qwen/Qwen-Image-2.1` (`bench/results/parity/qwen-image-2.1/full-vs-bf16.json`), and the GPU decoder gives the same bits as the CPU reference on all 225 (`bench/results/parity/qwen-image-2.1/kernel-vs-reference.json`). Its final latents match, bit for bit, those of - the same transformer streaming its BF16 weights one block at a time, and the images match pixel for pixel. That check + the same transformer streaming its BF16 weights one block at a time (against the uncompiled BF16 path; see "Generate + a Qwen-Image-2.1 image"), and the images match pixel for pixel. That check ran at 1024² for 4 steps at guidance 4.0. The negative prompt was a single space, so classifier-free guidance ran (record: `bench/results/identity/qwen-image-2.1-1024/compare.json`). `scripts/verify_image.py` runs both sides without mflux's step compilation. @@ -145,6 +146,23 @@ Released in 0.2.0: was measured. - `mlx-dfloat selftest` checks the GPU decoder on your Mac (see "Check the GPU decoder on your Mac"). +Released in 0.2.1: + +- Generation runs on mflux 0.23 and requires it. The `mflux` extra and the benchmark group now ask for + `mflux>=0.23,<0.24` and `pillow<13`, because mflux 0.23.0 fails to import with Pillow 13. +- FLUX.1-schnell, Z-Image, distilled FLUX.2-klein-4B and Qwen-Image-2.1 were run again on mflux 0.23.0. For each, the + final latents match, bit for bit, those of the same transformer streaming its BF16 weights one block at a time, and + the images match pixel for pixel. Each check ran 4 steps at 1024² (`bench/results/identity/flux1-schnell-1024/compare.json`, + `z-image-1024/compare.json`, `flux2-klein-4b-1024/compare.json` and `qwen-image-2.1-1024/compare.json`). For + Qwen-Image-2.1 both sides run without mflux's step compilation, so "bit for bit" there means against the BF16 path run the same way; stock + mflux's compiled path can differ in the last bits. The distilled FLUX.2-klein-4B had no such record before; its check used guidance + 1.0. +- ERNIE-Image, ERNIE-Image-Turbo, Krea 2 Turbo and Raw, FLUX.2-klein-base-4B, base-9B and 9B, Z-Image-Turbo, FLUX.1-dev + and FLUX.1-Krea-dev, and the stopped 16 GB and 24 GB rows were not run again. Their numbers are from mflux 0.20.0, and + the last column of the tier table says so row by row. +- Each model's BF16 base repository is pinned to a fixed revision, and a base checkpoint that lacks required weights is + refused with an error that names the missing tensors. + Earlier step-time numbers, measured with a smaller MLX cache limit, are in the 0.1.0 entry of the changelog. ## Install @@ -155,9 +173,13 @@ mlx-dfloat is for Macs with Apple Silicon and needs Python 3.11 or newer. Instal pip install "mlx-dfloat[mflux]" ``` -Without the extra you get the checkpoint reader and the decoder. The `mflux` extra installs mflux 0.20, which the -`mlx-dfloat generate` command needs for FLUX.1, Z-Image, FLUX.2 Klein, Qwen-Image-2.1, ERNIE-Image and Krea 2 -generation. +Without the extra you get the checkpoint reader and the decoder. The `mflux` extra installs mflux 0.23 (`>=0.23,<0.24`) +and `pillow<13`, which the `mlx-dfloat generate` command needs for FLUX.1, Z-Image, FLUX.2 Klein, Qwen-Image-2.1, +ERNIE-Image and Krea 2 generation. This version requires mflux 0.23, and the `mflux` extra's pin excludes 0.20. + +Each model's DFloat11 checkpoint is pinned to a fixed revision on the Hub. Each model's BF16 base repository (text +encoders, VAE, tokenizer) is pinned to a fixed revision too, the snapshot the recorded runs used; a `--df11` or `--base` +you pass is used as given. The parity scripts and the benchmark live in the repository, not in the package, so run them from a checkout with [uv](https://docs.astral.sh/uv/): @@ -261,26 +283,27 @@ returns; a wired limit your process already set is left alone. Only the command on a Mac near its memory limit. Measured on an M1 Max at 1024², one image per model, with the commands under "Measured numbers": FLUX.1-schnell (4 -steps) peaked at 19.78 GiB, FLUX.1-dev (20 steps, guidance 3.5) at 19.96 GiB and FLUX.1-Krea-dev (20 steps, guidance -3.5) at 20.04 GiB (`bench/results/tiers/schnell-1024.json`, `dev-1024.json` and `krea-dev-1024.json`). The machine's -budget, its recommended working set minus a 2 GiB reserve, is 22.96 GiB, so all three stay under it. Those records -hold no wall time. Earlier single runs on 2026-09-28 and 29, not committed, took 1 minute 54 seconds for -FLUX.1-schnell, imports included, and about 7 minutes each for FLUX.1-dev and FLUX.1-Krea-dev. +steps, macOS 27.0.1, mflux 0.23.0) peaked at 19.77 GiB, FLUX.1-dev (20 steps, guidance 3.5) at 19.96 GiB and +FLUX.1-Krea-dev (20 steps, guidance 3.5) at 20.04 GiB, the last two on mflux 0.20.0 and not re-run on 0.23 +(`bench/results/tiers/schnell-1024.json`, `dev-1024.json` and `krea-dev-1024.json`). The machine's budget, its +recommended working set minus a 2 GiB reserve, is 22.96 GiB, so all three stay under it. The schnell record holds its +wall time, 106 seconds (`elapsed_seconds`); the dev and Krea-dev records hold none. `scripts/verify_image.py` checked FLUX.1-schnell's output against the same transformer streaming its BF16 shards block by block instead of the compressed set. After 4 steps at 1024², the final latents matched bit for bit and the two -images matched pixel for pixel (record: `bench/results/identity/flux1-schnell-1024/compare.json`). FLUX.1-dev and +images matched pixel for pixel on mflux 0.23.0 (record: `bench/results/identity/flux1-schnell-1024/compare.json`). +FLUX.1-dev and FLUX.1-Krea-dev were not checked this way, because their BF16 transformers are gated and were not on disk to compare against. Generating a second image in the same process pays a reload. On a 32 GB Mac every call drops the compressed -transformer before the VAE decode, whatever the image size. At 1024², holding it next to the decode measured 23.29 GiB, -over the budget (a run on 2026-09-28, not committed). Smaller sizes have not been measured, so the estimate assumes the -decode needs at least as much there. The next call loads the transformer again, which took 27 to 28 seconds in the -three committed runs (`lifecycle.set_seconds`). A Mac with a larger budget keeps the transformer loaded between calls. +transformer before the VAE decode at 1024². Holding it next to the decode measured 23.29 GiB, over the budget (mflux +0.20.0, a run on 2026-09-28, not committed). Smaller sizes have not been measured, so the estimate assumes the decode +needs at least as much there and drops the transformer at every size. The next call loads the transformer again, which +took 25.7 seconds for FLUX.1-schnell on mflux 0.23 and 27.4 to 28.2 seconds for FLUX.1-dev and FLUX.1-Krea-dev on mflux 0.20 in the committed runs (`lifecycle.set_seconds`). A Mac with a larger budget keeps the transformer loaded between calls. -Even a repeated prompt pays the reload. In a test on 2026-09-28, not committed, three consecutive 4-step FLUX.1-schnell -calls with the same prompt each took about 103 seconds and peaked between 19.77 and 19.98 GiB. A different prompt also +Even a repeated prompt pays the reload. In a test on 2026-09-28 (mflux 0.20.0, not committed), three consecutive 4-step +FLUX.1-schnell calls with the same prompt each paid the reload and peaked between 19.77 and 19.98 GiB. A different prompt also reloads the text encoders, because they and the transformer are never resident together. `model.encode(*prompts)` pays that encoder reload once for several prompts: it encodes all of them before the compressed transformer loads. @@ -336,14 +359,17 @@ runs the step without that compilation on every chip, because it unpacks each bl runs, and a compiled step cannot pause between blocks to do that. What that costs against stock mflux was not measured: the only Mac measured here is an M1 Max, and no Z-Image step time is reported. -Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0) at 1024², seed 42, one image per model: -Z-Image-Turbo (9 steps) peaked at 13.11 GiB and Z-Image (50 steps, guidance 4.0) at 13.15 GiB. In both runs MLX's own -memory peaked during the VAE decode. The peak varies between runs of the same command: six runs of the Turbo command on -2026-10-08 peaked between 13.11 and 14.26 GiB, and the difference came from the VAE decode. The five runs besides the -tier row's are in `bench/results/repeats/z-image-turbo-1024/`. +Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2) at 1024², seed 42, one image per model: Z-Image-Turbo (9 +steps, mflux 0.20.0, not re-run on 0.23) peaked at 13.11 GiB and Z-Image (50 steps, guidance 4.0, mflux 0.23.0) at +14.40 GiB. In both runs MLX's own memory peaked during the VAE decode. The peak varies between runs of the same +command: six runs of the Turbo command on 2026-10-08 (mflux 0.20.0) peaked between 13.11 and 14.26 GiB. The spread is in +the process footprint the OS reports during the VAE decode, not in MLX's own peak: in each of the five repeats MLX's +peak in that phase was 13.35 GiB (the tier row's run read 12.20 GiB). The five runs besides the tier row's are in +`bench/results/repeats/z-image-turbo-1024/`. What was checked differs between the two models and is listed under "Status". For Z-Image, the final latents and the -image match the same transformer running on its BF16 weights (`bench/results/identity/z-image-1024/compare.json`). +image match the same transformer running on its BF16 weights, re-checked on mflux 0.23.0 +(`bench/results/identity/z-image-1024/compare.json`). Z-Image-Turbo has no BF16 original to compare with. On a 16 GB Mac, Z-Image-Turbo is at risk. No 16 GB Mac was measured: the run behind that row is CAPPED, made on the 32 @@ -414,16 +440,20 @@ model, just as mflux's Python API does not. As with Z-Image, this path runs the denoising step without mflux's step compilation (see "Generate a Z-Image image" for why), and what that costs against stock mflux was not measured. -Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0) at 1024², seed 42, one image per model, with -the commands under "Measured numbers": FLUX.2-klein-4B (4 steps) peaked at 11.88 GiB in 35 s, FLUX.2-klein-base-4B -(50 steps, guidance 4.0) at 12.78 GiB in 8.9 minutes, FLUX.2-klein-9B (4 steps) at 18.40 GiB in 1.3 minutes, and -FLUX.2-klein-base-9B (50 steps, guidance 4.0) at 17.60 GiB in 19.9 minutes. The times are the `elapsed_seconds` of -each run's report, from `bench/results/tiers/flux2-klein-4b-1024.json` to `flux2-klein-base-9b-1024.json`. In all -four runs MLX's own memory peaked during the VAE decode, and no run had to drop the compressed transformer before it. - -What was checked differs per model and is listed under "Status": only FLUX.2-klein-base-4B had its final latents -compared with the same transformer running on its BF16 weights -(`bench/results/identity/flux2-klein-base-4b-1024/compare.json`). +Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2) at 1024², seed 42, one image per model, with the commands +under "Measured numbers": FLUX.2-klein-4B (4 steps, mflux 0.23.0) peaked at 11.85 GiB in 34 s (the `elapsed_seconds` of +`bench/results/tiers/flux2-klein-4b-1024.json`). The other three ran on mflux 0.20.0 and were not re-run on 0.23: +FLUX.2-klein-base-4B (50 steps, guidance 4.0) peaked at 12.78 GiB, FLUX.2-klein-9B (4 steps) at 18.40 GiB, and +FLUX.2-klein-base-9B (50 steps, guidance 4.0) at 17.60 GiB (`flux2-klein-base-4b-1024.json`, `flux2-klein-9b-1024.json` +and `flux2-klein-base-9b-1024.json` in the same folder). In all four runs MLX's own memory peaked during the VAE decode, +and no run had to drop the compressed transformer before it. + +What was checked differs per model and is listed under "Status": FLUX.2-klein-base-4B and FLUX.2-klein-4B had their +final latents compared with the same transformer running on its BF16 weights, and in both the latents are equal and +the images identical pixel for pixel (`bench/results/identity/flux2-klein-base-4b-1024/compare.json`, run on mflux +0.20.0, and `bench/results/identity/flux2-klein-4b-1024/compare.json`, run on mflux 0.23.0, 4 steps at 1024² and +guidance 1.0). +The two 9B models have no such comparison. No 16 GB or 24 GB Mac was measured. The smaller-Mac rows are CAPPED runs on the 32 GB Mac under the caps mlx-dfloat installs on those Macs. FLUX.2-klein-base-4B (50 steps, guidance 4.0) passed under both: it peaked at 8.15 GiB under a @@ -484,13 +514,19 @@ image.save("lighthouse.png") mflux compiles Qwen-Image-2.1's denoising step. As with Z-Image and FLUX.2 Klein, this path runs the step without that compilation (see "Generate a Z-Image image" for why), and what that costs against stock mflux was not -measured. +measured. It also runs mflux's packed projections one block at a time. + +Both sides of the Qwen-Image-2.1 identity check run without mflux's step compilation: one decodes the DF11 weights, the other streams the +BF16 weights one block at a time through a copy of mflux's own loop body. Stock mflux's compiled step can differ in +the last bits from the same graph run without that compilation, so an image from stock mflux's default path can differ +in the last bits from this package's, which matches the BF16 path run without it bit for bit. The DFloat11 path always +runs without mflux's step compilation. -Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0) at 1024², seed 42, 40 steps at guidance 1.0, -with the command under "Measured numbers": the run peaked at 20.42 GiB and took 6.9 minutes (the `elapsed_seconds` of +Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.23.0) at 1024², seed 42, 40 steps at guidance 1.0, +with the command under "Measured numbers": the run peaked at 20.37 GiB and took 6.6 minutes (the `elapsed_seconds` of `bench/results/tiers/qwen-image-2.1-1024.json`). MLX's own memory peaked during the VAE decode. On this 32 GB Mac the -compressed transformer stays loaded through the VAE decode: the fit estimate predicts 21.2 GiB for that step (the -report's `fit.phases.vae`), under the 22.96 GiB budget. +compressed transformer stays loaded through the VAE decode. The report's estimate for that step is 21.2 GiB +(`fit.phases.vae`), made before this release raised the estimate to 21.3 GiB; both are under the 22.96 GiB budget. On a 16 GB or 24 GB Mac, Qwen-Image-2.1 is at risk. No such Mac was measured: the runs behind those rows are CAPPED, made on the 32 GB M1 Max under the caps mlx-dfloat installs on a Mac of that size (`--tier 16` and `--tier 24`). The fit @@ -562,8 +598,8 @@ without that compilation on every chip (see "Generate a Z-Image image" for why). timestep conditioning in float32, so the hidden state passed from block to block is float32, not BF16; this path runs mflux's own block code, so it does the same. No ERNIE-Image step time is reported, against stock mflux or otherwise. -Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0) at 1024², seed 42, with the commands under -"Measured numbers": ERNIE-Image-Turbo (8 steps at guidance 1.0) peaked at 17.34 GiB and took 2.9 minutes, and +Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0, not re-run on 0.23) at 1024², seed 42, with the +commands under "Measured numbers": ERNIE-Image-Turbo (8 steps at guidance 1.0) peaked at 17.34 GiB and took 2.9 minutes, and ERNIE-Image (50 steps at guidance 4.0) peaked at 16.79 GiB and took 29.9 minutes (the `elapsed_seconds` of `bench/results/tiers/ernie-image-turbo-1024.json` and `ernie-image-1024.json`). In both runs MLX's own memory peaked during the VAE decode, and the compressed transformer stayed loaded through it. The fit estimate, mlx-dfloat's @@ -613,7 +649,7 @@ mlx-dfloat generate --model krea-2-raw \ The defaults are 8 steps for Turbo and 25 for Raw, both at guidance 1.0 with the `er_sde` scheduler. Turbo's are those of mflux's own Krea 2 command. That command does not run Raw, so Raw's 25 steps are mflux's fallback for a model without a step default of its own.[^krea-defaults] At guidance 1.0 each step calls the transformer once. Raw's model card runs 52 steps at guidance 3.5 in every example. -To use that recipe, pass `--steps 52 --guidance 3.5` and raise the time limit. A single run took 56.2 minutes here, +To use that recipe, pass `--steps 52 --guidance 3.5` and raise the time limit. A single run on mflux 0.20.0 took 56.2 minutes here, close to the one hour after which the command's watchdog stops a run unless you pass `--wall-budget`: ``` @@ -678,8 +714,8 @@ different development commits). Most of that 1.4 GiB comes from emptying the cac 0.14 GiB on its own.[^krea-nonblock] The price is about 1.3 GB more decoded weights resident at the start of each call, decoded again every call; that decode time was not measured on its own. -Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0) at 1024², seed 42, with the commands under -"Measured numbers", one run each: Krea 2 Turbo (8 steps at guidance 1.0) peaked at 21.17 GiB and took 5.1 minutes, +Measured on an M1 Max (32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0, not re-run on 0.23) at 1024², seed 42, with the +commands under "Measured numbers", one run each: Krea 2 Turbo (8 steps at guidance 1.0) peaked at 21.17 GiB and took 5.1 minutes, and Krea 2 Raw (25 steps at guidance 1.0) peaked at 21.15 GiB in 14.1 minutes. With the card's recipe (52 steps at guidance 3.5) Raw peaked at 21.15 GiB in 56.2 minutes (`bench/results/recipes/krea-2-raw-1024-52steps-guidance3.5.json`; it is not a row of the table below, which shows no @@ -689,7 +725,7 @@ The fit estimate is mlx-dfloat's prediction of each step's peak memory, checked the denoising step at 22.58 GiB, under this Mac's 22.96 GiB budget. That is about 1.4 GiB above the measured peaks, because the estimate counts a full buffer cache that the emptied cache never reaches. The estimate implies that holding the compressed transformer through the VAE decode would go over the budget. So on a 32 GB Mac every call -drops it before the decode and loads it again at the next call (31 s in the Turbo run). +drops it before the decode and loads it again at the next call (31 s in the Turbo run, on mflux 0.20.0). On a 16 GB or 24 GB Mac, both models are at risk. No such Mac was measured: the runs behind those rows are CAPPED, made on the 32 GB M1 Max under the caps mlx-dfloat installs on a Mac of that size (`--tier 16` and `--tier 24`). The fit @@ -868,41 +904,43 @@ two Krea 2 Raw rows would look the same. In the last four the watchdog stops the `--no-fit-check` the last four are refused with exit 2 (`bench/results/refusals/krea-2-1024-tier16.json`, `krea-2-raw-1024-tier16.json`, `krea-2-1024-tier24.json` and `krea-2-raw-1024-tier24.json`). -Measured on an Apple M1 Max, 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0. The `git` fields in the committed result -files name development commits that were folded into a release; they tell the runs apart, but they cannot be fetched -from the public repository. +Measured on an Apple M1 Max, 32 GB, macOS 27.0.1, mlx 0.32.2. The last column, `mflux`, gives the mflux version of each +row. FLUX.1-schnell, Z-Image, FLUX.2-klein-4B and Qwen-Image-2.1 were run on mflux 0.23.0 and show `0.23`. Every other +completed row shows `0.20, not re-run on 0.23`. The rows the watchdog stopped carry no version and date from 0.2.0. +The `git` fields in the committed result files name development commits that were folded into a release; they tell the +runs apart. -| Mac | Working set − reserve | Model | DF11 size | Peak (watched) | Peak footprint | Peak MLX (sampled active + cache, or exact phase peak) | Label | Status | Limits | Result | -|---|---|---|---|---|---|---|---|---|---|---| -| 32 GB | 22.96 GiB | FLUX.1-dev | 15.21 GiB | 19.96 GiB | 19.96 GiB | 19.19 GiB | MEASURED | target | host caps | `bench/results/tiers/dev-1024.json` | -| 32 GB | 22.96 GiB | ERNIE-Image | 10.17 GiB | 16.79 GiB | 16.79 GiB | 15.21 GiB | MEASURED | target | host caps | `bench/results/tiers/ernie-image-1024.json` | -| 24 GB | 14.50 GiB | ERNIE-Image-Turbo | 10.17 GiB | 14.39 GiB | 14.39 GiB | 13.35 GiB | CAPPED | target | mlx-dfloat caps for the tier | `bench/results/tiers/ernie-image-turbo-1024-tier24.json` | -| 32 GB | 22.96 GiB | ERNIE-Image-Turbo | 10.17 GiB | 17.34 GiB | 17.34 GiB | 16.16 GiB | MEASURED | target | host caps | `bench/results/tiers/ernie-image-turbo-1024.json` | -| 32 GB | 22.96 GiB | FLUX.2-klein-4B | 4.90 GiB | 11.88 GiB | 11.88 GiB | 11.90 GiB | MEASURED | target | host caps | `bench/results/tiers/flux2-klein-4b-1024.json` | -| 32 GB | 22.96 GiB | FLUX.2-klein-9B | 11.47 GiB | 18.40 GiB | 18.40 GiB | 17.57 GiB | MEASURED | target | host caps | `bench/results/tiers/flux2-klein-9b-1024.json` | -| 16 GB | 9.17 GiB | FLUX.2-klein-base-4B | 4.90 GiB | 8.15 GiB | 8.15 GiB | 7.61 GiB | CAPPED | target | mlx-dfloat caps for the tier | `bench/results/tiers/flux2-klein-base-4b-1024-tier16.json` | -| 24 GB | 14.50 GiB | FLUX.2-klein-base-4B | 4.90 GiB | 12.29 GiB | 12.29 GiB | 10.91 GiB | CAPPED | target | mlx-dfloat caps for the tier | `bench/results/tiers/flux2-klein-base-4b-1024-tier24.json` | -| 32 GB | 22.96 GiB | FLUX.2-klein-base-4B | 4.90 GiB | 12.78 GiB | 12.78 GiB | 11.91 GiB | MEASURED | target | host caps | `bench/results/tiers/flux2-klein-base-4b-1024.json` | -| 32 GB | 22.96 GiB | FLUX.2-klein-base-9B | 11.47 GiB | 17.60 GiB | 17.60 GiB | 16.58 GiB | MEASURED | target | host caps | `bench/results/tiers/flux2-klein-base-9b-1024.json` | -| 32 GB | 22.96 GiB | Krea 2 Turbo | 16.34 GiB | 21.17 GiB | 21.17 GiB | 20.64 GiB | MEASURED | target | host caps | `bench/results/tiers/krea-2-1024.json` | -| 32 GB | 22.96 GiB | Krea 2 Raw | 16.34 GiB | 21.15 GiB | 21.15 GiB | 20.64 GiB | MEASURED | target | host caps | `bench/results/tiers/krea-2-raw-1024.json` | -| 32 GB | 22.96 GiB | FLUX.1-Krea-dev | 15.21 GiB | 20.04 GiB | 20.04 GiB | 19.19 GiB | MEASURED | target | host caps | `bench/results/tiers/krea-dev-1024.json` | -| 32 GB | 22.96 GiB | Qwen-Image-2.1 | 9.06 GiB | 20.42 GiB | 20.42 GiB | 19.72 GiB | MEASURED | target | host caps | `bench/results/tiers/qwen-image-2.1-1024.json` | -| 32 GB | 22.96 GiB | FLUX.1-schnell | 15.19 GiB | 19.78 GiB | 19.78 GiB | 19.03 GiB | MEASURED | target | host caps | `bench/results/tiers/schnell-1024.json` | -| 32 GB | 22.96 GiB | Z-Image | 7.80 GiB | 13.15 GiB | 13.15 GiB | 12.65 GiB | MEASURED | target | host caps | `bench/results/tiers/z-image-1024.json` | -| 32 GB | 22.96 GiB | Z-Image-Turbo | 7.80 GiB | 13.11 GiB | 13.11 GiB | 12.20 GiB | MEASURED | target | host caps | `bench/results/tiers/z-image-turbo-1024.json` | -| 16 GB | 9.17 GiB | ERNIE-Image | not recorded | at least 9.44 GiB (stopped after 23.3 s) | at least 9.44 GiB | at least 9.03 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/ernie-image-1024-tier16.json` | -| 24 GB | 14.50 GiB | ERNIE-Image | not recorded | at least 14.80 GiB (stopped after 20.6 s) | at least 14.80 GiB | at least 14.25 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/ernie-image-1024-tier24.json` | -| 16 GB | 9.17 GiB | ERNIE-Image-Turbo | not recorded | at least 9.22 GiB (stopped after 18.6 s) | at least 9.22 GiB | at least 8.75 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/ernie-image-turbo-1024-tier16.json` | -| 24 GB | 14.50 GiB | FLUX.2-klein-base-9B | not recorded | at least 14.64 GiB (stopped after 28.5 s) | at least 14.64 GiB | at least 14.26 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/flux2-klein-base-9b-1024-tier24.json` | -| 16 GB | 9.17 GiB | Krea 2 Turbo | not recorded | at least 9.50 GiB (stopped after 17.1 s) | at least 9.50 GiB | at least 9.39 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-1024-tier16.json` | -| 24 GB | 14.50 GiB | Krea 2 Turbo | not recorded | at least 15.05 GiB (stopped after 29.7 s) | at least 15.05 GiB | at least 14.95 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-1024-tier24.json` | -| 16 GB | 9.17 GiB | Krea 2 Raw | not recorded | at least 9.49 GiB (stopped after 17.7 s) | at least 9.49 GiB | at least 9.39 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-raw-1024-tier16.json` | -| 24 GB | 14.50 GiB | Krea 2 Raw | not recorded | at least 14.91 GiB (stopped after 32.0 s) | at least 14.91 GiB | at least 14.41 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-raw-1024-tier24.json` | -| 16 GB | 9.17 GiB | Qwen-Image-2.1 | not recorded | at least 9.28 GiB (stopped after 4.4 s) | at least 8.90 GiB | at least 9.28 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/qwen-image-2.1-1024-tier16.json` | -| 24 GB | 14.50 GiB | Qwen-Image-2.1 | not recorded | at least 14.59 GiB (stopped after 4.8 s) | at least 14.59 GiB | at least 14.31 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/qwen-image-2.1-1024-tier24.json` | -| 16 GB | 9.17 GiB | Z-Image-Turbo | not recorded | at least 9.26 GiB (stopped after 4.7 s) | at least 9.26 GiB | at least 8.81 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/z-image-turbo-1024-tier16.json` | +| Mac | Working set − reserve | Model | DF11 size | Peak (watched) | Peak footprint | Peak MLX (sampled active + cache, or exact phase peak) | Label | Status | Limits | Result | mflux | +|---|---|---|---|---|---|---|---|---|---|---|---| +| 32 GB | 22.96 GiB | FLUX.1-dev | 15.21 GiB | 19.96 GiB | 19.96 GiB | 19.19 GiB | MEASURED | target | host caps | `bench/results/tiers/dev-1024.json` | 0.20, not re-run on 0.23 | +| 32 GB | 22.96 GiB | ERNIE-Image | 10.17 GiB | 16.79 GiB | 16.79 GiB | 15.21 GiB | MEASURED | target | host caps | `bench/results/tiers/ernie-image-1024.json` | 0.20, not re-run on 0.23 | +| 24 GB | 14.50 GiB | ERNIE-Image-Turbo | 10.17 GiB | 14.39 GiB | 14.39 GiB | 13.35 GiB | CAPPED | target | mlx-dfloat caps for the tier | `bench/results/tiers/ernie-image-turbo-1024-tier24.json` | 0.20, not re-run on 0.23 | +| 32 GB | 22.96 GiB | ERNIE-Image-Turbo | 10.17 GiB | 17.34 GiB | 17.34 GiB | 16.16 GiB | MEASURED | target | host caps | `bench/results/tiers/ernie-image-turbo-1024.json` | 0.20, not re-run on 0.23 | +| 32 GB | 22.96 GiB | FLUX.2-klein-4B | 4.90 GiB | 11.85 GiB | 11.85 GiB | 10.90 GiB | MEASURED | target | host caps | `bench/results/tiers/flux2-klein-4b-1024.json` | 0.23 | +| 32 GB | 22.96 GiB | FLUX.2-klein-9B | 11.47 GiB | 18.40 GiB | 18.40 GiB | 17.57 GiB | MEASURED | target | host caps | `bench/results/tiers/flux2-klein-9b-1024.json` | 0.20, not re-run on 0.23 | +| 16 GB | 9.17 GiB | FLUX.2-klein-base-4B | 4.90 GiB | 8.15 GiB | 8.15 GiB | 7.61 GiB | CAPPED | target | mlx-dfloat caps for the tier | `bench/results/tiers/flux2-klein-base-4b-1024-tier16.json` | 0.20, not re-run on 0.23 | +| 24 GB | 14.50 GiB | FLUX.2-klein-base-4B | 4.90 GiB | 12.29 GiB | 12.29 GiB | 10.91 GiB | CAPPED | target | mlx-dfloat caps for the tier | `bench/results/tiers/flux2-klein-base-4b-1024-tier24.json` | 0.20, not re-run on 0.23 | +| 32 GB | 22.96 GiB | FLUX.2-klein-base-4B | 4.90 GiB | 12.78 GiB | 12.78 GiB | 11.91 GiB | MEASURED | target | host caps | `bench/results/tiers/flux2-klein-base-4b-1024.json` | 0.20, not re-run on 0.23 | +| 32 GB | 22.96 GiB | FLUX.2-klein-base-9B | 11.47 GiB | 17.60 GiB | 17.60 GiB | 16.58 GiB | MEASURED | target | host caps | `bench/results/tiers/flux2-klein-base-9b-1024.json` | 0.20, not re-run on 0.23 | +| 32 GB | 22.96 GiB | Krea 2 Turbo | 16.34 GiB | 21.17 GiB | 21.17 GiB | 20.64 GiB | MEASURED | target | host caps | `bench/results/tiers/krea-2-1024.json` | 0.20, not re-run on 0.23 | +| 32 GB | 22.96 GiB | Krea 2 Raw | 16.34 GiB | 21.15 GiB | 21.15 GiB | 20.64 GiB | MEASURED | target | host caps | `bench/results/tiers/krea-2-raw-1024.json` | 0.20, not re-run on 0.23 | +| 32 GB | 22.96 GiB | FLUX.1-Krea-dev | 15.21 GiB | 20.04 GiB | 20.04 GiB | 19.19 GiB | MEASURED | target | host caps | `bench/results/tiers/krea-dev-1024.json` | 0.20, not re-run on 0.23 | +| 32 GB | 22.96 GiB | Qwen-Image-2.1 | 9.06 GiB | 20.37 GiB | 20.37 GiB | 19.72 GiB | MEASURED | target | host caps | `bench/results/tiers/qwen-image-2.1-1024.json` | 0.23 | +| 32 GB | 22.96 GiB | FLUX.1-schnell | 15.19 GiB | 19.77 GiB | 19.77 GiB | 19.08 GiB | MEASURED | target | host caps | `bench/results/tiers/schnell-1024.json` | 0.23 | +| 32 GB | 22.96 GiB | Z-Image | 7.80 GiB | 14.40 GiB | 14.40 GiB | 13.79 GiB | MEASURED | target | host caps | `bench/results/tiers/z-image-1024.json` | 0.23 | +| 32 GB | 22.96 GiB | Z-Image-Turbo | 7.80 GiB | 13.11 GiB | 13.11 GiB | 12.20 GiB | MEASURED | target | host caps | `bench/results/tiers/z-image-turbo-1024.json` | 0.20, not re-run on 0.23 | +| 16 GB | 9.17 GiB | ERNIE-Image | not recorded | at least 9.44 GiB (stopped after 23.3 s) | at least 9.44 GiB | at least 9.03 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/ernie-image-1024-tier16.json` | not recorded | +| 24 GB | 14.50 GiB | ERNIE-Image | not recorded | at least 14.80 GiB (stopped after 20.6 s) | at least 14.80 GiB | at least 14.25 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/ernie-image-1024-tier24.json` | not recorded | +| 16 GB | 9.17 GiB | ERNIE-Image-Turbo | not recorded | at least 9.22 GiB (stopped after 18.6 s) | at least 9.22 GiB | at least 8.75 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/ernie-image-turbo-1024-tier16.json` | not recorded | +| 24 GB | 14.50 GiB | FLUX.2-klein-base-9B | not recorded | at least 14.64 GiB (stopped after 28.5 s) | at least 14.64 GiB | at least 14.26 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/flux2-klein-base-9b-1024-tier24.json` | not recorded | +| 16 GB | 9.17 GiB | Krea 2 Turbo | not recorded | at least 9.50 GiB (stopped after 17.1 s) | at least 9.50 GiB | at least 9.39 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-1024-tier16.json` | not recorded | +| 24 GB | 14.50 GiB | Krea 2 Turbo | not recorded | at least 15.05 GiB (stopped after 29.7 s) | at least 15.05 GiB | at least 14.95 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-1024-tier24.json` | not recorded | +| 16 GB | 9.17 GiB | Krea 2 Raw | not recorded | at least 9.49 GiB (stopped after 17.7 s) | at least 9.49 GiB | at least 9.39 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-raw-1024-tier16.json` | not recorded | +| 24 GB | 14.50 GiB | Krea 2 Raw | not recorded | at least 14.91 GiB (stopped after 32.0 s) | at least 14.91 GiB | at least 14.41 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/krea-2-raw-1024-tier24.json` | not recorded | +| 16 GB | 9.17 GiB | Qwen-Image-2.1 | not recorded | at least 9.28 GiB (stopped after 4.4 s) | at least 8.90 GiB | at least 9.28 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/qwen-image-2.1-1024-tier16.json` | not recorded | +| 24 GB | 14.50 GiB | Qwen-Image-2.1 | not recorded | at least 14.59 GiB (stopped after 4.8 s) | at least 14.59 GiB | at least 14.31 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/qwen-image-2.1-1024-tier24.json` | not recorded | +| 16 GB | 9.17 GiB | Z-Image-Turbo | not recorded | at least 9.26 GiB (stopped after 4.7 s) | at least 9.26 GiB | at least 8.81 GiB | CAPPED | stopped by the watchdog | mlx-dfloat caps for the tier | `bench/results/tiers/aborts/z-image-turbo-1024-tier16.json` | not recorded | The overhead block times one 1024² denoise step, five timed steps after two warm-up steps in each of three rounds, with @@ -919,8 +957,9 @@ with per-block evaluation over mflux's own transformer quantised to 8 bits, with under the scenario's 2.5 GB cache limit, like every other condition here; mflux's own generate sets no cache limit, so this is not quite mflux's q8 as you would run it. A q8 step changes the weights and the output; a DF11 step does not. -On FLUX.1-dev a step costs 3.52 % more with per-block evaluation (19.24 s against the control's 18.59 s) and 4.61 % more -with the depth-2 run-ahead. On FLUX.1-schnell the two figures are 4.05 % (19.28 s against 18.53 s) and 4.24 %. Round 1 +These step numbers are from mflux 0.20.0, and the FLUX.1-dev run skipped the charging gate. On FLUX.1-dev a step costs +3.52 % more with per-block evaluation (19.24 s against the control's 18.59 s) and 4.61 % more with the depth-2 +run-ahead. On FLUX.1-schnell the two figures are 4.05 % (19.28 s against 18.53 s) and 4.24 %. Round 1 of the schnell run saw GPU load from outside the bench; the same pooled medians over rounds 2 and 3 alone give +4.2 % for per-block evaluation. The dev depth-2 figure pools three rounds that disagree. Taken one round at a time (each round's DF11 median over its own control's), they read −1.6 %, +5.9 % and +6.5 %, because round 1's depth-2 control ran @@ -928,7 +967,7 @@ slow at 19.50 s against 18.77 s and 18.69 s later; rounds 2 and 3 alone give +6. help on either model. Evaluating after every block cost 0.41 s per step on schnell and nothing measurable on dev, where the per-block control came out 0.07 s faster, less than either condition moved between rounds. Both sides of the overhead comparison pay that cost, so it is not part of the decode overhead. A DF11 step takes 1.26 times as long as -mflux's q8 step on dev (15.26 s) and 1.35 times on schnell (14.25 s). Treat both ratios as rough: the q8 step's median +mflux's q8 step on dev (15.26 s) and 1.35 times on schnell (14.25 s), both on mflux 0.20.0. Treat both ratios as rough: the q8 step's median moved by about 2 s from one round to the next in both scenarios, while the DF11 step's moved by 0.7 s at most. The q8 step also peaked lower, at 14.8–14.9 GiB against DF11's 19.7–20.0 GiB. What DF11 keeps and q8 gives up is the exact BF16 output. @@ -940,7 +979,7 @@ whose BF16 weights were all resident (`bench/results/validation/control-vs-bf16- FLUX.1-dev, 1024², per-block evaluation: +3.5 % (depth-2: +4.6 %); eval policy cost -0.07 s/step; DF11 (per-block) over the mflux q8 step (one eval per step, same cache limit): 1.26× -Command: `uv run --group bench python -m scripts.bench_flux1 bench/scenarios/flux1-dev-1024.toml` (preflight skipped: not_charging) +Command: `uv run --group bench python -m scripts.bench_flux1 bench/scenarios/flux1-dev-1024.toml --skip-preflight` (preflight skipped: not_charging) FLUX.1-schnell, 1024², per-block evaluation: +4.0 % (depth-2: +4.2 %); eval policy cost 0.41 s/step; DF11 (per-block) over the mflux q8 step (one eval per step, same cache limit): 1.35×