import types from typing import List, Optional import os import torch from torch import nn from utils.scheduler import SchedulerInterface, FlowMatchScheduler from wan_5b.modules.tokenizers import HuggingfaceTokenizer from wan_5b.modules.model import WanModel from wan_5b.modules.vae2_2 import _video_vae from wan_5b.modules.t5 import umt5_xxl from wan_5b.modules.causal_model import CausalWanModel class WanTextEncoder(torch.nn.Module): def __init__(self) -> None: super().__init__() self.text_encoder = umt5_xxl( encoder_only=True, return_tokenizer=False, dtype=torch.float32, device=torch.device('cpu') ).eval().requires_grad_(False) self.text_encoder.load_state_dict( torch.load("wan_models/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth", map_location='cpu', weights_only=False) ) # Move text encoder to GPU if available if torch.cuda.is_available(): self.text_encoder = self.text_encoder.cuda() self.tokenizer = HuggingfaceTokenizer( name="wan_models/Wan2.2-TI2V-5B/google/umt5-xxl/", seq_len=512, clean='whitespace') @property def device(self): # Assume we are always on GPU return torch.cuda.current_device() def forward(self, text_prompts: List[str]) -> dict: ids, mask = self.tokenizer( text_prompts, return_mask=True, add_special_tokens=True) ids = ids.to(self.device) mask = mask.to(self.device) seq_lens = mask.gt(0).sum(dim=1).long() context = self.text_encoder(ids, mask) for u, v in zip(context, seq_lens): u[v:] = 0.0 # set padding to 0.0 return { "prompt_embeds": context } class WanVAEWrapper(torch.nn.Module): def __init__(self): super().__init__() mean = [ -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667, ] std = [ 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744, ] self.mean = torch.tensor(mean, dtype=torch.float32) self.std = torch.tensor(std, dtype=torch.float32) # init model self.model = _video_vae( pretrained_path="wan_models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth", ).eval().requires_grad_(False) def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor: # pixel: [batch_size, num_channels, num_frames, height, width] device, dtype = pixel.device, pixel.dtype scale = [self.mean.to(device=device, dtype=dtype), 1.0 / self.std.to(device=device, dtype=dtype)] output = [ self.model.encode(u.unsqueeze(0), scale).float().squeeze(0) for u in pixel ] output = torch.stack(output, dim=0) # from [batch_size, num_channels, num_frames, height, width] # to [batch_size, num_frames, num_channels, height, width] output = output.permute(0, 2, 1, 3, 4) return output def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torch.Tensor: # from [batch_size, num_frames, num_channels, height, width] # to [batch_size, num_channels, num_frames, height, width] zs = latent.permute(0, 2, 1, 3, 4) if use_cache: assert latent.shape[0] == 1, "Batch size must be 1 when using cache" device, dtype = latent.device, latent.dtype scale = [self.mean.to(device=device, dtype=dtype), 1.0 / self.std.to(device=device, dtype=dtype)] if use_cache: decode_function = self.model.cached_decode else: decode_function = self.model.decode output = [] for u in zs: output.append(decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)) output = torch.stack(output, dim=0) # from [batch_size, num_channels, num_frames, height, width] # to [batch_size, num_frames, num_channels, height, width] output = output.permute(0, 2, 1, 3, 4) return output def decode_to_pixel_chunk(self, latent: torch.Tensor, use_cache: bool = False, chunk_size: int = 1) -> torch.Tensor: """ Decode latent frames to pixel space. Args: latent: Latent tensor with shape [batch_size, num_frames, num_channels, height, width] use_cache: Whether to use cached decoding (for streaming) chunk_size: Number of latent frames to decode at once (default 240 to avoid OOM) Returns: Decoded video tensor with shape [batch_size, num_frames, num_channels, height, width] """ # latent shape: [batch_size, num_frames, num_channels, height, width] # zs shape after permute: [batch_size, num_channels, num_frames, height, width] zs = latent.permute(0, 2, 1, 3, 4) if use_cache: assert latent.shape[0] == 1, "Batch size must be 1 when using cache" device, dtype = latent.device, latent.dtype scale = [self.mean.to(device=device, dtype=dtype), 1.0 / self.std.to(device=device, dtype=dtype)] if use_cache: decode_function = self.model.cached_decode else: decode_function = self.model.decode output = [] for u in zs: num_frames = u.shape[1] if num_frames <= chunk_size: # Decode short clips in one pass. if use_cache: # Start this segment from a clean cache. self.model.clear_cache() decoded = decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0) decoded = decoded.cpu() if use_cache: # Clear after this segment so it cannot affect the next video. self.model.clear_cache() else: # Decode longer clips in temporal chunks. decoded_chunks = [] if use_cache: # Clear once at the segment start; later chunks share the # internal cache. self.model.clear_cache() for start_idx in range(0, num_frames, chunk_size): end_idx = min(start_idx + chunk_size, num_frames) chunk = u[:, start_idx:end_idx, :, :] # [C, chunk_frames, H, W] decoded_chunk = decode_function(chunk.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0) decoded_chunks.append(decoded_chunk.cpu()) del decoded_chunk torch.cuda.empty_cache() decoded = torch.cat(decoded_chunks, dim=1) if use_cache: # Clear the cache after the full segment. self.model.clear_cache() output.append(decoded) output = torch.stack(output, dim=0) output = output.permute(0, 2, 1, 3, 4) return output class WanDiffusionWrapper(torch.nn.Module): def __init__( self, model_name="Wan2.2-TI2V-5B", timestep_shift=8.0, is_causal=False, local_attn_size=-1, sink_size=0, num_frame_per_block=1, t_scale=1.0, rope_method="linear", original_seq_len=None, ): super().__init__() if is_causal: self.model = CausalWanModel.from_pretrained( f"wan_models/{model_name}/", local_attn_size=local_attn_size, sink_size=sink_size, num_frame_per_block=num_frame_per_block) else: self.model = WanModel.from_pretrained(f"wan_models/{model_name}/") self.model.eval() self.model.t_scale = t_scale self.model.rope_method = rope_method self.model.original_seq_len = original_seq_len # For non-causal diffusion, all frames share the same timestep self.uniform_timestep = not is_causal self.scheduler = FlowMatchScheduler( shift=timestep_shift, sigma_min=0.0, extra_one_step=True ) self.scheduler.set_timesteps(1000, training=True) self.seq_len = 28160 # [1, 32, 48, 44, 80] self.post_init() self._compiled_model_call = None def enable_gradient_checkpointing(self) -> None: self.model.enable_gradient_checkpointing() def configure_torch_compile( self, *, backend: str = "inductor", mode: str | None = "max-autotune-no-cudagraphs", fullgraph: bool = False, dynamic: bool | None = False, options: dict | None = None, suppress_errors: bool = True, ) -> bool: from utils.torch_compile_utils import configure_module_call_torch_compile self._compiled_model_call = configure_module_call_torch_compile( self.model, name="WanDiffusionWrapper5B.model", backend=backend, mode=mode, fullgraph=fullgraph, dynamic=dynamic, options=options, suppress_errors=suppress_errors, ) return self._compiled_model_call is not None def _call_model(self, *args, **kwargs): # iter-39 v2: publish kv_cache scalars BEFORE entering the compiled # graph. The earlier version (iter-39 v1) published them inside # `_forward_inference`, but that function IS compiled, so each # `.item()` triggered a graph break. Moving the reads to this eager # wrapper keeps the dict lookups in the compiled attention forward # free of `.item()` syncs without adding any graph break. kv_cache = kwargs.get("kv_cache", None) if kv_cache is not None and len(kv_cache) > 0: try: from wan_5b.modules.causal_model import _CURRENT_GRID_META first_block_cache = kv_cache[0] _CURRENT_GRID_META["global_end_index"] = int( first_block_cache["global_end_index"].item() ) _CURRENT_GRID_META["local_end_index"] = int( first_block_cache["local_end_index"].item() ) _ps = first_block_cache.get("pinned_start", None) if _ps is not None and hasattr(_ps, "item"): _CURRENT_GRID_META["pinned_start"] = int(_ps.item()) _CURRENT_GRID_META["pinned_len"] = int( first_block_cache["pinned_len"].item() ) else: _CURRENT_GRID_META["pinned_start"] = -1 _CURRENT_GRID_META["pinned_len"] = 0 except (KeyError, AttributeError, ImportError): pass defer_kv_updates = ( os.environ.get("LLV2_DEFER_KV_UPDATES", "0") == "1" and kv_cache is not None ) if defer_kv_updates: kwargs["defer_cache_updates"] = True if self._compiled_model_call is not None: # iter-25: signal cudagraph allocator that a new "step" starts. # Required for mode=reduce-overhead when modules cache state # (KV cache rolling buffers, fp4-quant scale tensors) so the # cudagraph pool knows it can safely reuse step-N memory now # that step-(N+1) is starting. mark_step = getattr(torch.compiler, "cudagraph_mark_step_begin", None) if mark_step is not None: mark_step() result = self._compiled_model_call(*args, **kwargs) else: result = self.model(*args, **kwargs) if defer_kv_updates: if not isinstance(result, tuple) or len(result) != 2: raise RuntimeError( "LLV2_DEFER_KV_UPDATES expected model to return " "(output, cache_update_infos)." ) output, cache_update_infos = result if cache_update_infos: self.model._apply_cache_updates(kv_cache, cache_update_infos) return output return result def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: """ Convert flow matching's prediction to x0 prediction. flow_pred: the prediction with shape [B, C, H, W] xt: the input noisy data with shape [B, C, H, W] timestep: the timestep with shape [B] pred = noise - x0 x_t = (1-sigma_t) * x0 + sigma_t * noise we have x0 = x_t - sigma_t * pred see derivations https://chatgpt.com/share/67bf8589-3d04-8008-bc6e-4cf1a24e2d0e """ # use higher precision for calculations original_dtype = flow_pred.dtype flow_pred, xt, sigmas, timesteps = map( lambda x: x.double().to(flow_pred.device), [flow_pred, xt, self.scheduler.sigmas, self.scheduler.timesteps] ) timestep_id = torch.argmin( (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) x0_pred = xt - sigma_t * flow_pred return x0_pred.to(original_dtype) @staticmethod def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: """ Convert x0 prediction to flow matching's prediction. x0_pred: the x0 prediction with shape [B, C, H, W] xt: the input noisy data with shape [B, C, H, W] timestep: the timestep with shape [B] pred = (x_t - x_0) / sigma_t """ # use higher precision for calculations original_dtype = x0_pred.dtype x0_pred, xt, sigmas, timesteps = map( lambda x: x.double().to(x0_pred.device), [x0_pred, xt, scheduler.sigmas, scheduler.timesteps] ) timestep_id = torch.argmin( (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) flow_pred = (xt - x0_pred) / sigma_t return flow_pred.to(original_dtype) def forward( self, noisy_image_or_video: torch.Tensor, conditional_dict: dict, timestep: torch.Tensor, kv_cache: Optional[List[dict]] = None, crossattn_cache: Optional[List[dict]] = None, current_start: Optional[int] = None, classify_mode: Optional[bool] = False, concat_time_embeddings: Optional[bool] = False, clean_x: Optional[torch.Tensor] = None, aug_t: Optional[torch.Tensor] = None, cache_start: Optional[int] = None, rope_temporal_offset: Optional[torch.Tensor] = None, ) -> torch.Tensor: prompt_embeds = conditional_dict["prompt_embeds"] # [B, F] -> [B] if self.uniform_timestep: input_timestep = timestep[:, 0] else: input_timestep = timestep logits = None rope_offset_was_set = ( rope_temporal_offset is not None and hasattr(self.model, "rope_temporal_offset") ) if rope_offset_was_set: prev_rope_temporal_offset = self.model.rope_temporal_offset self.model.rope_temporal_offset = rope_temporal_offset # X0 prediction if kv_cache is not None: flow_pred = self._call_model( noisy_image_or_video.permute(0, 2, 1, 3, 4), t=input_timestep, context=prompt_embeds, seq_len=self.seq_len, kv_cache=kv_cache, crossattn_cache=crossattn_cache, current_start=current_start, cache_start=cache_start ).permute(0, 2, 1, 3, 4) else: if clean_x is not None: # teacher forcing flow_pred = self._call_model( noisy_image_or_video.permute(0, 2, 1, 3, 4), t=input_timestep, context=prompt_embeds, seq_len=self.seq_len, clean_x=clean_x.permute(0, 2, 1, 3, 4), aug_t=aug_t, ).permute(0, 2, 1, 3, 4) else: if classify_mode: flow_pred, logits = self._call_model( noisy_image_or_video.permute(0, 2, 1, 3, 4), t=input_timestep, context=prompt_embeds, seq_len=self.seq_len, classify_mode=True, register_tokens=self._register_tokens, cls_pred_branch=self._cls_pred_branch, gan_ca_blocks=self._gan_ca_blocks, concat_time_embeddings=concat_time_embeddings ) flow_pred = flow_pred.permute(0, 2, 1, 3, 4) else: flow_pred = self._call_model( noisy_image_or_video.permute(0, 2, 1, 3, 4), t=input_timestep, context=prompt_embeds, seq_len=self.seq_len ).permute(0, 2, 1, 3, 4) if rope_offset_was_set: self.model.rope_temporal_offset = prev_rope_temporal_offset pred_x0 = self._convert_flow_pred_to_x0( flow_pred=flow_pred.flatten(0, 1), xt=noisy_image_or_video.flatten(0, 1), timestep=timestep.flatten(0, 1) ).unflatten(0, flow_pred.shape[:2]) if logits is not None: return flow_pred, pred_x0, logits return flow_pred, pred_x0 def get_scheduler(self) -> SchedulerInterface: """ Update the current scheduler with the interface's static method """ scheduler = self.scheduler scheduler.convert_x0_to_noise = types.MethodType( SchedulerInterface.convert_x0_to_noise, scheduler) scheduler.convert_noise_to_x0 = types.MethodType( SchedulerInterface.convert_noise_to_x0, scheduler) scheduler.convert_velocity_to_x0 = types.MethodType( SchedulerInterface.convert_velocity_to_x0, scheduler) self.scheduler = scheduler return scheduler def post_init(self): """ A few custom initialization steps that should be called after the object is created. Currently, the only one we have is to bind a few methods to scheduler. We can gradually add more methods here if needed. """ self.get_scheduler() _MG_LIGHTVAE_DEFAULT_PATHS = { "mg_lightvae": os.path.join("wan_models", "Matrix-Game-3.0", "MG-LightVAE.pth"), "mg_lightvae_v2": os.path.join("wan_models", "Matrix-Game-3.0", "MG-LightVAE_v2.pth"), } def build_vae_5b(args): """Return the 5B VAE wrapper requested by args.vae_type.""" vae_type = str(getattr(args, "vae_type", "wan")).lower().strip() if vae_type in ("wan", "wan2.2", ""): return WanVAEWrapper() if vae_type in _MG_LIGHTVAE_DEFAULT_PATHS: from utils.lightvae_5b_wrapper import LightVAE5BWrapper return LightVAE5BWrapper(vae_path=_MG_LIGHTVAE_DEFAULT_PATHS[vae_type]) raise ValueError( f"Unknown vae_type '{vae_type}'. " "Expected one of: wan, mg_lightvae, mg_lightvae_v2." )