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 InteractiveCausalInferencePipeline(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) self.apr_enabled = getattr(args, 'apr_enabled', False) self.apr_alpha_max = float(max(0.0, min(0.8, getattr(args, 'apr_alpha_max', 0.8)))) self.apr_d_window = getattr(args, 'apr_d_window', None) self.apr_blend_sink = getattr(args, 'apr_blend_sink', False) def _recache_after_switch(self, output, current_start_frame, new_conditional_dict): sink_tok_size = None sink_backup = None sink_norm_before = None if self.global_sink and current_start_frame > 0: inner_tmp = self.generator.model sink_tok_size = inner_tmp.blocks[0].self_attn.sink_size * self.frame_seq_length sink_backup = [{'k': self.kv_cache1[i]['k'][:, :sink_tok_size].clone(), 'v': self.kv_cache1[i]['v'][:, :sink_tok_size].clone()} for i in range(self.num_transformer_blocks)] sink_norm_before = torch.norm(self.kv_cache1[0]['k'][:, :sink_tok_size]).float().item() print(f'[Recache-Inf] switch@frame={current_start_frame}, global_sink=True, sink_tok={sink_tok_size}, sink_norm(block0) BEFORE = {sink_norm_before:.4f}') elif not self.global_sink: print(f'[Recache-Inf] switch@frame={current_start_frame}, global_sink=False (sink will be recached with new prompt)') apr_old_cache = None if self.apr_enabled and current_start_frame > 0: apr_old_cache = [{'k': self.kv_cache1[i]['k'].clone(), 'v': self.kv_cache1[i]['v'].clone()} for i in range(self.num_transformer_blocks)] 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 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, num_recache_frames], device=device, dtype=torch.int64) * self.args.context_noise self.generator.model.block_mask = block_mask _has_memory = getattr(inner, 'query_memory_encoder', None) is not None or getattr(inner, 'sink_memory', None) is not None if _has_memory: frame_seqlen = self.frame_seq_length sink_frames = inner.blocks[0].self_attn.sink_size max_attn = inner.blocks[0].self_attn.max_attention_size recent_window_frames = (max_attn - sink_frames * frame_seqlen) // frame_seqlen inner._ei_prev_window_start = max(sink_frames, current_start_frame - recent_window_frames) with torch.no_grad(): self.generator(noisy_image_or_video=frames_to_recache, conditional_dict=new_conditional_dict, timestep=context_timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=recache_start_frame * self.frame_seq_length, sink_recache_after_switch=not self.global_sink) if sink_backup is not None: for i in range(self.num_transformer_blocks): self.kv_cache1[i]['k'][:, :sink_tok_size] = sink_backup[i]['k'] self.kv_cache1[i]['v'][:, :sink_tok_size] = sink_backup[i]['v'] sink_norm_after = torch.norm(self.kv_cache1[0]['k'][:, :sink_tok_size]).float().item() delta = abs(sink_norm_after - sink_norm_before) status = 'PRESERVED ✓' if delta < 0.001 else f'MISMATCH ✗ (delta={delta:.4e})' print(f'[Recache-Inf] global_sink=True: sink_norm AFTER = {sink_norm_after:.4f} [{status}]') del sink_backup for blk in self.crossattn_cache: blk['k'].zero_() blk['v'].zero_() blk['is_init'] = False if self.apr_enabled and apr_old_cache is not None: self._apply_apr_blend(old_cache=apr_old_cache, recache_start_frame=recache_start_frame, current_start_frame=current_start_frame, num_recache_frames=num_recache_frames) del apr_old_cache if getattr(inner, 'query_memory_encoder', None) is not None: enc = inner.query_memory_encoder if getattr(enc, 'memory_recache', False): enc.reset(batch_size=batch_size, device=frames_to_recache.device, dtype=torch.bfloat16) last_blk = self.num_transformer_blocks - 1 cache = self.kv_cache1[last_blk] frame_seqlen = self.frame_seq_length sink_tok = inner.blocks[0].self_attn.sink_size * frame_seqlen local_end = cache['local_end_index'].item() if local_end > sink_tok: recache_k = cache['k'][:, sink_tok:local_end].clone() recache_v = cache['v'][:, sink_tok:local_end].clone() sink_k = cache['k'][:, :sink_tok].clone() if sink_tok > 0 else None sink_v = cache['v'][:, :sink_tok].clone() if sink_tok > 0 else None enc.update(recache_k, recache_v, sink_k, sink_v) if getattr(inner, 'sink_memory', None) is not None: sm = inner.sink_memory sm.reset() sink_frames_count = inner.blocks[0].self_attn.sink_size sink_output = output[:, :sink_frames_count] if sink_output.device.type == 'cpu': sink_output = sink_output.to(device) sink_ts = torch.ones([batch_size, sink_frames_count], device=device, dtype=torch.int64) * self.args.context_noise n_heads = inner.blocks[0].self_attn.num_heads head_dim = inner.blocks[0].self_attn.head_dim sink_cache_size = sink_frames_count * self.frame_seq_length temp_kv = [{'k': torch.zeros([batch_size, sink_cache_size, n_heads, head_dim], device=device, dtype=torch.bfloat16), 'v': torch.zeros([batch_size, sink_cache_size, n_heads, head_dim], device=device, dtype=torch.bfloat16), 'global_end_index': torch.zeros(1, dtype=torch.long, device=device), 'local_end_index': torch.zeros(1, dtype=torch.long, device=device)} for _ in range(self.num_transformer_blocks)] temp_crossattn = [{'k': torch.zeros_like(self.crossattn_cache[0]['k']), 'v': torch.zeros_like(self.crossattn_cache[0]['v']), 'is_init': False} for _ in range(self.num_transformer_blocks)] with torch.no_grad(): self.generator(noisy_image_or_video=sink_output, conditional_dict=new_conditional_dict, timestep=sink_ts, kv_cache=temp_kv, crossattn_cache=temp_crossattn, current_start=0) inner._ei_prev_window_start = max(sink_frames, current_start_frame - recent_window_frames) print(f'[SinkMem] reset + re-captured on prompt switch at frame {current_start_frame}') def _apply_apr_blend(self, *, old_cache, recache_start_frame, current_start_frame, num_recache_frames): inner = self.generator.model frame_seqlen = self.frame_seq_length sink_size = inner.blocks[0].self_attn.sink_size sink_tok = sink_size * frame_seqlen D = self.apr_d_window if self.apr_d_window is not None else num_recache_frames D = max(1, int(D)) alpha_max = self.apr_alpha_max n_frames = num_recache_frames with torch.no_grad(): for blk_idx in range(self.num_transformer_blocks): new_k = self.kv_cache1[blk_idx]['k'] new_v = self.kv_cache1[blk_idx]['v'] old_k = old_cache[blk_idx]['k'] old_v = old_cache[blk_idx]['v'] if self.apr_blend_sink and (not self.global_sink): for f in range(sink_size): d_t = current_start_frame - f alpha = min(alpha_max, max(0.0, 1.0 - d_t / D)) s0, s1 = (f * frame_seqlen, (f + 1) * frame_seqlen) new_k[:, s0:s1] = (1 - alpha) * old_k[:, s0:s1] + alpha * new_k[:, s0:s1] new_v[:, s0:s1] = (1 - alpha) * old_v[:, s0:s1] + alpha * new_v[:, s0:s1] base = sink_tok for f in range(n_frames): absolute_frame = recache_start_frame + f d_t = current_start_frame - absolute_frame alpha = min(alpha_max, max(0.0, 1.0 - d_t / D)) s0, s1 = (base + f * frame_seqlen, base + (f + 1) * frame_seqlen) new_k[:, s0:s1] = (1 - alpha) * old_k[:, s0:s1] + alpha * new_k[:, s0:s1] new_v[:, s0:s1] = (1 - alpha) * old_v[:, s0:s1] + alpha * new_v[:, s0:s1] if not dist.is_initialized() or dist.get_rank() == 0: print(f'[APR] switch@{current_start_frame}: blended {n_frames} frames (α_max={alpha_max}, D_window={D})') def inference(self, noise: torch.Tensor, *, text_prompts_list: List[List[str]], switch_frame_indices: List[int], return_latents: bool=False, low_memory: bool=False): batch_size, num_output_frames, num_channels, height, width = noise.shape assert len(text_prompts_list) >= 1, 'text_prompts_list must not be empty' assert len(switch_frame_indices) == len(text_prompts_list) - 1, 'length of switch_frame_indices should be one less than text_prompts_list' assert num_output_frames % self.num_frame_per_block == 0 num_blocks = num_output_frames // self.num_frame_per_block print(text_prompts_list) cond_list = [self.text_encoder(text_prompts=p) for p in text_prompts_list] 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) 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 segment_idx = 0 next_switch_pos = switch_frame_indices[segment_idx] if segment_idx < len(switch_frame_indices) else None for current_num_frames in all_num_frames: if next_switch_pos is not None and current_start_frame >= next_switch_pos: segment_idx += 1 self._recache_after_switch(output, current_start_frame, cond_list[segment_idx]) next_switch_pos = switch_frame_indices[segment_idx] if segment_idx < len(switch_frame_indices) else None print(f'segment_idx: {segment_idx}') print(f'text_prompts_list[segment_idx]: {text_prompts_list[segment_idx]}') cond_in_use = cond_list[segment_idx] noisy_input = noise[:, current_start_frame:current_start_frame + current_num_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