akhaliq/MOSS-Image-IDM-2-Steps

MOSS-Image-IDM-LoRA

Community implementations of Intrinsic Distribution Matching (IDM) applied to MOSS-Image-SFT — culminating in a purpose-trained 2-step adapter (repo root), the first open model trained with IDM's loss combination.

Lineage: OpenMOSS-Team/MOSS-Image-SFT (50 steps, CFG 3) → lodestones 8-step (data-space DMD, 2600 updates) → darask0 4-step (DMAD, 500 iters) → akhaliq 2-step (IDM, 1000 iters, export 750 shipped).

🚀 The 2-step IDM adapter (repo root)

lora.safetensors + time_embedding.safetensors + dataspace.safetensors at the repo root = MOSS-Image-SFT in 2 network evaluations, no CFG, with legible text rendering in English and Chinese.

	DiT forwards / 1024² image
MOSS-Image-SFT (50 steps, CFG 3)	100
This adapter (2 steps, no CFG)	2

Showcase — 30 diverse prompts at 2 steps (comics, calligraphy, posters + anime, fantasy, editorial, landscapes; prompts in showcase/ and showcase2/):

Method (IDM, community implementation of arXiv:2610.39633): fake-score network trained with flow matching on generated samples + DMD direction sg[v_real − v_fake] + L_pop EMA feature moment matching (λ=0.5, mean/std of last-block image-token features vs teacher reference, FIFO queues) — no adversarial losses, no discriminator. To our knowledge this is the first open model trained with IDM's loss combination (the official repo is inference-only) and the first purpose-trained 2-step MOSS-Image model. Deviations from the paper: the base is the data-space fixed-point generator (arXiv:2610.12102) rather than RCGM; LoRA r128 rather than full-parameter; 249 prompts rather than 5.6M images.

Training: 1000 iterations (~2 h on 1×H200), warm-started from darask0's 4-step adapter restricted to 2 step codes, 249 dense-text prompts (moss-textboost-data), EMA 0.998. Shipped checkpoint = EMA export at iteration 750, selected by side-by-side eval (below). Trainer: training/rapid_moss_train_idm.py (extends darask0's rapid_moss, Apache-2.0). All exports, sample sheets every 100 iterations, and logs: runs/idm/.

Usage (Python)

pip install "diffusers>=0.38" "transformers>=5.3" accelerate safetensors huggingface_hub pillow
python inference.py --prompt 'A cafe chalkboard menu titled "MORNING BREW" with lines: "Espresso $3", "Latte $5"' --steps 2 --out out.png

