STATE.md

HIDREAM-O1-MLX-LAB — STATE

Last updated: 2026-05-09 (session that landed Q8 + Phosphene integration)

TL;DR — where we are

• Q6 is the new sweet spot. 1.30 s/step at 1024×1024, ~36 s per image, ~8.5 GB RAM. 2× faster than Q8 with equivalent quality.
• Q8 still works (2.36 s/step, 11.5 GB RAM) — keep it for deterministic upper-bound RAM use cases.
• Q4 deleted from disk: ships dark, no reason to keep around (regenerable in 5 min if needed).
• Backbone sizes: Q6 backbone 7.95 GB, Q8 backbone 9.96 GB. Custom heads 75 MB.
• Shipped to Phosphene dev as kind="hidream" in agent/image_engine.py (commits 45cad69, 962b353). Default model on Phosphene side will be updated to Q6.
• Showcase battery + A/B vs mflux Z-Image-Turbo done at Q8. At Q6, HiDream is now ~2× faster than Z-Image-Turbo (36s vs 80s) AND has lower deterministic RAM (8.5 GB vs 5.9–29.4 GB variable).
• 19+ sample images in sample_outputs/.
• Lab branch: perf-lab-hidream-o1-mlx, no remote.

What's been done

Date	Work	Commit
2026-05-09	Initial scaffolding (Path B chosen)	746efe9
2026-05-09	Wire mlx-vlm Qwen3VLModel directly (4D mask path)	53eb605
2026-05-09	First working images (mushroom 512, cat/beach/portrait 1024 at Q4)	d944a31
2026-05-09	Q8 conversion + samples (dark aesthetic was Q4, not the model)	2bf029a
2026-05-09	Showcase battery + evaluation + Phosphene plan	0bac049
2026-05-09	Phosphene integration shipped to dev	phos 45cad69
2026-05-09	A/B vs mflux Z-Image-Turbo on 3 prompts	2761ad8
2026-05-09	Phosphene IMAGE_GEN_RESEARCH doc updated	phos 962b353
2026-05-09	--blend-seams post-process (opt-in, below-threshold at Q8)	0583356
2026-05-09	Q6 = sweet spot (2× faster than Q8, same quality) — Phosphene default switched	f4fb0ba + phos 8a48953
2026-05-09	Q6 verified across 10-prompt showcase battery	4d3f18c
2026-05-09	Edit/multi-ref scaffold (WIP — runs but output degenerate)	525b7ec
2026-05-09	BF16 default — Q4/Q6/Q8 all show patch-grid at non-square dims	(next)
2026-05-09	Phosphene default switched to BF16	phos af94bd0
2026-05-09	OSS release prep: HF model card, LICENSE, requirements, gitignore	(next)

Known characteristics (not bugs)

• Patch grid in flat regions — architectural (PATCH_SIZE=32 with no overlap). Mild at Q8. --blend-seams 1 is opt-in but doesn't visibly help.
• Text rendering — short, structured signs work ("BLOOM CAFE"). Long text falls apart.
• Deterministic per-prompt RAM — 11.5 GB at 1024 Q8 regardless of prompt complexity. Z-Image-Turbo varies wildly (5.9–29.4 GB).

Open work / next session candidates

Pick from these, listed roughly cheapest-first:

