rapid-moss: MOSS-Image-SFT 4-step adapter (data-space DMD → DMAD)

English summary. A 4-step, CFG-free adapter for OpenMOSS-Team/MOSS-Image-SFT (50 steps × 2 forwards with CFG 3). It continues lodestones/MOSS-Image-SFT-8-steps-dmd (data-space DMD, arXiv 2610.12102) restricted to 4 iterations, trained with DMAD (arXiv 2610.02188): a LoRA critic on the frozen teacher separates teacher samples from student samples, and the student minimises the negative critic logit. DiT time per 1024² image on one H100: 16.2 s (teacher) → 0.81 s. Compared with the lodestones checkpoint run for 4 steps, images have more detail and livelier composition, with some regressions (see Results). Same file format and sampler as the lodestones checkpoint; inference.py runs it with the release pipeline. Training code is in code/.

MOSS-Image-SFT(50 ステップ・CFG 3 で 1 枚 100 回の DiT 計算)を、4 回の計算・CFG なしで画像を出せるようにする高速化アダプタです。 lodestones の 8 ステップ版(data-space DMD、update 2600)を出発点に、4 ステップに絞って DMAD で追加学習しました。

	DiT 時間 / 1024² 1 枚(H100)	DiT の計算回数
MOSS-Image-SFT(50 step、CFG 3、負側 block 22 skip)	16.2 s	100
lodestones 8 step	1.54 s	8
このアダプタ(4 step)	0.81 s	4

(DiT のみの計測。テキストエンコーダと VAE は含まない。実装は code/rapid_moss/model.py。)

使い方

pip install "diffusers>=0.38" "transformers>=5.3" accelerate safetensors huggingface_hub pillow
python inference.py --prompt "an anime girl in a yellow raincoat at a rainy bus stop at night" --width 1024 --height 1024 --out out.png

ファイルは lodestones と同じ形式です。

ファイル	中身
`lora.safetensors`	fp32 LoRA、rank 128 / alpha 128、全 42 ブロックの `proj.{q,k,v,o,ffn_up_gate,ffn_down}`。kohya 形式のキー(`diffusion_model.<module>.lora_down.weight` / `lora_up.weight` / `alpha`)
`time_embedding.safetensors`	fp32 の `time_embedding.*` と `shared_block_modulation.bias` / `video_head_modulation.bias` の置き換え
`dataspace.safetensors`	`latent_mix.Q` [128, 128] と `step_codes` [1000, 968, 926, 871]

サンプラー(ノイズの再注入もタイムステップのスケジュールもなく、同じ z を使い続ける):

x0 = torch.zeros_like(z)                                  # z: fp32 noise [1, 128, H/16, W/16]
for r in range(4):                                        # 1〜3 回で止めても画像になる(各段階を学習している)
    x = z + torch.einsum("oc,bchw->bohw", Q, x0)
    x0 = z - transformer(x, codes[r], prompt_embeds, prompt_mask).sample.float()
# x0 を release の VAE でデコード(BN の逆正規化 + 2x2 unpatchify)

LoRA の適用方法、CFG、解像度などの決まりは、下の「注意点」を読んでください。

ComfyUI

カスタムノード comfyui/rapid-moss/ を同梱しています。ComfyUI 標準の LoraLoaderModelOnly は LoRA を重みにマージするので、このアダプタには使えません。このノードを使ってください。

ComfyUI 標準ローダーにも対応しています。 comfyui/rapid-moss/example_workflows/rapid-moss-4step-native.json を使い、Load Diffusion Model で moss_image_sft_bf16.safetensors、Load CLIP で Qwen3.5-9B (type: moss_image)、Load VAE で FLUX.2 VAE を選んでください。QKV が結合されたモデルで出ていた ModuleDict ... has no attribute q を修正しました。この構成では OpenMOSS の HF ローダーノードは不要です。

従来の HF ローダーを使う場合:

MOSS-Image 公式ノード(OpenMOSS/MOSS-Image の comfyui/moss-image/)を入れます。
このリポジトリの comfyui/rapid-moss/ を ComfyUI/custom_nodes/rapid-moss/ にコピーします。
上の 3 ファイルを ComfyUI/models/rapid_moss/rapid-moss-4step/ に置きます。
comfyui/rapid-moss/example_workflows/rapid-moss-4step.json を画面にドラッグ&ドロップし、HF Loader の model_path を MOSS-Image-SFT のフォルダにして実行します。

ノード MOSS-Image Few-Step Adapter (rapid-moss) は、次の処理をして MODEL と、専用の SAMPLER・SIGMAS を返します。

LoRA をマージせず fp32 で足す
time path を fp32 に置き換える

これを BasicGuider(CFG なし)と SamplerCustomAdvanced につなぎます。lodestones の 8 ステップ版も、同じ 3 ファイルを別フォルダに置けば同じノードで使えます。

詳しい手順と注意点は comfyui/rapid-moss/README.md にあります。GPU での動作確認では、学習コードで作った画像とほぼ同じ絵になりました(下の図)。1024² 1 枚は 4.5〜6.9 秒(RTX PRO 6000、テキストエンコード・VAE 込み)です。

