---
license: apache-2.0
base_model:
- OpenMOSS-Team/MOSS-Image-SFT
- lodestones/MOSS-Image-SFT-8-steps-dmd
library_name: diffusers
pipeline_tag: text-to-image
language:
- ja
- en
tags:
- moss-image
- lora
- distillation
- few-step
- dmad
- dmd
---

# 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](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT)
(50 steps × 2 forwards with CFG 3). It continues [lodestones/MOSS-Image-SFT-8-steps-dmd](https://huggingface.co/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`。)

## 使い方

```bash
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 を使い続ける):

```python
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 本体も MOSS-Image に対応していますが、専用サンプラーと time path の置き換えはなく、標準の `LoraLoaderModelOnly` は LoRA を bf16 の重みに足し込みます(当たる差分の相対誤差 81%)。標準ノードだけでも絵は出ますが、学習した生成器とは別物になります(下の比較)。このノードを使ってください。

**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 を選んでください。この構成では OpenMOSS の HF ローダーノードは不要です。新しい ComfyUI で出ていた `aimdo memory compile error`(2026-10-19 修正)と、QKV が結合されたモデルで出ていた `ModuleDict ... has no attribute q` は、ノードを更新すると直ります。

![ComfyUI: custom node vs stock nodes](images/comfyui-native-vs-node.jpg)

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

1. MOSS-Image 公式ノード([OpenMOSS/MOSS-Image](https://github.com/OpenMOSS/MOSS-Image) の `comfyui/moss-image/`)を入れます。
2. このリポジトリの `comfyui/rapid-moss/` を `ComfyUI/custom_nodes/rapid-moss/` にコピーします。
3. 上の 3 ファイルを `ComfyUI/models/rapid_moss/rapid-moss-4step/` に置きます。
4. `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 を重みにマージせず、差分として足す(標準ローダーでは学習コードのサンプラーと同じ bf16、HF ローダーでは fp32)
- time path を fp32 に置き換える(標準ローダーでは変調と残差も fp32)

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

詳しい手順と注意点は [comfyui/rapid-moss/README.md](comfyui/rapid-moss/README.md) にあります。HF ローダーでの動作確認では、学習コードで作った画像とほぼ同じ絵になりました(下の図)。標準ローダーでは、2026-10-19 の ComfyUI で 1024² 1 枚約 1.1 秒(2 枚目以降、RTX PRO 6000、VAE 込み)です。

![ComfyUI check](images/comfyui-check.jpg)

## 注意点

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

- **元のモデルは MOSS-Image-SFT です。**
  - 学習の teacher も初期値も SFT です。MOSS-Image-Base と組み合わせたときの動作は確かめていません。
- **LoRA はマージせず、差分として足してください。**
  - 差分が bf16 の丸め幅より小さいため、bf16 の重みにマージすると効果がほぼ消えます(当たる差分の相対誤差 81%)。差分そのものは bf16 で計算して構いません(学習コードのサンプラーと ComfyUI ノードは bf16、`inference.py` は fp32)。
  - そのため、マージ方式のローダーでは効きません。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 など)は、学習した経路ではありません。ComfyUI の標準ノード(LoRA ローダー + KSampler euler)でも絵は出ますが、全体が紫がかり、線が柔らかくなります。
- **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 だとエラーにしています。

### おすすめの設定

- **解像度:** 16 の倍数で、約 100 万画素が目安です。
  - 学習した解像度は 1024×1024、896×1152、832×1216、768×1344 と、それぞれの横向きです。これより大きい・小さい解像度での品質は確かめていません。
- **プロンプト:** MOSS-Image と同じです(自然文、タグ列、JSON レイアウト)。
  - 描かせたい文字は `"..."` で囲みます。
  - 日本語の文字描画は、付箋・看板・ポスター・吹き出しで確認しました。
- **seed:** 初期ノイズは CPU の `torch.randn`(fp32)です。`inference.py` と ComfyUI のノードは、同じ seed なら同じノイズになります。
  - ただしテキストエンコーダの実装(transformers と ComfyUI)の違いで、細部は少し変わります。
- **VRAM:** `inference.py` と ComfyUI の HF ローダーでは、base に加えて約 1.5 GB 使います。fp32 の LoRA が 1.19 GB、fp32 の time path と modulation が約 0.34 GB です。ComfyUI の標準ローダーでは約 0.65 GB(bf16 の LoRA と fp32 の time path)です。

### 分かっている弱点

- 厚塗り・絵画調の絵に、頼んでいない題字や漢字が入ることがあります。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`、学習には未使用)。

![compare-1](images/compare-1.jpg)
![compare-2](images/compare-2.jpg)
![compare-3](images/compare-3.jpg)

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

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

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

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

| init(lodestones) | 300 it | 500 it(公開版) | 800 it |
|---|---|---|---|
| ![](images/training/init.jpg) | ![](images/training/000300.jpg) | ![](images/training/000500.jpg) | ![](images/training/000800.jpg) |

- 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](https://huggingface.co/datasets/BootsofLagrangian/danbooru-multitier-captions-202606)(CC-BY-4.0、Giyeong Oh)、[CaptionEmporium/anime-caption-danbooru-2021-sfw-5m-hq](https://huggingface.co/datasets/CaptionEmporium/anime-caption-danbooru-2021-sfw-5m-hq)(CC-BY-SA-4.0) |
| タグ列 | 上記の Danbooru タグ、[FredZhang7/anime-prompts-180K](https://huggingface.co/datasets/FredZhang7/anime-prompts-180K)(CreativeML OpenRAIL-M)。絵師名は除去 |
| 一般 | [ProGamerGov/synthetic-dataset-1m-dalle3-high-quality-captions](https://huggingface.co/datasets/ProGamerGov/synthetic-dataset-1m-dalle3-high-quality-captions)(MIT)、[data-is-better-together/open-image-preferences-v1](https://huggingface.co/datasets/data-is-better-together/open-image-preferences-v1)(Apache-2.0)、[nateraw/parti-prompts](https://huggingface.co/datasets/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](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT)(Apache-2.0)、初期値: [lodestones/MOSS-Image-SFT-8-steps-dmd](https://huggingface.co/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、[コード](https://github.com/Yzmblog/DMAD))
  - Adversarial Distribution Matching / DMDX(arXiv 2507.18569)
- 前作: [Anima の DMDX・DMD2・PCM 蒸留の記録](https://note.com/darask/n/neef98f000087)