WanMoveSVI_FLF_v2_node.py

import torch
import torch.nn.functional as F
import comfy.model_management
import comfy.utils
import comfy.latent_formats
import node_helpers
from comfy_api.latest import io

# --- WAN-MOVE HELPER FUNCTIONS ---

def create_pos_embeddings(
    pred_tracks: torch.Tensor, 
    pred_visibility: torch.Tensor, 
    downsample_ratios: list[int], 
    height: int, 
    width: int, 
    track_num: int = -1, 
    t_down_strategy: str = "sample"
):
    assert t_down_strategy in ["sample", "average"], "Invalid strategy for downsampling time dimension."

    t, n, _ = pred_tracks.shape
    t_down, h_down, w_down = downsample_ratios
    track_pos = - torch.ones(n, (t-1) // t_down + 1, 2, dtype=torch.long)

    if track_num == -1:
        track_num = n

    tracks_idx = torch.randperm(n)[:track_num]
    tracks = pred_tracks[:, tracks_idx]
    visibility = pred_visibility[:, tracks_idx]

    for t_idx in range(0, t, t_down):
        if t_down_strategy == "sample" or t_idx == 0:
            cur_tracks = tracks[t_idx] 
            cur_visibility = visibility[t_idx] 
        else:
            cur_tracks = tracks[t_idx:t_idx+t_down].mean(dim=0)
            cur_visibility = torch.any(visibility[t_idx:t_idx+t_down], dim=0)

        for i in range(track_num):
            if not cur_visibility[i] or cur_tracks[i][0] < 0 or cur_tracks[i][1] < 0 or cur_tracks[i][0] >= width or cur_tracks[i][1] >= height:
                continue
            x, y = cur_tracks[i]
            x, y = int(x // w_down), int(y // h_down)
            track_pos[i, t_idx // t_down, 0], track_pos[i, t_idx // t_down, 1] = y, x

    return track_pos 

def replace_feature(
    vae_feature: torch.Tensor,  
    track_pos: torch.Tensor,    
    strength: float = 1.0
) -> torch.Tensor:
    b, _, t, h, w = vae_feature.shape
    assert b == track_pos.shape[0], "Batch size mismatch."
    n = track_pos.shape[1]

    track_pos = track_pos[:, torch.randperm(n), :, :]

    current_pos = track_pos[:, :, 1:, :] 
    mask = (current_pos[..., 0] >= 0) & (current_pos[..., 1] >= 0) 

    valid_indices = mask.nonzero(as_tuple=False) 
    num_valid = valid_indices.shape[0]

    if num_valid == 0:
        return vae_feature

    batch_idx = valid_indices[:, 0]
    track_idx = valid_indices[:, 1]
    t_rel = valid_indices[:, 2]
    t_target = t_rel + 1  

    h_target = current_pos[batch_idx, track_idx, t_rel, 0].long()  
    w_target = current_pos[batch_idx, track_idx, t_rel, 1].long()
    h_source = track_pos[batch_idx, track_idx, 0, 0].long()
    w_source = track_pos[batch_idx, track_idx, 0, 1].long()

    src_features = vae_feature[batch_idx, :, 0, h_source, w_source]
    dst_features = vae_feature[batch_idx, :, t_target, h_target, w_target]

    vae_feature[batch_idx, :, t_target, h_target, w_target] = dst_features + (src_features - dst_features) * strength

    return vae_feature

# --- MAIN NODE ---

class WanMoveSVI_FLF_v2(io.ComfyNode):
    """
    Integrates Wan-Move trajectory tracking with Wan SVI and First-Last-Frame (FLF) control. Version 2.
    """
    
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="WanMoveSVI_FLF_v2",
            category="WanMoveSVI_FLF_v2",
            inputs=[
                io.Conditioning.Input("positive"),
                io.Conditioning.Input("negative"),
                io.Vae.Input("vae"),
                io.ClipVision.Input("clip_vision", optional=True),
                io.Image.Input("first_image", optional=True, tooltip="First image to anchor the generation."),
                io.Image.Input("last_image", optional=True, tooltip="Optional last image sequence."),
                io.Image.Input("anchor_image", optional=True, display_name="anchor_image (optional)", tooltip="Optional image used as the feature tracking source (anchor latent)."),
                io.Tracks.Input("tracks", optional=True),
                io.Latent.Input("prev_samples", optional=True, tooltip="Previous latents for SVI."),
                io.Mask.Input("custom_mask", optional=True, tooltip="Override the generated mask with a custom sequence."),
                io.Float.Input("move_strength", default=2.0, min=0.0, max=100.0, step=0.01),
                io.Int.Input("svi_latent_count", default=1, min=0, max=128, step=1, tooltip="How many previous latents for SVI to inject."),
                io.Int.Input("width", default=512, min=16, max=8192, step=16),
                io.Int.Input("height", default=512, min=16, max=8192, step=16),
                io.Int.Input("length", default=41, min=1, max=8192, step=4),
            ],
            outputs=[
                io.Conditioning.Output(display_name="positive"),
                io.Conditioning.Output(display_name="negative"),
                io.Latent.Output(display_name="latent"),
                io.Mask.Output(display_name="mask"),
            ],
        )

    @classmethod
    def execute(cls, positive, negative, vae, first_image, svi_latent_count, 
                move_strength, width, height, length,
                clip_vision=None, last_image=None, 
                prev_samples=None, tracks=None, anchor_image=None, custom_mask=None) -> io.NodeOutput:
        
        device = comfy.model_management.intermediate_device()
        batch_size = 1

        if first_image is None:
            raise ValueError("WanMove SVI Integration requires 'first_image'.")

        clip_vision_output = None
        if clip_vision is not None:
            clip_vision_output = clip_vision.encode_image(first_image)

        # 1. Resolve First Latent (temporal t=0)
        first_image_resized = comfy.utils.common_upscale(first_image[:1].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
        first_latent = vae.encode(first_image_resized[:, :, :, :3])

        if anchor_image is not None:
            anchor_image_resized = comfy.utils.common_upscale(anchor_image[:1].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
            anchor_latent = vae.encode(anchor_image_resized[:, :, :, :3])
        else:
            anchor_latent = first_latent.clone()

        B, C, _, H_latent, W_latent = first_latent.shape
        total_latents = ((length - 1) // 4) + 1
        empty_latent = torch.zeros([batch_size, 16, total_latents, H_latent, W_latent], device=device)

        # 2. Resolve Motion Latent (SVI)
        if prev_samples is not None and svi_latent_count > 0:
            t_prev = prev_samples["samples"].shape[2]
            start_idx = max(t_prev - svi_latent_count, 0)
            
            # CRUCIAL: Never take Latent 0 from prev_samples because it is a 1-frame latent,
            # and first_image is already taking the Latent 0 spot for the new sequence.
            if start_idx == 0:
                start_idx = 1
                
            if start_idx < t_prev:
                svi_latent = prev_samples["samples"][:, :, start_idx:].clone()
                svi_injected_count = svi_latent.shape[2]
            else:
                svi_latent = None
                svi_injected_count = 0
        else:
            svi_latent = None
            svi_injected_count = 0

        # 3. Generate Padding
        frames_to_pad_svi = total_latents - 1 - svi_injected_count
        frames_to_pad_wan = total_latents - 1

        # VAE Encode a 50% Gray video (Native Wan-Move padding)
        gray_video = torch.ones((length, height, width, 3), device=device, dtype=first_latent.dtype) * 0.5
        gray_latent = vae.encode(gray_video)
        
        svi_padding = gray_latent[:, :, -frames_to_pad_svi:] if frames_to_pad_svi > 0 else None
        wan_padding = gray_latent[:, :, -frames_to_pad_wan:] if frames_to_pad_wan > 0 else None

        # 4. Construct the Environments using first_latent at index 0
        if svi_padding is not None and svi_latent is not None:
            svi_base = torch.cat([first_latent, svi_latent, svi_padding], dim=2)
        elif svi_latent is not None:
            svi_base = torch.cat([first_latent, svi_latent], dim=2)[:, :, :total_latents]
        else:
            svi_base = torch.cat([first_latent, wan_padding], dim=2)

        wan_base = torch.cat([first_latent, wan_padding], dim=2)

        # 5. Apply Wan-Move Tracking using anchor_latent as the feature source
        if tracks is not None and move_strength > 0.0:
            tracks_path = tracks["track_path"][:length]
            num_tracks = tracks_path.shape[-2]
            track_visibility = tracks.get("track_visibility", torch.ones((length, num_tracks), dtype=torch.bool, device=device))

            track_pos = create_pos_embeddings(tracks_path, track_visibility, [4, 8, 8], height, width, track_num=num_tracks)
            track_pos = track_pos.unsqueeze(0)

            # We build the tracking base with anchor_latent at t=0 so replace_feature reads features from it
            wan_base_for_tracking = torch.cat([anchor_latent, wan_padding], dim=2)
            wan_tracked = replace_feature(wan_base_for_tracking, track_pos, move_strength)
            
            # Restore first_latent at t=0 so the temporal start remains consistent with first_image
            wan_tracked[:, :, 0] = first_latent[:, :, 0]
        else:
            wan_tracked = wan_base.clone()

        # 6. Apply Time-based Alpha Blend (Optimized for svi_blend_length = 1)
        blend_weights = torch.ones(total_latents, device=device, dtype=anchor_latent.dtype)
        if svi_injected_count > 0:
            # Latents bound to SVI remain untouched (0.0). All future latents jump to Wan-Move (1.0).
            blend_weights[1:svi_injected_count + 1] = 0.0

        blend_weights = blend_weights.view(1, 1, total_latents, 1, 1)

        # Crossfade
        concat_latent_image_pos = svi_base * (1.0 - blend_weights) + wan_tracked * blend_weights
        concat_latent_image_neg = svi_base * (1.0 - blend_weights) + wan_base * blend_weights

        # 7. Apply FLF-style Overwrite to Target End Frames
        last_t_fix = 0
        if last_image is not None:
            num_frames = last_image.shape[0]
            
            if num_frames == 1:
                # Single frame bypass: Encode directly
                last_image_to_encode = last_image
            else:
                # Animated batch: Prepend a repeat of the first frame to act as the dummy anchor
                dummy_anchor = last_image[0:1]
                last_image_to_encode = torch.cat([dummy_anchor, last_image], dim=0)

            last_image_resized = comfy.utils.common_upscale(last_image_to_encode.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
            last_latent = vae.encode(last_image_resized[:, :, :, :3])

            if (last_latent.shape[1] == C and 
                last_latent.shape[3] == H_latent and 
                last_latent.shape[4] == W_latent):
                
                if num_frames == 1:
                    valid_last_latents = last_latent
                else:
                    # Animated batch: Discard Latent 0 (the dummy anchor).
                    if last_latent.shape[2] > 1:
                        valid_last_latents = last_latent[:, :, 1:]
                    else:
                        valid_last_latents = last_latent

                T_valid = valid_last_latents.shape[2]
                
                # Ensure we don't overwrite the global anchor (Index 0 of main sequence)
                last_t_fix = min(T_valid, total_latents - 1)

                if last_t_fix > 0:
                    concat_latent_image_pos[:, :, -last_t_fix:] = valid_last_latents[:, :, -last_t_fix:]
                    concat_latent_image_neg[:, :, -last_t_fix:] = valid_last_latents[:, :, -last_t_fix:]
            else:
                last_t_fix = 0

       # 8. Apply Conditioning Mask
        mask = torch.ones((1, 1, total_latents, H_latent, W_latent), device=anchor_latent.device, dtype=anchor_latent.dtype)
        
        # Lock anchor frame and SVI latents based on svi_latent_count (minimum 1 to ensure anchor is always locked)
        svi_lock_count = max(1, svi_latent_count)
        mask[:, :, :svi_lock_count] = 0.0 

        # Default FLF behavior (Hard cut)
        if last_t_fix > 0:
            mask[:, :, -last_t_fix:] = 0.0

        # 9. Override with Custom Mask if provided
        if custom_mask is not None:
            c_mask = custom_mask.clone()
            if c_mask.dim() == 2:
                c_mask = c_mask.unsqueeze(0)
            c_mask = c_mask.unsqueeze(0).unsqueeze(0)
            
            mask = F.interpolate(c_mask, size=(total_latents, H_latent, W_latent), mode='trilinear', align_corners=False)
            mask = mask.to(device=anchor_latent.device, dtype=anchor_latent.dtype)

        out_mask = mask.squeeze(0).squeeze(0)

        # Set values to Conditioning API
        positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": concat_latent_image_pos, "concat_mask": mask})
        negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": concat_latent_image_neg, "concat_mask": mask})

        if clip_vision_output is not None:
            positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output})
            negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output})

        out_latent = {"samples": empty_latent}

        return io.NodeOutput(positive, negative, out_latent, out_mask)