1. 2048×2048 Q8 generation pass — DONE 2026-05-09. sample_outputs/v4_2048_alice_q8.png — 276 s (9.86 s/step), peak RAM 10.8 GB. Q8 at 2048 is slower per step than Q4 (10s vs 8.4s) due to bandwidth. Output is showcase-grade: detailed cybernetic dress, holographic Cheshire cat, near-legible neon signs.
2. Test the Phosphene integration through the dev panel UI (port 8199). Generate one shot via the Image Studio dropdown, confirm pill goes green, the PNG lands.
3. Edit / multi-reference path — SCAFFOLD LANDED, NEEDS DEBUGGING.
   • build_edit_text_sample, resize_pilimage, calculate_dimensions, patchify_ref_image all ported from upstream pipeline.py + utils.py.
   • --ref-images flag wired in generate_hidream_o1_mlx.py.
   • precompute_text_embeds_with_vision precomputes the text+vision embeds once before the loop (since they don't change with timestep) — a meaningful perf win.
   • Smoke test (synthesized two-color ref, K=1, 28 steps, Q6) runs end-to-end without errors but output is uniform tan/khaki. T2I path with same prompt+seed produces a vibrant abstract correctly, so the model and weights are fine.
   • Debugging done so far (see scripts/hidream_o1/_edit_diag.py and _precompute_diag.py):
     • All shapes verified correct (input_ids 174 with 144 image-placeholders, vision tower outputs 144 features, vinput_mask = 256 tgt + 256 ref, position_ids 686 covering all spans).
     • Vision feature scatter verified mathematically correct — at image_token positions combined equals image_features exactly (diff=0); at text positions combined equals embed_tokens(input_ids) exactly (diff=0). Vision features are well-behaved (mean ~0, std ~0.4).
     • Position_ids structure looks right — text positions are sequential, target span gets fix_point=4096 base (per upstream), ref diffusion span continues sequentially.
   • Remaining suspects (in order of likelihood):
     • Mask construction: maybe text-row causal needs to ALSO see the K image_placeholder positions inside proc.input_ids? Upstream _run_decoder_flash has special handling — the non-flash 4D mask path may treat text positions as needing to see embedded vision features. Worth re-reading qwen3_vl_transformers.py:1486-1520.
     • Position_ids semantic alignment: my appended-vinputs at positions [174..686) get mrope codes from input_ids_pad's vision_tokens portion, but maybe these need to match the appended embedding ORDER not just their positions in input_ids_pad.
     • bf16 underflow in attention with the larger 686-token sequence vs T2I's 268.
   • Samples: sample_outputs/v6_edit_smoke.png (degenerate, synthesized 2-color ref), sample_outputs/v6_edit_cat_real.png (degenerate, real cat photo as ref), sample_outputs/v6_edit_t2i_baseline.png (T2I works fine same prompt+seed).
   • This is the single biggest open item. Would let HiDream replace mflux qwen-edit functionally.
4. Promote Phosphene integration to main after the user has tested on dev panel.
5. Quality-aware post-process — try a cheap learned upscaler instead of the seam blend (e.g. SeedVR2 via mflux's mflux-upscale-seedvr2 to take 1024 → 2048).
6. Text-cache reuse across denoising steps — fork mlx-vlm's Qwen3VLModel to cache the text-portion KV across the 28 denoising calls. ~2-5% speedup max but a real architectural improvement.

Hard stop conditions (still relevant)

• Q4 ships dark — established. Use Q8.
• mx.compile = 0% gain — established. Inference loop is at the floor.
• Splitting safetensors mid-read zeroed weights — fixed in converter; don't re-introduce.

How to ramp up fast (next session)

1. cd /Users/salo/HIDREAM-O1-MLX-LAB-active
2. cat README.md CLAUDE.md STATE.md docs/EVALUATION.md (in that order)
3. git log --oneline | head -10 to see where we are
4. ls sample_outputs/ to see what's been generated
5. To regenerate or extend: see the commands in CLAUDE.md

Disk situation snapshot

As of 2026-05-09 the data volume /dev/disk3s5 had ~45 GB free of 926 GB after the user's mid-session cleanup (deleted phosphene-model-lab.git and comfy.git, freed ~83 GB). The lab itself is ~16 GB on disk (10 GB Q8 + 6 GB Q4 models + 1.5 GB venv + samples + lab code). Do not re-download the HiDream HF source unless mlx_models/hidream-o1-dev-q4/ AND mlx_models/hidream-o1-dev-q8/ both go missing — both can be regenerated from one HF download.