ComfyUI native single-file loader fails on all Comfy-Org checkpoints: _Block expects split q/k/v but checkpoints ship fused proj.qkv (AttributeError: 'NoneType' object has no attribute 'dtype')

Bắt đầu tại comfyui/moss-image/model.py, so sánh _Block với _TextBlock và kiểm tra cách loader gốc ánh xạ các khóa checkpoint. Đối chiếu các khóa hợp nhất proj.qkv được báo cáo với các kỳ vọng về proj.q/k/v đã tách. Sau đó, tải checkpoint Comfy-Org dạng một tệp và kiểm tra xem nó có đi đến bước lấy mẫu mà không thiếu trọng số phép chiếu hay không. Công việc hoàn tất khi loader gốc xử lý thành công các checkpoint được báo cáo.

Summary

MOSSImageUNETLoader (the native single-file loader) cannot load any of the four Comfy-Org/MOSS-Image single-file checkpoints. It fails immediately with:

AttributeError: 'NoneType' object has no attribute 'dtype'

Root cause: _Block in comfyui/moss-image/model.py declares separate proj.q / proj.k / proj.v / proj.o linears, but every Comfy-Org repackaged checkpoint stores a fused proj.qkv.weight [3*dim, dim]. Nothing populates proj.q/k/v, so under ComfyUI's DynamicVRAM lazy init they stay None and line 96 blows up.

Note this is inconsistent with _TextBlock in the same file, which already uses fused self.qkv — so it looks like a simple oversight, not an intentional design difference.

Environment

ComfyUI	0.39.0 (d49e888)
comfy-aimdo	present → DynamicVRAM enabled (comfy.memory_management.aimdo_enabled = True)
Custom node	comfyui/moss-image copied to custom_nodes/moss-image/
Python	3.12
GPU	24 GB VRAM laptop (CUDA 13.x, torch 2.14.1+cu130)
Repo revision	95559c9 (docs: use official arXiv BibTeX citation)

Steps to reproduce

1. Download Comfy-Org/MOSS-Image → diffusion_models/moss_image_base_int8_convrot.safetensors (5.41 GB), text_encoders/qwen3.5_9b_w4a8_gptq.safetensors, vae/flux2-vae.safetensors into the standard ComfyUI/models/{diffusion_models,text_encoders,vae}/ folders.
2. Load the node MOSS-Image UNET Loader (MOSSImageUNETLoader) with unet_name = moss_image_base_int8_convrot.safetensors.
3. Queue any prompt.

Result: node fails during model load, before sampling starts.

Evidence — checkpoint key layout

All four Comfy-Org single-file DiT checkpoints use fused qkv, zero split q/k/v:

Checkpoint	proj.qkv.weight	proj.{q,k,v}.weight
moss_image_base_int8_convrot.safetensors	42	0
moss_image_sft_int8_convrot.safetensors	42	0
moss_image_base_bf16.safetensors	42	0
moss_image_sft_bf16.safetensors	42	0

Example keys from moss_image_base_int8_convrot.safetensors (layers.0.*, 42 blocks total):

layers.0.proj.qkv.weight        [9216, 3072]     # dim = 3072 -> 3*dim
layers.0.proj.qkv.weight_scale  [9216, 1]
layers.0.proj.o.weight          [3072, 3072]
layers.0.proj.ffn_up_gate.weight[16384, 3072]
layers.0.proj.ffn_down.weight   [3072, 8192]

For comparison, the original diffusers repo OpenMOSS-Team/MOSS-Image-Base stores split keys — transformer/diffusion_pytorch_model.safetensors.index.json maps 126 layers.N.proj.{q,k,v}.weight entries (42 blocks × 3; 60 in shard -00001-of-00003, 66 in -00002-of-00003) — which is exactly what the current _Block expects. That is why the HF Loader path works and only the native single-file path breaks.

Root cause in code

comfyui/moss-image/model.py (revision 95559c9), _Block.__init__ lines 84–89:

