From 366f740d28d972415c23b6c073b7b26b432daa30 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 30 Oct 2025 23:10:06 +0200
Subject: [PATCH 01/14] Update readme.md

---
 readme.md | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/readme.md b/readme.md
index 5f8add0..024219d 100644
--- a/readme.md
+++ b/readme.md
@@ -74,6 +74,13 @@ HuMo: https://github.com/Phantom-video/HuMo
 
 WanAnimate: https://github.com/Wan-Video/Wan2.2/tree/main/wan/modules/animate
 
+Lynx: https://github.com/bytedance/lynx
+
+
+Not exactly Wan model, but close enough to work with the code base:
+
+LongCat-Video: https://meituan-longcat.github.io/LongCat-Video/
+
 
 Examples:
 ---

From ce6e7b501d372d07da0b40faa3701527a8bde42a Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 30 Oct 2025 23:28:38 +0200
Subject: [PATCH 02/14] Fix unmerged LoRA application for certain LoRAs

---
 custom_linear.py | 6 +++---
 gguf/gguf.py     | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/custom_linear.py b/custom_linear.py
index a4bc6d3..6e8cce2 100644
--- a/custom_linear.py
+++ b/custom_linear.py
@@ -89,14 +89,14 @@ class CustomLinear(nn.Linear):
     def set_lora_diffs(self, lora_diffs, device=torch.device("cpu")):
         self.lora_diffs = []
         for i, diff in enumerate(lora_diffs):
-            if isinstance(diff, tuple):
+            if len(diff) > 1:
                 self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device))
                 self.register_buffer(f"lora_diff_{i}_1", diff[1].to(device))
                 setattr(self, f"lora_diff_{i}_2", diff[2])
                 self.lora_diffs.append((f"lora_diff_{i}_0", f"lora_diff_{i}_1", f"lora_diff_{i}_2"))
             else:
-                self.register_buffer(f"lora_diff_{i}", diff.to(device))
-                self.lora_diffs.append(f"lora_diff_{i}")
+                self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device))
+                self.lora_diffs.append(f"lora_diff_{i}_0")
 
     def forward(self, input):
         if self.bias is not None:
diff --git a/gguf/gguf.py b/gguf/gguf.py
index a7d860b..0cf760f 100644
--- a/gguf/gguf.py
+++ b/gguf/gguf.py
@@ -114,14 +114,14 @@ class GGUFLinear(nn.Linear):
     def set_lora_diffs(self, lora_diffs, device=torch.device("cpu")):
         self.lora_diffs = []
         for i, diff in enumerate(lora_diffs):
-            if isinstance(diff, tuple):
+            if len(diff) > 1:
                 self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device))
                 self.register_buffer(f"lora_diff_{i}_1", diff[1].to(device))
                 setattr(self, f"lora_diff_{i}_2", diff[2])
                 self.lora_diffs.append((f"lora_diff_{i}_0", f"lora_diff_{i}_1", f"lora_diff_{i}_2"))
             else:
-                self.register_buffer(f"lora_diff_{i}", diff.to(device))
-                self.lora_diffs.append(f"lora_diff_{i}")
+                self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device))
+                self.lora_diffs.append(f"lora_diff_{i}_0")
 
     def apply_lora(self, weight):
         for lora_diff_names, lora_strength in zip(self.lora_diffs, self.lora_strengths):

From 5da8a6b169e91b8a439530d8e57a3af46d33c8c2 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 31 Oct 2025 16:50:00 +0200
Subject: [PATCH 03/14] Fix MultiTalk on some models

---
 multitalk/multitalk.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/multitalk/multitalk.py b/multitalk/multitalk.py
index 220bb98..d017278 100644
--- a/multitalk/multitalk.py
+++ b/multitalk/multitalk.py
@@ -198,7 +198,7 @@ class AudioProjModel(nn.Module):
         context_tokens = self.proj3(audio_embeds_c).reshape(batch_size_c*N_t, self.context_tokens, self.output_dim)
 
         # normalization and reshape
-        context_tokens = self.norm(context_tokens)
+        context_tokens = self.norm(context_tokens.to(self.norm.weight.dtype)).to(context_tokens.dtype)
         context_tokens = rearrange(context_tokens, "(bz f) m c -> bz f m c", f=video_length)
 
         return context_tokens

From 5f4020b12d12983eba73ceec59755b56c396ab5d Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 31 Oct 2025 17:24:24 +0200
Subject: [PATCH 04/14] Fix a possible issue with Ovi audio model loading

---
 nodes_sampler.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/nodes_sampler.py b/nodes_sampler.py
index 6bf72b8..bcc7dd7 100644
--- a/nodes_sampler.py
+++ b/nodes_sampler.py
@@ -189,6 +189,11 @@ class WanVideoSampler:
         vae_upscale_factor = 16 if is_5b else 8
 
         # Load weights
