import gc import random import time from typing import List, Optional import numpy as np from omegaconf import DictConfig, OmegaConf import torch import torch.nn.functional as F from torch.nn.utils import clip_grad_norm_ from tqdm import tqdm from optimize_utils import transpose_dict_2d from tensor_utils import ( build_anisotropic_gaussian_from_mask, combine_gaussian_maps, combine_masks_or, normalize_tensor_to_match_tensor, resize_tensor, warp_tensor, warp_tensor_sequence, ) from utils.wan_wrapper import ( WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper, ) from demo_utils.memory import ( gpu, get_cuda_free_memory_gb, DynamicSwapInstaller, move_model_to_device_with_memory_preservation, ) IMAGE_HEIGHT = 480.0 def split_trajectories_segments( trajectories: list[dict[str, bool | list[torch.Tensor]]], translation_step: float, rotation_step: float, ) -> List[list[dict[str, bool | list[torch.Tensor]]]]: """ Split drag trajectories into evenly spaced intermediate segments for progressive (coarse-to-fine) optimization. Given N trajectories (each with per-frame deltas), this function: 1. Determines the maximum number of segments needed across all trajectories based on the magnitude of their deltas and the provided step sizes. 2. Divides every trajectory's deltas uniformly into that many segments. 3. Produces a list of cumulative intermediate trajectory snapshots, where segment k contains deltas scaled by (k / max_segments). :param trajectories: N x trajectory dicts. Each dict has keys: - 'is_rotation' (bool): Whether this trajectory is a rotation. - 'deltas' (list): Per-frame displacement values. For translation: each delta is a 2D vector (dx, dy). For rotation: each delta is a scalar angle. - 'start_point': The starting pixel coordinate of the drag. - 'rotation_center' (only if is_rotation): The center of rotation. :param translation_step: The pixel distance that defines one segment for translation trajectories. Larger values produce fewer segments. :param rotation_step: The angle (in the same units as deltas) that defines one segment for rotation trajectories. Larger values produce fewer segments. :returns: segment_num x N x trajectory dicts. A list of length `max_segment_number`, where each element is a list of N trajectory dicts. The k-th element (1-indexed) contains trajectories whose deltas are scaled to (k / max_segment_number) of the original deltas — i.e., cumulative intermediate waypoints. """ # ------------------------------------------------------------------------- # Phase 1: Convert raw deltas to torch tensors (ensure uniform type) # ------------------------------------------------------------------------- for trajectory in trajectories: trajectory["deltas"] = [torch.tensor(delta, device="cpu") for delta in trajectory["deltas"]] # ------------------------------------------------------------------------- # Phase 2: Determine the maximum number of segments across all trajectories. # - For rotations: segment count = |angle_delta| // rotation_step # - For translations: segment count = ||displacement_delta||₂ // translation_step # - We take the global maximum so every trajectory is split into the # same number of segments (ensuring synchronized progressive steps). # ------------------------------------------------------------------------- max_segment_number = 1 # at least one segment for trajectory in trajectories: print(f"{trajectory['is_rotation'] = }") for delta in trajectory["deltas"]: if trajectory["is_rotation"]: magnitude = abs(delta) step = rotation_step else: magnitude = abs(torch.norm(delta)) step = translation_step segment_number = int(magnitude // step) print(f"{delta = } {magnitude = } {segment_number = }") max_segment_number = max(max_segment_number, segment_number) print(f"{max_segment_number = }") # ------------------------------------------------------------------------- # Phase 3: Compute per-segment step sizes for each trajectory. # Each trajectory's deltas are divided by max_segment_number to get # the uniform per-segment increment. # ------------------------------------------------------------------------- split_trajectory_steps = [] for trajectory in trajectories: print(f"{trajectory['is_rotation'] = }") # Divide each frame's delta by the total number of segments trajectory_steps = [delta / float(max_segment_number) for delta in trajectory["deltas"]] print(f"{trajectory_steps = }") # Build the per-trajectory step metadata split_trajectory_step = { "is_rotation": trajectory["is_rotation"], "steps": trajectory_steps, # per-segment increment per frame "start_point": trajectory["start_point"], } if trajectory["is_rotation"]: split_trajectory_step["rotation_center"] = trajectory["rotation_center"] split_trajectory_steps.append(split_trajectory_step) # ------------------------------------------------------------------------- # Phase 4: Build cumulative intermediate trajectory lists. # For segment_index k (1-indexed from 1 to max_segment_number): # delta_k = step * k # This produces progressively larger displacements, enabling the # optimizer to move features gradually toward the final target. # ------------------------------------------------------------------------- new_trajectories_list = [] for segment_index in range(max_segment_number): segment_index += 1 # 1-indexed: cumulative scale factor new_trajectories = [] for trajectory_step in split_trajectory_steps: new_trajectory = { "is_rotation": trajectory_step["is_rotation"], "deltas": [step * segment_index for step in trajectory_step["steps"]], "start_point": trajectory_step["start_point"], } if trajectory_step["is_rotation"]: new_trajectory["rotation_center"] = trajectory_step["rotation_center"] new_trajectories.append(new_trajectory) print(f"{new_trajectories = }") new_trajectories_list.append(new_trajectories) # Return: list of length max_segment_number, each containing N trajectory dicts # with cumulatively scaled deltas (segment 1 = smallest, last = full original delta) return new_trajectories_list class CausalInferencePipeline(torch.nn.Module): def __init__( self, args, device, generator=None, text_encoder=None, vae=None, ): super().__init__() # Step 1: Initialize all models self.generator = ( WanDiffusionWrapper( **getattr(args, "model_kwargs", {}), is_causal=True, ) if generator is None else generator ) self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder self.vae = WanVAEWrapper() if vae is None else vae # Step 2: Initialize all causal hyperparmeters self.scheduler = self.generator.get_scheduler() self.denoising_step_list = torch.tensor(args.denoising_step_list, dtype=torch.long) if args.warp_denoising_step: timesteps = torch.cat( ( self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32), ) ) self.denoising_step_list = timesteps[1000 - self.denoising_step_list] self.num_transformer_blocks = 30 self.frame_seq_length = 1560 self.kv_cache1 = None self.args = args self.num_frame_per_block = getattr(args, "num_frame_per_block", 1) self.independent_first_frame = args.independent_first_frame self.local_attn_size = self.generator.model.local_attn_size print(f"KV inference with {self.num_frame_per_block} frames per block") if self.num_frame_per_block > 1: self.generator.model.num_frame_per_block = self.num_frame_per_block def stack_features( self, record_features: dict[int, torch.Tensor], # dict[block_index] -> Tensor[1, 3, 30, 52, 1536] ): """ Reorder and stack per-block attention features into one tensor. Input: record_features: Mapping `block_index -> Tensor` with shape `[1, 3, 30, 52, 1536]`. Returns: Tensor with shape `[1536 * L, 3, 30, 52]`, where `L = len(record_features)`. """ attention_values = { k: v.permute(0, 4, 1, 2, 3).contiguous() for k, v in record_features.items() } # [ [1, 1536, 3, 30, 52] ] features = [] for fet in attention_values.values(): # [1536, 3, 30, 52] features.append(fet.squeeze(0).contiguous()) del fet # Concatenate all features feature = torch.cat(features, dim=0) # [1536 * L, 3, 30, 52], L is the number of blocks return feature def generate_features( self, latents, conditional_dict, timestep, kv_cache, crossattn_cache, current_start, # below are for drag optimization model_config: DictConfig = None, ): """ Run one generator forward pass and return prediction + stacked features. Notes: - KV cache is deep-cloned/detached before forward to avoid mutating the caller's cache during optimization. - Forward runs under CUDA bfloat16 autocast. - Returned `record_features` are converted via `stack_features(...)`. Returns: denoised_pred: Model denoised prediction tensor. record_features: Dict `variant_key -> Tensor[1536 * L, 3, 30, 52]`. L is the number of blocks. """ temp_kv_cache = [ { "k": kv_cache[block_index]["k"].clone().detach(), "v": kv_cache[block_index]["v"].clone().detach(), "global_end_index": kv_cache[block_index]["global_end_index"].clone().detach(), "local_end_index": kv_cache[block_index]["local_end_index"].clone().detach(), } for block_index in range(self.num_transformer_blocks) ] # print(f"{temp_kv_cache[0]['k'].shape = }") # Forward pass through the transformer with user-specified autocast dtype with torch.autocast(device_type="cuda", dtype=torch.bfloat16): # Extract features during forward pass _, denoised_pred = self.generator( noisy_image_or_video=latents, conditional_dict=conditional_dict, timestep=timestep, kv_cache=temp_kv_cache, crossattn_cache=crossattn_cache, current_start=current_start, # below are for drag optimization model_config=model_config, ) denoised_pred, record_features = denoised_pred # record_features: dict[block_index] -> dict[key -> Tensor] [1, 3, 30, 52, 1536] record_features = transpose_dict_2d(record_features) # record_features: Dict[key] -> Dict[block_index -> Tensor] [1, 3, 30, 52, 1536] record_features = {k: self.stack_features(v) for k, v in record_features.items()} # record_features: Dict[key] -> Tensor [1536 * L, 3, 30, 52] return denoised_pred, record_features @staticmethod def precompute_conditions( trajectories: list[dict[str, bool | list[torch.Tensor]]], # N x trajectory masks: list[np.ndarray], # N x mask dtype: torch.dtype, device: torch.device, model_config: DictConfig, downsample_movable_mask: torch.Tensor, previous_record_features: dict[int, dict[str, torch.Tensor]], latent_spatial_size: tuple[int, int], # (H_lat, W_lat) ): with torch.no_grad(): _any_prev_feat = next(iter(previous_record_features.values())) feat_spatial_size = tuple(_any_prev_feat.shape[-2:]) # (Hf, Wf) # 1) Warped masks in image space warped_masks = [ warp_tensor_sequence( tensor=torch.tensor(mask, device=device).float(), is_rotation=trajectory["is_rotation"], deltas=trajectory["deltas"], rotation_center=trajectory.get("rotation_center", None), original_height=IMAGE_HEIGHT, mode="nearest", cumulative=False, ) for trajectory, mask in zip(trajectories, masks) ] # N x frame x [H_img, W_img] # print(f"{len(warped_masks) = }, {len(warped_masks[0]) = }, {warped_masks[0][0].shape = }") # 2) Downsampled warped masks in feature space down_warp_masks = [ [ resize_tensor( warped_mask.detach(), size=feat_spatial_size, mode="nearest", ).detach() for warped_mask in traj_warped_masks ] for traj_warped_masks in warped_masks ] # N x frame x [Hf, Wf] # 3) Gaussian heatmaps per trajectory gaussian_heatmaps_per_traj = [ [ build_anisotropic_gaussian_from_mask( warped_mask, padding_scale=model_config.drag_optim_config.gradient_gaussian_padding, sigma_scale=model_config.drag_optim_config.gradient_gaussian_sigma, ).detach() for warped_mask in traj_warped_masks ] for traj_warped_masks in warped_masks ] # N x frame x [H_img, W_img] # 4) Combined downsampled movable mask (OR of all warped + original) all_down_warp_masks = [ dwm for traj_down_masks in down_warp_masks for dwm in traj_down_masks ] all_down_warp_masks.append(downsample_movable_mask.clone()) combined_downsample_movable_mask = combine_masks_or(all_down_warp_masks) # 5) Precompute warped attention values per variant warped_for_prev: dict[str | int, List[List[torch.Tensor]]] = { key: [ warp_tensor_sequence( tensor=prev_feat.to(dtype=dtype, device=device), is_rotation=trajectory["is_rotation"], deltas=[ d.to( dtype=dtype, device=device, ) for d in trajectory["deltas"] ], rotation_center=trajectory.get("rotation_center", None), original_height=IMAGE_HEIGHT, mode="nearest", cumulative=False, ) for trajectory in trajectories ] for key, prev_feat in previous_record_features.items() } # warped_for_prev: dict[key] -> list[traj_index] -> list[frame_index] -> Tensor [1536 * L, 30 * scaling, 52 * scaling] # 6) Combined Gaussian heatmaps in latent space combined_gaussian_heatmaps = None num_frames = len(gaussian_heatmaps_per_traj[0]) num_trajs = len(gaussian_heatmaps_per_traj) combined_gaussian_heatmaps = torch.stack( [ combine_gaussian_maps( [ gaussian_heatmaps_per_traj[traj_idx][frame_idx] for traj_idx in range(num_trajs) ] ) # [H_img, W_img] for frame_idx in range(num_frames) ], dim=0, ).to( device=device ) # [F, H_img, W_img] combined_gaussian_heatmaps = resize_tensor( combined_gaussian_heatmaps, size=latent_spatial_size, mode="bilinear", ).detach() # [F, H_lat, W_lat] combined_gaussian_heatmaps = combined_gaussian_heatmaps.to(dtype=dtype) return ( warped_masks, down_warp_masks, gaussian_heatmaps_per_traj, combined_downsample_movable_mask, warped_for_prev, combined_gaussian_heatmaps, ) def optimize_latent( self, latents, conditional_dict, timestep, kv_cache, crossattn_cache, current_start, # below are for drag optimization trajectories: list[dict[str, bool | list[torch.Tensor]]], # N x trajectory, # trajectory has keys 'is_rotation' 'deltas' 'start_point' # if is_rotation: trajectory also has 'rotation_center' masks: list[np.ndarray], # N x mask movable_mask: np.ndarray, clean_previous_record_feature: dict[int, dict[str, torch.Tensor]], # dict[block_index] -> dict[key -> Tensor] [1, 3, 30, 52, 1536] noisy_previous_record_feature: dict[int, dict[str, torch.Tensor]], # dict[block_index] -> dict[key -> Tensor] [1, 3, 30, 52, 1536] model_config: DictConfig, optimize_target_latent_index: int = -1, ): """ :param trajectories: N x trajectory, trajectory has keys 'is_rotation' 'deltas' 'start_point' if is_rotation: trajectory also has 'rotation_center' :param masks: N x mask """ assert isinstance(model_config.drag_optim_config.optimize_iter, int) assert isinstance(model_config.drag_optim_config.optimize_lr, float) assert len(model_config.drag_optim_config.record_feature_block_indexes) > 0 assert len(trajectories) == len(masks) if len(trajectories) == 0: return latents print(f"{trajectories = }") print(f"{len(masks) = }") original_latents = latents.clone().detach() original_denoised_pred = self.generate_features( latents=latents, conditional_dict=conditional_dict, timestep=timestep, kv_cache=kv_cache, crossattn_cache=crossattn_cache, current_start=current_start, # below are for drag optimization model_config=model_config, ) original_denoised_pred, _ = original_denoised_pred # Check if optimization is enabled if model_config.drag_optim_config.optimize_iter <= 0: return latents # Cast latents to configured dtype for optimization latent_original_dtype = latents.dtype latents = latents.to(dtype=torch.bfloat16) timestep_original_dtype = timestep.dtype timestep = timestep.to(dtype=latents.dtype) for param in self.generator.parameters(): generator_original_dtype = param.dtype break self.generator = self.generator.to(dtype=latents.dtype) # self.generator.train(True) for param in self.generator.parameters(): param.requires_grad = False split_trajectories_list = split_trajectories_segments( trajectories=trajectories, translation_step=model_config.drag_optim_config.translation_step, rotation_step=model_config.drag_optim_config.rotation_step, ) # split_trajectories_list: list[segment_index] -> list[trajectory_index] -> trajectory dict def _select_variant(feat): if isinstance(feat, dict): keys = list(feat.keys()) if not keys: raise ValueError("Empty feature dict provided.") non_orig = [k for k in keys if str(k) != "original"] key = random.choice(keys) print(f"Selected feature variant {key = } from {keys = }") # key = "original" if "original" in feat else keys[0] return feat[key] return feat def get_previous_last( prev: dict[int, dict[str, torch.Tensor]], # dict[block_index] -> dict[key -> Tensor] [1, 3, 30, 52, 1536] ) -> dict[str, torch.Tensor]: out = transpose_dict_2d(prev) # out: dict[key] -> dict[block_index -> Tensor] [1, 3, 30, 52, 1536] out = {k: self.stack_features(v) for k, v in out.items()} # out: dict[key] -> Tensor [1536 * L, 3, 30, 52] out = {k: v[:, -1, ...].detach() for k, v in out.items()} # out: dict[key] -> Tensor [1536 * L, 30, 52] out = { k: resize_tensor( v.detach(), scale_factor=model_config.drag_optim_config.feature_scaling_factor, mode="bilinear", ).detach() for k, v in out.items() } # out: dict[key] -> Tensor [1536 * L, 30 * scaling, 52 * scaling] return out previous_record_features: dict[str, torch.Tensor] = get_previous_last( noisy_previous_record_feature ) # previous_record_features: dict[key] -> Tensor [1536 * L, 30 * scaling, 52 * scaling] movable_mask_torch = torch.tensor(movable_mask, device=latents.device).float() downsample_movable_mask = resize_tensor( movable_mask_torch.detach(), size=tuple(original_denoised_pred.shape[-2:]), mode="nearest", ).detach() # print(f"{downsample_movable_mask.shape = }") # [60, 104] with torch.enable_grad(): latents.requires_grad_(True) optimizer = torch.optim.AdamW( [latents], lr=model_config.drag_optim_config.optimize_lr, ) for split_traj_idx, split_trajectories in enumerate(split_trajectories_list): # split_trajectories: N x trajectory, list[trajectory_index] -> trajectory dict ( warped_masks, down_warp_masks, gaussian_heatmaps_per_traj, combined_downsample_movable_mask, warped_previous_record_features, combined_gaussian_heatmaps, ) = CausalInferencePipeline.precompute_conditions( trajectories=split_trajectories, masks=masks, dtype=latents.dtype, device=latents.device, model_config=model_config, downsample_movable_mask=downsample_movable_mask, previous_record_features=previous_record_features, latent_spatial_size=tuple(latents.shape[-2:]), ) # ------------------------- # Optimization iterations (reuse precomputed items) # ------------------------- for optimize_iter_idx in range(model_config.drag_optim_config.optimize_iter): print(f"{optimize_iter_idx = }") print(f"{latents.mean((0, 2, 3, 4)) = }") print(f"{latents.std((0, 2, 3, 4)) = }") denoised_pred, record_features = self.generate_features( latents=latents, conditional_dict=conditional_dict, timestep=timestep, kv_cache=kv_cache, crossattn_cache=crossattn_cache, current_start=current_start, model_config=model_config, ) # denoised_pred.shape [1, 3, 16, 60, 104] # record_features: dict[key] -> Tensor [1536 * L, 3, 30, 52] print(f"selecting features for optimization") record_features_selected = _select_variant( record_features, ) # print(f"{record_features_selected.shape = }") # record_features_selected: Tensor [1536 * L, 3, 30, 52] del record_features if optimize_target_latent_index >= 0: record_features_selected = record_features_selected[ :, optimize_target_latent_index : optimize_target_latent_index + 1, ] record_features_selected = resize_tensor( record_features_selected, scale_factor=model_config.drag_optim_config.feature_scaling_factor, mode="bilinear", ) # print( # f"{record_features_selected.shape = }" # ) # [1536 * L, 3, 30 * scaling, 52 * scaling] print(f"selecting warped previous features") warped_previous_feature_selected = _select_variant( warped_previous_record_features, ) # list[traj_index] -> list[frame_index] -> Tensor [1536 * L, 30 * scaling, 52 * scaling] loss = 0 loss_cnt = 0 # Iterate over each trajectory point for trajectory_index, trajectory in enumerate(split_trajectories): assert record_features_selected.shape[1] == len(trajectory["deltas"]) for frame_index in range(len(trajectory["deltas"])): warped_attention_values = warped_previous_feature_selected[ trajectory_index ][frame_index] downsample_warped_mask = down_warp_masks[trajectory_index][frame_index] pixel_wise_loss = F.mse_loss( warped_attention_values * downsample_warped_mask, record_features_selected[:, frame_index] * downsample_warped_mask, reduction="none", ).mean(dim=0) # print(f"{pixel_wise_loss.shape = }") # [60, 104] # Add weighted loss loss = loss + (downsample_warped_mask * pixel_wise_loss).sum() loss_cnt += downsample_warped_mask.sum() print(f"{loss = } / {loss_cnt = }") loss = loss / max(1e-8, loss_cnt) print(f"{loss = }") unchanged_mask = 1.0 - combined_downsample_movable_mask unchanged_loss = F.mse_loss( denoised_pred * unchanged_mask.detach(), original_denoised_pred.detach() * unchanged_mask.detach(), ) print(f"{unchanged_loss = }") loss = loss + unchanged_loss * 1.0 # Update latents self.generator.zero_grad() optimizer.zero_grad() if loss_cnt > 0: loss.backward() assert ( combined_gaussian_heatmaps.shape[0] == 1 or combined_gaussian_heatmaps.shape[0] == latents.shape[-4] ) assert combined_gaussian_heatmaps.shape[-2:] == latents.shape[-2:] latents.grad.mul_(combined_gaussian_heatmaps[:, None, :, :]) # Clip gradients clip_grad_norm_( [latents], max_norm=1.0, norm_type=2, ) optimizer.step() if model_config.drag_optim_config.normalize_latent_after_drag_optimize: print(f"Normalizing latents after optimize iteration") latents = ( normalize_tensor_to_match_tensor( latents.detach().clone(), dim=(0, 3, 4), reference_tensor=original_latents.to(dtype=latents.dtype), ) .detach() .clone() ) # latents = latents.clamp( # min=latents_min, # max=latents_max, # ).detach().clone() latents.requires_grad_(True) optimizer = torch.optim.AdamW( [latents], lr=model_config.drag_optim_config.optimize_lr, ) # Clean up to save memory gc.collect() torch.cuda.empty_cache() latents = latents.detach().requires_grad_(False) if model_config.drag_optim_config.normalize_latent_after_post_merge: latents = ( normalize_tensor_to_match_tensor( latents, dim=None, reference_tensor=original_latents, ) .detach() .clone() ) # Convert back to original dtype self.generator = self.generator.to(dtype=generator_original_dtype) self.generator.train(False) latents = latents.to(dtype=latent_original_dtype) timestep = timestep.to(dtype=timestep_original_dtype) # Detach latents and remove gradient latents = latents.detach().requires_grad_(False) return latents def inference( self, noise: torch.Tensor, text_prompts: List[str], initial_latent: Optional[torch.Tensor] = None, return_latents: bool = False, profile: bool = False, low_memory: bool = False, do_not_decode_video: bool = False, do_not_recompute_initial_latents: bool = False, # below are for drag optimization model_config: DictConfig = None, previous_record_feature_list: dict[int, dict[int, dict[str, torch.Tensor]]] = None, # dict[denoising_step] -> dict[block_index] -> dict[key -> Tensor] [1, 3, 30, 52, 1536] is_drag_optimize: bool = False, block_trajectories: list[list[dict[str, bool | list[torch.Tensor]]]] = [], masks: list[np.ndarray] = [], movable_mask: np.ndarray = None, drag_optimize_target_latent_index: int = -1, ) -> torch.Tensor: """ Perform inference on the given noise and text prompts. Inputs: noise (torch.Tensor): The input noise tensor of shape (batch_size, num_output_frames, num_channels, height, width). text_prompts (List[str]): The list of text prompts. initial_latent (torch.Tensor): The initial latent tensor of shape (batch_size, num_input_frames, num_channels, height, width). If num_input_frames is 1, perform image to video. If num_input_frames is greater than 1, perform video extension. return_latents (bool): Whether to return the latents. :param block_trajectories: block_num x N x trajectory, trajectory has keys 'is_rotation' 'deltas' 'start_point' if is_rotation: trajectory also has 'rotation_center' :param masks: N x mask Outputs: video (torch.Tensor): The generated video tensor of shape (batch_size, num_output_frames, num_channels, height, width). It is normalized to be in the range [0, 1]. """ batch_size, num_frames, num_channels, height, width = noise.shape if not self.independent_first_frame or ( self.independent_first_frame and initial_latent is not None ): # If the first frame is independent and the first frame is provided, then the number of frames in the # noise should still be a multiple of num_frame_per_block assert num_frames % self.num_frame_per_block == 0 num_blocks = num_frames // self.num_frame_per_block else: # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning assert (num_frames - 1) % self.num_frame_per_block == 0 num_blocks = (num_frames - 1) // self.num_frame_per_block num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0 num_output_frames = num_frames + num_input_frames # add the initial latent frames conditional_dict = self.text_encoder(text_prompts=text_prompts) if low_memory: gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5 move_model_to_device_with_memory_preservation( self.text_encoder, target_device=gpu, preserved_memory_gb=gpu_memory_preservation, ) output = torch.zeros( [batch_size, num_output_frames, num_channels, height, width], device=noise.device, dtype=noise.dtype, ) # Set up profiling if requested if profile: init_start = torch.cuda.Event(enable_timing=True) init_end = torch.cuda.Event(enable_timing=True) diffusion_start = torch.cuda.Event(enable_timing=True) diffusion_end = torch.cuda.Event(enable_timing=True) vae_start = torch.cuda.Event(enable_timing=True) vae_end = torch.cuda.Event(enable_timing=True) block_times = [] block_start = torch.cuda.Event(enable_timing=True) block_end = torch.cuda.Event(enable_timing=True) init_start.record() # Step 1: Initialize KV cache to all zeros if self.kv_cache1 is None: self._initialize_kv_cache(batch_size=batch_size, dtype=noise.dtype, device=noise.device) self._initialize_crossattn_cache( batch_size=batch_size, dtype=noise.dtype, device=noise.device ) else: if do_not_recompute_initial_latents: pass else: print(f"Resetting caches") self._reset_crossattn_cache() self._reset_kv_cache() # Step 2: Cache context feature current_start_frame = 0 if initial_latent is not None: timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0 if self.independent_first_frame: # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks assert (num_input_frames - 1) % self.num_frame_per_block == 0 num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block output[:, :1] = initial_latent[:, :1] if do_not_recompute_initial_latents: pass else: print(f"Recompute KV cache based on Initial Latents") self.generator( noisy_image_or_video=initial_latent[:, :1], conditional_dict=conditional_dict, timestep=timestep * 0, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length, ) current_start_frame += 1 else: # Assume num_input_frames is self.num_frame_per_block * num_input_blocks assert num_input_frames % self.num_frame_per_block == 0 num_input_blocks = num_input_frames // self.num_frame_per_block for _ in range(num_input_blocks): current_ref_latents = initial_latent[ :, current_start_frame : current_start_frame + self.num_frame_per_block, ] output[ :, current_start_frame : current_start_frame + self.num_frame_per_block, ] = current_ref_latents if do_not_recompute_initial_latents: pass else: print(f"Recompute KV cache based on Initial Latents") self.generator( noisy_image_or_video=current_ref_latents, conditional_dict=conditional_dict, timestep=timestep * 0, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length, ) current_start_frame += self.num_frame_per_block if profile: init_end.record() torch.cuda.synchronize() diffusion_start.record() # Step 3: Temporal denoising loop all_num_frames = [self.num_frame_per_block] * num_blocks if self.independent_first_frame and initial_latent is None: all_num_frames = [1] + all_num_frames for current_chunk_index, current_num_frames in enumerate( tqdm(all_num_frames), start=num_input_blocks if initial_latent is not None else 0, ): print(f"\n{current_chunk_index = } ; {current_start_frame = }") if profile: block_start.record() noisy_input = noise[ :, current_start_frame - num_input_frames : current_start_frame + current_num_frames - num_input_frames, ] if model_config is not None and OmegaConf.select( model_config, "drag_optim_config.record_feature_block_indexes" ): record_attention_values_list = {} # Step 3.1: Spatial denoising loop for time_step_index, current_timestep in enumerate(self.denoising_step_list): print(f"{time_step_index = } ; {current_timestep = }") # set current timestep timestep = ( torch.ones( [batch_size, current_num_frames], device=noise.device, dtype=torch.int64, ) * current_timestep ) if ( is_drag_optimize and time_step_index in model_config.drag_optim_config.optimize_denoising_steps_indexes ): noisy_input = self.optimize_latent( latents=noisy_input, conditional_dict=conditional_dict, timestep=timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length, # below are for drag optimization trajectories=block_trajectories[current_chunk_index - num_input_blocks], masks=masks, movable_mask=movable_mask, clean_previous_record_feature=previous_record_feature_list[-1], noisy_previous_record_feature=previous_record_feature_list[time_step_index], model_config=model_config, optimize_target_latent_index=drag_optimize_target_latent_index, ) print(f"{noisy_input.mean() = }") print(f"{noisy_input.std() = }") _, denoised_pred = self.generator( noisy_image_or_video=noisy_input, conditional_dict=conditional_dict, timestep=timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length, model_config=model_config, ) if model_config is not None and OmegaConf.select( model_config, "drag_optim_config.record_feature_block_indexes", ): denoised_pred, record_features = denoised_pred if ( time_step_index in model_config.drag_optim_config.optimize_denoising_steps_indexes ): record_attention_values_list[time_step_index] = record_features if ( model_config is not None and OmegaConf.select( model_config, "drag_optim_config.dynamic_chunk_normalization_block_number", default=0, ) > 0 ): num_norm_blocks = ( model_config.drag_optim_config.dynamic_chunk_normalization_block_number ) # Exclude the first chunk (independent first frame latent) by starting no earlier than num_frame_per_block dynamic_normalize_start_frame_index = max( self.num_frame_per_block, (current_chunk_index - num_norm_blocks) * self.num_frame_per_block, ) # print(f"{dynamic_normalize_start_frame_index = }") if dynamic_normalize_start_frame_index < current_start_frame: reference_tensor = torch.cat( [ output[ :, dynamic_normalize_start_frame_index:current_start_frame, ], denoised_pred, ], dim=1, ) denoised_pred = normalize_tensor_to_match_tensor( denoised_pred, dim=None, reference_tensor=reference_tensor, ) # print(f"{denoised_pred.mean() = }") # print(f"{denoised_pred.std() = }") if time_step_index < len(self.denoising_step_list) - 1: next_timestep = self.denoising_step_list[time_step_index + 1] noisy_input = self.scheduler.add_noise( denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), next_timestep * torch.ones( [batch_size * current_num_frames], device=noise.device, dtype=torch.long, ), ).unflatten(0, denoised_pred.shape[:2]) # Step 3.2: record the model's output output[ :, current_start_frame : current_start_frame + current_num_frames, ] = denoised_pred # Step 3.3: rerun with timestep zero to update KV cache using clean context context_timestep = torch.ones_like(timestep) * self.args.context_noise _, denoised_pred = self.generator( noisy_image_or_video=denoised_pred, conditional_dict=conditional_dict, timestep=context_timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length, model_config=model_config, ) if model_config is not None and OmegaConf.select( model_config, "drag_optim_config.record_feature_block_indexes" ): denoised_pred, record_features = denoised_pred record_attention_values_list[-1] = record_features if profile: block_end.record() torch.cuda.synchronize() block_time = block_start.elapsed_time(block_end) block_times.append(block_time) # Step 3.4: update the start and end frame indices current_start_frame += current_num_frames if profile: # End diffusion timing and synchronize CUDA diffusion_end.record() torch.cuda.synchronize() diffusion_time = diffusion_start.elapsed_time(diffusion_end) init_time = init_start.elapsed_time(init_end) vae_start.record() # Step 4: Decode the output if not do_not_decode_video: start_decode_time = time.time() video = self.vae.decode_to_pixel(output, use_cache=False) video = (video * 0.5 + 0.5).clamp(0, 1) print( f"{self.__class__.__name__}.inference() VAE decode time: {time.time() - start_decode_time:.2f} seconds" ) if profile: # End VAE timing and synchronize CUDA vae_end.record() torch.cuda.synchronize() vae_time = vae_start.elapsed_time(vae_end) total_time = init_time + diffusion_time + vae_time print("Profiling results:") print( f" - Initialization/caching time: {init_time:.2f} ms ({100 * init_time / total_time:.2f}%)" ) print( f" - Diffusion generation time: {diffusion_time:.2f} ms ({100 * diffusion_time / total_time:.2f}%)" ) for i, block_time in enumerate(block_times): print( f" - Block {i} generation time: {block_time:.2f} ms ({100 * block_time / diffusion_time:.2f}% of diffusion)" ) print(f" - VAE decoding time: {vae_time:.2f} ms ({100 * vae_time / total_time:.2f}%)") print(f" - Total time: {total_time:.2f} ms") return_values = [] if not do_not_decode_video: return_values.append(video) if return_latents: return_values.append(output) if model_config is not None and OmegaConf.select( model_config, "drag_optim_config.record_feature_block_indexes" ): return_values.append(record_attention_values_list) if len(return_values) == 0: return elif len(return_values) == 1: return return_values[0] else: return tuple(return_values) def _initialize_kv_cache( self, batch_size, dtype, device, ): """ Initialize a Per-GPU KV cache for the Wan model. """ print( f""" {type(self).__name__}._initialize_kv_cache {batch_size = } {dtype = } {device = } """ ) kv_cache1 = [] if self.local_attn_size != -1: print(f"use {self.local_attn_size = }") # Use the local attention size to compute the KV cache size kv_cache_size = self.local_attn_size * self.frame_seq_length else: # Use the default KV cache size kv_cache_size = 32760 print(f"{kv_cache_size = }") for _ in range(self.num_transformer_blocks): kv_cache1.append( { "k": torch.zeros( [batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device, ), "v": torch.zeros( [batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device, ), "global_end_index": torch.tensor([0], dtype=torch.long, device=device), "local_end_index": torch.tensor([0], dtype=torch.long, device=device), } ) self.kv_cache1 = kv_cache1 # always store the clean cache def _initialize_crossattn_cache( self, batch_size, dtype, device, ): """ Initialize a Per-GPU cross-attention cache for the Wan model. """ print( f""" {type(self).__name__}._initialize_crossattn_cache {batch_size = } {dtype = } {device = } """ ) crossattn_cache = [] for _ in range(self.num_transformer_blocks): crossattn_cache.append( { "k": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device), "v": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device), "is_init": False, } ) self.crossattn_cache = crossattn_cache def _reset_crossattn_cache(self): # reset cross attn cache print(f"{type(self).__name__}._reset_crossattn_cache") for block_index in range(self.num_transformer_blocks): self.crossattn_cache[block_index]["is_init"] = False def _reset_kv_cache(self): # reset kv cache print(f"{type(self).__name__}._reset_kv_cache") for block_index in range(len(self.kv_cache1)): self.kv_cache1[block_index]["global_end_index"] = torch.tensor( [0], dtype=torch.long, device=self.kv_cache1[block_index]["global_end_index"].device, ) self.kv_cache1[block_index]["local_end_index"] = torch.tensor( [0], dtype=torch.long, device=self.kv_cache1[block_index]["local_end_index"].device, ) def is_kv_cache_initialized(self): return hasattr(self, "kv_cache1") and self.kv_cache1 is not None def is_crossattn_cache_initialized(self): return hasattr(self, "crossattn_cache") and self.crossattn_cache is not None