Fix rapid-moss compatibility with native ComfyUI MOSS models

Support fused QKV projections and preserve fp32 modulation/residuals in native blocks. Add native UI/API workflows and regression coverage. Validated 1024x1024, 4-step generation with native and HF loaders; 8 tests pass.

README.md

@@ -71,6 +71,10 @@ LoRA の適用方法、CFG、解像度などの決まりは、下の「[注意
 
 カスタムノード `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 ローダーを使う場合:
+
 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/` に置きます。

code/tests/test_comfy_adapter.py

@@ -0,0 +1,132 @@
+"""COMFYUI_ROOT=/path/to/ComfyUI python -m pytest code/tests/test_comfy_adapter.py"""
+import importlib.util
+import os
+from pathlib import Path
+import sys
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+sys.path.insert(0, os.environ['COMFYUI_ROOT'])
+from comfy.cli_args import args
+args.cpu = True
+import comfy.ops
+from comfy.ldm.moss_image.model import MOSSImage
+from comfy.model_patcher import ModelPatcher
+
+spec = importlib.util.spec_from_file_location('rapid_nodes', Path(__file__).parents[2] / 'comfyui/rapid-moss/nodes.py')
+rapid = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(rapid)
+
+
+@pytest.fixture
+def adapter(tmp_path, monkeypatch):
+    torch.manual_seed(123)
+    dim, ffn, rank = 128, 192, 3
+    lora = {}
+    for name in rapid.PROJ:
+        inp, out = (ffn, dim) if name == 'ffn_down' else (dim, 2 * ffn if name == 'ffn_up_gate' else dim)
+        prefix = 'diffusion_model.layers.0.proj.' + name
+        lora[prefix + '.lora_down.weight'] = torch.randn(rank, inp) * 0.03
+        lora[prefix + '.lora_up.weight'] = torch.randn(out, rank) * 0.03
+        lora[prefix + '.alpha'] = torch.tensor(float(rank))
+    time = {
+        'time_embedding.0.weight': torch.randn(dim, 256) * .01,
+        'time_embedding.0.bias': torch.randn(dim) * .01,
+        'time_embedding.2.weight': torch.randn(dim, dim) * .01,
+        'time_embedding.2.bias': torch.randn(dim) * .01,
+        'shared_block_modulation.bias': torch.randn(6 * dim) * .01,
+        'video_head_modulation.bias': torch.randn(2 * dim) * .01,
+    }
+    save_file(lora, tmp_path / 'lora.safetensors')
+    save_file(time, tmp_path / 'time_embedding.safetensors')
+    save_file({'latent_mix.Q': torch.eye(128), 'step_codes': torch.tensor([1000, 968, 926, 871])},
+              tmp_path / 'dataspace.safetensors')
+    monkeypatch.setattr(rapid, '_adapter_dirs', lambda: {'test': str(tmp_path)})
+    return lora
+
+
+def make_model(fused):
+    holder = torch.nn.Module()
+    holder.diffusion_model = MOSSImage(dim=128, ffn_dim=192, num_heads=1, num_layers=1,
+                                     fused_qkv=fused, dtype=torch.float32, operations=comfy.ops.manual_cast)
+    for p in holder.parameters():
+        torch.nn.init.normal_(p, std=.02)
+    return ModelPatcher(holder, torch.device('cpu'), torch.device('cpu'))
+
+
+@pytest.mark.parametrize('fused', [False, True])
+def test_lora_apply_and_restore(adapter, fused):
+    model = make_model(fused)
+    layer = model.model.diffusion_model.layers[0]
+    base_weights = {k: m.weight.detach().clone() for k, m in layer.proj.items()}
+    patched, _, sigmas = rapid.RapidMOSSAdapter().apply(model, 'test', 4, .7)
+    assert sigmas.tolist() == pytest.approx([1, .968, .926, .871, 0])
+    assert not model.object_patches
+    patched.patch_model(load_weights=False)
+    for name, module in layer.proj.items():
+        x = torch.randn(1, 5, module.weight.shape[1])
+        expected = torch.nn.functional.linear(x, base_weights[name])
+        names = ('q', 'k', 'v') if name == 'qkv' else (name,)
+        deltas = []
+        for proj in names:
+            prefix = 'diffusion_model.layers.0.proj.' + proj
+            deltas.append(.7 * (x @ adapter[prefix + '.lora_down.weight'].T) @ adapter[prefix + '.lora_up.weight'].T)
+        expected += torch.cat(deltas, dim=-1)
+        torch.testing.assert_close(module(x), expected)
+        torch.testing.assert_close(module.weight, base_weights[name], rtol=0, atol=0)
+    patched.unpatch_model()
+    for name, module in layer.proj.items():
+        x = torch.randn(1, 2, module.weight.shape[1])
+        torch.testing.assert_close(module(x), torch.nn.functional.linear(x, base_weights[name]))
+
+
+def test_sampler_matches_dataspace_recurrence():
+    torch.manual_seed(4)
+    z = torch.randn(1, 3, 2, 2)
+    q = torch.randn(3, 3) * .1
+    codes = torch.tensor([1000, 968, 926, 871])
+    sigmas = torch.cat((codes.float() / 1000, torch.zeros(1)))
+    calls = []
+    def denoise(x, sigma):
+        calls.append(sigma.clone())
+        return x - sigma[:, None, None, None] * (x * .3 + sigma[:, None, None, None])
+    expected = torch.zeros_like(z)
+    for code in codes:
+        x = z + torch.einsum('oc,bchw->bohw', q, expected)
+        expected = z - (x * .3 + code / 1000)
+    actual = rapid.sample_dataspace(denoise, z, sigmas, Q=q, codes=codes)
+    torch.testing.assert_close(actual, expected)
+    assert len(calls) == 4
+
+
+def test_reject_too_many_steps(adapter):
+    with pytest.raises(ValueError, match='4 step codes'):
+        rapid.RapidMOSSAdapter().apply(make_model(True), 'test', 5, 1.)
+
+
+@pytest.mark.parametrize('dtype', [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize('device', ['cpu', pytest.param('cuda', marks=pytest.mark.skipif(
+    os.environ.get('RAPID_MOSS_TEST_CUDA') != '1', reason='set RAPID_MOSS_TEST_CUDA=1 for GPU checks'))])
+def test_native_forward_fused_matches_split(adapter, dtype, device):
+    model = make_model(False)
+    fused = make_model(True)
+    sd = model.model.state_dict()
+    prefix = 'diffusion_model.layers.0.proj.'
+    sd[prefix + 'qkv.weight'] = torch.cat([sd.pop(prefix + n + '.weight') for n in ('q', 'k', 'v')])
+    fused.model.load_state_dict(sd)
+    x = torch.randn(1, 128, 2, 3, device=device, dtype=dtype)
+    ctx = torch.randn(1, 4, 4096, device=device, dtype=dtype)
+    t = torch.tensor([.968], device=device)
+    outputs = []
+    for m in (model, fused):
+        m.model.to(device=device, dtype=dtype)
+        patched, _, _ = rapid.RapidMOSSAdapter().apply(m, 'test', 4, 1.)
+        patched.patch_model(load_weights=False)
+        with torch.no_grad():
+            outputs.append(m.model.diffusion_model(x.clone(), t, ctx))
+        patched.unpatch_model()
+    assert all(torch.isfinite(out).all() for out in outputs)
+    assert all(out.abs().max() < 10 for out in outputs)
+    torch.testing.assert_close(*outputs, rtol=.01, atol=.001)

comfyui/rapid-moss/README.md

@@ -3,7 +3,23 @@
 MOSS-Image-SFT を **4 ステップ・CFG なし**で動かすカスタムノードです(rapid-moss の DMAD 4 ステップアダプタ。lodestones の 8 ステップ DMD も同じノードで使えます)。
 ComfyUI 標準の `SamplerCustomAdvanced` と組み合わせて使います。
 
-## 必要なもの
+## ComfyUI 標準ローダーを使う場合
+
+MOSS-Image 対応の最新 ComfyUI と、この `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` です。
+
+`AttributeError: 'ModuleDict' object has no attribute 'q'` が出る場合は、このノードの古い版が使われています。ComfyUI 用のモデルは Q/K/V が `qkv` にまとめられているため、この版では各 LoRA の差分を Q/K/V の対応する出力へ別々に適用します。元の HF 形式も引き続き使えます。
+
+また、標準モデルのブロックではアダプタの FP32 modulation と残差を保持します。BF16 用の fused RMS/AdaLN に FP32 modulation を渡すことによる黒画像を防ぎ、FFN の LoRA も必ず適用します。
+
+画像生成モデルは **SFT の BF16** を使ってください。INT8/FP8/GGUF の生成モデルはこのアダプタの検証対象外です。テキストエンコーダの量子化は別で、ComfyUI が対応する Qwen3.5-9B を選べます。
+
+**English:** Native ComfyUI loaders are supported. Update this node and use `rapid-moss-4step-native.json` with the SFT BF16 diffusion model, `CLIPLoader` type `moss_image`, and the FLUX.2 VAE. The OpenMOSS custom loader is optional. This fixes discussion #1: repacked models have a fused `proj.qkv`, while the adapter has separate Q/K/V LoRAs. Their unmerged fp32 updates are applied to the corresponding output slices. Native blocks also preserve fp32 modulation/residuals instead of feeding mixed dtypes into the fused RMS/AdaLN kernel, which can produce black images. The original HF-loader workflow remains supported.
+
+## HF ローダーを使う場合の必要なもの
 
 | もの | 入手先 |
 |---|---|
@@ -11,14 +27,14 @@ ComfyUI 標準の `SamplerCustomAdvanced` と組み合わせて使います。
 | 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/rapid-moss](https://huggingface.co/darask0/rapid-moss) の 3 ファイルを `ComfyUI/models/rapid_moss/rapid-moss-4step/` へ |
+| アダプタ | [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/rapid-moss lora.safetensors time_embedding.safetensors dataspace.safetensors \
+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/rapid-moss --include "comfyui/rapid-moss/*" --local-dir /tmp/rapid-moss
+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 \
@@ -57,7 +73,7 @@ API で投げる場合は `example_workflows/rapid-moss-4step-api.json` を使
 
 ## 注意
 
-モデル全体の注意点(分かっている弱点、検証していないことなど)は、[モデルカードの「注意点」](https://huggingface.co/darask0/rapid-moss#注意点)にまとめています。ComfyUI で特に気をつけることは次のとおりです。
+モデル全体の注意点(分かっている弱点、検証していないことなど)は、[モデルカードの「注意点」](https://huggingface.co/darask0/MOSS-Image_4-step_adapter_lora#注意点)にまとめています。ComfyUI で特に気をつけることは次のとおりです。
 
 - **`model_path` は MOSS-Image-SFT のフォルダにしてください。** アダプタは SFT から学習しています。Base との組み合わせは確かめていません。
 - **このアダプタに `LoraLoaderModelOnly` は使えません。** LoRA を重みにマージすると、差分が bf16 の丸めで消えて効きません。必ずこのノードで適用してください。
@@ -72,17 +88,22 @@ API で投げる場合は `example_workflows/rapid-moss-4step-api.json` を使
 
 ## 動作確認(2026-10-17)
 
+- **標準ローダー対応の修正版:** ComfyUI `5cc4b870`、torch 2.13.0+cu130、RTX PRO 6000 で、標準ローダーと HF ローダーの両方から 1024×1024・4ステップ・seed 1234 の画像生成を確認しました。標準側は公式 SFT 重みの Q/K/V を連結した BF16 ファイルと、BF16 Qwen3.5-9B を使用しています。両実装の計算の違いにより画像の細部は異なります。
+- **回帰テスト:** Q/K/V 分離・結合形式の LoRA 計算と解除、専用サンプラー、ステップ数上限、CPU/CUDA の FP32/BF16 forward の計8件が成功しました。`COMFYUI_ROOT=/path/to/ComfyUI RAPID_MOSS_TEST_CUDA=1 python -m pytest code/tests/test_comfy_adapter.py` で実行できます。
+
+以下は初版の 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/rapid-moss/resolve/main/images/comfyui-check.jpg)
+![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). Install the
-official MOSS-Image nodes, copy this folder to `ComfyUI/custom_nodes/rapid-moss/`, and put `lora.safetensors`,
+Custom node for 4-step, CFG-free MOSS-Image-SFT with the rapid-moss adapter (also runs lodestones' 8-step DMD). For the
+HF-loader workflow below, install the official MOSS-Image nodes. Copy this folder to `ComfyUI/custom_nodes/rapid-moss/`, and put `lora.safetensors`,
 `time_embedding.safetensors`, `dataspace.safetensors` in `ComfyUI/models/rapid_moss/<name>/`. Load
 `example_workflows/rapid-moss-4step.json`, set the HF loader's `model_path` to MOSS-Image-SFT, and run. Use `BasicGuider`
 (no CFG / negative), no `ModelSamplingSD3` or SLG, `weight_dtype` default/bf16/fp32 (fp8 unsupported). The node applies the

comfyui/rapid-moss/example_workflows/rapid-moss-4step-native-api.json

@@ -0,0 +1,144 @@
+{
+  "1": {
+    "class_type": "UNETLoader",
+    "inputs": {
+      "unet_name": "moss_image_sft_bf16.safetensors",
+      "weight_dtype": "default"
+    }
+  },
+  "2": {
+    "class_type": "RapidMOSSAdapter",
+    "_meta": {
+      "title": "MOSS-Image Few-Step Adapter (rapid-moss)"
+    },
+    "inputs": {
+      "model": [
+        "1",
+        0
+      ],
+      "adapter": "rapid-moss-4step",
+      "steps": 4,
+      "lora_strength": 1.0
+    }
+  },
+  "3": {
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "Prompt"
+    },
+    "inputs": {
+      "clip": [
+        "10",
+        0
+      ],
+      "text": "A high-quality anime illustration of a girl with long silver hair and teal eyes, wearing a yellow raincoat, standing at a rainy bus stop at night, neon reflections on wet asphalt"
+    }
+  },
+  "4": {
+    "class_type": "BasicGuider",
+    "_meta": {
+      "title": "BasicGuider (no CFG)"
+    },
+    "inputs": {
+      "model": [
+        "2",
+        0
+      ],
+      "conditioning": [
+        "3",
+        0
+      ]
+    }
+  },
+  "5": {
+    "class_type": "RandomNoise",
+    "_meta": {
+      "title": "RandomNoise"
+    },
+    "inputs": {
+      "noise_seed": 42
+    }
+  },
+  "6": {
+    "class_type": "EmptyFlux2LatentImage",
+    "_meta": {
+      "title": "Empty Flux 2 Latent"
+    },
+    "inputs": {
+      "width": 1024,
+      "height": 1024,
+      "batch_size": 1
+    }
+  },
+  "7": {
+    "class_type": "SamplerCustomAdvanced",
+    "_meta": {
+      "title": "SamplerCustomAdvanced"
+    },
+    "inputs": {
+      "noise": [
+        "5",
+        0
+      ],
+      "guider": [
+        "4",
+        0
+      ],
+      "sampler": [
+        "2",
+        1
+      ],
+      "sigmas": [
+        "2",
+        2
+      ],
+      "latent_image": [
+        "6",
+        0
+      ]
+    }
+  },
+  "8": {
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE Decode"
+    },
+    "inputs": {
+      "samples": [
+        "7",
+        0
+      ],
+      "vae": [
+        "11",
+        0
+      ]
+    }
+  },
+  "9": {
+    "class_type": "SaveImage",
+    "_meta": {
+      "title": "Save Image"
+    },
+    "inputs": {
+      "filename_prefix": "rapid-moss",
+      "images": [
+        "8",
+        0
+      ]
+    }
+  },
+  "10": {
+    "class_type": "CLIPLoader",
+    "inputs": {
+      "clip_name": "qwen3.5_9b_int8_convrot.safetensors",
+      "type": "moss_image",
+      "device": "default"
+    }
+  },
+  "11": {
+    "class_type": "VAELoader",
+    "inputs": {
+      "vae_name": "flux2-vae.safetensors"
+    }
+  }
+}

comfyui/rapid-moss/example_workflows/rapid-moss-4step-native.json

@@ -0,0 +1,546 @@
+{
+  "last_node_id": 11,
+  "last_link_id": 12,
+  "nodes": [
+    {
+      "id": 1,
+      "type": "UNETLoader",
+      "pos": [
+        40,
+        120
+      ],
+      "size": [
+        380,
+        100
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "MODEL",
+          "type": "MODEL",
+          "links": [
+            1
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "UNETLoader"
+      },
+      "widgets_values": [
+        "moss_image_sft_bf16.safetensors",
+        "default"
+      ],
+      "title": "Load MOSS-Image SFT (BF16)"
+    },
+    {
+      "id": 2,
+      "type": "RapidMOSSAdapter",
+      "pos": [
+        470,
+        60
+      ],
+      "size": [
+        380,
+        130
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "model",
+          "type": "MODEL",
+          "link": 1
+        }
+      ],
+      "outputs": [
+        {
+          "name": "model",
+          "type": "MODEL",
+          "links": [
+            3
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "sampler",
+          "type": "SAMPLER",
+          "links": [
+            7
+          ],
+          "slot_index": 1
+        },
+        {
+          "name": "sigmas",
+          "type": "SIGMAS",
+          "links": [
+            8
+          ],
+          "slot_index": 2
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "RapidMOSSAdapter"
+      },
+      "widgets_values": [
+        "rapid-moss-4step",
+        4,
+        1.0
+      ]
+    },
+    {
+      "id": 3,
+      "type": "CLIPTextEncode",
+      "pos": [
+        470,
+        250
+      ],
+      "size": [
+        420,
+        200
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "clip",
+          "type": "CLIP",
+          "link": 2
+        }
+      ],
+      "outputs": [
+        {
+          "name": "CONDITIONING",
+          "type": "CONDITIONING",
+          "links": [
+            4
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CLIPTextEncode"
+      },
+      "widgets_values": [
+        "A high-quality anime illustration of a girl with long silver hair and teal eyes, wearing a yellow raincoat, standing at a rainy bus stop at night, neon reflections on wet asphalt"
+      ],
+      "title": "Prompt"
+    },
+    {
+      "id": 4,
+      "type": "BasicGuider",
+      "pos": [
+        940,
+        60
+      ],
+      "size": [
+        260,
+        50
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "model",
+          "type": "MODEL",
+          "link": 3
+        },
+        {
+          "name": "conditioning",
+          "type": "CONDITIONING",
+          "link": 4
+        }
+      ],
+      "outputs": [
+        {
+          "name": "GUIDER",
+          "type": "GUIDER",
+          "links": [
+            6
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "BasicGuider"
+      },
+      "widgets_values": [],
+      "title": "BasicGuider (no CFG)"
+    },
+    {
+      "id": 5,
+      "type": "RandomNoise",
+      "pos": [
+        940,
+        160
+      ],
+      "size": [
+        260,
+        82
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "NOISE",
+          "type": "NOISE",
+          "links": [
+            5
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "RandomNoise"
+      },
+      "widgets_values": [
+        42,
+        "randomize"
+      ]
+    },
+    {
+      "id": 6,
+      "type": "EmptyFlux2LatentImage",
+      "pos": [
+        940,
+        290
+      ],
+      "size": [
+        260,
+        106
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "LATENT",
+          "type": "LATENT",
+          "links": [
+            9
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "EmptyFlux2LatentImage"
+      },
+      "widgets_values": [
+        1024,
+        1024,
+        1
+      ]
+    },
+    {
+      "id": 7,
+      "type": "SamplerCustomAdvanced",
+      "pos": [
+        1250,
+        60
+      ],
+      "size": [
+        300,
+        110
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "noise",
+          "type": "NOISE",
+          "link": 5
+        },
+        {
+          "name": "guider",
+          "type": "GUIDER",
+          "link": 6
+        },
+        {
+          "name": "sampler",
+          "type": "SAMPLER",
+          "link": 7
+        },
+        {
+          "name": "sigmas",
+          "type": "SIGMAS",
+          "link": 8
+        },
+        {
+          "name": "latent_image",
+          "type": "LATENT",
+          "link": 9
+        }
+      ],
+      "outputs": [
+        {
+          "name": "output",
+          "type": "LATENT",
+          "links": [
+            10
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "denoised_output",
+          "type": "LATENT",
+          "links": [],
+          "slot_index": 1
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "SamplerCustomAdvanced"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 8,
+      "type": "VAEDecode",
+      "pos": [
+        1250,
+        230
+      ],
+      "size": [
+        210,
+        46
+      ],
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "samples",
+          "type": "LATENT",
+          "link": 10
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": 11
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            12
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VAEDecode"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 9,
+      "type": "SaveImage",
+      "pos": [
+        1600,
+        60
+      ],
+      "size": [
+        520,
+        560
+      ],
+      "flags": {},
+      "order": 8,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 12
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "SaveImage"
+      },
+      "widgets_values": [
+        "rapid-moss"
+      ]
+    },
+    {
+      "id": 10,
+      "type": "CLIPLoader",
+      "title": "Load Qwen3.5 text encoder",
+      "pos": [
+        40,
+        310
+      ],
+      "size": [
+        380,
+        120
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "CLIP",
+          "type": "CLIP",
+          "links": [
+            2
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CLIPLoader"
+      },
+      "widgets_values": [
+        "qwen3.5_9b_int8_convrot.safetensors",
+        "moss_image",
+        "default"
+      ]
+    },
+    {
+      "id": 11,
+      "type": "VAELoader",
+      "title": "Load FLUX.2 VAE",
+      "pos": [
+        40,
+        510
+      ],
+      "size": [
+        380,
+        120
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "VAE",
+          "type": "VAE",
+          "links": [
+            11
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VAELoader"
+      },
+      "widgets_values": [
+        "flux2-vae.safetensors"
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      1,
+      0,
+      2,
+      0,
+      "MODEL"
+    ],
+    [
+      2,
+      10,
+      0,
+      3,
+      0,
+      "CLIP"
+    ],
+    [
+      3,
+      2,
+      0,
+      4,
+      0,
+      "MODEL"
+    ],
+    [
+      4,
+      3,
+      0,
+      4,
+      1,
+      "CONDITIONING"
+    ],
+    [
+      5,
+      5,
+      0,
+      7,
+      0,
+      "NOISE"
+    ],
+    [
+      6,
+      4,
+      0,
+      7,
+      1,
+      "GUIDER"
+    ],
+    [
+      7,
+      2,
+      1,
+      7,
+      2,
+      "SAMPLER"
+    ],
+    [
+      8,
+      2,
+      2,
+      7,
+      3,
+      "SIGMAS"
+    ],
+    [
+      9,
+      6,
+      0,
+      7,
+      4,
+      "LATENT"
+    ],
+    [
+      10,
+      7,
+      0,
+      8,
+      0,
+      "LATENT"
+    ],
+    [
+      11,
+      11,
+      0,
+      8,
+      1,
+      "VAE"
+    ],
+    [
+      12,
+      8,
+      0,
+      9,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {},
+  "version": 0.4
+}

comfyui/rapid-moss/nodes.py

@@ -1,6 +1,6 @@
 """Few-step MOSS-Image adapters (rapid-moss DMAD 4-step, lodestones 8-step DMD) for ComfyUI.
 
-Works on the MODEL from the MOSS-Image nodes (github.com/OpenMOSS/MOSS-Image, comfyui/moss-image).
+Works on native ComfyUI MOSS-Image models and the OpenMOSS HF loader.
 An adapter is a folder in models/rapid_moss/ holding lora.safetensors, time_embedding.safetensors and
 dataspace.safetensors. The node returns a patched MODEL plus the data-space SAMPLER and SIGMAS for
 SamplerCustomAdvanced (use BasicGuider: no CFG, no negative prompt).
@@ -72,6 +72,47 @@ def _lora_forward(module, down, up, scale):
     return forward
 
 
+def _qkv_lora_forward(module, adapters):
+    """Repacked ComfyUI weights fuse Q/K/V; each LoRA still has its own down projection."""
+    class_forward = type(module).forward
+    tensors = [_DeviceCopy(down, up) for down, up, _ in adapters]
+
+    def forward(x, *args, **kwargs):
+        out = class_forward(module, x, *args, **kwargs)
+        parts = out.chunk(3, dim=-1)
+        result = []
+        for part, cached, (_, _, scale) in zip(parts, tensors, adapters):
+            down, up = cached.get(x.device)
+            delta = ((x.float() @ down.t()) @ up.t()) * scale
+            result.append(part + delta.to(part.dtype))
+        return torch.cat(result, dim=-1)
+    return forward
+
+
+def _native_block_forward(block):
+    # Import only for native models; the standalone OpenMOSS plugin also runs on older ComfyUI.
+    from comfy.ldm.moss_image.model import attention, situ
+
+    def forward(x, mod, pe, transformer_options=None):
+        transformer_options = {} if transformer_options is None else transformer_options
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod
+        dtype = block.proj['ffn_up_gate'].weight.dtype
+        # The native fused RMS/AdaLN kernel assumes matching dtypes. The adapter's
+        # modulation is fp32, so use the module norm and keep the residual in fp32.
+        y = (block.norm1(x) * (1 + scale_a) + shift_a).to(dtype)
+        if block.fused_qkv:
+            q, k, v = block.proj['qkv'](y).chunk(3, dim=-1)
+        else:
+            q, k, v = (block.proj[n](y) for n in ('q', 'k', 'v'))
+        out = attention(q, k, v, block.q_norm, block.k_norm, pe, block.num_heads,
+                        block.comfy_attention, transformer_options)
+        x = x + block.proj['o'](out).float() * gate_a
+        y = (block.norm2(x) * (1 + scale_f) + shift_f).to(dtype)
+        hidden, gate = block.proj['ffn_up_gate'](y).chunk(2, dim=-1)
+        return x + block.proj['ffn_down'](situ(hidden, gate)).float() * gate_f
+    return forward
+
+
 @torch.no_grad()
 def sample_dataspace(model, x, sigmas, extra_args=None, callback=None, disable=None, Q=None, codes=None):
     """x0 = 0; for r: x = z + Q x0, v = net(x, code_r), x0 = z - v (one noise z, no schedule, no CFG)."""
@@ -132,13 +173,22 @@ class RapidMOSSAdapter:
         for name in ('shared_block_modulation', 'video_head_modulation'):
             weight = patched.get_model_object(DM + name).weight.detach().to('cpu', torch.float32)
             patched.add_object_patch(DM + name, _FP32Linear(weight, time_path[name + '.bias'].float()))
-        for i in range(len(patched.get_model_object(DM + 'layers'))):
+        for i, layer in enumerate(patched.get_model_object(DM + 'layers')):
+            if hasattr(layer, 'fused_qkv'):
+                patched.add_object_patch(f'{DM}layers.{i}.forward', _native_block_forward(layer))
+            qkv = []
             for proj in PROJ:
                 target = f'{DM}layers.{i}.proj.{proj}'
                 down = lora[target + '.lora_down.weight'].float()
                 up = lora[target + '.lora_up.weight'].float()
                 scale = lora_strength * float(lora[target + '.alpha']) / down.shape[0]
-                patched.add_object_patch(target + '.forward', _lora_forward(patched.get_model_object(target), down, up, scale))
+                if proj in ('q', 'k', 'v') and 'qkv' in layer.proj:
+                    qkv.append((down, up, scale))
+                else:
+                    patched.add_object_patch(target + '.forward', _lora_forward(patched.get_model_object(target), down, up, scale))
+            if qkv:
+                target = f'{DM}layers.{i}.proj.qkv'
+                patched.add_object_patch(target + '.forward', _qkv_lora_forward(patched.get_model_object(target), qkv))
 
         sampler = comfy.samplers.KSAMPLER(sample_dataspace, extra_options={'Q': data['latent_mix.Q'].float(), 'codes': codes})
         sigmas = torch.cat((codes[:steps] / 1000.0, torch.zeros(1)))