feat: Port Ideogram 4 (ideogram4) text-to-image with FP8 weights #2427

Part of #2348

Problem / Background

Ideogram 4 is a non-commercial, license-gated text-to-image model: a 34-layer single-stream DiT (hidden 4608) conditioned on 13 stacked hidden-state layers of a Qwen3-VL 8B text model, a separate unconditional transformer for classifier-free guidance, the FLUX.2 VAE, and logit-normal sampler presets. The public checkpoint ships FP8 weights with per-row scales. mlxcel cannot run it; #524 recorded it as out of scope and #2348 now brings image generation in scope. The FLUX.2 VAE it decodes with is already ported by #2341.

Upstream reference

• Model directory: https://github.com/Blaizzy/mlx-vlm/tree/main/mlx_vlm/models/ideogram4 (README). PRs: Add support for the Ideogram 4 image generation model Blaizzy/mlx-vlm#1281 (model), Add Ideogram 4 local prompt expansion support Blaizzy/mlx-vlm#1276 (local prompt expansion, out of scope here), Batched num_images generation for qwen_image, flux2, ideogram4 Blaizzy/mlx-vlm#2355 (batched num_images), Use model-based sampling defaults for image generation Blaizzy/mlx-vlm#2373 (sampling defaults), removing duplicate rotate_half and check_array_shape funcs Blaizzy/mlx-vlm#2414 (shared rotate_half).
• Transformer: transformer.py (Ideogram4Transformer: Ideogram4MRoPE with mrope_section (24, 20, 20) and theta 5,000,000, Ideogram4RMSNorm, Ideogram4Attention, Ideogram4MLP, AdaLN width 512, token-type indicator embedding with LLM_TOKEN_INDICATOR / OUTPUT_IMAGE_INDICATOR, segment_ids), config fields and defaults at config.py#L16-L48.
• Weights: weights.py (dequantize_fp8_weight_only: <name>.weight FP8 E4M3 bytes times <name>.weight_scale broadcast per output row, to bf16; other floats cast to bf16; text encoder from text_encoder/config.json text_config with language_model. stripped and only embed_tokens., layers., norm. kept; VAE through FLUX.2 load_vae). Two transformers: transformer/ (conditional) and unconditional_transformer/.
• Pipeline: pipeline.py#L127-L283: constants at L27-L33 (QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, ..., 33, 35), IMAGE_POSITION_OFFSET 65536, patch 2, VAE scale 8, latent dim 128, max 2048 text tokens); chat-template tokenization at L301-L316; inputs at L318-L362 (3-axis position ids, text positions (i, i, i), image positions offset by 65536); text features at L364-L385 (hidden_states[layer + 1] for each activation layer, stacked on a trailing axis and flattened hidden-major to [B, T, 4096 * 13], zero-padded over image tokens, f32); denoise loop runs steps in reverse over logit-normal intervals, both transformers every step, v = g_i * pos + (1 - g_i) * neg, z += v * (s - t); noise from mx.random.seed(seed) then normal((n, tokens, 128)) f32; decode at L387-L404 (latent norm shift and scale from latent_norm.py, unpatchify, FLUX.2 VAE, clip to [-1, 1]).
• Scheduler: scheduler.py: LogitNormalSchedule (logsnr clamp -15..18), presets V4_QUALITY_48 (48 steps, guidance 3 x3 then 7, std 1.5), V4_DEFAULT_20 (20 steps, 3 x2 then 7, std 1.75), V4_TURBO_12 (12 steps, 3 then 7, mu 0.5, std 1.75), get_schedule_for_resolution (mean shifted by 0.5 * ln(pixels / 512^2)), make_step_intervals. Defaults resolve to the preset's step count and last guidance value (model.py#L81-L84): 20 steps, guidance 7.0.
• Prompting: prompting.py#L127-L223: a JSON-object prompt passes through (with caption warnings); plain text is wrapped into a minimal JSON caption (_minimal_caption) unless auto_json_caption is false; the wrapped text is returned as revised_prompt.
• Checkpoint: ideogram-ai/ideogram-4-fp8 (gated, Ideogram 4 Non-Commercial license, about 27.5 GB: two 9.3 GB transformers, 8.8 GB text encoder, FLUX.2 VAE).

