Spaces:
Sleeping
Sleeping
File size: 7,690 Bytes
3e936b2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | 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
|