From 0414df00c9811fc0476941fef0df5bda452d1ebb Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 8 Nov 2026 20:19:39 +0000 Subject: [PATCH] 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 Claude-Session: https://claude.ai/code/session_01Tw9fjX7o2YfeGQE18MEtPu --- modules/loaders/lora/keys.py | 5 ++- .../loaders/tests/test_diffusers_h3_lora.py | 34 +++++++++++++++++++ modules/output/save/__init__.py | 2 +- 3 files changed, 39 insertions(+), 2 deletions(-) diff --git a/modules/loaders/lora/keys.py b/modules/loaders/lora/keys.py index 160041a3..149bf70b 100644 --- a/modules/loaders/lora/keys.py +++ b/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) diff --git a/modules/loaders/tests/test_diffusers_h3_lora.py b/modules/loaders/tests/test_diffusers_h3_lora.py index 56428f7e..889570c2 100644 --- a/modules/loaders/tests/test_diffusers_h3_lora.py +++ b/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) + diff --git a/modules/output/save/__init__.py b/modules/output/save/__init__.py index 44374943..4250448f 100644 --- a/modules/output/save/__init__.py +++ b/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"},