Spaces:
Running on Zero
Running on Zero
| import types | |
| from typing import List, Optional | |
| import torch | |
| from torch import nn | |
| from utils.scheduler import SchedulerInterface, FlowMatchScheduler | |
| from wan.modules.tokenizers import HuggingfaceTokenizer | |
| from wan.modules.vae import _video_vae | |
| from wan.modules.t5 import umt5_xxl | |
| from wan.modules.causal_model import CausalWanModel | |
| from utils.attention import validate_local_attn_size | |
| # ZeroGPU Space layout: base-model assets live under wan_models/Wan2.1-T2V-1.3B | |
| # (downloaded at startup from Wan-AI/Wan2.1-T2V-1.3B). | |
| WAN_DIR = "wan_models/Wan2.1-T2V-1.3B" | |
| 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(f"{WAN_DIR}/models_t5_umt5-xxl-enc-bf16.pth", | |
| map_location='cpu', weights_only=False) | |
| ) | |
| self.tokenizer = HuggingfaceTokenizer( | |
| name=f"{WAN_DIR}/google/umt5-xxl", seq_len=512, clean='whitespace') | |
| 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.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, | |
| 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 | |
| ] | |
| std = [ | |
| 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, | |
| 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 | |
| ] | |
| self.mean = torch.tensor(mean, dtype=torch.float32) | |
| self.std = torch.tensor(std, dtype=torch.float32) | |
| # init model | |
| self.model = _video_vae( | |
| pretrained_path=f"{WAN_DIR}/Wan2.1_VAE.pth", | |
| z_dim=16, | |
| ).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, | |
| output_device=None) -> 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: | |
| pixels = decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0) | |
| output.append(pixels.to(output_device) if output_device is not None else pixels) | |
| 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 | |
| class WanDiffusionWrapper(torch.nn.Module): | |
| def __init__( | |
| self, | |
| timestep_shift=5.0, | |
| is_causal=False, | |
| local_attn_size=21, | |
| sink_size=0, | |
| max_num_frames=21, | |
| frame_seq_length=1560, | |
| ): | |
| super().__init__() | |
| local_attn_size = validate_local_attn_size(local_attn_size) | |
| if not is_causal: | |
| raise ValueError("This Space uses the causal (AR) generator only") | |
| # Build the architecture directly (same params as | |
| # wan_models/Wan2.1-T2V-1.3B/config.json); the History-Forcing DMD | |
| # checkpoint loaded in app.py overwrites every generator weight, so the | |
| # 5.7 GB base safetensors download is skipped. | |
| self.model = CausalWanModel( | |
| model_type="t2v", | |
| patch_size=(1, 2, 2), | |
| text_len=512, | |
| in_dim=16, | |
| dim=1536, | |
| ffn_dim=8960, | |
| freq_dim=256, | |
| text_dim=4096, | |
| out_dim=16, | |
| num_heads=12, | |
| num_layers=30, | |
| local_attn_size=local_attn_size, | |
| sink_size=sink_size, | |
| ) | |
| self.model.eval() | |
| # 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) | |
| # One 480x832 latent frame has (60/2)*(104/2)=1560 patch tokens. | |
| # Set sequence capacity from the configured frame and token counts. | |
| self.max_num_frames = int(max_num_frames) | |
| self.frame_seq_length = int(frame_seq_length) | |
| if self.max_num_frames <= 0 or self.frame_seq_length <= 0: | |
| raise ValueError("max_num_frames and frame_seq_length must be positive") | |
| self.seq_len = self.max_num_frames * self.frame_seq_length | |
| self.post_init() | |
| def enable_gradient_checkpointing(self) -> None: | |
| self.model.enable_gradient_checkpointing() | |
| 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 | |
| """ | |
| # 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) | |
| 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, | |
| cache_start: Optional[int] = None, | |
| ) -> torch.Tensor: | |
| prompt_embeds = conditional_dict["prompt_embeds"] | |
| patch_h, patch_w = self.model.patch_size[1:] | |
| _, num_frames, _, height, width = noisy_image_or_video.shape | |
| actual_frame_seq = ((height + patch_h - 1) // patch_h) * ( | |
| (width + patch_w - 1) // patch_w | |
| ) | |
| required_seq_len = num_frames * actual_frame_seq | |
| if required_seq_len > self.seq_len: | |
| raise RuntimeError( | |
| "Input exceeds configured Wan sequence capacity: " | |
| f"{num_frames} frames * {actual_frame_seq} tokens/frame = " | |
| f"{required_seq_len} > {self.seq_len}. Set model_kwargs.max_num_frames " | |
| "and model_kwargs.frame_seq_length for this training resolution." | |
| ) | |
| # [B, F] -> [B] | |
| if self.uniform_timestep: | |
| input_timestep = timestep[:, 0] | |
| else: | |
| input_timestep = timestep | |
| # X0 prediction | |
| if kv_cache is not None: | |
| flow_pred = self.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: | |
| # diffusion forcing or bidirectional | |
| flow_pred = self.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) | |
| 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]) | |
| 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() | |