diff --git a/generate_infinitetalk.py b/generate_infinitetalk.py index e83daa8..77426ce 100644 --- a/generate_infinitetalk.py +++ b/generate_infinitetalk.py @@ -20,7 +20,10 @@ import wan from wan.configs import SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS from wan.utils.utils import str2bool, is_video, split_wav_librosa from wan.utils.multitalk_utils import save_video_ffmpeg -from kokoro import KPipeline +try: + from kokoro import KPipeline +except ImportError: # TTS stack unused when driving from a local audio file + KPipeline = None from transformers import Wav2Vec2FeatureExtractor from src.audio_analysis.wav2vec2 import Wav2Vec2Model from wan.utils.segvideo import shot_detect @@ -544,6 +547,89 @@ def generate(args): num_persistent_param_in_dit=args.num_persistent_param_in_dit ) + if os.environ.get("WAN_TRT") == "1": + # The engines own the transformer stack now: keep the PyTorch blocks off + # the GPU (~27GB) so the 18.6GB of engines fit alongside the rest. + import torch as _t + _n = len(wan_i2v.model.blocks) + wan_i2v.model.blocks = _t.nn.ModuleList([]) # the engines own these now + _t.cuda.empty_cache() + logging.info(f"WAN_TRT: dropped {_n} PyTorch blocks (~27GB); engines own the stack") + + if os.environ.get("CAPTURE") == "1": + import sys as _sys + _sys.path.insert(0, "/workspace/trt") + from capture_calib import install_capture + install_capture(wan_i2v.model) + + # --- FP8 (H100 native) + torch.compile toggles, applied after LoRA merge --- + if os.environ.get("WAN_FP8") == "1": + import torch as _t + from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig + _dit = wan_i2v.model + # Only the transformer blocks: the tiny embedders/head stay bf16. + _n = 0 + for _blk in _dit.blocks: + quantize_(_blk, Float8DynamicActivationFloat8WeightConfig()) + _n += 1 + logging.info(f"FP8: quantized {_n} transformer blocks (fp8 dynamic act + fp8 weight)") + + if os.environ.get("WAN_COMPILE") == "1": + import torch as _t + _dit = wan_i2v.model + for _i, _blk in enumerate(_dit.blocks): + _dit.blocks[_i] = _t.compile(_blk, dynamic=False) + logging.info(f"compiled {len(_dit.blocks)} transformer blocks") + + # --- profiling hook: set PROFILE=1 to time every DiT forward --- + if os.environ.get("PROFILE") == "1": + import time as _time, atexit as _atexit + import torch as _torch + _stats = {"n": 0, "t": 0.0} + _wall0 = _time.perf_counter() + _orig_fwd = wan_i2v.model.forward + + _deep = os.environ.get("PROFILE_DEEP") == "1" + + def _timed_fwd(*a, **kw): + _torch.cuda.synchronize() + _t0 = _time.perf_counter() + # kernel-level trace of a single steady-state forward (the 2nd) + if _deep and _stats["n"] == 1: + from torch.profiler import profile as _tp, ProfilerActivity as _PA + with _tp(activities=[_PA.CPU, _PA.CUDA], record_shapes=False) as _prof: + out = _orig_fwd(*a, **kw) + _torch.cuda.synchronize() + print("\n======= TOP CUDA KERNELS (one forward) =======") + print(_prof.key_averages().table( + sort_by="self_cuda_time_total", row_limit=28, + max_name_column_width=55)) + else: + out = _orig_fwd(*a, **kw) + _torch.cuda.synchronize() + _stats["t"] += _time.perf_counter() - _t0 + _stats["n"] += 1 + return out + + wan_i2v.model.forward = _timed_fwd + + def _report(): + n, t = _stats["n"], _stats["t"] + wall = _time.perf_counter() - _wall0 + peak = _torch.cuda.max_memory_allocated() / 1e9 + print("\n================ PROFILE ================") + print(f"DiT forwards : {n}") + print(f"DiT total time : {t:.1f} s") + if n: + print(f"DiT per forward : {t / n * 1000:.0f} ms") + print(f"wall (post-load) : {wall:.1f} s") + if wall > 0: + print(f"DiT share of wall : {t / wall * 100:.0f} %") + print(f"peak VRAM allocated : {peak:.1f} GB") + print("=========================================") + + _atexit.register(_report) + generated_list = [] with open(args.input_json, 'r', encoding='utf-8') as f: input_data = json.load(f) diff --git a/wan/modules/multitalk_model.py b/wan/modules/multitalk_model.py index 958e930..6eef6f3 100644 --- a/wan/modules/multitalk_model.py +++ b/wan/modules/multitalk_model.py @@ -21,6 +21,29 @@ try: except: USE_SAGEATTN = False +from torch.nn.attention import sdpa_kernel, SDPBackend + +# cuDNN's Hopper attention beat sageattn 44.6ms vs 58.5ms at our shape. +USE_CUDNN_SDPA = os.environ.get("WAN_CUDNN_ATTN", "1") == "1" + +# --- FP8 TensorRT block stack (WAN_TRT=1) ------------------------------------- +_TRT_STACK = None +_TRT_SEQ_LEN = int(os.environ.get("WAN_TRT_SEQ_LEN", "30576")) # 81 frames @ 448x832 + + +def _trt_stack(): + """Lazily load the engines; returns None when TRT is off.""" + global _TRT_STACK + if os.environ.get("WAN_TRT") != "1": + return None + if _TRT_STACK is None: + import sys + sys.path.insert(0, "/workspace/trt") + from trt_runner import TRTStack + _TRT_STACK = TRTStack() + return _TRT_STACK + + __all__ = ['WanModel'] @@ -50,30 +73,52 @@ def rope_params(max_seq_len, dim, theta=10000): return freqs +_ROPE_CACHE = {} + + +def _rope_cos_sin(grid_sizes, freqs, device, dtype): + """cos/sin RoPE table for this latent grid, built once and kept on-device.""" + f, h, w = (int(v) for v in grid_sizes[0].tolist()) + key = (f, h, w, str(device), dtype) + hit = _ROPE_CACHE.get(key) + if hit is not None: + return hit + + c = freqs.size(1) + fr = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + freqs_i = torch.cat([ + fr[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + fr[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + fr[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], dim=-1).reshape(f * h * w, 1, -1) # [L, 1, C/2] complex + + cos = freqs_i.real.to(device=device, dtype=dtype).contiguous() + sin = freqs_i.imag.to(device=device, dtype=dtype).contiguous() + _ROPE_CACHE[key] = (cos, sin) + return cos, sin + + @amp.autocast(enabled=False) def rope_apply(x, grid_sizes, freqs): - s, n, c = x.size(1), x.size(2), x.size(3) // 2 + """Real-valued RoPE. Mathematically identical to the complex fp64 version, + minus the per-layer host->device copy and the float64 traffic.""" + b, s, n, d = x.shape + cos, sin = _rope_cos_sin(grid_sizes, freqs, x.device, torch.float32) + L = cos.size(0) - freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + xf = x[:, :L].float().reshape(b, L, n, d // 2, 2) + x_r, x_i = xf[..., 0], xf[..., 1] - output = [] - for i, (f, h, w) in enumerate(grid_sizes.tolist()): - seq_len = f * h * w + cos_ = cos.unsqueeze(0) # [1, L, 1, C/2] + sin_ = sin.unsqueeze(0) - x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape( - s, n, -1, 2)) - freqs_i = torch.cat([ - freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), - freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), - freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) - ], - dim=-1).reshape(seq_len, 1, -1) - freqs_i = freqs_i.to(device=x_i.device) - x_i = torch.view_as_real(x_i * freqs_i).flatten(2) - x_i = torch.cat([x_i, x[i, seq_len:]]) + o_r = x_r * cos_ - x_i * sin_ + o_i = x_r * sin_ + x_i * cos_ + out = torch.stack([o_r, o_i], dim=-1).flatten(3) - output.append(x_i) - return torch.stack(output).float() + if L < s: # keep any padding untouched + out = torch.cat([out, x[:, L:].float()], dim=1) + return out.float() class WanRMSNorm(nn.Module): @@ -137,7 +182,7 @@ class WanSelfAttention(nn.Module): self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() - def forward(self, x, seq_lens, grid_sizes, freqs, ref_target_masks=None): + def forward(self, x, seq_lens, grid_sizes, freqs, ref_target_masks=None, human_num=None): b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim # query, key, value function @@ -151,7 +196,15 @@ class WanSelfAttention(nn.Module): q = rope_apply(q, grid_sizes, freqs) k = rope_apply(k, grid_sizes, freqs) - if USE_SAGEATTN: + if USE_CUDNN_SDPA: + # [B, L, H, D] -> [B, H, L, D] for SDPA, and back + qb = q.transpose(1, 2).to(torch.bfloat16) + kb = k.transpose(1, 2).to(torch.bfloat16) + vb = v.transpose(1, 2).to(torch.bfloat16) + with sdpa_kernel(SDPBackend.CUDNN_ATTENTION): + x = F.scaled_dot_product_attention(qb, kb, vb) + x = x.transpose(1, 2).type_as(v) + elif USE_SAGEATTN: x = sageattn(q.to(torch.bfloat16), k.to(torch.bfloat16), v, tensor_layout='NHD') else: x = flash_attention( @@ -165,9 +218,15 @@ class WanSelfAttention(nn.Module): # output x = x.flatten(2) x = self.o(x) - with torch.no_grad(): - x_ref_attn_map = get_attn_map_with_target(q.type_as(x), k.type_as(x), grid_sizes[0], - ref_target_masks=ref_target_masks) + # The ref-attn map only feeds SingleStreamMutiAttention's multi-speaker routing; + # with one speaker that branch short-circuits and never reads it, so skip building + # a [heads, seq, ref_seq] map (GBs) in every layer. + if human_num == 1: + x_ref_attn_map = None + else: + with torch.no_grad(): + x_ref_attn_map = get_attn_map_with_target(q.type_as(x), k.type_as(x), grid_sizes[0], + ref_target_masks=ref_target_masks) return x, x_ref_attn_map @@ -294,7 +353,7 @@ class WanAttentionBlock(nn.Module): # self-attention y, x_ref_attn_map = self.self_attn( (self.norm1(x).float() * (1 + e[1]) + e[0]).type_as(x), seq_lens, grid_sizes, - freqs, ref_target_masks=ref_target_masks) + freqs, ref_target_masks=ref_target_masks, human_num=human_num) with amp.autocast(dtype=torch.float32): x = x + y * e[2] @@ -757,7 +816,14 @@ class WanModel(ModelMixin, ConfigMixin): for block in self.blocks: x = block(x, **kwargs) self.previous_residual_uncond = x - ori_x + elif _trt_stack() is not None and x.shape[1] == _TRT_SEQ_LEN: + cos, sin = _rope_cos_sin(grid_sizes, self.freqs, x.device, torch.float32) + x = _trt_stack()(x, e0, context, audio_embedding.squeeze(0), cos, sin) else: + if os.environ.get("WAN_TRT") == "1": + raise RuntimeError( + f"WAN_TRT=1 but seq_len {x.shape[1]} != engine seq_len {_TRT_SEQ_LEN}. " + f"Engines are static; rebuild them for this frame count/resolution.") for block in self.blocks: x = block(x, **kwargs) diff --git a/wan/multitalk.py b/wan/multitalk.py index be7819e..1e9513b 100644 --- a/wan/multitalk.py +++ b/wan/multitalk.py @@ -1,6 +1,5 @@ # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. import gc -from inspect import ArgSpec import logging import json import math