Merge branch 'main' into self-refine-video

I'm not 100% sure on this, I initially tested this when I noticed the original code doesn't, but it's described in the paper... now I see the original code has added it too so it seems to be the intended way to use it.

LongCat/layers.py

@@ -2,6 +2,10 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torch
 import math
+from einops import rearrange
+
+from ..wanvideo.modules.model import WanRMSNorm, attention
+from ..multitalk.multitalk import RotaryPositionalEmbedding1D, normalize_and_scale
 
 class FeedForwardSwiGLU(nn.Module):
     def __init__(
@@ -22,7 +26,7 @@ class FeedForwardSwiGLU(nn.Module):
 
     def forward(self, x):
         return self.w2(F.silu(self.w1(x)) * self.w3(x))
-    
+
 class TimestepEmbedder(nn.Module):
     """
     Embeds scalar timesteps into vector representations.
@@ -62,4 +66,147 @@ class TimestepEmbedder(nn.Module):
         if t_freq.dtype != dtype:
             t_freq = t_freq.to(dtype)
         t_emb = self.mlp(t_freq)
-        return t_emb
+        return t_emb
+
+
+class SingleStreamAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        encoder_hidden_states_dim: int,
+        num_heads: int,
+        qkv_bias: bool,
+        qk_norm: bool,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        eps: float = 1e-6,
+        class_range: int = 24,
+        class_interval: int = 4,
+        attention_mode: str = "sdpa",
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, "dim should be divisible by num_heads"
+        self.dim = dim
+        self.encoder_hidden_states_dim = encoder_hidden_states_dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_linear = nn.Linear(dim, dim, bias=qkv_bias)
+        self.q_norm = WanRMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.kv_linear = nn.Linear(encoder_hidden_states_dim, dim * 2, bias=qkv_bias)
+        self.k_norm = WanRMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()
+
+        self.attention_mode = attention_mode
+
+        # multitalk related params
+        self.class_interval = class_interval
+        self.class_range = class_range
+        self.rope_h1  = (0, self.class_interval)
+        self.rope_h2  = (self.class_range - self.class_interval, self.class_range)
+        self.rope_bak = int(self.class_range // 2)
+        self.rope_1d = RotaryPositionalEmbedding1D(self.head_dim)
+
+    def _process_cross_attn(self, x, cond, frames_num=None, x_ref_attn_map=None):
+
+        N_t = frames_num
+        out_dtype = x.dtype
+        x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
+
+        # get q for hidden_state
+        B, N, C = x.shape
+        q = self.q_linear(x)
+        q_shape = (B, N, self.num_heads, self.head_dim)
+        q = q.view(q_shape).permute((0, 2, 1, 3)) # [B, H, N, D]
+        q = self.q_norm(q.to(self.q_norm.weight.dtype)).to(q.dtype)
+
+        # multitalk with rope1d pe
+        if x_ref_attn_map is not None:
+            max_values = x_ref_attn_map.max(1).values[:, None, None] 
+            min_values = x_ref_attn_map.min(1).values[:, None, None] 
+            max_min_values = torch.cat([max_values, min_values], dim=2) 
+            human1_max_value, human1_min_value = max_min_values[0, :, 0].max(), max_min_values[0, :, 1].min()
+            human2_max_value, human2_min_value = max_min_values[1, :, 0].max(), max_min_values[1, :, 1].min()
+
+            human1 = normalize_and_scale(x_ref_attn_map[0], (human1_min_value, human1_max_value), (self.rope_h1[0], self.rope_h1[1]))
+            human2 = normalize_and_scale(x_ref_attn_map[1], (human2_min_value, human2_max_value), (self.rope_h2[0], self.rope_h2[1]))
+            back   = torch.full((x_ref_attn_map.size(1),), self.rope_bak, dtype=human1.dtype).to(human1.device)
+            max_indices = x_ref_attn_map.argmax(dim=0)
+            normalized_map = torch.stack([human1, human2, back], dim=1)
+            normalized_pos = normalized_map[range(x_ref_attn_map.size(1)), max_indices] 
+
+            q = rearrange(q, "(B N_t) H S C -> B H (N_t S) C", N_t=N_t)
+            q = self.rope_1d(q, normalized_pos)
+            q = rearrange(q, "B H (N_t S) C -> (B N_t) H S C", N_t=N_t)
+
+        # get kv from encoder_hidden_states
+        _, N_a, _ = cond.shape
+        encoder_kv = self.kv_linear(cond)
+        encoder_kv_shape = (B, N_a, 2, self.num_heads, self.head_dim)
+        encoder_kv = encoder_kv.view(encoder_kv_shape).permute((2, 0, 3, 1, 4))
+
+        encoder_k, encoder_v = encoder_kv.unbind(0)
+        encoder_k = self.k_norm(encoder_k.to(self.k_norm.weight.dtype)).to(encoder_k.dtype)
+
+
+        # multitalk with rope1d pe
+        if x_ref_attn_map is not None:
+            per_frame = torch.zeros(N_a, dtype=encoder_k.dtype).to(encoder_k.device)
+            per_frame[:per_frame.size(0)//2] = (self.rope_h1[0] + self.rope_h1[1]) / 2
+            per_frame[per_frame.size(0)//2:] = (self.rope_h2[0] + self.rope_h2[1]) / 2
+            encoder_pos = torch.concat([per_frame]*N_t, dim=0)
+            encoder_k = rearrange(encoder_k, "(B N_t) H S C -> B H (N_t S) C", N_t=N_t)
+            encoder_k = self.rope_1d(encoder_k, encoder_pos)
+            encoder_k = rearrange(encoder_k, "B H (N_t S) C -> (B N_t) H S C", N_t=N_t)
+
+        # Input tensors must be in format ``[B, M, H, K]``, where B is the batch size, M \
+        # the sequence length, H the number of heads, and K the embeding size per head
+
+        q = rearrange(q, "B H M K -> B M H K")
+        encoder_k = rearrange(encoder_k, "B H M K -> B M H K")
+        encoder_v = rearrange(encoder_v, "B H M K -> B M H K")
+        x = attention(q, encoder_k, encoder_v, attention_mode=self.attention_mode)
+        x = rearrange(x, "B M H K -> B H M K")
+
+        # linear transform
+        x_output_shape = (B, N, C)
+        x = x.transpose(1, 2)
+        x = x.reshape(x_output_shape)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        # reshape x to origin shape
+        x = rearrange(x, "(B N_t) S C -> B (N_t S) C", N_t=N_t)
+
+        return x.type(out_dtype)
+
+    def forward(self, x, cond, num_latent_frames=None, num_cond_latents=None, x_ref_attn_map=None, human_num=None):
+
+        B, N, C = x.shape
+        if (num_cond_latents is None or num_cond_latents == 0):
+            # text to video
+            output = self._process_cross_attn(x, cond, num_latent_frames, x_ref_attn_map)
+            return None, output
+        elif num_cond_latents is not None and num_cond_latents > 0:
+            # image to video or video continuation
+            num_cond_latents_thw = num_cond_latents * (N // num_latent_frames)
+            x_noise = x[:, num_cond_latents_thw:]
+            cond = rearrange(cond, "(B N_t) M C -> B N_t M C", B=B)
+            cond = cond[:, num_cond_latents:]
+            cond = rearrange(cond, "B N_t M C -> (B N_t) M C")
+            frames_num = num_latent_frames - num_cond_latents
+            if human_num is not None and human_num == 2:
+                # multitalk mode
+                output_noise = self._process_cross_attn(x_noise, cond, frames_num, x_ref_attn_map)
+            else:
+                # singletalk mode
+                output_noise = self._process_cross_attn(x_noise, cond, frames_num)
+            output_cond = torch.zeros((B, num_cond_latents_thw, C), dtype=output_noise.dtype, device=output_noise.device)
+            return output_cond, output_noise
+        else:
+            raise NotImplementedError

LongCat/nodes.py

@@ -0,0 +1,120 @@
+import torch
+from ..utils import log
+import comfy.model_management as mm
+from comfy_api.latest import io
+
+device = mm.get_torch_device()
+offload_device = mm.unet_offload_device()
+
+
+class WanVideoLongCatAvatarExtendEmbeds(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="WanVideoLongCatAvatarExtendEmbeds",
+            category="WanVideoWrapper",
+            inputs=[
+                io.Latent.Input("prev_latents", tooltip="Full previous latents to be used to continue generation, continuation frames are selected based on 'overlap' parameter"),
+                io.Custom("MULTITALK_EMBEDS").Input("audio_embeds", tooltip="Full length audio embeddings"),
+                io.Int.Input("num_frames", default=93, min=1, max=256, step=1, tooltip="Number of new frames to generate"),
+                io.Int.Input("overlap", default=13, min=0, max=16, step=1, tooltip="Number of overlapping frames from previous latents for video continuation, set to 0 for T2V"),
+                io.Int.Input("frames_processed", default=0, min=0, max=10000, step=1, tooltip="Number of frames already processed in the video, used to select audio features"),
+                io.Combo.Input("if_not_enough_audio", ["pad_with_start", "mirror_from_end"], default="pad_with_start", tooltip="What to do if there are not enough frames in pose_images for the window"),
+                io.Int.Input("ref_frame_index", default=10, min=0, max=1000, step=1, tooltip="Values between 0 - 24 ensures better consistency, while selecting other ranges (e.g., -10 or 30) helps reduce repeated actions"),
+                io.Int.Input("ref_mask_frame_range", default=3, min=0, max=20, step=1, tooltip="Larger range can further help mitigate repeated actions, but excessively large values may introduce artifacts"),
+                io.Latent.Input("ref_latent", optional=True, tooltip="Reference latent used for consistency, generally should be either the init image, or first latent from first generation"),
+                io.Latent.Input("samples", optional=True, tooltip="For the sampler 'samples' input, used for slicing samples per window for vid2vid"),
+            ],
+            outputs=[
+                io.Custom("WANVIDIMAGE_EMBEDS").Output(display_name="image_embeds", tooltip="Embeds for WanVideo LongCat Avatar generation"),
+                io.Latent.Output(display_name="samples_slice", tooltip="Sliced latent samples for the new frames"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, prev_latents, audio_embeds, num_frames, overlap, if_not_enough_audio, frames_processed, ref_frame_index, ref_mask_frame_range, ref_latent=None, samples=None) -> io.NodeOutput:
+
+        new_audio_embed = audio_embeds.copy()
+
+        audio_features = torch.stack(new_audio_embed["audio_features"])
+        num_audio_features = audio_features.shape[1]
+        if audio_features.shape[1] < frames_processed + num_frames:
+            deficit = frames_processed + num_frames - audio_features.shape[1]
+            if if_not_enough_audio == "pad_with_start":
+                pad = audio_features[:, :1].repeat(1, deficit, 1, 1)
+                audio_features = torch.cat([audio_features, pad], dim=1)
+            elif if_not_enough_audio == "mirror_from_end":
+                to_add = audio_features[:, -deficit:, :].flip(dims=[1])
+                audio_features = torch.cat([audio_features, to_add], dim=1)
+            log.warning(f"Not enough audio features, padded with strategy '{if_not_enough_audio}' from {num_audio_features} to {audio_features.shape[1]} frames")
+
+        ref_target_masks = new_audio_embed.get("ref_target_masks", None)
+        if ref_target_masks is not None:
+            new_audio_embed["ref_target_masks"] = ref_target_masks[:, frames_processed:frames_processed+num_frames, :]
+
+        prev_samples = prev_latents["samples"].clone()
+        if overlap != 0:
+            latent_overlap = (overlap - 1) // 4 + 1
+            prev_samples = prev_samples[:, :, -latent_overlap:]
+
+        ref_sample = None
+        if ref_latent is not None:
+            ref_sample = ref_latent["samples"][0, :, :1].clone()
+            log.info(f"Previous latents shape: {prev_samples.shape}, using last {latent_overlap} latent frames for overlap.")
+
+        new_latent_frames = (num_frames - 1) // 4 + 1
+        target_shape = (16, new_latent_frames, prev_samples.shape[-2], prev_samples.shape[-1])
+
+        audio_stride = 2
+        indices = torch.arange(2 * 2 + 1) - 2
+
+        if frames_processed == 0:
+            audio_start_idx = 0
+        else:
+            audio_start_idx = (frames_processed - overlap) * audio_stride
+        audio_end_idx = audio_start_idx + num_frames * audio_stride
+
+        log.info(f"Extracting audio embeddings from index {audio_start_idx} to {audio_end_idx}")
+
+        audio_embs = []
+        for human_idx in range(len(audio_features)):
+            center_indices = torch.arange(audio_start_idx, audio_end_idx, audio_stride).unsqueeze(1) + indices.unsqueeze(0)
+            center_indices = torch.clamp(center_indices, min=0, max=audio_features[human_idx].shape[0] - 1)
+
+            audio_emb = audio_features[human_idx][center_indices].unsqueeze(0).to(device)
+            audio_embs.append(audio_emb)
+        audio_emb = torch.cat(audio_embs, dim=0)
+
+        new_audio_embed["audio_features"] = None
+        new_audio_embed["audio_emb_slice"] = audio_emb
+
+        longcat_avatar_options = {
+            "longcat_ref_latent": ref_sample,
+            "ref_frame_index": ref_frame_index,
+            "ref_mask_frame_range": ref_mask_frame_range,
+        }
+
+        embeds = {
+            "target_shape": target_shape,
+            "num_frames": num_frames,
+            "extra_latents": [{"samples": prev_samples, "index": 0}] if overlap != 0 else None,
+            "multitalk_embeds": new_audio_embed,
+            "longcat_avatar_options": longcat_avatar_options,
+        }
+
+        samples_slice = None
+        if samples is not None:
+            latent_start_index = (frames_processed - 1) // 4 + 1 if frames_processed > 0 else 0
+            latent_end_index = latent_start_index + new_latent_frames
+            samples_slice = samples.copy()
+            samples_slice["samples"] = samples["samples"][:, :, latent_start_index:latent_end_index].clone()
+
+        return io.NodeOutput(embeds, samples_slice)
+
+
+NODE_CLASS_MAPPINGS = {
+    "WanVideoLongCatAvatarExtendEmbeds": WanVideoLongCatAvatarExtendEmbeds,
+    }
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "WanVideoLongCatAvatarExtendEmbeds": "WanVideo LongCat Avatar Extend Embeds",
+    }

LongVie2/modules.py

@@ -0,0 +1,199 @@
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from ..wanvideo.modules.attention import attention
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
+    return (x * (1 + scale) + shift)
+
+
+def sinusoidal_embedding_1d(dim, position):
+    sinusoid = torch.outer(position.type(torch.float64), torch.pow(
+        10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x.to(position.dtype)
+
+
+def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 3d rope precompute
+    f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta)
+    h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    return f_freqs_cis, h_freqs_cis, w_freqs_cis
+
+
+def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 1d rope precompute
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)
+                   [: (dim // 2)].double() / dim))
+    freqs = torch.outer(torch.arange(end, device=freqs.device), freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def rope_apply(x, freqs, num_heads):
+    x = rearrange(x, "b s (n d) -> b s n d", n=num_heads)
+    x_out = torch.view_as_complex(x.to(torch.float64).reshape(
+        x.shape[0], x.shape[1], x.shape[2], -1, 2))
+    x_out = torch.view_as_real(x_out * freqs).flatten(2)
+    return x_out.to(x.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        dtype = x.dtype
+        return self.norm(x.float()).to(dtype) * self.weight
+
+
+class AttentionModule(nn.Module):
+    def __init__(self, num_heads, head_dim):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+
+    def forward(self, q, k, v):
+        b, n, d = q.size(0), self.num_heads, self.head_dim
+        x = attention(
+            q.view(b, -1, n, d),
+            k.view(b, -1, n, d),
+            v.view(b, -1, n, d)
+        )
+        return x.flatten(2)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+
+        self.attn = AttentionModule(self.num_heads, self.head_dim)
+
+    def forward(self, x, freqs):
+        q = self.norm_q(self.q(x))
+        k = self.norm_k(self.k(x))
+        v = self.v(x)
+        q = rope_apply(q, freqs, self.num_heads)
+        k = rope_apply(k, freqs, self.num_heads)
+        x = self.attn(q, k, v)
+        return self.o(x)
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, clip_fea: torch.Tensor = None):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+
+
+        self.k_img = nn.Linear(dim, dim)
+        self.v_img = nn.Linear(dim, dim)
+        self.norm_k_img = RMSNorm(dim, eps=eps)
+
+        self.attn = AttentionModule(self.num_heads, self.head_dim)
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor, clip_fea: torch.Tensor = None):
+        ctx = y
+        q = self.norm_q(self.q(x))
+        k = self.norm_k(self.k(ctx))
+        v = self.v(ctx)
+        x = self.attn(q, k, v)
+        if clip_fea is not None:
+            k_img = self.norm_k_img(self.k_img(clip_fea))
+            v_img = self.v_img(clip_fea)
+            y = self.attn(q, k_img, v_img)
+            x = x + y
+        return self.o(x)
+
+
+class GateModule(nn.Module):
+    def __init__(self,):
+        super().__init__()
+
+    def forward(self, x, gate, residual):
+        return x + gate * residual
+
+class DiTBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.ffn_dim = ffn_dim
+
+        self.self_attn = SelfAttention(dim, num_heads, eps)
+        self.cross_attn = CrossAttention(dim, num_heads, eps)
+        self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm3 = nn.LayerNorm(dim, eps=eps)
+        self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(
+            approximate='tanh'), nn.Linear(ffn_dim, dim))
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+        self.gate = GateModule()
+
+    def forward(self, x, context, t_mod, freqs, clip_fea=None):
+        has_seq = len(t_mod.shape) == 4
+        chunk_dim = 2 if has_seq else 1
+        # msa: multi-head self-attention  mlp: multi-layer perceptron
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=chunk_dim)
+        if has_seq:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2),
+                shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2),
+            )
+        input_x = modulate(self.norm1(x), shift_msa, scale_msa)
+        x = self.gate(x, gate_msa, self.self_attn(input_x, freqs))
+        x = x + self.cross_attn(self.norm3(x), context, clip_fea=clip_fea)
+        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
+        x = self.gate(x, gate_mlp, self.ffn(input_x))
+        return x
+
+
+class WanModelDualControl(torch.nn.Module):
+    def __init__(self, dim: int, ffn_dim: int, eps: float, num_heads: int, control_layers = 12):
+        super().__init__()
+        self.control_layers = control_layers
+        self.control_blocks_dense = nn.ModuleList([
+            DiTBlock(dim//2, num_heads//2, ffn_dim//2, eps)
+            for _ in range(self.control_layers)
+        ])
+
+        self.control_blocks_sparse = nn.ModuleList([
+            DiTBlock(dim//2, num_heads//2, ffn_dim//2, eps)
+            for _ in range(self.control_layers)
+        ])
+
+        self.control_initial_combine_linear_dense = torch.nn.Linear(dim, dim//2)
+        self.control_initial_combine_linear_sparse = torch.nn.Linear(dim, dim//2)
+
+        self.control_text_linear = torch.nn.Linear(dim, dim//2)
+        self.control_t_mod = torch.nn.Linear(dim, dim//2)
+
+        self.control_combine_linears = torch.nn.ModuleList([torch.nn.Linear(dim//2, dim) for _ in range(self.control_layers)])
+        head_dim = dim // num_heads
+        self.freqs = precompute_freqs_cis_3d(head_dim)

LongVie2/nodes.py

@@ -0,0 +1,88 @@
+import torch
+from ..utils import log
+import comfy.model_management as mm
+
+device = mm.get_torch_device()
+offload_device = mm.unet_offload_device()
+
+class WanVideoAddDualControlEmbeds:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                    "embeds": ("WANVIDIMAGE_EMBEDS",),
+                    "vae": ("WANVAE", {"tooltip": "VAE model"}),
+                    "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Strength of the reference embedding"}),
+                    "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Start percentage of the embedding application"}),
+                    "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "End percentage of the embedding application"}),
+                    "first_frame_noise_level": ("FLOAT", {"default": 0.925926, "min": 0.0, "max": 1.0, "step": 0.000001, "tooltip": "Noise level for the first frame when using previous frames"}),
+                },
+                "optional": {
+                    "dense": ("IMAGE", {"tooltip": "Dense control signal (depth) video input"}),
+                    "sparse": ("IMAGE", {"tooltip": "Sparse control signal (tracks) video input"}),
+                    "prev_images": ("IMAGE", {"tooltip": "Previous frames for temporal consistency, default is 8 frames"}),
+                }
+        }
+
+    RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
+    RETURN_NAMES = ("image_embeds",)
+    FUNCTION = "add"
+    CATEGORY = "WanVideoWrapper"
+
+    def add(self, embeds, vae, strength, start_percent, end_percent, first_frame_noise_level, dense=None, sparse=None, prev_images=None):
+        updated = dict(embeds)
+        updated.setdefault("dual_control", {})
+
+        if dense is None and sparse is None:
+            raise ValueError("At least one of dense or sparse inputs must be provided.")
+
+        num_frames = dense.shape[0] if dense is not None else sparse.shape[0]
+        height = dense.shape[1] if dense is not None else sparse.shape[1]
+        width = dense.shape[2] if dense is not None else sparse.shape[2]
+        msk = torch.ones(1, num_frames, height//8, width//8, device=device)
+        msk[:, 1:] = 0
+        msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1)
+        msk = msk.view(1, msk.shape[1] // 4, 4, height//8, width//8)
+        msk = msk.transpose(1, 2)
+
+        dense_input_latent = sparse_input_latent = None
+
+        vae.to(device)
+        if dense is not None:
+            dense_images = 1 - dense[..., :3] # Invert colors for depth to match the usual range in comfy
+            dense_images = dense_images.permute(3, 0, 1, 2) * 2 - 1
+            dense_video_latent = vae.encode([dense_images.to(device, vae.dtype)], device, tiled=False)
+            dense_first = (dense_images[:, :1]).to(device, vae.dtype)
+            vae_input_dense = torch.cat([dense_first, torch.zeros(3, num_frames-1, height, width, device=device, dtype=vae.dtype)], dim=1)
+            dense_concat_latent = vae.encode([vae_input_dense], device, tiled=False)
+            dense_concat_latent = torch.cat([msk, dense_concat_latent], dim=1)
+            dense_input_latent = torch.cat([dense_video_latent, dense_concat_latent],dim=1)
+        if sparse is not None:
+            sparse_images = sparse[..., :3].permute(3, 0, 1, 2) * 2 - 1
+            sparse_video_latent = vae.encode([sparse_images.to(device, vae.dtype)], device, tiled=False)
+            sparse_first = (sparse_images[:, :1]).to(device, vae.dtype)
+            vae_input_sparse = torch.cat([sparse_first, torch.zeros(3, num_frames-1, height, width, device=device, dtype=vae.dtype)], dim=1)
+            sparse_concat_latent = vae.encode([vae_input_sparse], device, tiled=False)
+            sparse_concat_latent = torch.cat([msk, sparse_concat_latent], dim=1)
+            sparse_input_latent = torch.cat([sparse_video_latent, sparse_concat_latent],dim=1)
+
+        if prev_images is not None:
+            prev_images = prev_images[..., :3].permute(3, 0, 1, 2) * 2 - 1
+            prev_video_latent = vae.encode([prev_images.to(device, vae.dtype)], device, tiled=False)
+            updated["dual_control"]["prev_latent"] = prev_video_latent[0]
+
+        vae.to(offload_device)
+        updated["dual_control"]["dense_input_latent"] = dense_input_latent
+        updated["dual_control"]["sparse_input_latent"] = sparse_input_latent
+        updated["dual_control"]["strength"] = strength
+        updated["dual_control"]["start_percent"] = start_percent
+        updated["dual_control"]["end_percent"] = end_percent
+        updated["dual_control"]["first_frame_noise_level"] = first_frame_noise_level
+        return (updated,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "WanVideoAddDualControlEmbeds": WanVideoAddDualControlEmbeds,
+    }
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "WanVideoAddDualControlEmbeds": "WanVideo Add Dual Control Embeds",
+    }

__init__.py

@@ -1,11 +1,11 @@
 try:
-    from .utils import check_duplicate_nodes, log
+    from .utils import check_duplicate_nodes, log, color_text
     duplicate_dirs = check_duplicate_nodes()
     if duplicate_dirs:
         warning_msg = f"WARNING:  Found {len(duplicate_dirs)} other WanVideoWrapper directories:\n"
         for dir_path in duplicate_dirs:
-            warning_msg += f"  - {dir_path}\n"
-        log.warning(warning_msg + "Please remove duplicates to avoid possible conflicts.")
+            warning_msg += f"  - {color_text(dir_path, 'yellow')}\n"
+        log.warning(color_text(warning_msg + "Please remove duplicates to avoid possible conflicts.", "red"))
 except:
     pass
 
@@ -48,6 +48,8 @@ OPTIONAL_MODULES = [
     (".onetoall.nodes", "OneToAll"),
     (".WanMove.nodes", "WanMove"),
     (".SCAIL.nodes", "SCAIL"),
+    (".LongCat.nodes", "LongCat"),
+    (".LongVie2.nodes", "LongVie2"),
 ]
 
 def register_nodes(module_path: str, name: str, optional: bool) -> None:
@@ -70,4 +72,4 @@ for module_path, name in REQUIRED_MODULES:
 for module_path, name in OPTIONAL_MODULES:
     register_nodes(module_path, name, optional=True)
 
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

custom_linear.py

@@ -56,7 +56,7 @@ def _replace_linear(model, compute_dtype, state_dict, prefix="", patches=None, s
         module_prefix = module_prefix.replace("_orig_mod.", "")
         _replace_linear(module, compute_dtype, state_dict, module_prefix, patches, scale_weights, compile_args, modules_to_not_convert)
 
-        if isinstance(module, nn.Linear) and "loras" not in module_prefix and name not in modules_to_not_convert:
+        if isinstance(module, nn.Linear) and "loras" not in module_prefix and "dual_controller" not in module_prefix and name not in modules_to_not_convert:
             weight_key = module_prefix + "weight"
             if weight_key not in state_dict:
                 continue

example_workflows/wanvideo_long_T2V_example_01.json

@@ -1,780 +0,0 @@
-{
-  "id": "8b7a9a57-2303-4ef5-9fc2-bf41713bd1fc",
-  "revision": 0,
-  "last_node_id": 46,
-  "last_link_id": 58,
-  "nodes": [
-    {
-      "id": 33,
-      "type": "Note",
-      "pos": [
-        227.3764190673828,
-        -205.28524780273438
-      ],
-      "size": [
-        351.70458984375,
-        88
-      ],
-      "flags": {},
-      "order": 0,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [],
-      "properties": {},
-      "widgets_values": [
-        "Models:\nhttps://huggingface.co/Kijai/WanVideo_comfy/tree/main"
-      ],
-      "color": "#432",
-      "bgcolor": "#653"
-    },
-    {
-      "id": 11,
-      "type": "LoadWanVideoT5TextEncoder",
-      "pos": [
-        224.15325927734375,
-        -34.481563568115234
-      ],
-      "size": [
-        377.1661376953125,
-        130
-      ],
-      "flags": {},
-      "order": 1,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [
-        {
-          "name": "wan_t5_model",
-          "type": "WANTEXTENCODER",
-          "slot_index": 0,
-          "links": [
-            15
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "LoadWanVideoT5TextEncoder"
-      },
-      "widgets_values": [
-        "umt5-xxl-enc-bf16.safetensors",
-        "bf16",
-        "offload_device",
-        "disabled"
-      ]
-    },
-    {
-      "id": 28,
-      "type": "WanVideoDecode",
-      "pos": [
-        1692.973876953125,
-        -404.8614501953125
-      ],
-      "size": [
-        315,
-        174
-      ],
-      "flags": {},
-      "order": 12,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "vae",
-          "type": "WANVAE",
-          "link": 43
-        },
-        {
-          "name": "samples",
-          "type": "LATENT",
-          "link": 33
-        }
-      ],
-      "outputs": [
-        {
-          "name": "images",
-          "type": "IMAGE",
-          "slot_index": 0,
-          "links": [
-            48
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "WanVideoDecode"
-      },
-      "widgets_values": [
-        true,
-        272,
-        272,
-        144,
-        128
-      ]
-    },
-    {
-      "id": 38,
-      "type": "WanVideoVAELoader",
-      "pos": [
-        1687.4093017578125,
-        -582.2750854492188
-      ],
-      "size": [
-        416.25482177734375,
-        82
-      ],
-      "flags": {},
-      "order": 2,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [
-        {
-          "name": "vae",
-          "type": "WANVAE",
-          "slot_index": 0,
-          "links": [
-            43
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "WanVideoVAELoader"
-      },
-      "widgets_values": [
-        "wanvideo\\Wan2_1_VAE_bf16.safetensors",
-        "bf16"
-      ]
-    },
-    {
-      "id": 42,
-      "type": "GetImageSizeAndCount",
-      "pos": [
-        1708.7301025390625,
-        -140.99705505371094
-      ],
-      "size": [
-        277.20001220703125,
-        86
-      ],
-      "flags": {},
-      "order": 13,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "image",
-          "type": "IMAGE",
-          "link": 48
-        }
-      ],
-      "outputs": [
-        {
-          "name": "image",
-          "type": "IMAGE",
-          "slot_index": 0,
-          "links": [
-            56
-          ]
-        },
-        {
-          "label": "832 width",
-          "name": "width",
-          "type": "INT",
-          "links": null
-        },
-        {
-          "label": "480 height",
-          "name": "height",
-          "type": "INT",
-          "links": null
-        },
-        {
-          "label": "257 count",
-          "name": "count",
-          "type": "INT",
-          "links": null
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "GetImageSizeAndCount"
-      },
-      "widgets_values": []
-    },
-    {
-      "id": 16,
-      "type": "WanVideoTextEncode",
-      "pos": [
-        675.8850708007812,
-        -36.032100677490234
-      ],
-      "size": [
-        420.30511474609375,
-        261.5306701660156
-      ],
-      "flags": {},
-      "order": 10,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "t5",
-          "type": "WANTEXTENCODER",
-          "link": 15
-        },
-        {
-          "name": "model_to_offload",
-          "shape": 7,
-          "type": "WANVIDEOMODEL",
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "text_embeds",
-          "type": "WANVIDEOTEXTEMBEDS",
-          "slot_index": 0,
-          "links": [
-            30
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "WanVideoTextEncode"
-      },
-      "widgets_values": [
-        "high quality nature video featuring a red panda balancing on a bamboo stem while a bird lands on it's head, on the background there is a waterfall",
-        "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
-        true
-      ]
-    },
-    {
-      "id": 30,
-      "type": "VHS_VideoCombine",
-      "pos": [
-        2127.120849609375,
-        -511.9014587402344
-      ],
-      "size": [
-        873.2135620117188,
-        840.2385864257812
-      ],
-      "flags": {},
-      "order": 14,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "images",
-          "type": "IMAGE",
-          "link": 56
-        },
-        {
-          "name": "audio",
-          "shape": 7,
-          "type": "AUDIO",
-          "link": null
-        },
-        {
-          "name": "meta_batch",
-          "shape": 7,
-          "type": "VHS_BatchManager",
-          "link": null
-        },
-        {
-          "name": "vae",
-          "shape": 7,
-          "type": "VAE",
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "Filenames",
-          "type": "VHS_FILENAMES",
-          "links": null
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "VHS_VideoCombine"
-      },
-      "widgets_values": {
-        "frame_rate": 16,
-        "loop_count": 0,
-        "filename_prefix": "WanVideo2_1_T2V",
-        "format": "video/h264-mp4",
-        "pix_fmt": "yuv420p",
-        "crf": 19,
-        "save_metadata": true,
-        "trim_to_audio": false,
-        "pingpong": false,
-        "save_output": true,
-        "videopreview": {
-          "hidden": false,
-          "paused": false,
-          "params": {
-            "filename": "WanVideo2_1_T2V_00412.mp4",
-            "subfolder": "",
-            "type": "output",
-            "format": "video/h264-mp4",
-            "frame_rate": 16,
-            "workflow": "WanVideo2_1_T2V_00412.png",
-            "fullpath": "N:\\AI\\ComfyUI\\output\\WanVideo2_1_T2V_00412.mp4"
-          }
-        }
-      }
-    },
-    {
-      "id": 37,
-      "type": "WanVideoEmptyEmbeds",
-      "pos": [
-        1305.26708984375,
-        -571.7843627929688
-      ],
-      "size": [
-        315,
-        106
-      ],
-      "flags": {},
-      "order": 3,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [
-        {
-          "name": "image_embeds",
-          "type": "WANVIDIMAGE_EMBEDS",
-          "links": [
-            42
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "WanVideoEmptyEmbeds"
-      },
-      "widgets_values": [
-        832,
-        480,
-        257
-      ]
-    },
-    {
-      "id": 35,
-      "type": "WanVideoTorchCompileSettings",
-      "pos": [
-        193.47103881835938,
-        -614.6900024414062
-      ],
-      "size": [
-        390.5999755859375,
-        178
-      ],
-      "flags": {},
-      "order": 4,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [
-        {
-          "name": "torch_compile_args",
-          "type": "WANCOMPILEARGS",
-          "slot_index": 0,
-          "links": []
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "WanVideoTorchCompileSettings"
-      },
-      "widgets_values": [
-        "inductor",
-        false,
-        "default",
-        false,
-        64,
-        true
-      ]
-    },
-    {
-      "id": 45,
-      "type": "WanVideoTeaCache",
-      "pos": [
-        931.4036865234375,
-        -792.5159912109375
-      ],
-      "size": [
-        315,
-        154
-      ],
-      "flags": {},
-      "order": 5,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [
-        {
-          "name": "cache_args",
-          "type": "CACHEARGS",
-          "links": [
-            58
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "WanVideoTeaCache"
-      },
-      "widgets_values": [
-        0.10000000000000002,
-        1,
-        -1,
-        "offload_device",
-        true
-      ]
-    },
-    {
-      "id": 36,
-      "type": "Note",
-      "pos": [
-        796.0189208984375,
-        -521.5020751953125
-      ],
-      "size": [
-        298.2554016113281,
-        108.62744140625
-      ],
-      "flags": {},
-      "order": 6,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [],
-      "properties": {},
-      "widgets_values": [
-        "sdpa should work too, haven't tested flaash\n\nfp8_fast seems to cause huge quality degradation"
-      ],
-      "color": "#432",
-      "bgcolor": "#653"
-    },
-    {
-      "id": 46,
-      "type": "Note",
-      "pos": [
-        937.9556274414062,
-        -940.750244140625
-      ],
-      "size": [
-        297.4364013671875,
-        88
-      ],
-      "flags": {},
-      "order": 7,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [],
-      "properties": {},
-      "widgets_values": [
-        "TeaCache with context windows is VERY experimental and lower values than normal should be used."
-      ],
-      "color": "#432",
-      "bgcolor": "#653"
-    },
-    {
-      "id": 27,
-      "type": "WanVideoSampler",
-      "pos": [
-        1315.2401123046875,
-        -401.48028564453125
-      ],
-      "size": [
-        315,
-        574.1923217773438
-      ],
-      "flags": {},
-      "order": 11,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "model",
-          "type": "WANVIDEOMODEL",
-          "link": 29
-        },
-        {
-          "name": "text_embeds",
-          "type": "WANVIDEOTEXTEMBEDS",
-          "link": 30
-        },
-        {
-          "name": "image_embeds",
-          "type": "WANVIDIMAGE_EMBEDS",
-          "link": 42
-        },
-        {
-          "name": "samples",
-          "shape": 7,
-          "type": "LATENT",
-          "link": null
-        },
-        {
-          "name": "feta_args",
-          "shape": 7,
-          "type": "FETAARGS",
-          "link": null
-        },
-        {
-          "name": "context_options",
-          "shape": 7,
-          "type": "WANVIDCONTEXT",
-          "link": 57
-        },
-        {
-          "name": "cache_args",
-          "shape": 7,
-          "type": "CACHEARGS",
-          "link": 58
-        },
-        {
-          "name": "flowedit_args",
-          "shape": 7,
-          "type": "FLOWEDITARGS",
-          "link": null
-        },
-        {
-          "name": "slg_args",
-          "shape": 7,
-          "type": "SLGARGS",
-          "link": null
-        },
-        {
-          "name": "loop_args",
-          "shape": 7,
-          "type": "LOOPARGS",
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "samples",
-          "type": "LATENT",
-          "slot_index": 0,
-          "links": [
-            33
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "WanVideoSampler"
-      },
-      "widgets_values": [
-        30,
-        6,
-        5,
-        1057359483639288,
-        "fixed",
-        true,
-        "unipc",
-        0,
-        1,
-        "",
-        "comfy"
-      ]
-    },
-    {
-      "id": 43,
-      "type": "WanVideoContextOptions",
-      "pos": [
-        1307.9542236328125,
-        -855.8865356445312
-      ],
-      "size": [
-        315,
-        226
-      ],
-      "flags": {},
-      "order": 8,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "vae",
-          "shape": 7,
-          "type": "WANVAE",
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "context_options",
-          "type": "WANVIDCONTEXT",
-          "slot_index": 0,
-          "links": [
-            57
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "WanVideoContextOptions"
-      },
-      "widgets_values": [
-        "uniform_standard",
-        81,
-        4,
-        16,
-        true,
-        false,
-        6,
-        2
-      ]
-    },
-    {
-      "id": 22,
-      "type": "WanVideoModelLoader",
-      "pos": [
-        620.3950805664062,
-        -357.8426818847656
-      ],
-      "size": [
-        477.4410095214844,
-        226.43276977539062
-      ],
-      "flags": {},
-      "order": 9,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "compile_args",
-          "shape": 7,
-          "type": "WANCOMPILEARGS",
-          "link": null
-        },
-        {
-          "name": "block_swap_args",
-          "shape": 7,
-          "type": "BLOCKSWAPARGS",
-          "link": null
-        },
-        {
-          "name": "lora",
-          "shape": 7,
-          "type": "WANVIDLORA",
-          "link": null
-        },
-        {
-          "name": "vram_management_args",
-          "shape": 7,
-          "type": "VRAM_MANAGEMENTARGS",
-          "link": null
-        }
-      ],
-      "outputs": [
-        {
-          "name": "model",
-          "type": "WANVIDEOMODEL",
-          "slot_index": 0,
-          "links": [
-            29
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "WanVideoModelLoader"
-      },
-      "widgets_values": [
-        "WanVideo\\wan2.1_t2v_1.3B_fp16.safetensors",
-        "fp16",
-        "disabled",
-        "offload_device",
-        "sdpa"
-      ]
-    }
-  ],
-  "links": [
-    [
-      15,
-      11,
-      0,
-      16,
-      0,
-      "WANTEXTENCODER"
-    ],
-    [
-      29,
-      22,
-      0,
-      27,
-      0,
-      "WANVIDEOMODEL"
-    ],
-    [
-      30,
-      16,
-      0,
-      27,
-      1,
-      "WANVIDEOTEXTEMBEDS"
-    ],
-    [
-      33,
-      27,
-      0,
-      28,
-      1,
-      "LATENT"
-    ],
-    [
-      42,
-      37,
-      0,
-      27,
-      2,
-      "WANVIDIMAGE_EMBEDS"
-    ],
-    [
-      43,
-      38,
-      0,
-      28,
-      0,
-      "VAE"
-    ],
-    [
-      48,
-      28,
-      0,
-      42,
-      0,
-      "IMAGE"
-    ],
-    [
-      56,
-      42,
-      0,
-      30,
-      0,
-      "IMAGE"
-    ],
-    [
-      57,
-      43,
-      0,
-      27,
-      5,
-      "WANVIDCONTEXT"
-    ],
-    [
-      58,
-      45,
-      0,
-      27,
-      6,
-      "TEACACHEARGS"
-    ]
-  ],
-  "groups": [],
-  "config": {},
-  "extra": {
-    "ds": {
-      "scale": 0.8140274938684471,
-      "offset": [
-        -122.25834160503663,
-        993.5739491626379
-      ]
-    },
-    "node_versions": {
-      "ComfyUI-WanVideoWrapper": "5a2383621a05825d0d0437781afcb8552d9590fd",
-      "ComfyUI-KJNodes": "a5bd3c86c8ed6b83c55c2d0e7a59515b15a0137f",
-      "ComfyUI-VideoHelperSuite": "0a75c7958fe320efcb052f1d9f8451fd20c730a8"
-    },
-    "VHS_latentpreview": true,
-    "VHS_latentpreviewrate": 0,
-    "VHS_MetadataImage": true,
-    "VHS_KeepIntermediate": true
-  },
-  "version": 0.4
-}

multitalk/multitalk.py

@@ -42,41 +42,43 @@ def rotate_half(x):
     x = torch.stack((-x2, x1), dim=-1)
     return rearrange(x, "... d r -> ... (d r)")
 
-def calculate_x_ref_attn_map(visual_q, ref_k, ref_target_masks, mode='mean', attn_bias=None):
-    
-    ref_k = ref_k.to(visual_q.dtype).to(visual_q.device)
+def calculate_x_ref_attn_map(visual_q, ref_k, ref_target_masks, split_num=4):
     scale = 1.0 / visual_q.shape[-1] ** 0.5
-    visual_q = visual_q * scale
-    visual_q = visual_q.transpose(1, 2)
-    ref_k = ref_k.transpose(1, 2)
-    attn = visual_q @ ref_k.transpose(-2, -1)
-
-    if attn_bias is not None:
-        attn = attn + attn_bias
-
-    x_ref_attn_map_source = attn.softmax(-1) # B, H, x_seqlens, ref_seqlens
+    visual_q = visual_q.transpose(1, 2) * scale
 
+    B, H, x_seqlens, K = visual_q.shape
 
     x_ref_attn_maps = []
-    ref_target_masks = ref_target_masks.to(visual_q.dtype)
-    x_ref_attn_map_source = x_ref_attn_map_source.to(visual_q.dtype)
-
     for class_idx, ref_target_mask in enumerate(ref_target_masks):
-        ref_target_mask = ref_target_mask[None, None, None, ...]
-        x_ref_attnmap = x_ref_attn_map_source * ref_target_mask
-        x_ref_attnmap = x_ref_attnmap.sum(-1) / ref_target_mask.sum() # B, H, x_seqlens, ref_seqlens --> B, H, x_seqlens
-        x_ref_attnmap = x_ref_attnmap.permute(0, 2, 1) # B, x_seqlens, H
-       
-        if mode == 'mean':
-            x_ref_attnmap = x_ref_attnmap.mean(-1) # B, x_seqlens
-        elif mode == 'max':
-            x_ref_attnmap = x_ref_attnmap.max(-1) # B, x_seqlens
-        
-        x_ref_attn_maps.append(x_ref_attnmap)
-    
-    del attn, x_ref_attn_map_source
+        ref_target_mask = ref_target_mask.view(1, 1, 1, -1)
 
-    return torch.concat(x_ref_attn_maps, dim=0)
+        x_ref_attnmap = torch.zeros(B, H, x_seqlens, device=visual_q.device, dtype=visual_q.dtype)
+        chunk_size = min(max(x_seqlens // split_num, 1), x_seqlens)
+
+        for i in range(0, x_seqlens, chunk_size):
+            end_i = min(i + chunk_size, x_seqlens)
+
+            attn_chunk = visual_q[:, :, i:end_i] @ ref_k.permute(0, 2, 3, 1)  # B, H, chunk, ref_seqlens
+
+            # Apply softmax
+            attn_max = attn_chunk.max(dim=-1, keepdim=True).values
+            attn_chunk = (attn_chunk - attn_max).exp()
+            attn_sum = attn_chunk.sum(dim=-1, keepdim=True)
+            attn_chunk = attn_chunk / (attn_sum + 1e-8)
+
+            # Apply mask and sum
+            masked_attn = attn_chunk * ref_target_mask
+            x_ref_attnmap[:, :, i:end_i] = masked_attn.sum(-1) / (ref_target_mask.sum() + 1e-8)
+
+            del attn_chunk, masked_attn
+
+        # Average across heads
+        x_ref_attnmap = x_ref_attnmap.mean(dim=1)  # B, x_seqlens
+        x_ref_attn_maps.append(x_ref_attnmap)
+
+    del visual_q, ref_k
+
+    return torch.cat(x_ref_attn_maps, dim=0)
 
 def get_attn_map_with_target(visual_q, ref_k, shape, ref_target_masks=None, split_num=2):
     """Args:
@@ -129,27 +131,30 @@ class RotaryPositionalEmbedding1D(nn.Module):
             query with the same shape as input.
         """
         freqs_cis = self.precompute_freqs_cis_1d(pos_indices)
-
-        x_ = x.float()
+        in_dtype = x.dtype
+        x = x.float()
 
         freqs_cis = freqs_cis.float().to(x.device)
-        cos, sin = freqs_cis.cos(), freqs_cis.sin()
-        cos, sin = rearrange(cos, 'n d -> 1 1 n d'), rearrange(sin, 'n d -> 1 1 n d')
-        x_ = (x_ * cos) + (rotate_half(x_) * sin)
+        cos = rearrange(freqs_cis.cos(), 'n d -> 1 1 n d')
+        sin = rearrange(freqs_cis.sin(), 'n d -> 1 1 n d')
 
-        return x_.type_as(x)
+        # In-place rotation to save memory
+        x_rotated = rotate_half(x)
+        x.mul_(cos).add_(x_rotated * sin)
+
+        return x.to(in_dtype)
 
 class AudioProjModel(nn.Module):
     def __init__(
         self,
         seq_len=5,
-        seq_len_vf=12,
-        blocks=12,  
-        channels=768, 
+        seq_len_vf=8,
+        blocks=12,
+        channels=768,
         intermediate_dim=512,
         output_dim=768,
         context_tokens=32,
-        norm_output_audio=False,
+        norm_output_audio=True,
     ):
         super().__init__()
 
@@ -273,9 +278,9 @@ class SingleStreamMultiAttention(SingleStreamAttention):
     def __init__(
         self,
         dim: int,
-        encoder_hidden_states_dim: int,
         num_heads: int,
-        qkv_bias: bool,
+        qkv_bias: bool = True,
+        encoder_hidden_states_dim: int = 768,
         class_range: int = 24,
         class_interval: int = 4,
         attention_mode: str = 'sdpa',

multitalk/multitalk_loop.py

@@ -0,0 +1,569 @@
+import torch
+import os
+import gc
+from PIL import Image
+import numpy as np
+from ..latent_preview import prepare_callback
+from ..wanvideo.schedulers import get_scheduler
+from .multitalk import timestep_transform, add_noise
+from ..utils import log, print_memory, temporal_score_rescaling, offload_transformer, init_blockswap, match_and_blend_colors
+from comfy.utils import load_torch_file
+from ..nodes_model_loading import load_weights
+from ..HuMo.nodes import get_audio_emb_window
+import comfy.model_management as mm
+from tqdm import tqdm
+import copy
+
+VAE_STRIDE = (4, 8, 8)
+PATCH_SIZE = (1, 2, 2)
+vae_upscale_factor = 8
+script_directory = os.path.dirname(os.path.abspath(__file__))
+
+device = mm.get_torch_device()
+offload_device = mm.unet_offload_device()
+
+def multitalk_loop(self, **kwargs):
+    # Unpack kwargs into local variables
+    (latent, total_steps, steps, start_step, end_step, shift, cfg, denoise_strength,
+     sigmas, weight_dtype, transformer, patcher, block_swap_args, model, vae, dtype,
+     scheduler, scheduler_step_args, text_embeds, image_embeds, multitalk_embeds,
+     multitalk_audio_embeds, unianim_data, dwpose_data, unianimate_poses, uni3c_embeds,
+     humo_image_cond, humo_image_cond_neg, humo_audio, humo_reference_count,
+     add_noise_to_samples, audio_stride, use_tsr, tsr_k, tsr_sigma, fantasy_portrait_input,
+     noise, timesteps, force_offload, add_cond, control_latents, audio_proj,
+     control_camera_latents, samples, masks, seed_g, gguf_reader, predict_func
+    ) = (kwargs.get(k) for k in (
+        'latent', 'total_steps', 'steps', 'start_step', 'end_step', 'shift', 'cfg',
+        'denoise_strength', 'sigmas', 'weight_dtype', 'transformer', 'patcher',
+        'block_swap_args', 'model', 'vae', 'dtype', 'scheduler', 'scheduler_step_args',
+        'text_embeds', 'image_embeds', 'multitalk_embeds', 'multitalk_audio_embeds',
+        'unianim_data', 'dwpose_data', 'unianimate_poses', 'uni3c_embeds',
+        'humo_image_cond', 'humo_image_cond_neg', 'humo_audio', 'humo_reference_count',
+        'add_noise_to_samples', 'audio_stride', 'use_tsr', 'tsr_k', 'tsr_sigma',
+        'fantasy_portrait_input', 'noise', 'timesteps', 'force_offload', 'add_cond',
+        'control_latents', 'audio_proj', 'control_camera_latents', 'samples', 'masks',
+        'seed_g', 'gguf_reader', 'predict_with_cfg'
+    ))
+
+    mode = image_embeds.get("multitalk_mode", "multitalk")
+    if mode == "auto":
+        mode = transformer.multitalk_model_type.lower()
+    elif mode == "skyreelsv3":
+        num_pseudo_frames = 5
+        pseudo_frames = reference_keyframes = None
+        keyframe_index = 0
+        reference_video = image_embeds.get("reference_video", None)
+    log.info(f"Multitalk mode: {mode}")
+    drop_frames = image_embeds.get("drop_frames", 0)
+    cond_frame = None
+    offload = image_embeds.get("force_offload", False)
+    offloaded = False
+    tiled_vae = image_embeds.get("tiled_vae", False)
+    frame_num = clip_length = image_embeds.get("frame_window_size", 81)
+
+    clip_embeds = image_embeds.get("clip_context", None)
+    if clip_embeds is not None:
+        clip_embeds = clip_embeds.to(dtype)
+    colormatch = image_embeds.get("colormatch", "disabled")
+    motion_frame = image_embeds.get("motion_frame", 25)
+    target_w = image_embeds.get("target_w", None)
+    target_h = image_embeds.get("target_h", None)
+    original_images = image_embeds.get("multitalk_start_image", None)
+    cond_image = original_images.clone() if original_images is not None else None
+    original_color_reference = cond_image.clone() if cond_image is not None else None
+    if original_images is None:
+        original_images = torch.zeros([noise.shape[0], 1, target_h, target_w], device=device)
+
+    output_path = image_embeds.get("output_path", "")
+    img_counter = 0
+
+    if len(multitalk_embeds['audio_features'])==2 and (multitalk_embeds['ref_target_masks'] is None):
+        face_scale = 0.1
+        x_min, x_max = int(target_h * face_scale), int(target_h * (1 - face_scale))
+        lefty_min, lefty_max = int((target_w//2) * face_scale), int((target_w//2) * (1 - face_scale))
+        righty_min, righty_max = int((target_w//2) * face_scale + (target_w//2)), int((target_w//2) * (1 - face_scale) + (target_w//2))
+        human_mask1, human_mask2 = (torch.zeros([target_h, target_w]) for _ in range(2))
+        human_mask1[x_min:x_max, lefty_min:lefty_max] = 1
+        human_mask2[x_min:x_max, righty_min:righty_max] = 1
+        background_mask = torch.where((human_mask1 + human_mask2) > 0, torch.tensor(0), torch.tensor(1))
+        human_masks = [human_mask1, human_mask2, background_mask]
+        ref_target_masks = torch.stack(human_masks, dim=0)
+        multitalk_embeds['ref_target_masks'] = ref_target_masks
+
+    gen_video_list = []
+    is_first_clip = True
+    arrive_last_frame = False
+    cur_motion_frames_num = 1
+    audio_start_idx = iteration_count = step_iteration_count = 0
+    audio_end_idx = (audio_start_idx + clip_length) * audio_stride
+    indices = (torch.arange(4 + 1) - 2) * 1
+    current_condframe_index = 0
+
+    audio_embedding = multitalk_audio_embeds
+    human_num = len(audio_embedding)
+    audio_embs = None
+
+    uni3c_data = None
+    if uni3c_embeds is not None:
+        transformer.controlnet = uni3c_embeds["controlnet"]
+        uni3c_data = uni3c_embeds.copy()
+
+    encoded_silence = None
+
+    try:
+        silence_path = os.path.join(script_directory, "encoded_silence.safetensors")
+        encoded_silence = load_torch_file(silence_path)["audio_emb"].to(dtype)
+    except:
+            log.warning("No encoded silence file found, padding with end of audio embedding instead.")
+
+    total_frames = len(audio_embedding[0])
+    estimated_iterations = total_frames // (frame_num - motion_frame - drop_frames) + 1
+    callback = prepare_callback(patcher, estimated_iterations)
+
+    # If reference_video is provided, extract keyframes from it
+    if mode == "skyreelsv3" and reference_video is not None:
+        ref_video_length = reference_video.shape[1]  # (C, T, H, W)
+        if colormatch == "reinhard_torch":
+            reference_video = match_and_blend_colors(reference_video, original_color_reference, 1.0)
+
+        if ref_video_length >= total_frames:
+            # Reference is long enough - extract keyframes at the expected positions
+            segment_interval = frame_num - motion_frame - drop_frames
+            generate_idx = []
+            current_idx = frame_num - 1
+            while current_idx < total_frames:
+                generate_idx.append(min(current_idx, ref_video_length - 1))
+                current_idx += segment_interval
+        else:
+            # Calculate target indices then map to reference video
+            audio_length = total_frames
+            generate_idx_target = [0]
+            segment_interval = frame_num - motion_frame - drop_frames
+            current_idx = frame_num - 1
+            while current_idx < audio_length - 1:
+                generate_idx_target.append(current_idx)
+                current_idx += segment_interval
+            if generate_idx_target[-1] != audio_length - 1:
+                generate_idx_target.append(audio_length - 1)
+
+            # Map target indices to reference video
+            generate_idx_target = np.array(generate_idx_target, dtype=np.int16)
+            original_max = generate_idx_target[-1]
+            original_min = generate_idx_target[0]
+            if original_max > original_min:
+                generate_idx_float = (generate_idx_target.astype(np.float64) - original_min) * (ref_video_length - 1) / (original_max - original_min)
+                generate_idx = np.clip(np.round(generate_idx_float), 0, ref_video_length - 1).astype(np.int32).tolist()
+            else:
+                generate_idx = [0]
+
+            generate_idx = generate_idx[1:]
+            log.info(f"Reference video ({ref_video_length} frames) mapped to target ({total_frames} frames). Keyframe indices: {generate_idx}")
+
+        # Extract keyframes from reference video
+        # reference_video shape: (C, T, H, W) from nodes.py processing
+        # Select keyframes and add batch dimension: (C, num_keyframes, H, W) -> (1, C, num_keyframes, H, W)
+        selected_keyframes = reference_video[:, generate_idx]  # (C, num_keyframes, H, W)
+        reference_keyframes = selected_keyframes.unsqueeze(0).cpu()  # (1, C, num_keyframes, H, W)
+        log.info(f"Extracted {len(generate_idx)} keyframes from provided reference video at indices {generate_idx}, shape: {reference_keyframes.shape}")
+        log.info(f"Reference video total frames: {reference_video.shape[1]}, will generate {total_frames} total frames with {estimated_iterations} windows")
+
+    if frame_num >= total_frames:
+        arrive_last_frame = True
+        estimated_iterations = 1
+
+    log.info(f"Sampling {total_frames} frames in {estimated_iterations} windows, at {latent.shape[3]*vae_upscale_factor}x{latent.shape[2]*vae_upscale_factor} with {steps} steps")
+
+    while True: # start video generation iteratively
+        self.cache_state = [None, None]
+
+        if mode == "skyreelsv3" and reference_keyframes is not None:
+            clamped_index = min(keyframe_index, reference_keyframes.shape[2] - 1) # Clamp keyframe_index to reuse last keyframe if we run out
+            pseudo_frames = reference_keyframes[:, :, clamped_index:clamped_index+1].repeat(1, 1, num_pseudo_frames, 1, 1) # Use one keyframe and repeat it 5 times
+            log.info(f"Window {iteration_count}: using keyframe {clamped_index}/{reference_keyframes.shape[2]-1} for pseudo frames.")
+            keyframe_index += 1
+        else:
+            pseudo_frames = None
+
+        cur_motion_frames_latent_num = int(1 + (cur_motion_frames_num-1) // 4)
+        if mode == "infinitetalk":
+            cond_image = original_images[:, :, current_condframe_index:current_condframe_index+1] if cond_image is not None else None
+        if multitalk_embeds is not None:
+            audio_embs = []
+            # split audio with window size
+            for human_idx in range(human_num):
+                center_indices = torch.arange(audio_start_idx, audio_end_idx, audio_stride).unsqueeze(1) + indices.unsqueeze(0)
+                center_indices = torch.clamp(center_indices, min=0, max=audio_embedding[human_idx].shape[0]-1)
+                audio_emb = audio_embedding[human_idx][center_indices].unsqueeze(0).to(device)
+                audio_embs.append(audio_emb)
+            audio_embs = torch.cat(audio_embs, dim=0).to(dtype)
+
+        h, w = (cond_image.shape[-2], cond_image.shape[-1]) if cond_image is not None else (target_h, target_w)
+        lat_h, lat_w = h // VAE_STRIDE[1], w // VAE_STRIDE[2]
+        latent_frame_num = (frame_num - 1) // 4 + 1
+
+        noise = torch.randn(16, latent_frame_num, lat_h, lat_w, dtype=torch.float32, device=torch.device("cpu"), generator=seed_g).to(device)
+
+        # Calculate the correct latent slice based on current iteration
+        if is_first_clip:
+            latent_start_idx = 0
+            latent_end_idx = noise.shape[1]
+        else:
+            new_frames_per_iteration = frame_num - motion_frame
+            new_latent_frames_per_iteration = ((new_frames_per_iteration - 1) // 4 + 1)
+            latent_start_idx = iteration_count * new_latent_frames_per_iteration
+            latent_end_idx = latent_start_idx + noise.shape[1]
+
+        if samples is not None:
+            noise_mask = samples.get("noise_mask", None)
+            input_samples = samples["samples"]
+            if input_samples is not None:
+                input_samples = input_samples.squeeze(0).to(noise)
+                # Check if we have enough frames in input_samples
+                if latent_end_idx > input_samples.shape[1]:
+                    # We need more frames than available - pad the input_samples at the end
+                    pad_length = latent_end_idx - input_samples.shape[1]
+                    last_frame = input_samples[:, -1:].repeat(1, pad_length, 1, 1)
+                    input_samples = torch.cat([input_samples, last_frame], dim=1)
+                input_samples = input_samples[:, latent_start_idx:latent_end_idx]
+                if noise_mask is not None:
+                    original_image = input_samples.to(device)
+
+                assert input_samples.shape[1] == noise.shape[1], f"Slice mismatch: {input_samples.shape[1]} vs {noise.shape[1]}"
+
+                if add_noise_to_samples:
+                    latent_timestep = timesteps[0]
+                    noise = noise * latent_timestep / 1000 + (1 - latent_timestep / 1000) * input_samples
+                else:
+                    noise = input_samples
+
+            # diff diff prep
+            if noise_mask is not None:
+                if len(noise_mask.shape) == 4:
+                    noise_mask = noise_mask.squeeze(1)
+                if audio_end_idx > noise_mask.shape[0]:
+                    noise_mask = noise_mask.repeat(audio_end_idx // noise_mask.shape[0], 1, 1)
+                noise_mask = noise_mask[audio_start_idx:audio_end_idx]
+                noise_mask = torch.nn.functional.interpolate(
+                    noise_mask.unsqueeze(0).unsqueeze(0),  # Add batch and channel dims [1,1,T,H,W]
+                    size=(noise.shape[1], noise.shape[2], noise.shape[3]),
+                    mode='trilinear',
+                    align_corners=False
+                ).repeat(1, noise.shape[0], 1, 1, 1)
+
+                thresholds = torch.arange(len(timesteps), dtype=original_image.dtype) / len(timesteps)
+                thresholds = thresholds.reshape(-1, 1, 1, 1, 1).to(device)
+                masks = (1-noise_mask.repeat(len(timesteps), 1, 1, 1, 1).to(device)) > thresholds
+
+        # zero padding and vae encode for img cond
+        if cond_image is not None or cond_frame is not None:
+            cond_ = cond_image if (is_first_clip or humo_image_cond is None) else cond_frame
+            cond_frame_num = cond_.shape[2]
+
+            # Prepare pseudo frames if enabled and available from reference_video
+            if mode == "skyreelsv3" and pseudo_frames is not None:
+                video_frames = torch.zeros(1, 3, frame_num-cond_frame_num-num_pseudo_frames, target_h, target_w, device=device, dtype=vae.dtype)
+                padding_frames_pixels_values = torch.cat([cond_.to(device, vae.dtype), video_frames, pseudo_frames.to(device, vae.dtype)], dim=2)
+            else:
+                video_frames = torch.zeros(1, 3, frame_num-cond_frame_num, target_h, target_w, device=device, dtype=vae.dtype)
+                padding_frames_pixels_values = torch.cat([cond_.to(device, vae.dtype), video_frames], dim=2)
+
+            # encode
+            vae.to(device)
+            y = vae.encode(padding_frames_pixels_values, device=device, tiled=tiled_vae, pbar=False).to(dtype)[0]
+
+            if mode == "infinitetalk":
+                cond_ = cond_image if is_first_clip else cond_frame
+                latent_motion_frames = vae.encode(cond_.to(device, vae.dtype), device=device, tiled=tiled_vae, pbar=False).to(dtype)[0]
+            else:
+                latent_motion_frames = y[:, :cur_motion_frames_latent_num] # C T H W
+
+            vae.to(offload_device)
+
+            #motion_frame_index = cur_motion_frames_latent_num if mode == "infinitetalk" else 1
+            if mode == "skyreelsv3" and pseudo_frames is not None:
+                # create mask in pixel space, then transform
+                msk_pixel = torch.ones(1, frame_num, lat_h, lat_w, device=device)
+                msk_pixel[:, cur_motion_frames_num : -num_pseudo_frames] = 0
+                msk_pixel = torch.cat([
+                    torch.repeat_interleave(msk_pixel[:, 0:1], repeats=4, dim=1),
+                    msk_pixel[:, 1:],
+                ], dim=1)
+                msk_pixel = msk_pixel.view(1, msk_pixel.shape[1] // 4, 4, lat_h, lat_w)
+                msk = msk_pixel.transpose(1, 2).squeeze(0).to(dtype)  # 4 T H W
+            else:
+                msk = torch.zeros(4, latent_frame_num, lat_h, lat_w, device=device, dtype=dtype)
+                msk[:, :1] = 1
+            y = torch.cat([msk, y]) # 4+C T H W
+            mm.soft_empty_cache()
+        else:
+            y = None
+            latent_motion_frames = noise[:, :1]
+
+        partial_humo_cond_input = partial_humo_cond_neg_input = partial_humo_audio = partial_humo_audio_neg = None
+        if humo_image_cond is not None:
+            partial_humo_cond_input = humo_image_cond[:, :latent_frame_num]
+            partial_humo_cond_neg_input = humo_image_cond_neg[:, :latent_frame_num]
+            if y is not None:
+                partial_humo_cond_input[:, :1] = y[:, :1]
+            if humo_reference_count > 0:
+                partial_humo_cond_input[:, -humo_reference_count:] = humo_image_cond[:, -humo_reference_count:]
+                partial_humo_cond_neg_input[:, -humo_reference_count:] = humo_image_cond_neg[:, -humo_reference_count:]
+
+        if humo_audio is not None:
+            if is_first_clip:
+                audio_embs = None
+
+            partial_humo_audio, _ = get_audio_emb_window(humo_audio, frame_num, frame0_idx=audio_start_idx)
+            #zero_audio_pad = torch.zeros(humo_reference_count, *partial_humo_audio.shape[1:], device=partial_humo_audio.device, dtype=partial_humo_audio.dtype)
+            partial_humo_audio[-humo_reference_count:] = 0
+            partial_humo_audio_neg = torch.zeros_like(partial_humo_audio, device=partial_humo_audio.device, dtype=partial_humo_audio.dtype)
+
+        if scheduler == "multitalk":
+            timesteps = list(np.linspace(1000, 1, steps, dtype=np.float32))
+            timesteps.append(0.)
+            timesteps = [torch.tensor([t], device=device) for t in timesteps]
+            timesteps = [timestep_transform(t, shift=shift, num_timesteps=1000) for t in timesteps]
+        else:
+            if isinstance(scheduler, dict):
+                sample_scheduler = copy.deepcopy(scheduler["sample_scheduler"])
+                timesteps = scheduler["timesteps"]
+            else:
+                sample_scheduler, timesteps,_,_ = get_scheduler(scheduler, total_steps, start_step, end_step, shift, device, transformer.dim, denoise_strength, sigmas=sigmas)
+            timesteps = [torch.tensor([float(t)], device=device) for t in timesteps] + [torch.tensor([0.], device=device)]
+
+        # sample videos
+        latent = noise
+
+        # injecting motion frames
+        if not is_first_clip and mode != "infinitetalk":
+            latent_motion_frames = latent_motion_frames.to(latent.dtype).to(device)
+            motion_add_noise = torch.randn(latent_motion_frames.shape, device=torch.device("cpu"), generator=seed_g).to(device).contiguous()
+            add_latent = add_noise(latent_motion_frames, motion_add_noise, timesteps[0])
+            latent[:, :add_latent.shape[1]] = add_latent
+            del motion_add_noise, add_latent
+
+        if offloaded:
+            # Load weights
+            if transformer.patched_linear 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)
+            elif 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)
+            #blockswap init
+            init_blockswap(transformer, block_swap_args, model)
+
+        # Use the appropriate prompt for this section
+        if len(text_embeds["prompt_embeds"]) > 1:
+            prompt_index = min(iteration_count, len(text_embeds["prompt_embeds"]) - 1)
+            positive = [text_embeds["prompt_embeds"][prompt_index]]
+            log.info(f"Using prompt index: {prompt_index}")
+        else:
+            positive = text_embeds["prompt_embeds"]
+
+        # uni3c slices
+        if uni3c_embeds is not None:
+            vae.to(device)
+            # Pad original_images if needed
+            num_frames = original_images.shape[2]
+            if audio_end_idx > num_frames:
+                pad_len = audio_end_idx - num_frames
+                last_frame = original_images[:, :, -1:].repeat(1, 1, pad_len, 1, 1)
+                padded_images = torch.cat([original_images, last_frame], dim=2)
+            else:
+                padded_images = original_images
+            render_latent = vae.encode(
+                padded_images[:, :, audio_start_idx:audio_end_idx].to(device, vae.dtype),
+                device=device, tiled=tiled_vae
+            ).to(dtype)
+
+            vae.to(offload_device)
+            uni3c_data['render_latent'] = render_latent
+
+        # unianimate slices
+        partial_unianim_data = None
+        if unianim_data is not None:
+            partial_dwpose = dwpose_data[:, :, latent_start_idx:latent_end_idx]
+            partial_unianim_data = {
+                "dwpose": partial_dwpose,
+                "random_ref": unianim_data["random_ref"],
+                "strength": unianimate_poses["strength"],
+                "start_percent": unianimate_poses["start_percent"],
+                "end_percent": unianimate_poses["end_percent"]
+            }
+
+        # fantasy portrait slices
+        partial_fantasy_portrait_input = None
+        if fantasy_portrait_input is not None:
+            adapter_proj = fantasy_portrait_input["adapter_proj"]
+            if latent_end_idx > adapter_proj.shape[1]:
+                pad_len = latent_end_idx - adapter_proj.shape[1]
+                last_frame = adapter_proj[:, -1:, :, :].repeat(1, pad_len, 1, 1)
+                padded_proj = torch.cat([adapter_proj, last_frame], dim=1)
+            else:
+                padded_proj = adapter_proj
+            partial_fantasy_portrait_input = fantasy_portrait_input.copy()
+            partial_fantasy_portrait_input["adapter_proj"] = padded_proj[:, latent_start_idx:latent_end_idx]
+
+        mm.soft_empty_cache()
+        gc.collect()
+        # sampling loop
+        sampling_pbar = tqdm(total=len(timesteps)-1, desc=f"Sampling audio indices {audio_start_idx}-{audio_end_idx}", position=0, leave=True)
+        for i in range(len(timesteps)-1):
+            timestep = timesteps[i]
+            latent_model_input = latent.to(device)
+            if mode == "infinitetalk":
+                if humo_image_cond is None or not is_first_clip:
+                    latent_model_input[:, :cur_motion_frames_latent_num] = latent_motion_frames
+
+            noise_pred, _, self.cache_state = predict_func(
+                latent_model_input, cfg[min(i, len(timesteps)-1)], positive, text_embeds["negative_prompt_embeds"],
+                timestep, i, y, clip_embeds, control_latents, None, partial_unianim_data, audio_proj, control_camera_latents, add_cond,
+                cache_state=self.cache_state, multitalk_audio_embeds=audio_embs, fantasy_portrait_input=partial_fantasy_portrait_input,
+                humo_image_cond=partial_humo_cond_input, humo_image_cond_neg=partial_humo_cond_neg_input, humo_audio=partial_humo_audio, humo_audio_neg=partial_humo_audio_neg,
+                uni3c_data = uni3c_data)
+
+            if callback is not None:
+                callback_latent = (latent_model_input.to(device) - noise_pred.to(device) * timestep.to(device) / 1000).detach().permute(1,0,2,3)
+                callback(step_iteration_count, callback_latent, None, estimated_iterations*(len(timesteps)-1))
+                del callback_latent
+
+            sampling_pbar.update(1)
+            step_iteration_count += 1
+
+            # update latent
+            if use_tsr:
+                noise_pred = temporal_score_rescaling(noise_pred, latent, timestep, tsr_k, tsr_sigma)
+            if scheduler == "multitalk":
+                noise_pred = -noise_pred
+                dt = (timesteps[i] - timesteps[i + 1]) / 1000
+                latent = latent + noise_pred * dt[:, None, None, None]
+            else:
+                latent = sample_scheduler.step(noise_pred.unsqueeze(0), timestep, latent.unsqueeze(0).to(noise_pred.device), **scheduler_step_args)[0].squeeze(0)
+            del noise_pred, latent_model_input, timestep
+
+            # differential diffusion inpaint
+            if masks is not None:
+                if i < len(timesteps) - 1:
+                    image_latent = add_noise(original_image.to(device), noise.to(device), timesteps[i+1])
+                    mask = masks[i].to(latent)
+                    latent = image_latent * mask + latent * (1-mask)
+
+            # injecting motion frames
+            if not is_first_clip and mode != "infinitetalk":
+                latent_motion_frames = latent_motion_frames.to(latent.dtype).to(device)
+                motion_add_noise = torch.randn(latent_motion_frames.shape, device=torch.device("cpu"), generator=seed_g).to(device).contiguous()
+                add_latent = add_noise(latent_motion_frames, motion_add_noise, timesteps[i+1])
+                latent[:, :add_latent.shape[1]] = add_latent
+                del motion_add_noise, add_latent
+            elif mode == "infinitetalk":
+                if humo_image_cond is None or not is_first_clip:
+                    latent[:, :cur_motion_frames_latent_num] = latent_motion_frames
+
+        del noise, latent_motion_frames
+        if offload:
+            offload_transformer(transformer, remove_lora=False)
+            offloaded = True
+        if humo_image_cond is not None and humo_reference_count > 0:
+            latent = latent[:,:-humo_reference_count]
+
+        vae.to(device)
+        videos = vae.decode(latent.unsqueeze(0).to(device, vae.dtype), device=device, tiled=tiled_vae, pbar=False)[0].cpu()
+        vae.to(offload_device)
+
+        sampling_pbar.close()
+
+        # crop drop_frames from end if enabled
+        if mode == "skyreelsv3" and drop_frames > 0 and not arrive_last_frame:
+            videos = videos[:, :-drop_frames]
+
+        # optional color correction (less relevant for InfiniteTalk)
+        if colormatch != "disabled":
+            if colormatch == "reinhard_torch":
+                videos = match_and_blend_colors(videos, original_color_reference, 1.0)
+            else:
+                videos = videos.permute(1, 2, 3, 0).float().numpy()
+                from color_matcher import ColorMatcher
+                cm = ColorMatcher()
+                cm_result_list = []
+                for img in videos:
+                    if mode == "infinitetalk":
+                        cm_result = cm.transfer(src=img, ref=cond_image[0].permute(1, 2, 3, 0).squeeze(0).cpu().float().numpy(), method=colormatch)
+                    else:
+                        cm_result = cm.transfer(src=img, ref=original_images[0].permute(1, 2, 3, 0).squeeze(0).cpu().float().numpy(), method=colormatch)
+                    cm_result_list.append(torch.from_numpy(cm_result).to(vae.dtype))
+
+                videos = torch.stack(cm_result_list, dim=0).permute(3, 0, 1, 2)
+
+        # optionally save generated samples to disk
+        if output_path:
+            video_np = videos.clamp(-1.0, 1.0).add(1.0).div(2.0).mul(255).cpu().float().numpy().transpose(1, 2, 3, 0).astype('uint8')
+            num_frames_to_save = video_np.shape[0] if is_first_clip else video_np.shape[0] - cur_motion_frames_num
+            log.info(f"Saving {num_frames_to_save} generated frames to {output_path}")
+            start_idx = 0 if is_first_clip else cur_motion_frames_num
+            for i in range(start_idx, video_np.shape[0]):
+                im = Image.fromarray(video_np[i])
+                im.save(os.path.join(output_path, f"frame_{img_counter:05d}.png"))
+                img_counter += 1
+        else:
+            gen_video_list.append(videos if is_first_clip else videos[:, cur_motion_frames_num:])
+
+        current_condframe_index += 1
+        iteration_count += 1
+
+        # decide whether is done
+        if arrive_last_frame:
+            break
+
+        # update next condition frames
+        is_first_clip = False
+        cur_motion_frames_num = motion_frame
+
+        cond_ = videos[:, -cur_motion_frames_num:].unsqueeze(0)
+        if mode == "infinitetalk":
+            cond_frame = cond_
+        else:
+            cond_image = cond_
+
+        del videos, latent
+
+        # Repeat audio emb
+        if multitalk_embeds is not None:
+            audio_start_idx += (frame_num - cur_motion_frames_num - humo_reference_count - drop_frames)
+            audio_end_idx = audio_start_idx + clip_length
+            if audio_end_idx >= len(audio_embedding[0]):
+                arrive_last_frame = True
+                miss_lengths = []
+                source_frames = []
+                for human_inx in range(human_num):
+                    source_frame = len(audio_embedding[human_inx])
+                    source_frames.append(source_frame)
+                    if audio_end_idx >= len(audio_embedding[human_inx]):
+                        log.warning(f"Audio embedding for subject {human_inx} not long enough: {len(audio_embedding[human_inx])}, need {audio_end_idx}, padding...")
+                        miss_length = audio_end_idx - len(audio_embedding[human_inx]) + 3
+                        log.warning(f"Padding length: {miss_length}")
+                        if encoded_silence is not None:
+                            add_audio_emb = encoded_silence[-1*miss_length:]
+                        else:
+                            add_audio_emb = torch.flip(audio_embedding[human_inx][-1*miss_length:], dims=[0])
+                        audio_embedding[human_inx] = torch.cat([audio_embedding[human_inx], add_audio_emb.to(device, dtype)], dim=0)
+                        miss_lengths.append(miss_length)
+                    else:
+                        miss_lengths.append(0)
+            if mode == "infinitetalk" and current_condframe_index >= original_images.shape[2]:
+                last_frame = original_images[:, :, -1:, :, :]
+                miss_length   = 1
+                original_images = torch.cat([original_images, last_frame.repeat(1, 1, miss_length, 1, 1)], dim=2)
+
+    if not output_path:
+        gen_video_samples = torch.cat(gen_video_list, dim=1)
+    else:
+        gen_video_samples = torch.zeros(3, 1, 64, 64) # dummy output
+
+    if force_offload:
+        if not model["auto_cpu_offload"]:
+            offload_transformer(transformer)
+    try:
+        print_memory(device)
+        torch.cuda.reset_peak_memory_stats(device)
+    except:
+        pass
+    return {"video": gen_video_samples.permute(1, 2, 3, 0), "output_path": output_path},

multitalk/nodes.py

@@ -7,6 +7,8 @@ from ..utils import log, set_module_tensor_to_device
 import os
 import json
 import datetime
+import scipy.signal as ss
+import numpy as np
 
 script_directory = os.path.dirname(os.path.abspath(__file__))
 folder_paths.add_model_folder_path("wav2vec2", os.path.join(folder_paths.models_dir, "wav2vec2"))
@@ -134,6 +136,15 @@ def loudness_norm(audio_array, sr=16000, lufs=-23):
         return audio_array
     normalized_audio = pyloudnorm.normalize.loudness(audio_array, loudness, lufs)
     return normalized_audio
+
+def _add_noise_floor(audio, noise_db=-45):
+        noise_amp = 10 ** (noise_db / 20)
+        noise = np.random.randn(len(audio)) * noise_amp
+        return audio + noise
+
+def _smooth_transients(audio, sr=16000):
+        b, a = ss.butter(3, 3000 / (sr/2))
+        return ss.lfilter(b, a, audio)
     
 class MultiTalkWav2VecEmbeds:
     @classmethod
@@ -153,6 +164,8 @@ class MultiTalkWav2VecEmbeds:
                 "audio_3": ("AUDIO",),
                 "audio_4": ("AUDIO",),
                 "ref_target_masks": ("MASK", {"tooltip": "Per-speaker semantic mask(s) in pixel space. Supply one mask per speaker (plus optional background) to guide mouth assignment"}),
+                "add_noise_floor": ("BOOLEAN", {"default": False, "tooltip": "Add a low-level noise floor to the audio to reduce silent gaps"}),
+                "smooth_transients": ("BOOLEAN", {"default": False, "tooltip": "Apply a low-pass filter to the audio to smooth out transients"}),
             }
         }
 
@@ -161,7 +174,8 @@ class MultiTalkWav2VecEmbeds:
     FUNCTION = "process"
     CATEGORY = "WanVideoWrapper"
 
-    def process(self, wav2vec_model, normalize_loudness, fps, num_frames, audio_1, audio_scale, audio_cfg_scale, multi_audio_type, audio_2=None, audio_3=None, audio_4=None, ref_target_masks=None):
+    def process(self, wav2vec_model, normalize_loudness, fps, num_frames, audio_1, audio_scale, audio_cfg_scale, multi_audio_type, audio_2=None, audio_3=None, audio_4=None,
+                ref_target_masks=None, add_noise_floor=False, smooth_transients=False):
         model_type = wav2vec_model["model_type"]
         if not "tencent" in model_type.lower():
             raise ValueError("Only tencent wav2vec2 models supported by MultiTalk")
@@ -207,6 +221,10 @@ class MultiTalkWav2VecEmbeds:
 
             if normalize_loudness:
                 audio_segment = loudness_norm(audio_segment, sr=sr)
+            if add_noise_floor:
+                audio_segment = _add_noise_floor(audio_segment, noise_db=-45)
+            if smooth_transients:
+                audio_segment = _smooth_transients(audio_segment, sr=sr)
 
             audio_feature = np.squeeze(
                 wav2vec2_feature_extractor(audio_segment, sampling_rate=sr).input_values
@@ -443,13 +461,100 @@ class WanVideoImageToVideoMultiTalk:
         }
 
         return (image_embeds, output_path)
-    
+
+class WanVideoImageToVideoSkyreelsv3_audio:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "vae": ("WANVAE",),
+            "width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "Width of the generation"}),
+            "height": ("INT", {"default": 480, "min": 64, "max": 29048, "step": 8, "tooltip": "Height of the generation"}),
+            "frame_window_size": ("INT", {"default": 81, "min": 1, "max": 10000, "step": 4, "tooltip": "The number of frames to process at once, should be a value the model is generally good at."}),
+            "motion_frame": ("INT", {"default": 5, "min": 1, "max": 10000, "step": 1, "tooltip": "Driven frame length used in the long video generation. Basically the overlap length."}),
+            "drop_frames": ("INT", {"default": 12, "min": 0, "max": 10000, "step": 1, "tooltip": "Additional frames to drop when advancing the audio window. Higher values = less overlap = faster generation but potentially less smooth transitions."}),
+            "tiled_vae": ("BOOLEAN", {"default": False, "tooltip": "Use tiled VAE encoding for reduced memory use"}),
+            "force_offload": ("BOOLEAN", {"default": False, "tooltip": "Whether to force offload the model within the loop for VAE operations, enable if you encounter memory issues."}),
+            "colormatch": (
+            [
+                'disabled',
+                'reinhard_torch',
+                'mkl',
+                'hm',
+                'reinhard',
+                'mvgd',
+                'hm-mvgd-hm',
+                'hm-mkl-hm',
+            ], {
+               "default": 'disabled', "tooltip": "Color matching method to use between the windows"
+            },),
+            },
+            "optional": {
+                "start_image": ("IMAGE", {"tooltip": "Images to encode"}),
+                "reference_video": ("IMAGE", {"tooltip": "Optional: Pre-generated reference video to use for keyframes instead of extracting from first generation. Should be color-matched to source image."}),
+                "clip_embeds": ("WANVIDIMAGE_CLIPEMBEDS", {"tooltip": "Clip vision encoded image"}),
+                "output_path": ("STRING", {"default": "", "tooltip": "If set, will save each window's resulting frames to this folder, also DISABLES returning the final video tensor to save memory"}),
+            }
+        }
+
+    RETURN_TYPES = ("WANVIDIMAGE_EMBEDS", "STRING",)
+    RETURN_NAMES = ("image_embeds", "output_path")
+    FUNCTION = "process"
+    CATEGORY = "WanVideoWrapper"
+    DESCRIPTION = "Enables Multi/InfiniteTalk long video generation sampling method, the video is created in windows with overlapping frames. Not compatible or necessary to be used with context windows and many other features besides Multi/InfiniteTalk."
+
+    def process(self, vae, width, height, frame_window_size, motion_frame, drop_frames, force_offload, colormatch, start_image=None,
+                tiled_vae=False, clip_embeds=None, mode="multitalk", output_path="", reference_video=None):
+
+        H, W = height, width
+        num_frames = ((frame_window_size - 1) // 4) * 4 + 1
+
+        # Resize and rearrange the input image dimensions
+        if start_image is not None:
+            resized_start_image = common_upscale(start_image.movedim(-1, 1), W, H, "lanczos", "disabled").movedim(0, 1)
+            resized_start_image = resized_start_image * 2 - 1
+            resized_start_image = resized_start_image.unsqueeze(0)
+
+        target_shape = (16, (num_frames - 1) // 4 + 1, height // 8, width // 8)
+
+        if output_path:
+            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+            output_path = os.path.join(output_path, f"{timestamp}_{mode}_output")
+            os.makedirs(output_path, exist_ok=True)
+
+        processed_reference_video = None
+        if reference_video is not None:
+            processed_reference_video = common_upscale(reference_video.movedim(-1, 1), W, H, "lanczos", "disabled").movedim(0, 1)
+            processed_reference_video = processed_reference_video * 2 - 1
+
+        image_embeds = {
+            "multitalk_sampling": True,
+            "multitalk_start_image": resized_start_image if start_image is not None else None,
+            "frame_window_size": num_frames,
+            "motion_frame": motion_frame,
+            "drop_frames": drop_frames,
+            "use_pseudo_frames": True,
+            "reference_video": processed_reference_video,
+            "target_h": H,
+            "target_w": W,
+            "tiled_vae": tiled_vae,
+            "force_offload": force_offload,
+            "vae": vae,
+            "target_shape": target_shape,
+            "clip_context": clip_embeds.get("clip_embeds", None) if clip_embeds is not None else None,
+            "colormatch": colormatch,
+            "multitalk_mode": "skyreelsv3",
+            "output_path": output_path
+        }
+
+        return (image_embeds, output_path)
+
 NODE_CLASS_MAPPINGS = {
     "MultiTalkModelLoader": MultiTalkModelLoader,
     "MultiTalkWav2VecEmbeds": MultiTalkWav2VecEmbeds,
     "WanVideoImageToVideoMultiTalk": WanVideoImageToVideoMultiTalk,
     "Wav2VecModelLoader": Wav2VecModelLoader,
     "MultiTalkSilentEmbeds": MultiTalkSilentEmbeds,
+    "WanVideoImageToVideoSkyreelsv3_audio": WanVideoImageToVideoSkyreelsv3_audio,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -458,4 +563,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "WanVideoImageToVideoMultiTalk": "WanVideo Long I2V Multi/InfiniteTalk",
     "Wav2VecModelLoader": "Wav2vec2 Model Loader",
     "MultiTalkSilentEmbeds": "MultiTalk Silent Embeds",
+    "WanVideoImageToVideoSkyreelsv3_audio": "WanVideo Long SkyReelsV3 A2V",
 }

nodes.py

@@ -2,11 +2,13 @@ import os, gc, math
 import torch
 import torch.nn.functional as F
 import hashlib
+from tqdm import tqdm
 
 from .utils import(log, clip_encode_image_tiled, add_noise_to_reference_video, set_module_tensor_to_device)
 from .taehv import TAEHV
 
 from comfy import model_management as mm
+from comfy_api.latest import io
 from comfy.utils import ProgressBar, common_upscale
 from comfy.clip_vision import clip_preprocess, ClipVisionModel
 import folder_paths
@@ -366,10 +368,18 @@ class WanVideoTextEncode:
             cast_dtype = encoder.dtype
 
         params_to_keep = {'norm', 'pos_embedding', 'token_embedding'}
-        for name, param in encoder.model.named_parameters():
+        if hasattr(encoder, 'state_dict'):
+            model_state_dict = encoder.state_dict
+        else:
+            model_state_dict = encoder.model.state_dict()
+
+        params_list = list(encoder.model.named_parameters())
+        pbar = tqdm(params_list, desc="Loading T5 parameters", leave=True)
+        for name, param in pbar:
             dtype_to_use = dtype if any(keyword in name for keyword in params_to_keep) else cast_dtype
-            value = encoder.state_dict[name] if hasattr(encoder, 'state_dict') else encoder.model.state_dict()[name]
+            value = model_state_dict[name]
             set_module_tensor_to_device(encoder.model, name, device=device_to, dtype=dtype_to_use, value=value)
+        del model_state_dict
         if hasattr(encoder, 'state_dict'):
             del encoder.state_dict
             mm.soft_empty_cache()
@@ -550,6 +560,9 @@ class WanVideoApplyNAG:
             "nag_tau": ("FLOAT", {"default": 2.5, "min": 0.0, "max": 10.0, "step": 0.1}),
             "nag_alpha": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}),
             },
+            "optional": {
+                "inplace": ("BOOLEAN", {"default": True, "tooltip": "If true, modifies tensors in place to save memory. Leads to different numerical results which may change the output slightly."}),
+            }
         }
 
     RETURN_TYPES = ("WANVIDEOTEXTEMBEDS", )
@@ -558,7 +571,7 @@ class WanVideoApplyNAG:
     CATEGORY = "WanVideoWrapper"
     DESCRIPTION = "Adds NAG prompt embeds to original prompt embeds: 'https://github.com/ChenDarYen/Normalized-Attention-Guidance'"
 
-    def process(self, original_text_embeds, nag_text_embeds, nag_scale, nag_tau, nag_alpha):
+    def process(self, original_text_embeds, nag_text_embeds, nag_scale, nag_tau, nag_alpha, inplace=True):
         prompt_embeds_dict_copy = original_text_embeds.copy()
         prompt_embeds_dict_copy.update({
                 "nag_prompt_embeds": nag_text_embeds["prompt_embeds"],
@@ -566,6 +579,7 @@ class WanVideoApplyNAG:
                     "nag_scale": nag_scale,
                     "nag_tau": nag_tau,
                     "nag_alpha": nag_alpha,
+                    "inplace": inplace,
                 }
             })
         return (prompt_embeds_dict_copy,)
@@ -888,6 +902,86 @@ class WanVideoAddMTVMotion:
         updated["mtv_crafter_motion"] = new_entry
         return (updated,)
 
+class WanVideoAddStoryMemLatents:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                    "vae": ("WANVAE",),
+                    "embeds": ("WANVIDIMAGE_EMBEDS",),
+                    "memory_images": ("IMAGE",),
+                    "rope_negative_offset": ("BOOLEAN", {"default": False, "tooltip": "Use positive RoPE frequency offset for the memory latents"}),
+                    "rope_negative_offset_frames": ("INT", {"default": 5, "min": 0, "max": 100, "step": 1, "tooltip": "RoPE frequency offset for the memory latents"}),
+                }
+        }
+
+    RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
+    RETURN_NAMES = ("image_embeds",)
+    FUNCTION = "add"
+    CATEGORY = "WanVideoWrapper"
+
+    def add(self, vae, embeds, memory_images, rope_negative_offset, rope_negative_offset_frames):
+        updated = dict(embeds)
+        story_mem_latents, = WanVideoEncodeLatentBatch().encode(vae, memory_images)
+        updated["story_mem_latents"] = story_mem_latents["samples"].squeeze(2).permute(1, 0, 2, 3)  # [C, T, H, W]
+        updated["rope_negative_offset_frames"] = rope_negative_offset_frames if rope_negative_offset else 0
+        return (updated,)
+
+
+class WanVideoSVIProEmbeds:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                    "anchor_samples": ("LATENT", {"tooltip": "Initial start image encoded"}),
+                    "num_frames": ("INT", {"default": 81, "min": 1, "max": 10000, "step": 4, "tooltip": "Number of frames to encode"}),
+                },
+                "optional": {
+                    "prev_samples": ("LATENT", {"tooltip": "Last latent from previous generation"}),
+                    "motion_latent_count": ("INT", {"default": 1, "min": 0, "max": 100, "step": 1, "tooltip": "Number of latents used to continue"}),
+                }
+        }
+
+    RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
+    RETURN_NAMES = ("image_embeds",)
+    FUNCTION = "add"
+    CATEGORY = "WanVideoWrapper"
+
+    def add(self, anchor_samples, num_frames, prev_samples=None, motion_latent_count=1):
+
+        anchor_latent = anchor_samples["samples"][0].clone()
+
+        C, T, H, W = anchor_latent.shape
+
+        total_latents = (num_frames - 1) // 4 + 1
+        device = anchor_latent.device
+        dtype = anchor_latent.dtype
+
+        if prev_samples is None or motion_latent_count == 0:
+            padding_size = total_latents - anchor_latent.shape[1]
+            padding = torch.zeros(C, padding_size, H, W, dtype=dtype, device=device)
+            y = torch.concat([anchor_latent, padding], dim=1)
+        else:
+            prev_latent = prev_samples["samples"][0].clone()
+            motion_latent = prev_latent[:, -motion_latent_count:]
+            padding_size = total_latents - anchor_latent.shape[1] - motion_latent.shape[1]
+            padding = torch.zeros(C, padding_size, H, W, dtype=dtype, device=device)
+            y = torch.concat([anchor_latent, motion_latent, padding], dim=1)
+
+        msk = torch.ones(1, num_frames, H, W, device=device, dtype=dtype)
+        msk[:, 1:] = 0
+        msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1)
+        msk = msk.view(1, msk.shape[1] // 4, 4, H, W)
+        msk = msk.transpose(1, 2)[0]
+
+        image_embeds = {
+            "image_embeds": y,
+            "num_frames": num_frames,
+            "lat_h": H,
+            "lat_w": W,
+            "mask": msk
+        }
+
+        return (image_embeds,)
+
 #region I2V encode
 class WanVideoImageToVideoEncode:
     @classmethod
@@ -1826,33 +1920,7 @@ class WanVideoContextOptions:
         }
 
         return (context_options,)
-    
-    
-class WanVideoFlowEdit:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-                "source_embeds": ("WANVIDEOTEXTEMBEDS", ),
-                "skip_steps": ("INT", {"default": 4, "min": 0}),
-                "drift_steps": ("INT", {"default": 0, "min": 0}),
-                "drift_flow_shift": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 30.0, "step": 0.01}),
-                "source_cfg": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 30.0, "step": 0.01}),
-                "drift_cfg": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 30.0, "step": 0.01}),
-            },
-            "optional": {
-                "source_image_embeds": ("WANVIDIMAGE_EMBEDS", ),
-            }
-        }
 
-    RETURN_TYPES = ("FLOWEDITARGS", )
-    RETURN_NAMES = ("flowedit_args",)
-    FUNCTION = "process"
-    CATEGORY = "WanVideoWrapper"
-    DESCRIPTION = "Flowedit options for WanVideo"
-
-    def process(self, **kwargs):
-        return (kwargs,)
-    
 class WanVideoLoopArgs:
     @classmethod
     def INPUT_TYPES(s):
@@ -2002,6 +2070,76 @@ class WanVideoAddTTMLatents:
 
         return (updated,)
 
+#region self-refine-video
+class WanVideoSelfRefineVideo(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        # Default values for each range
+        default_ranges = [
+            (2, 5, 3),   # Range 1
+            (6, 14, 1),  # Range 2
+            (6, 14, 1),   # Range 3
+            (6, 14, 1),   # Range 4
+            (6, 14, 1),   # Range 5
+        ]
+
+        options = []
+        for num_ranges in range(1, 6):  # 1 to 5 ranges
+            range_inputs = []
+            for i in range(1, num_ranges + 1):
+                start_default, end_default, steps_default = default_ranges[i - 1]
+                range_inputs.extend([
+                    io.Int.Input(f"start_step{i}", default=start_default, min=0, max=999, step=1, tooltip=f"Start step for range {i}"),
+                    io.Int.Input(f"end_step{i}", default=end_default, min=0, max=999, step=1, tooltip=f"End step for range {i}"),
+                    io.Int.Input(f"steps_{i}", default=steps_default, min=1, max=100, step=1, tooltip=f"Number of P&P steps for range {i}"),
+                ])
+            options.append(io.DynamicCombo.Option(
+                key=str(num_ranges),
+                inputs=range_inputs
+            ))
+
+        return io.Schema(
+            node_id="WanVideoSelfRefineVideo",
+            category="WanVideoWrapper",
+            description="https://github.com/agwmon/self-refine-video - Configure stochastic plan for Perturb-and-Project sampling",
+            inputs=[
+                io.Custom("WANVIDIMAGE_EMBEDS").Input("embeds", tooltip="Image embeddings to update"),
+                io.Float.Input(
+                    "uncertainty_threshold",
+                    default=0.25, min=0.0, max=1.0, step=0.01,
+                    tooltip="Lower values make it harder for regions to be considered \"certain\", meaning more pixels will continue being refined. Higher values make it easier to lock in pixels early."
+                ),
+                io.Float.Input("certain_percentage", default=0.999, min=0.0, max=1.0, step=0.001, tooltip="Higher values = stricter requirement = fewer early stops = more iterations"),
+                io.DynamicCombo.Input("num_ranges", options=options, display_name="Number of Ranges", tooltip="Number of step ranges to configure for the stochastic plan"),
+            ],
+            outputs=[
+                io.Custom("WANVIDIMAGE_EMBEDS").Output(display_name="image_embeds", tooltip="Updated image embeddings with self-refine parameters"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, embeds, uncertainty_threshold, certain_percentage, num_ranges) -> io.NodeOutput:
+        updated = dict(embeds)
+        updated["self_refine_uncertainty_threshold"] = uncertainty_threshold
+        updated["self_refine_certain_percentage"] = certain_percentage
+
+        # Build stochastic plan from the dynamic inputs in list format: [(start, end, steps), ...]
+        stochastic_plan = []
+        range_keys = sorted([k for k in num_ranges.keys() if k.startswith('start_step')])
+
+        for start_key in range_keys:
+            i = start_key.replace('start_step', '')
+            start = num_ranges.get(f"start_step{i}")
+            end = num_ranges.get(f"end_step{i}")
+            steps = num_ranges.get(f"steps_{i}")
+
+            if start is not None and end is not None and steps is not None:
+                stochastic_plan.append((start, end, steps))
+
+        updated["stochastic_plan"] = stochastic_plan
+
+        return io.NodeOutput(updated)
+
 #region VideoDecode
 class WanVideoDecode:
     @classmethod
@@ -2046,7 +2184,7 @@ class WanVideoDecode:
             video.clamp_(-1.0, 1.0)
             video.add_(1.0).div_(2.0)
             return video.cpu().float(),
-        latents = samples["samples"]
+        latents = samples["samples"].clone()
         end_image = samples.get("end_image", None)
         has_ref = samples.get("has_ref", False)
         drop_last = samples.get("drop_last", False)
@@ -2122,7 +2260,7 @@ class WanVideoEncodeLatentBatch:
     CATEGORY = "WanVideoWrapper"
     DESCRIPTION = "Encodes a batch of images individually to create a latent video batch where each video is a single frame, useful for I2V init purposes, for example as multiple context window inits"
 
-    def encode(self, vae, images, enable_vae_tiling, tile_x, tile_y, tile_stride_x, tile_stride_y, latent_strength=1.0):
+    def encode(self, vae, images, enable_vae_tiling=False, tile_x=272, tile_y=272, tile_stride_x=144, tile_stride_y=128, latent_strength=1.0):
         vae.to(device)
 
         images = images.clone()
@@ -2228,7 +2366,6 @@ NODE_CLASS_MAPPINGS = {
     "WanVideoEnhanceAVideo": WanVideoEnhanceAVideo,
     "WanVideoContextOptions": WanVideoContextOptions,
     "WanVideoTextEmbedBridge": WanVideoTextEmbedBridge,
-    "WanVideoFlowEdit": WanVideoFlowEdit,
     "WanVideoControlEmbeds": WanVideoControlEmbeds,
     "WanVideoSLG": WanVideoSLG,
     "WanVideoLoopArgs": WanVideoLoopArgs,
@@ -2255,6 +2392,9 @@ NODE_CLASS_MAPPINGS = {
     "TextImageEncodeQwenVL": TextImageEncodeQwenVL,
     "WanVideoUniLumosEmbeds": WanVideoUniLumosEmbeds,
     "WanVideoAddTTMLatents": WanVideoAddTTMLatents,
+    "WanVideoAddStoryMemLatents": WanVideoAddStoryMemLatents,
+    "WanVideoSVIProEmbeds": WanVideoSVIProEmbeds,
+    "WanVideoSelfRefineVideo": WanVideoSelfRefineVideo,
     }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -2270,7 +2410,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "WanVideoEnhanceAVideo": "WanVideo Enhance-A-Video",
     "WanVideoContextOptions": "WanVideo Context Options",
     "WanVideoTextEmbedBridge": "WanVideo TextEmbed Bridge",
-    "WanVideoFlowEdit": "WanVideo FlowEdit",
     "WanVideoControlEmbeds": "WanVideo Control Embeds",
     "WanVideoSLG": "WanVideo SLG",
     "WanVideoLoopArgs": "WanVideo Loop Args",
@@ -2296,4 +2435,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "WanVideoAddBindweaveEmbeds": "WanVideo Add Bindweave Embeds",
     "WanVideoUniLumosEmbeds": "WanVideo UniLumos Embeds",
     "WanVideoAddTTMLatents": "WanVideo Add TTMLatents",
+    "WanVideoAddStoryMemLatents": "WanVideo Add StoryMem Latents",
+    "WanVideoSVIProEmbeds": "WanVideo SVIPro Embeds",
 }

nodes_model_loading.py

@@ -13,7 +13,7 @@ from .wanvideo.wan_video_vae import WanVideoVAE, WanVideoVAE38
 from .custom_linear import _replace_linear
 
 from accelerate import init_empty_weights
-from .utils import set_module_tensor_to_device
+from .utils import set_module_tensor_to_device, get_module_memory_mb_per_device
 
 import folder_paths
 import comfy.model_management as mm
@@ -36,6 +36,9 @@ try:
 except:
     PromptServer = None
 
+attention_modes = ["sdpa", "flash_attn_2", "flash_attn_3", "sageattn", "sageattn_3", "radial_sage_attention", "sageattn_compiled",
+                    "sageattn_ultravico", "comfy"]
+
 #from city96's gguf nodes
 def update_folder_names_and_paths(key, targets=[]):
     # check for existing key
@@ -49,17 +52,6 @@ def update_folder_names_and_paths(key, targets=[]):
         log.warning(f"Unknown file list already present on key {key}: {base}")
 update_folder_names_and_paths("unet_gguf", ["diffusion_models", "unet"])
 
-class WanVideoModel(comfy.model_base.BaseModel):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.pipeline = {}
-
-    def __getitem__(self, k):
-        return self.pipeline[k]
-
-    def __setitem__(self, k, v):
-        self.pipeline[k] = v
-
 try:
     from comfy.latent_formats import Wan21, Wan22
     latent_format = Wan21
@@ -68,16 +60,27 @@ except: #for backwards compatibility
     from comfy.latent_formats import HunyuanVideo
     latent_format = HunyuanVideo
 
+class WanVideoModel(torch.nn.Module):
+    def __init__(self, model_config, transformer, device=None):
+        super().__init__()
+        self.latent_format = model_config.latent_format
+        self.model_config = model_config
+        self.device = device
+        self.current_patcher = None
+        self.diffusion_model = transformer
+        self.pipeline = {}
+
+    def __getitem__(self, k):
+        return self.pipeline[k]
+
+    def __setitem__(self, k, v):
+        self.pipeline[k] = v
+
 class WanVideoModelConfig:
-    def __init__(self, dtype, latent_format=latent_format):
+    def __init__(self, latent_format=latent_format):
         self.unet_config = {}
         self.unet_extra_config = {}
         self.latent_format = latent_format
-        #self.latent_format.latent_channels = 16
-        self.manual_cast_dtype = dtype
-        self.sampling_settings = {"multiplier": 1.0}
-        self.memory_usage_factor = 2.0
-        self.unet_config["disable_unet_model_creation"] = True
 
 def filter_state_dict_by_blocks(state_dict, blocks_mapping, layer_filter=[]):
     filtered_dict = {}
@@ -132,6 +135,7 @@ def standardize_lora_key_format(lora_sd):
             k = k.replace('vace_blocks.', 'diffusion_model.vace_blocks.')
         k = k.replace('.default.', '.')
         k = k.replace('.diff_m', '.modulation.diff')
+        k = k.replace('base_model.model.', 'diffusion_model.')
 
         # Fun LoRA format
         if k.startswith('lora_unet__'):
@@ -177,7 +181,7 @@ def standardize_lora_key_format(lora_sd):
 
                         new_key += f".{component}"
 
-                # Handle weight type - this is the critical fix
+                # Handle weight type
                 if weight_type:
                     if weight_type == 'alpha':
                         new_key += '.alpha'
@@ -208,12 +212,12 @@ def standardize_lora_key_format(lora_sd):
                 new_key = new_key.replace('time_embedding', 'time.embedding')
                 new_key = new_key.replace('time_projection', 'time.projection')
 
-                # Replace remaining underscores with dots, carefully
+                # Replace remaining underscores with dots
                 parts = new_key.split('.')
                 final_parts = []
                 for part in parts:
                     if part in ['img_emb', 'self_attn', 'cross_attn']:
-                        final_parts.append(part)  # Keep these intact
+                        final_parts.append(part)
                     else:
                         final_parts.append(part.replace('_', '.'))
                 new_key = '.'.join(final_parts)
@@ -273,6 +277,20 @@ def standardize_lora_key_format(lora_sd):
         new_sd[k] = v
     return new_sd
 
+def compensate_rs_lora_format(lora_sd):
+    rank = lora_sd["base_model.model.blocks.0.cross_attn.k.lora_A.weight"].shape[0]
+    alpha = torch.tensor(rank * rank // rank ** 0.5)
+    log.info(f"Detected rank stabilized peft lora format with rank {rank}, setting alpha to {alpha} to compensate.")
+    new_sd = {}
+    for k, v in lora_sd.items():
+        if k.endswith(".lora_A.weight"):
+            new_sd[k] = v
+            new_k = k.replace(".lora_A.weight", ".alpha")
+            new_sd[new_k] = alpha
+        else:
+            new_sd[k] = v
+    return new_sd
+
 class WanVideoBlockSwap:
     @classmethod
     def INPUT_TYPES(s):
@@ -363,7 +381,7 @@ class WanVideoLoraSelect:
             "required": {
                "lora": (folder_paths.get_filename_list("loras"),
                 {"tooltip": "LORA models are expected to be in ComfyUI/models/loras with .safetensors extension"}),
-                "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
+                "strength": ("FLOAT", {"default": 1.0, "min": -1000.0, "max": 1000.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
             },
             "optional": {
                 "prev_lora":("WANVIDLORA", {"default": None, "tooltip": "For loading multiple LoRAs"}),
@@ -755,6 +773,8 @@ class WanVideoSetLoRAs:
             lora_sd = load_torch_file(lora_path, safe_load=True)
             if "dwpose_embedding.0.weight" in lora_sd: #unianimate
                 raise NotImplementedError("Unianimate LoRA patching is not implemented in this node.")
+            if "base_model.model.blocks.0.cross_attn.k.lora_A.weight" in lora_sd: # assume rs_lora
+                lora_sd = compensate_rs_lora_format(lora_sd)
 
             lora_sd = standardize_lora_key_format(lora_sd)
             if l["blocks"]:
@@ -790,7 +810,6 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None,
                       "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())
     pbar = ProgressBar(param_count)
-    cnt = 0
     block_idx = vace_block_idx = None
 
     if gguf:
@@ -810,7 +829,7 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None,
             all_tensors.extend(r.tensors)
         for tensor in all_tensors:
             name = rename_fuser_block(tensor.name)
-            if "glob" not in name and "audio_proj" in name:
+            if "glob" not in name and "multitalk_audio_proj" not in name and "audio_proj" in name:
                 name = name.replace("audio_proj", "multitalk_audio_proj")
             load_device = device
             if "vace_blocks." in name:
@@ -844,7 +863,7 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None,
             )
             transformer.gguf_patched = True
     else:
-        log.info("Using accelerate to load and assign model weights to device...")
+        log.info("Loading and assigning model weights to device...")
     named_params = transformer.named_parameters()
 
     for name, param in tqdm(named_params,
@@ -900,12 +919,17 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None,
                     load_device = offload_device
         # Set tensor to device
         set_module_tensor_to_device(transformer, name, device=load_device, dtype=dtype_to_use, value=value)
-        cnt += 1
-        if cnt % 100 == 0:
-            pbar.update(100)
+        pbar.update(1)
 
     #[print(name, param.device, param.dtype) for name, param in transformer.named_parameters()]
+    memory_on_device = get_module_memory_mb_per_device(transformer)
+    log.info("-" * 25)
+    log.info("Transformer weights loaded:")
+    for dev, mem_mb in memory_on_device.items():
+        log.info(f"Device: {dev:8s} | Memory: {mem_mb:,.2f} MB")
 
+    if hasattr(pbar, "_last_sent_value"):
+        pbar._last_sent_value = -1
     pbar.update_absolute(0)
 
 def patch_control_lora(transformer, device):
@@ -966,7 +990,8 @@ def add_lora_weights(patcher, lora, base_dtype, merge_loras=False):
             from .unianimate.nodes import update_transformer
             log.info("Unianimate LoRA detected, patching model...")
             patcher.model.diffusion_model, unianimate_sd = update_transformer(patcher.model.diffusion_model, lora_sd)
-
+        if "base_model.model.blocks.0.cross_attn.k.lora_A.weight" in lora_sd: # assume rs_lora
+                lora_sd = compensate_rs_lora_format(lora_sd)
         lora_sd = standardize_lora_key_format(lora_sd)
 
         if l["blocks"]:
@@ -988,6 +1013,66 @@ def add_lora_weights(patcher, lora, base_dtype, merge_loras=False):
         del lora_sd
     return patcher, control_lora, unianimate_sd
 
+class WanVideoSetAttentionModeOverride:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("WANVIDEOMODEL", ),
+                "attention_mode": (attention_modes, {"default": "sdpa"}),
+                "start_step": ("INT", {"default": 0, "min": 0, "max": 10000, "step": 1, "tooltip": "Step to start applying the attention mode override"}),
+                "end_step": ("INT", {"default": 10000, "min": 1, "max": 10000, "step": 1, "tooltip": "Step to end applying the attention mode override"}),
+                "verbose": ("BOOLEAN", {"default": False, "tooltip": "Print verbose info about attention mode override during generation"}),
+            },
+            "optional": {
+                "blocks":("INT", {"forceInput": True} ),
+            }
+        }
+
+    RETURN_TYPES = ("WANVIDEOMODEL",)
+    RETURN_NAMES = ("model", )
+    FUNCTION = "getmodelpath"
+    CATEGORY = "WanVideoWrapper"
+    DESCRIPTION = "Override the attention mode for the model for specific step and/or block range"
+
+    def getmodelpath(self, model, attention_mode, start_step, end_step, verbose, blocks=None):
+        model_clone = model.clone()
+        attention_mode_override = {
+            "mode": attention_mode,
+            "start_step": start_step,
+            "end_step": end_step,
+            "verbose": verbose,
+        }
+        if blocks is not None:
+            attention_mode_override["blocks"] = blocks
+        model_clone.model_options['transformer_options']["attention_mode_override"] = attention_mode_override
+
+        return (model_clone,)
+
+
+class WanVideoUltraVicoSettings:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("WANVIDEOMODEL", ),
+                "alpha": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.001, "tooltip": "Alpha value for the decay, higher values mean slower decay"}),
+            },
+        }
+
+    RETURN_TYPES = ("WANVIDEOMODEL",)
+    RETURN_NAMES = ("model", )
+    FUNCTION = "getmodelpath"
+    CATEGORY = "WanVideoWrapper"
+    DESCRIPTION = "Set UltraVico parameters, attention mode still needs to be set to sageattn_ultravico, https://github.com/thu-ml/DiT-Extrapolation"
+
+    def getmodelpath(self, model, alpha):
+        model_clone = model.clone()
+        model_clone.model_options['transformer_options']["ultravico_alpha"] = alpha
+
+        return (model_clone,)
+
+
 #region Model loading
 class WanVideoModelLoader:
     @classmethod
@@ -1002,17 +1087,7 @@ class WanVideoModelLoader:
             "load_device": (["main_device", "offload_device"], {"default": "offload_device", "tooltip": "Initial device to load the model to, NOT recommended with the larger models unless you have 48GB+ VRAM"}),
             },
             "optional": {
-                "attention_mode": ([
-                    "sdpa",
-                    "flash_attn_2",
-                    "flash_attn_3",
-                    "sageattn",
-                    "sageattn_3",
-                    "radial_sage_attention",
-                    "sageattn_compiled",
-                    "sageattn_ultravico",
-                    "comfy"
-                    ], {"default": "sdpa"}),
+                "attention_mode": (attention_modes, {"default": "sdpa"}),
                 "compile_args": ("WANCOMPILEARGS", ),
                 "block_swap_args": ("BLOCKSWAPARGS", ),
                 "lora": ("WANVIDLORA", {"default": None}),
@@ -1128,7 +1203,7 @@ class WanVideoModelLoader:
         scale_weights = {}
         if "fp8" in quantization:
             for k, v in sd.items():
-                if k.endswith(".scale_weight"):
+                if k.endswith(".scale_weight") or k.endswith(".weight_scale"):
                     is_scaled_fp8 = True
                     break
 
@@ -1153,7 +1228,7 @@ class WanVideoModelLoader:
         # currently this can be VACE, MTV-Crafter, Lynx or Ovi-audio weights
         if extra_model is not None:
             for _model in extra_model:
-                print("Loading extra model: ", _model["path"])
+                log.info(f"Loading extra model: {_model['path']}")
                 if gguf:
                     if not _model["path"].endswith(".gguf"):
                         raise ValueError("With GGUF main model the extra model must also be GGUF quantized, if the main model already has VACE included, you can disconnect the extra module loader")
@@ -1217,9 +1292,7 @@ class WanVideoModelLoader:
             lynx_ip_layers = "lite"
 
         model_type = "t2v"
-        if "audio_injector.injector.0.k.weight" in sd:
-            model_type = "s2v"
-        elif not "text_embedding.0.weight" in sd:
+        if not "text_embedding.0.weight" in sd:
             model_type = "no_cross_attn" #minimaxremover
         elif "model_type.Wan2_1-FLF2V-14B-720P" in sd or "img_emb.emb_pos" in sd or "flf2v" in model.lower():
             model_type = "fl2v"
@@ -1229,6 +1302,8 @@ class WanVideoModelLoader:
             model_type = "t2v"
         elif "control_adapter.conv.weight" in sd:
             model_type = "t2v"
+        if "audio_injector.injector.0.k.weight" in sd:
+            model_type = "s2v"
 
         out_dim = 16
         if dim == 5120: #14B
@@ -1436,7 +1511,45 @@ class WanVideoModelLoader:
                     block.cross_attn.ip_adapter_single_stream_k_proj = nn.Linear(context_dim, dim, bias=False)
                     block.cross_attn.ip_adapter_single_stream_v_proj = nn.Linear(context_dim, dim, bias=False)
 
-        if multitalk_model is not None:
+        # LongCat Avatar
+        if "multitalk_audio_proj.proj1.weight" in sd and "blocks.0.audio_cross_attn.q_norm.weight" in sd:
+            log.info("MultiTalk/InfiniteTalk model detected, patching model...")
+            from .multitalk.multitalk import AudioProjModel
+            from .wanvideo.modules.model import WanLayerNorm
+            from .LongCat.layers import SingleStreamAttention
+
+
+            for block in transformer.blocks:
+                with init_empty_weights():
+                    if "blocks.0.audio_modulation.1.weight" in sd:
+                        block.audio_modulation = nn.Sequential(nn.SiLU(), nn.Linear(512, 3 * dim, bias=True))
+                    block.norm_x = WanLayerNorm(dim, transformer.eps, elementwise_affine=True)
+                    block.audio_cross_attn = SingleStreamAttention(
+                            dim=dim,
+                            encoder_hidden_states_dim=768,
+                            num_heads=num_heads,
+                        qkv_bias=True,
+                        qk_norm=True,
+                        class_range=24,
+                        class_interval=4,
+                        attention_mode=attention_mode,
+                    )
+                    multitalk_proj_model = AudioProjModel()
+            transformer.multitalk_audio_proj = multitalk_proj_model
+        # SkyreelsV3
+        elif "blocks.1.audio_cross_attn.kv_linear.weight" in sd and "audio_proj.proj1.weight" in sd:
+            sd = {k.replace("audio_proj", "multitalk_audio_proj"): v for k, v in sd.items()}
+            # init audio module
+            from .multitalk.multitalk import SingleStreamMultiAttention, AudioProjModel
+            from .wanvideo.modules.model import WanLayerNorm
+
+            for block in transformer.blocks:
+                with init_empty_weights():
+                    block.norm_x = WanLayerNorm(dim, transformer.eps, elementwise_affine=True)
+                    block.audio_cross_attn = SingleStreamMultiAttention(dim=dim, num_heads=num_heads, attention_mode=attention_mode)
+
+            transformer.multitalk_audio_proj = AudioProjModel()
+        elif multitalk_model is not None:
             multitalk_model_type = multitalk_model.get("model_type", "MultiTalk")
             log.info(f"{multitalk_model_type} detected, patching model...")
 
@@ -1453,15 +1566,7 @@ class WanVideoModelLoader:
             for block in transformer.blocks:
                 with init_empty_weights():
                     block.norm_x = WanLayerNorm(dim, transformer.eps, elementwise_affine=True)
-                    block.audio_cross_attn = SingleStreamMultiAttention(
-                            dim=dim,
-                            encoder_hidden_states_dim=768,
-                            num_heads=num_heads,
-                        qkv_bias=True,
-                        class_range=24,
-                        class_interval=4,
-                        attention_mode=attention_mode,
-                    )
+                    block.audio_cross_attn = SingleStreamMultiAttention(dim=dim, num_heads=num_heads, attention_mode=attention_mode)
             transformer.multitalk_audio_proj = multitalk_model["proj_model"]
             transformer.multitalk_model_type = multitalk_model_type
 
@@ -1504,11 +1609,7 @@ class WanVideoModelLoader:
             transformer.text_projection = nn.Sequential(nn.Linear(sd["text_projection.0.weight"].shape[1], text_dim), nn.GELU(approximate='tanh'), nn.Linear(text_dim, text_dim))
 
         latent_format=Wan22 if dim == 3072 else Wan21
-        comfy_model = WanVideoModel(
-            WanVideoModelConfig(base_dtype, latent_format=latent_format),
-            model_type=comfy.model_base.ModelType.FLOW,
-            device=device,
-        )
+        comfy_model = WanVideoModel(WanVideoModelConfig(latent_format=latent_format), device=device, transformer=transformer)
 
         # SteadyDancer
         if "condition_embedding_align.cross_attn.in_proj_bias" in sd:
@@ -1572,6 +1673,24 @@ class WanVideoModelLoader:
                     block.ref_attn_v_img = nn.Linear(in_features, out_features)
                     block.ref_attn_norm_k_img = WanRMSNorm(out_features, eps=1e-6)
 
+        if "blocks.0.control_blocks_dense.cross_attn.k.weight" in sd:
+            log.info("LongVie2 model detected, patching model...")
+            from .LongVie2.modules import WanModelDualControl
+            control_layers = 12
+            with init_empty_weights():
+                dual_controller = WanModelDualControl(dim=5120, ffn_dim=13824, eps=1e-06, num_heads=40, control_layers=control_layers)
+                for b in range(control_layers):
+                    transformer.blocks[b].control_blocks_dense = dual_controller.control_blocks_dense[b]
+                    transformer.blocks[b].control_blocks_sparse = dual_controller.control_blocks_sparse[b]
+                    transformer.blocks[b].control_combine_linears = dual_controller.control_combine_linears[b]
+                transformer.dual_controller = nn.Module()
+                transformer.dual_controller.control_initial_combine_linear_dense = dual_controller.control_initial_combine_linear_dense
+                transformer.dual_controller.control_initial_combine_linear_sparse = dual_controller.control_initial_combine_linear_sparse
+                transformer.dual_controller.control_t_mod = dual_controller.control_t_mod
+                transformer.dual_controller.control_text_linear = dual_controller.control_text_linear
+                transformer.dual_controller_freqs = dual_controller.freqs
+
+
         comfy_model.diffusion_model = transformer
         comfy_model.load_device = transformer_load_device
         patcher = comfy.model_patcher.ModelPatcher(comfy_model, device, offload_device)
@@ -1752,6 +1871,7 @@ class WanVideoVAELoader:
                 ),
                 "compile_args": ("WANCOMPILEARGS", ),
                 "use_cpu_cache": ("BOOLEAN", {"default": False, "tooltip": "Reduces VRAM usage, but slows the VAE down a lot"}),
+                "verbose": ("BOOLEAN", {"default": False, "tooltip": "Enables memory usage logging when using the model"}),
             }
         }
 
@@ -1761,7 +1881,7 @@ class WanVideoVAELoader:
     CATEGORY = "WanVideoWrapper"
     DESCRIPTION = "Loads Wan VAE model from 'ComfyUI/models/vae'"
 
-    def loadmodel(self, model_name, precision, compile_args=None, use_cpu_cache=False):
+    def loadmodel(self, model_name, precision, compile_args=None, use_cpu_cache=False, verbose=False):
         dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
         model_path = folder_paths.get_full_path_or_raise("vae", model_name)
         vae_sd = load_torch_file(model_path, safe_load=True)
@@ -1778,9 +1898,9 @@ class WanVideoVAELoader:
             pruning_rate = 0.0
 
         if vae_sd["model.conv2.weight"].shape[0] == 16:
-            vae = WanVideoVAE(dtype=dtype, pruning_rate=pruning_rate, cpu_cache=use_cpu_cache)
+            vae = WanVideoVAE(dtype=dtype, pruning_rate=pruning_rate, cpu_cache=use_cpu_cache, verbose=verbose)
         elif vae_sd["model.conv2.weight"].shape[0] == 48:
-            vae = WanVideoVAE38(dtype=dtype, pruning_rate=pruning_rate, cpu_cache=use_cpu_cache)
+            vae = WanVideoVAE38(dtype=dtype, pruning_rate=pruning_rate, cpu_cache=use_cpu_cache, verbose=verbose)
 
         vae.load_state_dict(vae_sd)
         del vae_sd
@@ -1992,6 +2112,8 @@ NODE_CLASS_MAPPINGS = {
     "WanVideoTorchCompileSettings": WanVideoTorchCompileSettings,
     "LoadWanVideoT5TextEncoder": LoadWanVideoT5TextEncoder,
     "LoadWanVideoClipTextEncoder": LoadWanVideoClipTextEncoder,
+    "WanVideoSetAttentionModeOverride": WanVideoSetAttentionModeOverride,
+    "WanVideoUltraVicoSettings": WanVideoUltraVicoSettings,
     }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -2010,4 +2132,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "WanVideoTorchCompileSettings": "WanVideo Torch Compile Settings",
     "LoadWanVideoT5TextEncoder": "WanVideo T5 Text Encoder Loader",
     "LoadWanVideoClipTextEncoder": "WanVideo CLIP Text Encoder Loader",
+    "WanVideoSetAttentionModeOverride": "WanVideo Set Attention Mode Override",
+    "WanVideoUltraVicoSettings": "WanVideo UltraVico Settings"
     }

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI-WanVideoWrapper"
 description = "ComfyUI wrapper nodes for WanVideo"
-version = "1.4.4"
+version = "1.4.7"
 license = {file = "LICENSE"}
 dependencies = ["accelerate >= 1.2.1", "diffusers >= 0.33.0", "peft >= 0.17.0", "ftfy", "gguf >= 0.17.1", "pyloudnorm"]
 

ultravico/sageattn/attn_qk_int8_per_block.py

@@ -38,7 +38,7 @@ def _attn_fwd_inner(acc, l_i, m_i, q, q_scale, kv_len, current_flag,
 
         qk = tl.dot(q, k).to(tl.float32) * q_scale * k_scale
 
-        window_th   =  1560 * 21 / 2
+        window_th   =  frame_tokens * window_width / 2
         dist2       = tl.abs(m - n).to(tl.int32)
         dist_mask   = dist2 <= window_th
 
@@ -46,7 +46,7 @@ def _attn_fwd_inner(acc, l_i, m_i, q, q_scale, kv_len, current_flag,
 
         qk = tl.where(dist_mask | negative_mask, qk, qk*multi_factor)
 
-        window3 = (m <= frame_tokens) & (n > 21*frame_tokens)
+        window3 = (m <= frame_tokens) & (n > window_width*frame_tokens)
         qk = tl.where(window3, -1e4, qk)
 
 

utils.py

@@ -4,17 +4,121 @@ import logging
 import math
 from tqdm import tqdm
 from pathlib import Path
-import os
+import gc
 import types, collections
 from comfy.utils import ProgressBar, copy_to_param, set_attr_param
-from comfy.model_patcher import get_key_weight, string_to_seed
+from comfy.model_patcher import get_key_weight
 from comfy.lora import calculate_weight
-from comfy.model_management import cast_to_device
+
+try:
+    from comfy.utils import string_to_seed
+except:
+    from comfy.model_patcher import string_to_seed
+
 from comfy.float import stochastic_rounding
+from .custom_linear import remove_lora_from_module
 import folder_paths
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 log = logging.getLogger(__name__)
 
+import comfy.model_management as mm
+device = mm.get_torch_device()
+offload_device = mm.unet_offload_device()
+
+try:
+    from .gguf.gguf import GGUFParameter
+except:
+    pass
+
+COLOR_CODES = {
+    "reset": "\033[0m",
+    "red": "\033[31m",
+    "green": "\033[32m",
+    "yellow": "\033[33m",
+    "blue": "\033[34m",
+    "magenta": "\033[35m",
+    "cyan": "\033[36m",
+    "white": "\033[37m",
+}
+
+def color_text(text, color):
+    try:
+        return f"{COLOR_CODES.get(color, COLOR_CODES['reset'])}{text}{COLOR_CODES['reset']}"
+    except Exception:
+        return text
+
+class MetaParameter(torch.nn.Parameter):
+    def __new__(cls, dtype, quant_type=None):
+        data = torch.empty(0, dtype=dtype)
+        self = torch.nn.Parameter(data, requires_grad=False)
+        self.quant_type = quant_type
+        return self
+
+def offload_transformer(transformer, remove_lora=True):
+    transformer.teacache_state.clear_all()
+    transformer.magcache_state.clear_all()
+    transformer.easycache_state.clear_all()
+
+    if transformer.patched_linear:
+        for name, param in transformer.named_parameters():
+            if "loras" in name or "controlnet" in name:
+                continue
+            module = transformer
+            subnames = name.split('.')
+            for subname in subnames[:-1]:
+                module = getattr(module, subname)
+            attr_name = subnames[-1]
+            if param.data.is_floating_point():
+                meta_param = torch.nn.Parameter(torch.empty_like(param.data, device='meta'), requires_grad=False)
+                setattr(module, attr_name, meta_param)
+            elif isinstance(param.data, GGUFParameter):
+                quant_type = getattr(param, 'quant_type', None)
+                setattr(module, attr_name, MetaParameter(param.data.dtype, quant_type))
+            else:
+                pass
+        if remove_lora:
+            remove_lora_from_module(transformer)
+    else:
+        transformer.to(offload_device)
+
+    for block in transformer.blocks:
+        block.kv_cache = None
+        if transformer.audio_model is not None and hasattr(block, 'audio_block'):
+            block.audio_block = None
+
+    mm.soft_empty_cache()
+    gc.collect()
+
+
+def init_blockswap(transformer, block_swap_args, model):
+    if not transformer.patched_linear:
+        if block_swap_args is not None:
+            for name, param in transformer.named_parameters():
+                if "block" not in name or "control_adapter" in name or "face" in name:
+                    param.data = param.data.to(device)
+                elif block_swap_args["offload_txt_emb"] and "txt_emb" in name:
+                    param.data = param.data.to(offload_device)
+                elif block_swap_args["offload_img_emb"] and "img_emb" in name:
+                    param.data = param.data.to(offload_device)
+
+            transformer.block_swap(
+                block_swap_args["blocks_to_swap"] - 1 ,
+                block_swap_args["offload_txt_emb"],
+                block_swap_args["offload_img_emb"],
+                vace_blocks_to_swap = block_swap_args.get("vace_blocks_to_swap", None),
+            )
+        elif model["auto_cpu_offload"]:
+            for module in transformer.modules():
+                if hasattr(module, "offload"):
+                    module.offload()
+                if hasattr(module, "onload"):
+                    module.onload()
+            for block in transformer.blocks:
+                block.modulation = torch.nn.Parameter(block.modulation.to(device))
+            transformer.head.modulation = torch.nn.Parameter(transformer.head.modulation.to(device))
+        else:
+            transformer.to(device)
+
 def check_device_same(first_device, second_device):
     if first_device.type != second_device.type:
         return False
@@ -91,9 +195,9 @@ def set_module_tensor_to_device(module, tensor_name, device, value=None, dtype=N
             device = device_quantization
         if is_buffer:
             module._buffers[tensor_name] = new_value
-        elif value is not None or not check_device_same(torch.device(device), module._parameters[tensor_name].device):
+        elif value is not None or not check_device_same(device, module._parameters[tensor_name].device):
             param_cls = type(module._parameters[tensor_name])
-            new_value = param_cls(new_value, requires_grad=False).to(device)
+            new_value = param_cls(new_value, requires_grad=False)
             module._parameters[tensor_name] = new_value
 
     #if device != "cpu":
@@ -109,10 +213,8 @@ def check_diffusers_version():
         raise AssertionError("diffusers is not installed.")
 
 def print_memory(device, process="Sampling"):
-    memory = torch.cuda.memory_allocated(device) / 1024**3
     max_memory = torch.cuda.max_memory_allocated(device) / 1024**3
     max_reserved = torch.cuda.max_memory_reserved(device) / 1024**3
-    log.info(f"[{process}] Allocated memory: {memory=:.3f} GB")
     log.info(f"[{process}] Max allocated memory: {max_memory=:.3f} GB")
     log.info(f"[{process}] Max reserved memory: {max_reserved=:.3f} GB")
     #memory_summary = torch.cuda.memory_summary(device=device, abbreviated=False)
@@ -125,6 +227,18 @@ def get_module_memory_mb(module):
             memory += param.nelement() * param.element_size()
     return memory / (1024 * 1024)  # Convert to MB
 
+def get_module_memory_mb_per_device(module):
+    memory_per_device = {}
+    memory = 0
+    for param in module.parameters():
+        if param.data is not None:
+            device = str(param.device)
+            memory += param.nelement() * param.element_size()
+            memory_per_device[device] = memory_per_device.get(device, 0) + memory
+
+    memory_per_device = {dev: mem / (1024 * 1024) for dev, mem in memory_per_device.items()}
+    return memory_per_device
+
 def get_tensor_memory(tensor):
     memory_bytes = tensor.element_size() * tensor.nelement()
     return f"{memory_bytes / (1024 * 1024):.2f} MB"
@@ -140,7 +254,7 @@ def patch_weight_to_device(self, key, device_to=None, inplace_update=False, back
         self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update)
 
     if device_to is not None:
-        temp_weight = cast_to_device(weight, device_to, torch.float32, copy=True)
+        temp_weight = mm.cast_to_device(weight, device_to, torch.float32, copy=True)
     else:
         temp_weight = weight.to(torch.float32, copy=True)
     if convert_func is not None:
@@ -584,9 +698,9 @@ def check_duplicate_nodes():
     """Check ComfyUI custom_nodes directory for duplicate installations"""
     custom_nodes_dir = Path(folder_paths.folder_names_and_paths["custom_nodes"][0][0])
     current_path = Path(__file__).parent
-    
+
     wanvideo_dirs = []
-    
+
     # Check all directories in custom_nodes
     for path in custom_nodes_dir.iterdir():
         if (path.is_dir() and 
@@ -594,7 +708,7 @@ def check_duplicate_nodes():
             'wanvideo' in path.name.lower() and
             'wrapper' in path.name.lower()):
             wanvideo_dirs.append(str(path))
-    
+
     return wanvideo_dirs
 
 #https://github.com/temporalscorerescaling/TSR/
@@ -609,3 +723,55 @@ def temporal_score_rescaling(model_output, sample, timestep, k=1.0, tsr_sigma=0.
     if not t == 1.0:
         model_output = (ratio * ((1-t) * model_output + sample) - sample) / (1 - t)
     return model_output
+
+def match_and_blend_colors(
+    source_chunk: torch.Tensor,  # (C, T, H, W), range [-1, 1]
+    reference_image: torch.Tensor,  # (C, 1, H, W), range [-1, 1]
+    strength: float,
+) -> torch.Tensor:
+    import kornia
+    if strength == 0.0:
+        return source_chunk
+    source_chunk = source_chunk.unsqueeze(0)  # (1, C, T, H, W)
+
+    # shapes
+    B, C, T, H, W = source_chunk.shape
+    input_dtype = source_chunk.dtype
+
+    # [-1,1] -> [0,1]
+    src_01 = (source_chunk + 1.0) * 0.5
+    ref_01 = (reference_image + 1.0) * 0.5
+
+    src32 = src_01.to(torch.float32)
+    ref32 = ref_01.to(torch.float32)
+
+    # (B, C, T, H, W) -> (B*T, C, H, W)
+    src_bt = src32.permute(0, 2, 1, 3, 4).contiguous().view(B * T, C, H, W)
+    ref_bchw = ref32[:, :, 0, :, :].contiguous()
+
+    # RGB->Lab
+    src_lab = kornia.color.rgb_to_lab(src_bt)  # (B*T, C, H, W)
+    ref_lab = kornia.color.rgb_to_lab(ref_bchw)  # (B,   C, H, W)
+
+    src_lab_flat = src_lab.view(B * T, C, -1)  # (B*T, C, HW)
+    ref_lab_flat = ref_lab.view(B, C, -1)  # (B,   C, HW)
+    src_std, src_mean = torch.std_mean(src_lab_flat, dim=-1, keepdim=True, unbiased=False)
+    ref_std, ref_mean = torch.std_mean(ref_lab_flat, dim=-1, keepdim=True, unbiased=False)
+    src_std = src_std.clamp_min_(1e-6)
+
+    ref_mean_bt = ref_mean.repeat_interleave(T, dim=0)  # (B*T, C, 1)
+    ref_std_bt = ref_std.repeat_interleave(T, dim=0)  # (B*T, C, 1)
+
+    corrected_lab_flat = (src_lab_flat - src_mean) * (ref_std_bt / src_std) + ref_mean_bt
+    corrected_lab = corrected_lab_flat.view(B * T, C, H, W)
+
+    # Lab->RGB
+    corrected_rgb_01 = kornia.color.lab_to_rgb(corrected_lab)  # (B*T, C, H, W)
+
+    blended_rgb_01 = (1.0 - strength) * src_bt + strength * corrected_rgb_01
+
+    # (B, C, T, H, W)
+    blended_rgb_01 = blended_rgb_01.view(B, T, C, H, W).permute(0, 2, 1, 3, 4).contiguous()
+
+    # [0,1] -> [-1,1]
+    return (blended_rgb_01 * 2.0 - 1.0)[0].to(dtype=input_dtype)

wanvideo/modules/attention.py

@@ -80,9 +80,9 @@ except:
 try:
     from ...ultravico.sageattn.core import sage_attention as sageattn_ultravico
     @torch.library.custom_op("wanvideo::sageattn_ultravico", mutates_args=())
-    def sageattn_func_ultravico(qkv: List[torch.Tensor], attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, multi_factor: float = 0.9
+    def sageattn_func_ultravico(qkv: List[torch.Tensor], attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, multi_factor: float = 0.9, frame_tokens: int = 1536
     ) -> torch.Tensor:
-        return sageattn_ultravico(qkv, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, multi_factor=multi_factor)
+        return sageattn_ultravico(qkv, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, multi_factor=multi_factor, frame_tokens=frame_tokens)
 
     @sageattn_func_ultravico.register_fake
     def _(qkv, attn_mask=None, dropout_p=0.0, is_causal=False, multi_factor=0.9):
@@ -94,7 +94,7 @@ except:
 
 def attention(q, k, v, q_lens=None, k_lens=None, max_seqlen_q=None, max_seqlen_k=None, dropout_p=0.,
     softmax_scale=None, q_scale=None, causal=False,  window_size=(-1, -1), deterministic=False, dtype=torch.bfloat16,
-    attention_mode='sdpa', attn_mask=None, multi_factor=0.9, heads=128):
+    attention_mode='sdpa', attn_mask=None, transformer_options={}, frame_tokens=1536, heads=128):
     if "flash" in attention_mode:
         return flash_attention(q, k, v, q_lens=q_lens, k_lens=k_lens, dropout_p=dropout_p, softmax_scale=softmax_scale,
             q_scale=q_scale, causal=causal, window_size=window_size, deterministic=deterministic, dtype=dtype, version=2 if attention_mode == 'flash_attn_2' else 3,
@@ -108,7 +108,7 @@ def attention(q, k, v, q_lens=None, k_lens=None, max_seqlen_q=None, max_seqlen_k
     elif attention_mode == 'sageattn':
         return sageattn_func(q, k, v, tensor_layout="NHD").contiguous()
     elif attention_mode == 'sageattn_ultravico':
-        return sageattn_func_ultravico([q, k, v], multi_factor=multi_factor).contiguous()
+        return sageattn_func_ultravico([q, k, v], multi_factor=transformer_options.get("ultravico_alpha", 0.9), frame_tokens=frame_tokens).contiguous()
     elif attention_mode == 'comfy':
         return optimized_attention(q.transpose(1,2), k.transpose(1,2), v.transpose(1,2), heads=heads, skip_reshape=True)
     else: # sdpa

wanvideo/modules/model.py

@@ -467,7 +467,7 @@ class WanSelfAttention(nn.Module):
         v = (self.v(x) + self.v_loras(x)).view(b, s, n, d)
         return q, k, v
 
-    def forward(self, q, k, v, seq_lens, lynx_ref_feature=None, lynx_ref_scale=1.0, attention_mode_override=None, onetoall_ref=None, onetoall_ref_scale=1.0):
+    def forward(self, q, k, v, seq_lens, transformer_options={}, attention_mode_override=None, lynx_ref_feature=None, lynx_ref_scale=1.0, onetoall_ref=None, onetoall_ref_scale=1.0, frame_tokens=1536):
         r"""
         Args:
             x(Tensor): Shape [B, L, num_heads, C / num_heads]
@@ -482,7 +482,7 @@ class WanSelfAttention(nn.Module):
         if self.ref_adapter is not None and lynx_ref_feature is not None:
             ref_x = self.ref_adapter(self, q, lynx_ref_feature)
 
-        x = attention(q, k, v, k_lens=seq_lens, attention_mode=attention_mode, heads=self.num_heads)
+        x = attention(q, k, v, k_lens=seq_lens, attention_mode=attention_mode, heads=self.num_heads, frame_tokens=frame_tokens, transformer_options=transformer_options)
 
         if self.ref_adapter is not None and lynx_ref_feature is not None:
             x = x.add(ref_x, alpha=lynx_ref_scale)
@@ -497,7 +497,7 @@ class WanSelfAttention(nn.Module):
         attention_mode = self.attention_mode
         if attention_mode_override is not None:
             attention_mode = attention_mode_override
-        
+
         # Concatenate main and IP keys/values for main attention
         full_k = torch.cat([k, k_ip], dim=1)
         full_v = torch.cat([v, v_ip], dim=1)
@@ -558,39 +558,53 @@ class WanSelfAttention(nn.Module):
         # output
         return self.o(x.flatten(2))
 
-    def normalized_attention_guidance(self, b, n, d, q, context, nag_context=None, nag_params={}):
+    def nag_attention(self, b, n, d, q, context, nag_context=None):
+        k_positive = self.norm_k(self.k(context).to(self.norm_k.weight.dtype)).view(b, -1, n, d).to(q.dtype)
+        v_positive = self.v(context).view(b, -1, n, d)
+        x_positive = attention(q, k_positive, v_positive, attention_mode=self.attention_mode, heads=self.num_heads)
+        del k_positive, v_positive
+
+        k_negative = self.norm_k(self.k(nag_context).to(self.norm_k.weight.dtype)).view(b, -1, n, d).to(q.dtype)
+        v_negative = self.v(nag_context).view(b, -1, n, d)
+        x_negative = attention(q, k_negative, v_negative, attention_mode=self.attention_mode, heads=self.num_heads)
+        del k_negative, v_negative
+
+        return x_positive.flatten(2), x_negative.flatten(2)
+
+    def normalized_attention_guidance(self, x_positive, x_negative,nag_params={}):
         # NAG text attention
-        context_positive = context
-        context_negative = nag_context
         nag_scale = nag_params['nag_scale']
         nag_alpha = nag_params['nag_alpha']
         nag_tau = nag_params['nag_tau']
+        inplace = nag_params.get('inplace', True)
 
-        k_positive = self.norm_k(self.k(context_positive).to(self.norm_k.weight.dtype)).view(b, -1, n, d).to(q.dtype)
-        v_positive = self.v(context_positive).view(b, -1, n, d)
-        k_negative = self.norm_k(self.k(context_negative).to(self.norm_k.weight.dtype)).view(b, -1, n, d).to(q.dtype)
-        v_negative = self.v(context_negative).view(b, -1, n, d)
+        if inplace:
+            nag_guidance = x_negative.mul_(nag_scale - 1).neg_().add_(x_positive, alpha=nag_scale)
+        else:
+            nag_guidance = x_positive * nag_scale - x_negative * (nag_scale - 1)
+        del x_negative
 
-        x_positive = attention(q, k_positive, v_positive, attention_mode=self.attention_mode, heads=self.num_heads)
-        x_positive = x_positive.flatten(2)
-
-        x_negative = attention(q, k_negative, v_negative, attention_mode=self.attention_mode, heads=self.num_heads)
-        x_negative = x_negative.flatten(2)
-
-        nag_guidance = x_positive * nag_scale - x_negative * (nag_scale - 1)
-        
         norm_positive = torch.norm(x_positive, p=1, dim=-1, keepdim=True)
         norm_guidance = torch.norm(nag_guidance, p=1, dim=-1, keepdim=True)
-        
+
         scale = norm_guidance / norm_positive
-        scale = torch.nan_to_num(scale, nan=10.0)
-        
+        torch.nan_to_num_(scale, nan=10.0)
         mask = scale > nag_tau
+        del scale
+
         adjustment = (norm_positive * nag_tau) / (norm_guidance + 1e-7)
-        nag_guidance = torch.where(mask, nag_guidance * adjustment, nag_guidance)
+        del norm_positive, norm_guidance
+
+        nag_guidance.mul_(torch.where(mask, adjustment, 1.0))
         del mask, adjustment
-        
-        return nag_guidance * nag_alpha + x_positive * (1 - nag_alpha)
+
+        if inplace:
+            nag_guidance.sub_(x_positive).mul_(nag_alpha).add_(x_positive)
+        else:
+            nag_guidance = nag_guidance * nag_alpha + x_positive * (1 - nag_alpha)
+        del x_positive
+
+        return nag_guidance
 
 class LoRALinearLayer(nn.Module):
     def __init__(
@@ -633,22 +647,25 @@ class WanT2VCrossAttention(WanSelfAttention):
     def forward(self, x, context, grid_sizes=None, clip_embed=None, audio_proj=None, audio_scale=1.0,
                 num_latent_frames=21, nag_params={}, nag_context=None, rope_func="comfy",
                 inner_t=None, inner_c=None, cross_freqs=None,
-                adapter_proj=None, adapter_attn_mask=None, ip_scale=1.0, orig_seq_len=None, lynx_x_ip=None, lynx_ip_scale=1.0, num_cond_latents=None, **kwargs):
+                adapter_proj=None, adapter_attn_mask=None, ip_scale=1.0, orig_seq_len=None, lynx_x_ip=None, lynx_ip_scale=1.0, longcat_num_cond_latents=None, **kwargs):
         b, n, d = x.size(0), self.num_heads, self.head_dim
         s = x.size(1)
         # compute query
         is_longcat = x.shape[-1] == 4096
 
         if is_longcat:
-            if num_cond_latents is not None and num_cond_latents > 0:
-                num_cond_latents_thw = num_cond_latents * (s // num_latent_frames)
+            if longcat_num_cond_latents is not None and longcat_num_cond_latents > 0:
+                num_cond_latents_thw = longcat_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:
             q = self.norm_q(self.q(x).to(self.norm_q.weight.dtype),num_chunks=2 if rope_func == "comfy_chunked" else 1).to(x.dtype).view(b, -1, n, d)
 
         if nag_context is not None:
-            x = self.normalized_attention_guidance(b, n, d, q, context, nag_context, nag_params)
+            x_positive, x_negative = self.nag_attention(b, n, d, q, context, nag_context)
+            del q
+            x = self.normalized_attention_guidance(x_positive, x_negative, nag_params)
+            del x_positive, x_negative
         else:
             if is_longcat:
                 k = self.norm_k(self.k(context).to(self.norm_k.weight.dtype).view(b, -1, n, d)).to(x.dtype)
@@ -712,7 +729,7 @@ class WanT2VCrossAttention(WanSelfAttention):
 
             x = x.add(target_x)
 
-        if is_longcat and num_cond_latents is not None and num_cond_latents > 0:
+        if is_longcat and longcat_num_cond_latents > 0:
             return torch.cat([torch.zeros((b, num_cond_latents_thw, x.shape[-1]), dtype=x.dtype, device=x.device), self.o(x)], dim=1).contiguous()
 
         return self.o(x)
@@ -752,7 +769,8 @@ class WanI2VCrossAttention(WanSelfAttention):
             k_img = self.norm_k_img(self.k_img(clip_embed).to(self.norm_k_img.weight.dtype)).view(b, -1, n, d).to(x.dtype)
             v_img = self.v_img(clip_embed).view(b, -1, n, d)
             img_x = attention(q, k_img, v_img, attention_mode=self.attention_mode, heads=self.num_heads).flatten(2)
-            x = x_text + img_x
+            x_text.add_(img_x)
+            x = x_text
         else:
             x = x_text
 
@@ -914,7 +932,7 @@ class WanAttentionBlock(nn.Module):
             from ...LongCat.layers import FeedForwardSwiGLU
             mlp_ratio = 4
             self.ffn = FeedForwardSwiGLU(dim=self.dim, hidden_dim=int(self.dim * mlp_ratio))
-        
+
         # modulation
         if not is_longcat:
             self.modulation = nn.Parameter(torch.randn(1, 6, out_features) / in_features**0.5)
@@ -1003,9 +1021,10 @@ class WanAttentionBlock(nn.Module):
         humo_audio_input=None, humo_audio_scale=1.0, #humo audio
         lynx_x_ip=None, lynx_ref_feature=None, lynx_ip_scale=1.0, lynx_ref_scale=1.0, #lynx
         x_ovi=None, e_ovi=None, freqs_ovi=None, context_ovi=None, seq_lens_ovi=None, grid_sizes_ovi=None,
-        num_cond_latents=None, #longcat image cond amount
+        longcat_num_cond_latents=0, longcat_avatar_options=None, #longcat image cond amount
         x_onetoall_ref=None, onetoall_freqs=None, onetoall_ref=None, onetoall_ref_scale=1.0, #one-to-all
         e_tr=None, tr_num=0, tr_start=0, #token replacement
+        attention_mode_override=None, frame_tokens=None, transformer_options={}
     ):
         r"""
         Args:
@@ -1015,6 +1034,11 @@ class WanAttentionBlock(nn.Module):
             grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
             freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
         """
+        input_dtype = x.dtype
+        B, N, C = x.shape
+        T = num_latent_frames
+        is_longcat = C == 4096
+
         zero_timestep = len(e) == 2
         if zero_timestep: #s2v zero timestep
             self.seg_idx = e[1]
@@ -1030,11 +1054,10 @@ class WanAttentionBlock(nn.Module):
             tr_end = tr_start + (tr_num or 0)
 
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.get_mod(e.to(x.device), self.modulation)
+        if multitalk_audio_embedding is not None and is_longcat:
+            audio_shift_mca, audio_scale_mca, audio_gate_mca = self.audio_modulation(e[:, longcat_num_cond_latents:]).unsqueeze(2).chunk(3, dim=-1)
         del e
-        input_dtype = x.dtype
-        B, N, C = x.shape
-        T = num_latent_frames
-        is_longcat = C == 4096
+
         if is_longcat:
             input_x = self.modulate(self.norm1(x.view(B, T, -1, C).to(shift_msa.dtype)), shift_msa, scale_msa, seg_idx=self.seg_idx).to(input_dtype).view(B, N, C)
         elif use_token_replace:
@@ -1146,6 +1169,10 @@ class WanAttentionBlock(nn.Module):
         if enhance_enabled:
             feta_scores = get_feta_scores(q, k)
 
+        if self.attention_mode == "sageattn_3" and attention_mode_override is None:
+            if current_step != 0 and not last_step:
+                attention_mode_override = "sageattn"
+
         #self-attention
         split_attn = (context is not None
                       and (context.shape[0] > 1 or (clip_embed is not None and clip_embed.shape[0] > 1))
@@ -1157,19 +1184,14 @@ class WanAttentionBlock(nn.Module):
             y = self.self_attn.forward_split(q, k, v, seq_lens, grid_sizes, seq_chunks)
         elif ref_target_masks is not None: #multi/infinite talk
             y, x_ref_attn_map = self.self_attn.forward_multitalk(q, k, v, seq_lens, grid_sizes, ref_target_masks)
-        elif self.attention_mode == "radial_sage_attention":
+        elif self.attention_mode == "radial_sage_attention" or attention_mode_override is not None and attention_mode_override == "radial_sage_attention":
             if self.dense_block or self.dense_timesteps is not None and current_step < self.dense_timesteps:
                 if self.dense_attention_mode == "sparse_sage_attn":
                     y = self.self_attn.forward_radial(q, k, v, dense_step=True)
                 else:
-                    y = self.self_attn.forward(q, k, v, seq_lens)
+                    y = self.self_attn.forward(q, k, v, seq_lens, attention_mode_override=attention_mode_override)
             else:
                 y = self.self_attn.forward_radial(q, k, v, dense_step=False)
-        elif self.attention_mode == "sageattn_3":
-            if current_step != 0 and not last_step:
-                y = self.self_attn.forward(q, k, v, seq_lens, attention_mode_override="sageattn_3")
-            else:
-                y = self.self_attn.forward(q, k, v, seq_lens, attention_mode_override="sageattn")
         elif x_ip is not None and self.kv_cache is None: #stand-in
             # First pass: cache IP keys/values and compute attention
             self.kv_cache = {"k_ip": k_ip.detach(), "v_ip": v_ip.detach()}
@@ -1180,23 +1202,69 @@ class WanAttentionBlock(nn.Module):
             v_ip = self.kv_cache["v_ip"]
             full_k = torch.cat([k, k_ip], dim=1)
             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 // num_latent_frames)
-            # process the condition tokens
-            x_cond = self.self_attn.forward(
-                q[:, :num_cond_latents_thw].contiguous(),
-                k[:, :num_cond_latents_thw].contiguous(),
-                v[:, :num_cond_latents_thw].contiguous(),
-                seq_lens)
-            # process the noise tokens
-            x_noise = self.self_attn.forward(q[:, num_cond_latents_thw:].contiguous(), k, v, seq_lens)
-            # merge x_cond and x_noise
-            y = torch.cat([x_cond, x_noise], dim=1).contiguous()
-        else:
-            y = self.self_attn.forward(q, k, v, seq_lens, lynx_ref_feature=lynx_ref_feature, lynx_ref_scale=lynx_ref_scale, onetoall_ref=onetoall_ref, onetoall_ref_scale=onetoall_ref_scale)
+            y = self.self_attn.forward(q, full_k, full_v, seq_lens, attention_mode_override=attention_mode_override)
+        elif is_longcat and longcat_num_cond_latents > 0:
+            if longcat_num_cond_latents == 1:
+                num_cond_latents_thw = longcat_num_cond_latents * (N // num_latent_frames)
+                # process the noise tokens
+                x_noise = self.self_attn.forward(q[:, num_cond_latents_thw:].contiguous(), k, v, seq_lens, attention_mode_override=attention_mode_override, transformer_options=transformer_options)
+                # process the condition tokens
+                x_cond = self.self_attn.forward(
+                    q[:, :num_cond_latents_thw].contiguous(),
+                    k[:, :num_cond_latents_thw].contiguous(),
+                    v[:, :num_cond_latents_thw].contiguous(),
+                    seq_lens, attention_mode_override=attention_mode_override, transformer_options=transformer_options)
+                # merge x_cond and x_noise
+                y = torch.cat([x_cond, x_noise], dim=1).contiguous()
+            elif longcat_num_cond_latents > 1: # video continuation
+                num_ref_latents_thw = (N // num_latent_frames)
+                num_cond_latents_thw = longcat_num_cond_latents * (N // num_latent_frames)
+                if not longcat_num_cond_latents == num_latent_frames:
+                    # process the noise tokens
+                    q_noise = q[:, num_cond_latents_thw:].contiguous()
+                    start_noise, end_noise, num_noisy_frames = 0, 0, num_latent_frames - longcat_num_cond_latents
+                    mask_frame_range = longcat_avatar_options["ref_mask_frame_range"]
+                    ref_img_index = longcat_avatar_options["ref_frame_index"]
+                    num_ref_latents = 1
+                    if mask_frame_range is not None and mask_frame_range > 0:
+                        start_noise = ref_img_index - mask_frame_range - longcat_num_cond_latents + num_ref_latents
+                        end_noise   = ref_img_index + mask_frame_range - longcat_num_cond_latents + num_ref_latents + 1
 
-        del q, k, v,
+                    if start_noise >= 0 and end_noise > start_noise and end_noise <= num_noisy_frames:
+                        # remove attention with the reference image in the target range, preventing repeated actions.
+
+                        start_pos = start_noise * (N // num_latent_frames)
+                        end_pos   = end_noise * (N // num_latent_frames)
+
+                        q_noise_front = q_noise[:, :start_pos].contiguous()
+                        q_noise_maskref = q_noise[:, start_pos:end_pos].contiguous()
+                        q_noise_back = q_noise[:, end_pos:].contiguous()
+                        k_non_ref = k[:, num_ref_latents_thw:].contiguous()
+                        v_non_ref = v[:, num_ref_latents_thw:].contiguous()
+
+                        x_noise_front = self.self_attn.forward(q_noise_front, k, v, seq_lens, attention_mode_override=attention_mode_override, transformer_options=transformer_options) # q_front has attention with ref + cond + noisy
+                        x_noise_back = self.self_attn.forward(q_noise_back, k, v, seq_lens, attention_mode_override=attention_mode_override, transformer_options=transformer_options) # q_back has attention with ref + cond + noisy
+                        x_noise_maskref = self.self_attn.forward(q_noise_maskref, k_non_ref, v_non_ref, seq_lens, attention_mode_override=attention_mode_override, transformer_options=transformer_options) # q_mask has attention with cond+noisy
+                        x_noise = torch.cat([x_noise_front, x_noise_maskref, x_noise_back], dim=1).contiguous()
+                    else:
+                        x_noise = self.self_attn.forward(q_noise, k, v, seq_lens, attention_mode_override=attention_mode_override, transformer_options=transformer_options)
+                # process the condition tokens
+                q_ref = q[:, :num_ref_latents_thw].contiguous()
+                k_ref = k[:, :num_ref_latents_thw].contiguous()
+                v_ref = v[:, :num_ref_latents_thw].contiguous()
+                q_cond = q[:, num_ref_latents_thw:num_cond_latents_thw].contiguous()
+                k_cond = k[:, num_ref_latents_thw:num_cond_latents_thw].contiguous()
+                v_cond = v[:, num_ref_latents_thw:num_cond_latents_thw].contiguous()
+                x_ref = self.self_attn.forward(q_ref, k_ref, v_ref, seq_lens, attention_mode_override=attention_mode_override, transformer_options=transformer_options)
+                x_cond = self.self_attn.forward(q_cond, k_cond, v_cond, seq_lens, attention_mode_override=attention_mode_override, transformer_options=transformer_options)
+
+                # merge x_cond and x_noise
+                y = torch.cat([x_ref, x_cond, x_noise], dim=1).contiguous()
+        else:
+            y = self.self_attn.forward(q, k, v, seq_lens, lynx_ref_feature=lynx_ref_feature, lynx_ref_scale=lynx_ref_scale,
+                                       onetoall_ref=onetoall_ref, onetoall_ref_scale=onetoall_ref_scale, attention_mode_override=attention_mode_override, transformer_options=transformer_options, frame_tokens=frame_tokens)
+
+        del q, k, v
 
         # FETA
         if enhance_enabled:
@@ -1230,7 +1298,7 @@ class WanAttentionBlock(nn.Module):
                     y[:, tr_end:] * gate_msa
                 ], dim=1).to(input_dtype)
             else:
-                x = x.addcmul(y, gate_msa)
+                x.addcmul_(y, gate_msa)
         del y, gate_msa
 
         # cross-attention & ffn function
@@ -1260,16 +1328,25 @@ class WanAttentionBlock(nn.Module):
                 x = self.split_cross_attn_ffn(x, context, shift_mlp, scale_mlp, gate_mlp, clip_embed, grid_sizes)
                 return x, x_ip, lynx_ref_feature, x_ovi
             else:
-                x = x + self.cross_attn(self.norm3(x.to(self.norm3.weight.dtype)).to(input_dtype), context, grid_sizes, clip_embed=clip_embed, audio_proj=audio_proj, audio_scale=audio_scale,
+                x += self.cross_attn(self.norm3(x.to(self.norm3.weight.dtype)).to(input_dtype), context, grid_sizes, clip_embed=clip_embed, audio_proj=audio_proj, audio_scale=audio_scale,
                                     num_latent_frames=num_latent_frames, nag_params=nag_params, nag_context=nag_context,
                                     rope_func=self.rope_func, inner_t=inner_t, inner_c=inner_c, cross_freqs=cross_freqs,
-                                    adapter_proj=adapter_proj, ip_scale=ip_scale, orig_seq_len=original_seq_len, lynx_x_ip=lynx_x_ip, lynx_ip_scale=lynx_ip_scale, num_cond_latents=num_cond_latents)
-                x = x.to(input_dtype)
+                                    adapter_proj=adapter_proj, ip_scale=ip_scale, orig_seq_len=original_seq_len, lynx_x_ip=lynx_x_ip, lynx_ip_scale=lynx_ip_scale, longcat_num_cond_latents=longcat_num_cond_latents).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.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)
+
+                    if is_longcat:
+                        audio_output_cond, x_audio = self.audio_cross_attn(self.norm_x(x.to(self.norm_x.weight.dtype)).to(input_dtype), multitalk_audio_embedding, num_latent_frames=num_latent_frames,
+                                                        num_cond_latents=longcat_num_cond_latents, x_ref_attn_map=x_ref_attn_map, human_num=human_num)
+                        x_audio = self.modulate(self.norm1(x_audio.view(B, T-longcat_num_cond_latents, -1, C).to(audio_shift_mca.dtype)), audio_shift_mca, audio_scale_mca, seg_idx=self.seg_idx).to(input_dtype).view(B, -1, C)
+                        x_audio = (x_audio.view(B, T-longcat_num_cond_latents, -1, C).float() * audio_gate_mca).to(input_dtype).view(B, -1, C)
+                        if audio_output_cond is not None:
+                            x_audio = torch.cat([audio_output_cond, x_audio], dim=1).contiguous()
+                    else:
+                        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.add_(x_audio, alpha=audio_scale)
+                    del x_audio
 
                 # MTV-Crafter Motion Attention
                 if self.use_motion_attn and mtv_motion_tokens is not None and mtv_motion_rotary_emb is not None:
@@ -1282,7 +1359,7 @@ class WanAttentionBlock(nn.Module):
 
 
         # ffn
-        if self.rope_func == "comfy_chunked":
+        if self.rope_func == "comfy_chunked" and not is_longcat and not use_token_replace and not zero_timestep:
             mod_x = torch.addcmul(shift_mlp, self.norm2(x.to(shift_mlp.dtype)), 1 + scale_mlp)
             x_ffn = self.ffn_chunked(mod_x)
         else:
@@ -1308,7 +1385,7 @@ class WanAttentionBlock(nn.Module):
                     mod_x = torch.addcmul(shift_mlp, self.norm2(x.to(shift_mlp.dtype)), 1 + scale_mlp)
 
                 del shift_mlp, scale_mlp
-                x_ffn = self.ffn_chunked(mod_x.to(input_dtype), num_chunks=1)
+                x_ffn = self.ffn_chunked(mod_x.to(input_dtype), num_chunks=2 if is_longcat else 1)
                 del mod_x
 
         # gate_mlp
@@ -1960,24 +2037,24 @@ class WanModel(torch.nn.Module):
     def block_swap(self, blocks_to_swap, offload_txt_emb=False, offload_img_emb=False, vace_blocks_to_swap=None, prefetch_blocks=0, block_swap_debug=False):
         # Clamp blocks_to_swap to valid range
         blocks_to_swap = max(0, min(blocks_to_swap, len(self.blocks)))
-        
+
         log.info(f"Swapping {blocks_to_swap} transformer blocks")
         self.blocks_to_swap = blocks_to_swap
         self.prefetch_blocks = prefetch_blocks
         self.block_swap_debug = block_swap_debug
-        
+
         self.offload_img_emb = offload_img_emb
         self.offload_txt_emb = offload_txt_emb
 
         total_offload_memory = 0
         total_main_memory = 0
-        
+
         # Calculate the index where swapping starts
         swap_start_idx = len(self.blocks) - blocks_to_swap
-       
+
         for b, block in tqdm(enumerate(self.blocks), total=len(self.blocks), desc="Initializing block swap"):
             block_memory = get_module_memory_mb(block)
-            
+
             if b < swap_start_idx:
                 block.to(self.main_device)
                 total_main_memory += block_memory
@@ -1992,13 +2069,13 @@ class WanModel(torch.nn.Module):
             # Clamp vace_blocks_to_swap to valid range
             vace_blocks_to_swap = max(0, min(vace_blocks_to_swap, len(self.vace_blocks)))
             self.vace_blocks_to_swap = vace_blocks_to_swap
-            
+
             # Calculate the index where VACE swapping starts
             vace_swap_start_idx = len(self.vace_blocks) - vace_blocks_to_swap
 
             for b, block in tqdm(enumerate(self.vace_blocks), total=len(self.vace_blocks), desc="Initializing vace block swap"):
                 block_memory = get_module_memory_mb(block)
-                
+
                 if b < vace_swap_start_idx:
                     block.to(self.main_device)
                     total_main_memory += block_memory
@@ -2009,13 +2086,13 @@ class WanModel(torch.nn.Module):
         mm.soft_empty_cache()
         gc.collect()
 
-        log.info("----------------------")
-        log.info(f"Block swap memory summary:")
+        log.info("-" * 25)
+        log.info("Block swap memory summary:")
         log.info(f"Transformer blocks on {self.offload_device}: {total_offload_memory:.2f}MB")
         log.info(f"Transformer blocks on {self.main_device}: {total_main_memory:.2f}MB")
         log.info(f"Total memory used by transformer blocks: {(total_offload_memory + total_main_memory):.2f}MB")
         log.info(f"Non-blocking memory transfer: {self.use_non_blocking}")
-        log.info("----------------------")
+        log.info("-" * 25)
 
     def forward_vace(
         self,
@@ -2128,7 +2205,8 @@ class WanModel(torch.nn.Module):
 
 
     def rope_encode_comfy(self, t, h, w, freq_offset=0, t_start=0, ref_frame_shape=None, pose_frame_shape=None,
-                          steps_t=None, steps_h=None, steps_w=None, ntk_alphas=[1,1,1], device=None, dtype=None):
+                          steps_t=None, steps_h=None, steps_w=None, ntk_alphas=[1,1,1], device=None, dtype=None,
+                          ref_frame_index=10, longcat_num_ref_latents=0, num_memory_frames=3, rope_negative_offset=5):
 
         patch_size = self.patch_size
         t_len = ((t + (patch_size[0] // 2)) // patch_size[0])
@@ -2144,7 +2222,27 @@ class WanModel(torch.nn.Module):
 
         # Main frames position IDs
         img_ids = torch.zeros((steps_t, steps_h, steps_w, 3), device=device, dtype=dtype)
-        img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.linspace(t_start+freq_offset, t_start+freq_offset + (t_len - 1), steps=steps_t, device=device, dtype=dtype).reshape(-1, 1, 1)
+
+        if longcat_num_ref_latents > 0:
+            # Create temporal grid with ref_frame_index prepended, followed by sequential frames
+            grid_t = torch.cat([
+                torch.tensor([ref_frame_index], dtype=dtype, device=device),
+                torch.arange(0, steps_t - longcat_num_ref_latents, dtype=dtype, device=device)
+            ], dim=0)
+            img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + grid_t.reshape(-1, 1, 1)
+        elif num_memory_frames > 0 and rope_negative_offset > 0:
+            # Negative RoPE shift for memory frames
+            # Memory frames get negative indices: {-f_m*S, -(f_m-1)*S, ..., -S}
+            # Current video frames start from 0: {0, 1, ..., f-1}
+            memory_indices = torch.arange(-num_memory_frames * rope_negative_offset, 0, rope_negative_offset, dtype=dtype, device=device)
+            current_indices = torch.arange(0, steps_t - num_memory_frames, dtype=dtype, device=device)
+            grid_t = torch.cat([memory_indices, current_indices], dim=0)
+            log.info(f"{num_memory_frames} memory frames, temporal rope positions: {grid_t}")
+            img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + grid_t.reshape(-1, 1, 1)
+        else:
+            # Standard temporal encoding
+            img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.linspace(t_start+freq_offset, t_start+freq_offset + (t_len - 1), steps=steps_t, device=device, dtype=dtype).reshape(-1, 1, 1)
+
         img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.linspace(freq_offset, freq_offset + (h_len - 1), steps=steps_h, device=device, dtype=dtype).reshape(1, -1, 1)
         img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.linspace(freq_offset, freq_offset + (w_len - 1), steps=steps_w, device=device, dtype=dtype).reshape(1, 1, -1)
         img_ids = img_ids.reshape(1, -1, img_ids.shape[-1])
@@ -2243,11 +2341,15 @@ class WanModel(torch.nn.Module):
         lynx_embeds=None,
         x_ovi=None, seq_len_ovi=None, ovi_negative_text_embeds=None,
         flashvsr_LQ_latent=None, flashvsr_strength=1.0,
-        num_cond_latents=None,
+        longcat_num_cond_latents=0, longcat_num_ref_latents=0, longcat_avatar_options=None, # for LongCat
         add_text_emb=None,
         sdancer_input=None,  # SteadyDancer
         one_to_all_input=None, one_to_all_controlnet_strength=0.0, # One-to-All
         scail_input=None,  # SCAIL pose
+        dual_control_input=None,  # LongVie2 dual controlnet
+        transformer_options={},
+        rope_negative_offset=0,
+        num_memory_frames=0,
     ):
         r"""
         Forward pass through the diffusion model
@@ -2474,6 +2576,16 @@ class WanModel(torch.nn.Module):
         x = [u.flatten(2).transpose(1, 2) for u in x]
         self.original_seq_len = x[0].shape[1]
 
+        prev_latent = None
+        if dual_control_input is not None:
+            prev_latent = dual_control_input.get("prev_latent", None)
+            if prev_latent is not None:
+                F += prev_latent.shape[2]
+                prev_x = [self.original_patch_embedding(u.unsqueeze(0).to(torch.float32)).to(x[0].dtype) for u in prev_latent]
+                prev_x = [u.flatten(2).transpose(1, 2).to(self.base_dtype) for u in prev_x]
+                seq_len += prev_x[0].shape[1]
+                x = [torch.cat([u, v], dim=1) for u, v in zip(prev_x, x)]
+
         # SCAIL pose
         if scail_input is not None:
             scail_pose_latents = scail_input.get("pose_latent", None)
@@ -2482,6 +2594,7 @@ class WanModel(torch.nn.Module):
                 scail_x = [u.flatten(2).transpose(1, 2) * scail_input.get("pose_strength", 1) for u in scail_x]
                 x = [torch.cat([u, v], dim=1) for u, v in zip(x, scail_x)]
                 seq_len += scail_x[0].shape[1]
+                del scail_x
                 pose_frame_shape = scail_pose_latents.shape
 
         seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.int32)
@@ -2559,6 +2672,9 @@ class WanModel(torch.nn.Module):
                 tuple(pose_frame_shape) if pose_frame_shape is not None else None,
                 self.rope_embedder.k,
                 tuple(ntk_alphas),
+                longcat_num_ref_latents,
+                rope_negative_offset,
+                num_memory_frames,
             )
 
             # Check cache using key comparison
@@ -2567,16 +2683,19 @@ class WanModel(torch.nn.Module):
                 self.cached_key == cache_key):
                 freqs = self.cached_freqs
             else:
-                log.info("Generating new RoPE frequencies")
                 freqs = self.rope_encode_comfy(
                     F, H, W,
                     freq_offset=freq_offset,
                     ntk_alphas=ntk_alphas,
                     ref_frame_shape=ref_frame_shape,
                     pose_frame_shape=pose_frame_shape,
+                    longcat_num_ref_latents=longcat_num_ref_latents,
+                    rope_negative_offset=rope_negative_offset,
+                    num_memory_frames=num_memory_frames,
                     device=x.device,
                     dtype=x.dtype
                 )
+                tqdm.write("Generated new RoPE frequencies")
 
                 if s2v_ref_latent is not None:
                     freqs_ref = self.rope_encode_comfy(
@@ -2641,8 +2760,8 @@ class WanModel(torch.nn.Module):
             if len(t.shape) == 1:
                 t = t.unsqueeze(1).expand(-1, F) # [B, T]
             self.time_embedding.to(torch.float32)
-            e = e0 = self.time_embedding(t.float().flatten(), dtype=torch.float32).reshape(1, F, -1)
-
+            e = e0 = self.time_embedding(t.float().flatten(), dtype=torch.float32)#.reshape(1, F, -1)
+            e = e0 = e0.reshape(1, F, -1)
 
         if self.audio_model is not None:
             #if t.dim() == 1:
@@ -2762,6 +2881,44 @@ class WanModel(torch.nn.Module):
             chunked_self_attention = False
             seq_chunks = 0
 
+        # dual control
+        if dual_control_input is not None and dual_control_input["start_percent"] <= current_step_percentage <= dual_control_input["end_percent"]:
+            dense_latent = dual_control_input["dense_input_latent"]
+            print("dense_latent shape:", dense_latent.shape)
+            sparse_latent = dual_control_input["sparse_input_latent"]
+            if dense_latent is None and sparse_latent is None:
+                raise ValueError("At least one of dense_input_latent or sparse_input_latent must be provided in dual_control_input")
+
+            if dense_latent is not None:
+                dense_x = [self.original_patch_embedding(u.unsqueeze(0).to(torch.float32)).to(x[0].dtype) for u in dense_latent]
+                dense_x = [u.flatten(2).transpose(1, 2).to(self.base_dtype) for u in dense_x]
+                dense = self.dual_controller.control_initial_combine_linear_dense(dense_x[0])
+
+            if sparse_latent is not None:
+                sparse_x = [self.original_patch_embedding(u.unsqueeze(0).to(torch.float32)).to(x[0].dtype) for u in sparse_latent]
+                sparse_x = [u.flatten(2).transpose(1, 2).to(self.base_dtype) for u in sparse_x]
+                sparse = self.dual_controller.control_initial_combine_linear_sparse(sparse_x[0])
+
+            if dense_latent is None:
+                dense = torch.zeros_like(sparse)
+            elif sparse_latent is None:
+                sparse = torch.zeros_like(dense)
+
+            control_context = clip_fea_control = None
+            if context != []:
+                control_context = self.dual_controller.control_text_linear(context)
+                if clip_embed is not None:
+                    clip_fea_control = self.dual_controller.control_text_linear(clip_embed)
+            control_t_mod = self.dual_controller.control_t_mod(e0)
+
+            control_freqs = torch.cat([
+                self.dual_controller_freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+                self.dual_controller_freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+                self.dual_controller_freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+            ], dim=-1).reshape(f * h * w, 1, -1).to(x.device)
+        else:
+            dual_control_input = None
+
         # MultiTalk
         if multitalk_audio is not None:
             self.multitalk_audio_proj.to(self.main_device)
@@ -2777,10 +2934,24 @@ class WanModel(torch.nn.Module):
             latter_middle_frame_audio_emb = latter_frame_audio_emb[:, :, 1:-1, middle_index:middle_index+1, ...] 
             latter_middle_frame_audio_emb = rearrange(latter_middle_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") 
             latter_frame_audio_emb_s = torch.concat([latter_first_frame_audio_emb, latter_middle_frame_audio_emb, latter_last_frame_audio_emb], dim=2) 
-            multitalk_audio_embedding = self.multitalk_audio_proj(first_frame_audio_emb_s, latter_frame_audio_emb_s) 
-            human_num = len(multitalk_audio_embedding)
-            multitalk_audio_embedding = torch.concat(multitalk_audio_embedding.split(1), dim=2).to(self.base_dtype)
+            multitalk_audio_embedding = self.multitalk_audio_proj(first_frame_audio_emb_s, latter_frame_audio_emb_s)
             self.multitalk_audio_proj.to(self.offload_device)
+            human_num = len(multitalk_audio_embedding)
+
+            # LongCat-Avatar specific
+            if longcat_num_ref_latents > 0:
+                audio_start_ref = multitalk_audio_embedding[:, [0], :, :] # padding
+                multitalk_audio_embedding = torch.cat([audio_start_ref, multitalk_audio_embedding], dim=1).contiguous()
+
+            if longcat_num_cond_latents > 0:
+                multitalk_audio_embedding = multitalk_audio_embedding[:, (-F // self.patch_size[0]):]
+
+            if ref_target_masks is not None:
+                multitalk_audio_embedding = torch.concat(multitalk_audio_embedding.split(1), dim=2).to(self.base_dtype)
+                multitalk_audio_embedding = multitalk_audio_embedding.squeeze(0)
+            else:
+                multitalk_audio_embedding = rearrange(multitalk_audio_embedding, "b t n c -> (b t) n c")
+
 
         # convert ref_target_masks to token_ref_target_masks
         token_ref_target_masks = None
@@ -2953,6 +3124,7 @@ class WanModel(torch.nn.Module):
                 camera_embed=camera_embed,
                 audio_proj=audio_proj,
                 num_latent_frames = F,
+                frame_tokens=x.shape[1] // F,
                 original_seq_len=self.original_seq_len,
                 enhance_enabled=enhance_enabled,
                 audio_scale=audio_scale,
@@ -2974,11 +3146,13 @@ class WanModel(torch.nn.Module):
                 lynx_x_ip=lynx_x_ip,
                 lynx_ip_scale=lynx_ip_scale,
                 lynx_ref_scale=lynx_ref_scale,
-                num_cond_latents=num_cond_latents,
+                longcat_num_cond_latents=longcat_num_cond_latents,
+                longcat_avatar_options=longcat_avatar_options,
                 onetoall_ref_scale=onetoall_ref_scale,
                 e_tr=e0_token_replace if use_token_replace else None,
                 tr_start=token_replace_start,
                 tr_num=replace_token_num,
+                transformer_options=transformer_options
             )
             if self.audio_model is not None:
                 kwargs['e_ovi'] = e0_ovi.to(self.base_dtype)
@@ -3038,8 +3212,22 @@ class WanModel(torch.nn.Module):
             if lynx_ref_buffer is None and lynx_ref_feature_extractor:
                 lynx_ref_buffer = {}
 
+            attn_override_blocks = attention_mode = None
+            attention_mode_override_active = False
+            attention_mode_override = transformer_options.get("attention_mode_override", None)
+            if attention_mode_override is not None:
+                attn_override_blocks = attention_mode_override.get("blocks", range(len(self.blocks)))
+                if attention_mode_override["start_step"] <= current_step < attention_mode_override["end_step"]:
+                    attention_mode_override_active = True
+                    if attention_mode_override["verbose"]:
+                        tqdm.write(f"Applying attention mode override: {attention_mode_override['mode']} at step {current_step} on blocks: {attn_override_blocks if attn_override_blocks is not None else 'all'}")
+
             for b, block in enumerate(self.blocks):
                 mm.throw_exception_if_processing_interrupted()
+                if attention_mode_override_active and b in attn_override_blocks:
+                    attention_mode = attention_mode_override['mode']
+                else:
+                    attention_mode = None
                 block_idx = f"{b:02d}"
                 if lynx_ref_buffer is not None and not lynx_ref_feature_extractor:
                     lynx_ref_feature = lynx_ref_buffer.get(block_idx, None)
@@ -3083,9 +3271,21 @@ class WanModel(torch.nn.Module):
                     x_onetoall_ref = onetoall_ref_block_samples[b // interval_ref]
 
                 # ---run block----#
-                x, x_ip, lynx_ref_feature, x_ovi = block(x, x_ip=x_ip, lynx_ref_feature=lynx_ref_feature, x_ovi=x_ovi, x_onetoall_ref=x_onetoall_ref, onetoall_freqs=onetoall_freqs, **kwargs)
+                x, x_ip, lynx_ref_feature, x_ovi = block(x, x_ip=x_ip, lynx_ref_feature=lynx_ref_feature, x_ovi=x_ovi, x_onetoall_ref=x_onetoall_ref, onetoall_freqs=onetoall_freqs, attention_mode_override=attention_mode, **kwargs)
                 # ---post block----#
 
+                # dual controlnet
+                if dual_control_input is not None and (hasattr(block, "control_blocks_dense") or hasattr(block, "control_blocks_sparse")):
+                    if dense_latent is not None and hasattr(block, "control_blocks_dense"):
+                        dense = block.control_blocks_dense(dense, control_context, control_t_mod, control_freqs, clip_fea=clip_fea_control)
+                    if sparse_latent is not None and hasattr(block, "control_blocks_sparse"):
+                        sparse = block.control_blocks_sparse(sparse, control_context, control_t_mod, control_freqs, clip_fea=clip_fea_control)
+
+                    if prev_latent is not None:
+                        x[:, -self.original_seq_len:] += block.control_combine_linears(dense + sparse) * dual_control_input["strength"]
+                    else:
+                        x += block.control_combine_linears(dense + sparse) * dual_control_input["strength"]
+
                 if self.audio_injector is not None and s2v_audio_input is not None:
                     x = self.audio_injector_forward(b, x, merged_audio_emb, scale=s2v_audio_scale) #s2v
                 if block.has_face_fuser_block and motion_vec is not None:
@@ -3188,8 +3388,10 @@ class WanModel(torch.nn.Module):
         #    x = x[:, :self.original_seq_len]
             #grid_sizes = torch.stack([torch.tensor([u[0] - 1, u[1], u[2]]) for u in grid_sizes]).to(grid_sizes.device)
 
-
-        x = x[:, :self.original_seq_len]
+        if prev_latent is not None:
+            x = x[:, -self.original_seq_len:]
+        else:
+            x = x[:, :self.original_seq_len]
 
         x = self.head(x, e.to(x.device), temp_length=F,
                       e_tr=e_token_replace.to(x.device) if use_token_replace else None, tr_start=token_replace_start, tr_num=replace_token_num)

wanvideo/schedulers/__init__.py

@@ -1,9 +1,11 @@
 import torch
+import numpy as np
 from .fm_solvers import (FlowDPMSolverMultistepScheduler)
 from .fm_solvers_unipc import FlowUniPCMultistepScheduler
 from .basic_flowmatch import FlowMatchScheduler
 from .flowmatch_pusa import FlowMatchSchedulerPusa
 from .flowmatch_res_multistep import FlowMatchSchedulerResMultistep
+from .ersde_scheduler import ERSDEScheduler
 from .scheduling_flow_match_lcm import FlowMatchLCMScheduler
 from .fm_sa_ode import FlowMatchSAODEStableScheduler
 from .fm_rcm import rCMFlowMatchScheduler
@@ -25,6 +27,7 @@ scheduler_list = [
     "deis",
     "lcm", "lcm/beta",
     "res_multistep",
+    "er_sde",
     "flowmatch_causvid",
     "flowmatch_distill",
     "flowmatch_pusa",
@@ -39,7 +42,7 @@ def _apply_custom_sigmas(sample_scheduler, sigmas, device):
     sample_scheduler.timesteps = (sample_scheduler.sigmas[:-1] * 1000).to(torch.int64).to(device)
     sample_scheduler.num_inference_steps = len(sample_scheduler.timesteps)
 
-def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer_dim=5120, flowedit_args=None, denoise_strength=1.0, sigmas=None, log_timesteps=False, **kwargs):
+def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer_dim=5120, denoise_strength=1.0, sigmas=None, log_timesteps=False, enhance_hf=False, **kwargs):
     timesteps = None
     if sigmas is not None:
         steps = len(sigmas) - 1
@@ -136,6 +139,12 @@ def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transfo
             sample_scheduler.set_timesteps(steps, denoising_strength=denoise_strength)
         else:
             _apply_custom_sigmas(sample_scheduler, sigmas, device)
+    elif scheduler == 'er_sde':
+        sample_scheduler = ERSDEScheduler(shift=shift)
+        if sigmas is None:
+            sample_scheduler.set_timesteps(steps, denoising_strength=denoise_strength)
+        else:
+            _apply_custom_sigmas(sample_scheduler, sigmas, device)
     elif "sa_ode_stable" in scheduler:
         sample_scheduler = FlowMatchSAODEStableScheduler(shift=shift, **kwargs)
         if sigmas is None:
@@ -152,6 +161,18 @@ def get_scheduler(scheduler, steps, start_step, end_step, shift, device, transfo
     if timesteps is None:
         timesteps = sample_scheduler.timesteps
 
+    if enhance_hf:
+        num_tail_uniform_steps = max(3, min(15, int(len(timesteps) * 0.2))) # Use 20% of steps for uniform tail (minimum 3, maximum 15)
+        tail_uniform_start = float(timesteps.max()) * 0.5 # Split at 50% of the timestep range
+        tail_uniform_end = 0
+
+        timesteps_uniform_tail = list(np.linspace(tail_uniform_start, tail_uniform_end, num_tail_uniform_steps, dtype=np.float32, endpoint=(tail_uniform_end != 0)))
+        timesteps_uniform_tail = [torch.tensor(t, device=device).unsqueeze(0) for t in timesteps_uniform_tail]
+        filtered_timesteps = [timestep.unsqueeze(0).to(device)  for timestep in timesteps if timestep > tail_uniform_start]
+        timesteps = torch.cat(filtered_timesteps + timesteps_uniform_tail)
+        sample_scheduler.timesteps = timesteps
+        sample_scheduler.sigmas = torch.cat([timesteps / 1000, torch.zeros(1, device=timesteps.device)])
+
     steps = len(timesteps)
     if (isinstance(start_step, int) and end_step != -1 and start_step >= end_step) or (not isinstance(start_step, int) and start_step != -1 and end_step >= start_step):
         raise ValueError("start_step must be less than end_step")

wanvideo/schedulers/ersde_scheduler.py

@@ -0,0 +1,154 @@
+import torch
+
+class ERSDEScheduler():
+    """Extended Reverse-Time SDE solver (VP ER-SDE-Solver-3).
+
+    Based on: arXiv: https://arxiv.org/abs/2309.06169
+    Code reference: https://github.com/QinpengCui/ER-SDE-Solver/blob/main/er_sde_solver.py
+    """
+
+    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0,
+                 sigma_max=1.0, sigma_min=0.003 / 1.002, max_stage=3, s_noise=1.0,
+                 num_integration_points=200):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.max_stage = max_stage
+        self.s_noise = s_noise
+        self.num_integration_points = num_integration_points
+        self.set_timesteps(num_inference_steps)
+        self.old_denoised = None
+        self.old_denoised_d = None
+        self.step_index = 0
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, sigmas=None):
+        """Generate the full sigma schedule (from max to min)."""
+        full_sigmas = torch.linspace(self.sigma_max, self.sigma_min, self.num_train_timesteps)
+        ss = len(full_sigmas) / num_inference_steps
+        if sigmas is None:
+            sigmas = []
+            for x in range(num_inference_steps):
+                idx = int(round(x * ss))
+                sigmas.append(float(full_sigmas[idx]))
+            sigmas.append(0.0)
+        self.sigmas = torch.FloatTensor(sigmas)
+        self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
+        self.timesteps = self.sigmas[:-1] * self.num_train_timesteps
+        self.step_index = 0
+        self.old_denoised = None
+        self.old_denoised_d = None
+
+    def default_er_sde_noise_scaler(self, x):
+        return x * ((x ** 0.3).exp() + 10.0)
+
+    def step(self, model_output, timestep, sample, generator):
+
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+
+        self.sigmas = self.sigmas.to(model_output.device)
+        self.timesteps = self.timesteps.to(model_output.device)
+
+        if timestep.ndim == 0:
+            timestep_id = torch.argmin((self.timesteps - timestep).abs(), dim=0)
+        else:
+            timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+
+        noise_scaler = self.default_er_sde_noise_scaler
+
+        # Get current and next sigma
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        if (timestep_id + 1 >= len(self.sigmas)).any():
+            sigma_next = torch.zeros_like(sigma)
+        else:
+            sigma_next = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
+
+        er_lambda_s = sigma
+        er_lambda_t = sigma_next
+
+        # Calculate alpha values
+        alpha_s = sigma / (er_lambda_s + 1e-10)
+        alpha_t = sigma_next / (er_lambda_t + 1e-10)
+        r_alpha = alpha_t / (alpha_s + 1e-10)
+
+        # Denoised prediction (x_0 estimate)
+        denoised = sample - sigma * model_output
+
+        # Determine which stage to use
+        stage_used = min(self.max_stage, self.step_index + 1)
+
+        if sigma_next == 0 or (sigma_next == 0.0).all():
+            # Final step - return denoised
+            x = denoised
+        else:
+            r = noise_scaler(er_lambda_t) / (noise_scaler(er_lambda_s) + 1e-10)
+
+            # Stage 1: Euler step
+            x = r_alpha * r * sample + alpha_t * (1 - r) * denoised
+
+            if stage_used >= 2 and self.old_denoised is not None:
+                dt = er_lambda_t - er_lambda_s
+                lambda_step_size = -dt / self.num_integration_points
+
+                # Create integration points
+                point_indice = torch.arange(0, self.num_integration_points,
+                                           dtype=torch.float32, device=sample.device)
+                lambda_pos = er_lambda_t + point_indice * lambda_step_size
+                scaled_pos = noise_scaler(lambda_pos)
+
+                # Stage 2: Second-order correction
+                s = torch.sum(1 / (scaled_pos + 1e-10)) * lambda_step_size
+
+                # Get previous sigma for derivative calculation
+                if timestep_id > 0:
+                    sigma_prev = self.sigmas[timestep_id - 1].reshape(-1, 1, 1, 1)
+                    er_lambda_prev = sigma_prev
+                else:
+                    er_lambda_prev = er_lambda_s
+
+                denoised_d = (denoised - self.old_denoised) / ((er_lambda_s - er_lambda_prev) + 1e-10)
+                x = x + alpha_t * (dt + s * noise_scaler(er_lambda_t)) * denoised_d
+
+                if stage_used >= 3 and self.old_denoised_d is not None:
+                    # Stage 3: Third-order correction
+                    s_u = torch.sum((lambda_pos - er_lambda_s) / (scaled_pos + 1e-10)) * lambda_step_size
+
+                    # Get sigma from two steps ago
+                    if timestep_id > 1:
+                        sigma_prev_prev = self.sigmas[timestep_id - 2].reshape(-1, 1, 1, 1)
+                        er_lambda_prev_prev = sigma_prev_prev
+                    else:
+                        er_lambda_prev_prev = er_lambda_prev
+
+                    denoised_u = (denoised_d - self.old_denoised_d) / (((er_lambda_s - er_lambda_prev_prev) / 2) + 1e-10)
+                    x = x + alpha_t * ((dt ** 2) / 2 + s_u * noise_scaler(er_lambda_t)) * denoised_u
+
+                self.old_denoised_d = denoised_d
+
+            # Add stochastic noise
+            if self.s_noise > 0:
+                noise_term = (er_lambda_t ** 2 - er_lambda_s ** 2 * r ** 2).sqrt()
+                noise_term = torch.nan_to_num(noise_term, nan=0.0)
+                noise = torch.randn(*x.shape, dtype=torch.float32, device=torch.device("cpu"), generator=generator).to(x)
+                x = x + alpha_t * noise * self.s_noise * noise_term
+
+        # Store current denoised for next iteration
+        self.old_denoised = denoised
+        self.step_index += 1
+
+        return x
+
+    def add_noise(self, original_samples, noise, timestep):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)

wanvideo/schedulers/flowmatch_res_multistep.py

@@ -35,9 +35,7 @@ class FlowMatchSchedulerResMultistep():
         self.sigmas = torch.FloatTensor(sigmas)
         self.sigmas = self.shift * self.sigmas / \
              (1 + (self.shift - 1) * self.sigmas)
-        self.timesteps = self.sigmas * self.num_train_timesteps
-        #print(f"Timesteps: {self.timesteps}, Sigmas: {self.sigmas}")
-
+        self.timesteps = self.sigmas[:-1] * self.num_train_timesteps
 
     def step(self, model_output, timestep, sample):
         if timestep.ndim == 2:
@@ -48,14 +46,14 @@ class FlowMatchSchedulerResMultistep():
             timestep_id = torch.argmin((self.timesteps - timestep).abs(), dim=0)
         else:
             timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
-        
+
         sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
         sigma_prev = self.sigmas[timestep_id - 1].reshape(-1, 1, 1, 1) if timestep_id > 0 else sigma
         if (timestep_id + 1 >= len(self.sigmas)).any():
             sigma_next = torch.tensor(0)
         else:
             sigma_next = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
-      
+
         x0_pred = (sample - sigma * model_output)
 
         if sigma_next == 0 or self.prev_model_output is None:
@@ -73,7 +71,7 @@ class FlowMatchSchedulerResMultistep():
         self.old_sigma_next = sigma_next
         self.prev_model_output = x0_pred
         return x
-        
+
 
     def add_noise(self, original_samples, noise, timestep):
         """

wanvideo/wan_video_vae.py

@@ -989,7 +989,8 @@ class VideoVAE_(nn.Module):
                  mean=None,
                  inv_std=None,
                  pruning_rate=0.0,
-                 cpu_cache=False):
+                 cpu_cache=False,
+                 verbose=False):
         super().__init__()
         self.dim = dim
         self.z_dim = z_dim
@@ -1000,6 +1001,7 @@ class VideoVAE_(nn.Module):
         self.temperal_upsample = temperal_downsample[::-1]
         self.mean = mean
         self.inv_std = inv_std
+        self.verbose = verbose
 
         # modules
         self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,
@@ -1085,12 +1087,13 @@ class VideoVAE_(nn.Module):
             std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
             eps = torch.randn_like(std)
             return mu + std * eps
-        try:
-            log.info(f"WanVAE encoded input:{input_shape} to {out.shape}")
-            print_memory(device, process="WanVAE encode")
-            torch.cuda.reset_peak_memory_stats(device)
-        except:
-            pass
+        if self.verbose:
+            try:
+                log.info(f"WanVAE encoded input:{input_shape} to {out.shape}")
+                print_memory(device, process="WanVAE encode")
+                torch.cuda.reset_peak_memory_stats(device)
+            except:
+                pass
         return mu
 
 
@@ -1137,7 +1140,7 @@ class VideoVAE_(nn.Module):
         except:
             pass
         x = self.conv2(z)
-        for i in range(iter_):
+        for i in tqdm(range(iter_), desc="WanVAE decoding frames", disable=not pbar):
             self._conv_idx = [0]
             if i == 0:
                 out = self.decoder(x[:, :, i:i + 1, :, :],
@@ -1154,12 +1157,13 @@ class VideoVAE_(nn.Module):
         if pbar:
             pbar.update_absolute(0)
         self.clear_cache()
-        try:
-            log.info(f"WanVAE decoded input:{input_shape} to {out.shape}")
-            print_memory(device, process="WanVAE decode")
-            torch.cuda.reset_peak_memory_stats(device)
-        except:
-            pass
+        if self.verbose:
+            try:
+                log.info(f"WanVAE decoded input:{input_shape} to {out.shape}")
+                print_memory(device, process="WanVAE decode")
+                torch.cuda.reset_peak_memory_stats(device)
+            except:
+                pass
         return out
 
     def reparameterize(self, mu, log_var):
@@ -1186,12 +1190,12 @@ class VideoVAE_(nn.Module):
 
 class WanVideoVAE(nn.Module):
 
-    def __init__(self, z_dim=16, dtype=torch.float32, pruning_rate=0.0, cpu_cache=False):
+    def __init__(self, z_dim=16, dtype=torch.float32, pruning_rate=0.0, cpu_cache=False, verbose=False):
         super().__init__()
 
         self.dtype = dtype
         self.cpu_cache = cpu_cache
-
+        self.verbose = verbose
         mean = [
             -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
             0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
@@ -1205,7 +1209,7 @@ class WanVideoVAE(nn.Module):
         self.z_dim = z_dim
 
         # init model
-        self.model = VideoVAE_(z_dim=z_dim, mean=self.mean, inv_std=self.inv_std, pruning_rate=pruning_rate, cpu_cache=self.cpu_cache).eval().requires_grad_(False)
+        self.model = VideoVAE_(z_dim=z_dim, mean=self.mean, inv_std=self.inv_std, pruning_rate=pruning_rate, cpu_cache=self.cpu_cache, verbose=self.verbose).eval().requires_grad_(False)
         self.upsampling_factor = 8
 
 
@@ -1431,7 +1435,8 @@ class VideoVAE38_(VideoVAE_):
                  mean=None,
                  inv_std=None,
                  pruning_rate=0.0,
-                 cpu_cache=False):
+                 cpu_cache=False,
+                 verbose=False):
         super(VideoVAE_, self).__init__()
         self.dim = dim
         self.z_dim = z_dim
@@ -1444,6 +1449,7 @@ class VideoVAE38_(VideoVAE_):
         self.mean = mean
         self.inv_std = inv_std
         self.cpu_cache = cpu_cache
+        self.verbose = verbose
 
         # modules
         self.encoder = Encoder3d_38(dim, z_dim * 2, dim_mult, num_res_blocks,
@@ -1481,12 +1487,13 @@ class VideoVAE38_(VideoVAE_):
         mu = self.conv1(out).chunk(2, dim=1)[0]
         mu = (mu - self.mean.to(mu)) * self.inv_std.to(mu)
         self.clear_cache()
-        try:
-            log.info(f"WanVAE decoded input:{input_shape} to {out.shape}")
-            print_memory(device, process="WanVAE decode")
-            torch.cuda.reset_peak_memory_stats(device)
-        except:
-            pass
+        if self.verbose:
+            try:
+                log.info(f"WanVAE decoded input:{input_shape} to {out.shape}")
+                print_memory(device, process="WanVAE decode")
+                torch.cuda.reset_peak_memory_stats(device)
+            except:
+                pass
         return mu
 
 
@@ -1519,18 +1526,19 @@ class VideoVAE38_(VideoVAE_):
                 pbar.update(1)
         out = unpatchify(out, patch_size=2)
         self.clear_cache()
-        try:
-            log.info(f"WanVAE decoded input:{input_shape} to {out.shape}")
-            print_memory(device, process="WanVAE decode")
-            torch.cuda.reset_peak_memory_stats(device)
-        except:
-            pass
+        if self.verbose:
+            try:
+                log.info(f"WanVAE decoded input:{input_shape} to {out.shape}")
+                print_memory(device, process="WanVAE decode")
+                torch.cuda.reset_peak_memory_stats(device)
+            except:
+                pass
         return out
 
 
 class WanVideoVAE38(WanVideoVAE):
 
-    def __init__(self, z_dim=48, dim=160, dtype=torch.bfloat16, pruning_rate=0.0, cpu_cache=False):
+    def __init__(self, z_dim=48, dim=160, dtype=torch.bfloat16, pruning_rate=0.0, cpu_cache=False, verbose=False):
         super(WanVideoVAE, self).__init__()
 
         mean = [
@@ -1554,7 +1562,8 @@ class WanVideoVAE38(WanVideoVAE):
         self.dtype = dtype
         self.z_dim = z_dim
         self.cpu_cache = cpu_cache
+        self.verbose = verbose
 
         # init model
-        self.model = VideoVAE38_(z_dim=z_dim, dim=dim, dtype=dtype, mean=self.mean, inv_std=self.inv_std, pruning_rate=pruning_rate, cpu_cache=cpu_cache).eval().requires_grad_(False)
-        self.upsampling_factor = 16
+        self.model = VideoVAE38_(z_dim=z_dim, dim=dim, dtype=dtype, mean=self.mean, inv_std=self.inv_std, pruning_rate=pruning_rate, cpu_cache=cpu_cache, verbose=verbose).eval().requires_grad_(False)
+        self.upsampling_factor = 16