inference.py loads the adapter from this repo by default (data-space sampler, fp32 LoRA hooks — same format as darask0's adapter). Keep the LoRA and time path in fp32; the base stays bf16.

Usage (ComfyUI)

This adapter is the same three-file format as darask0's 4-step adapter, so the rapid-moss custom node is the recommended way to run it — stock nodes can't express the data-space sampler (see the caveat below).

1. ComfyUI with MOSS-Image support (master from 2026-10-16 or later).
2. Base models from Comfy-Org/MOSS-Image: moss_image_sft_bf16.safetensors → models/diffusion_models/, the Qwen3.5-9B text encoder → models/text_encoders/, flux2-vae.safetensors → models/vae/. Use the bf16 SFT checkpoint — INT8/FP8/GGUF bases are untested with this adapter.
3. This adapter (3 files from this repo's root) → models/rapid_moss/akhaliq-idm-2step/:

hf download akhaliq/MOSS-Image-IDM-LoRA lora.safetensors time_embedding.safetensors dataspace.safetensors \
    --local-dir models/rapid_moss/akhaliq-idm-2step

4. The custom node (from darask0's repo, Apache-2.0 — the adapter format is identical) → custom_nodes/rapid-moss/, then restart ComfyUI:

hf download darask0/MOSS-Image_4-step_adapter_lora --include "comfyui/rapid-moss/*" --local-dir /tmp/rapid-moss
cp -r /tmp/rapid-moss/comfyui/rapid-moss custom_nodes/rapid-moss

5. Open the example workflow custom_nodes/rapid-moss/example_workflows/rapid-moss-4step-native.json and point the adapter widget at akhaliq-idm-2step. The node reads the step codes from dataspace.safetensors, so it runs 2 iterations automatically (our dataspace carries 2 codes where darask0's carries 4). In the Load CLIP node, set the type to moss_image.

Wiring (same as the 4-step example): MOSS-Image HF Loader → MOSS-Image Few-Step Adapter (rapid-moss) → BasicGuider → SamplerCustomAdvanced → VAE Decode, with the node's sampler and sigmas outputs feeding the sampler.

Why not just LoraLoaderModelOnly? Stock ComfyUI applies the LoRA by merging it into the bf16 weights — the adapter's delta is smaller than bf16 rounding (~81% relative error after merging, per darask0's measurements), and the stock path can't apply time_embedding.safetensors or the data-space sampler at all. You still get coherent images (KSampler, euler/simple, CFG 1, 4–8 steps), but it's a different, weaker generator than the one we trained. The custom node applies the LoRA unmerged in fp32, replaces the time path, and runs the trained sampler.

Image editing (experimental, with MOSS-Image-Edit-HQ)

The 2-step adapter can be stacked with the armV5 edit LoRA (merged into the base weights; the IDM LoRA applied unmerged in fp32, edit conditioning via reference tokens + ai-toolkit prompt embeds). Honest verdict after a CFG sweep — it works for attribute edits, not for structural ones:

edit	2 steps (no CFG)	2 steps CFG 2	2 steps CFG 3	armV5 @ 50 steps
Expression edit	identity drift	✅ works, clean	works, over-sharpened	✅ best
Hairstyle (color)	faint tint	✅ visible shift (no shape change)	oversaturated	✅ best
Sketch → photo	❌ passive	❌ passive	❌ passive	✅ works
Colorize B&W	❌ passive	❌ passive	❌ passive	✅ works

• Use 2 steps + CFG 2 for quick attribute drafts (expression, color tint). Reference implementation: training/edit_2step_eval.py (ai-toolkit conditioning + data-space sampler with a CFG branch).
• For real edits, use armV5 @ 50 steps CFG 3 — or the documented 8-step stack (armV5 + lodestones DMD in chained ComfyUI loaders) for fast drafts.
• Why it's limited: this adapter was distilled on T2I prompts only, CFG-free; the edit LoRA's instruction signal needs guidance (CFG ≈ 2) and more refinement passes for structural changes. The proper fix is edit-conditioned distillation (continuing IDM training with reference tokens + edit pairs in the loop) — planned as v2; the shared training assets in shared/ are the starting point.

Samples (2 steps, no CFG, seed 42)

prompt	2-step IDM (export 750)
Chalkboard menu — "MORNING BREW", three priced lines
Neon bar sign — "OPEN LATE", "Craft Cocktails & Vinyl"
慢生活咖啡 — wooden café sign, Chinese
Recipe card — "GRANDMA'S RAMEN" with ingredient list

All 8 held-out prompts × all checkpoints: eval/2step_idm/.

Eval — honest read

8 held-out dense-text prompts, seed 42. Columns: base @ 50 steps · darask0 4-step @ 2 iterations · our IDM exports at it 250/500/750/1000:

• Exports 250–750 are all competitive with darask0's adapter at 2 iterations — the purpose-trained 2-step works: legible text, correct small print in most rows, EN and ZH.
• On 2 of 8 prompts (recipe card, wooden café sign) the 2-step adapters beat the 50-step teacher.
• Export 750 is the strongest overall and is what's at the repo root.
• it-1000 degraded (people/dogs intruding into text scenes): the L_pop term destabilized after ~it 850 (visible in runs/idm/log.jsonl), so late exports were excluded. Improvement over the init (darask0 @ 2 iterations) is real but incremental — the value here is the purpose-trained 2-step path and the IDM training implementation, not a quality leap.
• The gain vs our earlier one-shot 50→2 attempt (v3, below) is night and day — the 8→4→2 curriculum is what makes 2-step tractable at this scale.

v3: OPD + EMA-DMD 2-step distillation — it trains, but one-shot 50→2 is not enough

Status (2026-10-31): honest null-ish result, published with all artifacts. v3 adds the DMD score loss on top of v2's OPD trajectory loss — 2 of the 3 IDM losses, the first public IDM-style training implementation — and the run is fully reproducible. The quality gain over v2 is real but small, and 2-step text is still illegible. Details:

Method	OPD trajectory loss + DMD score loss (EMA copy of the student as fake-score net, λ=1.0), no adversarial
Trainer	training/train_v3_idm_lora.py — extends DiffSynth-Studio's OPD trainer @ 543bf84; all fixes documented in its commit history
Init	warm-start from opd_8step/step-249 (v2)
Training	1992 steps, LoRA r64, lr 1e-4, 249 prompts, 1024px, ~3.9 h on 1×H200
Teacher / student	MOSS-Image-SFT @ 25 steps CFG 3 → 2 steps, sigmas [1.0, 0.024], CFG-free
Checkpoints	opd_2step_v3/step-250.safetensors … step-1992.safetensors

What the evals showed (grids in eval/v3/):

• Base @ 50 steps: crisp text on all 8 held-out prompts. v2 and v3 @ 2 steps: composition and color, but text still illegible.
• v3 vs v2: modestly better contrast/color; not a step-change.
• Checkpoint trajectory (step-500 vs 1000 vs 1992): nearly flat — the last 1000 steps bought little. A loss spike at steps 501–750 healed on its own and did no lasting damage.
• Diagnosis: one-shot 50→2-step distillation saturates quickly at LoRA/batch-1 scale. IDM's headline results used 15K steps, batch 64, full-parameter on 32×H100; the community 8-step MOSS distill needed ~2600 updates for a much easier target.

The decisive probe (eval/rapid_probe/): darask0's 4-step adapter run at 2 iterations already produces legible text — the data-space DMD lineage (lodestones 8-step → darask0 4-step) is simply further along than the flow-matching trajectory approach in this repo. That probe is what moved the 2-step work to the lineage above.

v2: 2-step OPD distillation — the paper's actual regime

The paper's headline is few-step text rendering (2 NFE). v2 attacks it directly: on-policy distillation (OPD) to 2 steps with our dense-text dataset, so the few-step generator itself learns text. Teacher: MOSS-Image-SFT @ 50 steps CFG 3. Student: 2 steps CFG 1, LoRA r=64. Checkpoints in opd_8step/ (folder name is a leftover — it holds the 2-step run).

Status (2026-10-30): 249 steps in, the 2-step student has learned composition/color/layout but text is still faint — see grid below (left: base @ 2 steps = smudge; mid: OPD @ 2 steps; right: 50-step reference). For reference, the community 8-step distillation used ~2600 updates. Superseded by v3 (above), which continues from step-249. Use opd_8step/step-249.safetensors at num_inference_steps=2, cfg_scale=1.0.

---

v1: SFT LoRA on the 50-step base — what it does — and an honest eval

Trained to improve legible text rendering (menus, posters, signs, labels, sticky notes — English and Chinese) on MOSS-Image-SFT.

v1 result: marginal — read this before using. Side-by-side evals (left: base, right: +LoRA, same seed, 12 held-out prompts):

At 50 steps the base model is already excellent at text (structured-prompt text rendering is MOSS-Image's core training spec), so the delta from this 100-image LoRA is small. At 8 steps (below), both base and LoRA garble text equally — few-step failure is a distillation problem, not something SFT data fixes. The IDM paper's big data-driven jump (78→93 LongText) started from a base that was weak at text; that gap doesn't exist here.

We're publishing it anyway: the checkpoints, data, and evals are all open, and a documented null result on a hyped paper's recipe is useful.

Training (v1)

Base	OpenMOSS-Team/MOSS-Image-SFT (Apache-2.0)
Method	LoRA r=64, all linear layers, plain flow-matching SFT
Data	100 dense-text images self-generated by MOSS-Image-SFT (30 steps, CFG 1.0) — menus/posters/signs/receipts/whiteboards, ~60% EN / 40% ZH: akhaliq/moss-textboost-data
Steps	500 (checkpoints: step-250.safetensors, step-500.safetensors)
LR	1e-4, AdamW, grad accum 4, bf16, 1024×1024
Trainer	DiffSynth-Studio stock MOSS-Image LoRA trainer
Cost	~15 min on 1×A100, trained 2026-10-30

Usage (DiffSynth) — v1 SFT LoRA

from diffsynth.pipelines.moss_image import MOSSImagePipeline, ModelConfig
import torch

pipe = MOSSImagePipeline.from_pretrained(
    torch_dtype=torch.bfloat16, device="cuda",
    model_configs=[
        ModelConfig(model_id="OpenMOSS-Team/MOSS-Image-SFT", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"),
        ModelConfig(model_id="OpenMOSS-Team/MOSS-Image-SFT", origin_file_pattern="text_encoder/model*.safetensors"),
        ModelConfig(model_id="OpenMOSS-Team/MOSS-Image-SFT", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", computation_dtype=torch.float32),
    ],
    tokenizer_config=ModelConfig(model_id="OpenMOSS-Team/MOSS-Image-SFT", origin_file_pattern="tokenizer/"),
)
pipe.load_lora(pipe.dit, ModelConfig(model_id="akhaliq/MOSS-Image-IDM-LoRA", origin_file_pattern="step-500.safetensors"), alpha=1)

image = pipe('A cafe chalkboard menu with the heading "MORNING COFFEE" and three lines: "Americano $4", "Latte $5", "Mocha $6".',
             height=1024, width=1024, seed=42, num_inference_steps=50, cfg_scale=3.0)

Limitations

• 2-step adapter: trained on 249 text-rendering prompts — the showcases above suggest good generalization, but general-purpose 2-step quality is not exhaustively evaluated. Improvement over darask0's adapter at 2 iterations is incremental. it-1000 exports degraded (L_pop instability late in training) and are kept in runs/idm/ only for analysis.
• Editing: attribute edits only at 2 steps + CFG 2 (see the editing section); structural edits need the 50-step path.
• v1 experiment: only 100 training images and 500 steps, and the eval shows no meaningful visible delta over the base at 50 steps.
• v2/v3 (2-step DiffSynth line): readable composition, illegible text — see v3 above.
• No LongText-Bench score; evaluation is visual (see eval/ and showcase/).

Citation

@misc{liu2026intrinsicdistributionmatching,
  title={Intrinsic Distribution Matching},
  author={Deyuan Liu and Yansen Han and Kaiwen Zheng and Jun Xie and Jiacheng Liu and Zhenglin Cheng and Peng Sun and Jintao Zhang and Jianguo Li and Tao Lin},
  year={2026}, eprint={2610.39633}, archivePrefix={arXiv}, primaryClass={cs.CV},
  url={https://arxiv.org/abs/2610.39633},
}

Lineage papers: data-space DMD (arXiv:2610.12102, lodestones 8-step) and DMAD (arXiv:2610.02188, darask0 4-step). MOSS-Image tech report: arXiv:2610.23374.