+        if transformer.audio_model is not None:
+            for block in transformer.blocks:
+                if hasattr(block, 'audio_block'):
+                    block.audio_block = None
+
         if not transformer.patched_linear and patcher.model["sd"] is not None and len(patcher.patches) != 0:
             transformer = _replace_linear(transformer, dtype, patcher.model["sd"])
             transformer.patched_linear = True

From 393fe78ec2d5a515ce4ff794067ea4d72c830643 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 31 Oct 2025 23:39:55 +0200
Subject: [PATCH 05/14] Update model.py

---
 wanvideo/modules/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index fc03124..dc8baca 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -1265,7 +1265,7 @@ class WanAttentionBlock(nn.Module):
                 x = x.to(input_dtype)
                 # MultiTalk
                 if multitalk_audio_embedding is not None and not isinstance(self, VaceWanAttentionBlock):
-                    x_audio = self.audio_cross_attn(self.norm_x(x), encoder_hidden_states=multitalk_audio_embedding,
+                    x_audio = self.audio_cross_attn(self.norm_x(x.to(self.norm_x.weight.dtype)).to(input_dtype), encoder_hidden_states=multitalk_audio_embedding,
                                                 shape=grid_sizes[0], x_ref_attn_map=x_ref_attn_map, human_num=human_num)
                     x = x.add(x_audio, alpha=audio_scale)
 

From 5eae7087fa12b061b042172df9df97dd6d6d7ad4 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 2 Nov 2025 01:24:50 +0200
Subject: [PATCH 06/14] Fix for S2V

---
 wanvideo/modules/model.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index dc8baca..2f12f78 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -29,14 +29,14 @@ from comfy import model_management as mm
 __all__ = ['WanModel']
 
 class AdaLayerNorm(nn.Module):
-    def __init__(self, embedding_dim, output_dim=None, norm_elementwise_affine=False, norm_eps=1e-5, dtype=None, device=None, operations=None):
+    def __init__(self, embedding_dim, output_dim=None, norm_elementwise_affine=False, norm_eps=1e-5):
         super().__init__()
 
         output_dim = output_dim or embedding_dim * 2
 
         self.silu = nn.SiLU()
-        self.linear = operations.Linear(embedding_dim, output_dim, dtype=dtype, device=device)
-        self.norm = operations.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine, dtype=dtype, device=device)
+        self.linear = nn.Linear(embedding_dim, output_dim)
+        self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)
 
     def forward(self, x, temb):
         temb = self.linear(self.silu(temb))
@@ -1613,7 +1613,7 @@ class AudioInjector_WAN(nn.Module):
         if enable_adain:
             self.injector_adain_layers = nn.ModuleList([
                 AdaLayerNorm(
-                    output_dim=dim * 2, embedding_dim=adain_dim, chunk_dim=1)
+                    output_dim=dim * 2, embedding_dim=adain_dim)
                 for _ in range(audio_injector_id)
             ])
             if need_adain_ont:

From 75109fdb79ce56e46b5a2bea85af3249c9e7f68a Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 3 Nov 2025 10:12:05 +0200
Subject: [PATCH 07/14] Fix custom sigmas with euler

---
 wanvideo/schedulers/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/wanvideo/schedulers/__init__.py b/wanvideo/schedulers/__init__.py
index e99fe83..7126fe6 100644
--- a/wanvideo/schedulers/__init__.py
+++ b/wanvideo/schedulers/__init__.py
@@ -60,7 +60,7 @@ def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transfo
             if flowedit_args: #seems to work better
                 timesteps, _ = retrieve_timesteps(sample_scheduler, device=device, sigmas=get_sampling_sigmas(steps, shift))
             else:
-                sample_scheduler.set_timesteps(steps, device=device, sigmas=sigmas)
+                sample_scheduler.set_timesteps(steps, device=device, sigmas=sigmas[:-1].tolist() if sigmas is not None else None)
     elif 'dpm' in scheduler:
         if 'sde' in scheduler:
             algorithm_type = "sde-dpmsolver++"

From 0d0d28569a59a92b93217b543125646a213956ce Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 3 Nov 2025 10:29:10 +0200
Subject: [PATCH 08/14] Fix cases where text encoder isn't used (eg. Minimax
 remover)