self.proj = nn.ModuleDict(
    {
        name: operations.Linear(dim, dim, bias=False, device=device, dtype=dtype)
        for name in ("q", "k", "v", "o")
    }
)

and _Block.forward lines 96–100:

dtype = self.proj["q"].weight.dtype                                    # <-- raises
q = self.q_norm(self.proj["q"](y).unflatten(-1, (self.heads, -1))).flatten(2)
k = self.k_norm(self.proj["k"](y).unflatten(-1, (self.heads, -1))).flatten(2)
x = x + self.proj["o"](_attention(q, k, self.proj["v"](y), rope, self.heads)) * gate_a[:, None]

load_state_dict(..., strict=False) confirms the mismatch:

missing_keys    : ['proj.k.weight', 'proj.q.weight', 'proj.v.weight']
unexpected_keys : ['proj.qkv.weight']

With comfy-aimdo enabled, comfy.ops.disable_weight_init.Linear.__init__ skips torch.nn.Linear.__init__ and leaves self.weight = None (it is materialised lazily in _load_from_state_dict, only for keys present in the checkpoint). Since proj.q.weight is absent, it stays None → .dtype on None → AttributeError.

Fix that works (verified)

Patch _Block to a fused qkv projection, mirroring _TextBlock:

self.proj = nn.ModuleDict(
    {
        "qkv": operations.Linear(dim, 3 * dim, bias=False, device=device, dtype=dtype),
        "o": operations.Linear(dim, dim, bias=False, device=device, dtype=dtype),
    }
)

def forward(self, x, rope, modulation):
    shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = modulation.unbind(1)
    y = self.norm1(x) * (1 + scale_a[:, None]) + shift_a[:, None]
    dtype = self.proj["qkv"].weight.dtype
    y = y.to(dtype)
    q, k, v = self.proj["qkv"](y).chunk(3, -1)
    q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2)
    k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2)
    x = x + self.proj["o"](_attention(q, k, v, rope, self.heads)) * gate_a[:, None]
    ...

After the patch, load_state_dict reports no missing/unexpected keys and generation succeeds.

Verified output with the patch (ComfyUI 0.39.0, int8 single-file)

Official recipe from the workflow (50 steps, CFG 3, Shift 5, SLG layer 22, 1024×1024, Euler/sgm_uniform):

Seed	CFG	Steps	Wall time	Peak VRAM
42	3.0	50	71.7 s	~10.5 GB / 24 GB
1234	3.5	50	48.8 s	~10.5 GB / 24 GB

≈1.15 it/s, no OOM, images are clean (not noise/black).

Suggestion

Since the fused layout is what Comfy-Org ships (and what HF diffusers actually uses for text_embedder), making _Block fused is likely the low-risk default. Options:

1. Preferred: switch _Block to fused qkv (matches _TextBlock and all Comfy-Org checkpoints).
2. Accept both by branching on the key layout.
3. Or add a key-remap in load_model_state() that splits proj.qkv.weight into proj.q/k/v when the split layout is detected — this would keep the original diffusers shards loading through the native loader too.

Also worth documenting: the native single-file loader currently only works with the diffusers-layout shards, not the Comfy-Org repackaged single files linked from the main README's day-0 support note.

Happy to send a PR for option 1 if that's the direction you prefer.

---

中文摘要

comfyui/moss-image/model.py 嘅 _Block 用分開 q/k/v/o，但 Comfy-Org 全部 4 個單檔 checkpoint 都係融合 proj.qkv（42 個 block、[9216,3072]），proj.q/k/v 完全冇 key → 開 DynamicVRAM 時 weight 保持 None → AttributeError: 'NoneType' object has no attribute 'dtype'。

同一份檔案嘅 _TextBlock 本身已經用融合 qkv，所以似係手誤。改成融合 qkv 之後實測正常出圖（1024²、50 步、48.8–71.7 秒、峰值 ~10.5GB VRAM）。而原始 diffusers 格式（OpenMOSS-Team/MOSS-Image-Base）確實係分開 q/k/v，所以只有 native 單檔 loader 配 Comfy-Org 檔先爆。