Current Behavior

• No image generation on main; feat: Add image generation runtime: generate-image CLI, /v1/images routes, shared VAE and scheduler #2341 adds the runtime, probe (class-name scan) and Flux2Vae.
• FP8 support in mlxcel is block-scaled only: src/models/fp8_block.rs handles <name>.weight_scale_inv per 128x128 block (FP8_BLOCK_SCALE_SUFFIX, L72) for Qwen and DeepSeek checkpoints. The E4M3 decoder it uses, f8_e4m3_to_f32 (src/models/sanitize.rs:101), is the reuse point; there is no per-row weight_scale path.
• The Qwen3 hidden-state conditioner comes from feat: Add image generation runtime: generate-image CLI, /v1/images routes, shared VAE and scheduler #2341 (contract below).

Proposed Solution

New module src/imagegen/ideogram4/ (config.rs, transformer.rs, weights.rs, scheduler.rs, latent_norm.rs, prompting.rs, pipeline.rs, mod.rs) implementing ImageGenerationModel with supports(Generate) only, family id ideogram4.

• weights.rs: dequantize_fp8_rowwise(weights) -> WeightMap decoding E4M3 bytes with f8_e4m3_to_f32 and multiplying by the per-row weight_scale, producing bf16, then the project's bf16 to f16 conversion at this boundary unless the precision check below fails. Unpaired weight_scale keys are dropped; a weight_scale whose shape is not [out] errors naming the tensor. Put it here rather than in fp8_block.rs, whose contract is block scales; add a one-line pointer comment in fp8_block.rs so the next per-row user finds it.
• transformer.rs: port of Ideogram4Transformer; both transformer/ and unconditional_transformer/ load into it (the unconditional one is called with no LLM features).
• Text features: the shared Qwen3 conditioner loaded from text_encoder/ (nested text_config, language_model. prefix, FP8 dequantized first), then hidden_states[layer + 1] for the 13 activation layers stacked hidden-major exactly as upstream (stack on a new last axis, reshape [B, T, 4096 * 13]). This order differs from FLUX.2's stacked_prompt_embeds; do not reuse that helper here.
• scheduler.rs, latent_norm.rs, prompting.rs: ports of the upstream files (NormalDist().inv_cdf needs an inverse normal CDF; use an implementation accurate to 1e-12 and test it against the upstream values). Prompt expansion (prompt_expansion_model) is out of scope; a request carrying it returns a clear unsupported error rather than silently ignoring it.
• pipeline.rs: generate as upstream; --gen-kwargs keys sampler_preset, auto_json_caption, guidance_schedule, mu, std; --steps and --guidance override the preset as upstream does (an explicit guidance disables the per-step schedule). revised_prompt in the server response carries the wrapped caption.
• Register ideogram4 in the feat: Add image generation runtime: generate-image CLI, /v1/images routes, shared VAE and scheduler #2341 FAMILIES table and the registry standalone family table.

Scope

In scope: the module above, docs (docs/image-generation.md row noting the license gate, presets and memory footprint; docs/supported-models.md).