---
 wanvideo/modules/model.py | 85 ++++++++++++++++++++-------------------
 1 file changed, 44 insertions(+), 41 deletions(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index 2f12f78..57db863 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -1279,39 +1279,39 @@ class WanAttentionBlock(nn.Module):
                     x = self.audio_cross_attn_wrapper(x, humo_audio_input, grid_sizes, humo_audio_scale)
 
 
-            # ffn
-            if self.rope_func == "comfy_chunked":
-                x_ffn = self.ffn_chunked(x, shift_mlp, scale_mlp)
-            else:
-                if zero_timestep:
-                    norm2_x = self.norm2(x)
-                    parts = []
-                    for i in range(2):
-                        parts.append(norm2_x[:, self.seg_idx[i]:self.seg_idx[i + 1]] *
-                                    (1 + scale_mlp[:, i:i + 1]) + shift_mlp[:, i:i + 1])
-                    norm2_x = torch.cat(parts, dim=1)
-                    x_ffn = self.ffn(norm2_x)
-                else:
-                    if not is_longcat:
-                        mod_x = torch.addcmul(shift_mlp, self.norm2(x.to(shift_mlp.dtype)), 1 + scale_mlp)
-                    else:
-                        mod_x = torch.addcmul(shift_mlp, self.norm2(x.view(B, -1, N//T, C).float()), 1 + scale_mlp).view(B, -1, C)
-                    x_ffn = self.ffn(mod_x.to(input_dtype))
-                del shift_mlp, scale_mlp
-            
-            # gate_mlp
+        # ffn
+        if self.rope_func == "comfy_chunked":
+            x_ffn = self.ffn_chunked(x, shift_mlp, scale_mlp)
+        else:
             if zero_timestep:
-                z = []
+                norm2_x = self.norm2(x)
+                parts = []
                 for i in range(2):
-                    z.append(x_ffn[:, self.seg_idx[i]:self.seg_idx[i + 1]] * gate_mlp[:, i:i + 1])
-                x_ffn = torch.cat(z, dim=1)
-                x = x.add(x_ffn)
+                    parts.append(norm2_x[:, self.seg_idx[i]:self.seg_idx[i + 1]] *
+                                (1 + scale_mlp[:, i:i + 1]) + shift_mlp[:, i:i + 1])
+                norm2_x = torch.cat(parts, dim=1)
+                x_ffn = self.ffn(norm2_x)
             else:
                 if not is_longcat:
-                    x = x.addcmul(x_ffn.to(gate_mlp.dtype), gate_mlp).to(input_dtype)
+                    mod_x = torch.addcmul(shift_mlp, self.norm2(x.to(shift_mlp.dtype)), 1 + scale_mlp)
                 else:
-                    x = x + (gate_mlp * x_ffn.view(B, -1, N//T, C).float()).to(input_dtype).view(B, -1, C)
-            del gate_mlp
+                    mod_x = torch.addcmul(shift_mlp, self.norm2(x.view(B, -1, N//T, C).float()), 1 + scale_mlp).view(B, -1, C)
+                x_ffn = self.ffn(mod_x.to(input_dtype))
+            del shift_mlp, scale_mlp
+        
+        # gate_mlp
+        if zero_timestep:
+            z = []
+            for i in range(2):
+                z.append(x_ffn[:, self.seg_idx[i]:self.seg_idx[i + 1]] * gate_mlp[:, i:i + 1])
+            x_ffn = torch.cat(z, dim=1)
+            x = x.add(x_ffn)
+        else:
+            if not is_longcat:
+                x = x.addcmul(x_ffn.to(gate_mlp.dtype), gate_mlp).to(input_dtype)
+            else:
+                x = x + (gate_mlp * x_ffn.view(B, -1, N//T, C).float()).to(input_dtype).view(B, -1, C)
+        del gate_mlp
 
         if x_ip is not None: #stand-in
             x_ip = x_ip.addcmul(y_ip, gate_msa_ip)
@@ -2566,7 +2566,17 @@ class WanModel(torch.nn.Module):
             
             e = e.to(self.offload_device, non_blocking=self.use_non_blocking)
 
-        
+        # clip vision embedding
+        clip_embed = None
+        if clip_fea is not None and hasattr(self, "img_emb"):
+            clip_fea = clip_fea.to(self.main_device)
+            if self.offload_img_emb:
+                self.img_emb.to(self.main_device)
+            clip_embed = self.img_emb(clip_fea)  # bs x 257 x dim
+            #context = torch.concat([context_clip, context], dim=1)
+            if self.offload_img_emb:
+                self.img_emb.to(self.offload_device, non_blocking=self.use_non_blocking)
+
         #context (text embedding)
         if hasattr(self, "text_embedding") and context != []:
             text_embed_dtype = self.text_embedding[0].weight.dtype
@@ -2606,18 +2616,13 @@ class WanModel(torch.nn.Module):
             
             if self.offload_txt_emb:
                 self.text_embedding.to(self.offload_device, non_blocking=self.use_non_blocking)
+
+            seq_chunks = max(context.shape[0], clip_embed.shape[0] if clip_embed is not None else 0)
+            chunked_self_attention = seq_chunks > 1 and current_step in self.video_attention_split_steps
         else:
             context = None
-
-        clip_embed = None
-        if clip_fea is not None and hasattr(self, "img_emb"):
-            clip_fea = clip_fea.to(self.main_device)
-            if self.offload_img_emb:
-                self.img_emb.to(self.main_device)
-            clip_embed = self.img_emb(clip_fea)  # bs x 257 x dim
-            #context = torch.concat([context_clip, context], dim=1)
-            if self.offload_img_emb:
-                self.img_emb.to(self.offload_device, non_blocking=self.use_non_blocking)
+            chunked_self_attention = False
+            seq_chunks = 0
 
         # MultiTalk
         if multitalk_audio is not None:
@@ -2795,8 +2800,6 @@ class WanModel(torch.nn.Module):
                     dwpose_emb = rearrange(unianim_data['dwpose'], 'b c f h w -> b (f h w) c').contiguous()
                     x.add_(dwpose_emb, alpha=unianim_data['strength'])
 
-            seq_chunks = max(context.shape[0], clip_embed.shape[0] if clip_embed is not None else 0)
-            chunked_self_attention = seq_chunks > 1 and current_step in self.video_attention_split_steps
             # arguments
             kwargs = dict(
                 e=e0,

From 8ce6916d72132f9d9866f35db8670c0d3cf74602 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 3 Nov 2025 10:36:55 +0200
Subject: [PATCH 09/14] Fix for some cases of using comfy_chunked rope

---
 wanvideo/modules/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index 57db863..e089733 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -1003,7 +1003,7 @@ class WanAttentionBlock(nn.Module):
             return torch.addcmul(shift_msa, norm_x, 1 + scale_msa)
     
     def ffn_chunked(self, x, shift_mlp, scale_mlp, num_chunks=4):
-        modulated_input = torch.addcmul(shift_mlp, self.norm2(x), 1 + scale_mlp)
+        modulated_input = torch.addcmul(shift_mlp, self.norm2(x.to(shift_mlp.dtype)), 1 + scale_mlp).to(x.dtype)
         
         result = torch.empty_like(x)
         seq_len = modulated_input.shape[1]

From 9fa4140159533f831af5de738143abe1a4e22a61 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 4 Nov 2025 01:34:51 +0200
Subject: [PATCH 10/14] Make lora torch.compile optional for unmerged lora
 application

This change has caused issues especially with LoRAs that have dynamic rank. Will now be disabled by default, to allow full graph with unmerged LoRAs the option to allow compile is available in the Torch Compile Settings -node
---
 custom_linear.py       | 65 ++++++++++++++++++++++++++----------------
 gguf/gguf.py           | 33 ++++++++++++++-------
 nodes_model_loading.py | 11 ++++---
 nodes_sampler.py       | 10 ++++---
 skyreels/nodes.py      |  2 +-
 5 files changed, 78 insertions(+), 43 deletions(-)

diff --git a/custom_linear.py b/custom_linear.py
index 6e8cce2..9e70ee8 100644
--- a/custom_linear.py
+++ b/custom_linear.py
@@ -3,15 +3,20 @@ import torch.nn as nn
 from accelerate import init_empty_weights
 
 #based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/quantizers/gguf/utils.py
-def _replace_linear(model, compute_dtype, state_dict, prefix="", patches=None, scale_weights=None):
+def _replace_linear(model, compute_dtype, state_dict, prefix="", patches=None, scale_weights=None, compile_args=None):
    
     has_children = list(model.children())
     if not has_children:
         return
+    
+    allow_compile = False
+
     for name, module in model.named_children():
+        if compile_args is not None:
+            allow_compile = compile_args.get("allow_unmerged_lora_compile", False)
         module_prefix = prefix + name + "."
         module_prefix = module_prefix.replace("_orig_mod.", "")
-        _replace_linear(module, compute_dtype, state_dict, module_prefix, patches, scale_weights)
+        _replace_linear(module, compute_dtype, state_dict, module_prefix, patches, scale_weights, compile_args)
 
         if isinstance(module, nn.Linear) and "loras" not in module_prefix:
             in_features = state_dict[module_prefix + "weight"].shape[1]
@@ -25,7 +30,8 @@ def _replace_linear(model, compute_dtype, state_dict, prefix="", patches=None, s
                     out_features,
                     module.bias is not None,
                     compute_dtype=compute_dtype,
-                    scale_weight=scale_weights.get(scale_key) if scale_weights else None
+                    scale_weight=scale_weights.get(scale_key) if scale_weights else None,
+                    allow_compile=allow_compile
                 )
             model._modules[name].source_cls = type(module)
             model._modules[name].requires_grad_(False)
@@ -77,7 +83,8 @@ class CustomLinear(nn.Linear):
         bias=False,
         compute_dtype=None,
         device=None,
-        scale_weight=None
+        scale_weight=None,
+        allow_compile=False
     ) -> None:
         super().__init__(in_features, out_features, bias, device)
         self.compute_dtype = compute_dtype
@@ -85,6 +92,10 @@ class CustomLinear(nn.Linear):
         self.step = 0
         self.scale_weight = scale_weight
         self.lora_strengths = []
+        self.allow_compile = allow_compile
+
+        if not allow_compile:
+            self._get_weight_with_lora = torch.compiler.disable()(self._get_weight_with_lora)
     
     def set_lora_diffs(self, lora_diffs, device=torch.device("cpu")):
         self.lora_diffs = []
@@ -98,25 +109,11 @@ class CustomLinear(nn.Linear):
                 self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device))
                 self.lora_diffs.append(f"lora_diff_{i}_0")
 
-    def forward(self, input):
-        if self.bias is not None:
-            bias = self.bias.to(input)
-        else:
-            bias = None
-        weight = self.weight.to(input)
-
-        if self.scale_weight is not None:
-            if weight.numel() < input.numel():
-                weight = weight * self.scale_weight
-            else:
-                input = input * self.scale_weight
-
-        if hasattr(self, f"lora_diff_0_0"):
-            weight = self.apply_lora(weight).to(self.compute_dtype)
-
-        return torch.nn.functional.linear(input, weight, bias)
-
-    def apply_lora(self, weight):
+    def _get_weight_with_lora(self, weight):
+        """Apply LoRA outside compiled region"""
+        if not hasattr(self, "lora_diff_0_0"):
+            return weight
+        
         for lora_diff_names, lora_strength in zip(self.lora_diffs, self.lora_strengths):
             if isinstance(lora_strength, list):
                 lora_strength = lora_strength[self.step]
@@ -131,7 +128,7 @@ class CustomLinear(nn.Linear):
                 patch_diff = torch.mm(
                     lora_diff_0.flatten(start_dim=1),
                     lora_diff_1.flatten(start_dim=1)
-                ).reshape(weight.shape) + 0
+                ).reshape(weight.shape)
                 alpha = lora_diff_2 / lora_diff_1.shape[0] if lora_diff_2 is not None else 1.0
                 scale = lora_strength * alpha
                 weight = weight.add(patch_diff, alpha=scale)
@@ -139,6 +136,26 @@ class CustomLinear(nn.Linear):
                 lora_diff = getattr(self, lora_diff_names)
                 weight = weight.add(lora_diff, alpha=lora_strength)
         return weight
+
+    def forward(self, input):
+        if self.bias is not None:
+            bias = self.bias.to(input)
+        else:
+            bias = None
+        weight = self.weight.to(input)
+
+        if self.scale_weight is not None:
+            if weight.numel() < input.numel():
+                weight = weight * self.scale_weight
+            else:
+                input = input * self.scale_weight
+
+        weight = self._get_weight_with_lora(weight)
+        
+        if self.compute_dtype is not None:
+            weight = weight.to(self.compute_dtype)
+
+        return torch.nn.functional.linear(input, weight, bias)
     
 def remove_lora_from_module(module):
     for name, submodule in module.named_modules():
diff --git a/gguf/gguf.py b/gguf/gguf.py
index 0cf760f..177e1e9 100644
--- a/gguf/gguf.py
+++ b/gguf/gguf.py
@@ -19,7 +19,7 @@ def load_gguf(model_path):
     return parsed_parameters, reader
 
 #based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/quantizers/gguf/utils.py
-def _replace_with_gguf_linear(model, compute_dtype, state_dict, prefix="", modules_to_not_convert=[], patches=None):
+def _replace_with_gguf_linear(model, compute_dtype, state_dict, prefix="", modules_to_not_convert=[], patches=None, compile_args=None):
     def _should_convert_to_gguf(state_dict, prefix):
         weight_key = prefix + "weight"
         return weight_key in state_dict and isinstance(state_dict[weight_key], GGUFParameter)
@@ -27,10 +27,14 @@ def _replace_with_gguf_linear(model, compute_dtype, state_dict, prefix="", modul
     has_children = list(model.children())
     if not has_children:
         return
+    
+    allow_compile = False
 
     for name, module in model.named_children():
+        if compile_args is not None:
+            allow_compile = compile_args.get("allow_unmerged_lora_compile", False)
         module_prefix = prefix + name + "."
-        _replace_with_gguf_linear(module, compute_dtype, state_dict, module_prefix, modules_to_not_convert, patches)
+        _replace_with_gguf_linear(module, compute_dtype, state_dict, module_prefix, modules_to_not_convert, patches, compile_args)
 
         if (
             isinstance(module, nn.Linear)
@@ -46,7 +50,8 @@ def _replace_with_gguf_linear(model, compute_dtype, state_dict, prefix="", modul
                     in_features,
                     out_features,
                     module.bias is not None,
-                    compute_dtype=compute_dtype
+                    compute_dtype=compute_dtype,
+                    allow_compile=allow_compile
                 )
             
             model._modules[name].source_cls = type(module)
@@ -95,19 +100,23 @@ class GGUFLinear(nn.Linear):
         bias=False,
         compute_dtype=None,
         device=None,
+        allow_compile=False
     ) -> None:
         super().__init__(in_features, out_features, bias, device)
         self.compute_dtype = compute_dtype
         self.lora_diffs = []
         self.lora_strengths = []
         self.step = 0
+        self.allow_compile = allow_compile
+
+        if not allow_compile:
+            self._get_weight_with_lora = torch.compiler.disable()(self._get_weight_with_lora)
 
     def forward(self, inputs):
         weight = dequantize_gguf_tensor(self.weight).to(self.compute_dtype)
         bias = self.bias.to(self.compute_dtype) if self.bias is not None else None
 
-        if hasattr(self, f"lora_diff_0_0"):
-            weight = self.apply_lora(weight).to(self.compute_dtype)
+        weight = self._get_weight_with_lora(weight)#.to(self.compute_dtype)
 
         return torch.nn.functional.linear(inputs, weight, bias)
     
@@ -115,15 +124,19 @@ class GGUFLinear(nn.Linear):
         self.lora_diffs = []
         for i, diff in enumerate(lora_diffs):
             if len(diff) > 1:
-                self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device))
-                self.register_buffer(f"lora_diff_{i}_1", diff[1].to(device))
+                self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device, self.compute_dtype))
+                self.register_buffer(f"lora_diff_{i}_1", diff[1].to(device, self.compute_dtype))
                 setattr(self, f"lora_diff_{i}_2", diff[2])
                 self.lora_diffs.append((f"lora_diff_{i}_0", f"lora_diff_{i}_1", f"lora_diff_{i}_2"))
             else:
-                self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device))
+                self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device, self.compute_dtype))
                 self.lora_diffs.append(f"lora_diff_{i}_0")
 
-    def apply_lora(self, weight):
+    def _get_weight_with_lora(self, weight):
+        """Apply LoRA outside compiled region"""
+        if not hasattr(self, "lora_diff_0_0"):
+            return weight
+        
         for lora_diff_names, lora_strength in zip(self.lora_diffs, self.lora_strengths):
             if isinstance(lora_strength, list):
                 lora_strength = lora_strength[self.step]
@@ -138,7 +151,7 @@ class GGUFLinear(nn.Linear):
                 patch_diff = torch.mm(
                     lora_diff_0.flatten(start_dim=1),
                     lora_diff_1.flatten(start_dim=1)
-                ).reshape(weight.shape) + 0
+                ).reshape(weight.shape)
                 alpha = lora_diff_2 / lora_diff_1.shape[0] if lora_diff_2 is not None else 1.0
                 scale = lora_strength * alpha
                 weight = weight.add(patch_diff, alpha=scale)
