Spaces:
Sleeping
Sleeping
| from typing import List, Optional | |
| import torch | |
| from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper | |
| from utils.memory import gpu, get_cuda_free_memory_gb, move_model_to_device_with_memory_preservation | |
| from pipeline.causal_inference import CausalInferencePipeline | |
| import torch.distributed as dist | |
| class SwitchCausalInferencePipeline(CausalInferencePipeline): | |
| def __init__(self, args, device, *, generator: WanDiffusionWrapper | None=None, text_encoder: WanTextEncoder | None=None, vae: WanVAEWrapper | None=None): | |
| super().__init__(args, device, generator=generator, text_encoder=text_encoder, vae=vae) | |
| self.global_sink = getattr(args, 'global_sink', False) | |
| def _recache_after_switch(self, output, current_start_frame, new_conditional_dict): | |
| if not self.global_sink: | |
| for block_idx in range(self.num_transformer_blocks): | |
| cache = self.kv_cache1[block_idx] | |
| cache['k'].zero_() | |
| cache['v'].zero_() | |
| for blk in self.crossattn_cache: | |
| blk['k'].zero_() | |
| blk['v'].zero_() | |
| blk['is_init'] = False | |
| inner = self.generator.model | |
| if getattr(inner, '_ei_prev_window_start', None) is not None: | |
| inner._ei_prev_window_start = None | |
| if current_start_frame == 0: | |
| return | |
| num_recache_frames = current_start_frame if self.local_attn_size == -1 else min(self.local_attn_size, current_start_frame) | |
| recache_start_frame = current_start_frame - num_recache_frames | |
| frames_to_recache = output[:, recache_start_frame:current_start_frame] | |
| if frames_to_recache.device.type == 'cpu': | |
| target_device = next(self.generator.parameters()).device | |
| frames_to_recache = frames_to_recache.to(target_device) | |
| batch_size = frames_to_recache.shape[0] | |
| print(f'num_recache_frames: {num_recache_frames}, recache_start_frame: {recache_start_frame}, current_start_frame: {current_start_frame}') | |
| device = frames_to_recache.device | |
| block_mask = self.generator.model._prepare_blockwise_causal_attn_mask(device=device, num_frames=num_recache_frames, frame_seqlen=self.frame_seq_length, num_frame_per_block=self.num_frame_per_block, local_attn_size=self.local_attn_size) | |
| context_timestep = torch.ones([batch_size, recompute_frames], device=device, dtype=torch.int64) * self.args.context_noise | |
| self.generator.model.block_mask = block_mask | |
| with torch.no_grad(): | |
| self.generator(noisy_image_or_video=frames_to_recompute, conditional_dict=new_conditional_dict, timestep=context_timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=recompute_start_frame * self.frame_seq_length) | |
| for blk in self.crossattn_cache: | |
| blk['k'].zero_() | |
| blk['v'].zero_() | |
| blk['is_init'] = False | |
| def inference(self, noise: torch.Tensor, *, text_prompts_first: List[str], text_prompts_second: List[str], switch_frame_index: int, initial_latent: Optional[torch.Tensor]=None, return_latents: bool=False, low_memory: bool=False): | |
| batch_size, num_output_frames, num_channels, height, width = noise.shape | |
| assert num_output_frames % self.num_frame_per_block == 0 | |
| num_blocks = num_output_frames // self.num_frame_per_block | |
| cond_first = self.text_encoder(text_prompts=text_prompts_first) | |
| cond_second = self.text_encoder(text_prompts=text_prompts_second) | |
| 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_device = torch.device('cpu') if low_memory else noise.device | |
| output = torch.zeros([batch_size, num_output_frames, num_channels, height, width], device=output_device, dtype=noise.dtype) | |
| local_attn_cfg = getattr(self.args.model_kwargs, 'local_attn_size', -1) | |
| kv_policy = '' | |
| if local_attn_cfg != -1: | |
| kv_cache_size = local_attn_cfg * self.frame_seq_length | |
| kv_policy = f'int->local, size={local_attn_cfg}' | |
| else: | |
| kv_cache_size = num_output_frames * self.frame_seq_length | |
| kv_policy = 'global (-1)' | |
| print(f'kv_cache_size: {kv_cache_size} (policy: {kv_policy}, frame_seq_length: {self.frame_seq_length}, num_output_frames: {num_output_frames})') | |
| self._initialize_kv_cache(batch_size, dtype=noise.dtype, device=noise.device, kv_cache_size_override=kv_cache_size) | |
| self._initialize_crossattn_cache(batch_size=batch_size, dtype=noise.dtype, device=noise.device) | |
| num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0 | |
| current_start_frame = 0 | |
| self.generator.model.local_attn_size = self.local_attn_size | |
| print(f'[inference] local_attn_size set on model: {self.generator.model.local_attn_size}') | |
| self._set_all_modules_max_attention_size(self.local_attn_size) | |
| all_num_frames = [self.num_frame_per_block] * num_blocks | |
| using_second = False | |
| for current_num_frames in all_num_frames: | |
| if not using_second and current_start_frame >= switch_frame_index: | |
| self._recache_after_switch(output, current_start_frame, cond_second) | |
| cond_in_use = cond_second | |
| using_second = True | |
| print('switch_frame_index', switch_frame_index) | |
| print('current_start_frame', current_start_frame) | |
| else: | |
| cond_in_use = cond_second if using_second else cond_first | |
| noisy_input = noise[:, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames] | |
| for index, current_timestep in enumerate(self.denoising_step_list): | |
| timestep = torch.ones([batch_size, current_num_frames], device=noise.device, dtype=torch.int64) * current_timestep | |
| if index < len(self.denoising_step_list) - 1: | |
| _, denoised_pred = self.generator(noisy_image_or_video=noisy_input, conditional_dict=cond_in_use, timestep=timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length) | |
| next_timestep = self.denoising_step_list[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]) | |
| else: | |
| _, denoised_pred = self.generator(noisy_image_or_video=noisy_input, conditional_dict=cond_in_use, timestep=timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length) | |
| output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred.to(output.device) | |
| context_timestep = torch.ones_like(timestep) * self.args.context_noise | |
| self.generator(noisy_image_or_video=denoised_pred, conditional_dict=cond_in_use, timestep=context_timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length) | |
| current_start_frame += current_num_frames | |
| video = self.vae.decode_to_pixel(output.to(noise.device), use_cache=False) | |
| video = (video * 0.5 + 0.5).clamp(0, 1) | |
| if return_latents: | |
| return (video, output) | |
| return video | |