Diffusers LoRA modules with no place are counted under any wrapper prefix; Watermark hint names the trimmed-clip exception

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Tw9fjX7o2YfeGQE18MEtPu

modules/loaders/lora/keys.py

@@ -33,6 +33,9 @@ _SUFFIX = (("attn.to_out.0", "attn.out_proj"), ("ff.net.0.proj", "mlp.fc1"), ("f
 _QKV = ("to_q", "to_k", "to_v")
 
 
+_BLOCKS = re.compile(r"(?:^|\.)(?:token_refiner\.)?(?:transformer_blocks|refiner_blocks)\.\d+\.")      # a Diffusers module under any training wrapper
+
+
 def from_diffusers(lora, fused_rows=None):
     """-> (lora, extra key map). A LoRA written against the Diffusers layout of the H3 transformer, renamed to the one
     ComfyUI loads, or the LoRA unchanged with {} when it is not in that layout.
@@ -129,7 +132,7 @@ def match(model, clip, lora):
     modules = 0
     sd = model.model.state_dict()
     if any(".attn.qkv_proj." in str(v) for v in key_map.values()):         # only the H3 transformer is renamed: other models use `transformer_blocks` as they are
-        modules = len({k.split(".lora_")[0].split(".alpha")[0] for k in lora if k.startswith("transformer.") and (".lora_" in k or k.endswith(".alpha"))})
+        modules = len({k.split(".lora_")[0].split(".alpha")[0] for k in lora if (k.startswith("transformer.") or _BLOCKS.search(k)) and (".lora_" in k or k.endswith(".alpha"))})
         lora, extra = from_diffusers(lora, lambda base: (lambda w: int(w.shape[0]) if w is not None and w.dim() else None)(sd.get(f"diffusion_model.{base}.weight")))
     key_map.update(extra)
     converted = comfy.lora_convert.convert_lora(lora)

modules/loaders/tests/test_diffusers_h3_lora.py

@@ -205,3 +205,37 @@ def test_modules_with_no_place_in_the_model_are_counted_and_alerted(patcher, mon
     monkeypatch.setattr(comfy.utils, "load_torch_file", lambda p, **kw: dict(lora))
     nodes.FunUILoraLoader.execute(model=patcher, lora_name="x.safetensors", strength_model=1.0)
     assert any("2 of its 8 modules have no place" in t for t in said), said
+
+
+# Modules with no place in H3 are counted whatever training wrapper sits in front of the Diffusers key names.
+import pytest as _pytest
+_RANK = 4
+
+@_pytest.fixture(scope="module")
+def wrapper_patcher():
+    import comfy.model_patcher as mp
+    import comfy.supported_models as sm
+    config = sm.MiniMaxH3(dict(
+        image_model="minimax_h3", hidden_size=64, num_layers=2, token_refiner_num_layers=1, num_attention_heads=2,
+        attention_head_dim=32, ffn_hidden_size=64, text_dim=48, timestep_input_dim=32, time_embed_hidden_size=64,
+        time_embed_dim=32, rope_inv_freq_len=4))
+    config.set_inference_dtype(torch.float32, None)
+    base = config.get_model({}, device=torch.device("cpu"))
+    return mp.ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu"))
+
+def _pair(o, i, g):
+    return torch.randn(_RANK, i, generator=g), torch.randn(o, _RANK, generator=g)
+
+@_pytest.mark.parametrize("wrapper", ["transformer.", "", "diffusion_model.", "base_model.model.", "dit."])
+def test_unplaced_counted_for_wrapper(wrapper_patcher, wrapper):
+    from modules.loaders.lora.keys import match
+    g = torch.Generator().manual_seed(0)
+    lora = {}
+    for name, (o, i) in {"attn.to_q": (64, 64), "attn.to_k": (64, 64), "attn.to_v": (64, 64), "attn.to_out.0": (64, 64),
+                         "ff.net.0.proj": (128, 64), "ff.net.2": (64, 64), "norm1.linear": (384, 64), "attn.to_add_out": (64, 64)}.items():
+        a, b = _pair(o, i, g)
+        lora[f"{wrapper}transformer_blocks.1.{name}.lora_A.weight"] = a
+        lora[f"{wrapper}transformer_blocks.1.{name}.lora_B.weight"] = b
+    patches, how, dropped, unplaced = match(wrapper_patcher, None, lora)
+    assert unplaced == 2, (wrapper, unplaced)
+

modules/output/save/__init__.py

@@ -45,7 +45,7 @@ SETTINGS = {
         "options": [{"value": "off", "label": "Off"}, {"value": "text", "label": "Text"}, {"value": "picture", "label": "Picture"}],
         "label": "Watermark",
         "hint": "A text or a picture from the media bin in one corner of the finished video. It is laid on last, so the settings above never change it. "
-                "It is in the saved file only: the next shot's start picture and the Preview's own frames do not carry it.",
+                "It is in the saved file only: the next shot's start picture and the Preview's own frames do not carry it (a clip trimmed short at its end starts the next shot from the saved file, so then it does).",
     },
     "watermark_text": {
         "type": "text", "default": "", "label": "Watermark text", "when": {"watermark": "text"},