diff --git a/nodes_model_loading.py b/nodes_model_loading.py
index a35fd9b..d644770 100644
--- a/nodes_model_loading.py
+++ b/nodes_model_loading.py
@@ -330,6 +330,7 @@ class WanVideoTorchCompileSettings:
             "optional": {
                 "dynamo_recompile_limit": ("INT", {"default": 128, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.recompile_limit"}),
                 "force_parameter_static_shapes": ("BOOLEAN", {"default": False, "tooltip": "torch._dynamo.config.force_parameter_static_shapes"}),
+                "allow_unmerged_lora_compile": ("BOOLEAN", {"default": False, "tooltip": "Allow LoRA application to be compiled with torch.compile to avoid graph breaks, causes issues with some LoRAs, mostly dynamic ones"}),
             },
         }
     RETURN_TYPES = ("WANCOMPILEARGS",)
@@ -338,7 +339,8 @@ class WanVideoTorchCompileSettings:
     CATEGORY = "WanVideoWrapper"
     DESCRIPTION = "torch.compile settings, when connected to the model loader, torch.compile of the selected layers is attempted. Requires Triton and torch > 2.7.0 is recommended"
 
-    def set_args(self, backend, fullgraph, mode, dynamic, dynamo_cache_size_limit, compile_transformer_blocks_only, dynamo_recompile_limit=128, force_parameter_static_shapes=True):
+    def set_args(self, backend, fullgraph, mode, dynamic, dynamo_cache_size_limit, compile_transformer_blocks_only, dynamo_recompile_limit=128,
+                 force_parameter_static_shapes=True, allow_unmerged_lora_compile=False):
 
         compile_args = {
             "backend": backend,
@@ -349,6 +351,7 @@ class WanVideoTorchCompileSettings:
             "dynamo_recompile_limit": dynamo_recompile_limit,
             "compile_transformer_blocks_only": compile_transformer_blocks_only,
             "force_parameter_static_shapes": force_parameter_static_shapes,
+            "allow_unmerged_lora_compile": allow_unmerged_lora_compile,
         }
 
         return (compile_args, )
@@ -782,7 +785,7 @@ def rename_fuser_block(name):
     return new_name
 
 def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None, 
-                 transformer_load_device=None, block_swap_args=None, gguf=False, reader=None, patcher=None):
+                 transformer_load_device=None, block_swap_args=None, gguf=False, reader=None, patcher=None, compile_args=None):
     params_to_keep = {"time_in", "patch_embedding", "time_", "modulation", "text_embedding", 
                       "adapter", "add", "ref_conv", "casual_audio_encoder", "cond_encoder", "frame_packer", "audio_proj_glob", "face_encoder", "fuser_block"}
     param_count = sum(1 for _ in transformer.named_parameters())
