Allow stand-in with VACE

Don't know how well it can actually work, has some effect

wanvideo/modules/model.py

@@ -19,17 +19,18 @@ except:
 
 from .attention import attention
 import numpy as np
-__all__ = ['WanModel']
 
 from tqdm import tqdm
 import gc
-from comfy import model_management as mm
+
 from ...utils import log, get_module_memory_mb
 from ...cache_methods.cache_methods import TeaCacheState, MagCacheState, EasyCacheState, relative_l1_distance
 from ...multitalk.multitalk import get_attn_map_with_target
 from ...echoshot.echoshot import rope_apply_z, rope_apply_c, rope_apply_echoshot
 
-from comfy.model_management import get_torch_device, get_autocast_device, get_offload_stream
+__all__ = ['WanModel']
+
+from comfy import model_management as mm
 from comfy.ldm.flux.math import apply_rope as apply_rope_comfy
 
 def apply_rope_comfy_chunked(xq, xk, freqs_cis, num_chunks=4):
@@ -156,7 +157,7 @@ def rope_params(max_seq_len, dim, theta=10000, L_test=25, k=0):
     freqs = torch.polar(torch.ones_like(freqs), freqs)
     return freqs
 
-@torch.autocast(device_type=get_autocast_device(get_torch_device()), enabled=False)
+@torch.autocast(device_type=mm.get_autocast_device(mm.get_torch_device()), enabled=False)
 @torch.compiler.disable()
 def rope_apply(x, grid_sizes, freqs):
     n, c = x.size(2), x.size(3) // 2
@@ -739,7 +740,7 @@ class WanAttentionBlock(nn.Module):
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.get_mod(e.to(x.device))
         input_x = self.modulate(self.norm1(x), shift_msa, scale_msa)
 
-        if e_ip is not None:
+        if x_ip is not None:
             shift_msa_ip, scale_msa_ip, gate_msa_ip, shift_mlp_ip, scale_mlp_ip, gate_mlp_ip = self.get_mod(e_ip.to(x.device))
             input_x_ip = self.modulate(self.norm1(x_ip), shift_msa_ip, scale_msa_ip)
             self.self_attn.cond_size = input_x_ip.shape[1]
@@ -862,9 +863,8 @@ class WanAttentionBlock(nn.Module):
             x_ip = x_ip.addcmul(y_ip, gate_msa_ip)
             y_ip = self.ffn(torch.addcmul(shift_mlp_ip, self.norm2(x_ip), 1 + scale_mlp_ip))
             x_ip = x_ip.addcmul(y_ip, gate_mlp_ip)
-            return x, x_ip
-        
-        return x
+
+        return x, x_ip
 
     
     def cross_attn_ffn(self, x, context, grid_sizes, shift_mlp, scale_mlp, gate_mlp, clip_embed, 
@@ -986,14 +986,14 @@ class BaseWanAttentionBlock(WanAttentionBlock):
         self.block_id = block_id
 
     def forward(self, x, vace_hints=None, vace_context_scale=[1.0], **kwargs):
-        x = super().forward(x, **kwargs)
+        x, x_ip = super().forward(x, **kwargs)
         if vace_hints is None:
-            return x
+            return x, x_ip
         
         if self.block_id is not None:
             for i in range(len(vace_hints)):
                 x.add_(vace_hints[i][self.block_id].to(x.device), alpha=vace_context_scale[i])
-        return x
+        return x, x_ip
 
 class Head(nn.Module):
 
@@ -1405,7 +1405,7 @@ class WanModel(torch.nn.Module):
             else:
                 c_processed = current_c
                 
-            c_processed = block.forward(c_processed, **kwargs)
+            c_processed, _ = block.forward(c_processed, **kwargs)
             
             # Store skip connection
             c_skip = block.after_proj(c_processed)
@@ -1963,10 +1963,7 @@ class WanModel(torch.nn.Module):
                     if b in self.slg_blocks and is_uncond:
                         if self.slg_start_percent <= current_step_percentage <= self.slg_end_percent:
                             continue
-                if x_ip is None:
-                    x = block(x,**kwargs)
-                else:
-                    x, x_ip = block(x, x_ip=x_ip, **kwargs)
+                x, x_ip = block(x, x_ip=x_ip, **kwargs)
                 if self.block_swap_debug:
                     compute_end = time.perf_counter()
                     compute_time = compute_end - compute_start