Mflux Model Porting

mflux model porting

Goal

Provide a repeatable, MLX-focused workflow for porting ML models (typically from diffusers repo located near mflux repo in the system) into mflux with correctness first, then refactor to mflux style.

Principles

• Match the reference implementation first; prove correctness before cleanup.
• Lock correctness with deterministic tests before refactoring.
• During the initial port, avoid premature performance work (e.g., mx.compile, kernel fusion tweaks, scheduler micro-optimizations); add optimizations only after correctness is locked.
• Refactor toward shared components and clean APIs once tests are green.
• PyTorch and MLX RNGs are different; for strict parity checks, export the exact initial noise/latents from the reference and load them in MLX instead of relying on matching integer seeds.

Workflow (checklist)

1. Scope and parity
   • Define target parity (outputs, speed, memory) and acceptable tolerances.
   • Identify reference files, configs, and checkpoints to mirror.
   • Draft a Cursor plan for the port and review it before starting implementation.
2. Port fast to reference
   • Add the model package skeleton and a variant class + initializer.
   • Follow standard mflux initializer/weight-loading style; review recent ports like z_image_turbo and flux2_klein for structure and naming.
   • Wire weight definitions/mappings early so loading is exercised (implement quantization in the initializer, but skip it during early runs).
   • Keep the first implementation simple and explicit; defer mx.compile and other speed-focused changes until deterministic parity is passing.
   • When defining explicit weight mappings, inspect actual tensor values from the model in the Hugging Face cache to confirm names and shapes.
   • Add a minimal hardcoded runner for quick iteration (two tiny scripts: one in the reference repo, one in mflux), seeded with diffusers-style defaults (e.g., 1024×1024, default prompt).
   • Add lightweight shape checks close to the code paths.
   • Use mx.save/mx.load at critical points; it is OK to add these to the reference (without changing logic) to export latents.
3. Port order (work backwards from image)
   • Typical image generation flow: prompt → text_encoder → transformer_loop → VAE → image.
   • For porting, invert the order so you can validate pixel space early.
   • Start with VAE decode/encode to validate output images quickly:
     • Export packed latents from the reference just before VAE decode.
     • Load latents inline and decode to an image for visual inspection.
     • Run an encode→decode roundtrip to sanity check reconstruction; a good-looking image reconstruction increases confidence in the implementation.
     • Expect small numeric diffs in tensor values; when it is not clear from the numbers alone, always generate images and rely on human visual inspection to judge whether the match is acceptable.
   • Then port the transformer loop and its schedulers with intermediate latent checks.
     • If the reference uses a novel scheduler, port it; otherwise, reuse the existing mflux scheduler.
   • Finish with the text encoder and tokenizer details.
   • After each major component is validated (e.g., VAE, transformer, text encoder), commit with a clear milestone message like "VAE done" to preserve progress.
   • Once the full port is working, remove any loaded tensors or debug artifacts so no traces remain.
4. Deterministic validation
   • Create a deterministic MLX test (image or tensor) that locks the output.
   • Run tests via MFLUX_PRESERVE_TEST_OUTPUT=1 uv run <test command>.
   • If MLX OOMs on sensible inputs (e.g., 1024×1024), assume a likely porting mistake and re-check shapes or memory-heavy ops.
5. Post-test refactor (explicit step)
   • Review commits after the first deterministic test to capture refactoring preferences.
   • Consolidate shared components into common modules.
   • Remove debug paths and one-off schedulers once validated.
   • Move configuration defaults into standard config/scheduler paths.
   • Simplify and decompose large files into focused modules once behavior is locked.
   • Prefer shared scheduler implementations when they already exist in mflux.
   • Ensure CLIs register callbacks via CallbackManager.register_callbacks(...) so shared features like --stepwise-image-output-dir work; pass a latent_creator that supports unpack_latents(...).
   • Keep running the deterministic image test during refactors to avoid regressions.
   • Align the variant class with recent ports (flux2_klein, z_image): prompt_cache, merged _predict, RoPE setup inside predict path, _decode_latents helper, no verbose comments/docstrings (see repo RULE.md).
   • Strip dead scaffolding (e.g. unused gradient-checkpointing flags) once training/inference paths are stable.