@@ -837,7 +840,7 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None,
 
         if not getattr(transformer, "gguf_patched", False):
             transformer = _replace_with_gguf_linear(
-                transformer, base_dtype, sd, patches=patcher.patches
+                transformer, base_dtype, sd, patches=patcher.patches, compile_args=compile_args
             )
             transformer.gguf_patched = True
     else:
@@ -1534,7 +1537,7 @@ class WanVideoModelLoader:
                 transformer.patched_linear = False
                 sd = None
             elif "scaled" in quantization or lora is not None:
-                transformer = _replace_linear(transformer, base_dtype, sd, scale_weights=scale_weights)
+                transformer = _replace_linear(transformer, base_dtype, sd, scale_weights=scale_weights, compile_args=compile_args)
                 transformer.patched_linear = True
 
         if "fast" in quantization:
diff --git a/nodes_sampler.py b/nodes_sampler.py
index bcc7dd7..5726335 100644
--- a/nodes_sampler.py
+++ b/nodes_sampler.py
@@ -194,14 +194,16 @@ class WanVideoSampler:
                 if hasattr(block, 'audio_block'):
                     block.audio_block = None
 
-        if not transformer.patched_linear and patcher.model["sd"] is not None and len(patcher.patches) != 0:
-            transformer = _replace_linear(transformer, dtype, patcher.model["sd"])
+        if not transformer.patched_linear and patcher.model["sd"] is not None and len(patcher.patches) != 0 and gguf_reader is None:
+            transformer = _replace_linear(transformer, dtype, patcher.model["sd"], model["compile_args"])
             transformer.patched_linear = True
         if patcher.model["sd"] is not None and gguf_reader is None:
-            load_weights(patcher.model.diffusion_model, patcher.model["sd"], weight_dtype, base_dtype=dtype, transformer_load_device=device, block_swap_args=block_swap_args)
+            load_weights(patcher.model.diffusion_model, patcher.model["sd"], weight_dtype, base_dtype=dtype, transformer_load_device=device, 
+                         block_swap_args=block_swap_args, compile_args=model["compile_args"])
 
         if gguf_reader is not None: #handle GGUF
-            load_weights(transformer, patcher.model["sd"], base_dtype=dtype, transformer_load_device=device, patcher=patcher, gguf=True, reader=gguf_reader, block_swap_args=block_swap_args)
+            load_weights(transformer, patcher.model["sd"], base_dtype=dtype, transformer_load_device=device, patcher=patcher, gguf=True,
+                         reader=gguf_reader, block_swap_args=block_swap_args, compile_args=model["compile_args"])
             set_lora_params_gguf(transformer, patcher.patches)
             transformer.patched_linear = True
         elif len(patcher.patches) != 0: #handle patched linear layers (unmerged loras, fp8 scaled)
diff --git a/skyreels/nodes.py b/skyreels/nodes.py
index e67aa25..690dec9 100644
--- a/skyreels/nodes.py
+++ b/skyreels/nodes.py
@@ -172,7 +172,7 @@ class WanVideoDiffusionForcingSampler:
 
         # Load weights
         if not transformer.patched_linear and patcher.model["sd"] is not None and len(patcher.patches) != 0:
-            transformer = _replace_linear(transformer, dtype, patcher.model["sd"])
+            transformer = _replace_linear(transformer, dtype, patcher.model["sd"], model["compile_args"])
             transformer.patched_linear = True
         if patcher.model["sd"] is not None and gguf_reader is None:
             load_weights(patcher.model.diffusion_model, patcher.model["sd"], weight_dtype, base_dtype=dtype, transformer_load_device=device, block_swap_args=block_swap_args)

From 509d6922f5fe760e4c5af8f4d5a06e137b484d53 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 4 Nov 2025 09:44:07 +0200
Subject: [PATCH 11/14] Update custom_linear.py

---
 custom_linear.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/custom_linear.py b/custom_linear.py
index 9e70ee8..541436a 100644
--- a/custom_linear.py
+++ b/custom_linear.py
@@ -101,12 +101,12 @@ class CustomLinear(nn.Linear):
         self.lora_diffs = []
         for i, diff in enumerate(lora_diffs):
             if len(diff) > 1:
-                self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device))
-                self.register_buffer(f"lora_diff_{i}_1", diff[1].to(device))
+                self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device, self.compute_dtype))
+                self.register_buffer(f"lora_diff_{i}_1", diff[1].to(device, self.compute_dtype))
                 setattr(self, f"lora_diff_{i}_2", diff[2])
                 self.lora_diffs.append((f"lora_diff_{i}_0", f"lora_diff_{i}_1", f"lora_diff_{i}_2"))
             else:
-                self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device))
+                self.register_buffer(f"lora_diff_{i}_0", diff[0].to(device, self.compute_dtype))
                 self.lora_diffs.append(f"lora_diff_{i}_0")
 
     def _get_weight_with_lora(self, weight):

From 9a588a42ec965521890b40b1fc3251a55f6e2259 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 4 Nov 2025 09:47:38 +0200
Subject: [PATCH 12/14] Fix some precision issues with unmerged lora

---
 custom_linear.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/custom_linear.py b/custom_linear.py
index 541436a..edd3379 100644
--- a/custom_linear.py
+++ b/custom_linear.py
@@ -151,9 +151,6 @@ class CustomLinear(nn.Linear):
                 input = input * self.scale_weight
 
         weight = self._get_weight_with_lora(weight)
-        
-        if self.compute_dtype is not None:
-            weight = weight.to(self.compute_dtype)
 
         return torch.nn.functional.linear(input, weight, bias)
     

From 8002d8a2f95aced3be5e49d8c9af2bc78ff97f9a Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 4 Nov 2025 09:58:40 +0200
Subject: [PATCH 13/14] This still needed for some reason too

---
 custom_linear.py | 2 +-
 gguf/gguf.py     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/custom_linear.py b/custom_linear.py
index edd3379..e6cd2bf 100644
--- a/custom_linear.py
+++ b/custom_linear.py
@@ -128,7 +128,7 @@ class CustomLinear(nn.Linear):
                 patch_diff = torch.mm(
                     lora_diff_0.flatten(start_dim=1),
                     lora_diff_1.flatten(start_dim=1)
-                ).reshape(weight.shape)
+                ).reshape(weight.shape) + 0
                 alpha = lora_diff_2 / lora_diff_1.shape[0] if lora_diff_2 is not None else 1.0
                 scale = lora_strength * alpha
                 weight = weight.add(patch_diff, alpha=scale)
diff --git a/gguf/gguf.py b/gguf/gguf.py
index 177e1e9..76dd87e 100644
--- a/gguf/gguf.py
+++ b/gguf/gguf.py
@@ -151,7 +151,7 @@ class GGUFLinear(nn.Linear):
                 patch_diff = torch.mm(
                     lora_diff_0.flatten(start_dim=1),
                     lora_diff_1.flatten(start_dim=1)
-                ).reshape(weight.shape)
+                ).reshape(weight.shape) + 0
                 alpha = lora_diff_2 / lora_diff_1.shape[0] if lora_diff_2 is not None else 1.0
                 scale = lora_strength * alpha
                 weight = weight.add(patch_diff, alpha=scale)

From 1d0516a2a97dd99d2f4ebc5eed2b0c8e5b48edf7 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 4 Nov 2025 10:26:37 +0200
Subject: [PATCH 14/14] Avoid graph break for LongCat

---
 wanvideo/modules/model.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index e089733..3358fc3 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -672,7 +672,7 @@ class WanT2VCrossAttention(WanSelfAttention):
 
         if is_longcat:
             if num_cond_latents is not None and num_cond_latents > 0:
-                num_cond_latents_thw = num_cond_latents * (s // grid_sizes[0][0])
+                num_cond_latents_thw = num_cond_latents * (s // num_latent_frames)
                 x = x[:, num_cond_latents_thw:]
             q = self.norm_q(self.q(x).view(b, -1, n, d))
         else:
@@ -1185,7 +1185,7 @@ class WanAttentionBlock(nn.Module):
             full_v = torch.cat([v, v_ip], dim=1)
             y = self.self_attn.forward(q, full_k, full_v, seq_lens)
         elif is_longcat and num_cond_latents is not None and num_cond_latents > 0:
-            num_cond_latents_thw = num_cond_latents * (N // grid_sizes[0][0])
+            num_cond_latents_thw = num_cond_latents * (N // num_latent_frames)
             # process the condition tokens
             x_cond = self.self_attn.forward(
                 q[:, :num_cond_latents_thw].contiguous(),