注意点

必ず守ること(守らないと効かない、または絵が崩れる)

元のモデルは MOSS-Image-SFT です。
- 学習の teacher も初期値も SFT です。MOSS-Image-Base と組み合わせたときの動作は確かめていません。
LoRA はマージせず、fp32 のまま足してください。
- 差分が bf16 の丸め幅より小さいため、bf16 の重みにマージすると効果がほぼ消えます。
- そのため、マージ方式のローダーでは効きません。ComfyUI の LoraLoaderModelOnly、diffusers の fuse_lora、マージ済みの bf16 重みを書き出すツールなどがこれに当たります。
- inference.py か、ComfyUI のカスタムノード(comfyui/rapid-moss/)を使ってください。
time_embedding.safetensors の 3 モジュールも置き換え、fp32 で計算します。
- LoRA だけを当てても正しく動きません。
- t=1 の時間ベクトルがバイアスに畳み込まれているので、3 つは必ずセットで使います。
専用のサンプラーを使います。
- x = z + Q·x0、x0 = z − v を 4 回繰り返し、timestep には step_codes を渡します。
- Euler などの標準サンプラーや、通常のスケジューラ(shift 5 など)では正しく動きません。
CFG とネガティブプロンプトは使いません。
- CFG の効果は蒸留で取り込み済みです。CFG を 1 より大きくすると計算が倍になり、効果が二重にかかって色や形が崩れる可能性が高いです(試していません)。
- 負側 block 22 skip や SLG も入れないでください。
ステップ数は 1〜4 です。
- 4 を推奨します。1〜3 回で止めても画像になります(途中の各段階も学習しています)が、その品質は確かめていません。5 回以上にはできません(step code が 4 つしかないため)。
base の重みは bf16 か fp32 で使ってください。
- fp8 などに量子化した重みでは、LoRA の差分が失われるか、time path の置き換えが正しくできません。ComfyUI のノードは fp8 だとエラーにしています。

分かっている弱点

厚塗り・絵画調の絵に、頼んでいない題字や漢字が入ることがあります。lodestones の 8 ステップ版から引き継いだ癖で、teacher は入れません。
暖簾・看板の小さな文字や、白黒漫画の吹き出しは、意味の通らない文字になりやすいです。
頼んでいない人物が増えることがあります。JSON レイアウトの 2 コマ漫画では、下のコマに 2 人目が出ました。
lodestones の 4 ステップより描き込みは増えますが、teacher(50 ステップ)の絵画的な質感や多様性には届きません。
学習を続けると、プロンプトへの従い方と文字が崩れていきました(800 イテレーション目)。公開しているのは崩れる前の 500 イテレーション目です。

検証していないこと

他の LoRA との組み合わせ: キャラクター LoRA や画風 LoRA を重ねたときの挙動は試していません。4 ステップ化で効きが弱まる可能性があります。
評価の範囲: 評価は 17 プロンプトの目視だけです。FID や好みの評価などの数値指標は取っていません。
学習データの規模: 学習プロンプトは 3,199 件で、teacher サンプルは 1 プロンプトにつき 1 枚です。
学習データの内容: 性的に露骨な内容は学習プロンプトから除いています。そうしたプロンプトでの挙動は確かめていません。

結果

左から teacher(50 step)/ lodestones 8 step 版を 4 step で使ったもの / このアダプタ(4 step)。同じ seed、評価用 17 プロンプト(eval/eval_prompts.jsonl、学習には未使用)。

lodestones の 4 step と比べた所見(目視のみ。数値指標では評価していません):

良くなった点
- 背景や小物の描き込みが増え、構図に動きが出た(teacher の雰囲気に近づいた)
- 縦書きポスターで「夏」が二重になる崩れが直った
- 付箋の 3 行、看板、吹き出しの日本語は正確なまま
悪くなった点
- ライトノベル表紙に頼んでいない英字が増えた
- 2 コマ漫画(JSON レイアウト)の下のコマに 2 人目が出た
- 暖簾などの小さな文字の崩れが増えた
変わらない点
- 厚塗りの侍の絵に頼んでいない題字が入る(lodestones から引き継いだ癖。teacher は入れない)
- 白黒漫画の吹き出しが意味の通らない文字になる

結論として、細部は良くなったが一長一短です。

学習の経過(EMA、4 step、同じ seed)

init(lodestones)	300 it	500 it(公開版)	800 it

300 it では、桜の女子高生の絵が顔の寄せ集めになる崩れが一度出ました。500 it では元に戻っています。
500 → 800 it ではプロンプトへの従い方(1girl, solo なのに 2 人)と文字が崩れていきました。841〜850 it で NaN になり、学習を止めました。teacher サンプルに非有限値はなく、原因は特定できていません。現在のコードは非有限の勾配で更新を飛ばします。
公開版は 500 it の EMA です。

手法

