Skip to content

Run Z-Image Turbo and Z-Image from DFloat11 checkpoints - #12

Merged
IonDen merged 4 commits into
mainfrom
feature/s2-zimage
Oct 8, 2026
Merged

IonDen merged 4 commits into
mainfrom
feature/s2-zimage

Conversation

@IonDen

@IonDen IonDen commented Oct 8, 2026 •

Copy link
Copy Markdown
Owner

mlx-dfloat could only run FLUX.1. This adds Z-Image-Turbo and Z-Image, two current image models with DFloat11 checkpoints, and the layer that lets more mflux model families follow.

The FLUX.1 integration hooks mflux's per-block methods, and Z-Image's transformer has none: it calls its blocks inline. The new layer swaps each block's class for a subclass that installs that block's weights from the compressed set just before it runs and puts placeholders back afterwards, so every parameter keeps the path mflux expects. Weight names come from mflux's own Z-Image weight mapping, including the renames outside the blocks. The one compressed group that is not a block, the caption embedder, is decoded once when the set loads. mflux compiles Z-Image's denoising step on M1/M2 Max and Ultra and on M3 or later, and evaluating after each block cannot happen inside a compiled function, so the adapter takes mflux's uncompiled path on every chip and refuses to run if mflux hands back a compiled step anyway. A model registry backs mlx-dfloat generate --model z-image-turbo / --model z-image and a new DFloatModel(name), which returns a ready model. FLUX.1 keeps its own hook and behaves as before.

Measured on an M1 Max with 32 GB (macOS 27.0.1, mlx 0.32.2, mflux 0.20.0):

  • Z-Image's DFloat11 checkpoint decodes to exactly its BF16 original: all 271 compressed matrices and the 250 tensors stored uncompressed.
  • Z-Image-Turbo's original is FP32, so there is no BF16 original to compare with. The GPU kernel matches the CPU reference on all 35 groups, and five sampled groups (32 of 271 matrices) equal the FP32 original rounded to BF16 (round to nearest, ties to even). Turbo's latents have not been compared with a reference yet.
  • Z-Image's final latents at 1024², 4 steps, guidance 4.0 are bit-identical to the same transformer streaming its BF16 weights, through the same uncompiled path on both sides.
  • At 1024², Turbo (9 steps) peaked at 13.11 GiB and Z-Image (50 steps, guidance 4.0) at 13.15 GiB, both during the VAE decode. The memory estimate is calibrated to reproduce those peaks in that phase.
  • A 16 GB Mac was not measured. On the 32 GB Mac under a 16 GB Mac's limits, with the fit check turned off (it predicted 12.2 GiB and refused the run), the watchdog stopped Turbo after 6.7 s while encoding the prompt, at 9.29 GiB against a 9.17 GiB ceiling, so that size is marked at risk.

FLUX.1 is unchanged: compared with the commit this branch started from, schnell's 1024² image is pixel-identical and the watched peaks are 0.019 GiB apart, and schnell's latents are still bit-identical to the BF16-streamed run.

Smaller changes users will see: the default Z-Image checkpoints are fetched at fixed Hub revisions; generate says on stderr when a model spec is fetched from the Hub, warns when a negative prompt is ignored because guidance is 1 or less, and records elapsed_seconds in its --report; a watchdog stop records the phase that was running.

Also in this PR: scripts/verify_remote_group.py --cast-fp32-to-bf16 for checkpoints whose original is FP32, Z-Image support in scripts/verify_image.py (Turbo is refused there, since it has no BF16 original), and tier-table rows for runs the watchdog stopped, shown as lower bounds with the time they ran.

@IonDen IonDen added the enhancement New feature or request label Oct 8, 2026
@IonDen IonDen self-assigned this Oct 8, 2026
@IonDen
IonDen force-pushed the feature/s2-zimage branch 2 times, most recently from a4ae127 to da1e1bf Compare October 8, 2026 01:04
Base automatically changed from feature/0029-canary-selftest to main October 8, 2026 08:39
IonDen added 4 commits October 8, 2026 10:41
Block seam for transformers without per-block hooks, name maps that rename and transform non-block extras, a resident non-block group, shared Hub and lifecycle helpers, a family registry with DFloatModel(name), and a generic generate CLI. The FLUX.1 modules keep their import paths through re-exports.
Adds the Z-Image name map, transformer build, memory phases measured at 1024 by 1024, and DFloatZImage with the shared encoder and set lifecycle.
The sampled check can compare against an FP32 original rounded to BF16, the image identity check covers Z-Image, and the tier table shows Z-Image and watchdog stops.
README and CHANGELOG describe Z-Image generation; the tier results hold the 1024 rows and the 16 GB tier abort.
@IonDen
IonDen force-pushed the feature/s2-zimage branch from da1e1bf to 7e29b69 Compare October 8, 2026 08:44
@IonDen
IonDen merged commit 14c1433 into main Oct 8, 2026
6 checks passed
@IonDen
IonDen deleted the feature/s2-zimage branch October 8, 2026 08:49
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