MOSS-Image-IDM-LoRA

A community implementation of the recipe behind Intrinsic Distribution Matching (IDM), applied to MOSS-Image-SFT.

This is not the official IDM model (IDM weights and training code are not released). v1 replicated the data recipe; v2/v3 are the first public implementations of IDM-style training (trajectory distillation + DMD score loss; v3 uses an EMA fake-score net), and the 2-step continuation adds L_pop moment matching — the full IDM loss stack. See v3 for the honest outcome of the DiffSynth line and where the work continues.

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.

Where it continued: we pivoted the 2-step goal to that lineage and trained a 4→2 continuation with the full IDM loss stack (DMD direction + L_pop EMA feature moment matching, no adversarial losses) — in this repo: trainer at training/rapid_moss_train_idm.py, checkpoints under runs/idm/, shared training assets under shared/. This is the first open model trained with IDM's loss combination (arXiv 2610.39633, community implementation).

This repo keeps: the v1/v2/v3 checkpoints, both trainers (first public IDM-style training code), all eval grids, and the dataset — as a documented, reproducible account of what worked and what didn't.

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

• 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): readable composition, illegible text — see v3 above for the full account and the 2-step continuation.
• No LongText-Bench score; evaluation is visual (see eval/).

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},
}