MOSS-Image-IDM-2-Steps

MOSS-Image-SFT in 2 network evaluations, no CFG — a single LoRA adapter. Community implementation of Intrinsic Distribution Matching (IDM): the first open model trained with IDM's loss combination, and the first purpose-trained 2-step adapter for MOSS-Image.

Lineage: OpenMOSS-Team/MOSS-Image-SFT (50 steps, CFG 3) → lodestones 8-step (data-space DMD, 2600 updates) → darask0 4-step (DMAD, 500 iters) → v1: IDM 2-step T2I (1000 iters) → v2: edit continuation (1500 iters, armV5 folded).

⚠️ Choose your checkpoint (read this first)

Use case	Checkpoint	Path
Text-to-image	v1 — IDM export 750 (recommended)	runs/idm/export/000750/
Editing (attribute edits, local)	v2 — folded adapter	repo root or adapter_v2/

Why: the v2 edit continuation (75% edit-prompt batches) visibly degraded general-purpose T2I versus v1 — an honest regression we confirmed after the grids were compared. Use v1 for T2I. The demo Space (akhaliq/moss-image-2steps) serves T2I on v1; editing runs locally via inference_edit.py with the folded adapter.

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

Showcase — v1 at 2 steps, no CFG

30-prompt set (showcase2/) and the original 16-prompt set (comics, calligraphy, photos, posters, anime, game UI; showcase/):

Samples (2 steps, seed 42, v1)

prompt	2-step IDM (v1)
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

T2I eval (8 held-out dense-text prompts vs 50-step teacher and darask0's 4-step adapter): on 2 of 8 prompts the 2-step adapters beat the 50-step teacher; exports 250–750 are competitive throughout. Full grid: eval/2step_idm/grid_2step_idm.jpg.

Edit samples (2 steps, experimental — v2 folded adapter)

Attribute-level edits work at 2 steps CFG-free, identity preserved:

instruction	input	2-step edit
"change her hair to long wavy auburn hair"
"make both of them laugh with big open-mouth smiles"

Editing at 2 steps — the honest verdict

We ran the full experiment matrix on instruction editing: attribute-level edits (expression, color changes) work at 2 steps; structural transformations (sketch→photo, colorization, big style transfers) do not — at 2 or even 4 steps, for any few-step MOSS adapter that exists.

The decisive grid — input | our regression-warmed 2-step model | darask0's 4-step adapter + armV5 | armV5 @ 50 steps:

Edit	reg warmup · 2 steps	darask0+armV5 · 4 steps	base · 50
Sketch→photo	❌ sketch + shadow	❌ sketch + hatching	✅ photo
Hairstyle (color)	✅	✅	✅ (also length)
Expression	✅ identity preserved	✅	✅
Colorize	❌	❌	✅

Why this is the real answer and not a recipe gap: five independent attempts hit the same wall — naive stacking, DMD distillation (3000 triplets), DMD + ground-truth anchoring, MT-Forcing-style regression warmup, and darask0's own DMAD-trained 4-step adapter. Structural transformation needs the teacher's trajectory; few-step samplers take a shortcut (copy the input, shade it, or frame it). A 12-type held-out eval scored ~1/12 instruction-following at 2 steps (eval/edit_wide/). If a paper solves few-step structural editing at this scale, we'll be glad to be wrong — the full training code and evals are here to beat.

For production editing today: armV5 @ 50 steps (quality) or the armV5 + lodestones 8-step stack (faster drafts). For attribute edits at 2 steps: the v2 folded adapter — see inference_edit.py (requires diffusers<1.0 and transformers ≤5.20.x for ai-toolkit compatibility).

Usage (Python)

pip install "diffusers>=0.38" "transformers>=5.3" accelerate safetensors huggingface_hub pillow
python inference.py --adapter-dir runs/idm/export/000750 \
    --prompt 'A steaming bowl of ramen on a wooden counter, soft window light' --steps 2 --out out.png

inference.py defaults to the repo-root adapter; pass --adapter-dir runs/idm/export/000750 for the recommended v1 T2I checkpoint. Keep the LoRA and time path in fp32; the base stays bf16.

Usage (ComfyUI)

The adapter is the same three-file format as darask0's 4-step adapter, so the rapid-moss custom node is the recommended path:

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 — quantized bases are untested.
3. The v1 T2I adapter → models/rapid_moss/akhaliq-idm-2step/:

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

4. The custom node (from darask0's repo, Apache-2.0) → custom_nodes/rapid-moss/, restart:

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 example_workflows/rapid-moss-4step-native.json, point the adapter widget at akhaliq-idm-2step. The node reads our 2 step codes from dataspace.safetensors automatically. Set Load CLIP type to moss_image.

Why not LoraLoaderModelOnly? Stock loaders merge the LoRA into bf16 weights — the delta is smaller than bf16 rounding (~81% relative error after merging), and they can't apply the fp32 time path or the data-space sampler. You get a different, weaker generator. The custom node applies everything unmerged in fp32.

Method

IDM stack (community implementation of arXiv:2610.39633): fake-score network trained with flow matching, DMD direction sg[v_real − v_fake], L_pop EMA feature moment matching (λ=0.5) — no adversarial losses, no discriminator. v2 adds: reference-image conditioning in the training loop (Kontext-style ref tokens), armV5 active in teacher and student (teacher = base+armV5 @ CFG 3), 3000 ConceptEdit triplets + 25% T2I, micro-batch 4 (paper-faithful moment-matching statistics). Editing experiments (all open in training/ + eval/): naive stacking, DMD + GT step-1 anchor, MT-Forcing-style regression warmup (runs/reg1/), darask0 4-step stacking. Deviations from the paper: data-space fixed-point generator (arXiv:2610.12102) rather than RCGM; LoRA r128 not full-parameter; 3.2K prompts not 5.6M images.

Limitations

• v2 edit continuation regressed general T2I — use v1 (runs/idm/export/000750) for text-to-image (see checkpoint guide above).
• Structural editing doesn't work at 2–4 steps (see verdict above) — attribute edits only.
• v1's it-1000 exports degraded (L_pop instability) — kept in runs/idm/ for analysis only; v1 ships export 750.
• No LongText-Bench score; evaluation is visual (eval/, 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. Edit LoRA: akhaliq/MOSS-Image-Edit-HQ (armV5). Edit data: inclusionAI/ConceptEdit-12M (Apache-2.0).