生成器(data-space iteration、arXiv 2610.12102)

lodestones の形式そのままです。x = z + Q·x0, v = net(x, code_r), x0 ← z − v。

学習対象は LoRA(rank 128)、time path(fp32)、Q の 3 つです。
lodestones update 2600 から初期化し、step codes は先頭 4 つを使います。
反復番号 r は 0〜3 を順番に回します。r 回目までは勾配なしで進め、最後の 1 回だけ勾配を流します。

識別器と損失(DMAD、arXiv 2610.02188)

critic の構成: 凍結した teacher(MOSS-Image-SFT の bf16)に、block 21〜41 の LoRA(rank 128)を付けます。block 41 の画像トークンにヘッドを付け、1 枚につき 1 つのロジットを出します。ヘッドは 4×4 stride-2 conv(3072→512)、GroupNorm、SiLU を 2 回重ね、空間平均を取ってから線形層に通します。
critic の損失: softplus(−d(teacher サンプル)) + softplus(d(生徒の出力))。ノイズ量は σ = 5u / (1 + 4u)、u ~ U[0.02, 0.98] で、teacher と生徒の組で同じ σ を使います。
生成器の損失: −d(ノイズを足した生徒の x0)。teacher の CFG 推論も fake score も、学習中は使いません。
teacher サンプル: 3,199 枚。1 プロンプトにつき 1 枚で、50 step・CFG 3・負側 block 22 skip で作りました。解像度は約 1MP で、正方形・縦長・横長を混ぜています。学習時のプロンプトはこの 3,199 件から選びます。
本物の画像を使うヘッドは使っていません(DMAD は 2 ヘッドですが、こちらは teacher ヘッドだけ)。

設定

項目	値
バッチ	16(8 GPU × 2、1 サンプルずつ順に処理)
生成器	AdamW lr 2e-5、betas (0, 0.99)、weight decay 0、EMA 0.998
critic	AdamW lr 4e-5、betas (0, 0.99)、weight decay 0.01、生成器を更新しない 50 it のウォームアップ
反復	500 it(生成器の更新 450 回)、1 it 約 5.7 秒
計算機	RunPod 8×H100 SXM。teacher サンプル生成約 1.9 時間、学習 500 it 約 47 分
費用	試運転と比較実験を含めて計約 $126

学習プロンプト

24,000 件(再配布はしていません)から、teacher サンプルを作った 3,199 件を学習に使いました。

種類	出典(ライセンス)
アニメの自然文キャプション	BootsofLagrangian/danbooru-multitier-captions-202606(CC-BY-4.0、Giyeong Oh)、CaptionEmporium/anime-caption-danbooru-2021-sfw-5m-hq(CC-BY-SA-4.0)
タグ列	上記の Danbooru タグ、FredZhang7/anime-prompts-180K(CreativeML OpenRAIL-M)。絵師名は除去
一般	ProGamerGov/synthetic-dataset-1m-dalle3-high-quality-captions(MIT)、data-is-better-together/open-image-preferences-v1(Apache-2.0)、nateraw/parti-prompts(Apache-2.0)
文字描画(日 60% / 英 25% / 中 15%)、JSON レイアウト、短文	テンプレートから生成

うまくいかなかったこと(DMDX)

最初は DMDX(arXiv 2507.18569)の ADM を data-space DMD に足す計画でした。ADM は、生成器の出力にノイズをかけた同じ x_τ から、teacher(CFG あり)と fake score の ODE を Δt = 1/64 だけ進め、その 2 つを識別器に見分けさせる方式です。

結果: 識別器は 130 回更新しても正解率 0.500〜0.506 のままで、学習しませんでした(logs/probe_adm_discriminator.jsonl)。2 つの入力の差は Δt·(v_real − v_fake) だけで、1 枚ずつの判定ではほぼ見分けられません。
比較: 同じヘッドで本物側を teacher サンプルにすると、60 回で 0.56 → 0.68 まで上がりました(logs/probe_teacher_discriminator.jsonl)。
この結果と、その後に出た DMAD を受けて、fake score も teacher の推論もいらない DMAD に切り替えました。

ライセンスとクレジット

重みとコード: Apache-2.0
元のモデル: OpenMOSS-Team/MOSS-Image-SFT(Apache-2.0)、初期値: lodestones/MOSS-Image-SFT-8-steps-dmd(Apache-2.0)
参考にした論文:
- Few-Step Generation via Data-Space Iteration(arXiv 2610.12102)
- DMAD: Distribution Matching as Adversarial Distillation(arXiv 2610.02188、コード)
- Adversarial Distribution Matching / DMDX(arXiv 2507.18569)
前作: Anima の DMDX・DMD2・PCM 蒸留の記録

Downloads last month: 62

Model tree for darask0/MOSS-Image_4-step_adapter_lora

Base model

OpenMOSS-Team/MOSS-Image-Base

Finetuned

OpenMOSS-Team/MOSS-Image-SFT

Adapter

(3)

this model

darask0
/

MOSS-Image_4-step_adapter_lora