Repository navigation
Run ERNIE-Image and ERNIE-Image-Turbo from their DFloat11 checkpoints - #16
Merged
Merged
Conversation
…oise 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.
The adapter, the registry and command entries, the memory constants measured at 1024², and the committed records they and the tests read.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This adds ERNIE-Image and ERNIE-Image-Turbo to
mlx-dfloat generateand the Python API (DFloatModel("ernie-image"),DFloatModel("ernie-image-turbo")). Both run from their DFloat11 checkpoints,mingyi456/ERNIE-Image-DF11andmingyi456/ERNIE-Image-Turbo-DF11, pinned to fixed revisions. The text encoder, VAE and tokenizer come from thebaidurepositories, also pinned; their transformer and the prompt-enhancer folder are not downloaded.ERNIE-Image calls its 36 blocks inline, so it uses the per-block hook from Z-Image: each block gets its decoded weights just before it runs, and the three compressed groups that are not blocks are decoded once when the weights load. mflux compiles the denoising step on M1 and M2 Max and Ultra and on M3 or later; this path runs it uncompiled on every chip. The defaults are mflux's own and match each model card: ERNIE-Image runs 50 steps at guidance 4.0, ERNIE-Image-Turbo 8 steps at guidance 1.0. The base runs classifier-free guidance as one batched transformer call per step, so each block is decoded once per step. Turbo refuses another guidance and ignores a negative prompt, as mflux's command does. An empty or blank prompt is refused.
Measured on an M1 Max with 32 GB (macOS 27.0.1, mlx 0.32.2, mflux 0.20.0), 1024², seed 42. Peak is the process's memory footprint as the OS counts it:
bench/results/parity/ernie-image/full-vs-bf16.jsonbench/results/parity/ernie-image-turbo/full-vs-bf16.json…/kernel-vs-reference.jsonin each folderbench/results/identity/ernie-image-1024/compare.jsonbench/results/tiers/ernie-image-1024.jsonbench/results/tiers/ernie-image-turbo-1024.jsonOn a 32 GB Mac both keep the compressed transformer loaded through the VAE decode. Smaller Macs were simulated on the 32 GB Mac under their caps, and none was measured. At 24 GB the fit check refuses ERNIE-Image-Turbo: it predicts 14.39 GiB for the denoising step against a 14.00 GiB budget. A forced run passed at 14.39 GiB, only 0.11 GiB under the watchdog ceiling, and an identical earlier run peaked at 14.03 GiB, so the estimate takes the higher of the two. ERNIE-Image at 24 GB and both models at 16 GB are refused, and forced runs were stopped by the watchdog. All of these sizes are marked at risk. The estimate's terms come from the runs in
bench/results/calibration/ernie-image{,-turbo}-1024.json, and a test checks every committed 1024² record against the current estimate.mflux runs ERNIE-Image's hidden state in float32, because its timestep embedding is float32. This path keeps that, which is why the output stays bit-identical to stock mflux. No speed is claimed.
Changes that reach every model:
encode()now runs under the same memory caps asgenerate_image. The wired limit is read once per process instead of at every call, and a wired cap MLX refuses (a system wired limit set lower) leaves the memory cap in place with one warning instead of raising.The FLUX.1, Z-Image, FLUX.2 Klein and Qwen-Image 2.1 model code is otherwise unchanged. Each of the four commits passes the test suite on its own; the README table check passes from the last one.