--- license: apache-2.0 base_model: lodestones/Chroma1-HD base_model_relation: quantized pipeline_tag: text-to-image library_name: mlx tags: - mlx - apple-silicon - text-to-image - chroma - flux - quantized --- # Chroma1-HD in MLX A pure-MLX port of [lodestones/Chroma1-HD](https://huggingface.co/lodestones/Chroma1-HD) for Apple Silicon, with 4-bit and 6-bit transformer checkpoints and an 8-bit T5-XXL text encoder. It runs on a 16 GB Mac. On a Mac mini M4 (16 GB) a 1024 × 1024 image with the model authors' official settings takes about 18 minutes and a 512 × 512 draft about 4.5 minutes. Against diffusers' `ChromaPipeline` on the same machine, the port is 3.0× faster per denoising step at 512 × 512 and 1.7× faster at 768 × 1024 with either transformer. It needs 41–44% less memory with the 4-bit transformer and 25–29% less with the 6-bit one (see [Measured speed-up over PyTorch](#measured-speed-up-over-pytorch)). Everything that runs at generation time is MLX code in `mlx_chroma/`: the T5-XXL encoder, the MMDiT transformer, the FLUX VAE decoder, the sampler, the pipeline and the command line. PyTorch is not needed. This is an unofficial port. Chroma1-HD is an 8.9B-parameter model by lodestones, based on FLUX.1-schnell. ![A red fox sitting in snow, generated with the 4-bit transformer at 1024 × 1024](images/fox_1024_q4.jpg) *q4, 1024 × 1024, 26 steps, CFG 3.8, seed 42: "a red fox sitting in fresh snow at golden hour, photograph". 17.7 minutes on a Mac mini M4.* ## Measured speed-up over PyTorch Before this port, the same Mac ran Chroma1-HD with diffusers 0.42 `ChromaPipeline`, using the bf16 T5 and the similar-size [GGUF Q6_K transformer](https://huggingface.co/silveroxides/Chroma1-HD-GGUF) (7.65 GB), both on MPS. Both pipelines got the same prompt (the official tiger example), 3 steps, CFG 3.8 and seed 0, on a Mac mini M4 with 16 GB. Wall times include all loading and decoding; memory is the whole process's peak footprint from `/usr/bin/time -l`. | | 512 × 512: PyTorch → MLX q4 | Gain | 768 × 1024: PyTorch → MLX q4 | Gain | |---|---|---|---|---| | Time per denoising step | 29.3 → 9.65 s | **3.0×** | 50.3 → 29.3 s | **1.7×** | | Whole 3-step run, loading included | 136.3 → 38.0 s | **3.6×** | 198.3 → 98.5 s | **2.0×** | | 26-step run (projected: 3-step run + 23 steps) | 13.5 → 4.3 min | **3.1×** | 22.6 → 12.9 min | **1.8×** | | Text encoding, tokenizer and T5 load included | 22 → about 5 s | about 4.4× | 18 → about 5 s | about 3.6× | | Transformer load (the PyTorch figure includes the VAE) | 20 → 2.7 s | 7.4× | 21 → 2.7 s | 7.8× | | VAE decode | 1–2 → 0.9 s | | 3–4 → 2.5 s | | | Peak memory footprint | 12.0 → 7.0 GiB | **−41%** | 13.7 → 7.7 GiB | **−44%** | | Swapped out during the run (system-wide) | 2.3 GiB → 0 | | 3.0 GiB → 0 | | Measured 26-step MLX runs agree with the projection: 248–254 s (4.1–4.2 min) for q4 at 512 × 512. q6 is as fast per step (9.67 and 29.3 s) and peaks at 9.0 and 9.7 GiB (−25% and −29%). Its 3-step runs took 45.9 and 106.7 s (3.0× and 1.9×), because it was quantized at load time in these runs (10.4 s); it now loads from its checkpoint in 3.5 s. Where the gain comes from: - **The implementation.** diffusers attends to 512 padded text tokens, while this port keeps only the real tokens plus one pad (114 and 118 tokens for the tiger prompt and the negative prompt). Given the same trimmed text embeddings, the PyTorch transformer still takes 26.0 s per step at 512 × 512, against 9.8 s for MLX q4: 2.7× from the implementation alone. - **Memory.** The PyTorch runs needed 12.0–13.7 GiB, and the system swapped out 2.3–3.0 GiB during them, which also costs time. The MLX runs stayed at 7.0–9.7 GiB without swapping. - **Accuracy.** q6 keeps the per-step speed-up (9.67 and 29.3 s per step) and is closer to the exact float32 model than the GGUF Q6_K transformer at every timestep (mean velocity error 4.13% against 4.77%). q4, used in the table for its lower memory, is less accurate (12.76%); see [Accuracy](#accuracy). ## Quick start You need an Apple Silicon Mac, Python 3.12 or newer, the `hf` command (`pip install -U huggingface_hub` or `uv tool install huggingface_hub`), and 19.3 GB of disk for the download. ```bash hf download tillknuesting/chroma1-hd-mlx --local-dir chroma1-hd-mlx cd chroma1-hd-mlx python3.12 -m venv .venv # or: uv venv --python 3.12 .venv/bin/pip install -e . # or: uv pip install -e . .venv/bin/python generate_mlx.py --prompt "a red fox sitting in fresh snow at golden hour, photograph" --seed 1 ``` This writes `outputs/mlx__.png`. The PNG's text chunks store the prompt, the seed, every setting and the transformer that was used, so a run can be reproduced. Add `--width 512 --height 512` for a 4.5-minute draft. transformers prints "PyTorch was not found" when it loads; only its tokenizer is used, so that is expected. The code finds the weights relative to its own directory, so install it with `-e` (editable) and keep the `models/` folder next to `mlx_chroma/`. Every path can also be passed explicitly (`--transformer`, `--t5`, `--tokenizer`, `--vae`). Run one generation at a time on a 16 GB Mac. The pipeline keeps one large model in memory at a time: the T5 encoder, then the transformer, then the VAE. ## Files | Path | Size | Contents | |---|---|---| | `models/mlx/transformer-q6` | 7.58 GB (7.06 GiB) | 6-bit transformer (MLX affine quantization, group size 64) | | `models/mlx/transformer-q4` | 5.43 GB (5.06 GiB) | 4-bit transformer | | `models/mlx/t5-q8` | 6.13 GB (5.71 GiB) | T5-XXL encoder, 8-bit; its FFN output projection (`wo`) and embedding stay bf16 | | `models/Chroma1-HD/tokenizer`, `models/Chroma1-HD/vae` | 0.8 MB, 168 MB | Unchanged from lodestones/Chroma1-HD | | `mlx_chroma/` | | The port: `t5.py`, `transformer.py`, `quantize.py`, `vae.py`, `sampling.py`, `pipeline.py` | | `generate_mlx.py` | | The command line (a thin wrapper around `mlx_chroma.pipeline.main`) | ## Choosing the transformer q6 and q4 run at the same speed, but q6 is about three times closer to the exact model (see [Accuracy](#accuracy)). q6 needs 2 GiB more memory, though. The default, `--transformer auto`, therefore decides at the start of each run: 1. Available memory = `vm_stat` free + inactive + speculative + purgeable pages. 2. The memory a whole run needs at the requested size: the weights + 1.66 GiB + 0.375 GiB per 1024 image tokens (image tokens = W/16 × H/16), and at least 6.3 GiB (the T5 stage). 3. Pick q6 if the available memory covers q6's need plus a 1 GiB margin, otherwise q4. If the kernel reports memory pressure (level 2 or higher), or `vm_stat` is unavailable, pick q4. 4. If even q4 does not fit, print a warning: macOS will compress and swap other apps' memory. | Size | Image tokens | q4 needs | q6 needs | `auto` picks q6 from | |---|---|---|---|---| | 512 × 512 | 1024 | 7.09 GiB | 9.09 GiB | 10.09 GiB available | | 768 × 1024 | 3072 | 7.84 GiB | 9.84 GiB | 10.84 GiB | | 1024 × 1024 | 4096 | 8.22 GiB | 10.22 GiB | 11.22 GiB | | 1152 × 1152 | 5184 | 8.61 GiB | 10.62 GiB | 11.62 GiB | The estimate is an upper envelope of the six whole-process peaks it was fitted to, from 256 × 256 to 1152 × 1152 (`/usr/bin/time -l` "peak memory footprint"; sampled at 256 × 256). Two later 26-step runs at 512 × 512 peaked about 0.01 GiB above it (7.10 and 9.11 GiB against 7.09 and 9.09), which the 1 GiB margin covers. On a 16 GB Mac with other apps open, `auto` usually picks q4 at 1024 × 1024. The decision and its numbers are printed after the seed and stored in the PNG metadata (`transformer_choice`): ``` transformer auto -> q4 at 256x256: 7.85 GiB available now < 9.81 GiB (q6's ~8.81 GiB + 1.00 GiB margin); q4 needs ~6.81 GiB; q6 needs ~8.81 GiB (vm_stat free 3.22 + inactive 3.84 + speculative 0.68 + purgeable 0.12 GiB; other apps' anonymous memory 4.40 GiB, compressor 0.50 GiB holding 9.73 GiB) ``` `--transformer q6` or `q4` forces a variant. ## Command line `--help` lists every flag. The main ones: | Flag | Default | Meaning | |---|---|---| | `--prompt` | (required) | The prompt. | | `--negative` | the official negative prompt | `""` means an empty negative prompt, which is still encoded. | | `--width`, `--height` | 1024, 1024 | Multiples of 16. | | `--steps` | 26 | | | `--cfg` | 3.8 | `<= 1` skips the negative pass. | | `--schedule` | `beta` | `beta` is the official ComfyUI workflow (shift 1, alpha = beta = 0.45). `diffusers` is diffusers' FlowMatchEuler schedule with shift 3. | | `--seed` | random (printed) | | | `--out` | `outputs/mlx__.png` | | | `--transformer` | `auto` | `auto`, `q4`, `q6`, or a checkpoint directory. `q3`, `q5`, `q8` and `bf16` need the original weights (see [Re-creating the checkpoints](#re-creating-the-checkpoints)). | | `--cfg-mode` | `separate` | `separate` makes two batch-1 calls per step; `batched` makes one batch-2 call with a text mask. | | `--vae-dtype` | `float32` | | | `--text-embeds`, `--save-text-embeds` | | Load or save the T5 embeddings as `.npz`. Loading skips T5. | | `--save-latents` | | Save the final packed latents as a float32 `.npy`. | | `--torch-noise` | off | Draw diffusers' torch CPU-generator noise, for cross-framework comparisons. Needs `torch`; it is the only flag that imports it. | | `--quiet` | off | Hide the per-step lines. | ## Python API ```python from mlx_chroma.pipeline import ChromaMLXPipeline pipe = ChromaMLXPipeline() # transformer "auto"; nothing is loaded yet image = pipe.generate("a red fox in snow", seed=1) # PIL.Image; settings in image.info res = pipe.run("a red fox in snow", seed=1, width=768, height=1024) # res: GenerationResult(image, latents (1, N, 64), seed, metadata, timings, peak_memory, text, sigmas, image_stats) # Several prompts: one T5 load, and the transformer stays loaded between images. texts = pipe.encode_many(["a red fox in snow", "a lighthouse at dusk"]) lats = [pipe.denoise(t, seed=s) for t, s in zip(texts, (1, 2))] pipe.release_transformer() imgs = [pipe.decode(l, 1024, 1024) for l in lats] ChromaMLXPipeline(transformer="q6") # force a variant ``` The module docstring of `mlx_chroma/pipeline.py` documents the stages, the memory model and every helper. ## Speed and memory Measured in October 2026 on a Mac mini M4 (10-core GPU, 16 GB, 11.84 GiB GPU working set), macOS 26.5, MLX 0.32.4, with the owner's usual apps open. "s/step" is one denoising step with CFG, which is two transformer calls. | Size | Variant | s/step | 26 steps | Peak footprint | |---|---|---|---|---| | 512 × 512 | q4 | 9.26–9.44 | 248–254 s | 7.10 GiB | | 512 × 512 | q6 | 9.63–9.94 | 262–276 s | 9.11 GiB | | 768 × 1024 | q4 / q6 | 29.3 | about 13 min (projected) | 7.67 / 9.67 GiB | | 1024 × 1024 | q4 | 40.5–41.0 | 1062–1077 s (17.7–18.0 min) | not measured (MLX peak 5.96 GiB) | | 1152 × 1152 | q4 | 53.8 | about 23.5 min (projected) | 8.61 GiB | Loading is quick: the T5 encoder loads in about 3 s, the q4 transformer in 2.6 s and q6 in 3.5 s. Encoding the prompt and the negative prompt usually takes about a second, and the VAE decode 0.9 s at 512 × 512. The comparison with PyTorch is in [Measured speed-up over PyTorch](#measured-speed-up-over-pytorch). The MLX transformer runs at about 3.46 TFLOP/s on the M4 (66.5 TFLOP per 1024 × 1024 forward pass). Fused projections, `mx.compile` and float16 activations were all measured, and none was faster. More speed needs fewer FLOPs: fewer steps or smaller images. ## Accuracy Each component was checked in float32 against diffusers or transformers, on the real weights unless noted: | Component | Check | Result | |---|---|---| | T5-XXL encoder | vs transformers, 14 prompts from empty to 512 tokens | cos 1.0, max relative error 2.0e-5 | | Tokenization | vs diffusers' tokenizer call, 21 cases plus truncation of a 991-token prompt | identical | | Transformer | full depth vs diffusers' `ChromaTransformer2DModel` | cos 1.000000000, max relative error 1.56e-5 | | VAE decoder | vs diffusers at 512 and 1024 | cos 1.0, max relative error 2.8e-5 | | Sampling | packing, ids, both schedules, Euler step, CFG, noise | 20/20; the diffusers schedule is bit-exact | | Pipeline | vs diffusers' `ChromaPipeline` on tiny random models | 10/10, max relative error 5.3e-7 | The shipped checkpoints trade a little of that for memory: - **T5 (`t5-q8`)**: cos 0.99982 and relative L2 error 1.9% against float32, over the same 14 prompts. - **Transformer**: relative L2 error against float32 per forward pass is 1.47% for bf16, 2.67% for q6 and 8.47% for q4. For the final latents of a 3-step run with CFG 3.8 it is 1.70%, 2.65% and 7.65%. To measure each variant's error at every timestep, every variant was evaluated on the latents of the float32 trajectory (10 steps, 512 × 512, CFG 3.8), and its guided velocity was compared with the float32 one: | Relative L2 error of the guided velocity | mean | worst step | |---|---|---| | MLX bf16 | 2.75% | 3.81% | | **MLX q6** | **4.13%** | **5.83%** | | PyTorch, GGUF Q6_K (diffusers on MPS) | 4.77% | 7.62% | | MLX q4 | 12.76% | 19.79% | q6 is closer to the exact model than the GGUF Q6_K reference at every timestep. q4 is clearly the least accurate, but its images are still coherent: ![q4 and q6 side by side at 512 × 512](images/q4_vs_q6_512.jpg) *512 × 512, 26 steps, CFG 3.8, seed 42; q4 on the left, q6 on the right. Top: the official tiger example prompt. Bottom: "a red fox sitting in fresh snow at golden hour, photograph".* 512 × 512 images are coherent but softer than 1024 × 1024; fur and whiskers lose detail. ## Differences from diffusers These are deliberate, and each one follows the model authors' official ComfyUI workflow: - **Text masking.** The transformer sees only the real prompt tokens, including EOS, plus one padding token, which is the same as hard masking. diffusers 0.42 casts its 0/1 mask to float, and PyTorch's SDPA adds a float mask to the logits, so the padding tokens leak into attention. On tiny models that changes the latents by 1.9e-2 relative; with a boolean mask diffusers matches this port to 5e-7. The prompt and the negative prompt are each encoded at their exact length. - **Default sampling.** The beta schedule (shift 1, alpha = beta = 0.45), 26 steps, CFG 3.8 and the official negative prompt. `--schedule diffusers` gives diffusers' schedule (shift 3) instead, bit-exact. - **Timesteps in float32.** diffusers casts the timestep to bf16 before multiplying it by 1000, which rounds 1000·t to steps of about 4 above 500 (about 2 between 250 and 500). This port keeps float32 there, as ComfyUI does. - **VAE in float32.** bf16 would decode 0.7 s faster at 1024 × 1024, but its relative error is 2.4–2.5% against 3.5e-6. ## How the checkpoints were made - **Transformer.** MLX affine quantization with group size 64 and one bit width for all quantized layers. `x_embedder`, `context_embedder`, the approximator (`distilled_guidance_layer`) and `proj_out` stay bf16 (0.54 GiB); a 4-bit approximator alone gives 24.8% error. At run time, quantized layers with at least 512 input rows are dequantized to a dense bf16 matmul, which is 10–13% faster for 0.16 GiB more memory. - **T5.** 8-bit with group size 64 for q, k, v, o, wi_0 and wi_1. The FFN output projection `wo` and the embedding stay bf16, because quantizing `wo` raises the empty prompt's error from 2.9% to 17.8%. q, k and v are computed on float32 activations. A 4-bit T5 is too lossy (cos 0.977). - **No q8 transformer.** q8 is more accurate (1.74% per forward pass against q6's 2.67%), but at 9.06 GiB it swapped in every run on the 16 GB Mac with other apps open. ### Re-creating the checkpoints To use `--transformer q8`, `q5`, `q3` or `bf16`, or to rebuild the transformer checkpoints, first download the original transformer (17.8 GB). To rebuild `t5-q8` as well, add `--include "text_encoder/*"` (9.5 GB more): ```bash hf download lodestones/Chroma1-HD --include "transformer/*" --local-dir models/Chroma1-HD ``` With them in place, `--transformer q8` (or `q5`, `q3`) quantizes the original weights at load time, and `bf16` streams them block by block (slow, but only 0.8 GiB of memory). To save a checkpoint instead (about 25 s for q8; the conversions stream module by module and peak at 1.2–1.6 GiB): ```bash .venv/bin/python -c "from mlx_chroma.quantize import quantize_and_save as q; q('models/Chroma1-HD', 'models/mlx/transformer-q8', 8)" ``` `--transformer q8` then loads `models/mlx/transformer-q8`. `auto` only ever picks q6 or q4. The T5 checkpoint was made with: ```python from mlx_chroma import t5 t5.save_quantized("models/Chroma1-HD/text_encoder", "models/mlx/t5-q8", bits=8, group_size=64, modules=("q", "k", "v", "o", "wi_0", "wi_1"), embedding=False) ``` Both converters refuse to run if less than 10 GiB of disk would remain afterwards. ## Limitations - Apple Silicon only. All measurements come from one machine: a Mac mini M4 with 16 GB. - Text to image only: no image-to-image, inpainting or LoRA loading. - On a 16 GB Mac, other apps decide how much memory a run gets. `auto` accounts for that at the start of a run, but apps that grow during an 18-minute run can still push macOS into swapping. - Text in images can be misspelled; one q4 run at 1024 × 1024 wrote "CHIROMA1-HD" for "CHROMA1-HD". - Like the original model, it has no content filter. - The code comments mention the development test suite (`tests/...`), which ran the parity checks above against diffusers and transformers. It is not part of this repository. ## License and credits - **Weights:** derived from [lodestones/Chroma1-HD](https://huggingface.co/lodestones/Chroma1-HD) (Apache-2.0), which is based on [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) (Apache-2.0). The T5-XXL encoder, the tokenizer and the VAE are the ones shipped with Chroma1-HD. The quantized weights here are distributed under the same Apache-2.0 license. - **Code:** Apache-2.0 (see `LICENSE` and `NOTICE`). It is an MLX port of Hugging Face code under the same license: `transformer.py`, `vae.py`, `sampling.py` and `pipeline.py` follow diffusers' `ChromaTransformer2DModel`, `AutoencoderKL` decoder, `FlowMatchEulerDiscreteScheduler` and `ChromaPipeline`, and `t5.py` follows transformers' `T5EncoderModel`. The default sampler reimplements the beta schedule ([Lee et al., 2024](https://arxiv.org/abs/2407.12173)) that the authors' official ComfyUI workflow uses; no ComfyUI code is included. - This port is not affiliated with lodestones or Black Forest Labs.