From ab9ff8d323a150a07ef7b12fb6887dca3ec45391 Mon Sep 17 00:00:00 2001 From: Denis Ineshin Date: Thu, 8 Oct 2026 16:03:03 +0200 Subject: [PATCH 1/4] Cap encode() as well as generate_image and let a family floor its denoise term The per-call caps remember the wired limit they or install_memory_caps set, so a call no longer probes it, and a wired cap MLX refuses leaves the memory cap installed with one warning. PhaseConstants gains an optional, size-independent floor for the denoise activation term. --- src/mlx_dfloat/_memory_caps.py | 30 +++++++++ src/mlx_dfloat/bench/capped.py | 12 +++- src/mlx_dfloat/mflux/_phases.py | 11 +++- src/mlx_dfloat/mflux/_pipeline.py | 47 +++++++++++--- src/mlx_dfloat/mflux/flux1/model.py | 1 + src/mlx_dfloat/mflux/flux2/model.py | 1 + src/mlx_dfloat/mflux/qwen21/model.py | 1 + src/mlx_dfloat/mflux/zimage/model.py | 1 + tests/_mlx_limits.py | 4 +- tests/test_bench_capped.py | 11 +++- tests/test_flux1_model.py | 19 ++++++ tests/test_flux2_model.py | 19 ++++++ tests/test_mflux_phases.py | 14 ++++ tests/test_mflux_pipeline.py | 97 ++++++++++++++++++++++++++++ tests/test_qwen21_model.py | 19 ++++++ tests/test_zimage_model.py | 19 ++++++ 16 files changed, 289 insertions(+), 17 deletions(-) diff --git a/src/mlx_dfloat/_memory_caps.py b/src/mlx_dfloat/_memory_caps.py index f7a7d7d..f339957 100644 --- a/src/mlx_dfloat/_memory_caps.py +++ b/src/mlx_dfloat/_memory_caps.py @@ -7,8 +7,17 @@ # The same cap rule as the sibling packages mlx-taef and mlx-quant-fidelity. +import weakref +from typing import Any + import mlx.core as mx +# The wired limit this package last installed or read, per limits object (``mlx.core`` itself, or a test double). +# mlx 0.32.2 has no getter: reading the limit means setting it to 0 and back, which the per-call caps avoid once they +# know it. A wired limit set directly with ``mx.set_wired_limit`` afterwards is not seen; ``forget_wired_limit`` clears +# the record. +_KNOWN_WIRED: "weakref.WeakKeyDictionary[Any, int]" = weakref.WeakKeyDictionary() + DESIRED_WIRED_GB = 20 DESIRED_MEMORY_GB = 22 HEADROOM_GB = 2 @@ -62,6 +71,21 @@ def device_string() -> str | None: return None +def known_wired_limit(limits: Any = mx) -> int | None: + """The wired limit in bytes this package last installed or read on ``limits``, or None when it does not know.""" + return _KNOWN_WIRED.get(limits) + + +def remember_wired_limit(limits: Any, wired_bytes: int) -> None: + """Record the wired limit in bytes now in force on ``limits``.""" + _KNOWN_WIRED[limits] = wired_bytes + + +def forget_wired_limit(limits: Any = mx) -> None: + """Drop the record for ``limits`` (after its wired limit was changed outside this package).""" + _KNOWN_WIRED.pop(limits, None) + + def install_memory_caps() -> tuple[int, int]: """Apply wired + memory caps for the current device. Idempotent; never raises. @@ -74,8 +98,11 @@ def install_memory_caps() -> tuple[int, int]: return (0, 0) try: mx.set_wired_limit(wired_gb * 1024**3) + remember_wired_limit(mx, wired_gb * 1024**3) except Exception: wired_gb = 0 + # MLX's default stays in force: recorded, so a call never probes it. + remember_wired_limit(mx, 0) try: mx.set_memory_limit(memory_gb * 1024**3) except Exception: @@ -87,5 +114,8 @@ def install_memory_caps() -> tuple[int, int]: "caps_for_recommended_bytes", "compute_safe_caps_gb", "device_string", + "forget_wired_limit", "install_memory_caps", + "known_wired_limit", + "remember_wired_limit", ] diff --git a/src/mlx_dfloat/bench/capped.py b/src/mlx_dfloat/bench/capped.py index ac40585..fd686ab 100644 --- a/src/mlx_dfloat/bench/capped.py +++ b/src/mlx_dfloat/bench/capped.py @@ -25,7 +25,7 @@ import mlx.core as mx -from mlx_dfloat._memory_caps import caps_for_recommended_bytes +from mlx_dfloat._memory_caps import caps_for_recommended_bytes, remember_wired_limit from mlx_dfloat.errors import DFloatUnsupportedError from mlx_dfloat.integrate.memory import budget_for @@ -125,12 +125,18 @@ def tier_limits( def apply(limits: TierLimits) -> dict[str, int]: - """Install the tier's memory, cache and wired limits; returns the previous values.""" - return { + """Install the tier's memory, cache and wired limits; returns the previous values. + + The wired limit is recorded (``_memory_caps.remember_wired_limit``), so a model's per-call caps see it without + probing it again. + """ + previous = { "memory": int(mx.set_memory_limit(limits.memory_limit_bytes)), "cache": int(mx.set_cache_limit(limits.cache_limit_bytes)), "wired": int(mx.set_wired_limit(limits.wired_limit_bytes)), } + remember_wired_limit(mx, limits.wired_limit_bytes) + return previous def current_limits() -> dict[str, int]: diff --git a/src/mlx_dfloat/mflux/_phases.py b/src/mlx_dfloat/mflux/_phases.py index 7664b79..ea8c919 100644 --- a/src/mlx_dfloat/mflux/_phases.py +++ b/src/mlx_dfloat/mflux/_phases.py @@ -23,6 +23,8 @@ class PhaseConstants: ``max_measured_pixels`` is the largest image the constants were measured at. ``encode_activation_bytes``, when set, is what a prompt encode measured to hold beyond the encoder's weights, and sizes the encode phase's activations; unset (None), the encode phase carries the denoise allowance instead. + ``denoise_activation_floor_bytes``, when set, is a size-independent lower bound on the denoise activation term (a + buffer a step holds whatever the image size); unset (None), the term is the scaled one alone. """ overhead_bytes: int @@ -34,6 +36,7 @@ class PhaseConstants: allowance_floor: int = 500_000_000 max_measured_pixels: int = 1024 * 1024 encode_activation_bytes: int | None = None + denoise_activation_floor_bytes: int | None = None @dataclass(frozen=True, slots=True, kw_only=True) @@ -117,10 +120,14 @@ def denoise_activation_bytes( ) -> int: """The activation volume a denoise step holds beyond the cache limit, linear in the token count. - Calibrated at ``c.reference_tokens``; every other size is a prediction. + Calibrated at ``c.reference_tokens``; every other size is a prediction. A family's + ``denoise_activation_floor_bytes`` bounds it from below. """ tokens = height * width // 256 + text_tokens - return int(c.denoise_activation_at_reference * tokens / c.reference_tokens) + scaled = int(c.denoise_activation_at_reference * tokens / c.reference_tokens) + if c.denoise_activation_floor_bytes is None: + return scaled + return max(scaled, c.denoise_activation_floor_bytes) def family_phases( diff --git a/src/mlx_dfloat/mflux/_pipeline.py b/src/mlx_dfloat/mflux/_pipeline.py index fe7dd42..2116625 100644 --- a/src/mlx_dfloat/mflux/_pipeline.py +++ b/src/mlx_dfloat/mflux/_pipeline.py @@ -8,13 +8,18 @@ import functools import logging import math +import warnings from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager from typing import Any, NoReturn, ParamSpec, TypeVar import mlx.core as mx -from mlx_dfloat._memory_caps import caps_for_recommended_bytes +from mlx_dfloat._memory_caps import ( + caps_for_recommended_bytes, + known_wired_limit, + remember_wired_limit, +) from mlx_dfloat._watchdog import phys_footprint from mlx_dfloat.errors import DFloatResourceError, DFloatUnsupportedError from mlx_dfloat.integrate.memory import CallPlan @@ -71,19 +76,29 @@ def refuse_construction_args( refuse("bake_lora", "LoRA of any kind") +_wired_refusal_warned = False + + @contextmanager def call_caps(limits: Any = mx) -> Iterator[tuple[int, int]]: """Run the block under the ``mlx-dfloat`` command's memory caps when no wired cap is in force; restore after. The Python API starts at MLX's default wired limit 0 (no cap), while every VAE term was measured under the caps the command installs (``install_memory_caps``: with wired 0 a FLUX.2 Klein VAE decode's MLX peak rose from 6.89 - to 9.49 GiB). A wired limit already set (the command's own caps, or a CAPPED tier's) is left untouched. mlx 0.32.2 - has no getters: the wired limit is read by setting it to 0 and back (``bench.capped.current_limits``), which - leaves nothing behind. ``limits`` is the object holding the setters and the device report (``mlx.core``). - Yields the (wired, memory) caps in bytes installed for the block, (0, 0) when none were. + to 9.49 GiB). A wired limit already set (the command's own caps, a CAPPED tier's, an outer call's) is left + untouched. mlx 0.32.2 has no getters: the wired limit is read by setting it to 0 and back, once per process; after + that the limit this package installed or read is remembered (``_memory_caps.known_wired_limit``), so a call with + the compressed set resident does not touch it. Where MLX refuses the wired cap (the system's wired limit was + lowered), the memory cap is still installed and one warning is issued, as ``install_memory_caps`` goes on without + it. ``limits`` is the object holding the setters and the device report (``mlx.core``). Yields the (wired, memory) + caps in bytes installed for the block, 0 for one that was not. """ - wired_now = int(limits.set_wired_limit(0)) - limits.set_wired_limit(wired_now) + global _wired_refusal_warned + wired_now = known_wired_limit(limits) + if wired_now is None: + wired_now = int(limits.set_wired_limit(0)) + limits.set_wired_limit(wired_now) + remember_wired_limit(limits, wired_now) if wired_now != 0: yield (0, 0) return @@ -97,11 +112,23 @@ def call_caps(limits: Any = mx) -> Iterator[tuple[int, int]]: return previous_memory = int(limits.set_memory_limit(memory)) try: - limits.set_wired_limit(wired) + try: + limits.set_wired_limit(wired) + except Exception as exc: # MLX refuses a wired limit above what the system allows + wired = 0 + if not _wired_refusal_warned: + _wired_refusal_warned = True + warnings.warn( + f"MLX refused the wired memory cap ({exc}); the call runs with the memory cap only", + stacklevel=4, + ) + remember_wired_limit(limits, wired) try: yield (wired, memory) finally: - limits.set_wired_limit(wired_now) + if wired: + limits.set_wired_limit(wired_now) + remember_wired_limit(limits, wired_now) finally: limits.set_memory_limit(previous_memory) @@ -111,7 +138,7 @@ def call_caps(limits: Any = mx) -> Iterator[tuple[int, int]]: def with_call_caps(method: Callable[_P, _R]) -> Callable[_P, _R]: - """Wrap a model's ``generate_image`` in ``call_caps``: the whole call, prelude to VAE decode, runs capped.""" + """Wrap a model's ``generate_image`` or ``encode`` in ``call_caps``: the whole call runs capped.""" @functools.wraps(method) def capped(*args: _P.args, **kwargs: _P.kwargs) -> _R: diff --git a/src/mlx_dfloat/mflux/flux1/model.py b/src/mlx_dfloat/mflux/flux1/model.py index 39063bd..b96762a 100644 --- a/src/mlx_dfloat/mflux/flux1/model.py +++ b/src/mlx_dfloat/mflux/flux1/model.py @@ -420,6 +420,7 @@ def plan_call(self, *, height: int, width: int) -> CallPlan: self._plan = plan return plan + @_pipeline.with_call_caps def encode(self, *prompts: str) -> None: """Encode prompts now, so several generations pay the encoder reload once (drops a resident set first).""" self._lifecycle.ensure_embeddings(*prompts) diff --git a/src/mlx_dfloat/mflux/flux2/model.py b/src/mlx_dfloat/mflux/flux2/model.py index 2ad24ba..bdda41b 100644 --- a/src/mlx_dfloat/mflux/flux2/model.py +++ b/src/mlx_dfloat/mflux/flux2/model.py @@ -337,6 +337,7 @@ def _budget(self) -> int: # --- the public surface ----------------------------------------------------------------------- + @_pipeline.with_call_caps def encode(self, *prompts: str) -> None: """Encode prompts now, so several generations pay the encoder reload once (drops a resident set first).""" self._lifecycle.ensure_embeddings(*dict.fromkeys(prompts)) diff --git a/src/mlx_dfloat/mflux/qwen21/model.py b/src/mlx_dfloat/mflux/qwen21/model.py index 346bf1f..548c81e 100644 --- a/src/mlx_dfloat/mflux/qwen21/model.py +++ b/src/mlx_dfloat/mflux/qwen21/model.py @@ -336,6 +336,7 @@ def _budget(self) -> int: # --- the public surface ----------------------------------------------------------------------- + @_pipeline.with_call_caps def encode(self, *prompts: str) -> None: """Encode prompts now, so several generations pay the encoder reload once (drops a resident set first).""" self._lifecycle.ensure_embeddings(*dict.fromkeys(prompts)) diff --git a/src/mlx_dfloat/mflux/zimage/model.py b/src/mlx_dfloat/mflux/zimage/model.py index 228cd13..b0820d4 100644 --- a/src/mlx_dfloat/mflux/zimage/model.py +++ b/src/mlx_dfloat/mflux/zimage/model.py @@ -427,6 +427,7 @@ def plan_call(self, *, height: int, width: int) -> CallPlan: self._plan = plan return plan + @_pipeline.with_call_caps def encode(self, *prompts: str) -> None: """Encode prompts now, so several generations pay the encoder reload once (drops a resident set first).""" self._lifecycle.ensure_embeddings(*prompts) diff --git a/tests/_mlx_limits.py b/tests/_mlx_limits.py index a3daf3e..0bc4b21 100644 --- a/tests/_mlx_limits.py +++ b/tests/_mlx_limits.py @@ -5,7 +5,7 @@ import mlx.core as mx -from mlx_dfloat._memory_caps import compute_safe_caps_gb +from mlx_dfloat._memory_caps import compute_safe_caps_gb, forget_wired_limit, remember_wired_limit from mlx_dfloat.bench.capped import current_limits GIB = 1024**3 @@ -28,8 +28,10 @@ def mlx_without_wired_cap() -> Iterator[dict[str, int]]: """ before = current_limits() mx.set_wired_limit(0) + forget_wired_limit(mx) # a fresh process knows nothing of its wired limit try: yield {**before, "wired": 0} finally: mx.set_memory_limit(before["memory"]) mx.set_wired_limit(before["wired"]) + remember_wired_limit(mx, before["wired"]) diff --git a/tests/test_bench_capped.py b/tests/test_bench_capped.py index 6168979..bb19caf 100644 --- a/tests/test_bench_capped.py +++ b/tests/test_bench_capped.py @@ -167,8 +167,17 @@ def test_apply_installs_the_three_limits_and_returns_the_previous_ones(monkeypat monkeypatch.setattr( capped.mx, "set_wired_limit", lambda b: (calls.__setitem__("wired", b), 33)[1] ) + from mlx_dfloat import _memory_caps + + known = _memory_caps.known_wired_limit(capped.mx) lim = tier_limits(16, host_ram_bytes=HOST_RAM, host_recommended_bytes=HOST_REC) - previous = apply(lim) + try: + previous = apply(lim) + # The per-call caps read the wired limit from this record instead of probing it (P2): a tier's 8 GiB. + assert _memory_caps.known_wired_limit(capped.mx) == 8_589_934_592 + finally: + if known is not None: + _memory_caps.remember_wired_limit(capped.mx, known) # The 16 GB tier's caps, worked by hand in test_a_16_gb_tier_gets_the_caps_...: wired 8 GiB, # memory 10 GiB, cache 10 GiB. assert calls == { diff --git a/tests/test_flux1_model.py b/tests/test_flux1_model.py index fc549f1..096346f 100644 --- a/tests/test_flux1_model.py +++ b/tests/test_flux1_model.py @@ -686,3 +686,22 @@ def test_a_python_api_call_runs_under_the_commands_caps_and_restores_mlxs_defaul after = current_limits() assert (seen["limits_at_entry"]["wired"], seen["limits_at_entry"]["memory"]) == (wired, memory) assert after == start + + +def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): + # Bug caught (X4): encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # phases, sized under the command's caps), or the caps left installed after it returns. + from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap + + model = _fake_model(tmp_path, monkeypatch) + seen = [] + monkeypatch.setattr( + model._lifecycle, "ensure_embeddings", lambda *prompts: seen.append(current_limits()) + ) + wired, memory = command_caps() + assert wired > 0 + with mlx_without_wired_cap() as start: + model.encode("p") + after = current_limits() + assert (seen[0]["wired"], seen[0]["memory"]) == (wired, memory) + assert after == start diff --git a/tests/test_flux2_model.py b/tests/test_flux2_model.py index 7b61512..a02ae91 100644 --- a/tests/test_flux2_model.py +++ b/tests/test_flux2_model.py @@ -871,3 +871,22 @@ def test_a_python_api_call_runs_under_the_commands_caps_and_restores_mlxs_defaul after = current_limits() assert (seen["limits_at_entry"]["wired"], seen["limits_at_entry"]["memory"]) == (wired, memory) assert after == start + + +def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): + # Bug caught (X4): encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # phases, sized under the command's caps), or the caps left installed after it returns. + from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap + + model = fake_model(tmp_path, monkeypatch) + seen = [] + monkeypatch.setattr( + model._lifecycle, "ensure_embeddings", lambda *prompts: seen.append(current_limits()) + ) + wired, memory = command_caps() + assert wired > 0 + with mlx_without_wired_cap() as start: + model.encode("p") + after = current_limits() + assert (seen[0]["wired"], seen[0]["memory"]) == (wired, memory) + assert after == start diff --git a/tests/test_mflux_phases.py b/tests/test_mflux_phases.py index 96691ce..54c1e89 100644 --- a/tests/test_mflux_phases.py +++ b/tests/test_mflux_phases.py @@ -70,6 +70,20 @@ def test_the_denoise_activation_scales_with_tokens_from_the_reference(): assert denoise_activation_bytes(c, height=512, width=512, text_tokens=512) == 333_333 +def test_a_denoise_activation_floor_holds_below_the_reference_and_is_off_by_default(): + # Bug caught (X3): a family whose step holds a size-independent buffer (ERNIE-Image's float32 copies of a block's + # weights) predicted below it at a small image, or the floor applied to every family (the default must leave the + # scaled term as it was). 512^2 + 512: 333_333 scaled, floored at 400_000; 1024^2 + 512: 1_000_000 above the floor. + floored = _constants( + denoise_activation_at_reference=1_000_000, denoise_activation_floor_bytes=400_000 + ) + assert denoise_activation_bytes(floored, height=512, width=512, text_tokens=512) == 400_000 + assert denoise_activation_bytes(floored, height=1024, width=1024, text_tokens=512) == 1_000_000 + assert _constants().denoise_activation_floor_bytes is None + plain = _constants(denoise_activation_at_reference=1_000_000) + assert denoise_activation_bytes(plain, height=512, width=512, text_tokens=512) == 333_333 + + SIZES = FamilySizes(compressed=5_000, extras=10, nonblock=300, encoders=8_000, vae=160) SMALL_LARGEST = {"a": 400, "b": 200} diff --git a/tests/test_mflux_pipeline.py b/tests/test_mflux_pipeline.py index 1c6e252..965b1be 100644 --- a/tests/test_mflux_pipeline.py +++ b/tests/test_mflux_pipeline.py @@ -273,6 +273,103 @@ def test_the_command_path_keeps_its_caps_through_a_call(): assert current_limits() == before +class RecordingLimits(FakeLimits): + """FakeLimits that records every wired-limit call; with ``refuse_wired`` it raises as MLX does for a wired limit + above what the system allows (``[metal::set_wired_limit] ... is not allowed``, a ValueError).""" + + def __init__(self, *, refuse_wired=False, **kwargs): + super().__init__(**kwargs) + self.wired_calls = [] + self.refuse_wired = refuse_wired + + def set_wired_limit(self, n): + self.wired_calls.append(n) + if self.refuse_wired and n > 0: + raise ValueError( + "[metal::set_wired_limit] Setting a wired limit larger than the maximum is not allowed." + ) + return super().set_wired_limit(n) + + +def test_a_wired_cap_mlx_refuses_still_installs_the_memory_cap_and_warns_once(monkeypatch): + # Bug caught (P1): a Python generate_image raising MLX's raw ValueError when the system's wired limit was lowered + # (sysctl iogpu.wired_limit_mb), where install_memory_caps keeps the memory cap and goes on; or a warning on every + # call. The memory cap (22 GiB here) is installed for the block and MLX's default restored after it. + monkeypatch.setattr(_pipeline, "_wired_refusal_warned", False) + limits = RecordingLimits(refuse_wired=True, wired=0, memory=DEFAULT_MEMORY) + import warnings + + inside = [] + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with _pipeline.call_caps(limits) as installed: + inside.append((installed, limits.wired, limits.memory)) + assert [str(w.message) for w in caught if "wired" in str(w.message)] != [] + assert inside == [((0, 23_622_320_128), 0, 23_622_320_128)] + assert (limits.wired, limits.memory) == (0, DEFAULT_MEMORY) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + with _pipeline.call_caps(limits) as again: + assert again == (0, 23_622_320_128) + + +def test_a_second_call_does_not_probe_the_wired_limit(): + # Bug caught (P2): the probe (set the wired limit to 0, read the previous value, set it back) run at every call, + # with a compressed set resident, when this process already knows the limit it left in force. First call: probe + # (0, then the 0 it read), install 20 GiB, restore 0; the second: install and restore only. + limits = RecordingLimits(wired=0, memory=DEFAULT_MEMORY) + with _pipeline.call_caps(limits): + pass + with _pipeline.call_caps(limits): + pass + assert limits.wired_calls == [0, 0, 21_474_836_480, 0, 21_474_836_480, 0] + + +def test_the_command_path_never_probes(monkeypatch): + # Bug caught (P2): a call probing the wired limit after the command installed its caps (install_memory_caps + # records what it installed): the only wired-limit call is the command's own 20 GiB. + from mlx_dfloat import _memory_caps + + limits = RecordingLimits(wired=0, memory=DEFAULT_MEMORY) + monkeypatch.setattr(_memory_caps, "mx", limits) + assert _memory_caps.install_memory_caps() == (20, 22) + with _pipeline.call_caps(limits) as installed: + assert installed == (0, 0) + assert limits.wired_calls == [21_474_836_480] + + +def test_a_wired_cap_the_command_could_not_install_is_recorded_so_calls_never_probe(monkeypatch): + # Bug caught (R2): install_memory_caps leaving the wired limit unknown when MLX refuses its cap, so the first call + # probes it after all (set to 0 and back). Recorded as 0: the call tries the cap once (refused again, one warning) + # and never writes 0. The only wired-limit calls are the two refused 20 GiB attempts. + import warnings + + from mlx_dfloat import _memory_caps + + monkeypatch.setattr(_pipeline, "_wired_refusal_warned", False) + limits = RecordingLimits(refuse_wired=True, wired=0, memory=DEFAULT_MEMORY) + monkeypatch.setattr(_memory_caps, "mx", limits) + assert _memory_caps.install_memory_caps() == (0, 22) + assert _memory_caps.known_wired_limit(limits) == 0 + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + with _pipeline.call_caps(limits) as installed: + assert installed == (0, 23_622_320_128) + assert limits.wired_calls == [21_474_836_480, 21_474_836_480] + + +def test_a_nested_call_leaves_the_outer_calls_caps_alone(): + # Bug caught: an encode() run inside generate_image (both capped) restoring MLX's defaults when it returns, so + # the rest of the generation runs uncapped. + limits = RecordingLimits(wired=0, memory=DEFAULT_MEMORY) + with _pipeline.call_caps(limits): + with _pipeline.call_caps(limits) as inner: + assert inner == (0, 0) + assert (limits.wired, limits.memory) == (21_474_836_480, 23_622_320_128) + assert (limits.wired, limits.memory) == (0, DEFAULT_MEMORY) + + # --- the phase tracker -------------------------------------------------------------------------------------------- diff --git a/tests/test_qwen21_model.py b/tests/test_qwen21_model.py index e3f99a1..1aa2220 100644 --- a/tests/test_qwen21_model.py +++ b/tests/test_qwen21_model.py @@ -1037,3 +1037,22 @@ def test_a_tiny_generation_gives_stock_mfluxs_latents_bit_for_bit( assert ours.shape == theirs.shape view = mx.uint32 if ours.dtype == mx.float32 else mx.uint16 assert np.array_equal(np.array(ours.view(view)), np.array(theirs.view(view))) + + +def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): + # Bug caught (X4): encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # phases, sized under the command's caps), or the caps left installed after it returns. + from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap + + model = fake_model(tmp_path, monkeypatch) + seen = [] + monkeypatch.setattr( + model._lifecycle, "ensure_embeddings", lambda *prompts: seen.append(current_limits()) + ) + wired, memory = command_caps() + assert wired > 0 + with mlx_without_wired_cap() as start: + model.encode("p") + after = current_limits() + assert (seen[0]["wired"], seen[0]["memory"]) == (wired, memory) + assert after == start diff --git a/tests/test_zimage_model.py b/tests/test_zimage_model.py index dc62f3a..2c1c7a7 100644 --- a/tests/test_zimage_model.py +++ b/tests/test_zimage_model.py @@ -757,3 +757,22 @@ def test_a_python_api_call_runs_under_the_commands_caps_and_restores_mlxs_defaul after = current_limits() assert (seen["limits_at_entry"]["wired"], seen["limits_at_entry"]["memory"]) == (wired, memory) assert after == start + + +def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): + # Bug caught (X4): encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # phases, sized under the command's caps), or the caps left installed after it returns. + from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap + + model = fake_model(tmp_path, monkeypatch) + seen = [] + monkeypatch.setattr( + model._lifecycle, "ensure_embeddings", lambda *prompts: seen.append(current_limits()) + ) + wired, memory = command_caps() + assert wired > 0 + with mlx_without_wired_cap() as start: + model.encode("p") + after = current_limits() + assert (seen[0]["wired"], seen[0]["memory"]) == (wired, memory) + assert after == start From 6d20002191df4c496a41afc2b9f4e361c5f0a22c Mon Sep 17 00:00:00 2001 From: Denis Ineshin Date: Thu, 8 Oct 2026 16:03:03 +0200 Subject: [PATCH 2/4] Run ERNIE-Image and ERNIE-Image-Turbo from their DFloat11 checkpoints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The adapter, the registry and command entries, the memory constants measured at 1024², and the committed records they and the tests read. --- .../results/calibration/ernie-image-1024.json | 70 ++ .../calibration/ernie-image-turbo-1024.json | 97 ++ .../ernie-image-1024/bf16-result.json | 57 ++ .../identity/ernie-image-1024/compare.json | 34 + .../ernie-image-1024/df11-result.json | 141 +++ .../ernie-image-turbo/full-vs-bf16.json | 112 +++ .../kernel-vs-reference.json | 67 ++ .../parity/ernie-image/full-vs-bf16.json | 67 ++ .../ernie-image/kernel-vs-reference.json | 67 ++ .../refusals/ernie-image-1024-tier16.json | 49 + .../refusals/ernie-image-1024-tier24.json | 49 + .../ernie-image-turbo-1024-tier16.json | 49 + .../ernie-image-turbo-1024-tier24.json | 49 + .../tiers/aborts/ernie-image-1024-tier16.json | 30 + .../tiers/aborts/ernie-image-1024-tier24.json | 30 + .../aborts/ernie-image-turbo-1024-tier16.json | 30 + bench/results/tiers/ernie-image-1024.json | 133 +++ .../tiers/ernie-image-turbo-1024-tier24.json | 131 +++ .../results/tiers/ernie-image-turbo-1024.json | 133 +++ src/mlx_dfloat/mflux/__init__.py | 12 +- src/mlx_dfloat/mflux/ernie/__init__.py | 1 + src/mlx_dfloat/mflux/ernie/init.py | 176 ++++ src/mlx_dfloat/mflux/ernie/memory.py | 184 ++++ src/mlx_dfloat/mflux/ernie/model.py | 664 +++++++++++++ src/mlx_dfloat/mflux/ernie/names.py | 89 ++ src/mlx_dfloat/mflux/ernie/transformer.py | 156 ++++ src/mlx_dfloat/mflux/families.py | 57 +- src/mlx_dfloat/mflux/generate.py | 25 +- tests/_ernie_tiny.py | 305 ++++++ tests/test_cli.py | 106 +++ tests/test_ernie_build.py | 197 ++++ tests/test_ernie_init.py | 339 +++++++ tests/test_ernie_memory.py | 566 +++++++++++ tests/test_ernie_model.py | 884 ++++++++++++++++++ tests/test_ernie_names.py | 276 ++++++ tests/test_ernie_tiny.py | 55 ++ tests/test_ernie_transformer.py | 320 +++++++ tests/test_mflux_families.py | 84 +- tests/test_public_api.py | 10 + 39 files changed, 5881 insertions(+), 20 deletions(-) create mode 100644 bench/results/calibration/ernie-image-1024.json create mode 100644 bench/results/calibration/ernie-image-turbo-1024.json create mode 100644 bench/results/identity/ernie-image-1024/bf16-result.json create mode 100644 bench/results/identity/ernie-image-1024/compare.json create mode 100644 bench/results/identity/ernie-image-1024/df11-result.json create mode 100644 bench/results/parity/ernie-image-turbo/full-vs-bf16.json create mode 100644 bench/results/parity/ernie-image-turbo/kernel-vs-reference.json create mode 100644 bench/results/parity/ernie-image/full-vs-bf16.json create mode 100644 bench/results/parity/ernie-image/kernel-vs-reference.json create mode 100644 bench/results/refusals/ernie-image-1024-tier16.json create mode 100644 bench/results/refusals/ernie-image-1024-tier24.json create mode 100644 bench/results/refusals/ernie-image-turbo-1024-tier16.json create mode 100644 bench/results/refusals/ernie-image-turbo-1024-tier24.json create mode 100644 bench/results/tiers/aborts/ernie-image-1024-tier16.json create mode 100644 bench/results/tiers/aborts/ernie-image-1024-tier24.json create mode 100644 bench/results/tiers/aborts/ernie-image-turbo-1024-tier16.json create mode 100644 bench/results/tiers/ernie-image-1024.json create mode 100644 bench/results/tiers/ernie-image-turbo-1024-tier24.json create mode 100644 bench/results/tiers/ernie-image-turbo-1024.json create mode 100644 src/mlx_dfloat/mflux/ernie/__init__.py create mode 100644 src/mlx_dfloat/mflux/ernie/init.py create mode 100644 src/mlx_dfloat/mflux/ernie/memory.py create mode 100644 src/mlx_dfloat/mflux/ernie/model.py create mode 100644 src/mlx_dfloat/mflux/ernie/names.py create mode 100644 src/mlx_dfloat/mflux/ernie/transformer.py create mode 100644 tests/_ernie_tiny.py create mode 100644 tests/test_ernie_build.py create mode 100644 tests/test_ernie_init.py create mode 100644 tests/test_ernie_memory.py create mode 100644 tests/test_ernie_model.py create mode 100644 tests/test_ernie_names.py create mode 100644 tests/test_ernie_tiny.py create mode 100644 tests/test_ernie_transformer.py diff --git a/bench/results/calibration/ernie-image-1024.json b/bench/results/calibration/ernie-image-1024.json new file mode 100644 index 0000000..159608a --- /dev/null +++ b/bench/results/calibration/ernie-image-1024.json @@ -0,0 +1,70 @@ +{ + "model": "ernie-image", + "purpose": "the 1024² run the ERNIE-Image memory constants were calibrated on (mlx_dfloat.mflux.ernie.memory); one record per CFG batch size", + "df11": "mingyi456/ERNIE-Image-DF11@c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "base": "baidu/ERNIE-Image@5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "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": 26, + "batch": 2, + "device": "Apple M1 Max, 32 GB", + "macos": "27.0.1", + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "git": "a771e58b2d12142d484a6564b12484a523dc6595", + "footprint_peaks_label": "OS phys_footprint in bytes, the peak of each phase; for the MEASURED and identity runs the run's watched peak, reached in the VAE phase", + "sizes_label": "FamilySizes as mlx_dfloat.mflux.ernie.memory.sizes_for computes them: encoders is the language model's evaluated layers (its last layer never runs), not the 6_858_012_672 B of every language-model tensor", + "constants_from": { + "denoise": "derisk-base-1" + }, + "runs": [ + { + "label": "derisk-base-1", + "date": "2026-10-08", + "run": "the de-risk run: one step at guidance 4.0 without a negative prompt (mflux's \" \" then the prompt, one batch-2 transformer call), at the planner's limit for the probe's allowance, the VAE decode on the resident set", + "steps": 1, + "cache_limit": 3600000000, + "footprint_peaks": { + "encode": 7257593880, + "denoise": 16361874888, + "vae": 19127199104 + }, + "after_build": { + "footprint": 576161232, + "mlx_active": 27961608, + "mlx_cache": 4806 + }, + "mlx_encode_held": 6655149770, + "sizes": { + "compressed": 10894750607, + "extras": 27961600, + "nonblock": 335544320, + "encoders": 6625216512, + "vae": 168120878 + } + }, + { + "label": "measured-50-steps", + "date": "2026-10-08", + "run": "the MEASURED row (mlx-dfloat generate --tier 32): 50 steps at guidance 4.0 (one batch-2 call per step), the VAE decode on the resident set; bench/results/tiers/ernie-image-1024.json", + "steps": 50, + "cache_limit": 3600000000, + "footprint_peaks": { + "vae": 18032961064 + } + }, + { + "label": "identity-4-steps-cfg", + "date": "2026-10-08", + "run": "the identity check's df11 side: 4 steps at guidance 4.0 (one batch-2 call per step), the VAE decode on the resident set; bench/results/identity/ernie-image-1024/df11-result.json", + "steps": 4, + "cache_limit": 3600000000, + "footprint_peaks": { + "vae": 18537964904 + } + } + ] +} diff --git a/bench/results/calibration/ernie-image-turbo-1024.json b/bench/results/calibration/ernie-image-turbo-1024.json new file mode 100644 index 0000000..dfe0639 --- /dev/null +++ b/bench/results/calibration/ernie-image-turbo-1024.json @@ -0,0 +1,97 @@ +{ + "model": "ernie-image-turbo", + "purpose": "the 1024² run the ERNIE-Image memory constants were calibrated on (mlx_dfloat.mflux.ernie.memory); one record per CFG batch size", + "df11": "mingyi456/ERNIE-Image-Turbo-DF11@27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "base": "baidu/ERNIE-Image-Turbo@bc68c81e2a1730a394d5fc9fae70713dee940140", + "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": 26, + "batch": 1, + "device": "Apple M1 Max, 32 GB", + "macos": "27.0.1", + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "git": "a771e58b2d12142d484a6564b12484a523dc6595", + "footprint_peaks_label": "OS phys_footprint in bytes, the peak of each phase; for the MEASURED, identity and CAPPED runs the run's watched peak, reached in the phase named", + "sizes_label": "FamilySizes as mlx_dfloat.mflux.ernie.memory.sizes_for computes them: encoders is the language model's evaluated layers (its last layer never runs), not the 6_858_012_672 B of every language-model tensor", + "overhead_from": "derisk-turbo-1", + "constants_from": { + "encode": "derisk-turbo-1", + "denoise": "capped-24-gb-2", + "vae": "derisk-turbo-1" + }, + "runs": [ + { + "label": "derisk-turbo-1", + "date": "2026-10-08", + "run": "the de-risk run: one step at guidance 1.0 (one prompt, batch 1), at the planner's limit for the probe's allowance, the VAE decode on the resident set", + "steps": 1, + "cache_limit": 2600000000, + "footprint_peaks": { + "encode": 7335614536, + "denoise": 14974035520, + "vae": 19163080232 + }, + "after_build": { + "footprint": 644400592, + "mlx_active": 27961608, + "mlx_cache": 4806 + }, + "mlx_encode_held": 6629405444, + "sizes": { + "compressed": 10894810183, + "extras": 27961600, + "nonblock": 335544320, + "encoders": 6625216512, + "vae": 168120878 + } + }, + { + "label": "measured-8-steps", + "date": "2026-10-08", + "run": "the MEASURED row (mlx-dfloat generate --tier 32): 8 steps at guidance 1.0, the VAE decode on the resident set; bench/results/tiers/ernie-image-turbo-1024.json", + "steps": 8, + "cache_limit": 2600000000, + "footprint_peaks": { + "vae": 18621424880 + } + }, + { + "label": "capped-24-gb-1", + "date": "2026-10-08", + "run": "the first CAPPED 24 GB run (mlx-dfloat generate --tier 24, git 5e2e7c2): 8 steps at guidance 1.0 under a 24 GB Mac's caps, the set dropped before the VAE decode; its watched peak, reached in the denoise phase", + "steps": 8, + "cache_limit": 2600000000, + "footprint_peaks": { + "denoise": 15069554288 + }, + "sizes": { + "compressed": 10894810183, + "extras": 27961600, + "nonblock": 335544320, + "encoders": 6625216512, + "vae": 168120878 + } + }, + { + "label": "capped-24-gb-2", + "date": "2026-10-08", + "run": "the second CAPPED 24 GB run (mlx-dfloat generate --tier 24 --no-fit-check, git f4cd325): 8 steps at guidance 1.0 under a 24 GB Mac's caps, the set dropped before the VAE decode; its watched peak, reached in the denoise phase", + "steps": 8, + "cache_limit": 2600000000, + "footprint_peaks": { + "denoise": 15446156600 + }, + "sizes": { + "compressed": 10894810183, + "extras": 27961600, + "nonblock": 335544320, + "encoders": 6625216512, + "vae": 168120878 + } + } + ] +} diff --git a/bench/results/identity/ernie-image-1024/bf16-result.json b/bench/results/identity/ernie-image-1024/bf16-result.json new file mode 100644 index 0000000..a3cf965 --- /dev/null +++ b/bench/results/identity/ernie-image-1024/bf16-result.json @@ -0,0 +1,57 @@ +{ + "exit_code": 0, + "reads": 1008, + "cfg_batch": 2, + "output_shape": [ + 1, + 128, + 64, + 64 + ], + "output_dtype": "mlx.core.bfloat16", + "footprint_peak_bytes": 7442367832, + "mlx_peak_bytes_process": 7766118326, + "key": { + "model": "ernie-image", + "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": 4.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-DF11/snapshots/c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "base": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image/snapshots/5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "source": "a6a225e0a58ab9db6fbb3040a74fb13befe7b7681804c8b1831ac37d1d7311f7", + "mlx": "0.32.2" + }, + "memory_caps_gb": [ + 20, + 22 + ], + "provenance": { + "git": "5e2e7c2dfb0559a567f644b1e57eba524a39d11a", + "source_hash": "a6a225e0a58ab9db6fbb3040a74fb13befe7b7681804c8b1831ac37d1d7311f7", + "mlx": "0.32.2", + "mflux": "0.20.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": 99 + }, + "date": "2026-10-08" + } +} diff --git a/bench/results/identity/ernie-image-1024/compare.json b/bench/results/identity/ernie-image-1024/compare.json new file mode 100644 index 0000000..bd010b8 --- /dev/null +++ b/bench/results/identity/ernie-image-1024/compare.json @@ -0,0 +1,34 @@ +{ + "exit_code": 0, + "equal": true, + "degenerate": [], + "df11_key": { + "model": "ernie-image", + "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": 4.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-DF11/snapshots/c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "base": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image/snapshots/5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "source": "a6a225e0a58ab9db6fbb3040a74fb13befe7b7681804c8b1831ac37d1d7311f7", + "mlx": "0.32.2" + }, + "bf16_key": { + "model": "ernie-image", + "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": 4.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-DF11/snapshots/c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "base": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image/snapshots/5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "source": "a6a225e0a58ab9db6fbb3040a74fb13befe7b7681804c8b1831ac37d1d7311f7", + "mlx": "0.32.2" + }, + "pixels_identical": true, + "df11_png_sha256": "bcca459a7625242e0763f249e048d7d15d9835694bf343789e60858b23ddd2ad", + "bf16_png_sha256": "931e5d736872e474024b874b66cee9b3055a1d30b8849350ecb03f3615d61de3" +} diff --git a/bench/results/identity/ernie-image-1024/df11-result.json b/bench/results/identity/ernie-image-1024/df11-result.json new file mode 100644 index 0000000..94c9b90 --- /dev/null +++ b/bench/results/identity/ernie-image-1024/df11-result.json @@ -0,0 +1,141 @@ +{ + "exit_code": 0, + "degenerate": [], + "report": { + "family": "ernie", + "model": "ernie-image", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-DF11/snapshots/c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "repo_id": null, + "revision": "c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "config_source": "config.json" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image/snapshots/5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "repo_id": null, + "revision": "5346b31d68c9c23758ba56ef8be5e9dc174c7f99" + }, + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_batch": 2, + "text_tokens": 26, + "predict": "uncompiled", + "nonblock_groups": [ + "adaLN_modulation.1", + "final_norm.linear", + "time_embedding" + ], + "cache_limit_in_force": 3600000000, + "drop_set_before_vae": false, + "fit": { + "label": "predicted", + "phases": { + "encode": 7271583948, + "denoise": 16361874888, + "vae": 19163020656 + }, + "peak_phase": "vae", + "peak_bytes": 19163020656, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 10894750607, + "extras": 27961600, + "nonblock": 335544320, + "encoders": 6625216512, + "vae": 168120878 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 642893288, + "mlx_peak": 6682854018, + "footprint_end": 6870603800, + "active_at_start": 27961748 + }, + "set_load": { + "footprint_start": 6845437976, + "mlx_peak": 11259611992, + "footprint_end": 11991377352, + "active_at_start": 28289424 + }, + "denoise": { + "footprint_start": 11991377352, + "mlx_peak": 14181491900, + "footprint_end": 15731778896, + "active_at_start": 11026621584 + }, + "vae": { + "footprint_start": 15731746128, + "mlx_peak": 17355161398, + "footprint_end": 12703294848, + "active_at_start": 11032978612 + } + }, + "decode_launches": 144, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 23.25415904104011, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + } + }, + "output_shape": [ + 1, + 128, + 64, + 64 + ], + "output_dtype": "mlx.core.bfloat16", + "footprint_peak_bytes": 18537964904, + "mlx_peak_bytes_max_over_phases": 17355161398, + "key": { + "model": "ernie-image", + "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": 4.0, + "eval_policy": "per-block", + "df11": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-DF11/snapshots/c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "base": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image/snapshots/5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "source": "a6a225e0a58ab9db6fbb3040a74fb13befe7b7681804c8b1831ac37d1d7311f7", + "mlx": "0.32.2" + }, + "memory_caps_gb": [ + 20, + 22 + ], + "provenance": { + "git": "5e2e7c2dfb0559a567f644b1e57eba524a39d11a", + "source_hash": "a6a225e0a58ab9db6fbb3040a74fb13befe7b7681804c8b1831ac37d1d7311f7", + "mlx": "0.32.2", + "mflux": "0.20.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-08" + } +} diff --git a/bench/results/parity/ernie-image-turbo/full-vs-bf16.json b/bench/results/parity/ernie-image-turbo/full-vs-bf16.json new file mode 100644 index 0000000..1880cf9 --- /dev/null +++ b/bench/results/parity/ernie-image-turbo/full-vs-bf16.json @@ -0,0 +1,112 @@ +{ + "df11_repo": "mingyi456/ERNIE-Image-Turbo-DF11", + "df11_revision": "27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "bf16_repo": "baidu/ERNIE-Image-Turbo", + "bf16_revision": "bc68c81e2a1730a394d5fc9fae70713dee940140", + "mode": "parity, by HTTP range reads, 13 processes of 3 groups", + "config_source": [ + "config.json" + ], + "control": [ + "bf16" + ], + "exit_codes": { + "u00": 0, + "u01": 0, + "u02": 0, + "u03": 0, + "u04": 0, + "u05": 0, + "u06": 0, + "u07": 0, + "u08": 0, + "u09": 0, + "u10": 0, + "u11": 0, + "u12": 0 + }, + "groups": [ + "adaLN_modulation.1", + "final_norm.linear", + "layers.0", + "layers.1", + "layers.10", + "layers.11", + "layers.12", + "layers.13", + "layers.14", + "layers.15", + "layers.16", + "layers.17", + "layers.18", + "layers.19", + "layers.2", + "layers.20", + "layers.21", + "layers.22", + "layers.23", + "layers.24", + "layers.25", + "layers.26", + "layers.27", + "layers.28", + "layers.29", + "layers.3", + "layers.30", + "layers.31", + "layers.32", + "layers.33", + "layers.34", + "layers.35", + "layers.4", + "layers.5", + "layers.6", + "layers.7", + "layers.8", + "layers.9", + "time_embedding" + ], + "compared": 256, + "mismatched": 0, + "per_group_status": { + "time_embedding": "equal", + "adaLN_modulation.1": "equal", + "final_norm.linear": "equal", + "layers.0": "equal", + "layers.1": "equal", + "layers.2": "equal", + "layers.3": "equal", + "layers.4": "equal", + "layers.5": "equal", + "layers.6": "equal", + "layers.7": "equal", + "layers.8": "equal", + "layers.9": "equal", + "layers.10": "equal", + "layers.11": "equal", + "layers.12": "equal", + "layers.13": "equal", + "layers.14": "equal", + "layers.15": "equal", + "layers.16": "equal", + "layers.17": "equal", + "layers.18": "equal", + "layers.19": "equal", + "layers.20": "equal", + "layers.21": "equal", + "layers.22": "equal", + "layers.23": "equal", + "layers.24": "equal", + "layers.25": "equal", + "layers.26": "equal", + "layers.27": "equal", + "layers.28": "equal", + "layers.29": "equal", + "layers.30": "equal", + "layers.31": "equal", + "layers.32": "equal", + "layers.33": "equal", + "layers.34": "equal", + "layers.35": "equal" + } +} diff --git a/bench/results/parity/ernie-image-turbo/kernel-vs-reference.json b/bench/results/parity/ernie-image-turbo/kernel-vs-reference.json new file mode 100644 index 0000000..c974788 --- /dev/null +++ b/bench/results/parity/ernie-image-turbo/kernel-vs-reference.json @@ -0,0 +1,67 @@ +{ + "key": { + "df11_revision": "27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "bf16_revision": "none", + "mode": "kernel-vs-reference", + "decoder": "metal", + "source": "32ff46e183df401598f7756946d4d72c51c5bbe1f3d921f49eb369fffde7a547", + "git": "76932203e838174006341008f6dd9f94cb6bb4f0", + "mlx": "0.32.2" + }, + "exit_code": 0, + "groups": 39, + "partial": false, + "selected_groups": [ + "adaLN_modulation.1", + "final_norm.linear", + "layers.0", + "layers.1", + "layers.2", + "layers.3", + "layers.4", + "layers.5", + "layers.6", + "layers.7", + "layers.8", + "layers.9", + "layers.10", + "layers.11", + "layers.12", + "layers.13", + "layers.14", + "layers.15", + "layers.16", + "layers.17", + "layers.18", + "layers.19", + "layers.20", + "layers.21", + "layers.22", + "layers.23", + "layers.24", + "layers.25", + "layers.26", + "layers.27", + "layers.28", + "layers.29", + "layers.30", + "layers.31", + "layers.32", + "layers.33", + "layers.34", + "layers.35", + "time_embedding" + ], + "compared": 256, + "expected": 256, + "extras_compared": 0, + "mismatched": 0, + "missing_originals": [], + "uncovered_originals": [], + "ignored_originals": [], + "memory_caps_gb": [ + 20, + 22 + ], + "seconds": 486.8804898750386 +} diff --git a/bench/results/parity/ernie-image/full-vs-bf16.json b/bench/results/parity/ernie-image/full-vs-bf16.json new file mode 100644 index 0000000..6f3dff6 --- /dev/null +++ b/bench/results/parity/ernie-image/full-vs-bf16.json @@ -0,0 +1,67 @@ +{ + "key": { + "df11_revision": "c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "bf16_revision": "5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "mode": "parity", + "decoder": "metal", + "source": "32ff46e183df401598f7756946d4d72c51c5bbe1f3d921f49eb369fffde7a547", + "git": "76932203e838174006341008f6dd9f94cb6bb4f0", + "mlx": "0.32.2" + }, + "exit_code": 0, + "groups": 39, + "partial": false, + "selected_groups": [ + "adaLN_modulation.1", + "final_norm.linear", + "layers.0", + "layers.1", + "layers.2", + "layers.3", + "layers.4", + "layers.5", + "layers.6", + "layers.7", + "layers.8", + "layers.9", + "layers.10", + "layers.11", + "layers.12", + "layers.13", + "layers.14", + "layers.15", + "layers.16", + "layers.17", + "layers.18", + "layers.19", + "layers.20", + "layers.21", + "layers.22", + "layers.23", + "layers.24", + "layers.25", + "layers.26", + "layers.27", + "layers.28", + "layers.29", + "layers.30", + "layers.31", + "layers.32", + "layers.33", + "layers.34", + "layers.35", + "time_embedding" + ], + "compared": 256, + "expected": 256, + "extras_compared": 153, + "mismatched": 0, + "missing_originals": [], + "uncovered_originals": [], + "ignored_originals": [], + "memory_caps_gb": [ + 20, + 22 + ], + "seconds": 25.481306125002448 +} diff --git a/bench/results/parity/ernie-image/kernel-vs-reference.json b/bench/results/parity/ernie-image/kernel-vs-reference.json new file mode 100644 index 0000000..007b1d4 --- /dev/null +++ b/bench/results/parity/ernie-image/kernel-vs-reference.json @@ -0,0 +1,67 @@ +{ + "key": { + "df11_revision": "c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "bf16_revision": "none", + "mode": "kernel-vs-reference", + "decoder": "metal", + "source": "32ff46e183df401598f7756946d4d72c51c5bbe1f3d921f49eb369fffde7a547", + "git": "76932203e838174006341008f6dd9f94cb6bb4f0", + "mlx": "0.32.2" + }, + "exit_code": 0, + "groups": 39, + "partial": false, + "selected_groups": [ + "adaLN_modulation.1", + "final_norm.linear", + "layers.0", + "layers.1", + "layers.2", + "layers.3", + "layers.4", + "layers.5", + "layers.6", + "layers.7", + "layers.8", + "layers.9", + "layers.10", + "layers.11", + "layers.12", + "layers.13", + "layers.14", + "layers.15", + "layers.16", + "layers.17", + "layers.18", + "layers.19", + "layers.20", + "layers.21", + "layers.22", + "layers.23", + "layers.24", + "layers.25", + "layers.26", + "layers.27", + "layers.28", + "layers.29", + "layers.30", + "layers.31", + "layers.32", + "layers.33", + "layers.34", + "layers.35", + "time_embedding" + ], + "compared": 256, + "expected": 256, + "extras_compared": 0, + "mismatched": 0, + "missing_originals": [], + "uncovered_originals": [], + "ignored_originals": [], + "memory_caps_gb": [ + 20, + 22 + ], + "seconds": 486.67997491697315 +} diff --git a/bench/results/refusals/ernie-image-1024-tier16.json b/bench/results/refusals/ernie-image-1024-tier16.json new file mode 100644 index 0000000..d694162 --- /dev/null +++ b/bench/results/refusals/ernie-image-1024-tier16.json @@ -0,0 +1,49 @@ +{ + "exit_code": 2, + "output": "b.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "ernie-image", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-DF11/snapshots/c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "base_path": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image/snapshots/5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true, + "budget_bytes": 9305762474 + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 16, + "label": "CAPPED", + "watchdog_ceiling_bytes": 9842633386, + "limits": { + "tier": { + "tier_gb": 16, + "ram_bytes": 17179869184, + "recommended_bytes": 11453246122, + "memory_limit_bytes": 10737418240, + "cache_limit_bytes": 10737418240, + "wired_limit_bytes": 8589934592, + "reserve_bytes": 1610612736, + "budget_bytes": 11453246122, + "ceiling_bytes": 9842633386, + "fit_budget_bytes": 9305762474, + "budget_source": "ratio", + "is_host": false, + "label": "CAPPED" + }, + "effective": { + "memory_limit_bytes": 10737418240, + "cache_limit_bytes": 10737418240, + "wired_limit_bytes": 8589934592 + }, + "applied": "tier-caps" + }, + "error": "DFloatResourceError: predicted peak 15.2 GiB in the denoise phase exceeds the budget 8.7 GiB (encode 6.8 GiB, denoise 15.2 GiB, vae 7.4 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 632374760, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 632374760, + "mlx_peak_bytes": 28483214, + "elapsed_seconds": 4.034 +} diff --git a/bench/results/refusals/ernie-image-1024-tier24.json b/bench/results/refusals/ernie-image-1024-tier24.json new file mode 100644 index 0000000..0f16c15 --- /dev/null +++ b/bench/results/refusals/ernie-image-1024-tier24.json @@ -0,0 +1,49 @@ +{ + "exit_code": 2, + "output": "b.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "ernie-image", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-DF11/snapshots/c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "base_path": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image/snapshots/5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true, + "budget_bytes": 15032385536 + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 24, + "label": "CAPPED", + "watchdog_ceiling_bytes": 15569256448, + "limits": { + "tier": { + "tier_gb": 24, + "ram_bytes": 25769803776, + "recommended_bytes": 17179869184, + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536, + "reserve_bytes": 1610612736, + "budget_bytes": 17179869184, + "ceiling_bytes": 15569256448, + "fit_budget_bytes": 15032385536, + "budget_source": "ratio", + "is_host": false, + "label": "CAPPED" + }, + "effective": { + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536 + }, + "applied": "tier-caps" + }, + "error": "DFloatResourceError: predicted peak 15.2 GiB in the denoise phase exceeds the budget 14.0 GiB (encode 6.8 GiB, denoise 15.2 GiB, vae 7.4 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 632554960, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 632554960, + "mlx_peak_bytes": 28483214, + "elapsed_seconds": 3.996 +} diff --git a/bench/results/refusals/ernie-image-turbo-1024-tier16.json b/bench/results/refusals/ernie-image-turbo-1024-tier16.json new file mode 100644 index 0000000..9a11d72 --- /dev/null +++ b/bench/results/refusals/ernie-image-turbo-1024-tier16.json @@ -0,0 +1,49 @@ +{ + "exit_code": 2, + "output": "t.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "ernie-image-turbo", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-Turbo-DF11/snapshots/27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "base_path": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image-Turbo/snapshots/bc68c81e2a1730a394d5fc9fae70713dee940140", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true, + "budget_bytes": 9305762474 + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 16, + "label": "CAPPED", + "watchdog_ceiling_bytes": 9842633386, + "limits": { + "tier": { + "tier_gb": 16, + "ram_bytes": 17179869184, + "recommended_bytes": 11453246122, + "memory_limit_bytes": 10737418240, + "cache_limit_bytes": 10737418240, + "wired_limit_bytes": 8589934592, + "reserve_bytes": 1610612736, + "budget_bytes": 11453246122, + "ceiling_bytes": 9842633386, + "fit_budget_bytes": 9305762474, + "budget_source": "ratio", + "is_host": false, + "label": "CAPPED" + }, + "effective": { + "memory_limit_bytes": 10737418240, + "cache_limit_bytes": 10737418240, + "wired_limit_bytes": 8589934592 + }, + "applied": "tier-caps" + }, + "error": "DFloatResourceError: predicted peak 13.9 GiB in the denoise phase exceeds the budget 8.7 GiB (encode 6.8 GiB, denoise 13.9 GiB, vae 7.4 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 707167720, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 707167720, + "mlx_peak_bytes": 28533646, + "elapsed_seconds": 3.644 +} diff --git a/bench/results/refusals/ernie-image-turbo-1024-tier24.json b/bench/results/refusals/ernie-image-turbo-1024-tier24.json new file mode 100644 index 0000000..2911d5f --- /dev/null +++ b/bench/results/refusals/ernie-image-turbo-1024-tier24.json @@ -0,0 +1,49 @@ +{ + "exit_code": 2, + "output": "t.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "ernie-image-turbo", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-Turbo-DF11/snapshots/27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "base_path": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image-Turbo/snapshots/bc68c81e2a1730a394d5fc9fae70713dee940140", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true, + "budget_bytes": 15032385536 + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 24, + "label": "CAPPED", + "watchdog_ceiling_bytes": 15569256448, + "limits": { + "tier": { + "tier_gb": 24, + "ram_bytes": 25769803776, + "recommended_bytes": 17179869184, + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536, + "reserve_bytes": 1610612736, + "budget_bytes": 17179869184, + "ceiling_bytes": 15569256448, + "fit_budget_bytes": 15032385536, + "budget_source": "ratio", + "is_host": false, + "label": "CAPPED" + }, + "effective": { + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536 + }, + "applied": "tier-caps" + }, + "error": "DFloatResourceError: predicted peak 14.4 GiB in the denoise phase exceeds the budget 14.0 GiB (encode 6.8 GiB, denoise 14.4 GiB, vae 7.4 GiB); pass fit_check=False to run anyway", + "footprint_peak_bytes": 714360344, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 714360344, + "mlx_peak_bytes": 28532366, + "elapsed_seconds": 3.653 +} diff --git a/bench/results/tiers/aborts/ernie-image-1024-tier16.json b/bench/results/tiers/aborts/ernie-image-1024-tier16.json new file mode 100644 index 0000000..bb161ee --- /dev/null +++ b/bench/results/tiers/aborts/ernie-image-1024-tier16.json @@ -0,0 +1,30 @@ +{ + "reason": "memory", + "footprint": 10132152768, + "peak_footprint": 10132152768, + "ceiling": 9842633386, + "elapsed": 23.339034333010204, + "budget": 1800.0, + "rss": 1318895616, + "mlx_active": 9691547984, + "mlx_cache": 0, + "verdict_memory": 10132152768, + "verdict_counter": "footprint", + "peak_watched": 10132152768, + "peak_mlx": 9691547984, + "context": { + "model": "ernie-image", + "height": 1024, + "width": 1024, + "seed": 42, + "steps": 50, + "tier_gb": 16, + "label": "CAPPED", + "limits": { + "memory": 10737418240, + "cache": 10737418240, + "wired": 8589934592 + }, + "phase": "set_load" + } +} diff --git a/bench/results/tiers/aborts/ernie-image-1024-tier24.json b/bench/results/tiers/aborts/ernie-image-1024-tier24.json new file mode 100644 index 0000000..ca265fb --- /dev/null +++ b/bench/results/tiers/aborts/ernie-image-1024-tier24.json @@ -0,0 +1,30 @@ +{ + "reason": "memory", + "footprint": 15889441960, + "peak_footprint": 15889441960, + "ceiling": 15569256448, + "elapsed": 20.59221274999436, + "budget": 1800.0, + "rss": 5346459648, + "mlx_active": 13831431540, + "mlx_cache": 1466704645, + "verdict_memory": 15889441960, + "verdict_counter": "footprint", + "peak_watched": 15889441960, + "peak_mlx": 15298136185, + "context": { + "model": "ernie-image", + "height": 1024, + "width": 1024, + "seed": 42, + "steps": 50, + "tier_gb": 24, + "label": "CAPPED", + "limits": { + "memory": 17179869184, + "cache": 16320875724, + "wired": 15032385536 + }, + "phase": "denoise" + } +} diff --git a/bench/results/tiers/aborts/ernie-image-turbo-1024-tier16.json b/bench/results/tiers/aborts/ernie-image-turbo-1024-tier16.json new file mode 100644 index 0000000..bf076f8 --- /dev/null +++ b/bench/results/tiers/aborts/ernie-image-turbo-1024-tier16.json @@ -0,0 +1,30 @@ +{ + "reason": "memory", + "footprint": 9896764056, + "peak_footprint": 9896764056, + "ceiling": 9842633386, + "elapsed": 18.56828275002772, + "budget": 1800.0, + "rss": 4193386496, + "mlx_active": 9396602444, + "mlx_cache": 0, + "verdict_memory": 9896764056, + "verdict_counter": "footprint", + "peak_watched": 9896764056, + "peak_mlx": 9396602444, + "context": { + "model": "ernie-image-turbo", + "height": 1024, + "width": 1024, + "seed": 42, + "steps": 8, + "tier_gb": 16, + "label": "CAPPED", + "limits": { + "memory": 10737418240, + "cache": 10737418240, + "wired": 8589934592 + }, + "phase": "set_load" + } +} diff --git a/bench/results/tiers/ernie-image-1024.json b/bench/results/tiers/ernie-image-1024.json new file mode 100644 index 0000000..f9b178d --- /dev/null +++ b/bench/results/tiers/ernie-image-1024.json @@ -0,0 +1,133 @@ +{ + "exit_code": 0, + "output": "ernie-image.png", + "memory_caps_gb": [ + 20, + 22 + ], + "model_kwargs": { + "model": "ernie-image", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-DF11/snapshots/c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "base_path": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image/snapshots/5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 32, + "label": "MEASURED", + "watchdog_ceiling_bytes": 30064771072, + "limits": { + "tier": { + "tier_gb": 32, + "ram_bytes": 34359738368, + "recommended_bytes": 26800603136, + "memory_limit_bytes": 32641751449, + "cache_limit_bytes": 25460572979, + "wired_limit_bytes": 0, + "reserve_bytes": 2147483648, + "budget_bytes": 26800603136, + "ceiling_bytes": 24653119488, + "fit_budget_bytes": 24653119488, + "budget_source": "device", + "is_host": true, + "label": "MEASURED" + }, + "effective": { + "memory_limit_bytes": 23622320128, + "cache_limit_bytes": 32641751449, + "wired_limit_bytes": 21474836480 + }, + "applied": "host-caps" + }, + "family": "ernie", + "model": "ernie-image", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-DF11/snapshots/c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "repo_id": null, + "revision": "c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "config_source": "config.json" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image/snapshots/5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + "repo_id": null, + "revision": "5346b31d68c9c23758ba56ef8be5e9dc174c7f99" + }, + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_batch": 2, + "text_tokens": 26, + "predict": "uncompiled", + "nonblock_groups": [ + "adaLN_modulation.1", + "final_norm.linear", + "time_embedding" + ], + "cache_limit_in_force": 3600000000, + "drop_set_before_vae": false, + "fit": { + "label": "predicted", + "phases": { + "encode": 7271583948, + "denoise": 16361874888, + "vae": 19163020656 + }, + "peak_phase": "vae", + "peak_bytes": 19163020656, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 10894750607, + "extras": 27961600, + "nonblock": 335544320, + "encoders": 6625216512, + "vae": 168120878 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 642647528, + "mlx_peak": 6683031170, + "footprint_end": 6920067144, + "active_at_start": 27962176 + }, + "set_load": { + "footprint_start": 6920067144, + "mlx_peak": 11259611992, + "footprint_end": 11994408512, + "active_at_start": 28289424 + }, + "denoise": { + "footprint_start": 11994408512, + "mlx_peak": 14181492084, + "footprint_end": 15734679008, + "active_at_start": 11026621584 + }, + "vae": { + "footprint_start": 15734646240, + "mlx_peak": 16335375078, + "footprint_end": 12168980008, + "active_at_start": 11032978796 + } + }, + "decode_launches": 1800, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 22.896248666977044, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "footprint_peak_bytes": 18032961064, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 18032961064, + "mlx_peak_bytes": 15532329081, + "elapsed_seconds": 1791.188 +} diff --git a/bench/results/tiers/ernie-image-turbo-1024-tier24.json b/bench/results/tiers/ernie-image-turbo-1024-tier24.json new file mode 100644 index 0000000..dba2431 --- /dev/null +++ b/bench/results/tiers/ernie-image-turbo-1024-tier24.json @@ -0,0 +1,131 @@ +{ + "exit_code": 0, + "output": "t.png", + "memory_caps_gb": null, + "model_kwargs": { + "model": "ernie-image-turbo", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-Turbo-DF11/snapshots/27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "base_path": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image-Turbo/snapshots/bc68c81e2a1730a394d5fc9fae70713dee940140", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": false, + "budget_bytes": 15032385536 + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 24, + "label": "CAPPED", + "watchdog_ceiling_bytes": 15569256448, + "limits": { + "tier": { + "tier_gb": 24, + "ram_bytes": 25769803776, + "recommended_bytes": 17179869184, + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536, + "reserve_bytes": 1610612736, + "budget_bytes": 17179869184, + "ceiling_bytes": 15569256448, + "fit_budget_bytes": 15032385536, + "budget_source": "ratio", + "is_host": false, + "label": "CAPPED" + }, + "effective": { + "memory_limit_bytes": 17179869184, + "cache_limit_bytes": 16320875724, + "wired_limit_bytes": 15032385536 + }, + "applied": "tier-caps" + }, + "family": "ernie", + "model": "ernie-image-turbo", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-Turbo-DF11/snapshots/27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "repo_id": null, + "revision": "27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "config_source": "config.json" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image-Turbo/snapshots/bc68c81e2a1730a394d5fc9fae70713dee940140", + "repo_id": null, + "revision": "bc68c81e2a1730a394d5fc9fae70713dee940140" + }, + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_batch": 1, + "text_tokens": 26, + "predict": "uncompiled", + "nonblock_groups": [ + "adaLN_modulation.1", + "final_norm.linear", + "time_embedding" + ], + "cache_limit_in_force": 2600000000, + "drop_set_before_vae": true, + "fit": { + "label": "predicted", + "phases": { + "encode": 7335614536, + "denoise": 15069554288, + "vae": 7932725729 + }, + "peak_phase": "denoise", + "peak_bytes": 15069554288, + "budget_bytes": 15032385536, + "fits": false + }, + "sizes": { + "compressed": 10894810183, + "extras": 27961600, + "nonblock": 335544320, + "encoders": 6625216512, + "vae": 168120878 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 645826024, + "mlx_peak": 6696780044, + "footprint_end": 6745577448, + "active_at_start": 27961748 + }, + "set_load": { + "footprint_start": 6745577448, + "mlx_peak": 11259596372, + "footprint_end": 11993785776, + "active_at_start": 28125580 + }, + "denoise": { + "footprint_start": 11993785776, + "mlx_peak": 13399152656, + "footprint_end": 14543480120, + "active_at_start": 11026605964 + }, + "vae": { + "footprint_start": 5498785032, + "mlx_peak": 7428292746, + "footprint_end": 1700071200, + "active_at_start": 32365064 + } + }, + "decode_launches": 288, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 21.516794833005406, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "footprint_peak_bytes": 15446156600, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 15446156600, + "mlx_peak_bytes": 14338222481, + "elapsed_seconds": 169.058 +} diff --git a/bench/results/tiers/ernie-image-turbo-1024.json b/bench/results/tiers/ernie-image-turbo-1024.json new file mode 100644 index 0000000..730cd62 --- /dev/null +++ b/bench/results/tiers/ernie-image-turbo-1024.json @@ -0,0 +1,133 @@ +{ + "exit_code": 0, + "output": "ernie-image-turbo.png", + "memory_caps_gb": [ + 20, + 22 + ], + "model_kwargs": { + "model": "ernie-image-turbo", + "df11_path": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-Turbo-DF11/snapshots/27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "base_path": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image-Turbo/snapshots/bc68c81e2a1730a394d5fc9fae70713dee940140", + "eval_policy": "per-block", + "cache_limit": null, + "fit_check": true + }, + "height": 1024, + "width": 1024, + "memory_ceiling_bytes": null, + "tier_gb": 32, + "label": "MEASURED", + "watchdog_ceiling_bytes": 30064771072, + "limits": { + "tier": { + "tier_gb": 32, + "ram_bytes": 34359738368, + "recommended_bytes": 26800603136, + "memory_limit_bytes": 32641751449, + "cache_limit_bytes": 25460572979, + "wired_limit_bytes": 0, + "reserve_bytes": 2147483648, + "budget_bytes": 26800603136, + "ceiling_bytes": 24653119488, + "fit_budget_bytes": 24653119488, + "budget_source": "device", + "is_host": true, + "label": "MEASURED" + }, + "effective": { + "memory_limit_bytes": 23622320128, + "cache_limit_bytes": 32641751449, + "wired_limit_bytes": 21474836480 + }, + "applied": "host-caps" + }, + "family": "ernie", + "model": "ernie-image-turbo", + "df11": { + "root": "~/.cache/huggingface/hub/models--mingyi456--ERNIE-Image-Turbo-DF11/snapshots/27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "repo_id": null, + "revision": "27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "config_source": "config.json" + }, + "base": { + "root": "~/.cache/huggingface/hub/models--baidu--ERNIE-Image-Turbo/snapshots/bc68c81e2a1730a394d5fc9fae70713dee940140", + "repo_id": null, + "revision": "bc68c81e2a1730a394d5fc9fae70713dee940140" + }, + "eval_policy": "per-block", + "eval_policy_status": "measured", + "cfg_batch": 1, + "text_tokens": 26, + "predict": "uncompiled", + "nonblock_groups": [ + "adaLN_modulation.1", + "final_norm.linear", + "time_embedding" + ], + "cache_limit_in_force": 2600000000, + "drop_set_before_vae": false, + "fit": { + "label": "predicted", + "phases": { + "encode": 7271583948, + "denoise": 14974035520, + "vae": 19163080232 + }, + "peak_phase": "vae", + "peak_bytes": 19163080232, + "budget_bytes": 24653119488, + "fits": true + }, + "sizes": { + "compressed": 10894810183, + "extras": 27961600, + "nonblock": 335544320, + "encoders": 6625216512, + "vae": 168120878 + }, + "peaks": { + "label": "sampled at phase boundaries", + "encode": { + "footprint_start": 643040744, + "mlx_peak": 6696780044, + "footprint_end": 6919411760, + "active_at_start": 27962176 + }, + "set_load": { + "footprint_start": 6919411760, + "mlx_peak": 11259596372, + "footprint_end": 11991475392, + "active_at_start": 28125580 + }, + "denoise": { + "footprint_start": 11991475392, + "mlx_peak": 13399152656, + "footprint_end": 14541530256, + "active_at_start": 11026605964 + }, + "vae": { + "footprint_start": 14541513872, + "mlx_peak": 17353028234, + "footprint_end": 12145059056, + "active_at_start": 11030845448 + } + }, + "decode_launches": 288, + "lifecycle": { + "encoder_loads": 0, + "encoder_seconds": 0.0, + "set_loads": 1, + "set_seconds": 21.93140412500361, + "forced_set_drops": 0 + }, + "versions": { + "mlx": "0.32.2", + "mflux": "0.20.0" + }, + "footprint_peak_bytes": 18621424880, + "footprint_peak_label": "OS phys_footprint, sampled every 0.05 s by the watchdog", + "watched_peak_bytes": 18621424880, + "mlx_peak_bytes": 16279286410, + "elapsed_seconds": 171.488 +} diff --git a/src/mlx_dfloat/mflux/__init__.py b/src/mlx_dfloat/mflux/__init__.py index 9fcf365..780efe7 100644 --- a/src/mlx_dfloat/mflux/__init__.py +++ b/src/mlx_dfloat/mflux/__init__.py @@ -21,8 +21,8 @@ def require_mflux() -> None: def __getattr__(name: str) -> Any: """``DFloatModel`` and the model classes are imported on first use. - The classes are ``DFloatFlux1``, ``DFloatZImage``, ``DFloatFlux2Klein`` and ``DFloatQwenImage21``. So ``import - mlx_dfloat.mflux`` needs neither mflux nor MLX. + The classes are ``DFloatFlux1``, ``DFloatZImage``, ``DFloatFlux2Klein``, ``DFloatQwenImage21`` and + ``DFloatErnieImage``. So ``import mlx_dfloat.mflux`` needs neither mflux nor MLX. """ if name == "DFloatModel": from mlx_dfloat.mflux.families import DFloatModel @@ -44,10 +44,14 @@ def __getattr__(name: str) -> Any: from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 return DFloatQwenImage21 + if name == "DFloatErnieImage": + from mlx_dfloat.mflux.ernie.model import DFloatErnieImage + + return DFloatErnieImage raise AttributeError(f"module {__name__!r} has no attribute {name!r}") # The model classes stay out of __all__: a star-import must not need mflux; `mlx_dfloat.mflux.DFloatFlux1`, -# `mlx_dfloat.mflux.DFloatZImage`, `mlx_dfloat.mflux.DFloatFlux2Klein` and `mlx_dfloat.mflux.DFloatQwenImage21` still -# work. +# `mlx_dfloat.mflux.DFloatZImage`, `mlx_dfloat.mflux.DFloatFlux2Klein`, `mlx_dfloat.mflux.DFloatQwenImage21` and +# `mlx_dfloat.mflux.DFloatErnieImage` still work. __all__ = ["DFloatModel", "require_mflux"] diff --git a/src/mlx_dfloat/mflux/ernie/__init__.py b/src/mlx_dfloat/mflux/ernie/__init__.py new file mode 100644 index 0000000..fce66f8 --- /dev/null +++ b/src/mlx_dfloat/mflux/ernie/__init__.py @@ -0,0 +1 @@ +"""ERNIE-Image and ERNIE-Image-Turbo on mflux from a DFloat11 transformer.""" diff --git a/src/mlx_dfloat/mflux/ernie/init.py b/src/mlx_dfloat/mflux/ernie/init.py new file mode 100644 index 0000000..2a7a66c --- /dev/null +++ b/src/mlx_dfloat/mflux/ernie/init.py @@ -0,0 +1,176 @@ +"""The ERNIE-Image base repository without its transformer: text encoder, VAE and tokenizer, resolved and loaded. + +mflux's ``ErnieImageInitializer.init`` wants every component under one path, so it would download and size the +transformer; this module loads the other components through mflux's own loader and applier with weight definitions +that name only them. The download patterns never name ``transformer/``, the prompt enhancer ``pe/`` (7.66 GB, which +mflux never uses) or its tokenizer ``pe_tokenizer/``. +""" + +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from mlx.utils import tree_flatten + +from mlx_dfloat.errors import DFloatFormatError +from mlx_dfloat.mflux import _hub, require_mflux +from mlx_dfloat.mflux._hub import ResolvedRepo as ResolvedRepo +from mlx_dfloat.mflux._hub import resolve as resolve + +DF11_PATTERNS: tuple[str, ...] = ("*.safetensors", "config.json") +ENCODER_PATTERNS: tuple[str, ...] = ("text_encoder/*.safetensors", "text_encoder/*.json") +VAE_PATTERNS: tuple[str, ...] = ("vae/*.safetensors", "vae/*.json") +# The tokenizer's files, named one by one (never "tokenizer/*" or mflux's "tokenizer/**", which would fetch anything a +# later revision drops in there). Both base snapshots hold the first two (2026-10-08); the other three are files +# transformers' AutoTokenizer reads when present, named so a later revision that adds them still loads. +TOKENIZER_PATTERNS: tuple[str, ...] = ( + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "tokenizer/special_tokens_map.json", + "tokenizer/added_tokens.json", + "tokenizer/chat_template.jinja", +) +BASE_PATTERNS: tuple[str, ...] = ENCODER_PATTERNS + VAE_PATTERNS + TOKENIZER_PATTERNS + +_WHAT = "ERNIE-Image base" +# Parameters no checkpoint holds: built at construction, so a base never replaces them. mflux 0.20.0 has none for +# either module (checked by building both and listing their parameters): the text encoder's RoPE frequencies are the +# underscore attribute ``_inv_freq`` (text_encoder.py:137-138), which MLX does not count as a parameter, and every +# other encoder parameter sits under ``language_model.model.``, the prefix the base file holds; the VAE's batch-norm +# statistics come from the file, as for FLUX.2 Klein (flux2_weight_mapping.py:202-209). +COMPUTED: dict[str, frozenset[str]] = { + "text_encoder": frozenset(), + "vae": frozenset(), +} +_SHOWN = 3 # missing names quoted in the error + + +@dataclass(slots=True, kw_only=True) +class ErnieComponents: + """The loaded base components: mflux modules (weights still lazy) and the tokenizers.""" + + vae: Any + text_encoder: Any + tokenizers: dict[str, Any] + + +def _definition(names: tuple[str, ...], patterns: tuple[str, ...]) -> type: + """An mflux weight definition holding only the named ERNIE-Image components.""" + require_mflux() + from mflux.models.ernie_image.weights.ernie_weight_definition import ErnieWeightDefinition + + return _hub.subset_definition(ErnieWeightDefinition, names, patterns) + + +def encoder_definition() -> type: + """The text encoder component of mflux's ERNIE-Image definition (its language-model prefix filter kept).""" + return _definition(("text_encoder",), ENCODER_PATTERNS) + + +def vae_definition() -> type: + """The VAE component of mflux's ERNIE-Image definition, with its download patterns.""" + return _definition(("vae",), VAE_PATTERNS) + + +def _shapes(module: Any) -> dict[str, tuple[int, ...]]: + """Every parameter's shape by path (read without evaluating anything).""" + flat: dict[str, Any] = dict(tree_flatten(module.parameters())) + return {name: tuple(p.shape) for name, p in flat.items()} + + +def _require_replaced(root: Path, name: str, before: dict[str, Any], module: Any) -> None: + """Require every parameter of ``module`` to have been replaced by the base since ``before`` was taken. + + mflux's applier updates with ``strict=False``: a tensor the base lacks (a lost shard, a renamed key) leaves that + parameter at its random float32 init without an error. ``before`` holds each parameter as built; one still the + same object after the load was not in the base (``COMPUTED`` names the exceptions). + + Raises: + DFloatFormatError: A parameter was not replaced by the base. + """ + after: dict[str, Any] = dict(tree_flatten(module.parameters())) + missing = [n for n, init in before.items() if n not in COMPUTED[name] and after.get(n) is init] + if missing: + more = len(missing) - _SHOWN + shown = ", ".join(missing[:_SHOWN]) + (f" and {more} more" if more > 0 else "") + raise DFloatFormatError( + f"{root}: {name} {shown} not in the base: a missing shard or a renamed tensor?" + ) + + +def load_text_encoder(root: Path, *, dims: Mapping[str, Any] | None = None) -> Any: + """A fresh ``ErnieMistralTextEncoder`` with the base's weights (lazy). + + ``dims`` are the encoder's constructor arguments (none: mflux's defaults, the published model). Only the language + model's tensors are loaded; the vision tower and projector in the same file never are. mflux's applier assigns + weights without a shape check and ignores a tensor the base lacks, so after the load every parameter's shape is + compared with the freshly built module's (another model's base is refused here, before any encode) and every + parameter must have been replaced by the base. + + Raises: + DFloatFormatError: The directory lacks the text encoder, a loaded parameter's shape differs from the model's + (another model's base), or a parameter is not in the base. + """ + require_mflux() + from mflux.models.ernie_image.model.ernie_text_encoder.text_encoder import ( + ErnieMistralTextEncoder, + ) + + encoder = ErnieMistralTextEncoder(**dict(dims or {})) + want = _shapes(encoder) + before: dict[str, Any] = dict(tree_flatten(encoder.parameters())) + _hub.load_into(root, encoder_definition(), {"text_encoder": encoder}, what=_WHAT) + for name, got in _shapes(encoder).items(): + if want.get(name) != got: + raise DFloatFormatError( + f"{root}: {name} has shape {got}, this model needs {want.get(name)}: another model's " + "base repository?" + ) + _require_replaced(root, "text_encoder", before, encoder) + return encoder + + +def load_vae(root: Path) -> Any: + """A fresh ``Flux2VAE`` (ERNIE-Image's VAE) with the base's weights (lazy). + + Raises: + DFloatFormatError: The directory lacks the VAE, or a parameter is not in the base. + """ + require_mflux() + from mflux.models.flux2.model.flux2_vae.vae import Flux2VAE + + vae = Flux2VAE() + before: dict[str, Any] = dict(tree_flatten(vae.parameters())) + _hub.load_into(root, vae_definition(), {"vae": vae}, what=_WHAT) + _require_replaced(root, "vae", before, vae) + return vae + + +def load_tokenizers(root: Path) -> dict[str, Any]: + """The ERNIE-Image tokenizer (one, keyed ``"ernie"``, read from ``tokenizer/``). + + Raises: + DFloatFormatError: The tokenizer files are missing or incomplete. + """ + require_mflux() + from mflux.models.common.tokenizer.tokenizer_loader import TokenizerLoader + from mflux.models.ernie_image.weights.ernie_weight_definition import ErnieWeightDefinition + + try: + return dict( + TokenizerLoader.load_all( + definitions=ErnieWeightDefinition.get_tokenizers(), model_path=str(root) + ) + ) + except (OSError, RuntimeError) as exc: # OSError: transformers' error for a partial tokenizer/ + raise DFloatFormatError(f"{root}: no usable ERNIE-Image tokenizer: {exc}") from exc + + +def load_base(root: Path) -> ErnieComponents: + """Text encoder, VAE and tokenizer of a base repository (weights lazy, every parameter from the base).""" + return ErnieComponents( + vae=load_vae(root), + text_encoder=load_text_encoder(root), + tokenizers=load_tokenizers(root), + ) diff --git a/src/mlx_dfloat/mflux/ernie/memory.py b/src/mlx_dfloat/mflux/ernie/memory.py new file mode 100644 index 0000000..b9d08bf --- /dev/null +++ b/src/mlx_dfloat/mflux/ernie/memory.py @@ -0,0 +1,184 @@ +"""ERNIE-Image memory rules: the constants per CFG batch, the text-token count and the text encoder's resident bytes. + +The phase arithmetic itself is the shared one (``mlx_dfloat.mflux._phases``). ERNIE-Image differs in three inputs. Its +text encoder's file also holds a vision tower and a projector that mflux never loads (see ``encoder_bytes_used``). +Its prompt length varies per call (up to the tokenizer's 2048 tokens), so the text tokens are counted per prompt +(``prompt_tokens``). And classifier-free guidance runs as one batch-2 transformer call, which doubles the activations +but not the decoded weights, so the constants are keyed by the call's batch size (1 or 2), not by the model. + +The transformer has one block kind, so the cache limit holds one decoded block next to the activations, and each batch +size carries its own activation allowance. One encode term serves both batch sizes: the larger of the two measured +encodes at 26 tokens, one prompt (batch 1) and two (batch 2). The VAE transient is mflux's Flux2VAE decode; re-measure it when an mflux +upgrade changes the VAE's dtype. +""" + +import re +from collections.abc import Sequence +from dataclasses import replace +from pathlib import Path +from typing import Any + +import mlx.core as mx + +from mlx_dfloat._safetensors import read_header +from mlx_dfloat.format import DF11Checkpoint +from mlx_dfloat.mflux._phases import FamilySizes, PhaseConstants +from mlx_dfloat.mflux._phases import sizes_for as family_sizes_for +from mlx_dfloat.mflux.ernie.names import NONBLOCK_GROUPS + +# The only text-encoder tensors mflux loads (its weight definition's prefix filter for the text encoder). +ENCODER_PREFIX = "language_model.model." +# The tokenizer's max_length (mflux's weight definition; tokenization truncates at it). +TEXT_MAX_LENGTH = 2048 +MAX_MEASURED_PIXELS = 1024 * 1024 # no run above 1024² on this path + +# Measured on 2026-10-08 (git a771e58, M1 Max 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0): two one-step de-risk runs at +# 1024², seed 42, the 26-token lighthouse prompt; one process each: build, encode, drop, set load, one step at the +# planner's cache limit for the block-cache probe's allowance, VAE decode on the resident set. ERNIE-Image-Turbo +# (mingyi456/ERNIE-Image-Turbo-DF11 @ 27f84b4 over baidu/ERNIE-Image-Turbo @ bc68c81) at guidance 1.0: batch 1, run +# "derisk-turbo-1" in bench/results/calibration/ernie-image-turbo-1024.json. ERNIE-Image (mingyi456/ERNIE-Image-DF11 +# @ c2dd30a over baidu/ERNIE-Image @ 5346b31) at guidance 4.0 without a negative prompt (mflux's " " then the prompt): +# batch 2, run "derisk-base-1" in bench/results/calibration/ernie-image-1024.json. Bytes throughout; re-measure before +# changing any. +# +# REFERENCE_TOKENS: the calibration point, 1024² (4096 image tokens) + the run's text tokens, 26 (the prompt's; the +# negative " " is 2, and a call is sized by its longest prompt). +REFERENCE_TOKENS = 4096 + 26 +# OVERHEAD: the footprint MLX's counters do not see after the build (footprint - MLX active - MLX cache), the larger of +# the two runs: Turbo 644_400_592 - 27_961_608 - 4_806 = 616_434_178 (the base's 548_194_818). One value for both +# batches: the build is the same. +OVERHEAD = 616_434_178 # 0.574 GiB +# VAE_TRANSIENT, by the overall-peak rule over every 1024² sample (the two de-risk runs so far): the highest measured +# VAE-phase peak on the resident set minus the VAE phase's other terms (compressed + extras + non-block + VAE file + +# OVERHEAD), from the Turbo run: 19_163_080_232 - (10_894_810_183 + 27_961_600 + 335_544_320 + 168_120_878 + +# 616_434_178). The base run's peak (19_127_199_104, with its own sizes) would give a term 35_821_552 B lower. A fit +# check must not under-predict, so the term covers the highest sample; later 1024² samples join the rule. One value +# for both batches: the VAE decode does not see the CFG batch. About 0.96x FLUX.2 Klein 4B's for the same Flux2VAE. +VAE_TRANSIENT = 7_120_209_073 # 6.631 GiB +# DENOISE_ACTIVATION per batch, by the overall-peak rule over every 1024² denoise-phase sample: the highest denoise +# footprint peak minus the denoise phase's other terms at the planner's cache limit (compressed + extras + non-block + +# in-flight block 436_207_616 + cache limit + OVERHEAD), at REFERENCE_TOKENS. Batch 1 (limit 2_600_000_000): the second +# CAPPED 24 GB run of ERNIE-Image-Turbo (label "capped-24-gb-2" in ernie-image-turbo-1024.json; git f4cd325), whose +# watched peak was in the denoise phase: 15_446_156_600 - (10_894_810_183 + 27_961_600 + 335_544_320 + 436_207_616 + +# 2_600_000_000 + 616_434_178) = 535_198_703. The first identical run ("capped-24-gb-1", git 5e2e7c2) peaked at +# 15_069_554_288 (term 158_596_391), 376_602_312 B (0.35 GiB) lower, and the de-risk run's 14_974_035_520 gave +# 63_077_623: the denoise phase varies that much between identical runs. A fit check must not under-predict, so the +# term covers the highest sample. Batch 2 (base, limit 3_600_000_000), its de-risk +# run, the only denoise-phase sample: 16_361_874_888 - (10_894_750_607 + 27_961_600 + 335_544_320 + 436_207_616 + +# 3_600_000_000 + 616_434_178) = 450_976_567. +DENOISE_ACTIVATION: dict[int, int] = {1: 535_198_703, 2: 450_976_567} +# ALLOWANCE per batch: the cache room for the activation buffers a block frees, at REFERENCE_TOKENS. The block-cache +# probe (one synthetic block at 1024² + 26 tokens, its decoded buffer allocated then released per iteration) saw the +# buffer reused from the second iteration at limits of 2.6e9 B (batch 1) and 3.6e9 B (batch 2, the smallest tried), with +# the steady cache drop of a clean reuse; the allowance is that limit minus one decoded block (436_207_616). +ALLOWANCE: dict[int, int] = {1: 2_163_792_384, 2: 3_163_792_384} +# DENOISE_ACTIVATION_FLOOR per batch: a size-independent lower bound on the denoise term. ERNIE-Image's hidden state is +# float32 (mflux computes the timestep conditioning in float32), and a float32 input times a BF16 weight makes MLX copy +# the weight to float32 for the matmul: the MLP's gate_proj and up_proj together, 2 x 12288 x 4096 x 4 = 402_653_184 B, +# whatever the image size. Batch 2 only: its 1024² term (450_976_567) is above the floor, so the calibration point is +# unchanged. Batch 1 carries none until a 512² run measures whether a one-prompt step holds the copies beyond the cache +# term (its 1024² term, 535_198_703, is above the floor anyway). +DENOISE_ACTIVATION_FLOOR: dict[int, int | None] = {1: None, 2: 402_653_184} +# ENCODE_ACTIVATION: what a prompt encode held beyond the language model's evaluated tensors by footprint (encode +# footprint peak - encoder_bytes_used - OVERHEAD), the larger of the two runs: Turbo, 7_335_614_536 - 6_625_216_512 - +# 616_434_178 = 93_963_846 (the base: 15_943_190). Taken by footprint so the encode phase never under-predicts a +# calibration run (by MLX's measure the holds were 4_188_932 and 29_933_258). The encoder's last layer never runs (mflux +# returns the second-to-last hidden state): counted with every language-model tensor (6_858_012_672) the term came out +# negative. Measured at 26 tokens with one prompt (batch 1) and two (batch 2, where both encoder passes are one lazy +# graph evaluated once); the larger serves both batches, unscaled for longer prompts. +ENCODE_ACTIVATION = 93_963_846 + + +def _constants(batch: int) -> PhaseConstants: + return PhaseConstants( + overhead_bytes=OVERHEAD, + vae_transient_bytes=VAE_TRANSIENT, + denoise_activation_at_reference=DENOISE_ACTIVATION[batch], + reference_tokens=REFERENCE_TOKENS, + allowance_at_reference=ALLOWANCE[batch], + allowance_reference_tokens=REFERENCE_TOKENS, + max_measured_pixels=MAX_MEASURED_PIXELS, + encode_activation_bytes=ENCODE_ACTIVATION, + denoise_activation_floor_bytes=DENOISE_ACTIVATION_FLOOR[batch], + ) + + +CONSTANTS: dict[int, PhaseConstants] = {1: _constants(1), 2: _constants(2)} +# What a prompt encode may hold by MLX's measure beyond the estimate's encode phase before the model warns (the bound +# is the evaluated language model + ENCODE_ACTIVATION + this slack, 6_987_615_814 B). 256 MiB: the runs held 358_210_370 +# (Turbo) and 332_466_044 (base) under the bound, and an encode that also loads the vision tower and projector +# (840_167_424 B more) is over it by 507_701_380 at every prompt length. +ENCODE_SLACK_BYTES = 256 * 1024**2 +# The prompt length ENCODE_ACTIVATION was measured at (the de-risk prompt's text tokens). +ENCODE_ACTIVATION_TOKENS = 26 +_LAYER = re.compile(r"^language_model\.model\.layers\.(\d+)\.") + + +def encode_peak_bound(encode_phase_bytes: int, constants: PhaseConstants) -> int: + """What a prompt encode may hold by MLX's measure before the model warns. + + That is the estimate's encode phase without its overhead (the language model and the encode activations) plus + ``ENCODE_SLACK_BYTES``. ``constants`` are the call's (``CONSTANTS[batch]``). + """ + return encode_phase_bytes - constants.overhead_bytes + ENCODE_SLACK_BYTES + + +def encode_peak_warning(held: int, bound: int, *, text_tokens: int) -> str | None: + """The warning for an encode that held ``held`` bytes against ``bound``, or None at or under it. + + Two causes fit an excess: mflux loading the vision tower stored in the encoder's file, or a prompt longer than the + one the encode term was measured on (the term is not scaled with the prompt length). The message names the call's + longest prompt in tokens so the reader can tell which. + """ + if held <= bound: + return None + return ( + f"the prompt encode held {held / 1024**3:.2f} GiB for a {text_tokens}-token prompt, over the " + f"{bound / 1024**3:.2f} GiB the fit estimate allows: mflux may now load the vision tower stored in the text " + f"encoder's file, or this prompt holds more than the encode term, which was measured on one " + f"{ENCODE_ACTIVATION_TOKENS}-token prompt and is not scaled with the prompt length; the estimate " + "under-predicts the encode phase either way" + ) + + +def encoder_bytes_used(base_root: Path) -> int: + """The text encoder's bytes an ERNIE-Image prompt encode makes resident, read from the safetensors headers. + + Only the language model's tensors (``language_model.model.*``) are loaded; the vision tower and the projector in + the same file are never loaded. The encoder returns its second-to-last hidden state, so under lazy evaluation its + last layer never runs and its weights are never read: of ``n`` numbered layers the first ``n - 1`` count (25 of the + published 26). Every language-model tensor outside the numbered layers (embeddings, final norm) counts in full. + + Raises: + DFloatFormatError: An encoder file's header cannot be read. + """ + tensors = [ + (name, info.nbytes) + for path in sorted((base_root / "text_encoder").glob("*.safetensors")) + for name, info in read_header(path).items() + if name.startswith(ENCODER_PREFIX) + ] + indices = [int(m.group(1)) for name, _n in tensors if (m := _LAYER.match(name)) is not None] + evaluated = max(indices, default=-1) # layers 0 .. n - 2 run; the last one, n - 1, does not + total = 0 + for name, nbytes in tensors: + layer = _LAYER.match(name) + if layer is None or int(layer.group(1)) < evaluated: + total += nbytes + return total + + +def prompt_tokens(tokenizer: Any, prompt: str) -> int: + """The text tokens mflux feeds the transformer for ``prompt``: its attention mask's ones, at most 2048.""" + return min(int(mx.sum(tokenizer.tokenize(prompt).attention_mask[0])), TEXT_MAX_LENGTH) + + +def text_tokens_for(tokenizer: Any, prompts: Sequence[str]) -> int: + """The text tokens of one call: mflux pads every prompt of a CFG batch to the longest.""" + return max(prompt_tokens(tokenizer, p) for p in prompts) + + +def sizes_for(ckpt: DF11Checkpoint, base_root: Path) -> FamilySizes: + """The family sizes of an ERNIE-Image checkpoint and base, the encoder counted by its language model.""" + sizes = family_sizes_for(ckpt, base_root, nonblock_groups=NONBLOCK_GROUPS) + return replace(sizes, encoders=encoder_bytes_used(base_root)) diff --git a/src/mlx_dfloat/mflux/ernie/model.py b/src/mlx_dfloat/mflux/ernie/model.py new file mode 100644 index 0000000..b113aee --- /dev/null +++ b/src/mlx_dfloat/mflux/ernie/model.py @@ -0,0 +1,664 @@ +"""``DFloatErnieImage``: mflux's ERNIE-Image pipeline over a DFloat11 transformer, decoded one block at a time. + +The class subclasses mflux's ``ErnieImage`` so its loop, scheduler, callbacks, VAE decode and image metadata run +unchanged; only construction and the prelude of ``generate_image`` differ. The text encoder and the compressed +transformer are never resident together (see ``mlx_dfloat.mflux.lifecycle``): a call's prompts are encoded first into +a cache of mflux's own text batches, which the loop reads without touching the encoder. +""" + +import json +import logging +import traceback +import warnings +from collections.abc import Callable, Sequence +from dataclasses import asdict +from importlib import metadata +from pathlib import Path +from typing import Any + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_flatten + +from mlx_dfloat import _metal_decode +from mlx_dfloat.decode import DecodeResult +from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError, DFloatUnsupportedError +from mlx_dfloat.format import DF11Checkpoint, MxGroup, open_checkpoint +from mlx_dfloat.integrate.coverage import load_resident_set +from mlx_dfloat.integrate.memory import CallPlan, budget_bytes, largest_decoded_bytes +from mlx_dfloat.integrate.names import NameMap +from mlx_dfloat.integrate.providers import DF11Provider +from mlx_dfloat.integrate.resident import clear_nonblock, decode_nonblock, install_nonblock +from mlx_dfloat.mflux import _pipeline, require_mflux +from mlx_dfloat.mflux import families as registry +from mlx_dfloat.mflux._compile import uncompiled +from mlx_dfloat.mflux._phases import FamilySizes +from mlx_dfloat.mflux.ernie import init as einit +from mlx_dfloat.mflux.ernie import memory as emem +from mlx_dfloat.mflux.ernie.init import ErnieComponents, ResolvedRepo +from mlx_dfloat.mflux.ernie.names import NONBLOCK_GROUPS, ernie_name_map +from mlx_dfloat.mflux.ernie.transformer import ErnieBuild, build_transformer +from mlx_dfloat.mflux.lifecycle import Lifecycle + +require_mflux() +from mflux.models.ernie_image.variants.txt2img.ernie_image import ErnieImage # noqa: E402 + +log = logging.getLogger("mlx_dfloat.mflux.ernie") + +# (DF11 repository, base repository), derived from the registry. +MODELS: dict[str, tuple[str, str]] = { + n: (e.df11_repo, e.base_repo) for n, e in registry.MODELS.items() if e.family == "ernie" +} +POLICIES: tuple[str, ...] = ("per-block", "depth2") +# The policies an ERNIE-Image run has measured; depth2 is accepted, and the report labels it unmeasured. +MEASURED_POLICIES: frozenset[str] = frozenset({"per-block"}) +# What else may stay active across calls besides the assembly baseline, the cached batches and the VAE. It covers +# mflux's ``ErnieTransformer._pos_cache`` (transformer.py:120-122): up to 64 (cos, sin, attn_mask) triples keyed by +# image size, text length and the prompts' lengths, grown by one entry per new (size, prompt lengths) and evicted +# oldest-first past 64; it is not part of the set. One 1024² batch-2 entry at the 2048-token maximum holds 6_316_032 B +# (cos and sin (2, 6144, 1, 128) bf16, the mask (2, 1, 1, 6144) bf16; measured on mlx 0.32.2 / mflux 0.20.0), so 64 +# entries hold at most 404_226_048 B (0.38 GiB), inside the slack. Above 1024² (only with fit_check=False) an entry grows +# with the image tokens, so the worst case grows with it. +RETAINED_SLACK_BYTES = 2 * 1024**3 + + +def _clear_frames(exc: BaseException) -> None: + """Release the locals of the frames ``exc`` and every exception in its chain (causes and contexts) passed through. + + The walk keeps a seen-set, so a chain that loops back on itself ends. + """ + seen: set[int] = set() + pending: list[BaseException] = [exc] + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + if current.__traceback__ is not None: + traceback.clear_frames(current.__traceback__) + pending.extend(e for e in (current.__cause__, current.__context__) if e is not None) + + +class _EncoderNotResident: + """Takes the text encoder's place while it is dropped: mflux calls the encoder only for a batch not cached.""" + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Refuse: every batch a call needs is encoded before mflux's loop starts. + + Raises: + DFloatIntegrationError: Always. + """ + del args, kwargs + raise DFloatIntegrationError( + "the text encoder is not resident: a prompt was not encoded before the call " + "(generate_image encodes the prompts it needs first; call encode() for any other use)" + ) + + +class DFloatErnieImage(ErnieImage): # type: ignore[misc] # mflux ships no type information + """ERNIE-Image and ERNIE-Image-Turbo generation from a DFloat11 transformer through mflux. + + The transformer stays compressed and each of its 36 blocks is decoded on the GPU as it runs; the three non-block + groups (the timestep MLP, the shared modulation, the final norm's projection) are decoded once when the set loads. + The step function runs uncompiled on every chip (the per-block evaluation cannot run inside a compiled function). + The decode is lossless. + + The text encoder and the compressed set are never resident together: a new prompt list after a generation drops + the set, reloads the encoder, encodes and reloads the set. + """ + + def __init__( + self, + model: str = "ernie-image-turbo", + *, + df11_path: str | None = None, + base_path: str | None = None, + eval_policy: str = "per-block", + cache_limit: int | None = None, + fit_check: bool = True, + budget_bytes: int | None = None, + quantize: int | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + bake_lora: bool = True, + ) -> None: + """Resolve the checkpoint and the base repository and build the model. + + The checkpoint header, the base's tokenizer and the transformer's extras are read; the encoder, VAE and + compressed weights stay lazy until used. The default checkpoint and base are pinned to the revisions this + package was checked against; a user's own paths are taken as given. ``budget_bytes`` replaces the device's + budget in every call's fit 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"``, 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 an ERNIE-Image DF11 checkpoint, its block count differs from the + model's, or the base lacks a component or is another model's. + DFloatAccessError: A gated repository this account may not read. + DFloatDependencyError: The ``mlx-dfloat[mflux]`` extra is not installed. + DFloatIntegrationError: mflux's own weight definition or mapping does not shape the way this adapter + expects. + """ + _pipeline.refuse_construction_args( + quantize=quantize, lora_paths=lora_paths, lora_scales=lora_scales, bake_lora=bake_lora + ) + _pipeline.check_eval_policy(eval_policy, POLICIES) + if model not in MODELS: + raise DFloatUnsupportedError(f"model {model!r}: this path runs {sorted(MODELS)}") + # The GPU decoder proves itself on the two packaged self-check groups before anything is resolved, downloaded or + # loaded, so a broken decoder refuses here and not minutes later at the first block. + _metal_decode.ensure_canary(force_direct=False) + from mflux.models.common.config.model_config import ModelConfig + + model_config = ModelConfig.from_name(model_name=model, base_model=None) + entry = registry.MODELS[model] + df11_default, base_default = MODELS[model] + pin = None if df11_path else entry.df11_revision + df11 = einit.resolve(df11_path or df11_default, patterns=einit.DF11_PATTERNS, revision=pin) + ckpt = open_checkpoint(df11.root) + base_pin = None if base_path else entry.base_revision + base = einit.resolve( + base_path or base_default, patterns=einit.BASE_PATTERNS, revision=base_pin + ) + # mflux builds the transformer with the model's overrides (ErnieImageInitializer._init_models). + build = build_transformer(ckpt, transformer_overrides=model_config.transformer_overrides) + components = einit.load_base(base.root) + self._assemble( + model=model, + model_config=model_config, + ckpt=ckpt, + df11=df11, + base=base, + components=components, + build=build, + sizes=emem.sizes_for(ckpt, base.root), + eval_policy=eval_policy, + cache_limit=cache_limit, + fit_check=fit_check, + budget_bytes=budget_bytes, + ) + + @classmethod + def _from_parts(cls, **parts: Any) -> "DFloatErnieImage": + """Assemble a model from already loaded parts (what ``__init__`` ends with).""" + self = cls.__new__(cls) + self._assemble(**parts) + return self + + def _assemble( + self, + *, + model: str = "custom", + model_config: Any, + ckpt: DF11Checkpoint, + df11: ResolvedRepo, + base: ResolvedRepo, + components: ErnieComponents, + build: ErnieBuild, + sizes: FamilySizes, + name_map: NameMap | None = None, + eval_policy: str = "per-block", + cache_limit: int | None = None, + fit_check: bool = True, + budget_bytes: int | None = None, + decode: Callable[[MxGroup], DecodeResult] | None = None, + ) -> None: + from mflux.models.ernie_image.ernie_image_initializer import ErnieImageInitializer + + _pipeline.check_eval_policy(eval_policy, POLICIES) + nn.Module.__init__(self) # type: ignore[attr-defined] # not ErnieImage.__init__: mflux's initializer loads the whole base + # prompt_cache, model_config, callbacks, tiling_config (ernie_image_initializer.py:40-45) + ErnieImageInitializer._init_config(self, model_config) + self.vae = components.vae + self.text_encoder = components.text_encoder + self.tokenizers = components.tokenizers + self.transformer = build.transformer + self.bits = None + self.lora_paths: list[str] = [] + self.lora_scales: list[float] = [] + self._model = model + self._constants = emem.CONSTANTS + self._ckpt, self._df11, self._base, self._sizes = ckpt, df11, base, sizes + self._shapes, self._nonblock = build.shapes, build.nonblock + self._decode = decode + self._names = ernie_name_map() if name_map is None else name_map + self._batch = 1 + self._text_tokens: int | None = None + self._predict_mode: str | None = None + self._largest = largest_decoded_bytes(ckpt, self._names, skip=frozenset(NONBLOCK_GROUPS)) + self._policy, self._cache_limit_override, self._fit_check = ( + eval_policy, + cache_limit, + fit_check, + ) + self._budget_override = budget_bytes + self._provider: DF11Provider | None = None + self._plan: CallPlan | None = None + self._tracker = _pipeline.PhaseTracker() + self._launches = 0 + self._baseline_active = int(mx.get_active_memory()) + # mflux's text batches, keyed by the call's prompt list (``_batch_key``); mflux's own prompt_cache stays empty. + self._batches: dict[str, tuple[mx.array, ...]] = {} + self._lifecycle = Lifecycle( + load_encoders=self._load_encoders, + unload_encoders=self._unload_encoders, + encode=self._encode, + load_set=self._load_set, + unload_set=self._unload_set, + prompt_cache=self._batches, + retained_bound=self._retained_bound, + encoders_loaded=True, + ) + self.callbacks.register( + _pipeline.VaePoolGuard( + plan=lambda: self._plan, + end_denoise=lambda: self._tracker.end("denoise"), + drop_set=lambda: self._lifecycle.drop_set(), + begin_vae=lambda: self._tracker.begin("vae"), + ) + ) + + # --- lifecycle callbacks ------------------------------------------------------------------------ + + @staticmethod + def _batch_key(prompts: Sequence[str]) -> str: + """The cache key of a call's prompt list (in mflux's batch order).""" + return json.dumps(list(prompts)) + + def _retained_bound(self) -> int: + """What may stay active after a drop: the assembly baseline, the cached batches, the VAE, slack. + + The whole VAE counts, its encoder included: an upper bound, since mflux's decode reads only part of it and the + rest stays unloaded under lazy loading, but which part depends on mflux's Flux2VAE layout. + """ + batches = sum(int(a.nbytes) for arrays in self._batches.values() for a in arrays) + flat: list[tuple[str, mx.array]] = list(tree_flatten(self.vae.parameters())) # type: ignore[arg-type] + vae = sum(int(v.nbytes) for _name, v in flat) + return self._baseline_active + batches + vae + RETAINED_SLACK_BYTES + + def _load_encoders(self) -> None: + self.text_encoder = einit.load_text_encoder(self._base.root) + + def _unload_encoders(self) -> None: + # Not None: mflux would call it on a cache miss and fail with a TypeError far from the cause. + self.text_encoder = _EncoderNotResident() + + def _encode(self, key: str) -> tuple[mx.array, ...]: + """Mflux's text batch for the prompt list ``key`` names, built by mflux's own ``build_text_batch``.""" + from mflux.models.ernie_image.model.ernie_text_encoder.prompt_encoder import ( + ErniePromptEncoder, + ) + + text_bth, text_lens = ErniePromptEncoder.build_text_batch( + prompts=json.loads(key), + tokenizer=self.tokenizers["ernie"], + text_encoder=self.text_encoder, + ) + return text_bth, text_lens + + def _load_set(self) -> None: + try: + self._install_set() + except BaseException as exc: + # The traceback's frames hold the compressed set (_install_set's ``resident``) and any decoded non-block + # weight for as long as the caller keeps the exception; the retry's memory check would then refuse. + _clear_frames(exc) + self._unload_set() # nothing half-installed: the lifecycle believes no set is resident + raise + + def _install_set(self) -> None: + resident = load_resident_set(self._ckpt) + nonblock = {g: self._ckpt.groups[g].matrix_names for g in NONBLOCK_GROUPS if g in resident} + install_nonblock( + self.transformer, + decode_nonblock( + {g: resident[g] for g in nonblock}, + nonblock, + self._nonblock, + self._names, + decode=self._decode, + ), + ) + blocks = {n: g for n, g in resident.items() if n not in NONBLOCK_GROUPS} + self._provider = DF11Provider( + blocks, + {n: self._ckpt.groups[n].matrix_names for n in blocks}, + self._names, + decode=self._decode, + ) + # One transformer call per step (classifier-free guidance is one batch-2 call); verified inside the call to + # keep every family's invariant. + self.transformer.attach( + self._provider, self._shapes, eval_policy=self._policy, verify_in_call=True + ) + + def _unload_set(self) -> None: + if self._provider is not None: + self._launches += self._provider.launches + self.transformer.detach() + clear_nonblock(self.transformer, self._nonblock) + self._provider = None + + # --- phases ----------------------------------------------------------------------------------- + + @property + def open_phase(self) -> str | None: + """The phase running now (``encode``, ``set_load``, ``denoise`` or ``vae``); None outside them.""" + return self._tracker.open_phase + + def _budget(self) -> int: + """The fit budget: the constructor's ``budget_bytes``, else the device's (read at each call).""" + return self._budget_override if self._budget_override is not None else budget_bytes() + + # --- the public surface ----------------------------------------------------------------------- + + @staticmethod + def cfg_prompts( + prompt: str, *, negative_prompt: str | None, guidance: float | None + ) -> 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 + classifier-free guidance with the negative prompt, a blank or missing one replaced by ``" "``. + """ + g = 0.0 if guidance is None else float(guidance) + if g <= 1.0: + return (prompt,) + negative = negative_prompt if negative_prompt and negative_prompt.strip() else " " + return (negative, prompt) + + def _entry(self) -> Any: + """The registry entry of the model this instance runs, or None for a model assembled under another name.""" + return registry.MODELS.get(self._model) + + def _resolve_guidance(self, guidance: float | None) -> float: + """``guidance``, or this model's default when None (mflux's per variant; 1.0 for an unregistered name).""" + if guidance is not None: + return guidance + entry = self._entry() + return ( + 1.0 + if entry is None or entry.default_guidance is None + else float(entry.default_guidance) + ) + + @staticmethod + def _refuse_blank(prompt: str) -> None: + """Refuse an empty or blank prompt as a user error. + + Raises: + DFloatUnsupportedError: ``prompt`` is empty or only whitespace. + """ + if not prompt.strip(): + raise DFloatUnsupportedError( + f"prompt {prompt!r}: an empty or blank prompt is refused as a user error (an empty prompt gives " + "ERNIE-Image's tokenizer no tokens at all)" + ) + + @_pipeline.with_call_caps + def encode( + self, prompt: str, *, negative_prompt: str | None = None, guidance: float | None = None + ) -> None: + """Encode a call's prompts now (drops a resident set first), so several generations pay the encoder once. + + ``guidance`` left as None takes this model's default, as ``generate_image`` does, so ``encode(p)`` prepares + the batch ``generate_image(seed, p)`` runs. Runs under the command's memory caps unless a wired limit is set. + + Raises: + DFloatUnsupportedError: An empty or blank ``prompt``. + """ + self._refuse_blank(prompt) + guidance = self._resolve_guidance(guidance) + self._lifecycle.ensure_embeddings( + self._batch_key( + self.cfg_prompts(prompt, negative_prompt=negative_prompt, guidance=guidance) + ) + ) + + def _encode_prompts( + self, *, prompt: str, negative_prompt: str | None, guidance: float + ) -> tuple[mx.array, mx.array]: + """Mflux's hook, answered from the batch cache (the encoder is not resident during the loop). + + Raises: + DFloatIntegrationError: The call's prompts were not encoded first. + """ + key = self._batch_key( + self.cfg_prompts(prompt, negative_prompt=negative_prompt, guidance=guidance) + ) + if key not in self._batches: + raise DFloatIntegrationError(f"prompts {key} were not encoded before the call") + text_bth, text_lens = self._batches[key] + return text_bth, text_lens + + def _predict( + self, transformer: Any, text_bth: mx.array, text_lens: mx.array, latents: mx.array + ) -> Any: + """Mflux's predict closure, uncompiled on every chip (the seam evaluates at block boundaries). + + Raises: + DFloatIntegrationError: mflux's factory still compiled it. + """ + fn = uncompiled(ErnieImage._predict, transformer, text_bth, text_lens, latents) + self._predict_mode = "uncompiled" + return fn + + def plan_call(self, *, height: int, width: int, text_tokens: int, batch: int) -> CallPlan: + """The cache limit, the fit estimate and the VAE strategy of a call (sizes rounded down to multiples of 16). + + ``text_tokens`` is the call's longest prompt in tokens (``memory.text_tokens_for``); ``batch`` is 2 when the + call runs classifier-free guidance (one batch-2 transformer call), else 1, and selects the constants. The set + stays resident through the VAE decode when the estimate allows it; otherwise the call drops it before decoding + and the next call reloads it. + + Raises: + DFloatResourceError: The size is above 1024² or the predicted peak exceeds the budget even with the set + dropped, and ``fit_check`` is on. + """ + plan = _pipeline.plan_call_for( + constants=self._constants[batch], + sizes=self._sizes, + largest=self._largest, + policy=self._policy, + cache_limit_override=self._cache_limit_override, + fit_check=self._fit_check, + budget=self._budget(), + height=height, + width=width, + text_tokens=text_tokens, + log=log, + ) + self._plan = plan + return plan + + @_pipeline.with_call_caps + def generate_image( + self, + seed: int, + prompt: str, + num_inference_steps: int | None = None, + height: int = 1024, + width: int = 1024, + guidance: float | None = None, + image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str | None = None, + negative_prompt: str | None = None, + pid_decode: bool = False, + pid_degrade_sigma: float = 0.0, + ) -> Any: + """Mflux's ``generate_image`` behind a prelude: refusals, cache limit, fit check, prompt encoding, set load. + + ``guidance``, ``num_inference_steps`` and ``scheduler`` left as None take this model's defaults, mflux's per + variant: ERNIE-Image 4.0 and 50 steps, which run classifier-free guidance; ERNIE-Image-Turbo 1.0 and 8 steps; + ``"linear"`` for both. The same call through mflux's own ``ErnieImage`` class defaults to 1.0 and 8 steps for + both. Above guidance 1.0 a call runs classifier-free guidance as one batch-2 transformer call, with the + negative prompt or, when it is missing or blank, mflux's ``" "``; unlike the ``mlx-dfloat`` command, the + Python API lets ERNIE-Image-Turbo run it too, as mflux's API does. Each call runs under the command's memory + caps unless a wired limit is already set. + + Raises: + DFloatUnsupportedError: An empty or blank ``prompt`` (a user error), ``image_path`` / + ``image_strength`` (img2img) or ``pid_decode``. + 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 step function despite the bypass. + """ + self._refuse_blank(prompt) + if image_path is not None or image_strength is not None: + _pipeline.refuse("image_path/image_strength", "img2img") + if pid_decode: + _pipeline.refuse( + "pid_decode", + "mflux's alternative image decoder (PiD) loads its own caption encoder next to the compressed set", + ) + # mflux's defaults per variant (the registry entry); a model assembled from parts under another name takes + # mflux's class defaults (1.0, 8 steps, "linear"). + entry = self._entry() + guidance = self._resolve_guidance(guidance) + if num_inference_steps is None: + num_inference_steps = 8 if entry is None else entry.default_steps + if scheduler is None: + scheduler = ( + "linear" + if entry is None or entry.default_scheduler is None + else entry.default_scheduler + ) + prompts = self.cfg_prompts(prompt, negative_prompt=negative_prompt, guidance=guidance) + batch = len(prompts) + text_tokens = emem.text_tokens_for(self.tokenizers["ernie"], prompts) + plan = self.plan_call(height=height, width=width, text_tokens=text_tokens, batch=batch) + self._batch, self._text_tokens = batch, text_tokens + self._tracker.begin("encode") + try: + self._lifecycle.ensure_embeddings(self._batch_key(prompts)) + finally: # a failed encode closes its phase: the watchdog's context must not name it afterwards + self._tracker.end("encode") + self._check_encode_peak(plan) + self._tracker.begin("set_load") + try: + self._lifecycle.ensure_set() + finally: + self._tracker.end("set_load") + self._tracker.begin("denoise") # before set_cache_limit: a raise here must change nothing + previous = mx.set_cache_limit(plan.cache_limit) + try: + return super().generate_image( + seed=seed, + prompt=prompt, + num_inference_steps=num_inference_steps, + height=height, + width=width, + guidance=guidance, + image_path=None, + image_strength=None, + scheduler=scheduler, + negative_prompt=negative_prompt, + pid_decode=False, + pid_degrade_sigma=pid_degrade_sigma, + ) + except BaseException as exc: + # The traceback keeps the seam frames, and with them decoded weights or the resident set, alive for as + # long as the caller holds the exception (and past drop_set's own gc.collect()). Clear them for any + # exception, and for the one it was raised while handling. + _clear_frames(exc) + if isinstance(exc, DFloatFormatError): + self._lifecycle.drop_set() # a corrupt block: the retry starts from a clean load + raise + finally: + mx.set_cache_limit(previous) + self._tracker.end("denoise") + self._tracker.end("vae") + + def _check_encode_peak(self, plan: CallPlan) -> None: + """Warn when the encode held more than the estimate's encode phase (language model + encode term) + slack. + + The estimate counts only the encoder's language-model tensors (``memory.encoder_bytes_used``); an mflux change + that also loads the vision tower would make it under-predict the encode phase, and so may a prompt much longer + than the one the encode term was measured on (the warning names the call's prompt length). The measure is + MLX's peak minus what was active when the phase began, so an encode that first dropped a resident set reads + low (no warning), never high. + """ + record = self._tracker.peaks.get("encode") + if record is None: + return + message = emem.encode_peak_warning( + record["mlx_peak"] - record["active_at_start"], + emem.encode_peak_bound(plan.estimate.phases["encode"], self._constants[self._batch]), + text_tokens=self._text_tokens or 0, + ) + if message is not None: + warnings.warn(message, stacklevel=3) + + def report(self) -> dict[str, Any]: + """What this model is and what its calls cost: repositories, policy, limits, estimate, peaks, counts, versions. + + Peaks are sampled at phase boundaries (labelled ``"sampled"``); the fit estimate is a prediction. + ``predict`` is how the last call's step function ran (``"uncompiled"``; None before any call). ``cfg_batch`` + and ``text_tokens`` are the last call's (2 when classifier-free guidance ran, as one batch-2 call; the longest + prompt's tokens). ``sizes["encoders"]`` is the text encoder's language-model bytes, not its file size. + ``eval_policy_status`` is ``"unmeasured"`` for depth2, which no ERNIE-Image run has measured. + """ + launches = self._launches + (self._provider.launches if self._provider is not None else 0) + fit = None if self._plan is None else self._plan.estimate + try: + mflux_version: str | None = metadata.version("mflux") + except metadata.PackageNotFoundError: + mflux_version = None + return { + "family": "ernie", + "model": self._model, + "df11": { + "root": str(self._df11.root), + "repo_id": self._df11.repo_id, + "revision": self._df11.revision, + "config_source": self._ckpt.config_source, + }, + "base": { + "root": str(self._base.root), + "repo_id": self._base.repo_id, + "revision": self._base.revision, + }, + "eval_policy": self._policy, + "eval_policy_status": "measured" if self._policy in MEASURED_POLICIES else "unmeasured", + "cfg_batch": self._batch, + "text_tokens": self._text_tokens, + "predict": self._predict_mode, + "nonblock_groups": sorted(NONBLOCK_GROUPS), + "cache_limit_in_force": None if self._plan is None else self._plan.cache_limit, + "drop_set_before_vae": None if self._plan is None else self._plan.drop_set_before_vae, + "fit": None + if fit is None + else { + "label": "predicted", + "phases": dict(fit.phases), + "peak_phase": fit.peak_phase, + "peak_bytes": fit.peak_bytes, + "budget_bytes": fit.budget_bytes, + "fits": fit.fits, + }, + "sizes": asdict(self._sizes), + "peaks": {"label": "sampled at phase boundaries", **self._tracker.peaks}, + "decode_launches": launches, + "lifecycle": self._lifecycle.counters.as_dict(), + "versions": {"mlx": mx.__version__, "mflux": mflux_version}, # type: ignore[attr-defined] + } + + def save_model(self, base_path: str) -> None: + """Refused: the DFloat11 checkpoint is the saved form; there is nothing of mflux's to write. + + Raises: + DFloatUnsupportedError: Always. + """ + del base_path + _pipeline.refuse("save_model", "the DFloat11 repository is the checkpoint") + + def freeze(self, **kwargs: Any) -> None: + """Freeze the components that are loaded (a dropped encoder is a stand-in, not a module).""" + del kwargs + for module in (self.vae, self.transformer, self.text_encoder): + if not isinstance(module, _EncoderNotResident): + module.freeze() diff --git a/src/mlx_dfloat/mflux/ernie/names.py b/src/mlx_dfloat/mflux/ernie/names.py new file mode 100644 index 0000000..1cc115b --- /dev/null +++ b/src/mlx_dfloat/mflux/ernie/names.py @@ -0,0 +1,89 @@ +"""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 +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. +""" + +from typing import Any + +from mlx_dfloat.errors import DFloatFormatError +from mlx_dfloat.integrate.names import StaticNameMap +from mlx_dfloat.mflux import require_mflux +from mlx_dfloat.mflux._mapping import name_map_from_mapping + +KIND = "layers" +RUN_ORDER: tuple[str, ...] = (KIND,) +MATRIX_SUBS: dict[str, tuple[str, ...]] = { + KIND: ( + "self_attention.to_q", + "self_attention.to_k", + "self_attention.to_v", + "self_attention.to_out.0", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.linear_fc2", + ) +} +NONBLOCK_GROUPS: dict[str, tuple[str, ...]] = { + "time_embedding": ("time_embedding.linear_1.weight", "time_embedding.linear_2.weight"), + "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)`` + + +def ernie_name_map() -> StaticNameMap: + """The map read from mflux's ``ErnieWeightMapping`` (imports mflux). + + Raises: + DFloatDependencyError: The optional ``mflux`` extra is not installed. + DFloatIntegrationError: mflux's mapping no longer shapes the way this adapter expects. + """ + require_mflux() + from mflux.models.ernie_image.weights.ernie_weight_mapping import ErnieWeightMapping + + return name_map_from_mapping( + ErnieWeightMapping.get_transformer_mapping(), matrix_subs=MATRIX_SUBS + ) + + +def check_ernie_groups(ckpt: Any) -> dict[str, int]: + """The block count: block groups contiguous from 0, each holding the seven matrices, plus non-block groups. + + Each non-block group (``time_embedding``, ``adaLN_modulation.1``, ``final_norm.linear``) is optional: a checkpoint + that stores it as BF16 extras loads it as plain weights. mflux builds the block list from its own count, so a hole + would leave a block on placeholders. + + Raises: + DFloatFormatError: A group of another name, a hole in the block indices, no block groups, a block whose + matrices differ from the seven, or a non-block group with other matrices. + """ + expected = MATRIX_SUBS[KIND] + indices: set[int] = set() + for name, group in ckpt.groups.items(): + if name in NONBLOCK_GROUPS: + if tuple(group.matrix_names) != NONBLOCK_GROUPS[name]: + raise DFloatFormatError( + f"{name}: holds {list(group.matrix_names)}, expected {list(NONBLOCK_GROUPS[name])}" + ) + continue + kind, _dot, idx = name.partition(".") + # A zero-padded index ("01") is refused: the seam asks for the group by its canonical name. + if kind != KIND or not (idx.isascii() and idx.isdigit() and idx == str(int(idx))): + raise DFloatFormatError( + f"{name}: not an ERNIE-Image group (blocks: {KIND}.; non-block groups accepted: " + f"{', '.join(NONBLOCK_GROUPS)})" + ) + if tuple(group.matrix_names) != tuple(f"{name}.{sub}.weight" for sub in expected): + raise DFloatFormatError( + f"{name}: holds {len(group.matrix_names)} matrices {list(group.matrix_names)}; an ERNIE-Image block " + f"holds {', '.join(expected)}" + ) + indices.add(int(idx)) + if not indices: + raise DFloatFormatError(f"no {KIND} groups in the checkpoint") + if indices != set(range(len(indices))): + raise DFloatFormatError(f"{KIND} groups are not contiguous: {sorted(indices)}") + return {KIND: len(indices)} diff --git a/src/mlx_dfloat/mflux/ernie/transformer.py b/src/mlx_dfloat/mflux/ernie/transformer.py new file mode 100644 index 0000000..7d86e2f --- /dev/null +++ b/src/mlx_dfloat/mflux/ernie/transformer.py @@ -0,0 +1,156 @@ +"""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`` +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``. + +The build installs placeholders on the seven block matrices and on the three non-block groups' four matrices +(``time_embedding``, ``adaLN_modulation.1``, ``final_norm.linear``; decoded once when the set loads) and loads every +other parameter from the checkpoint's extras (or a BF16 base's), the patch convolution transposed from torch layout. +""" + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from functools import cache +from pathlib import Path +from typing import Any + +import mlx.core as mx +from mlx.utils import tree_flatten + +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.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 +from mlx_dfloat.integrate.resident import NonBlockShapes, install_nonblock_placeholders +from mlx_dfloat.mflux import require_mflux +from mlx_dfloat.mflux.ernie.names import ( + DEFAULT_LAYERS, + KIND, + NONBLOCK_GROUPS, + RUN_ORDER, + check_ernie_groups, + ernie_name_map, +) +from mlx_dfloat.mflux.flux2.transformer import base_transformer_files_index + +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 +# 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() + + +@cache +def seam_transformer_class() -> type: + """``StepSeam`` composed in front of mflux's ``ErnieTransformer`` (imports mflux). + + Raises: + DFloatDependencyError: The optional ``mflux`` extra is not installed. + """ + require_mflux() + from mflux.models.ernie_image.model.ernie_transformer.transformer import ErnieTransformer + + return type("SeamErnieTransformer", (StepSeam, ErnieTransformer), {}) + + +def block_lists(transformer: Any) -> list[tuple[str, Sequence[Any]]]: + """The one block list the forward pass runs, ``layers``.""" + return [(kind, getattr(transformer, kind)) for kind in RUN_ORDER] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ErnieBuild: + """A built transformer, its block shapes, the non-block matrices' shapes and the block count built.""" + + transformer: Any + shapes: Shapes + nonblock: NonBlockShapes + counts: dict[str, int] + + +def build_transformer( + ckpt: DF11Checkpoint, + *, + transformer_overrides: Mapping[str, Any] | None = None, + name_map: NameMap | None = None, + n_layers: int | None = None, + extras: Mapping[str, tuple[Path, TensorInfo]] | None = None, + nonblock_from_extras: bool = False, + transformer_class: Callable[..., Any] | None = None, +) -> ErnieBuild: + """Construct the seamed transformer: block and non-block matrices as placeholders, every other parameter loaded. + + ``transformer_overrides`` are the model's (``ModelConfig.transformer_overrides``); the checkpoint's block count + must equal the depth they build (mflux's default 36 when they leave it out). ``n_layers`` then builds fewer + blocks (a partial build for tests). ``extras`` replaces the checkpoint's own extras (a BF16 base's tensors, via + ``base_extras``); with ``nonblock_from_extras`` the non-block matrices load from them as plain weights instead of + placeholders (the BF16 side of an identity check). A non-block group the checkpoint stores as extras loads as + plain weights either way. The active memory the build added is asserted. + + Raises: + DFloatFormatError: The groups are not ERNIE-Image's; the checkpoint's block count differs from the model's + depth; an extra is not BF16 or has the wrong shape. + DFloatIntegrationError: ``n_layers`` out of range, an uncovered parameter or an extra without one, or too + much active memory added. + DFloatDependencyError: The ``mflux`` extra is missing (default class and map only). + """ + names = ernie_name_map() if name_map is None else name_map + counts = check_ernie_groups(ckpt) + kwargs = dict(transformer_overrides or {}) + want = int(kwargs.get("num_layers", DEFAULT_LAYERS)) + if counts[KIND] != want: + raise DFloatFormatError(f"checkpoint has {counts[KIND]} blocks; this model builds {want}") + n = counts[KIND] if n_layers is None else n_layers + if not 0 <= n <= counts[KIND]: + raise DFloatIntegrationError(f"asked for {n} blocks; the checkpoint has {counts[KIND]}") + make = seam_transformer_class() if transformer_class is None else transformer_class + before = int(mx.get_active_memory()) + tf = make(**{**kwargs, "num_layers": n}) + lists = block_lists(tf) + shapes = install_placeholders(lists, names) + seam_blocks(lists, tf.seam_cell) + nonblock_matrices = ( + [] + if nonblock_from_extras + else [m for g in NONBLOCK_GROUPS if g in ckpt.groups for m in ckpt.groups[g].matrix_names] + ) + nonblock = install_nonblock_placeholders(tf, nonblock_matrices, names) + 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 + check_extras_cover(params, (p for p, _f, _i in plan), matrix_paths) + load_extras(tf, plan, names) + mx.eval(PLACEHOLDER) + added = int(mx.get_active_memory()) - before + if added >= MAX_BUILD_ACTIVE_BYTES: + raise DFloatIntegrationError(f"extras added {added / 1024**3:.2f} GiB of active memory") + return ErnieBuild(transformer=tf, shapes=shapes, nonblock=nonblock, counts=built) + + +def base_extras( + index: Mapping[str, tuple[Path, TensorInfo]], ckpt: DF11Checkpoint +) -> dict[str, tuple[Path, TensorInfo]]: + """The BF16 base's tensors minus the block groups' matrices (the non-block matrices stay, as plain weights).""" + block = { + m for name, g in ckpt.groups.items() if name not in NONBLOCK_GROUPS for m in g.matrix_names + } + return {name: entry for name, entry in index.items() if name not in block} + + +__all__ = [ + "COMPUTED_PARAMS", + "MAX_BUILD_ACTIVE_BYTES", + "ErnieBuild", + "base_extras", + "base_transformer_files_index", + "block_lists", + "build_transformer", + "seam_transformer_class", +] diff --git a/src/mlx_dfloat/mflux/families.py b/src/mlx_dfloat/mflux/families.py index 87308bd..155ae24 100644 --- a/src/mlx_dfloat/mflux/families.py +++ b/src/mlx_dfloat/mflux/families.py @@ -97,6 +97,18 @@ def _qwen21_names() -> NameMap: return qwen21_name_map() +def _ernie_class() -> type: + from mlx_dfloat.mflux.ernie.model import DFloatErnieImage + + return DFloatErnieImage + + +def _ernie_names() -> NameMap: + from mlx_dfloat.mflux.ernie.names import ernie_name_map + + return ernie_name_map() + + FAMILIES: dict[str, FamilySpec] = { "flux1": FamilySpec( name="flux1", @@ -125,6 +137,12 @@ def _qwen21_names() -> NameMap: load_model_class=_qwen21_class, load_name_map=_qwen21_names, ), + "ernie": FamilySpec( + name="ernie", + refused_flags={}, + load_model_class=_ernie_class, + load_name_map=_ernie_names, + ), } @@ -261,6 +279,43 @@ def _klein(name: str, label: str, steps: int, revision: str, *, base: bool) -> M # The snapshot the identity check and the MEASURED row ran on. base_revision="d26bb61231c349cf6b7896fa83353113880e1ba3", ), + # mflux 0.20.0: 8 steps (cli/defaults/defaults.py:36); guidance 1.0 only, any other is an error + # (models/ernie_image/cli/ernie_image_turbo_generate.py:42-45), and --negative-prompt is ignored (:10-12, 40); + # scheduler "linear" (ernie_image.py:64-65). + ModelEntry( + name="ernie-image-turbo", + family="ernie", + label="ERNIE-Image-Turbo", + df11_repo="mingyi456/ERNIE-Image-Turbo-DF11", + df11_revision="27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + base_repo="baidu/ERNIE-Image-Turbo", + default_steps=8, + default_guidance=1.0, + default_scheduler="linear", + cfg_two_calls=False, + uses_negative_prompt=False, + fixed_guidance=1.0, + # The snapshot the parity, de-risk and MEASURED runs used (text encoder, VAE, tokenizer). + base_revision="bc68c81e2a1730a394d5fc9fae70713dee940140", + ), + # mflux 0.20.0: 50 steps (cli/defaults/defaults.py:35); guidance 4.0 (models/ernie_image/cli/ + # ernie_image_generate.py:22, 30-31); scheduler "linear" (:33-34). Classifier-free guidance runs above 1.0 + # with or without a negative prompt, as one batch-2 transformer call (ernie_image.py:164-168, 239-250). + ModelEntry( + name="ernie-image", + family="ernie", + label="ERNIE-Image", + df11_repo="mingyi456/ERNIE-Image-DF11", + df11_revision="c2dd30ad7dd5a928df2309581b282337f7cdb41f", + base_repo="baidu/ERNIE-Image", + default_steps=50, + default_guidance=4.0, + default_scheduler="linear", + cfg_two_calls=False, # CFG runs as one batch-2 call (ernie_image.py:239-250) + uses_negative_prompt=True, + # The snapshot the parity, identity, de-risk and MEASURED runs used. + base_revision="5346b31d68c9c23758ba56ef8be5e9dc174c7f99", + ), ) } @@ -289,7 +344,7 @@ def DFloatModel(name: str, /, **kwargs: Any) -> Any: # noqa: N802 Imports mflux. ``DFloatModel("z-image-turbo")`` returns a ``DFloatZImage``; ``DFloatModel("schnell")`` a ``DFloatFlux1``; ``DFloatModel("flux2-klein-4b")`` a ``DFloatFlux2Klein``; ``DFloatModel("qwen-image-2.1")`` a - ``DFloatQwenImage21``. + ``DFloatQwenImage21``; ``DFloatModel("ernie-image")`` a ``DFloatErnieImage``. Args: name: A registered model name, such as ``"schnell"`` or ``"z-image-turbo"``. diff --git a/src/mlx_dfloat/mflux/generate.py b/src/mlx_dfloat/mflux/generate.py index c6d478f..7d6550d 100644 --- a/src/mlx_dfloat/mflux/generate.py +++ b/src/mlx_dfloat/mflux/generate.py @@ -119,7 +119,7 @@ def _add_arguments(p: argparse.ArgumentParser) -> None: type=int, default=None, help="denoise steps (default per model, as mflux: schnell 4, dev 25, z-image 50, z-image-turbo 9, " - "FLUX.2 Klein base 50, distilled 4, qwen-image-2.1 40)", + "FLUX.2 Klein base 50, distilled 4, qwen-image-2.1 40, ernie-image 50, ernie-image-turbo 8)", ) p.add_argument("--height", type=int, default=1024) p.add_argument("--width", type=int, default=1024) @@ -130,18 +130,19 @@ def _add_arguments(p: argparse.ArgumentParser) -> None: help="guidance (default per model, as mflux: FLUX.1 dev 3.5, ignored by schnell and z-image-turbo; " "z-image 0, its model card suggests about 4; FLUX.2 Klein 1.0, where the base models take another " "value such as 4 and the distilled ones refuse it; Qwen-Image 2.1 1.0, where classifier-free guidance " - "needs a value above 1 and --negative-prompt)", + "needs a value above 1 and --negative-prompt; ERNIE-Image 4.0 (classifier-free guidance above 1, with or " + "without --negative-prompt); ERNIE-Image-Turbo 1.0 only)", ) p.add_argument( "--scheduler", default=None, - help="scheduler (default per model: FLUX.1 and Qwen-Image 2.1 linear; FLUX.2 Klein always runs " + help="scheduler (default per model: FLUX.1, Qwen-Image 2.1 and ERNIE-Image linear; FLUX.2 Klein always runs " "flow_match_euler_discrete and refuses the flag)", ) p.add_argument( "--negative-prompt", default=None, - help="used by z-image and qwen-image-2.1 with --guidance above 1; accepted and ignored for the " + help="used by z-image, qwen-image-2.1 and ernie-image with --guidance above 1; accepted and ignored for the " "other models, as mflux does", ) p.add_argument( @@ -283,6 +284,20 @@ def ceiling_for( return (default_ceiling_bytes if limits.is_host else limits.ceiling_bytes), limits, limits.label +def empty_prompt_refusal(args: argparse.Namespace) -> str | None: + """Why an empty or blank ``--prompt`` is refused for this model, or None. + + ERNIE-Image refuses an empty or blank prompt as a user error (its tokenizer gives an empty prompt no tokens at all, + not even a start token); the other families' mflux pipelines encode an empty prompt as it is. + """ + if entry(args.model).family != "ernie" or args.prompt.strip(): + return None + return ( + f"--prompt {args.prompt!r}: an empty or blank prompt is refused as a user error (an empty prompt gives " + "ERNIE-Image's tokenizer no tokens at all)" + ) + + def _steps(args: argparse.Namespace) -> int: """The denoise steps this call runs: ``--steps``, else the model's mflux default.""" return int(args.steps) if args.steps is not None else entry(args.model).default_steps @@ -386,7 +401,7 @@ def run( between them). """ started = clock() - refused = refused_option(args) or fixed_guidance_refusal(args) + refused = refused_option(args) or fixed_guidance_refusal(args) or empty_prompt_refusal(args) if refused is not None: print(f"error: {refused}", file=sys.stderr) return EXIT_ERROR diff --git a/tests/_ernie_tiny.py b/tests/_ernie_tiny.py new file mode 100644 index 0000000..0bdd5b8 --- /dev/null +++ b/tests/_ernie_tiny.py @@ -0,0 +1,305 @@ +"""Tiny ERNIE-Image transformer helpers for the mflux lane. Every mflux import happens inside a function.""" + +from types import SimpleNamespace + +import mlx.core as mx +import mlx.nn as nn +import numpy as np +from mlx.utils import tree_flatten +from tests._df11_fixtures import random_bf16, write_checkpoint + +__all__ = [ + "PATTERNS", + "TINY", + "TINY_SEED", + "StubTextEncoder", + "StubTokenizer", + "StubVAE", + "fake_model", + "stub_components", + "tiny_groups", + "tiny_inputs", + "tiny_predict", + "tiny_seamed", + "write_tiny_checkpoint", +] + +# Head dim 16 = 32 / 2 heads = the sum of the three RoPE axes (4 + 6 + 6, each even). 128 channels: ErnieLatentCreator +# always makes 128-channel latents (mflux 0.20.0 ernie_latent_creator.py:7). Text width 3072: build_text_batch pads +# every prompt with zeros of width 3072, a literal default (prompt_encoder.py:27). +TINY = { + "hidden_size": 32, + "num_attention_heads": 2, + "num_layers": 2, + "ffn_hidden_size": 48, + "in_channels": 128, + "out_channels": 128, + "patch_size": 1, + "text_in_dim": 3072, + "rope_axes_dim": [4, 6, 6], +} + + +def tiny_seamed(n_layers=2): + """A seamed mflux ``ErnieTransformer`` at ``TINY`` size with placeholders installed; returns (tf, shapes).""" + from mlx_dfloat.integrate.blockseam import seam_blocks + from mlx_dfloat.integrate.placeholders import install_placeholders + from mlx_dfloat.mflux.ernie.names import ernie_name_map + from mlx_dfloat.mflux.ernie.transformer import block_lists, seam_transformer_class + + tf = seam_transformer_class()(**{**TINY, "num_layers": n_layers}) + shapes = install_placeholders(block_lists(tf), ernie_name_map()) + seam_blocks(block_lists(tf), tf.seam_cell) + return tf, shapes + + +def tiny_inputs(batch=1, seed=0): + """A 4 x 4 latent of 128 channels, ``batch`` prompts of 8 tokens at width 3072, and a sigma. + + ``batch`` 2 is a CFG batch (mflux runs it as one batch-2 transformer call, ernie_image.py:239-250). + """ + keys = mx.random.split(mx.random.key(seed), 2) + return { + "hidden_states": mx.random.normal((1, 128, 4, 4), key=keys[0]).astype(mx.bfloat16), + "text_bth": mx.random.normal((batch, 8, 3072), key=keys[1]).astype(mx.bfloat16), + "text_lens": mx.array([8] * batch, dtype=mx.int32), + "sigma": mx.array([0.5]), + } + + +def tiny_predict(tf, inputs, *, compiled=False, guidance=1.0): + """One call of mflux's ``ErnieImage._predict`` on ``tf`` with the chip check answering "not M1/M2". + + Unless ``compiled``, the factory runs through ``_compile.uncompiled`` (the adapter's bypass); with ``compiled`` it is + the stock factory, which compiles off M1/M2 (ernie_image.py:252-254). + """ + from mflux.models.ernie_image.variants.txt2img.ernie_image import ErnieImage + from mflux.utils.apple_silicon import AppleSiliconUtil + + from mlx_dfloat.mflux._compile import uncompiled + + original = AppleSiliconUtil.__dict__["is_m1_or_m2"] + setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: False)) # noqa: B010 + try: + args = (tf, inputs["text_bth"], inputs["text_lens"], inputs["hidden_states"]) + fn = ErnieImage._predict(*args) if compiled else uncompiled(ErnieImage._predict, *args) + finally: + setattr(AppleSiliconUtil, "is_m1_or_m2", original) # noqa: B010 + return fn( + inputs["hidden_states"], + inputs["sigma"], + inputs["text_bth"], + inputs["text_lens"], + guidance, + ) + + +def tiny_groups(shapes, rng): + """One real DF11 group per block of ``shapes``; ERNIE's sub-paths are its attribute paths, so the names match. + + Returns groups, names, sources (``tests/_family_fakes.py:compress_blocks``). + """ + from tests._family_fakes import compress_blocks + + return compress_blocks(shapes, rng) + + +# The DF11 config's pattern_dict (mingyi456/ERNIE-Image{,-Turbo}-DF11 config.json, byte-identical in both). +BLOCK_PATTERN = r"layers\.\d+" +PATTERNS = { + "time_embedding": ("linear_1", "linear_2"), + "adaLN_modulation.1": (), + BLOCK_PATTERN: ( + "self_attention.to_q", + "self_attention.to_k", + "self_attention.to_v", + "self_attention.to_out.0", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.linear_fc2", + ), + "final_norm.linear": (), +} +# mflux parameter -> checkpoint name, the inverse of ErnieWeightMapping's two renames (ernie_weight_mapping.py:12-19). +_CHECKPOINT_NAME = { + "adaln_modulation.weight": "adaLN_modulation.1.weight", + "adaln_modulation.bias": "adaLN_modulation.1.bias", +} +_NONBLOCK_PARAMS = { + "time_embedding": ("time_embedding.linear_1.weight", "time_embedding.linear_2.weight"), + "adaLN_modulation.1": ("adaln_modulation.weight",), + "final_norm.linear": ("final_norm.linear.weight",), +} + + +def write_tiny_checkpoint( + root, rng, *, n_layers=2, nonblock_as_groups=True, hidden_override=None, random_extras=None +): + """A DF11 checkpoint (config.json + safetensors) for mflux's ``ErnieTransformer`` at ``TINY`` size. + + One group per block in ``MATRIX_SUBS`` order; the three non-block groups (or, with ``nonblock_as_groups=False``, + BF16 extras); every other parameter a BF16 extra under its checkpoint name (``adaln_modulation.*`` stored as + ``adaLN_modulation.1.*``), filled with a distinct constant ``0x3F80 + i``, except ``x_embedder.proj.weight``: torch + layout ``(hidden, 128, 1, 1)`` holding the bit patterns ``0x5000 + flat index`` (every value distinct, so any axis + mix-up of the transpose shows). ``hidden_override`` builds the probe that wide. ``random_extras`` (a NumPy generator) + fills every extra, the patch conv included, with random BF16 instead: constant weights make the forward pass blind + to its inputs, which an end-to-end comparison cannot use; the constants' values are then the uint16 arrays in mflux's + layout (the conv transposed to (hidden, 1, 1, 128)). + + Returns ``(matrices, constants, conv_bits)``: ``{group: {sub-path or parameter: uint16 matrix}}``, + ``{mflux parameter: constant}`` and the uint16 conv array written. + """ + from mflux.models.ernie_image.model.ernie_transformer.transformer import ErnieTransformer + + from mlx_dfloat.integrate.placeholders import install_placeholders + from mlx_dfloat.mflux.ernie.names import ernie_name_map + from mlx_dfloat.mflux.ernie.transformer import block_lists + + dims = {**TINY, "num_layers": n_layers} + if hidden_override is not None: + dims["hidden_size"] = hidden_override + probe = ErnieTransformer(**dims) + probe_params = dict(tree_flatten(probe.parameters())) + shapes = install_placeholders(block_lists(probe), ernie_name_map()) + matrices, groups = {}, {} + for block, per in shapes.items(): + matrices[block] = {attr: random_bf16(rng, shape) for attr, shape in per.items()} + groups[block] = [matrices[block][sub] for sub in PATTERNS[BLOCK_PATTERN]] + taken = {f"{b}.{a}.weight" for b, per in shapes.items() for a in per} + patterns = {BLOCK_PATTERN: PATTERNS[BLOCK_PATTERN]} + if nonblock_as_groups: + for group, params in _NONBLOCK_PARAMS.items(): + matrices[group] = {p: random_bf16(rng, tuple(probe_params[p].shape)) for p in params} + groups[group] = [matrices[group][p] for p in params] + patterns[group] = PATTERNS[group] + taken |= set(params) + hidden = dims["hidden_size"] + conv_bits = (0x5000 + np.arange(hidden * 128, dtype=np.uint16)).reshape(hidden, 128, 1, 1) + constants, extras = {}, {} + for i, (param, array) in enumerate(sorted(probe_params.items())): + if param in taken: + continue + if random_extras is not None: + if param == "x_embedder.proj.weight": + torch_layout = random_bf16(random_extras, (hidden, 128, 1, 1)) + extras[param] = torch_layout + constants[param] = torch_layout.transpose(0, 2, 3, 1) + else: + bits = random_bf16(random_extras, tuple(array.shape)) + extras[_CHECKPOINT_NAME.get(param, param)] = constants[param] = bits + continue + if param == "x_embedder.proj.weight": + extras[param] = conv_bits + continue + constants[param] = 0x3F80 + i + extras[_CHECKPOINT_NAME.get(param, param)] = np.full( + array.shape, 0x3F80 + i, dtype=np.uint16 + ) + write_checkpoint(root, groups=groups, patterns=patterns, extras=extras) + return matrices, constants, conv_bits + + +TINY_SEED = 5 # ``fake_model`` writes its checkpoint with this seed: a test regenerates the source matrices from it + + +class StubTokenizer: + """mflux's tokenizer protocol as ``ErniePromptEncoder`` uses it (mflux 0.20.0 prompt_encoder.py:15-17). + + ``tokenize(prompt, max_length=None)`` gives ids that depend on the prompt (``1 + ord(c) % 7`` per character, a + blank prompt one ``"_"``) and an all-ones attention mask; ``lengths`` overrides a prompt's token count (the ids then + cycle through its characters). + """ + + def __init__(self, lengths=None): + self.lengths = dict(lengths or {}) + + def tokenize(self, prompt, max_length=None): + del max_length + chars = prompt or "_" + n = self.lengths.get(prompt, len(chars)) + ids = [1 + ord(chars[i % len(chars)]) % 7 for i in range(n)] + return SimpleNamespace( + input_ids=mx.array([ids], dtype=mx.int32), + attention_mask=mx.ones((1, n), dtype=mx.int32), + ) + + +class StubTextEncoder(nn.Module): + """Stands in for the Mistral text stack: ``__call__(input_ids, attention_mask)`` -> (1, L, 3072) bf16 hidden states, + each position the weight times its id, so a wrong or swapped prompt shows. + + 3072 is the width ``build_text_batch`` pads with (prompt_encoder.py:27); mflux calls it positionally (:16). + """ + + def __init__(self): + super().__init__() + self.weight = mx.full((1, 3072), 0.01, dtype=mx.bfloat16) + self.calls = 0 + + def __call__(self, input_ids, attention_mask=None): + self.calls += 1 + return self.weight[None] * input_ids[..., None].astype(mx.bfloat16) + + +class StubVAE(nn.Module): + """One weight (so the retained bound can size it) and the decode ``ErnieImage._decode_latents`` calls + (ernie_image.py:193): a zero image of 16 pixels per latent cell.""" + + def __init__(self): + super().__init__() + self.weight = mx.zeros((4,), dtype=mx.bfloat16) + + def decode_packed_latents(self, latents, tiling_config=None): + del tiling_config + return mx.zeros((1, 3, 16 * latents.shape[-2], 16 * latents.shape[-1])) + + +def stub_components(tokenizer=None): + """Base components as ``ErnieComponents`` with the stub encoder, VAE and tokenizer.""" + from mlx_dfloat.mflux.ernie.init import ErnieComponents + + return ErnieComponents( + vae=StubVAE(), + text_encoder=StubTextEncoder(), + tokenizers={"ernie": tokenizer or StubTokenizer()}, + ) + + +def fake_model(tmp_path, monkeypatch, *, model="ernie-image-turbo", tokenizer=None, **overrides): + """A ``DFloatErnieImage`` over the tiny real transformer (2 blocks, config.json checkpoint) and stub base parts. + + No weights, no network. As the other families' fakes: the budget is fixed at 23 GiB, the encoder reload builds a + fresh stub, every decode runs on the NumPy reference. + """ + from functools import partial + + from mflux.models.common.config.model_config import ModelConfig + + from mlx_dfloat.decode import decode_group + from mlx_dfloat.format import open_checkpoint + from mlx_dfloat.mflux._hub import ResolvedRepo + from mlx_dfloat.mflux._phases import FamilySizes + from mlx_dfloat.mflux.ernie import init as einit + from mlx_dfloat.mflux.ernie import model as model_module + from mlx_dfloat.mflux.ernie.transformer import build_transformer + + write_tiny_checkpoint(tmp_path / "df11", np.random.default_rng(TINY_SEED)) + ckpt = open_checkpoint(tmp_path / "df11") + build = build_transformer(ckpt, transformer_overrides=TINY) + (tmp_path / "base").mkdir(exist_ok=True) + monkeypatch.setattr(einit, "load_text_encoder", lambda root, **kwargs: StubTextEncoder()) + # A fixed budget: the CI runner's device reports a 7 GiB working set; the tests are about the arithmetic. + monkeypatch.setattr(model_module, "budget_bytes", lambda: 23 * 1024**3) + parts = { + "model": model, + "model_config": ModelConfig.from_name(model_name=model, base_model=None), + "ckpt": ckpt, + "df11": ResolvedRepo(root=tmp_path / "df11", repo_id="mingyi456/tiny", revision="a" * 40), + "base": ResolvedRepo(root=tmp_path / "base", repo_id="baidu/tiny", revision="b" * 40), + "components": stub_components(tokenizer), + "build": build, + "sizes": FamilySizes(compressed=1_000, extras=100, nonblock=10, encoders=1_000, vae=10), + "decode": partial(decode_group, backend="reference"), + } + parts.update(overrides) + return model_module.DFloatErnieImage._from_parts(**parts) diff --git a/tests/test_cli.py b/tests/test_cli.py index 82d2688..5ee3f44 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1271,3 +1271,109 @@ def test_qwen_refuses_the_common_flags_before_anything_loads(flag, tmp_path, cap ) assert (code, watchdogs, log) == (2, [], []) assert flag[0] in capsys.readouterr().err + + +# --- ERNIE-Image ------------------------------------------------------------------------------------------------ + + +def test_ernie_steps_follow_mflux_per_variant(): + # Bug caught: an ERNIE model on another model's steps (mflux 0.20.0 cli/defaults/defaults.py:35-36: ernie-image 50, + # ernie-image-turbo 8). + steps = { + m: gen._steps(gen.build_parser().parse_args(["--model", m, "--prompt", "p"])) + for m in ("ernie-image", "ernie-image-turbo") + } + assert steps == {"ernie-image": 50, "ernie-image-turbo": 8} + + +def test_turbo_guidance_other_than_1_is_refused_before_the_build(tmp_path, capsys): + # Bug caught (Review Focus 1, command half): --guidance 4 reaching ERNIE-Image-Turbo (mflux's turbo command errors, + # ernie_image_turbo_generate.py:42-45), or the refusal coming after the caps or the build. + built = [] + args = gen.build_parser().parse_args( + [ + "--model", + "ernie-image-turbo", + "--prompt", + "p", + "--guidance", + "4", + "--output", + str(tmp_path / "o.png"), + ] + ) + code = gen.run( + args, + model_factory=lambda **kw: built.append(kw), + install_caps=lambda: pytest.fail("caps installed before the refusal"), + watchdog_factory=lambda *a, **k: pytest.fail("watchdog started before the refusal"), + resolve_output=_resolve, + **_host_kwargs(), + ) + assert (code, built) == (2, []) + assert ( + "error: --guidance: ERNIE-Image-Turbo is a distilled model and runs at guidance 1.0 only" + in capsys.readouterr().err + ) + + +def test_turbo_negative_prompt_warns_it_is_ignored(tmp_path, capsys): + # Bug caught: a negative_prompt keyword sent to Turbo (mflux's turbo command ignores it with a warning, + # ernie_image_turbo_generate.py:10-12, 40), or the warning missing. + call = _generate_kwargs("ernie-image-turbo", tmp_path, "--negative-prompt", "blurry") + assert "negative_prompt" not in call + assert (call["guidance"], call["num_inference_steps"], call["scheduler"]) == (1.0, 8, "linear") + assert ( + "warning: --negative-prompt is ignored: this model has no negative branch" + in capsys.readouterr().err + ) + + +@pytest.mark.parametrize( + ("extra", "negative"), [([], None), (["--negative-prompt", "blurry"], "blurry")] +) +def test_ernie_base_passes_guidance_4_and_the_negative_prompt_by_default( + tmp_path, capsys, extra, negative +): + # Bug caught: the base run at mflux's class default 1.0 (no CFG: not mflux's ernie-image command, + # ernie_image_generate.py:22, 30-31), the negative prompt dropped, or a Qwen-style "no CFG without + # --negative-prompt" warning (ERNIE runs CFG with mflux's " " negative). + call = _generate_kwargs("ernie-image", tmp_path, *extra) + assert (call["guidance"], call["num_inference_steps"], call["scheduler"]) == (4.0, 50, "linear") + assert call["negative_prompt"] == negative + assert "warning" not in capsys.readouterr().err + + +def test_ernie_base_negative_prompt_at_guidance_1_warns_with_its_default_4(tmp_path, capsys): + # Bug caught: the "no effect" warning quoting another model's default (Z-Image's 0, Qwen's 1.0). + _generate_kwargs("ernie-image", tmp_path, "--guidance", "1", "--negative-prompt", "blurry") + err = capsys.readouterr().err + assert "warning: --negative-prompt has no effect" in err + assert "(the default is 4.0)" in err + + +@pytest.mark.parametrize("model", ["ernie-image", "ernie-image-turbo"]) +@pytest.mark.parametrize("prompt", ["", " "]) +def test_an_empty_ernie_prompt_is_refused_before_anything_loads(model, prompt, tmp_path, capsys): + # Bug caught: an empty or blank prompt reaching ERNIE-Image (its tokenizer gives "" no tokens: a zero-length text + # through the transformer), or the refusal coming after the caps, the watchdog or the build. + log = [] + code, watchdogs = _run( + ["--model", model, "--prompt", prompt, "--output", str(tmp_path / "o.png")], + log, + tmp_path, + factory=lambda **kw: pytest.fail("built"), + install_caps=lambda: pytest.fail("caps installed before the refusal"), + ) + assert (code, watchdogs, log) == (2, [], []) + assert "an empty or blank prompt is refused as a user error" in capsys.readouterr().err + + +def test_an_empty_prompt_is_still_passed_to_the_other_families(tmp_path): + # Bug caught: the ERNIE refusal applied to every model (mflux encodes an empty prompt for the others). + log = [] + code, _ = _run( + ["--model", "schnell", "--prompt", "", "--output", str(tmp_path / "o.png")], log, tmp_path + ) + assert code == 0 + assert log[0][1]["prompt"] == "" diff --git a/tests/test_ernie_build.py b/tests/test_ernie_build.py new file mode 100644 index 0000000..708d13c --- /dev/null +++ b/tests/test_ernie_build.py @@ -0,0 +1,197 @@ +"""The ERNIE-Image build from a DF11 checkpoint: base extras offline, the real tiny build (mflux lane).""" + +from functools import partial +from pathlib import Path +from types import SimpleNamespace + +import mlx.core as mx +import numpy as np +import pytest +from mlx.utils import tree_flatten +from tests._ernie_tiny import TINY, write_tiny_checkpoint + +from mlx_dfloat.decode import decode_group +from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError +from mlx_dfloat.format import open_checkpoint +from mlx_dfloat.mflux.ernie.transformer import base_extras, build_transformer + +# The seven block matrices at TINY size (hidden 32, ffn 48), weights (out, in): ErnieAttention's four Linear(32, 32) +# (mflux 0.20.0 attention.py:15-25), ErnieFeedForward's gate_proj/up_proj Linear(32, 48) and linear_fc2 Linear(48, 32) +# (feed_forward.py:8-13). +SEVEN = { + "self_attention.to_q": (32, 32), + "self_attention.to_k": (32, 32), + "self_attention.to_v": (32, 32), + "self_attention.to_out.0": (32, 32), + "mlp.gate_proj": (48, 32), + "mlp.up_proj": (48, 32), + "mlp.linear_fc2": (32, 48), +} +# The four non-block matrices at hidden 32 (transformer.py:36-76): time_embedding.linear_{1,2} Linear(32, 32), +# adaln_modulation Linear(32, 6 x 32), final_norm.linear Linear(32, 2 x 32). +NONBLOCK = { + "time_embedding.linear_1.weight": (32, 32), + "time_embedding.linear_2.weight": (32, 32), + "adaln_modulation.weight": (192, 32), + "final_norm.linear.weight": (64, 32), +} + + +def test_base_extras_drop_the_block_matrices_and_keep_the_nonblock_ones(): + # Bug caught: the norms dropped with the block (an extras-coverage refusal on the BF16 side) or the block matrices + # kept as extras (the BF16 side would load them as plain weights next to its decoded ones), or the non-block + # matrices dropped (the BF16 side loads them as plain weights). + groups = { + "layers.0": SimpleNamespace( + matrix_names=("layers.0.mlp.up_proj.weight", "layers.0.self_attention.to_q.weight") + ), + "adaLN_modulation.1": SimpleNamespace(matrix_names=("adaLN_modulation.1.weight",)), + } + index = { + name: (Path(name), None) + for name in ( + "layers.0.mlp.up_proj.weight", + "layers.0.adaLN_sa_ln.weight", + "adaLN_modulation.1.weight", + "x_embedder.proj.weight", + ) + } + assert sorted(base_extras(index, SimpleNamespace(groups=groups))) == [ + "adaLN_modulation.1.weight", + "layers.0.adaLN_sa_ln.weight", + "x_embedder.proj.weight", + ] + + +def _open(tmp_path, **kwargs): + root = tmp_path / "df11" + matrices, constants, conv = write_tiny_checkpoint(root, np.random.default_rng(3), **kwargs) + return open_checkpoint(root), matrices, constants, conv + + +def _params(build): + return dict(tree_flatten(build.transformer.parameters())) + + +@pytest.mark.mflux +def test_the_build_installs_placeholders_and_loads_every_extra_by_name(tmp_path): + # Bug caught: the modulation rename lost (an extras-coverage refusal), a compressed matrix left live (memory), or + # an extra loaded onto the wrong parameter (a constant in the wrong place). + ckpt, _m, constants, _conv = _open(tmp_path) + build = build_transformer(ckpt, transformer_overrides=TINY) + params = _params(build) + assert build.nonblock == NONBLOCK + for name in NONBLOCK: + assert params[name].size == 0, name + assert build.shapes == {f"layers.{i}": SEVEN for i in range(2)} + for block in build.shapes: + for attr in SEVEN: + assert params[f"{block}.{attr}.weight"].size == 0, (block, attr) + # Every parameter outside the compressed groups and the patch conv (mflux 0.20.0 ErnieTransformer at 2 layers). + assert sorted(constants) == [ + "adaln_modulation.bias", + "final_linear.bias", + "final_linear.weight", + "final_norm.linear.bias", + "layers.0.adaLN_mlp_ln.weight", + "layers.0.adaLN_sa_ln.weight", + "layers.0.self_attention.norm_k.weight", + "layers.0.self_attention.norm_q.weight", + "layers.1.adaLN_mlp_ln.weight", + "layers.1.adaLN_sa_ln.weight", + "layers.1.self_attention.norm_k.weight", + "layers.1.self_attention.norm_q.weight", + "text_proj.weight", + "time_embedding.linear_1.bias", + "time_embedding.linear_2.bias", + "x_embedder.proj.bias", + ] + for name, value in constants.items(): + assert np.all(np.array(params[name].view(mx.uint16)) == value), name + assert build.counts == {"layers": 2} + + +@pytest.mark.mflux +def test_the_patch_conv_is_transposed_from_torch_layout(tmp_path): + # Bug caught: the transform dropped (a shape refusal) or applied on the wrong axes ((0, 3, 2, 1) lands (3, 5) + # elsewhere; every value is distinct, so no axis mix-up can match). Bits by hand: 0x5000 + o * 128 + i. + ckpt, _m, _c, _conv = _open(tmp_path) + weight = _params(build_transformer(ckpt, transformer_overrides=TINY))["x_embedder.proj.weight"] + assert weight.shape == (32, 1, 1, 128) + for (o, i), bits in {(0, 0): 20480, (3, 5): 20869, (31, 127): 24575}.items(): + assert int(weight[o, 0, 0, i].view(mx.uint16).item()) == bits, (o, i) + + +@pytest.mark.mflux +def test_nonblock_groups_stored_as_extras_load_as_plain_weights(tmp_path): + # Bug caught: a checkpoint without the non-block groups refused, or its weights left on placeholders (a zero-size + # matmul before the first block). + ckpt, _m, constants, _conv = _open(tmp_path, nonblock_as_groups=False) + assert "adaLN_modulation.1.weight" in ckpt.extras + build = build_transformer(ckpt, transformer_overrides=TINY) + weight = _params(build)["adaln_modulation.weight"] + assert build.nonblock == {} + assert weight.shape == (192, 32) + assert np.all(np.array(weight.view(mx.uint16)) == constants["adaln_modulation.weight"]) + + +@pytest.mark.mflux +def test_a_checkpoint_with_other_block_counts_is_refused_naming_both(tmp_path): + # Bug caught: a 3-block checkpoint built as 2 (a shape error at block 2, or the block silently dropped). + ckpt, _m, _c, _conv = _open(tmp_path, n_layers=3) + with pytest.raises(DFloatFormatError, match=r"checkpoint has 3 blocks; this model builds 2\b"): + build_transformer(ckpt, transformer_overrides=TINY) + + +@pytest.mark.mflux +def test_the_count_check_defaults_to_mfluxs_depth(tmp_path): + # Bug caught: overrides without num_layers checked against something other than mflux's default 36 + # (transformer.py:41): the published model's overrides carry only rope_axes_dim. + ckpt, _m, _c, _conv = _open(tmp_path) + no_depth = {k: v for k, v in TINY.items() if k != "num_layers"} + with pytest.raises(DFloatFormatError, match=r"checkpoint has 2 blocks; this model builds 36\b"): + build_transformer(ckpt, transformer_overrides=no_depth) + + +@pytest.mark.mflux +def test_n_layers_builds_fewer_blocks_and_is_range_checked(tmp_path): + # Bug caught: a partial build loading block 1's norms into a transformer without block 1 (an extra without a + # parameter), or a request beyond the checkpoint built from nothing. + ckpt, _m, _c, _conv = _open(tmp_path) + build = build_transformer(ckpt, transformer_overrides=TINY, n_layers=1) + assert build.counts == {"layers": 1} + assert len(build.transformer.layers) == 1 + with pytest.raises(DFloatIntegrationError, match=r"asked for 3 blocks; the checkpoint has 2"): + build_transformer(ckpt, transformer_overrides=TINY, n_layers=3) + + +@pytest.mark.mflux +def test_a_checkpoint_of_another_width_is_refused_by_the_extras_shape_check(tmp_path): + # Bug caught: the same block count at another hidden width (a hand-made checkpoint) accepted. + ckpt, _m, _c, _conv = _open(tmp_path, hidden_override=16) + with pytest.raises(DFloatFormatError, match="shape"): + build_transformer(ckpt, transformer_overrides=TINY) + + +@pytest.mark.mflux +def test_the_decoded_block_matrices_land_in_pattern_order(tmp_path): + # Bug caught: the order swapped between the checkpoint, the reader and the map (gate_proj and up_proj have the + # same shape, so only their values show it: the block would run up(x) * gelu(gate(x)) on the wrong halves). + from mlx_dfloat.integrate.coverage import load_resident_set + from mlx_dfloat.integrate.providers import DF11Provider + from mlx_dfloat.mflux.ernie.names import NONBLOCK_GROUPS, ernie_name_map + + ckpt, matrices, _c, _conv = _open(tmp_path) + build = build_transformer(ckpt, transformer_overrides=TINY) + blocks = [g for g in ckpt.groups if g not in NONBLOCK_GROUPS] + provider = DF11Provider( + load_resident_set(ckpt, blocks), + {g: ckpt.groups[g].matrix_names for g in blocks}, + ernie_name_map(), + decode=partial(decode_group, backend="reference"), + ) + got = provider.weights_for("layers.1", build.shapes["layers.1"]) + provider.verify() + for attr in SEVEN: + want = matrices["layers.1"][attr] + assert np.array_equal(np.array(got[attr].view(mx.uint16)), want), attr diff --git a/tests/test_ernie_init.py b/tests/test_ernie_init.py new file mode 100644 index 0000000..ff727a6 --- /dev/null +++ b/tests/test_ernie_init.py @@ -0,0 +1,339 @@ +"""ERNIE-Image base components (text encoder, VAE, tokenizer) without the transformer or the prompt enhancer.""" + +import os +import re +from fnmatch import fnmatch +from pathlib import Path + +import mlx.core as mx +import pytest +from mlx.utils import tree_flatten + +from mlx_dfloat.errors import DFloatFormatError +from mlx_dfloat.mflux import _hub +from mlx_dfloat.mflux.ernie.init import BASE_PATTERNS, DF11_PATTERNS, TOKENIZER_PATTERNS + +# baidu/ERNIE-Image @ 5346b31d68c9c23758ba56ef8be5e9dc174c7f99, every file on the Hub (listed 2026-10-08; the same +# paths as baidu/ERNIE-Image-Turbo @ bc68c81e). +BASE_FILES = ( + ".gitattributes", + "LICENSE", + "README.md", + "model_index.json", + "pe/chat_template.jinja", + "pe/config.json", + "pe/generation_config.json", + "pe/model.safetensors", + "pe/tokenizer.json", + "pe/tokenizer_config.json", + "pe_tokenizer/chat_template.jinja", + "pe_tokenizer/tokenizer.json", + "pe_tokenizer/tokenizer_config.json", + "scheduler/scheduler_config.json", + "text_encoder/config.json", + "text_encoder/model.safetensors", + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "transformer/config.json", + "transformer/diffusion_pytorch_model-00001-of-00002.safetensors", + "transformer/diffusion_pytorch_model-00002-of-00002.safetensors", + "transformer/diffusion_pytorch_model.safetensors.index.json", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", +) +# The de-risk prompt and its token count (mflux 0.20.0 TokenizerLoader on both base snapshots, 2026-10-08). +PROMPT = ( + "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the " + "water" +) +PROMPT_TOKENS = 26 + + +def _hits(files): + return sorted(f for f in files if any(fnmatch(f, p) for p in BASE_PATTERNS)) + + +def test_base_patterns_fetch_the_encoder_vae_and_tokenizer_and_never_the_transformer_or_the_prompt_enhancer(): + # Bug caught: a pattern pulling the 16 GB BF16 transformer, the 7.7 GB prompt enhancer (pe/) or its tokenizer + # (pe_tokenizer/, whose tokenizer.json a loose "*tokenizer*" pattern would match), or one that misses a file of + # the encoder, the VAE or the tokenizer. + assert _hits(BASE_FILES) == [ + "text_encoder/config.json", + "text_encoder/model.safetensors", + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", + ] + + +def test_the_tokenizer_pattern_fetches_tokenizer_files_only(): + # Bug caught: the tokenizer pattern widened back to the directory ("tokenizer/*" or mflux's "tokenizer/**"), + # fetching whatever a later revision drops there (binaries, nested copies). The listing is the snapshot's two + # files plus files a later revision could add. + listing = ( + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "tokenizer/extra/tokenizer.json", + "tokenizer/model.bin", + "tokenizer/README.md", + "tokenizer/special_tokens_map.json", + ) + assert _hits(listing) == [ + "tokenizer/special_tokens_map.json", + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + ] + assert all(not p.endswith("*") for p in TOKENIZER_PATTERNS) + + +def test_df11_patterns_fetch_the_checkpoint_file_and_its_config(): + # Bug caught: the DF11 repository resolved without its config (the reader could not split the groups) or with + # its README. mingyi456/ERNIE-Image-DF11 @ c2dd30ad holds these four files (Hub, 2026-10-08). + files = (".gitattributes", "README.md", "config.json", "model.safetensors") + assert [f for f in files if any(fnmatch(f, p) for p in DF11_PATTERNS)] == [ + "config.json", + "model.safetensors", + ] + + +# --- mflux lane ----------------------------------------------------------------------------------------------- + + +@pytest.mark.mflux +def test_the_definitions_name_only_their_component_and_the_tokenizer_subdir(): + # Bug caught: a subset keeping the transformer component (mflux's loader would read transformer/), the encoder + # losing its language-model prefix filter (the vision tower would load into nothing), or the tokenizer looked up + # anywhere but tokenizer/ (mflux 0.20.0 ernie_weight_definition.py:26-47). + from mlx_dfloat.mflux.ernie.init import encoder_definition, vae_definition + + (encoder,) = encoder_definition().get_components() + assert (encoder.name, encoder.weight_prefix_filters) == ( + "text_encoder", + ["language_model.model."], + ) + assert [c.name for c in vae_definition().get_components()] == ["vae"] + (tokenizer,) = encoder_definition().get_tokenizers() + assert (tokenizer.name, tokenizer.hf_subdir, tokenizer.max_length) == ( + "ernie", + "tokenizer", + 2048, + ) + + +# A two-layer Mistral text stack (tiny widths; mflux 0.20.0 ErnieMistralTextEncoder's constructor arguments). +TINY_ENCODER = { + "vocab_size": 100, + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 32, + "intermediate_size": 128, +} + + +@pytest.mark.mflux +def test_the_computed_exemptions_are_exactly_what_the_base_files_do_not_hold(): + # Bug caught: a computed buffer missing from COMPUTED (every real base refused), or a loaded weight listed there (a + # missing tensor never noticed). mflux 0.20.0: the encoder's RoPE frequencies are the underscore attribute + # _inv_freq (text_encoder.py:137-138), not a parameter, so every encoder parameter sits under the + # language_model.model. prefix the base file holds; the VAE maps every parameter (flux2_weight_mapping.py). + from mflux.models.ernie_image.model.ernie_text_encoder.text_encoder import ( + ErnieMistralTextEncoder, + ) + from mflux.models.flux2.model.flux2_vae.vae import Flux2VAE + from mflux.models.flux2.weights.flux2_weight_mapping import Flux2WeightMapping + + from mlx_dfloat.mflux.ernie.init import COMPUTED + + encoder = {n for n, _ in tree_flatten(ErnieMistralTextEncoder(**TINY_ENCODER).parameters())} + assert {n for n in encoder if not n.startswith("language_model.model.")} == set() + # The VAE targets are templated ("...resnets.{i}.conv1.weight"): an index token stands for any number. + targets = [ + re.compile(re.sub(r"\\\{[a-z_]+\\\}", r"[0-9]+", re.escape(t.to_pattern)) + "$") + for t in Flux2WeightMapping.get_vae_mapping() + ] + vae = {n for n, _ in tree_flatten(Flux2VAE().parameters())} + assert {n for n in vae if not any(t.match(n) for t in targets)} == set() + assert sorted(COMPUTED) == ["text_encoder", "vae"] + assert [len(v) for v in COMPUTED.values()] == [0, 0] + + +def _write_tiny_encoder(root, *, drop=()): + """The tiny encoder under the base file's names (mflux's paths: ``language_model.model.*``), plus a vision tower + and a projector tensor as the real file carries (2026-10-08 header: 218 and 4 tensors), which mflux's prefix filter + leaves out.""" + from mflux.models.ernie_image.model.ernie_text_encoder.text_encoder import ( + ErnieMistralTextEncoder, + ) + + params = dict(tree_flatten(ErnieMistralTextEncoder(**TINY_ENCODER).parameters())) + flat = {n: mx.full(p.shape, 0.5, dtype=mx.bfloat16) for n, p in params.items() if n not in drop} + flat["vision_tower.transformer.layers.0.attention.q_proj.weight"] = mx.zeros( + (4, 4), dtype=mx.bfloat16 + ) + flat["multi_modal_projector.linear_1.weight"] = mx.zeros((4, 4), dtype=mx.bfloat16) + (root / "text_encoder").mkdir(parents=True, exist_ok=True) + mx.save_safetensors(str(root / "text_encoder" / "model.safetensors"), flat) + + +@pytest.mark.mflux +def test_an_encoder_tensor_missing_from_the_base_is_refused_by_name(tmp_path): + # Bug caught (Review Focus 5): mflux's applier updates with strict=False, so a base missing one tensor (a lost + # shard, a renamed key) loads without error and that layer keeps its random float32 init: a wrong image, no error. + from mlx_dfloat.mflux.ernie.init import load_text_encoder + + missing = "language_model.model.layers.1.self_attn.q_proj.weight" + _write_tiny_encoder(tmp_path, drop={missing}) + with pytest.raises(DFloatFormatError, match=re.escape(missing) + ".*not in the base"): + load_text_encoder(tmp_path, dims=TINY_ENCODER) + + +@pytest.mark.mflux +def test_a_complete_encoder_loads_from_the_language_model_tensors_only(tmp_path): + # Bug caught: the vision tower or the projector in the same file breaking the load, or the file's weights not + # landing (still the float32 init). + from mlx_dfloat.mflux.ernie.init import load_text_encoder + + _write_tiny_encoder(tmp_path) + encoder = load_text_encoder(tmp_path, dims=TINY_ENCODER) + q = encoder.language_model.model.layers[1].self_attn.q_proj.weight + assert q.dtype == mx.bfloat16 + assert float(q[0, 0]) == 0.5 + + +@pytest.mark.mflux +def test_a_base_with_another_encoder_width_is_refused_before_any_encode(monkeypatch, tmp_path): + # Bug caught: a wrong base (a Qwen3-4B-wide embedding) loading silently: mflux's applier assigns weights without a + # shape check, so the failure would be a matmul error at the first encode, not a refusal. + from mlx_dfloat.mflux.ernie.init import load_text_encoder + + def wrong_width(root, definition, models, *, what): + model = models["text_encoder"].language_model.model + model.embed_tokens.weight = mx.zeros((151936, 2560), dtype=mx.bfloat16) + + monkeypatch.setattr(_hub, "load_into", wrong_width) + with pytest.raises( + DFloatFormatError, match=r"embed_tokens\.weight has shape \(151936, 2560\).*\(100, 64\)" + ): + load_text_encoder(tmp_path, dims=TINY_ENCODER) + + +def _write_vae(root, *, drop=()): + """A VAE file under the base's diffusers names, one-element tensors (the applier checks no shape). + + ERNIE's VAE is FLUX.2's (mflux 0.20.0 ernie_weight_mapping.py:111-113): mflux's parameter paths, except that the + two mid-block attentions' output projection is ``to_out.0`` on disk (flux2_weight_mapping.py:319-325, 369-375), + plus ``bn.num_batches_tracked``, which mflux does not map. + """ + from mflux.models.flux2.model.flux2_vae.vae import Flux2VAE + + names = [n.replace(".to_out.", ".to_out.0.") for n, _ in tree_flatten(Flux2VAE().parameters())] + names.append("bn.num_batches_tracked") + flat = {n: mx.zeros((1,), dtype=mx.bfloat16) for n in names if n not in drop} + (root / "vae").mkdir(parents=True, exist_ok=True) + mx.save_safetensors(str(root / "vae" / "diffusion_pytorch_model.safetensors"), flat) + + +@pytest.mark.mflux +@pytest.mark.parametrize("dropped", ["decoder.conv_out.weight", "bn.running_mean"]) +def test_a_vae_tensor_missing_from_the_base_is_refused_by_name(tmp_path, dropped): + # Bug caught (Review Focus 5): the VAE loaded through the same strict=False applier with no coverage check (a + # lost decoder conv keeps its random init and the decode returns noise), or the batch-norm statistics exempted as + # computed buffers (they come from the file: flux2_weight_mapping.py:202-209). + from mlx_dfloat.mflux.ernie.init import load_vae + + _write_vae(tmp_path, drop={dropped}) + with pytest.raises(DFloatFormatError, match=re.escape(dropped) + ".*not in the base"): + load_vae(tmp_path) + + +@pytest.mark.mflux +def test_a_complete_vae_loads(tmp_path): + # Bug caught: a complete VAE refused (a check that compares on-disk names such as to_out.0 instead of the + # parameters the applier replaced). + from mlx_dfloat.mflux.ernie.init import load_vae + + _write_vae(tmp_path) + vae = load_vae(tmp_path) + assert vae.bn.running_mean.dtype == mx.bfloat16 # the file's tensor (the init is float32) + + +@pytest.mark.mflux +def test_a_directory_without_the_text_encoder_or_the_vae_is_a_format_error(tmp_path): + # Bug caught: a DF11 checkpoint passed as --base crashing inside mflux's loader with a bare FileNotFoundError. + from mlx_dfloat.mflux.ernie.init import load_text_encoder, load_vae + + with pytest.raises(DFloatFormatError, match="ERNIE-Image base"): + load_text_encoder(tmp_path) + with pytest.raises(DFloatFormatError, match="ERNIE-Image base"): + load_vae(tmp_path) + + +@pytest.mark.mflux +def test_missing_tokenizer_files_are_a_format_error(tmp_path): + # Bug caught: an empty base directory surfacing mflux's own tokenizer error instead of the package's. + from mlx_dfloat.mflux.ernie.init import load_tokenizers + + with pytest.raises(DFloatFormatError, match="no usable ERNIE-Image tokenizer"): + load_tokenizers(tmp_path) + + +@pytest.mark.mflux +def test_a_partial_tokenizer_directory_is_a_format_error(monkeypatch, tmp_path): + # Bug caught: transformers' plain OSError for a tokenizer/ missing one of its files escaping as a raw error + # instead of the package's format error naming the base. + from mflux.models.common.tokenizer.tokenizer_loader import TokenizerLoader + + from mlx_dfloat.mflux.ernie.init import load_tokenizers + + def partial(**_kwargs): + raise OSError("Can't load tokenizer for 'tokenizer'") + + monkeypatch.setattr(TokenizerLoader, "load_all", staticmethod(partial)) + with pytest.raises( + DFloatFormatError, match=r"no usable ERNIE-Image tokenizer.*Can't load tokenizer" + ): + load_tokenizers(tmp_path) + + +def test_load_base_returns_the_three_components(monkeypatch, tmp_path): + # Bug caught: load_base wiring the encoder into the vae slot (or the reverse), or the tokenizers dropped. + from mlx_dfloat.mflux.ernie import init as einit + + monkeypatch.setattr(einit, "load_vae", lambda root: ("vae", root)) + monkeypatch.setattr(einit, "load_text_encoder", lambda root: ("encoder", root)) + monkeypatch.setattr(einit, "load_tokenizers", lambda root: {"ernie": root}) + parts = einit.load_base(tmp_path) + assert (parts.vae, parts.text_encoder, parts.tokenizers) == ( + ("vae", tmp_path), + ("encoder", tmp_path), + {"ernie": tmp_path}, + ) + + +@pytest.mark.mflux +@pytest.mark.slow +def test_the_real_tokenizer_loads_from_the_pattern_files_alone(tmp_path): + # Bug caught: the tokenizer subdir not where mflux looks, or the narrowed list missing a file mflux's + # TokenizerLoader (transformers' AutoTokenizer) needs, so a user's first fetch would load no tokenizer. Links only + # the snapshot's tokenizer/ files a pattern selects. Reads a local base snapshot (no network): gated on --run-slow + # and MLX_DFLOAT_ERNIE_BASE, like the real-blocks tests, so CI never runs it; run it by hand before a release. + base = os.environ.get("MLX_DFLOAT_ERNIE_BASE") + if not base: + pytest.skip( + "MLX_DFLOAT_ERNIE_BASE is not set (an ERNIE-Image base snapshot with tokenizer/)" + ) + from mlx_dfloat.mflux.ernie.init import load_tokenizers + + (tmp_path / "tokenizer").mkdir() + for path in sorted(Path(base, "tokenizer").iterdir()): + name = f"tokenizer/{path.name}" + if any(fnmatch(name, p) for p in BASE_PATTERNS): + (tmp_path / name).symlink_to(path.resolve()) + linked = sorted(f"tokenizer/{p.name}" for p in (tmp_path / "tokenizer").iterdir()) + assert linked == ["tokenizer/tokenizer.json", "tokenizer/tokenizer_config.json"] + tokenizer = load_tokenizers(tmp_path)["ernie"] + assert tokenizer.tokenize("a").input_ids.tolist() == [[1, 1097]] + assert int(mx.sum(tokenizer.tokenize(PROMPT).attention_mask[0])) == PROMPT_TOKENS diff --git a/tests/test_ernie_memory.py b/tests/test_ernie_memory.py new file mode 100644 index 0000000..39d07e6 --- /dev/null +++ b/tests/test_ernie_memory.py @@ -0,0 +1,566 @@ +"""ERNIE-Image's memory rules: the encoder's resident bytes, the text-token count, the sizes, the measured constants. + +Measured inputs (2026-10-08, git a771e58, M1 Max 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0): two one-step de-risk +runs at 1024², seed 42, the 26-token lighthouse prompt, one process each (build, encode, drop, set load, one step at the +planner's cache limit for the block-cache probe's allowance, VAE decode on the resident set): ERNIE-Image-Turbo at +guidance 1.0 (batch 1) and ERNIE-Image at guidance 4.0 (mflux's " " negative then the prompt: batch 2). Their phase +peaks, sizes, encode holds and post-build counters are committed in +``bench/results/calibration/ernie-image-turbo-1024.json`` and ``bench/results/calibration/ernie-image-1024.json`` +(``constants_from`` names the run each term came from), and the tests below read them from there. +""" + +import json +import logging +from pathlib import Path + +import mlx.core as mx +import numpy as np +import pytest +from tests._df11_fixtures import random_bf16, write_checkpoint + +from mlx_dfloat.bench.capped import tier_limits +from mlx_dfloat.errors import DFloatResourceError +from mlx_dfloat.format import open_checkpoint +from mlx_dfloat.integrate import memory as imem +from mlx_dfloat.mflux._phases import FamilySizes, PhaseConstants, activation_allowance, fit_for +from mlx_dfloat.mflux._pipeline import plan_call_for +from mlx_dfloat.mflux.ernie import memory as emem + +# The de-risk prompt's token count (mflux 0.20.0 TokenizerLoader on both base snapshots, 2026-10-08). +DERISK_TOKENS = 26 + + +def _bf16(*shape): + return mx.zeros(shape, dtype=mx.bfloat16) + + +def _write_encoder(root, *, layers=3): + """An ERNIE text encoder file the way the base stores it: the language model under ``language_model.model.`` + (``layers`` numbered layers), the vision tower and the projector beside it (the real file's prefixes, 2026-10-08). + + bf16 bytes: embed_tokens 10 x 4 = 80, one up_proj 4 x 4 = 32 per layer, the final norm (4,) = 8, the vision tensor + 8 x 4 = 64, the projector 6 x 4 = 48. + """ + enc = root / "text_encoder" + enc.mkdir(parents=True) + tensors = { + "language_model.model.embed_tokens.weight": _bf16(10, 4), + "language_model.model.norm.weight": _bf16(4), + "vision_tower.transformer.layers.0.attention.q_proj.weight": _bf16(8, 4), + "multi_modal_projector.linear_1.weight": _bf16(6, 4), + } + for i in range(layers): + tensors[f"language_model.model.layers.{i}.mlp.up_proj.weight"] = _bf16(4, 4) + mx.save_safetensors(str(enc / "model.safetensors"), tensors) + + +def test_encoder_bytes_count_the_language_model_layers_mflux_evaluates(tmp_path): + # Bug caught: the encoder's file size counted (as for FLUX.1 and Z-Image), the vision tower or the projector + # counted (mflux's prefix filter loads language_model.model.* only, ernie_weight_definition.py:31), or the last + # layer counted: mflux returns the second-to-last hidden state (text_encoder.py:217-221), so under lazy evaluation + # the last layer never runs and its weights are never read. Embeddings 80 + norm 8 + layers 0-1 2 x 32 = 152. + _write_encoder(tmp_path) + assert emem.encoder_bytes_used(tmp_path) == 152 + + +def test_the_encoder_counts_25_of_26_layers(tmp_path): + # Bug caught: the cut hard-coded or off by one against the published 26-layer encoder (25 evaluated, layers + # 0-24): 80 + 8 + 25 x 32 = 888 (26 layers would give 920, 24 give 856). + _write_encoder(tmp_path, layers=26) + assert emem.encoder_bytes_used(tmp_path) == 888 + + +def test_sizes_for_replaces_the_encoder_file_size_with_the_evaluated_language_model_bytes(tmp_path): + # Bug caught: the sizes still carrying the file's size (vision tower included: the encode phase over-predicted), + # or the DF11 / non-block / VAE terms lost when the encoder term is replaced, or a non-block group not counted. + rng = np.random.default_rng(0) + write_checkpoint( + tmp_path / "df11", + groups={ + "layers.0": [random_bf16(rng, (8, 4))], + "adaLN_modulation.1": [random_bf16(rng, (16, 4))], + "final_norm.linear": [random_bf16(rng, (8, 4))], + }, + patterns={ + r"layers\.\d+": ("self_attention.to_q",), + r"adaLN_modulation\.1": (), + r"final_norm\.linear": (), + }, + extras={"text_proj.weight": np.zeros((3, 5), dtype=np.uint16)}, + single_file=True, + ) + _write_encoder(tmp_path / "base") + (tmp_path / "base" / "vae").mkdir() + (tmp_path / "base" / "vae" / "diffusion_pytorch_model.safetensors").write_bytes(b"\0" * 100) + + sizes = emem.sizes_for(open_checkpoint(tmp_path / "df11"), tmp_path / "base") + total = (tmp_path / "df11" / "model.safetensors").stat().st_size + assert sizes.encoders == 152 + assert sizes.extras == 30 # 3 x 5 BF16 + assert sizes.compressed == total - 30 + assert sizes.nonblock == 128 + 64 # adaLN_modulation.1 16 x 4 + final_norm.linear 8 x 4, BF16 + assert sizes.vae == 100 + + +class _Tokens: + def __init__(self, n): + self.attention_mask = mx.ones((1, n), dtype=mx.int32) + self.input_ids = mx.ones((1, n), dtype=mx.int32) + + +class _Tokenizer: + """mflux's tokenizer protocol as ErniePromptEncoder uses it: tokenize(prompt) -> ids and an attention mask.""" + + def __init__(self, lengths): + self.lengths = lengths + + def tokenize(self, prompt): + return _Tokens(self.lengths[prompt]) + + +def test_prompt_tokens_count_the_attention_mask_and_cap_at_2048(): + # Bug caught: the cap missing or off by one, or the padded length counted instead of the mask (mflux's num_valid, + # prompt_encoder.py:15-17). The cap is the tokenizer's max_length, 2048 (ernie_weight_definition.py:43). + tok = _Tokenizer({"abc": 3, "u": 2047, "at": 2048, "o": 2049, "big": 3000}) + assert [emem.prompt_tokens(tok, p) for p in ("abc", "u", "at", "o", "big")] == [ + 3, + 2047, + 2048, + 2048, + 2048, + ] + + +def test_padding_is_not_counted(): + # Bug caught: the padded input length counted (a CFG batch pads every prompt to the longest: the short prompt's + # mask has zeros). + class Padded: + attention_mask = mx.array([[1, 1, 0, 0, 0]], dtype=mx.int32) + input_ids = mx.ones((1, 5), dtype=mx.int32) + + tok = type("T", (), {"tokenize": lambda self, p: Padded()})() + assert emem.prompt_tokens(tok, "x") == 2 + + +def test_a_cfg_call_is_sized_by_its_longest_prompt(): + # Bug caught: the positive prompt only (the negative can be the longer one; build_text_batch pads every prompt to + # the longest, prompt_encoder.py:36). + tok = _Tokenizer({" ": 2, "long": 40}) + assert emem.text_tokens_for(tok, (" ", "long")) == 40 + assert emem.text_tokens_for(tok, ("long", " ")) == 40 + + +def test_the_encode_bound_is_the_phase_without_overhead_plus_the_slack(): + # Bug caught: the overhead kept in the bound, or the slack missing. 10e9 - 0.4e9 + 268_435_456, by hand. + c = PhaseConstants( + overhead_bytes=400_000_000, + vae_transient_bytes=0, + denoise_activation_at_reference=0, + reference_tokens=1, + ) + assert emem.encode_peak_bound(10_000_000_000, c) == 9_868_435_456 + + +def test_the_encode_warning_is_silent_at_the_bound_and_names_the_prompt_length_over_it(): + # Bug caught: `>=` for `>` at the bound (a normal encode warning), or a warning that blames mflux's loading alone + # when the prompt is longer than the one the encode term was measured on (26 tokens). + bound = 7_000_000_000 + assert emem.encode_peak_warning(bound, bound, text_tokens=2048) is None + message = emem.encode_peak_warning(bound + 1, bound, text_tokens=2048) + assert message is not None + assert "for a 2048-token prompt" in message + assert f"measured on one {DERISK_TOKENS}-token prompt" in message + assert "vision tower" in message + + +# --- the measured constants, against the committed calibration records ----------------------------------------- + +GIB = 1024**3 +REPO = Path(__file__).resolve().parents[1] +RECORDS = { + 1: json.loads((REPO / "bench/results/calibration/ernie-image-turbo-1024.json").read_text()), + 2: json.loads((REPO / "bench/results/calibration/ernie-image-1024.json").read_text()), +} +RUNS = {b: {run["label"]: run for run in rec["runs"]} for b, rec in RECORDS.items()} +# The phase peaks each batch's constants were derived from (constants_from names the run per phase). +MEASURED = { + b: { + phase: RUNS[b][label]["footprint_peaks"][phase] + for phase, label in rec["constants_from"].items() + } + for b, rec in RECORDS.items() +} +SIZES = { + b: FamilySizes(**RUNS[b][rec["constants_from"]["denoise"]]["sizes"]) + for b, rec in RECORDS.items() +} +CACHE_LIMITS = { + b: RUNS[b][rec["constants_from"]["denoise"]]["cache_limit"] for b, rec in RECORDS.items() +} +# Every 1024² VAE-phase sample of either model: the VAE term is one value for both batches (overall-peak rule). +VAE_SAMPLES = { + f"{b}:{label}": run["footprint_peaks"]["vae"] + for b in RUNS + for label, run in RUNS[b].items() + if "vae" in run["footprint_peaks"] +} +LARGEST = { + "layers": 436_207_616 +} # one block kind; 36 groups of 218_103_808 elements (4 x 4096² + 3 x 4096 x 12288) +TEXT_TOKENS = 26 # the de-risk prompt's (the base's negative " " is 2) +# The block-cache probe's allowances at 4096 + 26 tokens (the smallest limit at which each batch's decoded buffer is +# reused: 2.6e9 and 3.6e9 B, minus one decoded block). +ALLOWANCE = {1: 2_163_792_384, 2: 3_163_792_384} +# This Mac (M1 Max 32 GB): mx.device_info()["max_recommended_working_set_size"] and the RAM. +RECOMMENDED_32 = 26_800_603_136 +RAM_32 = 34_359_738_368 + + +def _fit(batch, *, text_tokens=TEXT_TOKENS, budget=24_653_119_488, vae_with_set=True): + c = emem.CONSTANTS[batch] + return fit_for( + c, + sizes=SIZES[batch], + largest=LARGEST, + policy="per-block", + cache_limit=CACHE_LIMITS[batch], + allowance=activation_allowance(c, height=1024, width=1024, text_tokens=text_tokens), + budget=budget, + height=1024, + width=1024, + text_tokens=text_tokens, + vae_with_set=vae_with_set, + ) + + +def _plan(batch, budget, *, fit_check=True, cache_limit_override=None, policy="per-block"): + return plan_call_for( + constants=emem.CONSTANTS[batch], + sizes=SIZES[batch], + largest=LARGEST, + policy=policy, + cache_limit_override=cache_limit_override, + fit_check=fit_check, + budget=budget, + height=1024, + width=1024, + text_tokens=TEXT_TOKENS, + log=logging.getLogger("test"), + ) + + +# The constants are calibrated on the two de-risk runs and, for batch 1's denoise term, the CAPPED 24 GB run (the +# overall-peak rule), so at their working point the estimate equals the measurement by construction: the in-sample +# checks are exact. No ERNIE-Image run away from 1024² and 26 tokens exists yet. + + +def test_the_calibration_records_hold_the_runs_the_constants_name(): + # Bug caught: a constant re-derived from a run that is not recorded, or a record of the wrong model or snapshot: + # each batch's denoise term from its own run, the shared encode and VAE terms each from exactly one record. + assert (RECORDS[1]["batch"], RECORDS[2]["batch"]) == (1, 2) + assert ( + RECORDS[1]["df11"] + == "mingyi456/ERNIE-Image-Turbo-DF11@27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b" + ) + assert RECORDS[1]["base"] == "baidu/ERNIE-Image-Turbo@bc68c81e2a1730a394d5fc9fae70713dee940140" + assert ( + RECORDS[2]["df11"] == "mingyi456/ERNIE-Image-DF11@c2dd30ad7dd5a928df2309581b282337f7cdb41f" + ) + assert RECORDS[2]["base"] == "baidu/ERNIE-Image@5346b31d68c9c23758ba56ef8be5e9dc174c7f99" + for b, rec in RECORDS.items(): + assert "denoise" in rec["constants_from"] + assert set(rec["constants_from"].values()) <= set(RUNS[b]), b + assert rec["text_tokens"] == TEXT_TOKENS + for phase in ("encode", "vae"): + assert sum(phase in rec["constants_from"] for rec in RECORDS.values()) == 1, phase + + +def test_the_record_sizes_count_the_evaluated_language_model(): + # Bug caught: the records carrying every language-model tensor (6_858_012_672 B, Task 0's header sum) while the + # code counts the evaluated layers (the encode term would come out negative: the runs held about 6.63e9 B), or the + # other way round. 6_858_012_672 - the last layer's 232_796_160 = 6_625_216_512. + assert {b: s.encoders for b, s in SIZES.items()} == {1: 6_625_216_512, 2: 6_625_216_512} + + +def test_the_overhead_is_the_larger_post_build_gap_of_the_two_runs(): + # Bug caught: the overhead taken from one run (the base's 548_194_818 would put every phase of the Turbo run + # 68_239_360 B low), or computed another way than footprint - MLX active - MLX cache after the build. + gaps = { + label: run["after_build"]["footprint"] + - run["after_build"]["mlx_active"] + - run["after_build"]["mlx_cache"] + for b in RUNS + for label, run in RUNS[b].items() + if "after_build" + in run # the one-step de-risk runs; the later samples record a VAE peak only + } + assert gaps == {"derisk-turbo-1": 616_434_178, "derisk-base-1": 548_194_818} + assert RECORDS[1]["overhead_from"] == "derisk-turbo-1" + assert {b: c.overhead_bytes for b, c in emem.CONSTANTS.items()} == { + 1: 616_434_178, + 2: 616_434_178, + } + + +@pytest.mark.parametrize("batch", [1, 2]) +def test_the_estimate_at_1024_equals_each_batchs_calibration_run_phase_for_phase(batch): + # Bug caught: a term dropped from the denoise phase or counted twice, REFERENCE_TOKENS left at another prompt's + # count (the activation scaled off the calibration point), the batch-1 activation used for batch 2 (the base's + # denoise 387_898_944 B low), or the denoise term derived at another cache limit than the planner's. + # The encode phase (one term for both batches): the evaluated language model 6_625_216_512 + the encode term + # 93_963_846 + overhead 616_434_178 = 7_335_614_536, the Turbo run's encode footprint (the term is taken by + # footprint, so the encode phase never under-predicts a calibration run; MLX's own measure gave 29_933_258). + fit = _fit(batch) + assert fit.peak_phase == "vae" + assert fit.phases["denoise"] == MEASURED[batch]["denoise"] + assert fit.phases["encode"] == 7_335_614_536 == MEASURED[1]["encode"] + assert CACHE_LIMITS[batch] == 436_207_616 + ALLOWANCE[batch] + + +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 sample (the base's de-risk sample, with + # its own sizes, is 35_821_552 B lower and puts the Turbo run's VAE phase that much low), or a sample dropped from a + # record. Five samples: the two de-risk runs, the two MEASURED rows and the base identity check's df11 side; the + # highest is still the Turbo de-risk run. One term serves both models, whose sets differ by 59_576 B, so the + # estimate equals the sample only for the model whose record holds it (Turbo). + assert len(VAE_SAMPLES) == 5 + for b in RUNS: + fit = _fit(b) + vae_peaks = [ + r["footprint_peaks"]["vae"] for r in RUNS[b].values() if "vae" in r["footprint_peaks"] + ] + assert all(peak <= fit.phases["vae"] for peak in vae_peaks), b + source = next(b for b, rec in RECORDS.items() if "vae" in rec["constants_from"]) + assert source == 1 + assert _fit(1).phases["vae"] == MEASURED[1]["vae"] == max(VAE_SAMPLES.values()) + denoise_only = [r for r in RUNS[1].values() if set(r["footprint_peaks"]) == {"denoise"}] + assert len(denoise_only) == 2 # the two CAPPED 24 GB runs: denoise samples, not VAE ones + assert _fit(2).phases["vae"] == 19_163_020_656 + + +@pytest.mark.parametrize(("batch", "limit"), [(1, 2_600_000_000), (2, 3_600_000_000)]) +def test_the_planners_1024_cache_limit_is_one_block_plus_the_batchs_allowance(batch, limit): + # Bug caught: the shared 1.5e9 allowance, two blocks budgeted for a one-kind family, or the batch-1 allowance at + # CFG (the probe released the decoded buffer below 3.6e9 at batch 2). 436_207_616 + 2_163_792_384 / 3_163_792_384. + assert _plan(batch, 24_653_119_488).cache_limit == limit + + +@pytest.mark.parametrize(("batch", "allowance"), [(1, 2_163_792_384), (2, 3_163_792_384)]) +def test_a_cache_limit_override_under_the_one_kind_minimum_warns_for_both_policies( + batch, allowance, caplog +): + # Bug caught: the depth2 minimum missing its look-ahead group, or the per-block minimum without the allowance. + per_block = 436_207_616 + allowance + depth2 = 872_415_232 + allowance + with caplog.at_level("WARNING", logger="test"): + _plan(batch, 24_653_119_488, cache_limit_override=per_block) + assert "below the derived minimum" not in caplog.text + with caplog.at_level("WARNING", logger="test"): + _plan(batch, 24_653_119_488, cache_limit_override=per_block - 1) + assert f"below the derived minimum {per_block}" in caplog.text + caplog.clear() + with caplog.at_level("WARNING", logger="test"): + assert _plan(batch, 24_653_119_488, policy="depth2").cache_limit == depth2 + assert "below the derived minimum" not in caplog.text + with caplog.at_level("WARNING", logger="test"): + _plan(batch, 24_653_119_488, policy="depth2", cache_limit_override=per_block) + assert f"below the derived minimum {depth2}" in caplog.text + assert "twice, for depth2's look-ahead" in caplog.text + + +@pytest.mark.parametrize(("batch", "vae"), [(1, 19_163_080_232), (2, 19_163_020_656)]) +def test_a_32_gb_mac_keeps_the_set_through_the_vae_decode(monkeypatch, batch, vae): + # The shared rule drops the set before the VAE decode iff the VAE phase with the set is over the budget: the + # recommended working set minus the 2 GiB reserve, 24_653_119_488 B here. The VAE phase (17.85 GiB for either batch) is + # 5_490_039_256 / 5_490_098_832 B (5.11 GiB) under it. Bug caught: the decision inverted, or a margin added on top of the + # reserve. + monkeypatch.setattr( + imem.mx, "device_info", lambda: {"max_recommended_working_set_size": RECOMMENDED_32} + ) + budget = imem.budget_bytes() + assert budget == 24_653_119_488 + plan = _plan(batch, budget) + assert plan.drop_set_before_vae is False + assert (plan.estimate.peak_phase, plan.estimate.peak_bytes) == ("vae", vae) + # The boundary: one byte less of budget than the VAE phase drops the set (`>`, not `>=`, at the budget). + assert _plan(batch, vae, fit_check=False).drop_set_before_vae is False + assert _plan(batch, vae - 1, fit_check=False).drop_set_before_vae is True + + +# --- the CAPPED tiers at 1024², the 26-token prompt, each model's default call (Turbo batch 1, base batch 2) ---------- +# The fit budget a real Mac of that size applies: its recommended working set (2/3 of its RAM) minus the 2 GiB reserve. +# 24 GB: 15_032_385_536 B (14.00 GiB); 16 GB: 9_305_762_474 B (8.67 GiB). The VAE phase with the set (19.16e9 B) is +# over both, so the set is dropped before the decode; the dropped-set VAE phase keeps the extras (they stay on the +# transformer): extras 27_961_600 + VAE file 168_120_878 + transient 7_120_209_073 + overhead 616_434_178 = +# 7_932_725_729 B. The encode phase is 7_335_614_536 B. The denoise phase decides: +# - Turbo (batch 1), denoise 15_446_156_600 B (14.39 GiB, the higher of two CAPPED 24 GB runs' peaks, by the +# overall-peak rule): over the 24 GB budget by 413_771_064 B (0.39 GiB) and over the 16 GB budget: refused in the +# denoise phase at both. +# - Base (batch 2), denoise 16_361_874_888 B (15.24 GiB): over both budgets (24 GB by 1_329_489_352 B): refused in the +# denoise phase. + + +def _tier_budget(tier): + return tier_limits( + tier, host_ram_bytes=RAM_32, host_recommended_bytes=RECOMMENDED_32 + ).fit_budget_bytes + + +@pytest.mark.parametrize( + ("batch", "size", "expected"), + [(2, 1024, 450_976_567), (2, 512, 402_653_184), (1, 512, 136_331_547)], +) +def test_the_batch_2_denoise_term_is_floored_by_the_float32_weight_copies(batch, size, expected): + # Bug caught (X3): a small CFG call predicted below the float32 copies a batch-2 step holds whatever the image size + # (a float32 hidden state times a BF16 weight makes MLX copy the weight to float32: gate_proj and up_proj, 2 x + # 12288 x 4096 x 4 = 402_653_184 B), or the floor changing the 1024² term it was calibrated on. 512² + 26 tokens: + # 450_976_567 x 1050 / 4122 = 114_877_582 scaled, floored; batch 1 carries no floor (not measured below 1024²): + # 535_198_703 x 1050 / 4122 = 136_331_547. + from mlx_dfloat.mflux._phases import denoise_activation_bytes + + c = emem.CONSTANTS[batch] + assert denoise_activation_bytes(c, height=size, width=size, text_tokens=TEXT_TOKENS) == expected + + +def test_the_tier_budgets_are_a_real_macs(): + # Bug caught: a tier judged against its watchdog ceiling or the host's budget instead of the fit budget. + assert (_tier_budget(16), _tier_budget(24)) == (9_305_762_474, 15_032_385_536) + + +def test_turbo_at_24_gb_is_refused_0_39_gib_over_the_budget(): + # Bug caught (X1): the batch-1 denoise term taken from one run instead of the highest denoise sample (the de-risk + # run's passed the call 58 MB under the budget; the first CAPPED run's, 37 MB over it; the second identical run + # peaked 0.35 GiB higher still), or the dropped-set VAE phase without the extras. + forced = _plan(1, _tier_budget(24), fit_check=False) + assert forced.estimate.peak_bytes - _tier_budget(24) == 413_771_064 + + +@pytest.mark.parametrize(("batch", "tier"), [(1, 16), (1, 24), (2, 16), (2, 24)]) +def test_every_capped_call_is_refused_in_the_denoise_phase(batch, tier): + # Bug caught: the refusal naming the VAE phase (the set drop not applied), or a call that cannot hold the + # compressed set and one step planned to fit. + budget = _tier_budget(tier) + with pytest.raises(DFloatResourceError, match="in the denoise phase"): + _plan(batch, budget) + forced = _plan(batch, budget, fit_check=False) + assert forced.drop_set_before_vae is True + assert forced.estimate.peak_phase == "denoise" + assert forced.estimate.peak_bytes == {1: 15_446_156_600, 2: 16_361_874_888}[batch] + assert forced.estimate.phases["vae"] == 7_932_725_729 + + +# What the prompt encode may hold before the model warns (model._check_encode_peak): MLX's encode peak minus what was +# active when the phase began, against the encode phase without its overhead (evaluated language model + encode term) +# + the slack. The encode term is per prompt, so the bound does not depend on the batch or on the prompt length. +NOT_LOADED = ( + 806_610_944 + 33_556_480 +) # the vision tower and the projector in the encoder's file (Task 0's header) + + +def _encode_bound(batch, text_tokens=TEXT_TOKENS): + return emem.encode_peak_bound( + _fit(batch, text_tokens=text_tokens).phases["encode"], emem.CONSTANTS[batch] + ) + + +@pytest.mark.parametrize(("batch", "spare"), [(1, 358_210_370), (2, 332_466_044)]) +def test_the_encode_bound_holds_each_measured_encode_with_its_spare(batch, spare): + # Bug caught: a bound under what a normal encode holds (every call would warn). The bound is the evaluated + # language model 6_625_216_512 + the footprint encode term 93_963_846 + the 256 MiB slack = 6_987_615_814; the + # runs held 6_629_405_444 (Turbo) and 6_655_149_770 (base) by MLX's measure. + run = RUNS[batch][next(iter(RUNS[batch]))] + assert _encode_bound(batch) - run["mlx_encode_held"] == spare + + +@pytest.mark.parametrize("text_tokens", [1, TEXT_TOKENS, emem.TEXT_MAX_LENGTH]) +def test_an_encode_that_also_loads_the_vision_tower_is_over_the_bound_at_any_prompt(text_tokens): + # Bug caught: the encode bound carrying the denoise allowance (it grows with the prompt), or a slack wide enough + # that an mflux change loading the vision tower and projector (+840_167_424 B) goes unnoticed. Over by + # 6_655_149_770 + 840_167_424 - 6_987_615_814 = 507_701_380 at every prompt length, for the base run's hold. + held = RUNS[2]["derisk-base-1"]["mlx_encode_held"] + for batch in (1, 2): + assert held + NOT_LOADED - _encode_bound(batch, text_tokens) == 507_701_380 + + +# The later 1024² samples in the calibration records are the committed runs' own watched peaks (each reached in the VAE +# phase with the set resident), never hand-copied: the record of each run, and the calibration label it is filed under. +OUT_OF_SAMPLE = [ + ("bench/results/tiers/ernie-image-turbo-1024.json", 1, "measured-8-steps"), + ("bench/results/tiers/ernie-image-1024.json", 2, "measured-50-steps"), + ("bench/results/identity/ernie-image-1024/df11-result.json", 2, "identity-4-steps-cfg"), +] + + +def _record_run(path): + """A committed record's report, its height and width, and its watched (footprint) peak.""" + record = json.loads((REPO / path).read_text()) + if ( + "report" in record + ): # verify_image's df11 side: the model's report under "report", the size in the key + size = record["key"]["size"] + return record["report"], size, size, record["footprint_peak_bytes"] + return record, record["height"], record["width"], record["watched_peak_bytes"] + + +@pytest.mark.parametrize(("path", "batch", "label"), OUT_OF_SAMPLE) +def test_the_later_vae_samples_are_the_committed_runs_watched_peaks(path, batch, label): + # Bug caught: a calibration sample that is not the committed run's peak (a typo, or a re-run whose record was not + # carried over), so the overall-peak rule would be checked against a number no run measured. + _report, _h, _w, watched = _record_run(path) + assert RUNS[batch][label]["footprint_peaks"] == {"vae": watched} + + +# The committed 1024² records beyond the one-step calibration runs, each re-estimated with the current constants from +# its own report (its sizes, the cache limit in force, its text tokens, its policy, its CFG batch and its VAE +# strategy), so a re-run that changes the record is checked against the code as it stands. The MEASURED rows and the +# identity df11 side peak in the VAE decode with the set resident, whose 1024² estimate is the calibration value +# itself: they test that the step count and the CFG batch do not grow the 1024² peak. The Turbo 24 GB CAPPED row runs +# with the set dropped before the VAE and peaks in the denoise phase; the batch-1 denoise term is the highest of its +# runs' peaks (the overall-peak rule: two identical runs peaked 0.35 GiB apart), so a re-run of that row is checked +# against the estimate the earlier runs set. The check is +# one-sided: a phase varies by up to about 1 GiB run to run, +# and a run under the estimate is not the bug. +STEPS_AND_CFG_RECORDS = [ + pytest.param( + "bench/results/tiers/ernie-image-turbo-1024.json", "vae", id="turbo-measured-8-steps" + ), + pytest.param( + "bench/results/tiers/ernie-image-1024.json", "vae", id="base-measured-50-steps-cfg" + ), + pytest.param( + "bench/results/identity/ernie-image-1024/df11-result.json", + "vae", + id="base-identity-df11-4-steps-cfg", + ), + pytest.param( + "bench/results/tiers/ernie-image-turbo-1024-tier24.json", "denoise", id="turbo-capped-24-gb" + ), +] + + +@pytest.mark.parametrize(("path", "phase"), STEPS_AND_CFG_RECORDS) +def test_a_committed_1024_record_peaks_within_the_current_estimate(path, phase): + # Bug caught: a 1024² peak that grows with the step count or the CFG batch while the estimate stays at the one-step + # value (a recorded run more than 0.25 GiB over the current prediction), the estimate peaking in another phase, or + # a record checked against the wrong batch's constants. 1 GiB less of VAE transient turns the Turbo MEASURED row + # and the identity side red (0.50 and 0.42 GiB over the lowered estimate); the base MEASURED row stays green (it sits + # 1.05 GiB under the estimate, 0.05 GiB under the lowered one; the in-sample test pins the term) and so does the + # CAPPED row (its peak is the denoise phase). + report, height, width, watched = _record_run(path) + c = emem.CONSTANTS[report["cfg_batch"]] + tokens = report["text_tokens"] + fit = fit_for( + c, + sizes=FamilySizes(**report["sizes"]), + largest=LARGEST, + policy=report["eval_policy"], + cache_limit=report["cache_limit_in_force"], + allowance=activation_allowance(c, height=height, width=width, text_tokens=tokens), + budget=report["fit"]["budget_bytes"], + height=height, + width=width, + text_tokens=tokens, + vae_with_set=not report["drop_set_before_vae"], + ) + assert fit.peak_phase == phase + assert watched <= fit.peak_bytes + GIB // 4 diff --git a/tests/test_ernie_model.py b/tests/test_ernie_model.py new file mode 100644 index 0000000..176694b --- /dev/null +++ b/tests/test_ernie_model.py @@ -0,0 +1,884 @@ +"""Tests for `DFloatErnieImage`: construction, refusals, the lifecycle and the set, the prelude and the report. + +Every test is `@pytest.mark.mflux`; the module imports mflux only inside helpers, so it collects without it. +""" + +import inspect +import re + +import mlx.core as mx +import mlx.nn as nn +import numpy as np +import pytest +from tests._ernie_tiny import TINY_SEED, StubTextEncoder, fake_model, write_tiny_checkpoint + +from mlx_dfloat.errors import DFloatIntegrationError, DFloatUnsupportedError +from mlx_dfloat.mflux.ernie import init as einit + +pytestmark = pytest.mark.mflux + +# mflux's parameter for each of the four non-block matrices (ernie_weight_mapping.py:12-19, 40-63) and its shape at +# TINY size (hidden 32): Linear(32, 32) x 2, Linear(32, 6 x 32), Linear(32, 2 x 32). +NONBLOCK = { + "time_embedding.linear_1.weight": ("time_embedding", (32, 32)), + "time_embedding.linear_2.weight": ("time_embedding", (32, 32)), + "adaln_modulation.weight": ("adaLN_modulation.1", (192, 32)), + "final_norm.linear.weight": ("final_norm.linear", (64, 32)), +} + + +def _nonblock_sources(tmp_path): + """The non-block matrices ``fake_model`` compressed: the same seed gives the same arrays.""" + matrices, _c, _conv = write_tiny_checkpoint( + tmp_path / "again", np.random.default_rng(TINY_SEED) + ) + return {param: matrices[group][param] for param, (group, _shape) in NONBLOCK.items()} + + +def _param(model, name): + module = model.transformer + for part in name.split(".")[:-1]: + module = getattr(module, part) + return module.weight + + +def test_every_attribute_upstreams_generate_image_reads_exists_on_the_model(tmp_path, monkeypatch): + # Bug caught: an mflux point release reading a new self. in ErnieImage.generate_image our assembly never sets + # (mflux's own initializer is not run). + from mflux.models.ernie_image.variants.txt2img.ernie_image import ErnieImage + + names = set(re.findall(r"self\.(\w+)", inspect.getsource(ErnieImage.generate_image))) + assert names >= {"model_config", "callbacks", "transformer", "bits", "lora_paths"} + model = fake_model(tmp_path, monkeypatch) + assert [n for n in sorted(names) if not hasattr(model, n)] == [] + assert isinstance(model, ErnieImage) + assert (model.bits, model.prompt_cache, model.lora_paths, model.lora_scales) == ( + None, + {}, + [], + [], + ) + assert model.tiling_config.vae_decode_tiles_per_dim is None + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"quantize": 8}, "quantize"), + ({"lora_paths": ["a"]}, "lora_paths"), + ({"lora_scales": [1.0]}, "lora_scales"), + ({"bake_lora": False}, "bake_lora"), + ({"eval_policy": "none"}, "none"), + ({"model": "ernie"}, r"'ernie'.*ernie-image.*ernie-image-turbo"), + ], +) +def test_construction_refuses_quantize_lora_and_unknown_names(monkeypatch, kwargs, match): + # Bug caught: quantize=8 silently honoured by mflux's applier (it would quantize the decoded weights), a LoRA or + # the "none" policy accepted, an unknown name not listing what this path runs, or any of them resolved (a 10.9 GB + # download) before the refusal. + from mlx_dfloat.mflux.ernie.model import DFloatErnieImage + + monkeypatch.setattr( + einit, "resolve", lambda *a, **k: pytest.fail("resolved before the refusal") + ) + with pytest.raises(DFloatUnsupportedError, match=match): + DFloatErnieImage(**{"model": "ernie-image-turbo", **kwargs}) + + +@pytest.mark.parametrize( + ("name", "df11", "base"), + [ + ( + "ernie-image-turbo", + ("mingyi456/ERNIE-Image-Turbo-DF11", "27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b"), + ("baidu/ERNIE-Image-Turbo", "bc68c81e2a1730a394d5fc9fae70713dee940140"), + ), + ( + "ernie-image", + ("mingyi456/ERNIE-Image-DF11", "c2dd30ad7dd5a928df2309581b282337f7cdb41f"), + ("baidu/ERNIE-Image", "5346b31d68c9c23758ba56ef8be5e9dc174c7f99"), + ), + ], +) +def test_the_default_checkpoint_and_base_resolve_at_their_pins_and_user_paths_unpinned( + tmp_path, monkeypatch, name, df11, base +): + # Bug caught: a pin not reaching the resolver (the checkpoint, encoder, VAE and tokenizer following whatever lands + # on main instead of the snapshot the recorded runs used), or a pin applied to a user's own path. + from mlx_dfloat import _metal_decode + from mlx_dfloat.mflux._hub import ResolvedRepo + from mlx_dfloat.mflux.ernie import model as model_module + from mlx_dfloat.mflux.ernie.model import DFloatErnieImage + + class StopError(Exception): + pass + + df11_calls, base_calls = [], [] + + def stop_at_df11(spec, *, patterns, revision=None): + df11_calls.append((spec, revision)) + raise StopError + + def stop_at_base(spec, *, patterns, revision=None): + if patterns == einit.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(einit, "resolve", stop_at_df11) + for kwargs in ({}, {"df11_path": "me/fork"}): + with pytest.raises(StopError): + DFloatErnieImage(name, **kwargs) + monkeypatch.setattr(einit, "resolve", stop_at_base) + for kwargs in ({}, {"base_path": "me/base"}): + with pytest.raises(StopError): + DFloatErnieImage(name, **kwargs) + assert df11_calls == [df11, ("me/fork", None)] + assert base_calls == [base, ("me/base", None)] + + +def test_the_set_load_installs_the_three_nonblock_groups_and_attaches_with_verify_in_call( + tmp_path, monkeypatch +): + # Bug caught: a non-block group decoded but not installed (a zero-size matmul in the timestep MLP, the modulation + # or the final norm), a matrix installed from another group, or the set attached without verify_in_call. + model = fake_model(tmp_path, monkeypatch) + model._lifecycle.ensure_set() + for param, want in _nonblock_sources(tmp_path).items(): + got = _param(model, param) + assert got.shape == NONBLOCK[param][1], param + assert np.array_equal(np.array(got.view(mx.uint16)), want), param + assert model.transformer.seam_cell.state.verify_in_call is True + + +def test_a_new_batch_after_a_generation_drops_the_set_and_reinstalls_the_nonblock_weights( + tmp_path, monkeypatch +): + # Bug caught (Review Focus 4): a non-block weight left on its placeholder after the reload a new batch forces (a + # zero-size matmul on adaln_modulation before the first block), or the encoder reload reading another base. + model = fake_model(tmp_path, monkeypatch) + loads = [] + monkeypatch.setattr( + einit, + "load_text_encoder", + lambda root, **kwargs: (loads.append(root), StubTextEncoder())[1], + ) + model._lifecycle.ensure_embeddings( + model._batch_key(("a",)) + ) # the encoder came with the assembly + model._lifecycle.ensure_set() + model._lifecycle.ensure_embeddings(model._batch_key((" ", "a"))) # CFG: a new batch + counters = model._lifecycle.counters + assert (counters.forced_set_drops, counters.encoder_loads) == (1, 1) + assert loads == [model._base.root] + assert not model._lifecycle.set_resident + assert [_param(model, p).size for p in NONBLOCK] == [0, 0, 0, 0] + model._lifecycle.ensure_set() + assert counters.set_loads == 2 + for param, want in _nonblock_sources(tmp_path).items(): + assert np.array_equal(np.array(_param(model, param).view(mx.uint16)), want), param + + +def test_batches_are_cached_under_their_prompt_list(tmp_path, monkeypatch): + # Bug caught: mflux's own prompt_cache used (its lookup would bypass the lifecycle), or a batch stored under a key + # other than its prompt list (a CFG call would miss its batch). + model = fake_model(tmp_path, monkeypatch) + model._lifecycle.ensure_embeddings(model._batch_key(("a",)), model._batch_key((" ", "a"))) + assert set(model._batches) == {'["a"]', '[" ", "a"]'} + assert model.prompt_cache == {} + for key, batch in (('["a"]', 1), ('[" ", "a"]', 2)): + text_bth, text_lens = model._batches[key] + # The stub tokenizer gives "a" and " " one token each; build_text_batch's width 3072. + assert (text_bth.shape, text_bth.dtype) == ((batch, 1, 3072), mx.bfloat16) + assert (text_lens.shape, text_lens.dtype) == ((batch,), mx.int32) + + +def test_a_dropped_encoder_refuses_with_a_named_error(tmp_path, monkeypatch): + # Bug caught: None left in the encoder's place (mflux would call it on a cache miss and fail with a TypeError far + # from the cause). + model = fake_model(tmp_path, monkeypatch) + model._lifecycle.ensure_embeddings(model._batch_key(("a",))) + with pytest.raises(DFloatIntegrationError, match="not resident"): + model.text_encoder(mx.zeros((1, 1), mx.int32)) + + +def test_the_retained_bound_counts_the_cached_batches_and_the_vae(tmp_path, monkeypatch): + # Bug caught: the bound reading mflux's prompt_cache (empty here: every drop after an encode would refuse when the + # batches are large) instead of the batch cache. + model = fake_model(tmp_path, monkeypatch) + empty = model._retained_bound() + model._lifecycle.ensure_embeddings(model._batch_key((" ", "a"))) + # text_bth (2, 1, 3072) bf16 = 12_288 B, text_lens (2,) int32 = 8 B. + assert model._retained_bound() - empty == 12_288 + 8 + # The absolute bound before any batch (T4): the assembly baseline + the stub VAE's one weight, 4 bf16 = 8 B, + the + # 2 GiB slack. Bug caught: the VAE left out of the bound (a drop after the VAE decode would refuse). + from mlx_dfloat.mflux.ernie.model import RETAINED_SLACK_BYTES + + assert empty == model._baseline_active + 8 + RETAINED_SLACK_BYTES + + +def test_the_pos_cache_worst_case_fits_inside_the_retained_slack(): + # Bug caught: a slack under what mflux's _pos_cache may legitimately keep active across calls (64 entries, + # transformer.py:120-122), so a drop after many prompt lengths would refuse. One 1024² batch-2 entry at the + # 2048-token maximum, by hand: cos + sin 2 x (2 x 6144 x 1 x 128) bf16 = 6_291_456, mask 2 x 6144 bf16 = 24_576. + from mflux.models.ernie_image.model.ernie_transformer.transformer import ErnieTransformer + + from mlx_dfloat.mflux.ernie.model import RETAINED_SLACK_BYTES + + tf = ErnieTransformer(num_layers=0, rope_axes_dim=[32, 48, 48]) # the published head dim 128 + entry = tf.get_pos_encoding( + B=2, H=64, W=64, T=2048, text_lens=mx.array([2048, 2048], dtype=mx.int32) + ) + mx.eval(*entry) + assert sum(int(a.nbytes) for a in entry) == 6_316_032 + assert RETAINED_SLACK_BYTES > 64 * 6_316_032 + + +# --- the CFG batch, the generate prelude, the plan and the report ------------------------------------------------ + +GIB = 1024**3 + + +def _cache_limit_in_force(): + previous = mx.set_cache_limit(0) + mx.set_cache_limit(previous) + return previous + + +def _patch_upstream_generate(monkeypatch, *, raise_before_loop=None): + """Replace ErnieImage.generate_image with a probe: records the limits in force and the kwargs, fires our + after-loop subscriber the way mflux's GenerationContext would, returns a sentinel.""" + from mflux.models.ernie_image.variants.txt2img.ernie_image import ErnieImage + from tests._mlx_limits import current_limits + + seen = {} + + def fake(self, **kwargs): + seen["limit_at_entry"] = _cache_limit_in_force() + seen["limits_at_entry"] = current_limits() + seen["set_resident"] = self._lifecycle.set_resident + seen["kwargs"] = kwargs + if raise_before_loop is not None: + raise raise_before_loop + for subscriber in self.callbacks.after_loop_callbacks(): + subscriber.call_after_loop( + seed=kwargs["seed"], prompt=kwargs["prompt"], latents=mx.zeros(1), config=None + ) + return "image" + + monkeypatch.setattr(ErnieImage, "generate_image", fake) + return seen + + +def _no_encoder(monkeypatch): + """Any encoder reload fails the test (the batch must come from the cache).""" + monkeypatch.setattr( + einit, "load_text_encoder", lambda *a, **k: pytest.fail("the encoder was reloaded") + ) + + +@pytest.mark.parametrize( + ("guidance", "negative", "expected"), + [ + (1.0, "n", ("p",)), + (None, "n", ("p",)), + (0.5, "n", ("p",)), + (1.0001, "n", ("n", "p")), + (4.0, None, (" ", "p")), + (4.0, "", (" ", "p")), + (4.0, " ", (" ", "p")), + (4.0, " x ", (" x ", "p")), + ], +) +def test_cfg_prompts_follow_mflux(guidance, negative, expected): + # Bug caught: `>=` for `>` (ernie_image.py:164: guidance <= 1.0 runs one prompt; Config stores None as 0.0, + # config.py:48), the positive first (mflux takes pred[:1] as the unconditional branch, :249: guidance would push + # toward the negative), or a blank negative not replaced by " " (:167). + from mlx_dfloat.mflux.ernie.model import DFloatErnieImage + + assert ( + DFloatErnieImage.cfg_prompts("p", negative_prompt=negative, guidance=guidance) == expected + ) + + +def _bits(array): + return np.array(array.view(mx.uint16)) if array.dtype == mx.bfloat16 else np.array(array) + + +@pytest.mark.parametrize(("negative", "guidance"), [("n", 4.0), (None, 1.0)]) +def test_the_override_returns_mfluxs_own_batch(tmp_path, monkeypatch, negative, guidance): + # Bug caught: the copied CFG rule drifting from mflux's own _encode_prompts (order, blank handling, threshold), or + # a batch stored under another call's key. The oracle is mflux's method on a throwaway object. + from types import SimpleNamespace + + from mflux.models.ernie_image.variants.txt2img.ernie_image import ErnieImage + + model = fake_model(tmp_path, monkeypatch) + model.encode("p", negative_prompt=negative, guidance=guidance) + got = model._encode_prompts(prompt="p", negative_prompt=negative, guidance=guidance) + oracle = SimpleNamespace( + prompt_cache={}, tokenizers=model.tokenizers, text_encoder=StubTextEncoder() + ) + want = ErnieImage._encode_prompts( + oracle, prompt="p", negative_prompt=negative, guidance=guidance + ) + assert np.array_equal(_bits(got[0]), _bits(want[0])) + assert np.array_equal(np.array(got[1].view(mx.uint32)), np.array(want[1].view(mx.uint32))) + + +def test_a_call_whose_batch_was_not_encoded_is_refused(tmp_path, monkeypatch): + # Bug caught: a cache miss falling through to mflux's encoder (dropped: a TypeError, or an encoder reload inside + # the loop next to the set). + model = fake_model(tmp_path, monkeypatch) + with pytest.raises(DFloatIntegrationError, match=r'\["p"\] were not encoded'): + model._encode_prompts(prompt="p", negative_prompt=None, guidance=1.0) + + +def test_blank_and_missing_negatives_share_one_batch(tmp_path, monkeypatch): + # Bug caught (Review Focus 2): None, "" and " " keyed apart (three encodes and set reloads for one mflux batch + # [" ", prompt]), or a blank negative encoded as itself. Real mflux loop, one step each. + model = fake_model(tmp_path, monkeypatch) + run = {"seed": 1, "prompt": "a", "num_inference_steps": 1, "height": 64, "width": 64} + model.generate_image(**run, guidance=4.0, negative_prompt=None) + loads = model._lifecycle.counters.encoder_loads + _no_encoder(monkeypatch) + for negative in ("", " "): + model.generate_image(**run, guidance=4.0, negative_prompt=negative) + assert list(model._batches) == ['[" ", "a"]'] + assert model._lifecycle.counters.encoder_loads == loads + assert model._lifecycle.counters.set_loads == 1 + + +def test_guidance_4_then_5_reuses_the_batch_and_1_does_not(tmp_path, monkeypatch): + # Bug caught: the key built from the raw guidance (a re-encode and a set reload on every guidance change), or a + # guidance-1 call reusing the CFG batch (a batch-2 call where mflux runs one prompt). + model = fake_model(tmp_path, monkeypatch) + run = {"seed": 1, "prompt": "a", "num_inference_steps": 1, "height": 64, "width": 64} + model.generate_image(**run, guidance=4.0) + counters = model._lifecycle.counters + model.generate_image(**run, guidance=5.0) + assert (counters.forced_set_drops, len(model._batches)) == (0, 1) + model.generate_image(**run, guidance=1.0) + assert counters.forced_set_drops == 1 + assert set(model._batches) == {'[" ", "a"]', '["a"]'} + + +def test_text_tokens_follow_the_longest_prompt_of_the_call(tmp_path, monkeypatch): + # Bug caught: the plan sized by the positive prompt only (build_text_batch pads every prompt to the longest, a + # longer negative's activations unplanned). + from tests._ernie_tiny import StubTokenizer + + model = fake_model( + tmp_path, monkeypatch, model="ernie-image", tokenizer=StubTokenizer({"p": 7, "n": 12}) + ) + seen = _patch_upstream_generate(monkeypatch) + model.generate_image(seed=1, prompt="p", height=64, width=64, negative_prompt="n") + assert model.report()["text_tokens"] == 12 + # T3: the negative prompt reaches mflux's generate_image (its own CFG rule and the image metadata read it). + assert seen["kwargs"]["negative_prompt"] == "n" + + +def _double(allowance): + from mlx_dfloat.mflux._phases import PhaseConstants + + return PhaseConstants( + overhead_bytes=0, + vae_transient_bytes=0, + denoise_activation_at_reference=0, + reference_tokens=4097, + allowance_at_reference=allowance, + allowance_reference_tokens=4097, + ) + + +@pytest.mark.parametrize( + ("guidance", "batch", "limit"), [(4.0, 2, 3_000_017_408), (1.0, 1, 1_000_017_408)] +) +def test_a_cfg_call_plans_with_the_batch_two_constants( + tmp_path, monkeypatch, guidance, batch, limit +): + # Bug caught (Review Focus 1, Python-API half): the constants looked up by model name (Turbo planned with batch-1 + # activations for a batch-2 call: 1_000_017_408 at guidance 4), or the batch counted from the model instead of the + # prompts. The doubles' reference is 4096 + 1 tokens ("a" and " " are one stub token each), so no scaling, no + # floor: the tiny decoded block 17_408 B + the allowance, by hand. + model = fake_model(tmp_path, monkeypatch) + model._constants = {1: _double(1_000_000_000), 2: _double(3_000_000_000)} + model.generate_image( + seed=1, prompt="a", num_inference_steps=1, height=1024, width=1024, guidance=guidance + ) + report = model.report() + assert (report["cfg_batch"], report["cache_limit_in_force"]) == (batch, limit) + + +@pytest.mark.parametrize( + ("name", "guidance", "steps", "batch"), + [("ernie-image", 4.0, 50, 2), ("ernie-image-turbo", 1.0, 8, 1)], +) +def test_the_python_api_defaults_follow_mflux_per_variant( + tmp_path, monkeypatch, name, guidance, steps, batch +): + # Bug caught: the base's Python default left at mflux's class default 1.0 (no CFG: another image than the command + # gives), or the defaults swapped between variants. Literals: ernie_image_generate.py:22, 30-31 (4.0), + # ernie_image_turbo_generate.py:42-43 (1.0), defaults.py:35-36 (50, 8), ernie_image.py:64-65 ("linear"). + model = fake_model(tmp_path, monkeypatch, model=name) + seen = _patch_upstream_generate(monkeypatch) + model.generate_image(seed=1, prompt="a", height=64, width=64) + kwargs = seen["kwargs"] + assert (kwargs["guidance"], kwargs["num_inference_steps"], kwargs["scheduler"]) == ( + guidance, + steps, + "linear", + ) + assert model.report()["cfg_batch"] == batch + + +def test_an_explicit_guidance_on_turbo_is_passed_through(tmp_path, monkeypatch): + # Bug caught: the Python API forcing Turbo's guidance to 1.0 (mflux's API allows CFG on Turbo; only the command + # refuses it). + model = fake_model(tmp_path, monkeypatch, model="ernie-image-turbo") + seen = _patch_upstream_generate(monkeypatch) + model.generate_image(seed=1, prompt="a", height=64, width=64, guidance=2.0) + assert seen["kwargs"]["guidance"] == 2.0 + assert model.report()["cfg_batch"] == 2 + + +@pytest.mark.parametrize("raised", [None, RuntimeError("in the loop")], ids=["returns", "raises"]) +def test_a_python_api_call_runs_under_the_commands_caps_and_restores_mlxs_defaults( + tmp_path, monkeypatch, raised +): + # Bug caught: generate_image not decorated (the Python API at MLX's default wired limit 0, while the VAE term is + # measured under the command's caps), or the caps left installed after a call that returns or raises. + from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap + + model = fake_model(tmp_path, monkeypatch) + seen = _patch_upstream_generate(monkeypatch, raise_before_loop=raised) + wired, memory = command_caps() + assert wired > 0 + with mlx_without_wired_cap() as start: + if raised is None: + model.generate_image(seed=1, prompt="p", height=256, width=256) + else: + with pytest.raises(RuntimeError, match="in the loop"): + model.generate_image(seed=1, prompt="p", height=256, width=256) + after = current_limits() + assert (seen["limits_at_entry"]["wired"], seen["limits_at_entry"]["memory"]) == (wired, memory) + assert after == start + + +def _fake_encode_peak(monkeypatch, model, above_start): + """Make the tracker record an encode phase whose MLX peak is ``above_start`` bytes over the phase's start.""" + real_end = model._tracker.end + + def end(name): + real_end(name) + if name == "encode": + record = model._tracker.peaks["encode"] + record["mlx_peak"] = record["active_at_start"] + above_start + + monkeypatch.setattr(model._tracker, "end", end) + + +@pytest.mark.parametrize(("over", "warns"), [(1, True), (0, False)]) +def test_an_encode_over_the_bound_warns_through_the_pure_helper(tmp_path, monkeypatch, over, warns): + # Bug caught: _check_encode_peak never called after the encode phase, measuring the wrong phase record, or a bound + # of its own that drifts from the helper's (`>` at the bound: silent). The bound comes from the plan through the + # helper Task 7 pins with literals. + import warnings + + from mlx_dfloat.mflux.ernie import memory as emem + + model = fake_model(tmp_path, monkeypatch) + plan = model.plan_call(height=256, width=256, text_tokens=1, batch=1) + bound = emem.encode_peak_bound(plan.estimate.phases["encode"], model._constants[1]) + _patch_upstream_generate(monkeypatch) + _fake_encode_peak(monkeypatch, model, bound + over) + if warns: + with pytest.warns(UserWarning, match=r"for a 1-token prompt"): + model.generate_image(seed=1, prompt="a", height=256, width=256) + else: + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert model.generate_image(seed=1, prompt="a", height=256, width=256) == "image" + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"image_path": "in.png"}, "image_path"), + ({"image_strength": 0.5}, "image_path"), + ({"pid_decode": True}, "pid_decode"), + ], +) +def test_img2img_and_pid_decode_are_refused_before_anything_loads( + tmp_path, monkeypatch, kwargs, match +): + # Bug caught: encoding the prompt (an encoder load) or loading the set before refusing img2img or mflux's PiD + # decoder (which loads its own caption encoder next to the set). + model = fake_model(tmp_path, monkeypatch) + monkeypatch.setattr(model._lifecycle, "ensure_embeddings", lambda *p: pytest.fail("encoded")) + with pytest.raises(DFloatUnsupportedError, match=match): + model.generate_image(seed=1, prompt="p", **kwargs) + assert model._lifecycle.counters.set_loads == 0 + assert model._lifecycle.counters.encoder_loads == 0 + + +@pytest.mark.parametrize("prompt", ["", " "]) +def test_an_empty_or_blank_prompt_is_refused_before_anything_loads(tmp_path, monkeypatch, prompt): + # Bug caught: an empty positive prompt reaching mflux (ERNIE-Image's tokenizer gives no tokens for "" — 0 tokens, + # ids [] on both base snapshots, 2026-10-08 — so the call would run a zero-length text row through the transformer), + # or the refusal coming after an encode or a set load. A blank negative stays mflux's (" "). + model = fake_model(tmp_path, monkeypatch) + monkeypatch.setattr(model._lifecycle, "ensure_embeddings", lambda *p: pytest.fail("encoded")) + with pytest.raises( + DFloatUnsupportedError, match="an empty or blank prompt is refused as a user error" + ): + model.generate_image(seed=1, prompt=prompt, height=64, width=64) + assert model._lifecycle.counters.set_loads == 0 + + +def test_keyboard_interrupt_restores_the_cache_limit_and_leaves_no_pending_words( + tmp_path, monkeypatch +): + # Bug caught: mflux's StopImageGenerationException (raised for a Ctrl-C in its loop, ernie_image.py:104-108) + # leaving status words pending (begin_step would refuse the next call), the cache limit not restored, or the + # interrupt dropping the set (a reload per Ctrl-C). Real mflux loop; the interrupt fires at the second block of + # the CFG step's one batch-2 call. + from mflux.utils.exceptions import StopImageGenerationException + + model = fake_model(tmp_path, monkeypatch) + before = _cache_limit_in_force() + model.encode("p", negative_prompt=None, guidance=4.0) + model._lifecycle.ensure_set() + provider = model._provider + real = provider.weights_for + calls = [] + interrupt = {"on": True} + + def interrupting(name, shapes): + calls.append(name) + if interrupt["on"] and len(calls) == 2: + raise KeyboardInterrupt + return real(name, shapes) + + monkeypatch.setattr(provider, "weights_for", interrupting) + run = { + "seed": 1, + "prompt": "p", + "num_inference_steps": 1, + "height": 64, + "width": 64, + "guidance": 4.0, + } + with pytest.raises(StopImageGenerationException): + model.generate_image(**run) + assert len(calls) == 2 + assert _cache_limit_in_force() == before + assert provider.pending == [] + assert model._lifecycle.set_resident + interrupt["on"] = False + assert model.generate_image(**run) is not None + assert model._lifecycle.counters.set_loads == 1 + + +@pytest.mark.parametrize( + ("policy", "status"), [("per-block", "measured"), ("depth2", "unmeasured")] +) +def test_the_report_names_the_family_the_batch_and_the_uncompiled_predict( + tmp_path, monkeypatch, policy, status +): + # Bug caught: the report naming no predict mode, or "uncompiled" before any call ran the bypass; cfg_batch not + # following the call (the bench table would mislabel a CFG run); a non-block group missing; the estimate labelled + # as a measurement; depth2 reported like the measured path. + model = fake_model(tmp_path, monkeypatch, eval_policy=policy) + before = model.report() + assert (before["family"], before["model"], before["predict"], before["fit"]) == ( + "ernie", + "ernie-image-turbo", + None, + None, + ) + assert before["sizes"] == { + "compressed": 1_000, + "extras": 100, + "nonblock": 10, + "encoders": 1_000, + "vae": 10, + } # fake_model's own FamilySizes + model.generate_image( + seed=1, prompt="a", num_inference_steps=1, height=64, width=64, guidance=4.0 + ) + report = model.report() + assert (report["cfg_batch"], report["predict"], report["text_tokens"]) == (2, "uncompiled", 1) + assert report["nonblock_groups"] == [ + "adaLN_modulation.1", + "final_norm.linear", + "time_embedding", + ] + assert report["fit"]["label"] == "predicted" + assert (report["eval_policy"], report["eval_policy_status"]) == (policy, status) + assert report["base"]["revision"] == "b" * 40 + + +def test_save_model_is_refused_and_freeze_skips_a_dropped_encoder(tmp_path, monkeypatch): + # Bug caught: freeze() calling .freeze() on the encoder stand-in (it is not a module), or save_model writing a + # checkpoint with zero-size placeholders as the transformer. + model = fake_model(tmp_path, monkeypatch) + model.encode("p") # drops the encoder + model.freeze() + with pytest.raises(DFloatUnsupportedError, match="save"): + model.save_model(str(tmp_path)) + + +def test_the_constructor_builds_with_the_models_overrides_and_the_language_model_bytes( + tmp_path, monkeypatch +): + # Bug caught (S3 review fix 9): __init__ wiring past the resolver untested: the transformer built without the + # model's overrides (mflux builds ErnieTransformer(**model_config.transformer_overrides), + # ernie_image_initializer.py:63; ModelConfig.ernie_image_turbo() sets rope_axes_dim [32, 48, 48]), the base + # components read from another directory, or the sizes taken from the encoder file (vision tower included) instead + # of the language model's tensor bytes. + from tests._df11_fixtures import write_bf16_original + from tests._ernie_tiny import TINY, stub_components + + from mlx_dfloat import _metal_decode + from mlx_dfloat.mflux.ernie import model as model_module + + write_tiny_checkpoint(tmp_path / "df11", np.random.default_rng(TINY_SEED)) + base = tmp_path / "base" + # The language model: embed_tokens (10, 4) 80 B, two layers' up_proj (4, 4) 32 B each, final norm (4,) 8 B; the + # last layer never runs (mflux returns the second-to-last hidden state). Never loaded: a vision tower tensor + # (8, 4) 64 B. 80 + 32 + 8 = 120. + write_bf16_original( + base / "text_encoder", + { + "language_model.model.embed_tokens.weight": np.zeros((10, 4), np.uint16), + "language_model.model.layers.0.mlp.up_proj.weight": np.zeros((4, 4), np.uint16), + "language_model.model.layers.1.mlp.up_proj.weight": np.zeros((4, 4), np.uint16), + "language_model.model.norm.weight": np.zeros((4,), np.uint16), + "vision_tower.transformer.layers.0.attention.q_proj.weight": np.zeros( + (8, 4), np.uint16 + ), + }, + ) + (base / "vae").mkdir() + (base / "vae" / "diffusion_pytorch_model.safetensors").write_bytes(b"\0" * 100) + + monkeypatch.setattr(_metal_decode, "ensure_canary", lambda **kwargs: None) + loaded = [] + monkeypatch.setattr( + einit, "load_base", lambda root: (loaded.append(root), stub_components())[1] + ) + built = [] + real_build = model_module.build_transformer + + def tiny_build(ckpt, **kwargs): + built.append(kwargs) + return real_build(ckpt, transformer_overrides=TINY) + + monkeypatch.setattr(model_module, "build_transformer", tiny_build) + model = model_module.DFloatErnieImage( + "ernie-image-turbo", df11_path=str(tmp_path / "df11"), base_path=str(base) + ) + assert built == [{"transformer_overrides": {"rope_axes_dim": [32, 48, 48]}}] + assert loaded == [base] + report = model.report() + assert report["sizes"]["encoders"] == 120 + assert report["sizes"]["vae"] == 100 + # The three non-block groups decoded at TINY size: (32, 32) x 2 + (192, 32) + (64, 32) BF16. + assert report["sizes"]["nonblock"] == (2 * 32 * 32 + 192 * 32 + 64 * 32) * 2 + assert report["df11"]["config_source"] == "config.json" + assert report["base"]["root"] == str(base) + + +def _innermost_locals(tb): + while tb.tb_next is not None: + tb = tb.tb_next + return tb.tb_frame.f_locals + + +def test_clear_frames_walks_the_whole_chain_of_causes_and_contexts(): + # Bug caught: only the raised exception and its __context__ cleared, so a __cause__ one link further (raise ... from + # an earlier error) keeps its frame's locals, in production a block's decoded weights, alive while the caller holds + # the exception; or a chain that loops back on itself never ending. + from mlx_dfloat.mflux.ernie.model import _clear_frames + + def fail(kind, message): + big = mx.zeros(1024) # noqa: F841 # stands in for a block's decoded weights + raise kind(message) + + try: + try: + try: + fail(ValueError, "first") + except ValueError as first: + raise KeyError("second") from first + except KeyError: + fail(RuntimeError, "third") + except RuntimeError as third: + caught = third + second = caught.__context__ + first = second.__cause__ + first.__context__ = caught # a loop: the walk must still end + _clear_frames(caught) + for exc in (caught, second, first): + assert "big" not in _innermost_locals(exc.__traceback__), exc + + +class _RecordingVAE(nn.Module): + """The stub VAE's decode, keeping the packed latents mflux hands it (the denoised result).""" + + def __init__(self): + super().__init__() + self.weight = mx.zeros((4,), dtype=mx.bfloat16) + self.seen = [] + + def decode_packed_latents(self, latents, tiling_config=None): + del tiling_config + self.seen.append(latents) + return mx.zeros((1, 3, 16 * latents.shape[-2], 16 * latents.shape[-1])) + + +class _VaryingTextEncoder(nn.Module): + """A text stack whose hidden states differ in direction from token to token (the stub's are one vector scaled by + the id, which a norm maps to the same vector for every prompt); width 3072, as build_text_batch pads.""" + + def __init__(self): + super().__init__() + self.freq = mx.arange(1, 3073, dtype=mx.float32) / 7 + + def __call__(self, input_ids, attention_mask=None): + return mx.sin(input_ids[..., None].astype(mx.float32) * self.freq).astype(mx.bfloat16) + + +def _stock_ernie(tokenizer, vae, matrices, extras): + """Stock mflux ``ErnieImage`` over a plain tiny ``ErnieTransformer`` holding the BF16 source of a tiny checkpoint + (``write_tiny_checkpoint``'s matrices and extras in mflux's layout).""" + from mflux.models.common.config.model_config import ModelConfig + from mflux.models.ernie_image.ernie_image_initializer import ErnieImageInitializer + from mflux.models.ernie_image.model.ernie_transformer.transformer import ErnieTransformer + from mflux.models.ernie_image.variants.txt2img.ernie_image import ErnieImage + from mlx.utils import tree_flatten, tree_unflatten + from tests._ernie_tiny import TINY + + transformer = ErnieTransformer(**TINY) + shapes = {name: tuple(p.shape) for name, p in tree_flatten(transformer.parameters())} + weights = [] + for group, subs in matrices.items(): + for sub, matrix in subs.items(): + # Block groups key their matrices by sub-path; the non-block groups by mflux's parameter name. + name = sub if sub.endswith(".weight") else f"{group}.{sub}.weight" + weights.append((name, mx.array(matrix).view(mx.bfloat16))) + weights += [(name, mx.array(bits).view(mx.bfloat16)) for name, bits in extras.items()] + assert sorted(name for name, _w in weights) == sorted(shapes) + assert all(shapes[name] == tuple(w.shape) for name, w in weights) + transformer.update(tree_unflatten(weights)) + stock = ErnieImage.__new__(ErnieImage) + nn.Module.__init__(stock) + ErnieImageInitializer._init_config( + stock, ModelConfig.from_name(model_name="ernie-image", base_model=None) + ) + stock.vae = vae + stock.text_encoder = _VaryingTextEncoder() + stock.tokenizers = {"ernie": tokenizer} + stock.transformer = transformer + stock.bits = None + stock.lora_paths, stock.lora_scales = [], [] + return stock + + +@pytest.mark.parametrize("guidance", [1.0, 4.0], ids=["guidance-1", "cfg"]) +def test_a_tiny_generation_gives_stock_mfluxs_latents_bit_for_bit(tmp_path, monkeypatch, guidance): + # Bug caught (T1-T3): anything between the checkpoint and mflux's loop changing the result: a decoded matrix on the + # wrong parameter, a non-block group's matrices swapped, an extra mis-routed, the negative prompt dropped from the + # CFG batch or the batch built in the other order (the negative prompt is 12 ids here, the prompt 8, so the text + # lengths differ too). The reference is stock mflux over the same BF16 weights, its step 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 mflux.utils.apple_silicon import AppleSiliconUtil + from tests._ernie_tiny import TINY, StubTokenizer, stub_components + + from mlx_dfloat.format import open_checkpoint + from mlx_dfloat.mflux.ernie.transformer import build_transformer + + rng = np.random.default_rng(11) + matrices, extras, _conv = write_tiny_checkpoint(tmp_path / "random", rng, random_extras=rng) + ckpt = open_checkpoint(tmp_path / "random") + tokenizer = StubTokenizer({"n": 12, "p": 8}) + ours_vae, stock_vae = _RecordingVAE(), _RecordingVAE() + components = stub_components(tokenizer) + components.vae, components.text_encoder = ours_vae, _VaryingTextEncoder() + model = fake_model( + tmp_path, + monkeypatch, + model="ernie-image", + ckpt=ckpt, + build=build_transformer(ckpt, transformer_overrides=TINY), + components=components, + ) + stock = _stock_ernie(tokenizer, stock_vae, matrices, extras) + kwargs = { + "seed": 3, + "prompt": "p", + "num_inference_steps": 2, + "height": 64, + "width": 64, + "guidance": guidance, + "negative_prompt": "n", + "scheduler": "linear", + } + model.generate_image(**kwargs) + # Stock mflux runs its step uncompiled only on a base or Pro M1/M2: answer that, so both sides run the same graph. + monkeypatch.setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: True)) + stock.generate_image(**kwargs) + (ours,), (theirs,) = ours_vae.seen, stock_vae.seen + assert model.report()["cfg_batch"] == (2 if guidance > 1.0 else 1) + assert bool(mx.all(mx.isfinite(theirs))) + assert (ours.dtype, ours.shape) == (theirs.dtype, theirs.shape) + view = mx.uint32 if ours.dtype == mx.float32 else mx.uint16 + assert np.array_equal(np.array(ours.view(view)), np.array(theirs.view(view))) + + +def test_encode_then_generate_on_the_base_reuses_the_batch(tmp_path, monkeypatch): + # Bug caught (C1): encode("p") keyed at guidance 1.0 while generate_image(seed, "p") runs the base's default 4.0 + # (CFG: the batch [" ", "p"]), so the call drops the set and reloads the encoder for a batch encode() was meant to + # prepare. + model = fake_model(tmp_path, monkeypatch, model="ernie-image") + _patch_upstream_generate(monkeypatch) + model.encode("p") + _no_encoder(monkeypatch) + model.generate_image(seed=1, prompt="p", height=64, width=64) + counters = model._lifecycle.counters + assert list(model._batches) == ['[" ", "p"]'] + assert (counters.encoder_loads, counters.forced_set_drops) == (0, 0) + + +@pytest.mark.parametrize("prompt", ["", " "]) +def test_encode_refuses_an_empty_or_blank_prompt(tmp_path, monkeypatch, prompt): + # Bug caught (C3): encode() taking a prompt generate_image refuses (an empty batch cached, an encoder load spent). + model = fake_model(tmp_path, monkeypatch) + monkeypatch.setattr(model._lifecycle, "ensure_embeddings", lambda *p: pytest.fail("encoded")) + with pytest.raises( + DFloatUnsupportedError, match="an empty or blank prompt is refused as a user error" + ): + model.encode(prompt) + + +def test_encode_runs_under_the_commands_caps_and_restores_mlxs_defaults(tmp_path, monkeypatch): + # Bug caught (X4): encode() from Python run at MLX's default wired limit 0 (a prompt encode is one of the measured + # phases, sized under the command's caps), or the caps left installed after it returns. + from tests._mlx_limits import command_caps, current_limits, mlx_without_wired_cap + + model = fake_model(tmp_path, monkeypatch) + seen = [] + monkeypatch.setattr( + model._lifecycle, "ensure_embeddings", lambda *prompts: seen.append(current_limits()) + ) + wired, memory = command_caps() + assert wired > 0 + with mlx_without_wired_cap() as start: + model.encode("p") + after = current_limits() + assert (seen[0]["wired"], seen[0]["memory"]) == (wired, memory) + assert after == start diff --git a/tests/test_ernie_names.py b/tests/test_ernie_names.py new file mode 100644 index 0000000..e1f7eed --- /dev/null +++ b/tests/test_ernie_names.py @@ -0,0 +1,276 @@ +"""ERNIE-Image naming: the map derived from mflux's weight mapping, and the DF11 group check.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +from tests._df11_fixtures import random_bf16, write_checkpoint + +from mlx_dfloat.errors import DFloatFormatError, DFloatIntegrationError +from mlx_dfloat.format import open_checkpoint +from mlx_dfloat.mflux._mapping import name_map_from_mapping +from mlx_dfloat.mflux.ernie.names import KIND, MATRIX_SUBS, NONBLOCK_GROUPS, check_ernie_groups + +SEVEN = ( + "self_attention.to_q", + "self_attention.to_k", + "self_attention.to_v", + "self_attention.to_out.0", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.linear_fc2", +) + + +def _conv(w): # stands in for WeightTransforms.transpose_conv2d_weight + return w.transpose(0, 2, 3, 1) + + +def _t(to, frm, transform=None): + return SimpleNamespace(to_pattern=to, from_pattern=[frm], transform=transform) + + +def _mapping(): + # A copy of ErnieWeightMapping.get_transformer_mapping() (mflux 0.20.0 ernie_weight_mapping.py:11-109). + targets = [ + _t("adaln_modulation.weight", "adaLN_modulation.1.weight"), + _t("adaln_modulation.bias", "adaLN_modulation.1.bias"), + _t("x_embedder.proj.weight", "x_embedder.proj.weight", _conv), + _t("x_embedder.proj.bias", "x_embedder.proj.bias"), + ] + targets += [ + _t(n, n) + for n in ( + "text_proj.weight", + "time_embedding.linear_1.weight", + "time_embedding.linear_1.bias", + "time_embedding.linear_2.weight", + "time_embedding.linear_2.bias", + "final_norm.linear.weight", + "final_norm.linear.bias", + "final_linear.weight", + "final_linear.bias", + ) + ] + for sub in ( + "adaLN_sa_ln", + "adaLN_mlp_ln", + "self_attention.to_q", + "self_attention.to_k", + "self_attention.to_v", + "self_attention.to_out.0", + "self_attention.norm_q", + "self_attention.norm_k", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.linear_fc2", + ): + src = f"layers.{{layer}}.{sub}.weight" + targets.append(_t(src, src)) + return targets + + +# The DF11 config's pattern_dict (mingyi456/ERNIE-Image{,-Turbo}-DF11 config.json, byte-identical in both). +PATTERN = { + "time_embedding": ("linear_1", "linear_2"), + "adaLN_modulation.1": (), + r"layers\.\d+": SEVEN, + "final_norm.linear": (), +} + + +def test_the_map_places_every_block_matrix_on_its_own_path_in_decoded_order(): + # Bug caught: MATRIX_SUBS out of the checkpoint's pattern_dict order (gate_proj and up_proj swapped: the block + # would run up(x) * gelu(up(x)) on the wrong halves), or a block sub-path renamed. + m = name_map_from_mapping(_mapping(), matrix_subs=MATRIX_SUBS) + assert m.attrs_of(KIND) == SEVEN + assert m.place("layers.35.mlp.up_proj.weight").attr == "mlp.up_proj" + assert m.param_name("layers.3.adaLN_sa_ln.weight") == "layers.3.adaLN_sa_ln.weight" + + +def test_adaln_modulation_is_renamed_and_the_patch_conv_keeps_its_transform(): + # Bug caught: the one non-block rename lost (adaLN_modulation.1.weight has no parameter of that name in mflux: + # an extras-coverage refusal), or the Conv2d transpose dropped (a shape refusal at load, or a wrong layout). + m = name_map_from_mapping(_mapping(), matrix_subs=MATRIX_SUBS) + assert m.param_name("adaLN_modulation.1.weight") == "adaln_modulation.weight" + assert m.param_name("adaLN_modulation.1.bias") == "adaln_modulation.bias" + assert m.transform_of("x_embedder.proj.weight") is _conv + + +def test_the_adapters_tables_are_the_checkpoints_matrix_names(): + # Bug caught: a sub-path missing, misspelt or reordered against the published checkpoint (the group check would + # refuse every real ERNIE checkpoint, or accept the matrices in the wrong order), or a non-block group's matrices + # not what the reader derives from the pattern ("adaLN_modulation.1": [] is the group's own weight). + from mlx_dfloat.format import matrix_names_for + + assert matrix_names_for("layers.7", PATTERN) == tuple( + f"layers.7.{s}.weight" for s in MATRIX_SUBS[KIND] + ) + assert NONBLOCK_GROUPS == { + "time_embedding": ("time_embedding.linear_1.weight", "time_embedding.linear_2.weight"), + "adaLN_modulation.1": ("adaLN_modulation.1.weight",), + "final_norm.linear": ("final_norm.linear.weight",), + } + for group, matrices in NONBLOCK_GROUPS.items(): + assert matrix_names_for(group, PATTERN) == matrices + + +def test_a_sub_path_mflux_does_not_load_is_refused(): + # Bug caught: MATRIX_SUBS naming a matrix mflux's mapping has no target for (it would stay a placeholder). + targets = [t for t in _mapping() if t.from_pattern != ["layers.{layer}.mlp.up_proj.weight"]] + with pytest.raises( + DFloatIntegrationError, match=r"mlp\.up_proj: a compressed matrix with no mflux target" + ): + name_map_from_mapping(targets, matrix_subs=MATRIX_SUBS) + + +# --- the group check, on written tiny checkpoints ------------------------------------------------------------------- + +_BLOCK = r"layers\.\d+" +_NONBLOCK_PATTERNS = { + r"time_embedding": ("linear_1", "linear_2"), + r"adaLN_modulation\.1": (), + r"final_norm\.linear": (), +} + + +def _write(root, *, blocks=(0, 1), subs=SEVEN, nonblock=True, extra_groups=()): + rng = np.random.default_rng(3) + groups = {f"layers.{b}": [random_bf16(rng, (4, 4)) for _ in subs] for b in blocks} + patterns = {_BLOCK: subs} + if nonblock: + groups["time_embedding"] = [random_bf16(rng, (4, 4)), random_bf16(rng, (4, 4))] + groups["adaLN_modulation.1"] = [random_bf16(rng, (8, 4))] + groups["final_norm.linear"] = [random_bf16(rng, (8, 4))] + patterns.update(_NONBLOCK_PATTERNS) + for name in extra_groups: + groups[name] = [random_bf16(rng, (4, 4))] + patterns[name.replace(".", r"\.")] = () + write_checkpoint(root, groups=groups, patterns=patterns) + return open_checkpoint(root) + + +@pytest.mark.parametrize("nonblock", [False, True]) +def test_the_block_count_is_returned_with_and_without_the_nonblock_groups(tmp_path, nonblock): + # Bug caught: an optional non-block group counted as a block (or refused when present or absent), or the count + # off by one. + assert check_ernie_groups(_write(tmp_path, nonblock=nonblock)) == {"layers": 2} + + +def test_a_hole_in_the_block_indices_is_refused(tmp_path): + # Bug caught: layers.0 and layers.2 accepted, layers.1 left on placeholders (a zero-size matmul at run time). + with pytest.raises(DFloatFormatError, match="layers groups are not contiguous"): + check_ernie_groups(_write(tmp_path, blocks=(0, 2))) + + +def test_a_zero_padded_block_index_is_refused(tmp_path): + # Bug caught: layers.01 counted as block 1 (int("01") == 1), then the seam asks the provider for layers.1, a + # group the checkpoint does not have: a failure at the first block after the encoder and set loads. + with pytest.raises(DFloatFormatError, match=r"layers\.01: not an ERNIE-Image group"): + check_ernie_groups(_write(tmp_path, blocks=("0", "01"))) + + +def test_a_checkpoint_without_blocks_is_refused(tmp_path): + # Bug caught: a non-block-only checkpoint accepted as zero blocks. + with pytest.raises(DFloatFormatError, match="no layers groups in the checkpoint"): + check_ernie_groups(_write(tmp_path, blocks=())) + + +def test_a_stray_group_is_refused_naming_the_accepted_nonblock_groups(tmp_path): + # Bug caught: an unknown group (another model's) silently ignored, its matrices never placed. + with pytest.raises( + DFloatFormatError, + match=r"cap_embedder: not an ERNIE-Image group .*time_embedding, adaLN_modulation\.1, final_norm\.linear", + ): + check_ernie_groups(_write(tmp_path, extra_groups=("cap_embedder",))) + + +def test_a_time_embedding_group_with_one_matrix_is_refused(tmp_path): + # Bug caught: a non-block group of the wrong matrices accepted; time_embedding.linear_2 never installed (a + # zero-size placeholder in the timestep MLP before the first block). + rng = np.random.default_rng(5) + write_checkpoint( + tmp_path, + groups={ + "layers.0": [random_bf16(rng, (4, 4)) for _ in SEVEN], + "time_embedding": [random_bf16(rng, (4, 4))], + }, + patterns={_BLOCK: SEVEN, "time_embedding": ("linear_1",)}, + ) + with pytest.raises(DFloatFormatError, match=r"time_embedding: holds"): + check_ernie_groups(open_checkpoint(tmp_path)) + + +def test_a_zimage_block_under_the_same_group_name_is_refused_naming_the_seven_matrices(tmp_path): + # Bug caught (Review Focus 3): a Z-Image DF11 checkpoint (its main blocks are also named layers.) passed for + # ERNIE and failing with a shape error at the first block instead of at build. + zimage = ( + "attention.to_q", + "attention.to_k", + "attention.to_v", + "attention.to_out.0", + "feed_forward.w1", + "feed_forward.w2", + "feed_forward.w3", + "adaLN_modulation.0", + ) + with pytest.raises( + DFloatFormatError, match=r"layers\.0: holds 8 matrices .*self_attention\.to_q" + ): + check_ernie_groups(_write(tmp_path, subs=zimage, nonblock=False)) + + +@pytest.mark.mflux +def test_the_real_mapping_matches_the_copy_and_every_matrix_lands_on_a_linear(): + # Bug caught: mflux's real mapping no longer matching the copy (a rename or transform added upstream), or a + # decoded matrix with no nn.Linear to land on in a real tiny ErnieTransformer. + import mlx.nn as nn + from mflux.models.ernie_image.model.ernie_transformer.transformer import ErnieTransformer + from mflux.models.ernie_image.weights.ernie_weight_mapping import ErnieWeightMapping + from tests._ernie_tiny import TINY + + from mlx_dfloat.integrate.placeholders import get_attr_path + from mlx_dfloat.mflux.ernie.names import ernie_name_map + + real = [ + (t.to_pattern, list(t.from_pattern), t.transform is not None) + for t in ErnieWeightMapping.get_transformer_mapping() + ] + copy = [(t.to_pattern, list(t.from_pattern), t.transform is not None) for t in _mapping()] + assert sorted(real) == sorted(copy) + m = ernie_name_map() + assert m.param_name("adaLN_modulation.1.weight") == "adaln_modulation.weight" + tf = ErnieTransformer(**TINY) + for attr in m.attrs_of(KIND): + assert isinstance(get_attr_path(tf.layers[0], attr), nn.Linear), attr + for matrices in NONBLOCK_GROUPS.values(): + for name in matrices: + module = get_attr_path(tf, m.param_name(name).removesuffix(".weight")) + assert isinstance(module, nn.Linear), name + + +@pytest.mark.mflux +def test_the_default_layer_count_is_mfluxs_constructor_default(): + # Bug caught: an mflux bump changing ErnieTransformer's default depth (what the published model is built with) + # while the count check keeps 36: every real checkpoint refused, or one of another depth accepted. + import inspect + + from mflux.models.ernie_image.model.ernie_transformer.transformer import ErnieTransformer + + from mlx_dfloat.mflux.ernie.names import DEFAULT_LAYERS + + default = inspect.signature(ErnieTransformer.__init__).parameters["num_layers"].default + assert default == 36 + assert DEFAULT_LAYERS == 36 + + +@pytest.mark.parametrize("module", ["mlx_dfloat.mflux.ernie", "mlx_dfloat.mflux.ernie.names"]) +def test_importing_the_adapter_does_not_import_mflux(module): + # Bug caught: a module-level mflux import in the adapter (a user without the extra gets a bare ImportError on + # import instead of the dependency error naming it, and every import pays for mflux). + import subprocess + import sys + + code = f"import sys, {module}; print('mflux' in sys.modules)" + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert out.stdout.strip() == "False" diff --git a/tests/test_ernie_tiny.py b/tests/test_ernie_tiny.py new file mode 100644 index 0000000..5a48e2e --- /dev/null +++ b/tests/test_ernie_tiny.py @@ -0,0 +1,55 @@ +"""One real tiny ERNIE-Image generation through mflux's own loop: our prelude, the uncompiled predict, the batch +cache, the seam, the VAE guard and mflux's image assembly composing. `@pytest.mark.mflux`.""" + +import mlx.core as mx +import numpy as np +import pytest +from tests._ernie_tiny import fake_model + +pytestmark = pytest.mark.mflux + + +def _cache_limit_in_force(): + previous = mx.set_cache_limit(0) + mx.set_cache_limit(previous) + return previous + + +@pytest.mark.parametrize(("guidance", "batch"), [(4.0, 2), (1.0, 1)], ids=["cfg", "guidance-1"]) +def test_one_tiny_generation_decodes_each_block_once_per_step( + tmp_path, monkeypatch, guidance, batch +): + # Bug caught: anything between our prelude and mflux's loop not composing (the uncompiled predict, the batch + # cache, the VAE guard), or CFG split into two transformer calls (8 decodes instead of 4). The chip check answers + # "not M1/M2" so a lost bypass is red on any host. Tiny transformer: 2 blocks; 2 steps: 2 x 2 = 4 decodes either + # way (CFG is one batch-2 call, ernie_image.py:239-250). + from mflux.utils.apple_silicon import AppleSiliconUtil + + monkeypatch.setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: False)) + model = fake_model(tmp_path, monkeypatch) + seen = {} + + class Probe: + def call_in_loop(self, t, seed, prompt, latents, config, time_steps): + seen.setdefault("limits", []).append(_cache_limit_in_force()) + + def call_after_loop(self, seed, prompt, latents, config): + seen["latents"] = np.array(latents.astype(mx.float32)) + + model.callbacks.register(Probe()) + before = _cache_limit_in_force() + image = model.generate_image( + seed=1, prompt="a", num_inference_steps=2, height=64, width=64, guidance=guidance + ) + assert np.asarray(image.image).shape == (64, 64, 3) + # The denoised latents at 64²: 128 channels x 4 x 4 (ernie_latent_creator.py:11-17), finite. + assert seen["latents"].shape == (1, 128, 4, 4) + assert np.isfinite(seen["latents"]).all() + report = model.report() + assert report["decode_launches"] == 4 + assert report["cfg_batch"] == batch + assert report["predict"] == "uncompiled" + assert seen["limits"] == [report["cache_limit_in_force"]] * 2 # in force in the loop + assert set(report["peaks"]) >= {"encode", "set_load", "denoise", "vae"} + assert model.prompt_cache == {} + assert _cache_limit_in_force() == before diff --git a/tests/test_ernie_transformer.py b/tests/test_ernie_transformer.py new file mode 100644 index 0000000..c35eca4 --- /dev/null +++ b/tests/test_ernie_transformer.py @@ -0,0 +1,320 @@ +"""The ERNIE-Image transformer: the class-swap seam and the compile bypass, on a real tiny mflux transformer.""" + +import importlib.util +import types +from types import SimpleNamespace + +import mlx.core as mx +import numpy as np +import pytest +from mlx.utils import tree_flatten, tree_unflatten +from tests._ernie_tiny import TINY, tiny_groups, tiny_inputs, tiny_predict, tiny_seamed + +from mlx_dfloat.decode import decode_group +from mlx_dfloat.errors import DFloatDependencyError +from mlx_dfloat.integrate import seam +from mlx_dfloat.integrate.placeholders import get_attr_path +from mlx_dfloat.integrate.providers import DF11Provider, ResidentProvider +from mlx_dfloat.mflux._compile import uncompiled +from mlx_dfloat.mflux.ernie.transformer import block_lists, seam_transformer_class + +COMPILED = type( + mx.compile(lambda x: x) +) # mlx.gc_func, a FunctionType subclass: only `type(f) is` tells them apart +BLOCKS = ["layers.0", "layers.1"] + + +@pytest.fixture +def not_m1(monkeypatch): + """mflux's chip check answers "not a base or Pro M1/M2": stock mflux compiles ERNIE's predict there. + + On this M1 Max the unpatched check is already False; CI's base M1/M2 runners answer True, where mflux would hand + back the plain function without our override and a bypass test would pass vacuously. + """ + from mflux.utils.apple_silicon import AppleSiliconUtil + + monkeypatch.setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: False)) + return AppleSiliconUtil + + +def _stock_args(inputs): + from mflux.models.ernie_image.model.ernie_transformer.transformer import ErnieTransformer + + tf = ErnieTransformer(**TINY) + return (tf, inputs["text_bth"], inputs["text_lens"], inputs["hidden_states"]) + + +def _random_weights(shapes, seed=11): + """Every block matrix its own random bfloat16 values (small, so the tiny forward stays finite).""" + rng = np.random.default_rng(seed) + return { + block: { + attr: mx.array((rng.standard_normal(shape) * 0.05).astype(np.float32)).astype( + mx.bfloat16 + ) + for attr, shape in per.items() + } + for block, per in shapes.items() + } + + +def _constant_weights(shapes): + """Block i's matrices filled with the constant i + 1.""" + return { + block: { + attr: mx.full(shape, float(i + 1), dtype=mx.bfloat16) for attr, shape in per.items() + } + for i, (block, per) in enumerate(shapes.items()) + } + + +def _call_inputs(inputs): + """``ErnieTransformer.__call__`` keyword inputs from ``tiny_inputs``: the timestep is sigma x 1000, as mflux.""" + return { + "hidden_states": inputs["hidden_states"], + "timestep": mx.broadcast_to(inputs["sigma"] * 1000, (inputs["text_bth"].shape[0],)), + "text_bth": inputs["text_bth"], + "text_lens": inputs["text_lens"], + } + + +def test_block_lists_is_the_one_list_the_forward_runs(): + # Bug caught: another attribute than the list __call__ iterates (transformer.py:157-159), so no block would run + # through the seam. + tf = SimpleNamespace(layers=[1, 2]) + assert block_lists(tf) == [("layers", [1, 2])] + + +def test_the_seamed_class_needs_the_mflux_extra(monkeypatch): + # Bug caught: the adapter importing mflux without the guard (a bare ImportError instead of the package's own + # dependency error naming the extra). The class factory is cached, so the cache is emptied around the call. + real = importlib.util.find_spec + monkeypatch.setattr( + importlib.util, + "find_spec", + lambda name, *a, **k: None if name == "mflux" else real(name, *a, **k), + ) + seam_transformer_class.cache_clear() + try: + with pytest.raises(DFloatDependencyError, match=r"mlx-dfloat\[mflux\]"): + seam_transformer_class() + finally: + seam_transformer_class.cache_clear() + + +@pytest.mark.mflux +def test_off_m1_m2_the_stock_factory_compiles_and_on_them_it_does_not(not_m1, monkeypatch): + # Bug caught (the control): the reason for the bypass gone (an mflux change: ernie_image.py:252-254 no longer + # compiling off M1/M2), or a test of the bypass that passes vacuously on a base M1/M2 CI runner (where the + # unpatched chip check is True). On this M1 Max the unpatched check is already False: stock mflux compiles here. + from mflux.models.ernie_image.variants.txt2img.ernie_image import ErnieImage + + args = _stock_args(tiny_inputs()) + assert type(ErnieImage._predict(*args)) is COMPILED + monkeypatch.setattr(not_m1, "is_m1_or_m2", classmethod(lambda cls: True)) + assert type(ErnieImage._predict(*args)) is types.FunctionType + + +@pytest.mark.mflux +def test_uncompiled_returns_mfluxs_plain_predict_on_any_chip(not_m1): + # Bug caught: the chip override not reaching ErnieImage._predict (the seam's per-block mx.eval would then run + # inside mx.compile), or leaking past the factory call (every later mflux factory would run uncompiled). + from mflux.models.ernie_image.variants.txt2img.ernie_image import ErnieImage + + fn = uncompiled(ErnieImage._predict, *_stock_args(tiny_inputs())) + assert type(fn) is types.FunctionType + assert not_m1.is_m1_or_m2() is False + + +@pytest.mark.mflux +def test_the_compiled_predict_cannot_run_the_seam(): + # Bug caught: none in our code; it pins why the bypass exists (if the compiled predict could run per-block evals, + # the bypass and its cost on Max/Ultra and M3+ chips would be unjustified). + tf, shapes = tiny_seamed() + tf.attach(ResidentProvider(_constant_weights(shapes)), shapes) + with pytest.raises(ValueError, match=r"\[eval\] Attempting to eval"): + mx.eval(tiny_predict(tf, tiny_inputs(), compiled=True)) + + +@pytest.mark.mflux +def test_the_swap_keeps_parameter_paths_and_nothing_is_computed(): + # Bug caught: the class swap breaking on ErnieTransformerBlock or moving parameters under another path; a + # parameter mflux computes at construction added upstream (no checkpoint holds it: every build would be refused by + # the extras coverage), or a stale exemption hiding a parameter no checkpoint fills. + from mflux.models.ernie_image.model.ernie_transformer.transformer import ErnieTransformer + from mflux.models.ernie_image.model.ernie_transformer.transformer_block import ( + ErnieTransformerBlock, + ) + from mflux.models.ernie_image.weights.ernie_weight_mapping import ErnieWeightMapping + + from mlx_dfloat.integrate.blockseam import seam_blocks + from mlx_dfloat.mflux.ernie.transformer import COMPUTED_PARAMS + + 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) == ("layers.0", "layers.1") + after = sorted(k for k, _v in tree_flatten(tf.parameters())) + assert after == before + assert isinstance(tf.layers[1], ErnieTransformerBlock) + + params = {k for k, _v in tree_flatten(ErnieTransformer(**TINY).parameters())} + targets = set() + for t in ErnieWeightMapping.get_transformer_mapping(): + if "{layer}" in t.to_pattern: + targets |= {t.to_pattern.replace("{layer}", str(i)) for i in range(2)} + else: + targets.add(t.to_pattern) + # mflux 0.20.0 computes nothing at construction: every parameter is a mapping target. + assert params - targets == set() + assert len(COMPUTED_PARAMS) == 0 + + +def _eager_reference(tf, shapes, weights): + """Stock mflux holding tf's non-block parameters and ``weights`` on its blocks.""" + from mflux.models.ernie_image.model.ernie_transformer.transformer import ErnieTransformer + + plain = ErnieTransformer(**TINY) + matrices = {f"{b}.{a}.weight" for b, per in shapes.items() for a in per} + plain.update( + tree_unflatten([(k, v) for k, v in tree_flatten(tf.parameters()) if k not in matrices]) + ) + for block, per in weights.items(): + idx = int(block.partition(".")[2]) + for attr, w in per.items(): + get_attr_path(plain.layers[idx], attr).weight = w + return plain + + +@pytest.mark.mflux +def test_a_forward_runs_each_block_with_its_own_weights_and_ends_on_placeholders(monkeypatch): + # Bug caught: a block run on another block's weights, or a matrix landing on another attribute of its own block + # (to_q and to_k swapped: same shape, so only per-matrix values show it), either way the output differs from + # stock mflux holding the same weights; a block output not evaluated at its boundary; or the placeholders not + # restored afterwards. + tf, shapes = tiny_seamed() + weights = _random_weights(shapes) + evals = [] + real_eval = seam._eval + + def recording_eval(*args): + evals.append(args[0]) + real_eval(*args) + + monkeypatch.setattr(seam, "_eval", recording_eval) + tf.attach(ResidentProvider(weights), shapes) + inputs = _call_inputs(tiny_inputs()) + out = tf(**inputs) + mx.eval(out) + tf.verify_step() + want = _eager_reference(tf, shapes, weights)(**inputs) + mx.eval(want) + + assert len(evals) == 2 + # Each block returns one array: the joint [16 image | 8 text] sequence at hidden size 32. + assert [tuple(e.shape) for e in evals] == [(1, 24, 32), (1, 24, 32)] + assert out.shape == (1, 128, 4, 4) + assert bool(mx.isfinite(out).all().item()) + # float32: the tiny transformer's non-block parameters are mflux's float32 init (the published ones are BF16). + assert out.dtype == want.dtype == mx.float32 + assert mx.array_equal(out.view(mx.uint32), want.view(mx.uint32)).item() + for block, per in shapes.items(): + idx = int(block.partition(".")[2]) + for attr in per: + assert get_attr_path(tf.layers[idx], attr).weight.size == 0, (block, attr) + + +@pytest.mark.mflux +def test_depth2_drains_the_last_block(monkeypatch): + # Bug caught: end_step not draining the tail (the last block's output never evaluated inside the step), or the + # 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() + events = [] + real_eval, real_async = seam._eval, seam._async_eval + + def recording_eval(*args): + events.append(("eval", args[0])) + real_eval(*args) + + def recording_async(*args): + events.append(("async", args[0])) + real_async(*args) + + monkeypatch.setattr(seam, "_eval", recording_eval) + monkeypatch.setattr(seam, "_async_eval", recording_async) + tf.attach(ResidentProvider(_constant_weights(shapes)), shapes, eval_policy="depth2") + mx.eval(tf(**_call_inputs(tiny_inputs()))) + tf.verify_step() + + outputs = [] + for _kind, obj in events: + if not any(obj is seen for seen in outputs): + outputs.append(obj) + order = [(kind, next(i for i, o in enumerate(outputs) if o is obj)) for kind, obj in events] + assert order == [("async", 0), ("async", 1), ("eval", 0), ("eval", 1)] + + +@pytest.mark.mflux +def test_a_cfg_batch_decodes_each_block_once(): + # Bug caught: CFG split into two transformer calls (four decodes, as Klein and Qwen do), or a block decoded twice + # in one call. mflux runs a CFG step as one batch-2 call (ernie_image.py:239-250): 2 decodes per step. + from mlx_dfloat.mflux.ernie.names import ernie_name_map + + tf, shapes = tiny_seamed() + groups, names, _src = tiny_groups(shapes, np.random.default_rng(9)) + calls = [] + + def count(group): + calls.append(group.name) + return decode_group(group, backend="reference") + + tf.attach( + DF11Provider(groups, names, ernie_name_map(), decode=count), shapes, verify_in_call=True + ) + out = tiny_predict(tf, tiny_inputs(batch=2), guidance=4.0) + mx.eval(out) + assert calls == BLOCKS + assert out.shape == (1, 128, 4, 4) + assert bool(mx.isfinite(out).all().item()) + + +def test_importing_the_transformer_module_does_not_import_mflux(): + # Bug caught: a module-level mflux import in the transformer module (a user without the extra gets a bare + # ImportError on import). + import subprocess + import sys + + code = "import sys, mlx_dfloat.mflux.ernie.transformer; print('mflux' in sys.modules)" + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert out.stdout.strip() == "False" + + +@pytest.mark.mflux +def test_a_cfg_batch_through_the_seam_gives_stock_mfluxs_prediction_bit_for_bit(monkeypatch): + # Bug caught (T2): the batch-2 call wrong through the seam where batch 1 is right: the two prompts' rows mixed, + # their unequal text lengths (5 and 8 tokens, masked differently) mishandled, or the guidance combination fed the + # wrong half. The reference is mflux's own predict over a stock transformer holding the same weights, both + # uncompiled; equality is bit for bit. + from mflux.models.ernie_image.variants.txt2img.ernie_image import ErnieImage + from mflux.utils.apple_silicon import AppleSiliconUtil + + tf, shapes = tiny_seamed() + weights = _random_weights(shapes) + tf.attach(ResidentProvider(weights), shapes, verify_in_call=True) + inputs = {**tiny_inputs(batch=2), "text_lens": mx.array([5, 8], dtype=mx.int32)} + ours = tiny_predict(tf, inputs, guidance=4.0) + mx.eval(ours) + plain = _eager_reference(tf, shapes, weights) + monkeypatch.setattr(AppleSiliconUtil, "is_m1_or_m2", classmethod(lambda cls: True)) + predict = ErnieImage._predict( + plain, inputs["text_bth"], inputs["text_lens"], inputs["hidden_states"] + ) + assert type(predict) is types.FunctionType # stock, uncompiled like ours + theirs = predict( + inputs["hidden_states"], inputs["sigma"], inputs["text_bth"], inputs["text_lens"], 4.0 + ) + mx.eval(theirs) + assert ours.shape == theirs.shape == (1, 128, 4, 4) + assert ours.dtype == theirs.dtype + view = mx.uint32 if ours.dtype == mx.float32 else mx.uint16 + assert mx.array_equal(ours.view(view), theirs.view(view)).item() diff --git a/tests/test_mflux_families.py b/tests/test_mflux_families.py index b28e7c4..956ecb9 100644 --- a/tests/test_mflux_families.py +++ b/tests/test_mflux_families.py @@ -12,7 +12,7 @@ def test_importing_the_registry_does_not_import_mflux(): # Bug caught: a family module imported at registry import time (the CLI's --help would need mflux). code = "import sys, mlx_dfloat.mflux.families as f; assert 'mflux' not in sys.modules; print(len(f.MODELS))" out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) - assert out.stdout.strip() == "10" + assert out.stdout.strip() == "12" def test_resolving_an_adapter_module_and_importing_the_package_loads_no_mlx(): @@ -27,6 +27,19 @@ def test_resolving_an_adapter_module_and_importing_the_package_loads_no_mlx(): assert out.stdout.strip() == "False" +def test_resolving_the_ernie_coverage_specs_loads_no_mlx(): + # Bug caught: an eager import in mlx_dfloat/mflux/ernie/__init__.py. The mflux lane's + # --cov=mlx_dfloat.mflux.ernie.model / .init / .transformer resolve those specs (importing the ernie package) when + # coverage starts; MLX loaded there registers its native types twice and the lane aborts (exit 134). + code = ( + "import importlib.util, sys; " + "[importlib.util.find_spec(f'mlx_dfloat.mflux.ernie.{m}') for m in ('model', 'init', 'transformer')]; " + "import mlx_dfloat.mflux; print('mlx.core' in sys.modules)" + ) + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert out.stdout.strip() == "False" + + def test_resolving_the_qwen21_coverage_specs_loads_no_mlx(): # Bug caught: an eager import in mlx_dfloat/mflux/qwen21/__init__.py. The mflux lane's # --cov=mlx_dfloat.mflux.qwen21.model / .init resolve those specs (importing the qwen21 package) when coverage @@ -87,6 +100,12 @@ def test_an_unknown_name_is_refused_with_the_known_ones(): # cli/parser/parsers.py:180 (scheduler "linear"); qwen_image_21.py:89 (a second call per step above 1.0 with # a negative prompt, which mflux takes from --negative-prompt, parsers.py:177). ("qwen-image-2.1", "qwen21", 40, 1.0, "linear", True, True), + # mflux 0.20.0 cli/defaults/defaults.py:35-36 (50 and 8 steps); ernie_image_generate.py:22 (guidance 4.0) and + # ernie_image_turbo_generate.py:42-45 (1.0 only); "linear" (ernie_image_generate.py:33-34, + # ernie_image.py:64-65). CFG is one batch-2 call (ernie_image.py:239-250), so never a second call; Turbo's + # command ignores --negative-prompt (ernie_image_turbo_generate.py:10-12). + ("ernie-image", "ernie", 50, 4.0, "linear", False, True), + ("ernie-image-turbo", "ernie", 8, 1.0, "linear", False, False), ], ) def test_the_entries_carry_the_mflux_defaults( @@ -181,6 +200,9 @@ def test_only_the_distilled_klein_models_fix_their_guidance(): "flux2-klein-base-9b": None, "flux2-klein-9b": 1.0, "qwen-image-2.1": None, + # mflux 0.20.0 ernie_image_turbo_generate.py:42-45: any guidance but 1.0 is an error for Turbo. + "ernie-image-turbo": 1.0, + "ernie-image": None, } @@ -202,15 +224,17 @@ def test_the_qwen_entry_names_its_published_repos_and_pinned_checkpoint(): ) -def test_the_qwen_base_is_pinned_to_the_revision_the_measured_runs_used_and_no_other_base_is(): - # Bug caught: the Qwen base (text encoder, VAE, tokenizer) following whatever lands on main instead of the snapshot - # the identity check and the MEASURED row ran on (d26bb61, bench/results/tiers/qwen-image-2.1-1024.json), or a pin - # set on a family whose base no recorded run names. - assert ( - families.entry("qwen-image-2.1").base_revision == "d26bb61231c349cf6b7896fa83353113880e1ba3" - ) - assert {n for n, e in families.MODELS.items() if e.base_revision is not None} == { - "qwen-image-2.1" +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 de-risk 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", } @@ -297,3 +321,43 @@ def test_the_family_loaders_return_the_real_classes_and_maps(): assert families.FAMILIES["qwen21"].load_model_class() is DFloatQwenImage21 # mflux 0.20.0 Qwen21Transformer._forward (qwen21_transformer.py:96-100): one block list. assert families.FAMILIES["qwen21"].load_name_map().kinds == ("transformer_blocks",) + + +def test_the_ernie_entries_name_their_published_repos_and_pinned_checkpoints(): + # Bug caught: a swapped repository or a pin copied onto the wrong model (the checkpoint verified here swapped for + # whatever lands on main). + got = { + n: (e.label, e.df11_repo, e.df11_revision, e.base_repo) + for n, e in families.MODELS.items() + if e.family == "ernie" + } + assert got == { + "ernie-image": ( + "ERNIE-Image", + "mingyi456/ERNIE-Image-DF11", + "c2dd30ad7dd5a928df2309581b282337f7cdb41f", + "baidu/ERNIE-Image", + ), + "ernie-image-turbo": ( + "ERNIE-Image-Turbo", + "mingyi456/ERNIE-Image-Turbo-DF11", + "27f84b44a3b78fcfaaadbaeeeae7cc7f7d75153b", + "baidu/ERNIE-Image-Turbo", + ), + } + + +@pytest.mark.mflux +def test_the_ernie_family_loads_its_class_and_one_kind_map(): + # Bug caught: the ernie loader pointing at another family's class or map, or the class's own MODELS drifting + # from the registry. + from mlx_dfloat.mflux.ernie.model import MODELS as ERNIE + from mlx_dfloat.mflux.ernie.model import DFloatErnieImage + + assert families.FAMILIES["ernie"].load_model_class() is DFloatErnieImage + assert families.FAMILIES["ernie"].load_name_map().kinds == ("layers",) + assert families.family_of("ernie-image") is families.FAMILIES["ernie"] + assert ERNIE == { + "ernie-image": ("mingyi456/ERNIE-Image-DF11", "baidu/ERNIE-Image"), + "ernie-image-turbo": ("mingyi456/ERNIE-Image-Turbo-DF11", "baidu/ERNIE-Image-Turbo"), + } diff --git a/tests/test_public_api.py b/tests/test_public_api.py index 6594ee0..87b83e3 100644 --- a/tests/test_public_api.py +++ b/tests/test_public_api.py @@ -91,3 +91,13 @@ def test_dfloat_qwen_image_21_is_reachable_lazily(): from mlx_dfloat.mflux.qwen21.model import DFloatQwenImage21 assert adapters.DFloatQwenImage21 is DFloatQwenImage21 + + +@pytest.mark.mflux +def test_dfloat_ernie_image_is_reachable_lazily(): + # Bug caught: DFloatErnieImage missing from the lazy attribute hook (AttributeError on first use), or added to + # __all__ (a star-import would then need mflux; the star-import test pins __all__). + import mlx_dfloat.mflux as adapters + from mlx_dfloat.mflux.ernie.model import DFloatErnieImage + + assert adapters.DFloatErnieImage is DFloatErnieImage From 0ebe02dd6632838ce99d6aefa9c04526e9f43672 Mon Sep 17 00:00:00 2001 From: Denis Ineshin Date: Thu, 8 Oct 2026 16:03:04 +0200 Subject: [PATCH 3/4] Extend the identity check, the real-blocks test and the CI lane to ERNIE-Image --- .github/workflows/ci.yml | 6 +- pyproject.toml | 2 + scripts/verify_image.py | 224 ++++++++++++++++++++++++++++++-- tests/test_ci_config.py | 28 +++- tests/test_ernie_real_blocks.py | 198 ++++++++++++++++++++++++++++ tests/test_verify_image.py | 127 ++++++++++++++++++ 6 files changed, 568 insertions(+), 17 deletions(-) create mode 100644 tests/test_ernie_real_blocks.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 98d163f..90763bc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,7 +54,7 @@ jobs: - 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'))" - # No combined threshold: init.py's tests run in the default lane (its live ones skip without a token); the next four steps gate each model.py on its own. + # No combined threshold: init.py's tests run in the default lane (its live ones skip without a token); the next five steps gate each model.py on its own. - name: Integration tests (mflux lane) env: # Every FLUX.1 base repository is gated on the Hub (schnell auto-approved), so the three @@ -62,7 +62,7 @@ jobs: # Only this step gets it: the actions, the sync and the collection and coverage steps have # no use for it. The secret is a read-only Hub token; the default lane runs without one. HF_TOKEN: ${{ secrets.HF_TOKEN }} - run: uv run --extra mflux --group dev pytest -m mflux --run-network --cov=mlx_dfloat.mflux.flux1.model --cov=mlx_dfloat.mflux.flux1.init --cov=mlx_dfloat.mflux.zimage.model --cov=mlx_dfloat.mflux.flux2.model --cov=mlx_dfloat.mflux.qwen21.model --cov=mlx_dfloat.mflux.qwen21.init --cov=mlx_dfloat.mflux.qwen21.transformer --cov-config=.coveragerc-integration --cov-report=term-missing -q -p no:cacheprovider -rs + run: uv run --extra mflux --group dev pytest -m mflux --run-network --cov=mlx_dfloat.mflux.flux1.model --cov=mlx_dfloat.mflux.flux1.init --cov=mlx_dfloat.mflux.zimage.model --cov=mlx_dfloat.mflux.flux2.model --cov=mlx_dfloat.mflux.qwen21.model --cov=mlx_dfloat.mflux.qwen21.init --cov=mlx_dfloat.mflux.qwen21.transformer --cov=mlx_dfloat.mflux.ernie.model --cov=mlx_dfloat.mflux.ernie.init --cov=mlx_dfloat.mflux.ernie.transformer --cov-config=.coveragerc-integration --cov-report=term-missing -q -p no:cacheprovider -rs - name: model.py coverage on its own (mflux lane) run: uv run --extra mflux --group dev coverage report --include='*/mflux/flux1/model.py' --fail-under=80 - name: zimage/model.py coverage on its own (mflux lane) @@ -71,6 +71,8 @@ jobs: run: uv run --extra mflux --group dev coverage report --include='*/mflux/flux2/model.py' --fail-under=80 - name: qwen21/model.py coverage on its own (mflux lane) run: uv run --extra mflux --group dev coverage report --include='*/mflux/qwen21/model.py' --fail-under=80 + - name: ernie/model.py coverage on its own (mflux lane) + run: uv run --extra mflux --group dev coverage report --include='*/mflux/ernie/model.py' --fail-under=80 - name: The lane collected mflux tests run: test "$(uv run --extra mflux --group dev pytest -m mflux --co -q -p no:cacheprovider | grep -c '::')" -ge 2 build: diff --git a/pyproject.toml b/pyproject.toml index 543f1aa..1e005f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -176,6 +176,8 @@ omit = [ "src/mlx_dfloat/mflux/flux2/model.py", # the same for Qwen-Image 2.1's model class: every test is @pytest.mark.mflux, gated in the integration lane. "src/mlx_dfloat/mflux/qwen21/model.py", + # the same for ERNIE-Image's model class: every test is @pytest.mark.mflux, gated in the integration lane. + "src/mlx_dfloat/mflux/ernie/model.py", ] [tool.coverage.report] diff --git a/scripts/verify_image.py b/scripts/verify_image.py index fd07d52..49799b7 100644 --- a/scripts/verify_image.py +++ b/scripts/verify_image.py @@ -1,6 +1,7 @@ """Image identity: DFloat11 latents against the same seam streaming the BF16 weights, for every BF16 family. -The ``df11`` side generates through ``DFloatFlux1``, ``DFloatZImage``, ``DFloatFlux2Klein`` or ``DFloatQwenImage21`` +The ``df11`` side generates through ``DFloatFlux1``, ``DFloatZImage``, ``DFloatFlux2Klein``, ``DFloatQwenImage21`` or +``DFloatErnieImage`` (the prompt encoded by its own encoders; the final latents captured by an after-loop callback and saved with the embeddings and the image). The ``bf16`` side builds the same seamed transformer with its extras read from the base's BF16 transformer shards, drives it with a ``StreamingBF16Provider`` (each block's matrices read from @@ -18,7 +19,9 @@ For Qwen-Image 2.1, ``--base`` must hold the BF16 transformer (``transformer/``: its index and both shards) for the bf16 side and ``--orchestrate``; the df11 side needs only the text encoder, the VAE and the DF11 file. Both Qwen sides run the forward pass uncompiled, and the bf16 side's copy of mflux's Qwen loop body is measurement glue in the same -sense as Klein's. +sense as Klein's. ERNIE-Image is checked the same way (``--base`` holds its two-shard BF16 transformer for the bf16 side; +both sides take mflux's step function uncompiled); its bf16 loop body is measurement glue too, so a drift fails as a +false mismatch, not a false pass. ERNIE-Image-Turbo runs at guidance 1.0 only and takes no ``--negative-prompt``. Usage (from the repository root of a synced checkout, ``--group bench``): uv run python -m scripts.verify_image --orchestrate --model schnell --df11 DIR --base DIR --out DIR \ [--prompt "..."] [--seed 42] [--steps 4] [--size 1024] [--guidance 3.5] [--negative-prompt "..."] [--eval-policy per-block] @@ -53,7 +56,7 @@ from mlx_dfloat._memory_caps import install_memory_caps from mlx_dfloat.format import open_checkpoint from mlx_dfloat.integrate.providers import StreamingBF16Provider - from mlx_dfloat.mflux.families import MODELS, entry + from mlx_dfloat.mflux.families import MODELS, ModelEntry, entry from mlx_dfloat.mflux.flux1 import init as base_init from mlx_dfloat.mflux.flux1.names import flux_name_map from mlx_dfloat.mflux.flux1.transformer import ( @@ -124,16 +127,12 @@ def klein_base_problem(model: str, base: Path) -> str | None: return None -def qwen21_base_problem(model: str, base: Path) -> str | None: - """Why ``base`` cannot feed a Qwen-Image 2.1 bf16 side, or None (and None for every other family). +def _sharded_transformer_problem(e: ModelEntry, base: Path) -> str | None: + """Why ``base/transformer`` cannot feed a bf16 side that streams a sharded BF16 transformer, or None. - The bf16 side streams the BF16 transformer from ``base/transformer``: its ``*.safetensors.index.json`` and every - shard the index names must be there, each named by a plain file name (as the streaming reader requires). The df11 - side never reads it. + Its ``*.safetensors.index.json`` and every shard the index names must be there, each named by a plain file name + (as the streaming reader requires). """ - e = entry(model) - if e.family != "qwen21": - return None root = base / "transformer" need = ( f"--base {base}: the bf16 side streams the BF16 transformer from transformer/ " @@ -156,6 +155,31 @@ def qwen21_base_problem(model: str, base: Path) -> str | None: return None +def qwen21_base_problem(model: str, base: Path) -> str | None: + """Why ``base`` cannot feed a Qwen-Image 2.1 bf16 side, or None (and None for every other family). + + The bf16 side streams the BF16 transformer from ``base/transformer``: its ``*.safetensors.index.json`` and every + shard the index names must be there, each named by a plain file name (as the streaming reader requires). The df11 + side never reads it. + """ + e = entry(model) + if e.family != "qwen21": + return None + return _sharded_transformer_problem(e, base) + + +def ernie_base_problem(model: str, base: Path) -> str | None: + """Why ``base`` cannot feed an ERNIE-Image bf16 side, or None (and None for every other family). + + The same check as Qwen-Image 2.1's: the bf16 side streams the two-shard BF16 transformer from ``base/transformer``; + the df11 side never reads it. + """ + e = entry(model) + if e.family != "ernie": + return None + return _sharded_transformer_problem(e, base) + + def parse_args(argv: list[str] | None = None) -> argparse.Namespace: """Parse the command line: one side (``df11``, ``bf16`` or ``compare``), or ``--orchestrate``.""" p = argparse.ArgumentParser( @@ -179,12 +203,13 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: type=float, default=None, help="guidance (default per model, as mflux: FLUX.1 3.5; Z-Image unset: the model's own rule; " - "FLUX.2 Klein 1.0, base models with CFG at 4; Qwen-Image 2.1 1.0, CFG above 1 with --negative-prompt)", + "FLUX.2 Klein 1.0, base models with CFG at 4; Qwen-Image 2.1 1.0, CFG above 1 with --negative-prompt; " + "ERNIE-Image 4.0, CFG above 1; ERNIE-Image-Turbo 1.0 only)", ) p.add_argument( "--negative-prompt", default=None, - help="the negative prompt (used by z-image, a base model, and qwen-image-2.1 above guidance 1)", + help="the negative prompt (used by z-image, a base model, qwen-image-2.1 and ernie-image above guidance 1)", ) p.add_argument("--eval-policy", choices=("per-block", "depth2"), default="per-block") p.add_argument( @@ -193,6 +218,20 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: args = p.parse_args(argv) if args.model in NO_BF16_ORIGINAL: p.error(f"--model {args.model}: {NO_BF16_ORIGINAL[args.model]}") + e = entry(args.model) + if e.family == "ernie" and e.fixed_guidance is not None and args.negative_prompt is not None: + p.error( + f"--negative-prompt: {args.model} runs at guidance 1.0 only; there is no negative branch" + ) + if ( + e.family == "ernie" + and e.fixed_guidance is not None + and args.guidance is not None + and args.guidance != e.fixed_guidance + ): + p.error( + f"--guidance: {args.model} runs at guidance {e.fixed_guidance} only, as its mflux command does" + ) if entry(args.model).family == "flux2" and args.negative_prompt is not None: p.error( f"--negative-prompt: {args.model} takes none (mflux's FLUX.2 Klein encodes its own blank negative)" @@ -202,7 +241,8 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: p.error(f"--model {args.model}: {problem}") # Only the runs that stream the BF16 transformer need it; the df11 side and compare never read it. if args.orchestrate or args.side == "bf16": - problem = qwen21_base_problem(args.model, Path(args.base).expanduser()) + base = Path(args.base).expanduser() + problem = qwen21_base_problem(args.model, base) or ernie_base_problem(args.model, base) if problem is not None: p.error(f"--model {args.model}: {problem}") if args.guidance is None: # resolved before run_key, so a stored FLUX.1 side keeps its 3.5 @@ -1020,11 +1060,167 @@ def run_bf16_qwen21(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, A } +# --- the ERNIE-Image sides (mflux-touching; exercised by the identity run, not unit-tested here) --------------- + + +def run_df11_ernie(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, Any]: + """Generate through ``DFloatErnieImage``, capture the final latents, and save the image, latents and embeddings. + + The embeddings file holds mflux's text batch for the call (``text_bth`` and ``text_lens``): one prompt, or the + negative then the prompt when classifier-free guidance ran (one batch-2 call per step). + + Raises: + VerifyImageError: The after-loop callback never ran (mflux's own loop never called it). + """ + from mflux.models.ernie_image.latent_creator.ernie_latent_creator import ErnieLatentCreator + + from mlx_dfloat.mflux.ernie.model import DFloatErnieImage + + out_dir = args.out / "df11" + out_dir.mkdir(parents=True, exist_ok=True) + model = DFloatErnieImage( + args.model, df11_path=args.df11, base_path=args.base, eval_policy=args.eval_policy + ) + capture = _LatentCapture() + model.callbacks.register(capture) + image = model.generate_image( + args.seed, + args.prompt, + num_inference_steps=args.steps, + height=args.size, + width=args.size, + guidance=args.guidance, + negative_prompt=args.negative_prompt, + ) + captured = capture.latents + if captured is None: + raise VerifyImageError("the after-loop callback never ran; no latents were captured") + mx.eval(captured) + image.save(str(out_dir / "image.png"), export_json_metadata=False, overwrite=True) + mx.save_safetensors(str(out_dir / "latents.safetensors"), {"latents": captured}) + key = model._batch_key( + model.cfg_prompts(args.prompt, negative_prompt=args.negative_prompt, guidance=args.guidance) + ) + text_bth, text_lens = model._batches[key] + mx.save_safetensors( + str(out_dir / "embeds.safetensors"), {"text_bth": text_bth, "text_lens": text_lens} + ) + noise = ErnieLatentCreator.create_noise(args.seed, args.size, args.size) + degenerate = nondegenerate(captured, noise) + report = model.report() + return { + "exit_code": EXIT_OK if not degenerate else EXIT_ERROR, + "degenerate": degenerate, + "report": report, + "output_shape": list(captured.shape), + "output_dtype": str(captured.dtype), + "footprint_peak_bytes": max(watchdog.peak_footprint, phys_footprint()), + "mlx_peak_bytes_max_over_phases": max_phase_peak(report["peaks"]), + } + + +def run_bf16_ernie(args: argparse.Namespace, watchdog: Watchdog) -> dict[str, Any]: + """Stream the base's BF16 transformer shards through the seam and mflux's ERNIE-Image loop body, then the VAE. + + The three non-block groups load from the base as plain weights; the blocks stream as they run. The step function + is mflux's own ``ErnieImage._predict``, taken uncompiled as on the df11 side. + + Raises: + VerifyImageError: The df11 side's saved embeddings are not present yet. + """ + from mflux.models.common.config.config import Config + from mflux.models.common.config.model_config import ModelConfig + from mflux.models.common.vae.tiling_config import TilingConfig + from mflux.models.ernie_image.latent_creator.ernie_latent_creator import ErnieLatentCreator + from mflux.models.ernie_image.variants.txt2img.ernie_image import ErnieImage + from mflux.utils.image_util import ImageUtil + + from mlx_dfloat.mflux._compile import uncompiled + from mlx_dfloat.mflux.ernie import init as einit + from mlx_dfloat.mflux.ernie import transformer as etf + from mlx_dfloat.mflux.ernie.names import ernie_name_map + + out_dir = args.out / "bf16" + out_dir.mkdir(parents=True, exist_ok=True) + embeds_path = args.out / "df11" / "embeds.safetensors" + if not embeds_path.exists(): + raise VerifyImageError(f"{embeds_path}: run the df11 side first") + + df11_root = Path(args.df11).expanduser() + base_root = Path(args.base).expanduser() + ckpt = open_checkpoint(df11_root) + index = etf.base_transformer_files_index(base_root / "transformer") + model_config = ModelConfig.from_name(model_name=args.model, base_model=None) + build = etf.build_transformer( + ckpt, + transformer_overrides=model_config.transformer_overrides, + extras=etf.base_extras(index, ckpt), + nonblock_from_extras=True, + ) + provider = StreamingBF16Provider( + index, {n: ckpt.groups[n].matrix_names for n in build.shapes}, ernie_name_map() + ) + transformer = build.transformer + transformer.attach(provider, build.shapes, eval_policy=args.eval_policy, verify_in_call=True) + + embeds = mx.load(str(embeds_path)) + text_bth, text_lens = embeds["text_bth"], embeds["text_lens"] + config = Config( + model_config=model_config, + num_inference_steps=args.steps, + height=args.size, + width=args.size, + guidance=args.guidance, + scheduler="linear", + ) + # mflux ernie_image.py: txt2img noise (ernie_latent_creator.py:11-17). + latents = ErnieLatentCreator.create_noise(args.seed, args.size, args.size) + mx.eval(latents, text_bth, text_lens) + predict = uncompiled(ErnieImage._predict, transformer, text_bth, text_lens, latents) + + for t in config.time_steps: + # mflux ernie_image.py:90-102's step body: predict (one batch-2 call when CFG runs), step, eval. + sigma_t = config.scheduler.sigmas[t].reshape((1,)) + noise = predict( + latents=latents, + sigma=sigma_t, + text_bth=text_bth, + text_lens=text_lens, + guidance=config.guidance, + ) + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) + mx.eval(latents) + transformer.verify_step() + + mx.save_safetensors(str(out_dir / "latents.safetensors"), {"latents": latents}) + + vae = einit.load_vae(base_root) + mx.clear_cache() + mx.set_cache_limit(0) # the transformer-shaped buffers cannot serve the decoder + decoded = vae.decode_packed_latents( + latents, tiling_config=TilingConfig(vae_decode_tiles_per_dim=None) + ) + ImageUtil.to_pil(decoded).save( + str(out_dir / "image.png") + ) # ImageUtil.to_image's own conversion + + return { + "exit_code": EXIT_OK, + "reads": provider.reads, + "cfg_batch": int(text_bth.shape[0]), + "output_shape": list(latents.shape), + "output_dtype": str(latents.dtype), + "footprint_peak_bytes": max(watchdog.peak_footprint, phys_footprint()), + "mlx_peak_bytes_process": int(mx.get_peak_memory()), # never reset on this side + } + + SIDE_RUNNERS: dict[str, tuple[Callable[..., dict[str, Any]], Callable[..., dict[str, Any]]]] = { "flux1": (run_df11, run_bf16), "zimage": (run_df11_zimage, run_bf16_zimage), "flux2": (run_df11_flux2, run_bf16_flux2), "qwen21": (run_df11_qwen21, run_bf16_qwen21), + "ernie": (run_df11_ernie, run_bf16_ernie), } diff --git a/tests/test_ci_config.py b/tests/test_ci_config.py index 48125a6..6a7869f 100644 --- a/tests/test_ci_config.py +++ b/tests/test_ci_config.py @@ -87,7 +87,8 @@ def test_the_mflux_lane_installs_the_locked_extra_and_proves_it_ran_tests(): # Every `uv run` names the extra and the group, so no step depends on uv's sync mode keeping # what `uv sync --extra mflux` installed. uv_runs = [run for run in runs if "uv run" in run] - assert len(uv_runs) == 7 + # sanity + lane + five per-model coverage gates + the collection check + assert len(uv_runs) == 8 assert all("uv run --extra mflux --group dev " in run for run in uv_runs) pytest_runs = [ shlex.split(run) @@ -195,6 +196,31 @@ def test_the_mflux_lane_measures_and_gates_the_qwen21_model_module_on_its_own(): assert "src/mlx_dfloat/mflux/qwen21/model.py" in pyproject["tool"]["coverage"]["run"]["omit"] +def test_the_mflux_lane_measures_and_gates_the_ernie_model_module_on_its_own(): + # Bug caught: ernie/model.py (every test @pytest.mark.mflux) omitted from the required job's gate but measured by + # no job either, so its coverage could fall to zero with every check green; or its gate diluted by the other model + # modules' lines, or run before the lane's pytest wrote the data; or ERNIE's init.py and transformer.py not + # reported in the lane beside its model (measured for consistency, not gated). + import tomllib + + runs = [step.get("run", "") for step in _integration_job()["steps"]] + lane = next( + i + for i, run in enumerate(runs) + if run.startswith("uv run --extra mflux --group dev pytest") + and "--co" not in shlex.split(run) + ) + args = shlex.split(runs[lane]) + assert "--cov=mlx_dfloat.mflux.ernie.model" in args + assert "--cov=mlx_dfloat.mflux.ernie.init" in args + assert "--cov=mlx_dfloat.mflux.ernie.transformer" in args + report = "uv run --extra mflux --group dev coverage report --include='*/mflux/ernie/model.py' --fail-under=80" + assert report in runs + assert lane < runs.index(report) < lane + 6 + pyproject = tomllib.loads((CI.parents[2] / "pyproject.toml").read_text()) + assert "src/mlx_dfloat/mflux/ernie/model.py" in pyproject["tool"]["coverage"]["run"]["omit"] + + def _config() -> dict: return yaml.safe_load(CI.read_text()) diff --git a/tests/test_ernie_real_blocks.py b/tests/test_ernie_real_blocks.py new file mode 100644 index 0000000..4ff847f --- /dev/null +++ b/tests/test_ernie_real_blocks.py @@ -0,0 +1,198 @@ +"""The real ERNIE-Image DF11 checkpoint: its groups, Metal decode against the BF16 base, and a one-block build run. + +Env-gated: every test skips unless the snapshot directories are named by ``MLX_DFLOAT_ERNIE_DF11`` (the +``mingyi456/ERNIE-Image-DF11`` snapshot: ``config.json`` and ``model.safetensors``) and, for the base comparisons, +``MLX_DFLOAT_ERNIE_BASE`` (the ``baidu/ERNIE-Image`` snapshot with its ``transformer/`` index and two shards). Expected +counts are literals from the checkpoint's header (39 groups: ``layers.0..35`` and the three non-block groups; 153 BF16 +extras, read 2026-10-08) and from mflux 0.20.0 (``ErnieTransformer`` builds 36 blocks, transformer.py:36-76); the +expected bits come from the base's tensors, never from the code under test. +""" + +import gc +import os +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import mlx.core as mx +import numpy as np +import pytest + +from mlx_dfloat.decode import check, decode_group, split_matrices +from mlx_dfloat.format import DF11Checkpoint, load_group_mx, open_checkpoint +from mlx_dfloat.integrate.coverage import load_resident_set +from mlx_dfloat.integrate.providers import DF11Provider, StreamingBF16Provider, read_bf16 +from mlx_dfloat.integrate.resident import decode_nonblock, install_nonblock +from mlx_dfloat.mflux.ernie.names import KIND, NONBLOCK_GROUPS, check_ernie_groups + +DF11_ENV, BASE_ENV = "MLX_DFLOAT_ERNIE_DF11", "MLX_DFLOAT_ERNIE_BASE" +# Every block's split positions, by hand: to_q, to_k, to_v, to_out.0 (4096 x 4096 = 16777216 each), then gate_proj and +# up_proj (12288 x 4096 = 50331648 each); linear_fc2 (4096 x 12288) ends the group at 218103808. +BLOCK_SPLITS = (16777216, 33554432, 50331648, 67108864, 117440512, 167772160) +SEVEN = ( + "self_attention.to_q", + "self_attention.to_k", + "self_attention.to_v", + "self_attention.to_out.0", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.linear_fc2", +) +# ModelConfig.ernie_image().transformer_overrides (mflux 0.20.0 model_config.py:698-731). +OVERRIDES = {"rope_axes_dim": [32, 48, 48]} + + +def _env(name: str) -> Path: + value = os.environ.get(name) + if not value: + pytest.skip(f"{name} not set") + return Path(value) + + +def _bits(array: mx.array) -> np.ndarray: + """The array's raw bits through a dtype-sized unsigned view (equality of bits, never of floats).""" + view = {2: mx.uint16, 4: mx.uint32}[array.dtype.size] + return np.array(array.view(view)) + + +def _base_index(base: Path) -> dict[str, Any]: + from mlx_dfloat.mflux.ernie.transformer import base_transformer_files_index + + return base_transformer_files_index(base / "transformer") + + +@pytest.mark.slow +def test_the_real_checkpoint_opens_with_ernies_groups() -> None: + # Bug caught: a group or extra the adapter does not expect (the build would refuse at the first user's call): a + # group count other than the header's 39, a block holding other than the seven matrices, a non-block group + # missing, or an extra count other than 153 (four per block plus nine non-block tensors). + ckpt = open_checkpoint(_env(DF11_ENV)) + assert ckpt.config_source == "config.json" + assert check_ernie_groups(ckpt) == {KIND: 36} + assert len(ckpt.groups) == 39 + assert sorted(g for g in ckpt.groups if not g.startswith("layers.")) == [ + "adaLN_modulation.1", + "final_norm.linear", + "time_embedding", + ] + assert tuple(ckpt.groups["layers.35"].matrix_names) == tuple( + f"layers.35.{sub}.weight" for sub in SEVEN + ) + assert len(ckpt.extras) == 153 + assert sorted({info.dtype for _path, info in ckpt.extras.values()}) == ["BF16"] + + +@pytest.mark.slow +@pytest.mark.metal +def test_four_real_groups_decode_on_metal_bit_exact_against_the_base() -> None: + # Bug caught: a group decoded against the wrong base tensor (a stored matrix order off in the first or last + # block: gate_proj and up_proj have the same shape, so only their values show it), or a non-block group mis-split + # (time_embedding's two matrices cut at the wrong element, adaLN_modulation.1 (24576 x 4096) in the wrong layout). + ckpt = open_checkpoint(_env(DF11_ENV)) + index = _base_index(_env(BASE_ENV)) + for name in ("layers.0", "layers.35", "adaLN_modulation.1", "time_embedding"): + group = ckpt.groups[name] + loaded = load_group_mx(group) + if name.startswith("layers."): + assert loaded.split_positions == BLOCK_SPLITS, name + result = decode_group(loaded, backend="metal") + check(result, name=name) + parts = split_matrices(result.bits, loaded.split_positions) + assert len(parts) == len(group.matrix_names), name + for matrix, part in zip(group.matrix_names, parts, strict=True): + path, info = index[matrix] + want = np.array(read_bf16(path, info).view(mx.uint16)).ravel() + assert np.array_equal(np.array(part.view(mx.uint16)).ravel(), want), matrix + del loaded, result, parts + gc.collect() + mx.clear_cache() + + +def _inputs(*, size: int, text_tokens: int, seed: int = 1) -> dict[str, Any]: + """One call's keyword inputs: a (1, 128, size/16, size/16) latent, ``text_tokens`` of width 3072, one timestep.""" + side = size // 16 + keys = mx.random.split(mx.random.key(seed), 2) + return { + "hidden_states": mx.random.normal((1, 128, side, side), key=keys[0]).astype(mx.bfloat16), + "timestep": mx.array([500.0]), + "text_bth": mx.random.normal((1, text_tokens, 3072), key=keys[1]).astype(mx.bfloat16), + "text_lens": mx.array([text_tokens], dtype=mx.int32), + } + + +def _run(transformer: Any, inputs: Mapping[str, Any]) -> mx.array: + out = transformer(**inputs) + mx.eval(out) + return out + + +def df11_one_block(ckpt: DF11Checkpoint, inputs: Mapping[str, Any]) -> mx.array: + """Build one block from the checkpoint, decode and install the non-block groups, attach the DF11 decode, run.""" + from mlx_dfloat.mflux.ernie.names import ernie_name_map + from mlx_dfloat.mflux.ernie.transformer import build_transformer + + names = ernie_name_map() + build = build_transformer(ckpt, transformer_overrides=OVERRIDES, n_layers=1) + blocks = list(build.shapes) + assert blocks == ["layers.0"] + groups = load_resident_set(ckpt, [*blocks, *NONBLOCK_GROUPS]) + matrix_names = {n: ckpt.groups[n].matrix_names for n in groups} + nonblock = decode_nonblock( + {g: groups[g] for g in NONBLOCK_GROUPS}, + {g: matrix_names[g] for g in NONBLOCK_GROUPS}, + build.nonblock, + names, + ) + assert sorted(nonblock) == [ + "adaln_modulation.weight", + "final_norm.linear.weight", + "time_embedding.linear_1.weight", + "time_embedding.linear_2.weight", + ] + install_nonblock(build.transformer, nonblock) + provider = DF11Provider( + {n: groups[n] for n in blocks}, {n: matrix_names[n] for n in blocks}, names + ) + build.transformer.attach(provider, build.shapes, verify_in_call=True) + return _run(build.transformer, inputs) + + +def bf16_one_block( + ckpt: DF11Checkpoint, index: Mapping[str, Any], inputs: Mapping[str, Any] +) -> mx.array: + """The same one-block build with the base's BF16 extras, the non-block matrices plain weights, blocks streamed.""" + from mlx_dfloat.mflux.ernie.names import ernie_name_map + from mlx_dfloat.mflux.ernie.transformer import base_extras, build_transformer + + names = ernie_name_map() + build = build_transformer( + ckpt, + transformer_overrides=OVERRIDES, + n_layers=1, + extras=base_extras(index, ckpt), + nonblock_from_extras=True, + ) + assert build.nonblock == {} + provider = StreamingBF16Provider( + index, {n: ckpt.groups[n].matrix_names for n in build.shapes}, names + ) + build.transformer.attach(provider, build.shapes, verify_in_call=True) + return _run(build.transformer, inputs) + + +@pytest.mark.slow +@pytest.mark.metal +@pytest.mark.mflux +def test_a_one_block_build_matches_the_bf16_stream_bit_for_bit() -> None: + # Bug caught: a non-block rename or a block matrix misplaced on one side (the adaLN_modulation.1 rename, the patch + # conv's transpose, a non-block group left on its placeholder: a zero-size matmul before the first block), or any + # block-seam difference between the DF11 and the streamed BF16 paths (the decoder itself is the test above's). + # 1024²: a 64 x 64 latent of 128 channels, 32 text tokens of mflux's 3072-wide text input. + ckpt = open_checkpoint(_env(DF11_ENV)) + index = _base_index(_env(BASE_ENV)) + inputs = _inputs(size=1024, text_tokens=32) + out_a = df11_one_block(ckpt, inputs) + out_b = bf16_one_block(ckpt, index, inputs) + assert out_a.shape == out_b.shape == (1, 128, 64, 64) + assert out_a.dtype == out_b.dtype + assert np.array_equal(_bits(out_a), _bits(out_b)) diff --git a/tests/test_verify_image.py b/tests/test_verify_image.py index bfb1fa3..0d7d93a 100644 --- a/tests/test_verify_image.py +++ b/tests/test_verify_image.py @@ -750,3 +750,130 @@ def test_qwen_embeds_hold_the_positive_pair_and_the_negative_one_when_cfg_ran(): assert two["negative_prompt_embeds"] is neg[0] assert two["negative_prompt_mask"] is neg[1] assert two["prompt_embeds"] is pos[0] + + +# --- ERNIE-Image ---------------------------------------------------------------------------------------------- + +ERNIE = "ernie-image" +ERNIE_TURBO = "ernie-image-turbo" + + +def _parse_ernie(tmp_path, base, *extra, model=ERNIE, side=("--side", "df11")): + return vi.parse_args( + [ + *side, + "--model", + model, + "--df11", + str(tmp_path), + "--base", + str(base), + "--out", + str(tmp_path / "out"), + *extra, + ] + ) + + +@pytest.mark.parametrize("model", [ERNIE, ERNIE_TURBO]) +def test_runner_for_dispatches_ernie_to_its_sides(model): + # Bug caught: ERNIE-Image routed to the FLUX.1 sides (a wrong-model identity verdict), or no ernie row (a KeyError + # after the df11 side already ran for minutes). + assert vi.runner_for(model, "df11") is vi.run_df11_ernie + assert vi.runner_for(model, "bf16") is vi.run_bf16_ernie + + +@pytest.mark.parametrize("layout", ["no_dir", "empty_dir", "no_index"]) +def test_ernie_base_problem_names_the_missing_transformer(tmp_path, layout): + # Bug caught: the identity launched on a base without its BF16 transformer, failing after the df11 side's minutes. + base = tmp_path / "qbase" + base.mkdir() + if layout != "no_dir": + (base / "transformer").mkdir() + if layout == "no_index": + (base / "transformer" / "diffusion_pytorch_model-00001-of-00002.safetensors").write_bytes( + b"\0" + ) + problem = vi.ernie_base_problem(ERNIE, base) + assert problem is not None + assert f"--base {base}: the bf16 side streams the BF16 transformer from transformer/" in problem + assert "download baidu/ERNIE-Image's transformer/* first" in problem + + +def test_ernie_base_problem_is_none_for_an_index_with_its_shards(tmp_path): + # Bug caught: a complete transformer/ refused (the identity check could never run). + assert vi.ernie_base_problem(ERNIE, _qwen_base(tmp_path)) is None + + +@pytest.mark.parametrize("shard", ["../outside.safetensors", "sub/inner.safetensors"]) +def test_ernie_base_problem_keeps_the_shard_name_check(tmp_path, shard): + # Bug caught: the extraction of the shared check losing the plain-file-name guard (an index naming a path outside + # transformer/ accepted, so the bf16 side streams another file's tensors). + base = _qwen_base(tmp_path) + target = base / "transformer" / shard + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(b"\0") + index = base / "transformer" / "diffusion_pytorch_model.safetensors.index.json" + index.write_text(json.dumps({"weight_map": {"a.weight": shard}})) + problem = vi.ernie_base_problem(ERNIE, base) + assert problem is not None + assert f"not a plain file name: {shard!r}" in problem + + +def test_ernie_base_problem_is_none_for_the_other_families(tmp_path): + # Bug caught: the ERNIE check applied to other families' bases (Qwen keeps its own; FLUX.1 needs no transformer/). + assert vi.ernie_base_problem("schnell", tmp_path) is None + assert vi.ernie_base_problem("qwen-image-2.1", tmp_path) is None + + +@pytest.mark.parametrize("side", [("--side", "bf16"), ("--orchestrate",)]) +def test_an_ernie_run_that_streams_the_bf16_side_is_refused_without_its_transformer( + tmp_path, capsys, side +): + # Bug caught: the check not wired into parse_args for the runs that stream the BF16 transformer. + base = tmp_path / "qbase" + (base / "transformer").mkdir(parents=True) + with pytest.raises(SystemExit) as info: + _parse_ernie(tmp_path, base, side=side) + assert info.value.code == 2 + assert "transformer/" in capsys.readouterr().err + + +def test_ernie_turbo_negative_prompt_is_refused(tmp_path, capsys): + # Bug caught: a negative prompt accepted for Turbo (it runs at guidance 1.0 only: no negative branch), so the key + # would carry a setting that changes nothing, or the bf16 side would run CFG the df11 side never ran. + with pytest.raises(SystemExit) as info: + _parse_ernie(tmp_path, tmp_path, "--negative-prompt", "blurry", model=ERNIE_TURBO) + assert info.value.code == 2 + assert "runs at guidance 1.0 only; there is no negative branch" in capsys.readouterr().err + + +def test_ernie_negative_prompt_is_accepted_and_keyed_with_mfluxs_guidance(tmp_path, monkeypatch): + # Bug caught: --negative-prompt refused for the base (its CFG uses it), dropped from the key (a stored side reused + # after it changed) or from the children; or the guidance left None instead of mflux's 4.0 + # (ernie_image_generate.py:22), which the bf16 side's Config would read as 0.0 (no CFG). + monkeypatch.setattr(vi, "source_hash", lambda: "src") + args = _parse_ernie( + tmp_path, _qwen_base(tmp_path), "--negative-prompt", "blurry", side=("--orchestrate",) + ) + assert vi.run_key(args)["negative_prompt"] == "blurry" + assert args.guidance == 4.0 + command = vi.child_command(args, "bf16") + assert command[command.index("--negative-prompt") + 1] == "blurry" + assert command[command.index("--guidance") + 1] == "4.0" + + +@pytest.mark.parametrize("guidance", ["4", "1.5"]) +def test_ernie_turbo_guidance_other_than_1_is_refused(tmp_path, capsys, guidance): + # Bug caught (C2): an identity run of Turbo at a guidance its command refuses (mflux's turbo command errors on any + # but 1.0, ernie_image_turbo_generate.py:42-45), so the check would cover a path no user runs. + with pytest.raises(SystemExit) as info: + _parse_ernie(tmp_path, tmp_path, "--guidance", guidance, model=ERNIE_TURBO) + assert info.value.code == 2 + assert "runs at guidance 1.0 only" in capsys.readouterr().err + + +def test_ernie_turbo_guidance_1_is_accepted(tmp_path): + # Bug caught: the refusal firing on Turbo's own 1.0 (or on the default, which resolves to it). + assert _parse_ernie(tmp_path, tmp_path, "--guidance", "1", model=ERNIE_TURBO).guidance == 1.0 + assert _parse_ernie(tmp_path, tmp_path, model=ERNIE_TURBO).guidance == 1.0 From b3b94f2703960f2419edae7a979650d2623c3d49 Mon Sep 17 00:00:00 2001 From: Denis Ineshin Date: Thu, 8 Oct 2026 16:03:04 +0200 Subject: [PATCH 4/4] Document ERNIE-Image and ERNIE-Image-Turbo --- CHANGELOG.md | 31 ++++++++- README.md | 173 +++++++++++++++++++++++++++++++++++++++++++++------ 2 files changed, 185 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b84ddc..df4560c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -53,6 +53,28 @@ All notable changes to this project are documented here. The format follows decode. Under a 16 GB and a 24 GB Mac's caps the fit check refuses it, predicting 14.6 GiB while the prompt is encoded (`bench/results/refusals/qwen-image-2.1-1024-tier16.json` and `qwen-image-2.1-1024-tier24.json`), and forced runs were stopped by the watchdog in that phase, so both sizes are at risk. +- ERNIE-Image and ERNIE-Image-Turbo generation through mflux: `mlx-dfloat generate --model + ernie-image|ernie-image-turbo` and the `DFloatErnieImage` class, also returned by `DFloatModel`. The text encoder, VAE and tokenizer come from + `baidu/ERNIE-Image` and `baidu/ERNIE-Image-Turbo`, pinned to the revisions the measured runs used; their BF16 + transformers and prompt enhancers are not downloaded. The defaults follow mflux and the model cards: 50 steps at + guidance 4.0 for ERNIE-Image, which runs classifier-free guidance as one batched call per step with or without a + negative prompt, and 8 steps at guidance 1.0 for ERNIE-Image-Turbo. An empty or blank prompt is refused as a user + error. +- Every compressed matrix and stored tensor of the ERNIE-Image transformer (256 and 153) equals its BF16 original, and + the GPU decoder matches the CPU reference on all 256 (both records in `bench/results/parity/ernie-image/`). In a + 4-step, 1024² run with classifier-free guidance on, the final latents match bit for bit those of the same + transformer streaming its BF16 weights one block at a time + (`bench/results/identity/ernie-image-1024/compare.json`). ERNIE-Image-Turbo's 256 compressed matrices equal its BF16 + original too, read from the Hub one piece at a time, and decode on the GPU to the CPU reference's bits (both records + in `bench/results/parity/ernie-image-turbo/`); its latents have not been compared with a reference. +- The measured 32 GB runs peak at 16.79 GiB (ERNIE-Image, 50 steps) and 17.34 GiB (ERNIE-Image-Turbo, 8 steps) at + 1024² (`bench/results/tiers/ernie-image-1024.json` and `ernie-image-turbo-1024.json`), and keep the compressed + transformer loaded through the VAE decode. Under a 24 GB Mac's caps the fit check refuses ERNIE-Image-Turbo, + predicting 14.39 GiB against the 14.00 GiB budget; a forced run passed at 14.39 GiB + (`bench/results/tiers/ernie-image-turbo-1024-tier24.json`), 0.11 GiB under the watchdog ceiling, so that size is at + risk. The fit check also refuses ERNIE-Image at 24 GB and both models at 16 GB (`bench/results/refusals/`), and forced + runs were stopped by the watchdog, so those sizes are at risk too. +- An ERNIE-Image base that lacks a text-encoder or VAE tensor is refused with `DFloatFormatError`. - A FLUX.2 Klein base that lacks a text-encoder or VAE tensor is refused with `DFloatFormatError` instead of leaving that weight at its random initial value. - A FLUX.2 Klein call warns when the prompt encode holds more memory than the fit estimate allows for it. The @@ -69,6 +91,10 @@ All notable changes to this project are documented here. The format follows - `mlx-dfloat generate --model` also accepts the four FLUX.2 Klein names. A distilled FLUX.2 Klein model refuses a `--guidance` other than 1.0, and every FLUX.2 Klein model refuses `--scheduler`, matching mflux's own Klein command. The image identity check covers FLUX.2 Klein. +- `mlx-dfloat generate --model` also accepts `ernie-image` and `ernie-image-turbo`. ERNIE-Image-Turbo refuses a + `--guidance` other than 1.0 and ignores `--negative-prompt` with a warning, as mflux's Turbo command does. The image + identity check covers ERNIE-Image and refuses `--negative-prompt` and a guidance other than 1.0 for + ERNIE-Image-Turbo. - `mlx-dfloat generate --report` records `elapsed_seconds`, and the watchdog's `abort.json` names the phase that was running (`build`, `encode`, `set_load`, `denoise` or `vae`). A checkpoint or base taken as a Hub id (the defaults included) prints one line saying it is resolved on the Hub, from the local cache or by download. On the Z-Image base @@ -82,7 +108,10 @@ All notable changes to this project are documented here. The format follows 8-bit paths, and a Python example. - From Python, every model's `generate_image` runs under the memory caps `mlx-dfloat generate` installs and restores MLX's limits when it returns. Before, a Python caller ran at MLX's default limits, with no wired cap, while the fit - estimates' VAE terms were measured under the caps. A wired limit the process already set is left alone. + estimates' VAE terms were measured under the caps. A wired limit the process already set is left alone. Each model's + `encode()` runs under the same caps. The wired limit is read once per process instead of at every call, and a wired + cap MLX refuses (a system wired limit lowered below it) leaves the memory cap in place with one warning instead of + raising. - When a call drops the compressed transformer before the VAE decode, the fit estimate for that step now counts the transformer's uncompressed BF16 tensors, which stay loaded. The estimate grows by 6 MB to 137 MB depending on the model; at 1024² none of the 16, 24 or 32 GB verdicts changes. diff --git a/README.md b/README.md index 67e74f5..fb471be 100644 --- a/README.md +++ b/README.md @@ -7,15 +7,15 @@ Run [DFloat11](https://github.com/LeanModels/DFloat11) checkpoints on a Mac with [MLX](https://github.com/ml-explore/mlx). The weights stay about 30 % smaller than BF16 in memory, and the GPU decodes each block back to exactly the same bits just before it runs. Today that means FLUX.1 image generation on a 32 GB Mac; -Z-Image, FLUX.2 Klein and Qwen-Image-2.1 arrive in the next release. +Z-Image, FLUX.2 Klein, Qwen-Image-2.1 and ERNIE-Image arrive in the next release. DFloat11 is lossless compression for BF16 model weights. Each weight is 16 bits. DFloat11 stores the 8 exponent bits as a short variable-length code, the same idea a zip file uses, and keeps the other 8 bits (the sign and the fraction) exactly as they are. The DFloat11 authors report models at about 70% of their BF16 size with output that is bit for bit the same as the original.[^size] Their decoder runs on NVIDIA GPUs only. This project is an independent reader and decoder for Apple Silicon. The weights stay compressed in memory, and a Metal kernel decodes each block -on the GPU as it runs. That is how it generates FLUX.1 images through mflux; Z-Image, FLUX.2 Klein and Qwen-Image-2.1 -follow in the next release. +on the GPU as it runs. That is how it generates FLUX.1 images through mflux; Z-Image, FLUX.2 Klein, Qwen-Image-2.1 and +ERNIE-Image follow in the next release. ``` one BF16 weight, 16 bits: s eeeeeeee mmmmmmm @@ -48,9 +48,13 @@ of the three FLUX.1 models; no other Mac has been measured. The numbers are unde | 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, from the next release. Peak at 1024²: 20.42 GiB. 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, from the next release. 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". | +| A 24 GB Mac, ERNIE-Image | CAPPED: at risk. The fit check refuses it, and a forced run was stopped by the watchdog. See "Generate an ERNIE-Image image". | +| A 16 GB Mac, ERNIE-Image or ERNIE-Image-Turbo | CAPPED: at risk. The fit check refuses both, and forced runs were stopped by the watchdog. See "Generate an ERNIE-Image image". | | A Mac with room for the BF16 transformer next to everything else | Plain mflux in BF16. There is nothing to decode. | | A Mac with less than 32 GB | Not measured. The rows marked CAPPED are runs on the 32 GB Mac under the caps mlx-dfloat installs on a 16 GB or 24 GB Mac; the tables below make no other claim for those Macs. | -| Another model: an LLM, the original Qwen-Image or Qwen-Image-Edit, FLUX.2-dev, Krea-2 | Not wired up yet. The reader and the decoder handle all four published versions of the DFloat11 format, but generation in 0.1.0 covers FLUX.1 only. The next release adds Z-Image-Turbo, Z-Image, the four FLUX.2 Klein models and Qwen-Image-2.1. | +| Another model: an LLM, the original Qwen-Image or Qwen-Image-Edit, FLUX.2-dev, Krea-2 | Not wired up yet. The reader and the decoder handle all four published versions of the DFloat11 format, but generation in 0.1.0 covers FLUX.1 only. The next release adds Z-Image-Turbo, Z-Image, the four FLUX.2 Klein models, Qwen-Image-2.1, ERNIE-Image and ERNIE-Image-Turbo. | ## Status @@ -107,6 +111,17 @@ In the next release (not yet on PyPI): 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. +- ERNIE-Image and ERNIE-Image-Turbo make 1024² images through mflux, from the command line or from Python. For + ERNIE-Image, every one of the 256 compressed matrices and 153 other tensors equals the BF16 original in + `baidu/ERNIE-Image` (`bench/results/parity/ernie-image/full-vs-bf16.json`), and the GPU decoder gives the same bits as + the CPU reference on all 256 (`bench/results/parity/ernie-image/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 ran at 1024² for 4 steps at guidance 4.0, so classifier-free guidance ran (record: + `bench/results/identity/ernie-image-1024/compare.json`). For ERNIE-Image-Turbo, all 256 compressed matrices equal the + BF16 original in `baidu/ERNIE-Image-Turbo` as well, read from the Hub one piece at a time + (`bench/results/parity/ernie-image-turbo/full-vs-bf16.json`), and the GPU decoder matches the CPU reference on all + 256 (`bench/results/parity/ernie-image-turbo/kernel-vs-reference.json`). Turbo's latents were not compared with a + reference: that check needs its 16 GB BF16 transformer on disk. - `mlx-dfloat generate --tier 16` and `--tier 24` run under the memory caps mlx-dfloat installs on a Mac of that size. The 16 GB and 24 GB rows under "Measured numbers" are CAPPED runs on the 32 GB Mac under those caps; no smaller Mac was measured. @@ -123,7 +138,7 @@ 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 and Qwen-Image-2.1 generation. +`mlx-dfloat generate` command needs for FLUX.1, Z-Image, FLUX.2 Klein, Qwen-Image-2.1 and ERNIE-Image generation. 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/): @@ -175,10 +190,10 @@ result is compared with bits known in advance. It exits 0 when every check passe GPU decoder cannot run at all, for example because there is no Metal device. `--json` prints the same report as JSON. The same check runs automatically the first time a process decodes on the GPU, once for each write path, and when a -FLUX.1, Z-Image, FLUX.2 Klein or Qwen-Image-2.1 model or the `generate` command sets up its decoder, so a broken GPU -decoder refuses before any model loads. On an M1 Max the check itself took about 13 ms, measured after the kernel -pipelines were compiled; the one-time compile took about 0.4 s with a cold shader cache and 0.03 s with a warm one. It -does not check whether a real model fits in your memory or how fast it runs; the numbers below cover that. +FLUX.1, Z-Image, FLUX.2 Klein, Qwen-Image-2.1 or ERNIE-Image model or the `generate` command sets up its decoder, so a +broken GPU decoder refuses before any model loads. On an M1 Max the check itself took about 13 ms, measured after the +kernel pipelines were compiled; the one-time compile took about 0.4 s with a cold shader cache and 0.03 s with a warm +one. It does not check whether a real model fits in your memory or how fast it runs; the numbers below cover that. ### Generate a FLUX.1 image @@ -416,9 +431,12 @@ step calls the transformer once, with no classifier-free guidance. As in mflux, that no classifier-free guidance runs; pass a negative prompt at guidance 1 or below and it warns that the prompt has no effect. -The command takes the DFloat11 transformer from that checkpoint and the text encoder, VAE and tokenizer from -`Qwen/Qwen-Image-2.1`, without downloading that repository's BF16 transformer. Both default repositories are pinned to -the revisions the measured runs used; a `--df11` or `--base` you pass is used as given. +The command takes the DFloat11 transformer from that checkpoint (10.9 GB for either model, as listed on the Hub) and the +text encoder, VAE and tokenizer from the model's BF16 repository. It does not download that repository's BF16 +transformer (16.1 GB) or its prompt enhancer (`pe/`, 7.7 GB); mflux uses neither. The text-encoder file is 7.7 GB and +also holds a vision model that is never loaded. Both BF16 repositories serve the same text-encoder file, so it is stored +once. The VAE is 0.17 GB. These sizes are the files as listed on the Hub. The default checkpoints and BF16 repositories +are pinned to the revisions the measured runs used; a `--df11` or `--base` you pass is used as given. The same from Python: @@ -460,6 +478,90 @@ Quantisation, LoRA, img2img, mflux's alternative image decoder (`--pid-decode`) FLUX.1. Only text-to-image is on this path; Qwen-Image-2.1's image editing is not. Sizes above 1024² are refused unless you pass `--no-fit-check`, because none has been measured. +### Generate an ERNIE-Image image + +ERNIE-Image and ERNIE-Image-Turbo run through the same command. Their BF16 repositories (`baidu/ERNIE-Image` and +`baidu/ERNIE-Image-Turbo`) and the DFloat11 checkpoints (`mingyi456/ERNIE-Image-DF11` and +`mingyi456/ERNIE-Image-Turbo-DF11`) are ungated on the Hub, so no Hub login is needed. All four list the Apache-2.0 +licence on the Hub ([ERNIE-Image](https://huggingface.co/baidu/ERNIE-Image), +[ERNIE-Image-Turbo](https://huggingface.co/baidu/ERNIE-Image-Turbo)). + +``` +mlx-dfloat generate --model ernie-image-turbo \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report report.json +mlx-dfloat generate --model ernie-image \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report report.json +``` + +The defaults are mflux's own for each model, and they match the steps and guidance each model card recommends at +1024x1024 ([ERNIE-Image](https://huggingface.co/baidu/ERNIE-Image): 50 steps at guidance 4.0; +[ERNIE-Image-Turbo](https://huggingface.co/baidu/ERNIE-Image-Turbo): 8 steps at guidance 1.0). ERNIE-Image-Turbo refuses +any other `--guidance`, as mflux's Turbo command does, and ignores a `--negative-prompt` with a warning. Above guidance +1, ERNIE-Image runs classifier-free guidance whether or not you pass a negative prompt; without one, mflux uses a single +space. Each step runs the prompt and the negative prompt together, as one batched transformer call. Both models use the +`linear` scheduler. An empty or blank `--prompt` is refused for both, as a user error. + +The command takes the DFloat11 transformer from that checkpoint, 10.9 GB for either model, and the text encoder, VAE and +tokenizer from the model's BF16 repository. It does not download that repository's BF16 transformer (16.1 GB) or its +prompt enhancer (`pe/`, 7.7 GB), which mflux does not use. The text-encoder file is 7.7 GB and also holds a vision model +that is never loaded. Both BF16 repositories publish the same text-encoder file, and huggingface_hub 2.0's shared cache +keeps one copy of it for both models. The VAE is 0.17 GB. The default checkpoints and BF16 repositories are pinned to +the revisions the measured runs used; a `--df11` or `--base` you pass is used as given. + +The same from Python: + +```python +from mlx_dfloat.mflux import DFloatModel + +model = DFloatModel("ernie-image") +image = model.generate_image( + seed=42, + prompt="A stone lighthouse on a rocky shore at dawn", + height=1024, + width=1024, +) +image.save("lighthouse.png") +``` + +Left out, `guidance`, `num_inference_steps` and `scheduler` take the same per-model defaults as the command: 4.0 and +50 steps for ERNIE-Image, 1.0 and 8 steps for ERNIE-Image-Turbo. mflux's own `ErnieImage` class defaults to 1.0 and 8 +steps for both. Unlike the command, the Python class lets ERNIE-Image-Turbo run at another guidance, as mflux's class +does. An empty or blank prompt is refused here too. Each call runs under the memory caps the command installs, unless +the process has already set a wired limit. + +mflux compiles ERNIE-Image's denoising step on M1 and M2 Max and Ultra chips and on M3 and later. This path runs it +without that compilation on every chip (see "Generate a Z-Image image" for why). mflux also computes ERNIE-Image's +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 +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 +prediction of each step's peak memory that it checks before a call, puts the VAE decode at 17.85 GiB, under the 22.96 +GiB budget. + +Under a 24 GB Mac's caps (`--tier 24`, run on the 32 GB M1 Max), the fit check refuses ERNIE-Image-Turbo. It predicts +14.39 GiB for the denoising step against the 14.00 GiB fit budget. A forced run (`--no-fit-check`) passed with a 14.39 +GiB peak in 2.8 minutes, the compressed transformer dropped before the VAE decode +(`bench/results/tiers/ernie-image-turbo-1024-tier24.json`). That peak is only 0.11 GiB under the 14.50 GiB watchdog +ceiling. The run's report still shows the earlier 14.03 GiB estimate in its `fit` block, because the estimate was raised +from this run's peak. Two identical runs under these caps peaked 0.35 GiB apart, at 14.03 and 14.39 GiB, which is why +the estimate takes the higher one. ERNIE-Image-Turbo on a 24 GB Mac is at risk. No 24 GB Mac was measured. + +ERNIE-Image under a 24 GB Mac's caps, and both models under a 16 GB Mac's, are at risk too. The fit check refuses them: +it predicts 15.2 GiB for ERNIE-Image's denoising step and 14.39 GiB for Turbo's, against budgets of 14.00 GiB at 24 GB +and 8.67 GiB at 16 GB. Forced runs (`--no-fit-check`) were stopped by the watchdog. ERNIE-Image at 24 GB was stopped in +the denoising step after 20.6 s, at a 14.80 GiB footprint against the 14.50 GiB ceiling. At 16 GB both were stopped +while the compressed transformer was loading, against the 9.17 GiB ceiling: Turbo after 18.6 s at 9.22 GiB, ERNIE-Image +after 23.3 s at 9.44 GiB. + +Quantisation, LoRA, img2img, mflux's alternative image decoder (`--pid-decode`) and ControlNet are refused, as for +FLUX.1. Sizes above 1024² are refused unless you pass `--no-fit-check`, because none has been measured. + ## Measured numbers The blocks below are generated from the files under `bench/results/` by `scripts/bench_table.py`, and a test fails when @@ -471,12 +573,13 @@ and 16 GiB), so it sees the limits a user of that Mac runs under. PROOF marks a ceiling, made only to show that the watchdog works; it never appears as a tier row. The 32 GB rows are MEASURED, on one M1 Max. The 16 GB and 24 GB rows are CAPPED, run on that same Mac: -FLUX.2-klein-base-4B passed under both, while FLUX.2-klein-base-9B under a 24 GB Mac's caps, Z-Image-Turbo under a -16 GB Mac's and Qwen-Image-2.1 under both were stopped by the watchdog. The table makes no claim about any Mac it does -not list. Each row is one `mlx-dfloat generate` run at 1024² with seed 42 (4 steps for FLUX.1-schnell, FLUX.2-klein-4B -and FLUX.2-klein-9B, 20 for FLUX.1-dev and FLUX.1-Krea-dev, 9 for Z-Image-Turbo, 40 at guidance 1.0 for -Qwen-Image-2.1, 50 with guidance 4.0 for Z-Image and the two FLUX.2 Klein base models), with `--tier` and `--report` -writing the file in the last column. +FLUX.2-klein-base-4B passed under both and ERNIE-Image-Turbo under a 24 GB Mac's caps (a forced run: the fit check +refuses it), while FLUX.2-klein-base-9B under a 24 GB Mac's caps, Z-Image-Turbo under a 16 GB Mac's, ERNIE-Image-Turbo +under a 16 GB Mac's, and Qwen-Image-2.1 and ERNIE-Image under both were stopped by the watchdog. The table makes no +claim about any Mac it does not list. Each row is one `mlx-dfloat generate` run at 1024² with seed 42 (4 steps for +FLUX.1-schnell, FLUX.2-klein-4B and FLUX.2-klein-9B, 20 for FLUX.1-dev and FLUX.1-Krea-dev, 9 for Z-Image-Turbo, 40 at +guidance 1.0 for Qwen-Image-2.1, 8 at guidance 1.0 for ERNIE-Image-Turbo, 50 with guidance 4.0 for Z-Image, ERNIE-Image +and the two FLUX.2 Klein base models), with `--tier` and `--report` writing the file in the last column. The second column is the Mac's recommended GPU working set minus a reserve. On the 16 GB and 24 GB rows the working set is taken as two thirds of that Mac's memory. Two limits come off that working set, with different reserves. The @@ -565,12 +668,43 @@ In the last two the watchdog stops the run (exit 70) and writes its `abort.json` table's copies are in `bench/results/tiers/aborts/`. Without `--no-fit-check` both are refused with exit 2 (`bench/results/refusals/qwen-image-2.1-1024-tier16.json` and `qwen-image-2.1-1024-tier24.json`). +The ERNIE-Image rows come from these commands, with the same prompt: + +``` +mlx-dfloat generate --model ernie-image-turbo --tier 32 \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report bench/results/tiers/ernie-image-turbo-1024.json +mlx-dfloat generate --model ernie-image --tier 32 \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report bench/results/tiers/ernie-image-1024.json +mlx-dfloat generate --model ernie-image-turbo --tier 24 --no-fit-check \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 --report bench/results/tiers/ernie-image-turbo-1024-tier24.json +mlx-dfloat generate --model ernie-image-turbo --tier 16 --no-fit-check \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 +mlx-dfloat generate --model ernie-image --tier 16 --no-fit-check \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 +mlx-dfloat generate --model ernie-image --tier 24 --no-fit-check \ + --prompt "A stone lighthouse on a rocky shore at dawn, waves breaking below it and a small fishing boat far out on the water" \ + --seed 42 --height 1024 --width 1024 +``` + +In the last three the watchdog stops the run (exit 70) and writes its `abort.json` next to the output image; the +table's copies are in `bench/results/tiers/aborts/`. Without `--no-fit-check` the last four commands are refused with +exit 2 (`bench/results/refusals/ernie-image-turbo-1024-tier24.json`, `ernie-image-turbo-1024-tier16.json`, +`ernie-image-1024-tier16.json` and `ernie-image-1024-tier24.json`). + Measured on an Apple M1 Max, 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0. | 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` | @@ -582,6 +716,9 @@ Measured on an Apple M1 Max, 32 GB, macOS 27.0.1, mlx 0.32.2, mflux 0.20.0. | 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 | 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` |