# rapid-moss for ComfyUI

MOSS-Image-SFT を **4 ステップ・CFG なし**で動かすカスタムノードです(rapid-moss の DMAD 4 ステップアダプタ。lodestones の 8 ステップ DMD も同じノードで使えます)。
ComfyUI 標準の `SamplerCustomAdvanced` と組み合わせて使います。

## ComfyUI 標準ローダーを使う場合

MOSS-Image 対応の ComfyUI(2026-10-16 以降の master)と、この `rapid-moss` ノードだけで使えます。OpenMOSS の HF ローダーノードは不要です。

1. [Comfy-Org/MOSS-Image](https://huggingface.co/Comfy-Org/MOSS-Image) から `moss_image_sft_bf16.safetensors` を `models/diffusion_models/`、Qwen3.5-9B のテキストエンコーダを `models/text_encoders/`、`flux2-vae.safetensors` を `models/vae/` に置きます。
2. アダプタの3ファイルを `models/rapid_moss/rapid-moss-4step/` に置き、このフォルダを `custom_nodes/rapid-moss/` にコピーして再起動します。
3. `example_workflows/rapid-moss-4step-native.json` を開きます。`Load CLIP` の type は `moss_image` にします。API 版は `rapid-moss-4step-native-api.json` です。

**`aimdo memory compile error` が出る場合は、このノードの古い版が使われています。** 新しい ComfyUI は、DiT ブロックごとのメモリ確保を記録して再生する仕組み(comfy-aimdo のメモリコンパイラ)を使います。あるブロックで確保したテンソルを後のブロックで解放すると、この仕組みが止まります。古い版はブロックごとに新しい残差テンソルを返していたので、これに当たっていました(discussion #1)。この版は ComfyUI 標準のブロックと同じく、区間の前に確保した残差をインプレースで更新し、アダプタのテンソルも区間の前に GPU へ移します。

`AttributeError: 'ModuleDict' object has no attribute 'q'` も古い版で出るエラーです。ComfyUI 用のモデルは Q/K/V が `qkv` にまとめられていますが、この版はどちらの形式でも動きます。

画像生成モデルは **SFT の BF16** を使ってください。INT8/FP8/GGUF の生成モデルはこのアダプタの検証対象外です。テキストエンコーダの量子化は別で、ComfyUI が対応する Qwen3.5-9B を選べます。

## カスタムノードは必要か

ComfyUI 本体は MOSS-Image に対応しており、標準の `LoraLoaderModelOnly` も Q/K/V 結合ファイルに LoRA を正しく当てます。それでもこのアダプタには、本体にない次の 3 つが要ります。

- **専用サンプラー:** `x = z + Q·x0`、`x0 = z − v` を、学習した 4 つの step code(1000, 968, 926, 871)で繰り返します。KSampler の sampler・scheduler では表せません。
- **time path の置き換え:** `time_embedding.safetensors`(t=1 の時間ベクトルを変調バイアスに畳み込んだもの)。標準のローダーでは当てられません。
- **LoRA を重みに足し込まないこと:** 標準ローダーは LoRA を bf16 の重みに足します(dynamic VRAM でも同じ)。足し込むと、当たる差分の相対誤差は 81%(lodestones 版は 83%)になります。

標準ノードだけ(`LoraLoaderModelOnly` + `KSampler` euler / simple、CFG 1)でも、4〜8 ステップで破綻のない絵は出ます。lodestones の README が勧めているのはこの使い方です。ただし学習した生成器とは別物で、同じ seed で並べると、全体が紫がかり、線が柔らかくなります(下の図)。数値で比べると、学習コードの forward との 1 ステップ目の差は、このノードが 17%、標準ローダーが 57〜60% です。アダプタなしの base も 60% なので、標準ローダーでは学習した効果がほとんど残りません。このノードの 17% は、学習コード自身で SDPA のバックエンドや rope の精度を変えたときの揺れ(11〜23%)の範囲です。

![custom node vs stock nodes](https://huggingface.co/darask0/MOSS-Image_4-step_adapter_lora/resolve/main/images/comfyui-native-vs-node.jpg)

## HF ローダーを使う場合の必要なもの

| もの | 入手先 |
|---|---|
| ComfyUI | 新しめの版(`SamplerCustomAdvanced`、`BasicGuider`、`RandomNoise`、`EmptyFlux2LatentImage` があるもの。2026-10-16 の master で動作確認) |
| MOSS-Image 公式ノード | [OpenMOSS/MOSS-Image](https://github.com/OpenMOSS/MOSS-Image) の `comfyui/moss-image/` を `ComfyUI/custom_nodes/moss-image/` へコピー(ComfyUI Manager で「MOSS-Image」を検索してもよい) |
| MOSS-Image-SFT 本体 | [OpenMOSS-Team/MOSS-Image-SFT](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT) のフォルダ一式(`transformer/`、`text_encoder/`、`tokenizer/`、`vae/`) |
| このノード | このフォルダ(`comfyui/rapid-moss/`)を `ComfyUI/custom_nodes/rapid-moss/` へコピー |
| アダプタ | [darask0/MOSS-Image_4-step_adapter_lora](https://huggingface.co/darask0/MOSS-Image_4-step_adapter_lora) の 3 ファイルを `ComfyUI/models/rapid_moss/rapid-moss-4step/` へ |

```bash
cd ComfyUI
hf download OpenMOSS-Team/MOSS-Image-SFT --local-dir /path/to/MOSS-Image-SFT --exclude "assets/*"
hf download darask0/MOSS-Image_4-step_adapter_lora lora.safetensors time_embedding.safetensors dataspace.safetensors \
    --local-dir models/rapid_moss/rapid-moss-4step
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
# 任意: lodestones の 8 ステップ版も使う場合
hf download lodestones/MOSS-Image-SFT-8-steps-dmd lora.safetensors time_embedding.safetensors dataspace.safetensors \
    --local-dir models/rapid_moss/lodestones-8step
```

`models/rapid_moss/` の下の、3 ファイル(`lora.safetensors`、`time_embedding.safetensors`、`dataspace.safetensors`)を含むフォルダが、ノードの `adapter` の選択肢に出ます。

## 使い方

1. ComfyUI を再起動し、`example_workflows/rapid-moss-4step.json` を画面にドラッグ&ドロップします。
2. **MOSS-Image HF Loader (SFT)** の `model_path` に、MOSS-Image-SFT のフォルダを指定します。
3. **Prompt** を書き換えて **実行** します。

ノードのつなぎ方:

```
MOSS-Image HF Loader ─ MODEL ─▶ MOSS-Image Few-Step Adapter (rapid-moss) ─ model ──▶ BasicGuider ─▶ SamplerCustomAdvanced ─▶ VAE Decode ─▶ Save Image
        │                                                        ├ sampler ───────────────────────────▲ (sampler)
        ├ CLIP ─▶ CLIPTextEncode(プロンプト)─▶ BasicGuider        └ sigmas ────────────────────────────▲ (sigmas)
        └ VAE ─────────────────────────────────────────────────────────────▶ VAE Decode
RandomNoise ─▶ SamplerCustomAdvanced (noise)    EmptyFlux2LatentImage ─▶ SamplerCustomAdvanced (latent_image)
```

API で投げる場合は `example_workflows/rapid-moss-4step-api.json` を使います(`model_path` を書き換えてください)。

## ノードの設定

| 入力 | 説明 |
|---|---|
| `adapter` | `models/rapid_moss/` の下のフォルダ名 |
| `steps` | ネットワークの計算回数。rapid-moss は 1〜4(4 推奨)、lodestones は 1〜8 |
| `lora_strength` | 高速化 LoRA の強さ。1.0 推奨 |

出力の `model`・`sampler`・`sigmas` は必ずセットで `BasicGuider` と `SamplerCustomAdvanced` に渡してください。
`steps` をアダプタの step code の数より大きくはできません(4 ステップ版は 4 まで)。step code の間を補うスケジュールは学習していないためです。

## 注意

モデル全体の注意点(分かっている弱点、検証していないことなど)は、[モデルカードの「注意点」](https://huggingface.co/darask0/MOSS-Image_4-step_adapter_lora#注意点)にまとめています。ComfyUI で特に気をつけることは次のとおりです。

- **生成モデルは MOSS-Image-SFT にしてください。** アダプタは SFT から学習しています。Base との組み合わせは確かめていません。
- **このアダプタに `LoraLoaderModelOnly` は使えません。** 上の「カスタムノードは必要か」のとおり、LoRA の差分が bf16 の丸めで崩れます。必ずこのノードで適用してください。
- **CFG とネガティブプロンプトは使いません。** `BasicGuider` を使ってください。`CFGGuider` や `KSampler` で CFG を 1 より大きくすると、計算が倍になります。CFG の効果は蒸留で取り込み済みなので、二重にかかって崩れる可能性が高いです(試していません)。
- **このノードの MODEL は、ノードが出す SAMPLER と SIGMAS で回してください。** time path が専用サンプラーの step code 用に置き換わっているので、`KSampler` の sampler・scheduler(euler、sgm_uniform など)との組み合わせは想定していません。
- **`ModelSamplingSD3`(shift)と `SkipLayerGuidanceDiTSimple`(layer 22 skip)は入れないでください。** 通常の 50 ステップ生成用の設定です。なお MOSS-Image で shift を変えるときは `ModelSamplingSD3` ではなく `ModelSamplingAuraFlow` か `ModelSamplingFlux` を使います(`ModelSamplingSD3` だとノイズになります。ComfyUI issue #17119)。
- **解像度:** `EmptyFlux2LatentImage` で 16 の倍数、約 100 万画素がおすすめです。学習した解像度は 1024×1024、896×1152、832×1216、768×1344 と、それぞれの横向きです。
- **`weight_dtype`:** `default`(bf16)か `fp32` にしてください。fp8 には対応していません(エラーになります)。
- **仕組み:** 標準ローダーでは、MOSS-Image の forward 全体をアダプタ用に置き換えます(ComfyUI の `DIFFUSION_MODEL` ラッパー)。LoRA は重みに足し込まず、学習コードのサンプラーと同じ bf16 で別に計算して出力に足します。time path・変調・残差は fp32 で、timestep の sin/cos も bf16 に丸めません。追加の VRAM は約 0.65 GB(bf16 の LoRA 0.6 GB と fp32 の time path)です。HF ローダーでは従来どおり、LoRA を fp32 で足し、time path を fp32 に置き換えます(約 1.5 GB)。
- **seed:** `RandomNoise` の seed は、HF の `inference.py --seed` と同じ初期ノイズになります(どちらも CPU の `torch.randn`)。
- **他の LoRA との併用:** 画風 LoRA などは `LoraLoaderModelOnly` で重ねられます。このノードの前でも後でも構いません。ただし 4 ステップ化した状態での効き方は検証していません。

## 動作確認

### 2026-10-19(aimdo 対応版)

- **環境:** ComfyUI master `8f04fd6e`(2026-10-19)、comfy-aimdo 0.5.6(dynamic VRAM 有効)、torch 2.13.0+cu130、RTX PRO 6000。生成モデルは公式 SFT の重みを Q/K/V 結合の BF16 1 ファイルにしたもの(Comfy-Org の `moss_image_sft_bf16.safetensors` と同じ形式)、テキストエンコーダは BF16 の Qwen3.5-9B。
- **再現と修正:** 古い版は報告と同じ `aimdo memory compile error` で止まりました(`rogues: 26`)。この版は 1024²、896×1152、batch 2、lodestones 8 ステップのすべてで、`graph breaks: 0, rogues: 0` のまま最後まで動きました。
- **速度:** 1024² 4 ステップ 1 枚が約 1.1 秒(2 枚目以降、VAE デコード込み)。1 回の DiT forward は 222 ms で、アダプタなしの 183 ms、LoRA を fp32 で計算した場合の 306 ms と比べて速くなりました。
- **学習コードとの一致:** 同じノイズ・同じテキスト埋め込みで、学習コード(`code/rapid_moss/model.py`)の forward と比べました。1 ステップ目の相対誤差 17%、4 ステップ後 63% で、学習コード自身の揺れ(SDPA のバックエンドを変える: 11〜23% / 57〜61%、rope を fp32 にする: 14〜18% / 77%)と同じ範囲です。アダプタなしの ComfyUI 標準実装と公式実装の差は 1.5〜3.5% でした。
- **互換性:** 2026-10-17 の master(`5cc4b870`)でも、同じ入力に対して最新版とビット単位で同じ出力になりました。
- **回帰テスト:** 15 件が成功しました(学習コードと同じ式で書いた参照との一致を Q/K/V 分離・結合 × fp32/bf16 × CPU/CUDA で、残差を全ブロックでインプレース更新すること、ノードの再適用、HF ローダー用のパッチと解除、専用サンプラー、ステップ数上限)。`COMFYUI_ROOT=/path/to/ComfyUI RAPID_MOSS_TEST_CUDA=1 python -m pytest code/tests/test_comfy_adapter.py` で実行できます。

### 2026-10-17

- **標準ローダー対応の修正版:** ComfyUI `5cc4b870`、torch 2.13.0+cu130、RTX PRO 6000 で、標準ローダーと HF ローダーの両方から 1024×1024・4ステップ・seed 1234 の画像生成を確認しました。
- **初版(HF ローダー):** ComfyUI master(fc04da9、2026-10-16)、MOSS-Image 公式ノード(c1355e6)、torch 2.11 + CUDA 12.8、RTX PRO 6000。学習コードで作った正解画像(同じ seed 1234)とほぼ同じ絵になりました。lodestones 8 ステップも同じノードで動きました。1024² 1 枚あたり 4.5〜6.9 秒(テキストエンコード・VAE 込み、他のジョブと GPU を共用した状態)。

![ComfyUI check](https://huggingface.co/darask0/MOSS-Image_4-step_adapter_lora/resolve/main/images/comfyui-check.jpg)

## English

Custom node for 4-step, CFG-free MOSS-Image-SFT with the rapid-moss adapter (also runs lodestones' 8-step DMD).

**Native ComfyUI (recommended):** copy this folder to `ComfyUI/custom_nodes/rapid-moss/`, put `lora.safetensors`,
`time_embedding.safetensors`, `dataspace.safetensors` in `ComfyUI/models/rapid_moss/<name>/`, and load
`example_workflows/rapid-moss-4step-native.json` with the SFT BF16 diffusion model, `CLIPLoader` type `moss_image`, and the
FLUX.2 VAE. Use `BasicGuider` (no CFG / negative) and the node's SAMPLER and SIGMAS with `SamplerCustomAdvanced`; no
`ModelSamplingSD3` or SLG; `weight_dtype` default/bf16/fp32 (fp8 unsupported).

**`aimdo memory compile error` (discussion #1)** comes from older versions of this node. Current ComfyUI records each DiT
block's allocations and replays them (comfy-aimdo's memory compiler); a tensor allocated in one block and freed in a later one
aborts it, and the old node returned a new residual tensor from every block. The node now runs the whole MOSS-Image forward as
a `DIFFUSION_MODEL` wrapper that, like the native blocks, updates a residual allocated before the compiled region in place and
moves the adapter tensors to the GPU before that region. The LoRA stays unmerged and is computed in bf16 like the training
code's sampler; the time path, modulation and residual run in fp32 (the timestep sinusoid is not rounded to bf16). About
0.65 GB extra VRAM. Tested on ComfyUI `8f04fd6e` (2026-10-19, dynamic VRAM on) and `5cc4b870` (2026-10-17): no graph breaks
or rogues, ~1.1 s per 1024² 4-step image on an RTX PRO 6000.

**Is the node needed?** Native ComfyUI supports MOSS-Image and its LoRA loader handles the fused Q/K/V files, so
`LoraLoaderModelOnly` + `KSampler` (euler / simple, CFG 1, 4-8 steps) does give usable images, as lodestones' README
suggests. But it is not the trained generator: the stock loader adds the LoRA into bf16 weights (81-83% relative error on
the applied delta), keeps the base time path, and uses the flow sampler instead of the data-space one. Against the training
forward on the same inputs, the first-step velocity is 17% off with this node (inside the training code's own spread when
switching SDPA backend or rope precision, 11-23%) and 57-60% with the stock loader, about as far as the base model without
the adapter (60%). Side by side the stock-node images are tinted purple and softer (figure above).