wan/text2video.py

# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import gc
import logging
import math
import os
import random
import sys
import types
import cv2
from contextlib import contextmanager
from functools import partial
import torch.nn.functional as F
import torch
import torch.cuda.amp as amp
import torch.distributed as dist
from tqdm import tqdm
from .distributed.fsdp import shard_model
from .modules.model import WanModel
from .modules.t5 import T5EncoderModel
from .modules.vae import WanVAE
from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
                               get_sampling_sigmas, retrieve_timesteps)
from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
import numpy as np
from typing import Optional, Literal
import itertools

class WanT2V:

    def __init__(
        self,
        config,
        checkpoint_dir,
        device_id=0,
        rank=0,
        t5_fsdp=False,
        dit_fsdp=False,
        use_usp=False,
        t5_cpu=False,
    ):
        r"""
        Initializes the Wan text-to-video generation model components.

        Args:
            config (EasyDict):
                Object containing model parameters initialized from config.py
            checkpoint_dir (`str`):
                Path to directory containing model checkpoints
            device_id (`int`,  *optional*, defaults to 0):
                Id of target GPU device
            rank (`int`,  *optional*, defaults to 0):
                Process rank for distributed training
            t5_fsdp (`bool`, *optional*, defaults to False):
                Enable FSDP sharding for T5 model
            dit_fsdp (`bool`, *optional*, defaults to False):
                Enable FSDP sharding for DiT model
            use_usp (`bool`, *optional*, defaults to False):
                Enable distribution strategy of USP.
            t5_cpu (`bool`, *optional*, defaults to False):
                Whether to place T5 model on CPU. Only works without t5_fsdp.
        """
        self.device = torch.device(f"cuda:{device_id}")
        self.config = config
        self.rank = rank
        self.t5_cpu = t5_cpu

        self.num_train_timesteps = config.num_train_timesteps
        self.param_dtype = config.param_dtype

        shard_fn = partial(shard_model, device_id=device_id)
        self.text_encoder = T5EncoderModel(
            text_len=config.text_len,
            dtype=config.t5_dtype,
            device=torch.device('cpu'),
            checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
            tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
            shard_fn=shard_fn if t5_fsdp else None)

        self.vae_stride = config.vae_stride
        self.patch_size = config.patch_size
        self.vae = WanVAE(
            vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
            device=self.device)

        logging.info(f"Creating WanModel from {checkpoint_dir}")
        
        # condition
        self.model = WanModel.from_pretrained(checkpoint_dir)
        self.model.eval().requires_grad_(False)
        
        
        if use_usp:
            from xfuser.core.distributed import \
                get_sequence_parallel_world_size

            from .distributed.xdit_context_parallel import (usp_attn_forward,
                                                            usp_dit_forward)
            for block in self.model.blocks:
                block.self_attn.forward = types.MethodType(
                    usp_attn_forward, block.self_attn)
            self.model.forward = types.MethodType(usp_dit_forward, self.model)
            self.sp_size = get_sequence_parallel_world_size()
        else:
            self.sp_size = 1

        if dist.is_initialized():
            dist.barrier()
        if dit_fsdp:
            self.model = shard_fn(self.model)
        else:
            self.model.to(self.device)

        self.sample_neg_prompt = config.sample_neg_prompt


    def generate(self,
                 input_prompt,
                 size=(1280, 720),
                 frame_num=81,
                 shift=5.0,
                 sample_solver='unipc',
                 sampling_steps=50,
                 guide_scale=5.0,
                 n_prompt="",
                 seed=-1,
                 offload_model=True):
        r"""
        Generates video frames from text prompt using diffusion process.

        Args:
            input_prompt (`str`):
                Text prompt for content generation
            size (tupele[`int`], *optional*, defaults to (1280,720)):
                Controls video resolution, (width,height).
            frame_num (`int`, *optional*, defaults to 81):
                How many frames to sample from a video. The number should be 4n+1
            shift (`float`, *optional*, defaults to 5.0):
                Noise schedule shift parameter. Affects temporal dynamics
            sample_solver (`str`, *optional*, defaults to 'unipc'):
                Solver used to sample the video.
            sampling_steps (`int`, *optional*, defaults to 40):
                Number of diffusion sampling steps. Higher values improve quality but slow generation
            guide_scale (`float`, *optional*, defaults 5.0):
                Classifier-free guidance scale. Controls prompt adherence vs. creativity
            n_prompt (`str`, *optional*, defaults to ""):
                Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`
            seed (`int`, *optional*, defaults to -1):
                Random seed for noise generation. If -1, use random seed.
            offload_model (`bool`, *optional*, defaults to True):
                If True, offloads models to CPU during generation to save VRAM

        Returns:
            torch.Tensor:
                Generated video frames tensor. Dimensions: (C, N H, W) where:
                - C: Color channels (3 for RGB)
                - N: Number of frames (81)
                - H: Frame height (from size)
                - W: Frame width from size)
        """
        # preprocess
        F = frame_num
        target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1,
                        size[1] // self.vae_stride[1],
                        size[0] // self.vae_stride[2])

        seq_len = math.ceil((target_shape[2] * target_shape[3]) /
                            (self.patch_size[1] * self.patch_size[2]) *
                            target_shape[1] / self.sp_size) * self.sp_size

        if n_prompt == "":
            n_prompt = self.sample_neg_prompt
        seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
        seed_g = torch.Generator(device=self.device)
        seed_g.manual_seed(seed)

        if not self.t5_cpu:
            self.text_encoder.model.to(self.device)
            context = self.text_encoder([input_prompt], self.device)
            context_null = self.text_encoder([n_prompt], self.device)
            if offload_model:
                self.text_encoder.model.cpu()
        else:
            context = self.text_encoder([input_prompt], torch.device('cpu'))
            context_null = self.text_encoder([n_prompt], torch.device('cpu'))
            context = [t.to(self.device) for t in context]
            context_null = [t.to(self.device) for t in context_null]

        noise = [
            torch.randn(
                target_shape[0],
                target_shape[1],
                target_shape[2],
                target_shape[3],
                dtype=torch.float32,
                device=self.device,
                generator=seed_g)
        ]

        @contextmanager
        def noop_no_sync():
            yield

        no_sync = getattr(self.model, 'no_sync', noop_no_sync)

        # evaluation mode
        with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():

            if sample_solver == 'unipc':
                sample_scheduler = FlowUniPCMultistepScheduler(
                    num_train_timesteps=self.num_train_timesteps,
                    shift=1,
                    use_dynamic_shifting=False)
                sample_scheduler.set_timesteps(
                    sampling_steps, device=self.device, shift=shift)
                timesteps = sample_scheduler.timesteps
            elif sample_solver == 'dpm++':
                sample_scheduler = FlowDPMSolverMultistepScheduler(
                    num_train_timesteps=self.num_train_timesteps,
                    shift=1,
                    use_dynamic_shifting=False)
                sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)
                timesteps, _ = retrieve_timesteps(
                    sample_scheduler,
                    device=self.device,
                    sigmas=sampling_sigmas)
            else:
                raise NotImplementedError("Unsupported solver.")

            # sample videos
            latents = noise

            arg_c = {'context': context, 'seq_len': seq_len}
            arg_null = {'context': context_null, 'seq_len': seq_len}

            for _, t in enumerate(tqdm(timesteps)):
                latent_model_input = latents
                timestep = [t]

                timestep = torch.stack(timestep)

                self.model.to(self.device)
                noise_pred_cond, _ = self.model(
                    latent_model_input, t=timestep, **arg_c)
                noise_pred_uncond, _ = self.model(
                    latent_model_input, t=timestep, **arg_null)
                
                noise_pred_cond, noise_pred_uncond = noise_pred_cond[0], noise_pred_uncond[0]

                noise_pred = noise_pred_uncond + guide_scale * (
                    noise_pred_cond - noise_pred_uncond)

                temp_x0 = sample_scheduler.step(
                    noise_pred.unsqueeze(0),
                    t,
                    latents[0].unsqueeze(0),
                    return_dict=False,
                    generator=seed_g)[0]
                latents = [temp_x0.squeeze(0)]

            x0 = latents
            if offload_model:
                self.model.cpu()
            if self.rank == 0:
                videos = self.vae.decode(x0)

        del noise, latents
        del sample_scheduler
        if offload_model:
            gc.collect()
            torch.cuda.synchronize()
        if dist.is_initialized():
            dist.barrier()

        return videos[0] if self.rank == 0 else None
    

    def load_video_frames(self, video_path, size=(832, 480)):
        r"""
        Load video frames from the given path and preprocess them.

        Args:
            video_path (str): Path to the video file.
            size (tuple[`int`], *optional*, defaults to (1280,720)): Target resolution for resizing frames.

        Returns:
            torch.Tensor: Tensor of video frames with shape (frame_num, C, H, W).
        """
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"Cannot open video file: {video_path}")

        frames = []
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            # Convert BGR to RGB
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            # Resize frame to target size
            frame = cv2.resize(frame, size, interpolation=cv2.INTER_AREA)
            # Convert to tensor and normalize to [-1, 1]
            frame = torch.from_numpy(frame).float().permute(2, 0, 1) / 127.5 - 1.0
            # Convert to tensor and normailize to [0, 1]
            # frame = torch.from_numpy(frame).float().permute(2, 0, 1) / 255
            frames.append(frame)

        cap.release()
        if not frames:
            raise ValueError(f"No frames found in video: {video_path}")

        # Stack frames into a single tensor
        frames_tensor = torch.stack(frames).permute(1, 0, 2, 3).to(self.device)
        latents = self.vae.video_encode(frames_tensor)
        return latents # [C, F, H, W]


    def find_subtokens_range(self, source_tokens, target_tokens):
        """
        查找 target_tokens 在 source_tokens 中的位置，不包括 '
        返回起始索引和结束索引（类似 slice），找不到返回 None
        """
        valid_len = len(source_tokens)
        # 在 source 的有效范围内滑动窗口匹配 target
        for i in range(valid_len - len(target_tokens) + 1):
            if source_tokens[i:i + len(target_tokens)] == target_tokens:
                # return (i, i + len(target_tokens)-1)
                return list(range(*(i, i + len(target_tokens))))
            

        return None


    def create_binary_mask( # Renamed slightly for clarity
        self,
        attn_map: torch.Tensor,
        n: int, 
        pooling_mode: Literal['max', 'avg'] = 'max',
        threshold: Optional[float] = 0.5,
        threshold_method: Literal['fixed', 'otsu'] = 'fixed',
        normalize_per_slice=False
    ) -> torch.Tensor:

        # --- Preparation ---
        original_shape = attn_map.shape
        C, T, H, W = original_shape
        device = attn_map.device

        map_batched = attn_map.reshape(C * T, 1, H, W)

        padding = (n - 1) // 2

        if pooling_mode == 'max':
            smoothed_map_batched = torch.nn.functional.max_pool2d(map_batched, kernel_size=n, stride=1, padding=padding)
        elif pooling_mode == 'avg':
            smoothed_map_batched = torch.nn.functional.avg_pool2d(map_batched, kernel_size=n, stride=1, padding=padding)

        smoothed_map = smoothed_map_batched.squeeze(1).view(C, T, H, W)
        # smoothed_map has shape (C, T, H, W)

        map_to_binarize = smoothed_map_batched # Start with the smoothed map
        if normalize_per_slice:

            flat_map = map_to_binarize.view(C * T, -1)

            # Calculate min and max per slice (image in the batch C*T)
            min_vals = torch.min(flat_map, dim=1, keepdim=True)[0]
            max_vals = torch.max(flat_map, dim=1, keepdim=True)[0]

            # Calculate range, handle the case where min == max (flat slice)
            range_vals = max_vals - min_vals

            range_vals = torch.where(range_vals == 0,
                                    torch.tensor(1.0, device=device, dtype=map_to_binarize.dtype),
                                    range_vals)

            min_vals_b = min_vals.view(C * T, 1, 1, 1)
            range_vals_b = range_vals.view(C * T, 1, 1, 1)

            normalized_map_batched = (map_to_binarize - min_vals_b) / range_vals_b

            map_to_binarize = torch.clamp(normalized_map_batched, 0.0, 1.0)

        map_to_binarize = map_to_binarize.squeeze(1).view(C, T, H, W)
        binary_mask = torch.zeros_like(map_to_binarize, dtype=torch.bool, device=device)
        
        if threshold_method == 'fixed':
            if normalize_per_slice: 
                binary_mask = map_to_binarize > torch.mean(map_to_binarize).item()
            else:
                binary_mask = map_to_binarize > threshold

        return binary_mask.to(device=device, dtype=torch.float32)          
   

    def soften_mask_edges(
        self,
        binary_mask: torch.Tensor,
        decay_factor: float = 0.1
    ) -> torch.Tensor:
        """
        Softens the edges of a binary mask by assigning values to background pixels (0)
        based on their distance to the nearest foreground pixel (1).

        Pixels originally equal to 1 remain 1.
        Pixels originally equal to 0 get a value exp(-decay_factor * distance),
        where distance is the Euclidean distance to the nearest 1.
        This means pixels closer to the original mask edge get values closer to 1,
        and pixels farther away get values closer to 0.

        Args:
            binary_mask (torch.Tensor): The input binary mask. Expected to be a 4D
                                        tensor (C, T, H, W) with values 0.0 or 1.0.
            decay_factor (float): Controls how quickly the softened value decays
                                with distance. A larger value means a faster decay
                                (sharper transition near the edge), a smaller value
                                means a slower decay (softer, more spread-out transition).
                                Defaults to 0.1.

        Returns:
            torch.Tensor: A mask tensor with the same dimensions as the input,
                        where original mask areas are 1.0 and background areas
                        have softened values based on distance. Output is float32.

        Raises:
            TypeError: If binary_mask is not a PyTorch Tensor.
            ValueError: If binary_mask is not 4D.
            ImportError: If scipy is not installed.
        """
        from scipy.ndimage import distance_transform_edt
        # --- Input Validation ---
        if not isinstance(binary_mask, torch.Tensor):
            raise TypeError("binary_mask must be a PyTorch Tensor.")
        if binary_mask.ndim != 4:
            raise ValueError(f"Input binary_mask must be 4D (C, T, H, W), got {binary_mask.ndim}D")
        if not decay_factor > 0:
            raise ValueError("decay_factor must be positive.")

        # --- Preparation ---
        original_shape = binary_mask.shape
        C, T, H, W = original_shape
        device = binary_mask.device
        dtype = torch.float32 # Ensure output is float

        # Create an output tensor initialized with zeros
        softened_mask = torch.zeros_like(binary_mask, dtype=dtype, device=device)

        # --- Process each slice (C, T) independently ---
        for c in range(C):
            for t in range(T):
                # Extract the 2D slice
                mask_slice = binary_mask[c, t, :, :]

                # Move to CPU and convert to NumPy boolean array for SciPy
                # distance_transform_edt expects background (0) to be True
                # and foreground (1) to be False.
                inverted_mask_slice_np = (mask_slice == 0).cpu().numpy()

                # Compute Euclidean Distance Transform
                # distance_map_np contains the distance from each True pixel (background)
                # to the nearest False pixel (foreground/mask).
                # Pixels that were originally part of the mask (False in inverted) will have distance 0.
                distance_map_np = distance_transform_edt(inverted_mask_slice_np)

                # Convert distances to softened values (0 to 1) using exponential decay
                # exp(-k * distance). Larger distance -> smaller value.
                # Distance 0 (original mask) -> exp(0) = 1.
                # Add a small epsilon to distance before applying exp if you want to strictly avoid 1.0 in non-mask areas
                # but exp(-k*dist) naturally handles this.
                softened_values_np = np.exp(-decay_factor * distance_map_np)

                # Convert back to PyTorch tensor and move to the original device
                softened_values_slice = torch.from_numpy(softened_values_np).to(device=device, dtype=dtype)

                # --- Combine original mask and softened background ---
                # Use torch.where for clarity and efficiency:
                # Where the original mask was 1, keep 1.0.
                # Where the original mask was 0, use the calculated softened value.
                final_slice = torch.where(
                    mask_slice.bool(), # Condition: True where original mask was 1
                    torch.tensor(1.0, device=device, dtype=dtype), # Value if True
                    softened_values_slice # Value if False (use calculated softened value)
                )

                # Place the processed slice into the output tensor
                softened_mask[c, t, :, :] = final_slice

        return softened_mask                

    
    def generate_conflict_map( # Renamed slightly for clarity
        self,
        vsrc: torch.Tensor,
        vtar: torch.Tensor,
        normalize: bool = True, # Setting default to True as it simplifies tuning 'k' later
        norm_type: int = 2,
        epsilon: float = 1e-6
    ) -> torch.Tensor:
        """
        Generates a conflict map indicating the magnitude of difference between
        two velocity fields (vsrc and vtar) and returns it as a 4D tensor.

        The conflict at each spatial-temporal location is defined as the L-p norm
        (default L2, Euclidean distance) of the difference vector between vsrc and
        vtar along the channel dimension. The resulting scalar map (F, H, W) is
        then expanded to match the input shape (C, F, H, W) by repeating the
        scalar value across the channel dimension.

        Args:
            vsrc (torch.Tensor): The source velocity field tensor. Expected shape:
                                (Channels, Frames, Height, Width).
            vtar (torch.Tensor): The target velocity field tensor. Must have the
                                same shape and device as vsrc.
            normalize (bool, optional): If True, normalize the underlying 3D conflict
                                        map to the range [0, 1] using min-max scaling
                                        before expanding it to 4D. Recommended for
                                        easier tuning of 'k' in downstream functions.
                                        Defaults to True.
            norm_type (int, optional): The order of the norm (p-norm). Default is 2 (L2 norm).
                                    Use 1 for L1 norm (Manhattan distance), etc.
            epsilon (float, optional): A small value added to the denominator during
                                    normalization to prevent division by zero if
                                    all conflict values are identical. Defaults to 1e-6.

        Returns:
            torch.Tensor: The conflict map tensor, expanded to 4D.
                        Shape: (Channels, Frames, Height, Width). The value is
                        uniform across the channel dimension for each (F,H,W).

        Raises:
            TypeError: If inputs are not PyTorch tensors.
            ValueError: If input tensors do not have the same shape or are not 4D.
        """
        # --- Input Validation ---
        if not isinstance(vsrc, torch.Tensor) or not isinstance(vtar, torch.Tensor):
            raise TypeError("Inputs vsrc and vtar must be PyTorch tensors.")

        if vsrc.ndim != 4 or vtar.ndim != 4:
            raise ValueError(f"Input tensors must be 4D (C, F, H, W), "
                            f"got shapes {vsrc.shape} and {vtar.shape}")

        if vsrc.shape != vtar.shape:
            raise ValueError(f"Input tensors vsrc and vtar must have the same shape, "
                            f"got {vsrc.shape} and {vtar.shape}")

        if vsrc.device != vtar.device:
            logging.warning(f"Input tensors are on different devices ({vsrc.device}, {vtar.device}). "
                            f"Proceeding with calculations, but ensure this is intended.")

        # --- Conflict Calculation ---
        vsrc_float = vsrc.float()
        vtar_float = vtar.float()
        num_channels = vsrc.shape[0]

        difference = vtar_float - vsrc_float

        # Calculate the norm along the channel dimension -> shape (F, H, W)
        conflict_map_3d = torch.norm(difference, p=norm_type, dim=0)

        # --- Optional Normalization (applied to the 3D map) ---
        if normalize:
            # Avoid normalization issues on empty tensors
            if conflict_map_3d.numel() > 0:
                min_val = torch.min(conflict_map_3d)
                max_val = torch.max(conflict_map_3d)
                denominator = max_val - min_val

                if denominator < epsilon:
                    logging.warning(f"Conflict map values are nearly constant ({min_val.item()} to {max_val.item()}). "
                                    f"Normalization results in a map of all zeros.")
                    conflict_map_3d = torch.zeros_like(conflict_map_3d)
                else:
                    conflict_map_3d = (conflict_map_3d - min_val) / denominator
            else:
                logging.warning("Conflict map has zero elements, skipping normalization.")


        # --- Expand to 4D ---
        # Add channel dim and expand
        conflict_map_4d = conflict_map_3d.unsqueeze(0).expand(num_channels, -1, -1, -1)

        return conflict_map_4d


    def compute_dynamic_source_mask( # Renamed slightly for clarity
        self,
        conflict_map_4d: torch.Tensor,
        k: float,
        function_type: str = 'exponential_squared',
        clamp_output: bool = True
        # warn_threshold removed as normalization is best done in the generating function
    ) -> torch.Tensor:
        """
        Computes the Dynamic Source Mask M(p, t) based on a 4D conflict map input.

        Applies a chosen function element-wise to the conflict map values to generate
        the mask. High conflict should result in a mask value near 0, while low
        conflict should result in a value near 1. Assumes the input conflict map
        may have redundant values across the channel dimension.

        Args:
            conflict_map_4d (torch.Tensor): A tensor representing the conflict C(p, t)
                                            between Vsrc and Vtar. Expected shape:
                                            (Channels, Frames, Height, Width).
                                            Normalizing this input (e.g., via
                                            generate_conflict_map_4d with normalize=True)
                                            is recommended for easier tuning of 'k'.
            k (float): A positive sensitivity parameter. Controls how aggressively
                    the mask value decreases as conflict increases.
            function_type (str, optional): Specifies the function f(C) used to map
                                        conflict C to the mask value M element-wise.
                                        Options:
                                        - 'exponential_squared': M = exp(-k * C^2)
                                        - 'exponential': M = exp(-k * C)
                                        - 'inverse': M = 1 / (1 + k * C)
                                        - 'inverse_squared': M = 1 / (1 + k * C^2)
                                        Defaults to 'exponential_squared'.
            clamp_output (bool, optional): If True, clamps the output mask values
                                        strictly to the range [0.0, 1.0].
                                        Defaults to True.

        Returns:
            torch.Tensor: The dynamic source mask tensor M(p, t).
                        Shape: (Channels, Frames, Height, Width).

        Raises:
            TypeError: If conflict_map_4d is not a PyTorch tensor.
            ValueError: If conflict_map_4d is not 4D, k is not positive, or an
                        invalid function_type is provided.
        """
        # --- Input Validation ---
        if not isinstance(conflict_map_4d, torch.Tensor):
            raise TypeError("Input conflict_map_4d must be a PyTorch tensor.")

        if conflict_map_4d.ndim != 4:
            raise ValueError(f"Input conflict_map_4d must be 4D (C, F, H, W), "
                            f"got {conflict_map_4d.ndim}D shape {conflict_map_4d.shape}")

        # num_channels = conflict_map_4d.shape[0] # Get C if needed elsewhere

        if not isinstance(k, (int, float)) or k <= 0:
            raise ValueError(f"Parameter k must be a positive number, got {k}")

        supported_functions = ['exponential_squared', 'exponential', 'inverse', 'inverse_squared']
        if function_type not in supported_functions:
            raise ValueError(f"Invalid function_type '{function_type}'. "
                            f"Supported types are: {supported_functions}")

        # --- Mask Calculation (Applied element-wise on 4D tensor) ---
        # Note: conflict_map_4d already includes the (potentially redundant) channel dim
        conflict_map_float = conflict_map_4d.float()
        k_float = float(k)

        if function_type == 'exponential_squared':
            # Applies exp(-k * C^2) element-wise
            mask_4d = torch.exp(-k_float * torch.pow(conflict_map_float, 2))
        elif function_type == 'exponential':
            # Applies exp(-k * C) element-wise
            mask_4d = torch.exp(-k_float * conflict_map_float)
        elif function_type == 'inverse':
            # Applies 1 / (1 + k * C) element-wise
            mask_4d = 1.0 / (1.0 + k_float * conflict_map_float)
        elif function_type == 'inverse_squared':
            # Applies 1 / (1 + k * C^2) element-wise
            mask_4d = 1.0 / (1.0 + k_float * torch.pow(conflict_map_float, 2))
        else:
            raise NotImplementedError(f"Function type {function_type} calculation not implemented.")

        # --- Optional Clamping ---
        if clamp_output:
            mask_4d = torch.clamp(mask_4d, min=0.0, max=1.0)

        # --- Return the 4D Mask ---
        # No expansion needed, it's already 4D
        return mask_4d
    
                    
    def edit(self,
             target_prompt,
             size=(832, 480),
             frame_num=81,
             shift=5.0,
             sample_solver='unipc',
             sampling_steps=50,
             guide_scale=5.0,
             tar_guide_scale=10.0,
             n_prompt="",
             seed=-1,
             offload_model=True,
             source_video_path=None,
             source_prompt=None,
             nmax_step=50,
             nmin_step=0,     
             n_avg=5,
             worse_avg=3,
             omega=3,
             source_words=None,
             target_words=None,
             window_size=11,
             decay_factor=0.1,
             tmd_window_size=11,
             tmd_stride=8
             ):
        # preprocess
        F = frame_num
        W, H = size
        target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1,
                        size[1] // self.vae_stride[1], size[0] // self.vae_stride[2])
        seq_len = math.ceil((target_shape[2] * target_shape[3]) /
                            (self.patch_size[1] * self.patch_size[2]) * 
                            target_shape[1] / self.sp_size) * self.sp_size

        if n_prompt == "":
            n_prompt = self.sample_neg_prompt
        seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
        # seed_g = torch.Generator(device=self.device)
        torch.manual_seed(seed)

        # 加载源视频潜在表示和参考图像
        x_src = self.load_video_frames(source_video_path)
        C_latent, F_latent, H_latent, W_latent = x_src.shape

        # Validate TMD parameters
        if tmd_window_size > F_latent:
            logging.warning(f"tmd_window_size ({tmd_window_size}) > latent frames ({F_latent}). Using full sequence as one window.")
            tmd_window_size = F_latent
            tmd_stride = F_latent
        elif tmd_stride <= 0:
             logging.warning(f"Invalid tmd_stride ({tmd_stride}). Setting stride to window size / 2.")
             tmd_stride = max(1, tmd_window_size // 2)
             
             
        # 计算提示词相对位置
        source_words_idx = None
        target_words_idx = None 
        if source_words:
            tk1 = self.text_encoder.tokenizer.tokenizer.tokenize(source_prompt, add_special_tokens=True)
            tk2 = self.text_encoder.tokenizer.tokenizer.tokenize(source_words, add_special_tokens=True)
            source_words_idx = self.find_subtokens_range(tk1[:-1], tk2[:-1])
            
        if target_words:
            tk1 = self.text_encoder.tokenizer.tokenizer.tokenize(target_prompt, add_special_tokens=True)
            tk2 = self.text_encoder.tokenizer.tokenizer.tokenize(target_words, add_special_tokens=True)
            target_words_idx = self.find_subtokens_range(tk1[:-1], tk2[:-1])

        # 编码源和目标文本提示
        if not self.t5_cpu:
            self.text_encoder.model.to(self.device)
            context_src = self.text_encoder([source_prompt], self.device)
            context_tar = self.text_encoder([target_prompt], self.device)
            context_null = self.text_encoder([n_prompt], self.device)
            if offload_model:
                self.text_encoder.model.cpu()
        else:
            context_tar = self.text_encoder([target_prompt], torch.device('cpu'))
            context_src = self.text_encoder([source_prompt], torch.device('cpu'))
            context_null = self.text_encoder([n_prompt], torch.device('cpu'))
            context_src = [t.to(self.device) for t in context_src]
            context_tar = [t.to(self.device) for t in context_tar]
            context_null = [t.to(self.device) for t in context_null]

        arg_src_c = {'context': context_src, 'seq_len': seq_len, 'words_indices': source_words_idx, 'block_id': 18, 'type': 'src'}
        arg_tar_c = {'context': context_tar, 'seq_len': seq_len, 'words_indices': target_words_idx, 'block_id': 18, 'type': 'tar'}
        arg_unc = {'context': context_null, 'seq_len': seq_len}

        
        # 初始化编辑路径
        zt_edit = x_src.clone() # [16, 21, 60, 104]: C x Frames x H x W 
        conflict_mask = torch.ones_like(x_src)

        # 设置采样调度器
        if sample_solver == 'unipc':
            sample_scheduler = FlowUniPCMultistepScheduler(
                num_train_timesteps=self.num_train_timesteps, shift=1, use_dynamic_shifting=False, solver_order=2)
            sample_scheduler.set_timesteps(sampling_steps, device=self.device, shift=shift)
            timesteps = sample_scheduler.timesteps
        elif sample_solver == 'dpm++':
            sample_scheduler = FlowDPMSolverMultistepScheduler(
                num_train_timesteps=self.num_train_timesteps, shift=1, use_dynamic_shifting=False, solver_order=1)
            sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)
            timesteps, _ = retrieve_timesteps(sample_scheduler, device=self.device, sigmas=sampling_sigmas)
        else:
            raise NotImplementedError("Unsupported solver.")

        self.model.to(self.device)
        # 编辑过程
        with amp.autocast(dtype=self.param_dtype), torch.no_grad():
            for index, t in enumerate(tqdm(timesteps)):
                t_next = timesteps[timesteps.tolist().index(t) + 1] if t > timesteps[-1] else 0
                arg_src_c["timestep"] = t
                arg_tar_c["timestep"] = t
                arg_unc["timestep"] = t
                timestep = torch.tensor([t], device=self.device)
                relative_index = nmax_step - (sampling_steps - index)
                v_list = []

                if sampling_steps - (index + 1) >= nmax_step:
                    continue
                if sampling_steps - (index + 1) >= nmin_step:
                    t_i = t / 1000.0
                    t_im1 = t_next / 1000.0
                    v_delta_sum = torch.zeros_like(zt_edit)
                    v_worse = torch.zeros_like(zt_edit)
                    v_mask = torch.zeros_like(zt_edit)
                    for time in range(n_avg):
                        # --- Temporal MultiDiffusion Logic ---
                        V_delta_accumulator = torch.zeros_like(zt_edit)
                        window_counts = torch.zeros_like(zt_edit) # Use float for division later
                        window_mask_sum = torch.zeros_like(zt_edit)    
                        fwd_noise = torch.randn_like(zt_edit[:, 0:tmd_window_size, :, :], device=self.device)
                        # for f_start in range(0, F_latent, tmd_stride):
                        if tmd_window_size >= F_latent:
                            window_starts = [0]
                        else:
                            window_starts = list(range(0, F_latent - tmd_window_size, tmd_stride))
                            last_possible_start = F_latent - tmd_window_size
                            if not window_starts or window_starts[-1] < last_possible_start:
                                if last_possible_start >= 0:
                                    window_starts.append(last_possible_start)

                        for f_start in window_starts:
                            f_end = F_latent if tmd_window_size >= F_latent else f_start + tmd_window_size

                            # --- Calculate V_delta within the window ---
                            # Extract window slices
                            zt_edit_w = zt_edit[:, f_start:f_end, :, :]
                            x_src_w = x_src[:, f_start:f_end, :, :]  

                            # 计算 zt_src
                            zt_src = (1 - t_i) * x_src_w + t_i * fwd_noise
                            zt_tar = zt_edit_w + zt_src - x_src_w
                                                
                        
                            # 计算源和目标噪声预测
                            noise_pred_src, src_attn_map = self.model([zt_src], t=timestep, **arg_src_c)
                            noise_pred_tar, tar_attn_map = self.model([zt_tar], t=timestep, **arg_tar_c)
                            noise_pred_src, noise_pred_tar = noise_pred_src[0], noise_pred_tar[0]
                            # uncond
                            noise_pred_uncond_src, _ = self.model([zt_src], t=timestep, **arg_unc)
                            noise_pred_uncond_tar, _ = self.model([zt_tar], t=timestep, **arg_unc)
                            noise_pred_uncond_src, noise_pred_uncond_tar = noise_pred_uncond_src[0], noise_pred_uncond_tar[0]

                            # 计算引导后的噪声预测
                            noise_pred_src_guided = noise_pred_uncond_src + guide_scale * (noise_pred_src - noise_pred_uncond_src)
                            noise_pred_tar_guided = noise_pred_uncond_tar + tar_guide_scale * (noise_pred_tar - noise_pred_uncond_tar)
                            
                            sum_attn_mask = torch.zeros_like(zt_edit_w)
                            
                            conflict_map = self.generate_conflict_map(noise_pred_src_guided, noise_pred_tar_guided)
                            raw_mask = self.compute_dynamic_source_mask(conflict_map, 0.5)
                            clamped_mask = torch.clamp(raw_mask, min=0.0, max=1.0)
                            conflict_mask = 5.0 - 4.0 * clamped_mask
                                                        
                            if src_attn_map is not None:
                                src_attn_mask = self.create_binary_mask(src_attn_map,
                                                                        n=window_size,
                                                                        pooling_mode='avg',
                                                                        threshold=torch.mean(src_attn_map).item(),
                                                                        threshold_method='fixed')
                                
                                sum_attn_mask += src_attn_mask
                                
                            if tar_attn_map is not None:
                                tar_attn_mask = self.create_binary_mask(tar_attn_map,
                                                                        n=window_size,
                                                                        pooling_mode='avg',
                                                                        threshold=torch.mean(tar_attn_map).item(),
                                                                        threshold_method='fixed')
                                
                                sum_attn_mask += tar_attn_mask
                                
                            sum_attn_mask = torch.clamp(sum_attn_mask, min=0.0, max=1.0)
                        
                            V_delta = noise_pred_tar_guided - noise_pred_src_guided
                            
                            # Accumulate results
                            V_delta_accumulator[:, f_start:f_end, :, :] += V_delta
                            window_counts[:, f_start:f_end, :, :] += 1.0
                            window_mask_sum[:, f_start:f_end, :, :] += sum_attn_mask

                        V_delta_final = V_delta_accumulator / torch.clamp(window_counts, min=1.0) # Avoid division by zero
                        v_delta_sum += V_delta_final
                        v_list.append(V_delta_final)
                        v_mask += window_mask_sum
                        if time < worse_avg:
                            v_worse += V_delta_final

                    v_list = torch.stack(v_list, dim=0) 
                    V_delta_better = v_list.mean(dim=0)
                    
                    v_trend = []
                    for worse_set in itertools.combinations(v_list, worse_avg):
                        v_worse = torch.zeros_like(V_delta_better)
                        for worse in worse_set:
                            v_worse += worse
                        v_worse = v_worse / worse_avg
                        v_trend.append(V_delta_better - v_worse)
                    v_trend = torch.stack(v_trend, dim=0)
                    v_trend = v_trend.mean(dim=0)
                    
                    v_mask = torch.clamp(v_mask, min=0.0, max=1.0)
                    v_mask = self.soften_mask_edges(v_mask, decay_factor=decay_factor)
                    
                    V_delta_final = V_delta_better + (omega - 1) * v_trend
                    V_delta_final = V_delta_final * v_mask

                    zt_edit = zt_edit + (t_im1 - t_i) * V_delta_final                
                    
                else:
                    # 使用tar进行采样
                    noise_pred_uncond, _ = self.model([zt_edit], t=timestep, context=context_null, seq_len=seq_len)
                    noise_pred_cond, _ = self.model([zt_edit], t=timestep, context=context_tar, seq_len=seq_len)
                    noise_pred_cond, noise_pred_uncond = noise_pred_cond[0], noise_pred_uncond[0]
                    noise_pred = noise_pred_uncond + 6 * (
                        noise_pred_cond - noise_pred_uncond)

                    temp_x0 = sample_scheduler.step(
                        noise_pred.unsqueeze(0),
                        t,
                        zt_edit.unsqueeze(0),
                        return_dict=False)[0]
                    
                    zt_edit = temp_x0.squeeze(0)
             
        # 解码编辑结果
        if offload_model:
            self.model.cpu()
        if self.rank == 0:
            videos = self.vae.decode([zt_edit])
            return videos[0]
        return None