diff --git a/README.md b/README.md index 1650c714..88dd092a 100644 --- a/README.md +++ b/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; diff --git a/modules/loaders/lora/keys.py b/modules/loaders/lora/keys.py index 1e5b8e18..1a36d3cf 100644 --- a/modules/loaders/lora/keys.py +++ b/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())) diff --git a/modules/loaders/tests/test_diffusers_h3_lora.py b/modules/loaders/tests/test_diffusers_h3_lora.py index be05e9ba..076ab6b1 100644 --- a/modules/loaders/tests/test_diffusers_h3_lora.py +++ b/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