"""TAEHV-based Wan VAE wrapper used by offline inference with --use_taehv.""" from __future__ import annotations import logging import os import urllib.request from collections import namedtuple from pathlib import Path from types import SimpleNamespace import torch import torch.nn as nn import torch.nn.functional as F from models.model_interface import VAEInterface from models.wan.wan_wrapper import WanVAEWrapper DecoderResult = namedtuple("DecoderResult", ("frame", "memory")) TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index")) LOGGER = logging.getLogger(__name__) try: import tensorrt as trt except ImportError: trt = None def _resolve_project_root() -> Path: env_root = os.environ.get("STREAMDIFFUSIONV2_ROOT") if env_root: return Path(env_root).expanduser().resolve() repo_root = Path(__file__).resolve().parents[2] if (repo_root / "ckpts").exists(): return repo_root cwd = Path.cwd().resolve() if (cwd / "ckpts").exists(): return cwd return repo_root PROJECT_ROOT = _resolve_project_root() DEFAULT_TAEHV_CHECKPOINT = str(PROJECT_ROOT / "ckpts" / "taew2_1.pth") DEFAULT_TAEHV_URL = "https://github.com/madebyollin/taehv/raw/main/taew2_1.pth" def conv(n_in, n_out, **kwargs): return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) class Clamp(nn.Module): def forward(self, x): return torch.tanh(x / 3) * 3 class MemBlock(nn.Module): def __init__(self, n_in, n_out): super().__init__() self.conv = nn.Sequential( conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), ) self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() self.act = nn.ReLU(inplace=True) def forward(self, x, past): return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) class TPool(nn.Module): def __init__(self, n_f, stride): super().__init__() self.stride = stride self.conv = nn.Conv2d(n_f * stride, n_f, 1, bias=False) def forward(self, x): _nt, c, h, w = x.shape return self.conv(x.reshape(-1, self.stride * c, h, w)) class TGrow(nn.Module): def __init__(self, n_f, stride): super().__init__() self.stride = stride self.conv = nn.Conv2d(n_f, n_f * stride, 1, bias=False) def forward(self, x): _nt, c, h, w = x.shape x = self.conv(x) return x.reshape(-1, c, h, w) def apply_model_with_memblocks(model, x, parallel, show_progress_bar): assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor" n, t, c, h, w = x.shape if parallel: x = x.reshape(n * t, c, h, w) for block in model: if isinstance(block, MemBlock): nt, c, h, w = x.shape t = nt // n _x = x.reshape(n, t, c, h, w) mem = F.pad(_x, (0, 0, 0, 0, 0, 0, 1, 0), value=0)[:, :t].reshape(x.shape) x = block(x, mem) else: x = block(x) nt, c, h, w = x.shape t = nt // n x = x.view(n, t, c, h, w) else: out = [] work_queue = [TWorkItem(xt, 0) for xt in x.reshape(n, t * c, h, w).chunk(t, dim=1)] mem = [None] * len(model) while work_queue: xt, block_index = work_queue.pop(0) if block_index == len(model): out.append(xt) continue block = model[block_index] if isinstance(block, MemBlock): if mem[block_index] is None: xt_new = block(xt, xt * 0) mem[block_index] = xt else: xt_new = block(xt, mem[block_index]) mem[block_index].copy_(xt) work_queue.insert(0, TWorkItem(xt_new, block_index + 1)) elif isinstance(block, TPool): if mem[block_index] is None: mem[block_index] = [] mem[block_index].append(xt) if len(mem[block_index]) == block.stride: n, c, h, w = xt.shape xt = block(torch.cat(mem[block_index], 1).view(n * block.stride, c, h, w)) mem[block_index] = [] work_queue.insert(0, TWorkItem(xt, block_index + 1)) elif isinstance(block, TGrow): xt = block(xt) n_out, c_out, h_out, w_out = xt.shape batch_size = n_out // block.stride grown = xt.view(batch_size, block.stride * c_out, h_out, w_out) for xt_next in reversed(grown.chunk(block.stride, dim=1)): work_queue.insert(0, TWorkItem(xt_next, block_index + 1)) else: xt = block(xt) work_queue.insert(0, TWorkItem(xt, block_index + 1)) x = torch.stack(out, 1) return x class TAEHV(nn.Module): latent_channels = 16 image_channels = 3 def __init__(self, checkpoint_path=DEFAULT_TAEHV_CHECKPOINT, decoder_time_upscale=(True, True), decoder_space_upscale=(True, True, True)): super().__init__() self.encoder = nn.Sequential( conv(TAEHV.image_channels, 64), nn.ReLU(inplace=True), TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), TPool(64, 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), conv(64, TAEHV.latent_channels), ) n_f = [256, 128, 64, 64] self.frames_to_trim = 2 ** sum(decoder_time_upscale) - 1 self.decoder = nn.Sequential( Clamp(), conv(TAEHV.latent_channels, n_f[0]), nn.ReLU(inplace=True), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), conv(n_f[2], n_f[3], bias=False), nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels), ) self.load_state_dict(self.patch_tgrow_layers(torch.load(checkpoint_path, map_location="cpu", weights_only=True))) def patch_tgrow_layers(self, state_dict): new_state_dict = self.state_dict() for index, layer in enumerate(self.decoder): if isinstance(layer, TGrow): key = f"decoder.{index}.conv.weight" if state_dict[key].shape[0] > new_state_dict[key].shape[0]: state_dict[key] = state_dict[key][-new_state_dict[key].shape[0]:] return state_dict def encode_video(self, x, parallel=True, show_progress_bar=False): return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar) def decode_video(self, x, parallel=True, show_progress_bar=False): return apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar) class TAEHVParallelDecoderModule(nn.Module): """Shape-specialized decoder graph that is friendly to ONNX/TensorRT export.""" def __init__(self, decoder: nn.Module): super().__init__() self.decoder = decoder def forward(self, latent: torch.Tensor) -> torch.Tensor: return apply_model_with_memblocks( self.decoder, latent, parallel=True, show_progress_bar=False, ) class TAEHVTensorRTDecoder: """Cache TensorRT engines for the common fixed TAEHV decoder shapes.""" def __init__( self, decoder_module: nn.Module, cache_dir: str | Path, workspace_bytes: int, ) -> None: if trt is None: raise ImportError("TensorRT is not installed, but TensorRT decode was requested.") self.decoder_module = decoder_module.eval() self.cache_dir = Path(cache_dir) self.cache_dir.mkdir(parents=True, exist_ok=True) self.workspace_bytes = int(workspace_bytes) self.logger = trt.Logger(trt.Logger.WARNING) self.runtime = trt.Runtime(self.logger) self._engine_cache: dict[tuple[int, ...], tuple[object, object, str, str]] = {} self._stream_cache: dict[int, torch.cuda.Stream] = {} def _shape_tag(self, shape: tuple[int, ...]) -> str: return "x".join(str(dim) for dim in shape) def _engine_path(self, shape: tuple[int, ...]) -> Path: return self.cache_dir / f"taehv_decoder_trt_{self._shape_tag(shape)}.plan" def _onnx_path(self, shape: tuple[int, ...]) -> Path: return self.cache_dir / f"taehv_decoder_trt_{self._shape_tag(shape)}.onnx" def _build_engine(self, shape: tuple[int, ...], device: torch.device) -> bytes: onnx_path = self._onnx_path(shape) engine_path = self._engine_path(shape) sample = torch.randn(shape, device=device, dtype=torch.float16) with torch.no_grad(): torch.onnx.export( self.decoder_module, (sample,), onnx_path.as_posix(), input_names=["latent"], output_names=["video"], opset_version=18, do_constant_folding=True, ) builder = trt.Builder(self.logger) network = builder.create_network( 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) parser = trt.OnnxParser(network, self.logger) with open(onnx_path, "rb") as handle: if not parser.parse(handle.read()): errors = "\n".join(str(parser.get_error(i)) for i in range(parser.num_errors)) raise RuntimeError(f"Failed to parse ONNX for TAEHV TensorRT engine:\n{errors}") config = builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, self.workspace_bytes) config.set_flag(trt.BuilderFlag.FP16) serialized = builder.build_serialized_network(network, config) if serialized is None: raise RuntimeError("TensorRT failed to build the TAEHV decoder engine") engine_bytes = bytes(serialized) engine_path.write_bytes(engine_bytes) return engine_bytes def _load_engine(self, shape: tuple[int, ...], device: torch.device): cached = self._engine_cache.get(shape) if cached is not None: return cached engine_path = self._engine_path(shape) if engine_path.exists(): engine_bytes = engine_path.read_bytes() else: LOGGER.info("Building TensorRT engine for TAEHV decoder shape %s", shape) engine_bytes = self._build_engine(shape, device) engine = self.runtime.deserialize_cuda_engine(engine_bytes) if engine is None: raise RuntimeError(f"Failed to deserialize TensorRT engine for TAEHV shape {shape}") context = engine.create_execution_context() input_name = "" output_name = "" for index in range(engine.num_io_tensors): name = engine.get_tensor_name(index) mode = engine.get_tensor_mode(name) if mode == trt.TensorIOMode.INPUT: input_name = name elif mode == trt.TensorIOMode.OUTPUT: output_name = name if not input_name or not output_name: raise RuntimeError("Failed to discover TensorRT decoder tensor names") cached = (engine, context, input_name, output_name) self._engine_cache[shape] = cached return cached def decode(self, latent: torch.Tensor) -> torch.Tensor: latent = latent.contiguous().to(dtype=torch.float16) shape = tuple(int(dim) for dim in latent.shape) _engine, context, input_name, output_name = self._load_engine(shape, latent.device) if not context.set_input_shape(input_name, shape): raise RuntimeError(f"TensorRT rejected decoder input shape {shape}") output_shape = tuple(int(dim) for dim in context.get_tensor_shape(output_name)) output = torch.empty(output_shape, device=latent.device, dtype=latent.dtype) context.set_tensor_address(input_name, int(latent.data_ptr())) context.set_tensor_address(output_name, int(output.data_ptr())) device_index = latent.device.index if device_index is None: raise RuntimeError("TensorRT decode requires an explicit CUDA device index") stream = self._stream_cache.get(device_index) if stream is None: stream = torch.cuda.Stream(device=latent.device) self._stream_cache[device_index] = stream current_stream = torch.cuda.current_stream(device=latent.device) stream.wait_stream(current_stream) ok = context.execute_async_v3(stream.cuda_stream) if not ok: raise RuntimeError("TensorRT execution failed for the TAEHV decoder") current_stream.wait_stream(stream) return output class TAEHVWanVAEWrapper(VAEInterface): """Wan stream encoder with a TAEHV decoder for faster pixel reconstruction.""" def __init__( self, model_type: str = "T2V-1.3B", checkpoint_path: str | None = None, auto_download: bool = True, parallel_decode: bool = False, use_tensorrt: bool = False, tensorrt_cache_dir: str | None = None, tensorrt_workspace_bytes: int = 4 << 30, ): super().__init__() self.checkpoint_path = checkpoint_path or DEFAULT_TAEHV_CHECKPOINT self.use_tensorrt = use_tensorrt self.parallel_decode = parallel_decode or use_tensorrt self.model = SimpleNamespace(first_encode=True, first_decode=True) self.decode_context_latents = 3 self._decode_latent_cache: torch.Tensor | None = None self.encoder_vae = WanVAEWrapper(model_type=model_type) self.taehv = TAEHV(checkpoint_path=self._resolve_checkpoint(auto_download)) self._tensorrt_cache_dir = tensorrt_cache_dir or (PROJECT_ROOT / "ckpts" / "taehv_trt") self._tensorrt_workspace_bytes = int(tensorrt_workspace_bytes) self._tensorrt_decoder: TAEHVTensorRTDecoder | None = None self._tensorrt_failed = False def _resolve_checkpoint(self, auto_download: bool) -> str: if os.path.exists(self.checkpoint_path): return self.checkpoint_path if not auto_download: raise FileNotFoundError(f"TAEHV checkpoint not found: {self.checkpoint_path}") os.makedirs(os.path.dirname(self.checkpoint_path), exist_ok=True) urllib.request.urlretrieve(DEFAULT_TAEHV_URL, self.checkpoint_path) return self.checkpoint_path def to(self, *args, **kwargs): device = kwargs.get("device") dtype = kwargs.get("dtype") if args: if len(args) >= 1 and not isinstance(args[0], torch.dtype): device = args[0] if len(args) >= 2 and isinstance(args[1], torch.dtype): dtype = args[1] elif len(args) == 1 and isinstance(args[0], torch.dtype): dtype = args[0] encoder_kwargs = {} if device is not None: encoder_kwargs["device"] = device if dtype is not None: encoder_kwargs["dtype"] = dtype self.encoder_vae.to(**encoder_kwargs) taehv_kwargs = {"dtype": torch.float16} if device is not None: taehv_kwargs["device"] = device self.taehv.to(**taehv_kwargs) if self.use_tensorrt and not self._tensorrt_failed: self._tensorrt_decoder = TAEHVTensorRTDecoder( decoder_module=TAEHVParallelDecoderModule(self.taehv.decoder).to(**taehv_kwargs).eval(), cache_dir=self._tensorrt_cache_dir, workspace_bytes=self._tensorrt_workspace_bytes, ) return self def _pixels_to_unit_range(self, video: torch.Tensor) -> torch.Tensor: return (video * 0.5 + 0.5).clamp(0, 1) def _unit_range_to_pixels(self, video: torch.Tensor) -> torch.Tensor: return video.mul(2).sub(1).clamp(-1, 1) def _to_ntchw(self, video: torch.Tensor) -> torch.Tensor: return video.permute(0, 2, 1, 3, 4).contiguous() def _to_ncthw(self, video: torch.Tensor) -> torch.Tensor: return video.permute(0, 2, 1, 3, 4).contiguous() def _decode_video(self, latent: torch.Tensor) -> torch.Tensor: if self.use_tensorrt and self._tensorrt_decoder is not None and not self._tensorrt_failed: try: return self._tensorrt_decoder.decode(latent) except Exception as exc: self._tensorrt_failed = True self._tensorrt_decoder = None LOGGER.warning( "TAEHV TensorRT decode failed for shape %s, falling back to PyTorch parallel decode: %s", tuple(latent.shape), exc, ) return self.taehv.decode_video( latent, parallel=self.parallel_decode, show_progress_bar=False, ) def decode_to_pixel(self, latent: torch.Tensor) -> torch.Tensor: video = self._decode_video(latent) return self._unit_range_to_pixels(video) def decode(self, latent: torch.Tensor) -> torch.Tensor: return self.decode_to_pixel(latent) def stream_encode(self, video: torch.Tensor, is_scale=False) -> torch.Tensor: self.encoder_vae.model.first_encode = self.model.first_encode latent = self.encoder_vae.stream_encode(video, is_scale=is_scale) self.model.first_encode = self.encoder_vae.model.first_encode return latent def stream_decode_to_pixel(self, latent: torch.Tensor) -> torch.Tensor: model_dtype = next(self.taehv.parameters()).dtype latent = latent.to(dtype=model_dtype) # Match Self-Forcing's TAEHV usage: keep a short latent prefix so each # incremental decode sees recent temporal context, then trim the # already-emitted frames from the pixel output. if self.model.first_decode: self.model.first_decode = False self._decode_latent_cache = None decode_latent = latent emitted_latents = max(latent.shape[1] - 1, 0) else: context = self._decode_latent_cache decode_latent = latent if context is None else torch.cat([context.to(device=latent.device, dtype=model_dtype), latent], dim=1) emitted_latents = latent.shape[1] self._decode_latent_cache = decode_latent[:, -min(self.decode_context_latents, decode_latent.shape[1]):].detach().clone() video = self._decode_video(decode_latent) if emitted_latents: video = video[:, -emitted_latents * 4:, :, :, :] else: video = video[:, 0:0, :, :, :] return self._unit_range_to_pixels(video).to(dtype=latent.dtype)