6. Pre-merge polish (after core port works)
   • diffusers sanity check: run matched mflux + diffusers generations; use mflux-debugging latent injection if outputs disagree but you need to validate transformer/VAE.
   • Golden tests: pick prompt/seed/settings that are stable on target CI hardware; update reference PNGs only after explicit approval (see mflux-testing).
   • img2img: verify latent packing/normalization on the img2img path matches txt2img and training (especially when reusing a shared VAE from another model family). pack_latents must accept the 5D (B, C, 1, H, W) tensor that tiled VAE encode (vae_encode_tiled) returns — squeeze the singleton temporal axis first, as flux2/fibo/z_image latent creators do; a 4D-only unpack (or a bare passthrough pack_latents) breaks tiled img2img. This is reachable via --low-ram: MemorySaver sets tiling_config = TilingConfig() (vae_encode_tiled=True), so always test img2img with --low-ram, not just the default path (which is safe only because VAEUtil.encode squeezes 5D→4D when tiling is off).
   • Cross-model touch points: list every file outside models/<your_model>/; justify shared changes (memory_saver tiling guard, shared VAE tiling_config, training runner wiring). Drop unrelated edits (e.g. personal .gitignore entries).
   • README: follow an existing model README structure (e.g. Flux2): hero image, turbo + base examples, feature section (img2img), disk-size warning, Notes, Training. Measure on-disk sizes with du on HF cache and/or mflux-save + du -sh for quantized sizes.
   • Training: example JSON under models/common/training/_example/, un-ignore in .gitignore, fast unit tests for training-adapter preview defaults.
   • Re-run just lint, just test-fast, then slow golden tests before merge.
7. Finalize
   • Re-run tests and basic perf checks after polish.
   • Add CLI/pipeline defaults and completions later, once core output is stable.
   • Ensure the model is wired into the standard surfaces:
     • ModelConfig entry + aliases
     • Thin model CLI entrypoint that uses shared parser/config/callback patterns
     • README following the structure and tone of existing model READMEs
     • Python API example that matches the CLI/defaults
   • Document any new mapping rules, shape constraints, or tolerances.

Package layout (reference: flux2)

Use src/mflux/models/flux2/ as the canonical tree. Do not invent flat mlx-vlm-style roots (config.py, scheduler.py, fp8.py, layout.py, monolithic model/transformer.py). Aliases and defaults live in ModelConfig; checkpoint validation belongs in the initializer and/or *WeightDefinition, not a separate layout module.

{model}/
  {model}_initializer.py
  __init__.py                    # export variant + initializer
  README.md
  cli/
    {model}_generate.py          # (+ edit/turbo CLIs when applicable)
  latent_creator/
    {model}_latent_creator.py
  model/
    {model}_text_encoder/        # prompt_encoder.py, tokenizer pieces, text_encoder.py
    {model}_transformer/         # attention, blocks, rope, transformer.py (split files)
    {model}_vae/                 # or reuse shared VAE (e.g. flux2_vae) — document in README
    {model}_scheduler/           # only when not covered by models/common/schedulers
  variants/
    __init__.py                  # re-export public variant class(es)
    txt2img/
      __init__.py
      {model}.py                 # e.g. flux2_klein.py, ideogram4.py
    edit/                        # when the model supports image-conditioned generation
      __init__.py
      {model}_edit.py
  weights/
    __init__.py
    {model}_weight_definition.py # components, download patterns, tokenizers
    {model}_weight_mapping.py    # WeightTarget list / key transforms for base weights
    {model}_lora_mapping.py      # LoRA key aliases (diffusers, PEFT, kohya) — when LoRA is supported
  training_adapter/              # when mflux-train is supported
    {model}_training_adapter.py

Variants: always place txt2img classes under variants/txt2img/ (even for single-mode models). Use variants/edit/ for edit/img2img variants. Import from the full path in save.py and CLIs, e.g. variants.txt2img.flux2_klein.

Weights: *WeightDefinition is required for every port. Add *WeightMapping when diffusers/HF key names need explicit targets. Add *LoRAMapping when inference or training supports LoRA — wire --lora-paths / --lora-scales through the shared parser and add fast tests that community LoRA filenames map to non-zero keys (see integration checklist below).

Variant class style (post-refactor): match flux2_klein / z_image — prompt_cache, _predict / _decode_latents, thin generate_image, prompt encoding in {model}_text_encoder/prompt_encoder.py.

Skip when not applicable: variants/edit/, training_adapter/, {model}_lora_mapping.py, local VAE package (if reusing another model family’s VAE). Document omitted features in the model README.

Integration surfaces checklist (don’t forget)

Past closed PRs show the same wiring gaps recurring on every new model. Use this as a tick list alongside the workflow above — not every row applies to every model (e.g. skip vision-encoder rows for txt2img-only), but scan it before opening a port PR.

Repo wiring (required for every new model family)

Weights, tokenizer, download scope

LoRA

Inference CLI & shared features

Training (if supported)

Tests & CI

When fixing a gap for model N, ask whether the same gap exists for other recent models and whether a shared fix belongs in models/common/ (preferred over copy-paste per model).

Tooling expectations

• Use uv for running scripts and tests: uv run <command>.
• Prefer uv run python -m <module> for local modules.

Deliverables

• Deterministic MLX test that verifies correctness.
• Documented weight mapping, shape constraints, and any known tolerances.
• Cleaned, shared components aligned with mflux style.
• Standard mflux surfaces in place: config aliases, thin CLI, and a README/examples pass aligned with the final behavior.