Skip to content

Run ERNIE-Image and ERNIE-Image-Turbo from their DFloat11 checkpoints - #16

Merged
IonDen merged 4 commits into
mainfrom
feature/s5-ernie-image
Oct 8, 2026
Merged

IonDen merged 4 commits into
mainfrom
feature/s5-ernie-image

Conversation

@IonDen

@IonDen IonDen commented Oct 8, 2026

Copy link
Copy Markdown
Owner

This adds ERNIE-Image and ERNIE-Image-Turbo to mlx-dfloat generate and the Python API (DFloatModel("ernie-image"), DFloatModel("ernie-image-turbo")). Both run from their DFloat11 checkpoints, mingyi456/ERNIE-Image-DF11 and mingyi456/ERNIE-Image-Turbo-DF11, pinned to fixed revisions. The text encoder, VAE and tokenizer come from the baidu repositories, 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:

Check Result Record
ERNIE-Image: every compressed matrix (256) and stored tensor (153) vs the BF16 original equal bench/results/parity/ernie-image/full-vs-bf16.json
ERNIE-Image-Turbo: every compressed matrix (256) vs the BF16 original, by range reads equal bench/results/parity/ernie-image-turbo/full-vs-bf16.json
GPU decoder vs CPU reference, both models, all 256 matrices equal …/kernel-vs-reference.json in each folder
ERNIE-Image final latents vs the same transformer streaming BF16 weights (4 steps, guidance 4.0, classifier-free guidance on) bit-identical, images pixel-identical bench/results/identity/ernie-image-1024/compare.json
ERNIE-Image, 32 GB, 50 steps peak 16.79 GiB, 29.9 min bench/results/tiers/ernie-image-1024.json
ERNIE-Image-Turbo, 32 GB, 8 steps peak 17.34 GiB, 2.9 min bench/results/tiers/ernie-image-turbo-1024.json

On 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:

  • Each model's encode() now runs under the same memory caps as generate_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.
  • A family can floor its denoise memory term; only ERNIE-Image uses it, for the float32 weight copies its float32 hidden state makes.

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.

IonDen added 4 commits October 8, 2026 18:36
…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.
@IonDen IonDen self-assigned this Oct 8, 2026
@IonDen IonDen added the enhancement New feature or request label Oct 8, 2026
@IonDen
IonDen merged commit a4f10c8 into main Oct 8, 2026
8 of 9 checks passed
@IonDen
IonDen deleted the feature/s5-ernie-image branch October 8, 2026 17:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant