AstroAUmin/self-forcing
0
1from tqdm import tqdm2from typing import List, Optional3import torch4 5from wan.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps6from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler7from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper8 9 10class CausalDiffusionInferencePipeline(torch.nn.Module):11 def __init__(12 self,13 args,14 device,15 generator=None,16 text_encoder=None,17 vae=None18 ):19 super().__init__()20 # Step 1: Initialize all models21 self.generator = WanDiffusionWrapper(22 **getattr(args, "model_kwargs", {}), is_causal=True) if generator is None else generator23 self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder24 self.vae = WanVAEWrapper() if vae is None else vae25 26 # Step 2: Initialize scheduler27 self.num_train_timesteps = args.num_train_timestep28 self.sampling_steps = 5029 self.sample_solver = 'unipc'30 self.shift = args.timestep_shift31 32 self.num_transformer_blocks = 3033 self.frame_seq_length = 156034 35 self.kv_cache_pos = None36 self.kv_cache_neg = None37 self.crossattn_cache_pos = None38 self.crossattn_cache_neg = None39 self.args = args40 self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)41 self.independent_first_frame = args.independent_first_frame42 self.local_attn_size = self.generator.model.local_attn_size43 44 print(f"KV inference with {self.num_frame_per_block} frames per block")45 46 if self.num_frame_per_block > 1:47 self.generator.model.num_frame_per_block = self.num_frame_per_block48 49 def inference(50 self,51 noise: torch.Tensor,52 text_prompts: List[str],53 initial_latent: Optional[torch.Tensor] = None,54 return_latents: bool = False,55 start_frame_index: Optional[int] = 056 ) -> torch.Tensor:57 """58 Perform inference on the given noise and text prompts.59 Inputs:60 noise (torch.Tensor): The input noise tensor of shape61 (batch_size, num_output_frames, num_channels, height, width).62 text_prompts (List[str]): The list of text prompts.63 initial_latent (torch.Tensor): The initial latent tensor of shape64 (batch_size, num_input_frames, num_channels, height, width).65 If num_input_frames is 1, perform image to video.66 If num_input_frames is greater than 1, perform video extension.67 return_latents (bool): Whether to return the latents.68 start_frame_index (int): In long video generation, where does the current window start?69 Outputs:70 video (torch.Tensor): The generated video tensor of shape71 (batch_size, num_frames, num_channels, height, width). It is normalized to be in the range [0, 1].72 """73 batch_size, num_frames, num_channels, height, width = noise.shape74 if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):75 # If the first frame is independent and the first frame is provided, then the number of frames in the76 # noise should still be a multiple of num_frame_per_block77 assert num_frames % self.num_frame_per_block == 078 num_blocks = num_frames // self.num_frame_per_block79 elif self.independent_first_frame and initial_latent is None:80 # Using a [1, 4, 4, 4, 4, 4] model to generate a video without image conditioning81 assert (num_frames - 1) % self.num_frame_per_block == 082 num_blocks = (num_frames - 1) // self.num_frame_per_block83 num_input_frames = initial_latent.shape[1] if initial_latent is not None else 084 num_output_frames = num_frames + num_input_frames # add the initial latent frames85 conditional_dict = self.text_encoder(86 text_prompts=text_prompts87 )88 unconditional_dict = self.text_encoder(89 text_prompts=[self.args.negative_prompt] * len(text_prompts)90 )91 92 output = torch.zeros(93 [batch_size, num_output_frames, num_channels, height, width],94 device=noise.device,95 dtype=noise.dtype96 )97 98 # Step 1: Initialize KV cache to all zeros99 if self.kv_cache_pos is None:100 self._initialize_kv_cache(101 batch_size=batch_size,102 dtype=noise.dtype,103 device=noise.device104 )105 self._initialize_crossattn_cache(106 batch_size=batch_size,107 dtype=noise.dtype,108 device=noise.device109 )110 else:111 # reset cross attn cache112 for block_index in range(self.num_transformer_blocks):113 self.crossattn_cache_pos[block_index]["is_init"] = False114 self.crossattn_cache_neg[block_index]["is_init"] = False115 # reset kv cache116 for block_index in range(len(self.kv_cache_pos)):117 self.kv_cache_pos[block_index]["global_end_index"] = torch.tensor(118 [0], dtype=torch.long, device=noise.device)119 self.kv_cache_pos[block_index]["local_end_index"] = torch.tensor(120 [0], dtype=torch.long, device=noise.device)121 self.kv_cache_neg[block_index]["global_end_index"] = torch.tensor(122 [0], dtype=torch.long, device=noise.device)123 self.kv_cache_neg[block_index]["local_end_index"] = torch.tensor(124 [0], dtype=torch.long, device=noise.device)125 126 # Step 2: Cache context feature127 current_start_frame = start_frame_index128 cache_start_frame = 0129 if initial_latent is not None:130 timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0131 if self.independent_first_frame:132 # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks133 assert (num_input_frames - 1) % self.num_frame_per_block == 0134 num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block135 output[:, :1] = initial_latent[:, :1]136 self.generator(137 noisy_image_or_video=initial_latent[:, :1],138 conditional_dict=conditional_dict,139 timestep=timestep * 0,140 kv_cache=self.kv_cache_pos,141 crossattn_cache=self.crossattn_cache_pos,142 current_start=current_start_frame * self.frame_seq_length,143 cache_start=cache_start_frame * self.frame_seq_length144 )145 self.generator(146 noisy_image_or_video=initial_latent[:, :1],147 conditional_dict=unconditional_dict,148 timestep=timestep * 0,149 kv_cache=self.kv_cache_neg,150 crossattn_cache=self.crossattn_cache_neg,151 current_start=current_start_frame * self.frame_seq_length,152 cache_start=cache_start_frame * self.frame_seq_length153 )154 current_start_frame += 1155 cache_start_frame += 1156 else:157 # Assume num_input_frames is self.num_frame_per_block * num_input_blocks158 assert num_input_frames % self.num_frame_per_block == 0159 num_input_blocks = num_input_frames // self.num_frame_per_block160 161 for block_index in range(num_input_blocks):162 current_ref_latents = \163 initial_latent[:, cache_start_frame:cache_start_frame + self.num_frame_per_block]164 output[:, cache_start_frame:cache_start_frame + self.num_frame_per_block] = current_ref_latents165 self.generator(166 noisy_image_or_video=current_ref_latents,167 conditional_dict=conditional_dict,168 timestep=timestep * 0,169 kv_cache=self.kv_cache_pos,170 crossattn_cache=self.crossattn_cache_pos,171 current_start=current_start_frame * self.frame_seq_length,172 cache_start=cache_start_frame * self.frame_seq_length173 )174 self.generator(175 noisy_image_or_video=current_ref_latents,176 conditional_dict=unconditional_dict,177 timestep=timestep * 0,178 kv_cache=self.kv_cache_neg,179 crossattn_cache=self.crossattn_cache_neg,180 current_start=current_start_frame * self.frame_seq_length,181 cache_start=cache_start_frame * self.frame_seq_length182 )183 current_start_frame += self.num_frame_per_block184 cache_start_frame += self.num_frame_per_block185 186 # Step 3: Temporal denoising loop187 all_num_frames = [self.num_frame_per_block] * num_blocks188 if self.independent_first_frame and initial_latent is None:189 all_num_frames = [1] + all_num_frames190 for current_num_frames in all_num_frames:191 noisy_input = noise[192 :, cache_start_frame - num_input_frames:cache_start_frame + current_num_frames - num_input_frames]193 latents = noisy_input194 195 # Step 3.1: Spatial denoising loop196 sample_scheduler = self._initialize_sample_scheduler(noise)197 for _, t in enumerate(tqdm(sample_scheduler.timesteps)):198 latent_model_input = latents199 timestep = t * torch.ones(200 [batch_size, current_num_frames], device=noise.device, dtype=torch.float32201 )202 203 flow_pred_cond, _ = self.generator(204 noisy_image_or_video=latent_model_input,205 conditional_dict=conditional_dict,206 timestep=timestep,207 kv_cache=self.kv_cache_pos,208 crossattn_cache=self.crossattn_cache_pos,209 current_start=current_start_frame * self.frame_seq_length,210 cache_start=cache_start_frame * self.frame_seq_length211 )212 flow_pred_uncond, _ = self.generator(213 noisy_image_or_video=latent_model_input,214 conditional_dict=unconditional_dict,215 timestep=timestep,216 kv_cache=self.kv_cache_neg,217 crossattn_cache=self.crossattn_cache_neg,218 current_start=current_start_frame * self.frame_seq_length,219 cache_start=cache_start_frame * self.frame_seq_length220 )221 222 flow_pred = flow_pred_uncond + self.args.guidance_scale * (223 flow_pred_cond - flow_pred_uncond)224 225 temp_x0 = sample_scheduler.step(226 flow_pred,227 t,228 latents,229 return_dict=False)[0]230 latents = temp_x0231 print(f"kv_cache['local_end_index']: {self.kv_cache_pos[0]['local_end_index']}")232 print(f"kv_cache['global_end_index']: {self.kv_cache_pos[0]['global_end_index']}")233 234 # Step 3.2: record the model's output235 output[:, cache_start_frame:cache_start_frame + current_num_frames] = latents236 237 # Step 3.3: rerun with timestep zero to update KV cache using clean context238 self.generator(239 noisy_image_or_video=latents,240 conditional_dict=conditional_dict,241 timestep=timestep * 0,242 kv_cache=self.kv_cache_pos,243 crossattn_cache=self.crossattn_cache_pos,244 current_start=current_start_frame * self.frame_seq_length,245 cache_start=cache_start_frame * self.frame_seq_length246 )247 self.generator(248 noisy_image_or_video=latents,249 conditional_dict=unconditional_dict,250 timestep=timestep * 0,251 kv_cache=self.kv_cache_neg,252 crossattn_cache=self.crossattn_cache_neg,253 current_start=current_start_frame * self.frame_seq_length,254 cache_start=cache_start_frame * self.frame_seq_length255 )256 257 # Step 3.4: update the start and end frame indices258 current_start_frame += current_num_frames259 cache_start_frame += current_num_frames260 261 # Step 4: Decode the output262 video = self.vae.decode_to_pixel(output)263 video = (video * 0.5 + 0.5).clamp(0, 1)264 265 if return_latents:266 return video, output267 else:268 return video269 270 def _initialize_kv_cache(self, batch_size, dtype, device):271 """272 Initialize a Per-GPU KV cache for the Wan model.273 """274 kv_cache_pos = []275 kv_cache_neg = []276 if self.local_attn_size != -1:277 # Use the local attention size to compute the KV cache size278 kv_cache_size = self.local_attn_size * self.frame_seq_length279 else:280 # Use the default KV cache size281 kv_cache_size = 32760282 283 for _ in range(self.num_transformer_blocks):284 kv_cache_pos.append({285 "k": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),286 "v": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),287 "global_end_index": torch.tensor([0], dtype=torch.long, device=device),288 "local_end_index": torch.tensor([0], dtype=torch.long, device=device)289 })290 kv_cache_neg.append({291 "k": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),292 "v": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),293 "global_end_index": torch.tensor([0], dtype=torch.long, device=device),294 "local_end_index": torch.tensor([0], dtype=torch.long, device=device)295 })296 297 self.kv_cache_pos = kv_cache_pos # always store the clean cache298 self.kv_cache_neg = kv_cache_neg # always store the clean cache299 300 def _initialize_crossattn_cache(self, batch_size, dtype, device):301 """302 Initialize a Per-GPU cross-attention cache for the Wan model.303 """304 crossattn_cache_pos = []305 crossattn_cache_neg = []306 for _ in range(self.num_transformer_blocks):307 crossattn_cache_pos.append({308 "k": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),309 "v": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),310 "is_init": False311 })312 crossattn_cache_neg.append({313 "k": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),314 "v": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),315 "is_init": False316 })317 318 self.crossattn_cache_pos = crossattn_cache_pos # always store the clean cache319 self.crossattn_cache_neg = crossattn_cache_neg # always store the clean cache320 321 def _initialize_sample_scheduler(self, noise):322 if self.sample_solver == 'unipc':323 sample_scheduler = FlowUniPCMultistepScheduler(324 num_train_timesteps=self.num_train_timesteps,325 shift=1,326 use_dynamic_shifting=False)327 sample_scheduler.set_timesteps(328 self.sampling_steps, device=noise.device, shift=self.shift)329 self.timesteps = sample_scheduler.timesteps330 elif self.sample_solver == 'dpm++':331 sample_scheduler = FlowDPMSolverMultistepScheduler(332 num_train_timesteps=self.num_train_timesteps,333 shift=1,334 use_dynamic_shifting=False)335 sampling_sigmas = get_sampling_sigmas(self.sampling_steps, self.shift)336 self.timesteps, _ = retrieve_timesteps(337 sample_scheduler,338 device=noise.device,339 sigmas=sampling_sigmas)340 else:341 raise NotImplementedError("Unsupported solver.")342 return sample_scheduler343