LoRA: a wrapped Diffusers file reaches the fused qkv rows; a wrong-width qkv pair is dropped, not merged

README credits Karras et al. for the churn idea in Turbo. Tests for three wrappers and the width check.

README.md

@@ -168,6 +168,7 @@ I express my deepest gratitude to:
 - PTD few-step support for MiniMax H3 is ported from [ComfyUI-PTD-Compact](https://github.com/Nynxz/ComfyUI-PTD-Compact) by Nynxz (Apache-2.0), which builds on [ComfyUI-PTD](https://github.com/vita-epfl/PTD-inference) by EPFL VITA Lab (Apache-2.0; paper "Parametric Trajectory Distillation for Few-Step Video Generation", arXiv 2610.11498) and on ComfyUI-MiniMaxH3-PTD by taobao-content-ai (Apache-2.0); licence notice in `modules/sampling/ptd`;
 - NVIDIA and the NVIDIA Video Effects SDK (`nvidia-vfx`, NVIDIA's own licence) — RTX Super Resolution in Upscale and RTX Frame Generation in Interpolate run on it; FunUI installs it only when you switch the add-on on;
 - NVIDIA and the authors of Sol-H3 (arXiv 2609.35110, NVlabs/Sana, Apache-2.0) — the MXFP8 loader option (which blocks stay in bf16) and the Two stages setting (the cross-resolution schedule) follow its recipe;
+- Tero Karras, Miika Aittala, Timo Aila and Samuli Laine (Elucidating the Design Space of Diffusion-Based Generative Models, arXiv 2206.00364) — the churn idea (fresh noise added during sampling) behind the Turbo sampler's fresh-noise step;
 - Don Mitchell and Arun Netravali (the cubic filter family, SIGGRAPH 1988), Edwin Catmull and Raphael Rom (the Catmull-Rom spline, 1974) and Claude Duchon (the Lanczos window, 1979) — the Mitchell, Catmull-Rom and Lanczos choices when Two stages resizes the latent;
 - NVIDIA and the authors of Open-H3-Regenerate-2K (NVlabs/Sana) — the refiner LoRA's published layout, sampling recipe (a second stage at the full size) and strength convention that the LoRA Loader's Diffusers-layout reading and its "Apply to stage" option are built for; and Kijai, whose published conversion notes (fused attention, order of the gated MLP halves) were used to check our translation;
 - LBH-123-AI ([Comfyui_Minimax_h3_latent_Upscaler](https://github.com/LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler)) and TripleHeadedMonkey ([ComfyUI-MiniMaxH3_LatentUpscaler](https://github.com/TripleHeadedMonkey/ComfyUI-MiniMaxH3_LatentUpscaler)) — the H3 latent upscaler network (rebuilt from the published weights, not copied) and its use in Second pass and as the Two stages enlarge option;

modules/loaders/lora/keys.py

@@ -72,7 +72,8 @@ def from_diffusers(lora):
         if None in heights:
             continue
         start = sum(heights[: _QKV.index(part)])
-        extra[mod] = (f"diffusion_model.{base}.weight", (0, start, sizes[mod]))
+        plain = re.sub(r"^(?:base_model\.model\.|model\.diffusion_model\.|diffusion_model\.)+", "", base)       # the file's own wrapper is not part of the model's name
+        extra[mod] = (f"diffusion_model.{plain}.weight", (0, start, sizes[mod]))
     return out, extra
 
 
@@ -84,7 +85,9 @@ def _rekey(lora, prefix, replacement):
 
 def _misfit(sd, key, adapter):
     """True when a plain LoRA pair cannot make the weight it names (both dimensions, not the count)."""
-    w, target = getattr(adapter, "weights", None), sd.get(key)
+    w = getattr(adapter, "weights", None)
+    window = key[1] if isinstance(key, tuple) else None            # (dim, start, length): a LoRA aimed at some rows of a fused weight (Q, K, V of qkv_proj)
+    target = sd.get(key[0] if isinstance(key, tuple) else key)
     if not w or len(w) < 2 or target is None or target.dim() < 2:
         return False
     if (len(w) > 3 and w[3] is not None) or (len(w) > 5 and w[5] is not None):   # locon mid / padded target: not modelled
@@ -93,6 +96,9 @@ def _misfit(sd, key, adapter):
         got = (int(w[0].flatten(start_dim=1).shape[0]), int(w[1].flatten(start_dim=1).shape[-1]))
     except Exception:  # noqa: BLE001 -- not a plain pair (LoKr, LoHa, diff): left to ComfyUI
         return False
+    if window is not None:
+        dim, start, length = window
+        return got != (int(length), int(target.shape[1:].numel())) or start + length > int(target.shape[dim])
     return got != (int(target.shape[0]), int(target.shape[1:].numel()))
 
 

modules/loaders/tests/test_diffusers_h3_lora.py

@@ -108,3 +108,24 @@ def test_an_ltx_lora_with_transformer_blocks_names_is_left_as_it_is(tiny_ltx, mo
             name.replace(".weight", ".lora_B.weight"): torch.randn(out_dim, RANK)}
     patches, _how, dropped = keys.match(tiny_ltx.patcher, None, lora)
     assert len(patches) == 1 and dropped == 0
+
+
+@pytest.mark.parametrize("wrapper", ["diffusion_model.", "base_model.model.", "model.diffusion_model."])
+def test_a_wrapped_diffusers_file_still_reaches_the_fused_qkv_rows(patcher, wrapper):
+    from modules.loaders.lora.keys import match
+    lora, want = _file("transformer_blocks.1", torch.Generator().manual_seed(0))
+    lora = {wrapper + k[len("transformer."):]: v for k, v in lora.items()}
+    patches, how, dropped = match(patcher, None, lora)
+    clone = patcher.clone()
+    clone.add_patches(patches, 1.0)
+    got = _delta(clone, "diffusion_model.blocks.1.attn.qkv_proj.weight", (192, 64))
+    assert torch.allclose(got, torch.cat([want["attn.to_q"], want["attn.to_k"], want["attn.to_v"]], 0), atol=1e-5), (how, len(patches))
+
+
+def test_a_qkv_pair_of_the_wrong_width_is_dropped_not_merged_scrambled(patcher):
+    from modules.loaders.lora.keys import match
+    lora, _ = _file("transformer_blocks.1", torch.Generator().manual_seed(0))
+    k = "transformer.transformer_blocks.1.attn.to_q.lora_B.weight"
+    lora[k] = torch.randn(96, RANK)                                   # 96 rows where this model has 64 per part
+    patches, how, dropped = match(patcher, None, lora)
+    assert dropped >= 1 and len(patches) == 5