Out of scope: local prompt expansion through a second mlxcel model (Blaizzy/mlx-vlm#1276; follow-up), the ideogram-ai/ideogram-4-nf4* and GGUF variants, editing (not supported upstream).

Implementation Notes

• Reuse: f8_e4m3_to_f32; feat: Add image generation runtime: generate-image CLI, /v1/images routes, shared VAE and scheduler #2341 Flux2Vae, runtime and PNG output; the shared Qwen3 conditioner.
• Shared Qwen3 conditioner contract (owned by feat: Add image generation runtime: generate-image CLI, /v1/images routes, shared VAE and scheduler #2341 per its amendment comment; if feat: Add image generation runtime: generate-image CLI, /v1/images routes, shared VAE and scheduler #2341 merged without it, whichever image-generation issue lands first adds it at this exact path): src/imagegen/qwen3_text_encoder.rs Qwen3TextEncoder::load(dir) (top-level Qwen3 config or nested text_config; strips model. and language_model.; drops lm_head.* and non-text keys) and hidden_states(input_ids, attention_mask) -> Vec<MlxArray> with num_layers + 1 entries, index 0 the embeddings and index i decoder layer i output before the final norm, causal plus padding mask, positions 0..T. For this model it must accept already-dequantized weights from dequantize_fp8_rowwise (load from a WeightMap, not only from disk).
• Memory: after FP8 dequantization both transformers are bf16 (about 18.6 GB together) plus the 8 GB text encoder; drop the text encoder after encoding, as upstream does, before loading the transformers.
• Precision: apply the feat: Add image generation runtime: generate-image CLI, /v1/images routes, shared VAE and scheduler #2341 rule per component; upstream runs bf16 with f32 noise and f32 text features.
• Edge cases: dimensions not divisible by 16, outside 256..=2048, or with aspect ratio beyond 6:1 error; a prompt over 2048 tokens errors with the upstream message; an unknown preset errors listing the valid names; guidance_schedule length not matching steps errors.
• Testing without the gated checkpoint: real-checkpoint tests are #[ignore] and skip when models/ideogram-4-fp8 is absent; unit tests must not need it.

Acceptance Criteria

• mlxcel generate-image -m models/ideogram-4-fp8 -p '{"high_level_description":"A red square icon.","compositional_deconstruction":{"background":"A white background.","elements":[{"type":"obj","desc":"A centered red square."}]}}' --size 1024x1024 --seed 0 -o out.png writes a PNG with the V4_DEFAULT_20 defaults.
• Seed-fixed parity against mlx-vlm at 512x512 with --gen-kwargs '{"sampler_preset":"V4_TURBO_12"}', seed 0: initial noise bit-identical, 4x4-pooled final latents relative RMS at most 4.5e-2 (a per-case bound that rejects a wrong port), PNG PSNR at least 30 dB (#[ignore], committed fixture). It fails on current main.
• A plain-text prompt is wrapped and the server response's revised_prompt holds the caption; auto_json_caption: false passes the text through unchanged.
• Unit tests, each failing on current main: dequantize_fp8_rowwise on a synthetic 4x8 tensor against hand-computed values (including a negative and a subnormal E4M3 byte); LogitNormalSchedule and get_schedule_for_resolution values for all three presets at 512x512 and 1024x768 against upstream literals; inverse normal CDF at 1e-6, 0.5 and 0.999999; minimal-caption wrapping against upstream output.
• mlxcel arch --json lists ideogram4; docs updated.

Verification

cargo build --release --features metal,accelerate
cargo test --workspace --profile test-fast --features metal,accelerate ideogram4
cargo test --workspace --profile test-fast --features metal,accelerate imagegen
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
cargo fmt --check
cargo test --workspace --profile test-fast --features metal,accelerate ideogram4 -- --ignored

Dependencies

Blocked by #2341 (runtime, FLUX.2 VAE, probe, shared Qwen3 conditioner).

Acceptance revision (2026-10-28)

Epic #2348 replaced the final-latent relative RMS below 1e-2 bound with per-case latent bounds, PNG PSNR at least 30 dB against the CUDA oracle, and bit-identical initial noise; measured rel-RMS is recorded as drift in PR #2703. At 512x512 V4_TURBO_12 seed 0: token ids equal, initial noise bit-identical, 4x4-pooled latents 3.16e-2 (bound 4.5e-2), 8x8 image PSNR 43.72 dB (bound 38 dB), full PNG 37.13 dB. Full-latent rel-RMS 6.18e-2 and text-feature rel-RMS 6.40e-2 are recorded drift; on upstream's text features the latents reach 2.87e-2 and the PNG 45.95 dB, so the residual sits in the bf16 Qwen3 conditioner (follow-up #2666). Wrong-port controls: a wrong seed scores 4.2e-1 / 12.3 dB, and ignoring the preset guidance schedule scores 5.2e-2, which only the pooled-latent bound rejects.