Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. | |
| import gc | |
| import logging | |
| import math | |
| import os | |
| import random | |
| import sys | |
| import types | |
| from contextlib import contextmanager | |
| from functools import partial | |
| import torch | |
| import torch.cuda.amp as amp | |
| import torch.distributed as dist | |
| import torchvision.transforms.functional as TF | |
| from PIL import Image | |
| from tqdm import tqdm | |
| from .distributed.fsdp import shard_model | |
| from .distributed.sequence_parallel import sp_attn_forward, sp_dit_forward | |
| from .distributed.util import get_world_size | |
| from .modules.model import WanModel | |
| from .modules.t5 import T5EncoderModel | |
| from .modules.vae2_2 import Wan2_2_VAE | |
| from .utils.fm_solvers import ( | |
| FlowDPMSolverMultistepScheduler, | |
| get_sampling_sigmas, | |
| retrieve_timesteps, | |
| ) | |
| from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler | |
| from .utils.utils import best_output_size, masks_like | |
| class WanTI2V: | |
| def __init__( | |
| self, | |
| config, | |
| checkpoint_dir, | |
| device_id=0, | |
| rank=0, | |
| t5_fsdp=False, | |
| dit_fsdp=False, | |
| use_sp=False, | |
| t5_cpu=False, | |
| init_on_cpu=True, | |
| convert_model_dtype=False, | |
| ): | |
| r""" | |
| Initializes the Wan text-to-video generation model components. | |
| Args: | |
| config (EasyDict): | |
| Object containing model parameters initialized from config.py | |
| checkpoint_dir (`str`): | |
| Path to directory containing model checkpoints | |
| device_id (`int`, *optional*, defaults to 0): | |
| Id of target GPU device | |
| rank (`int`, *optional*, defaults to 0): | |
| Process rank for distributed training | |
| t5_fsdp (`bool`, *optional*, defaults to False): | |
| Enable FSDP sharding for T5 model | |
| dit_fsdp (`bool`, *optional*, defaults to False): | |
| Enable FSDP sharding for DiT model | |
| use_sp (`bool`, *optional*, defaults to False): | |
| Enable distribution strategy of sequence parallel. | |
| t5_cpu (`bool`, *optional*, defaults to False): | |
| Whether to place T5 model on CPU. Only works without t5_fsdp. | |
| init_on_cpu (`bool`, *optional*, defaults to True): | |
| Enable initializing Transformer Model on CPU. Only works without FSDP or USP. | |
| convert_model_dtype (`bool`, *optional*, defaults to False): | |
| Convert DiT model parameters dtype to 'config.param_dtype'. | |
| Only works without FSDP. | |
| """ | |
| self.device = torch.device(f"cuda:{device_id}") | |
| self.config = config | |
| self.rank = rank | |
| self.t5_cpu = t5_cpu | |
| self.init_on_cpu = init_on_cpu | |
| self.num_train_timesteps = config.num_train_timesteps | |
| self.param_dtype = config.param_dtype | |
| if t5_fsdp or dit_fsdp or use_sp: | |
| self.init_on_cpu = False | |
| shard_fn = partial(shard_model, device_id=device_id) | |
| self.text_encoder = T5EncoderModel( | |
| text_len=config.text_len, | |
| dtype=config.t5_dtype, | |
| device=torch.device('cpu'), | |
| checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint), | |
| tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), | |
| shard_fn=shard_fn if t5_fsdp else None) | |
| self.vae_stride = config.vae_stride | |
| self.patch_size = config.patch_size | |
| self.vae = Wan2_2_VAE( | |
| vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint), | |
| device=self.device) | |
| logging.info(f"Creating WanModel from {checkpoint_dir}") | |
| self.model = WanModel.from_pretrained(checkpoint_dir) | |
| self.model = self._configure_model( | |
| model=self.model, | |
| use_sp=use_sp, | |
| dit_fsdp=dit_fsdp, | |
| shard_fn=shard_fn, | |
| convert_model_dtype=convert_model_dtype) | |
| if use_sp: | |
| self.sp_size = get_world_size() | |
| else: | |
| self.sp_size = 1 | |
| self.sample_neg_prompt = config.sample_neg_prompt | |
| def _configure_model(self, model, use_sp, dit_fsdp, shard_fn, | |
| convert_model_dtype): | |
| """ | |
| Configures a model object. This includes setting evaluation modes, | |
| applying distributed parallel strategy, and handling device placement. | |
| Args: | |
| model (torch.nn.Module): | |
| The model instance to configure. | |
| use_sp (`bool`): | |
| Enable distribution strategy of sequence parallel. | |
| dit_fsdp (`bool`): | |
| Enable FSDP sharding for DiT model. | |
| shard_fn (callable): | |
| The function to apply FSDP sharding. | |
| convert_model_dtype (`bool`): | |
| Convert DiT model parameters dtype to 'config.param_dtype'. | |
| Only works without FSDP. | |
| Returns: | |
| torch.nn.Module: | |
| The configured model. | |
| """ | |
| model.eval().requires_grad_(False) | |
| if use_sp: | |
| for block in model.blocks: | |
| block.self_attn.forward = types.MethodType( | |
| sp_attn_forward, block.self_attn) | |
| model.forward = types.MethodType(sp_dit_forward, model) | |
| if dist.is_initialized(): | |
| dist.barrier() | |
| if dit_fsdp: | |
| model = shard_fn(model) | |
| else: | |
| if convert_model_dtype: | |
| model.to(self.param_dtype) | |
| if not self.init_on_cpu: | |
| model.to(self.device) | |
| return model | |
| def generate(self, | |
| input_prompt, | |
| img=None, | |
| size=(1280, 704), | |
| max_area=704 * 1280, | |
| frame_num=81, | |
| shift=5.0, | |
| sample_solver='unipc', | |
| sampling_steps=50, | |
| guide_scale=5.0, | |
| n_prompt="", | |
| seed=-1, | |
| offload_model=True): | |
| r""" | |
| Generates video frames from text prompt using diffusion process. | |
| Args: | |
| input_prompt (`str`): | |
| Text prompt for content generation | |
| img (PIL.Image.Image): | |
| Input image tensor. Shape: [3, H, W] | |
| size (`tuple[int]`, *optional*, defaults to (1280,704)): | |
| Controls video resolution, (width,height). | |
| max_area (`int`, *optional*, defaults to 704*1280): | |
| Maximum pixel area for latent space calculation. Controls video resolution scaling | |
| frame_num (`int`, *optional*, defaults to 81): | |
| How many frames to sample from a video. The number should be 4n+1 | |
| shift (`float`, *optional*, defaults to 5.0): | |
| Noise schedule shift parameter. Affects temporal dynamics | |
| sample_solver (`str`, *optional*, defaults to 'unipc'): | |
| Solver used to sample the video. | |
| sampling_steps (`int`, *optional*, defaults to 50): | |
| Number of diffusion sampling steps. Higher values improve quality but slow generation | |
| guide_scale (`float`, *optional*, defaults 5.0): | |
| Classifier-free guidance scale. Controls prompt adherence vs. creativity. | |
| n_prompt (`str`, *optional*, defaults to ""): | |
| Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` | |
| seed (`int`, *optional*, defaults to -1): | |
| Random seed for noise generation. If -1, use random seed. | |
| offload_model (`bool`, *optional*, defaults to True): | |
| If True, offloads models to CPU during generation to save VRAM | |
| Returns: | |
| torch.Tensor: | |
| Generated video frames tensor. Dimensions: (C, N H, W) where: | |
| - C: Color channels (3 for RGB) | |
| - N: Number of frames (81) | |
| - H: Frame height (from size) | |
| - W: Frame width from size) | |
| """ | |
| # i2v | |
| if img is not None: | |
| return self.i2v( | |
| input_prompt=input_prompt, | |
| img=img, | |
| max_area=max_area, | |
| frame_num=frame_num, | |
| shift=shift, | |
| sample_solver=sample_solver, | |
| sampling_steps=sampling_steps, | |
| guide_scale=guide_scale, | |
| n_prompt=n_prompt, | |
| seed=seed, | |
| offload_model=offload_model) | |
| # t2v | |
| return self.t2v( | |
| input_prompt=input_prompt, | |
| size=size, | |
| frame_num=frame_num, | |
| shift=shift, | |
| sample_solver=sample_solver, | |
| sampling_steps=sampling_steps, | |
| guide_scale=guide_scale, | |
| n_prompt=n_prompt, | |
| seed=seed, | |
| offload_model=offload_model) | |
| def t2v(self, | |
| input_prompt, | |
| size=(1280, 704), | |
| frame_num=121, | |
| shift=5.0, | |
| sample_solver='unipc', | |
| sampling_steps=50, | |
| guide_scale=5.0, | |
| n_prompt="", | |
| seed=-1, | |
| offload_model=True): | |
| r""" | |
| Generates video frames from text prompt using diffusion process. | |
| Args: | |
| input_prompt (`str`): | |
| Text prompt for content generation | |
| size (`tuple[int]`, *optional*, defaults to (1280,704)): | |
| Controls video resolution, (width,height). | |
| frame_num (`int`, *optional*, defaults to 121): | |
| How many frames to sample from a video. The number should be 4n+1 | |
| shift (`float`, *optional*, defaults to 5.0): | |
| Noise schedule shift parameter. Affects temporal dynamics | |
| sample_solver (`str`, *optional*, defaults to 'unipc'): | |
| Solver used to sample the video. | |
| sampling_steps (`int`, *optional*, defaults to 50): | |
| Number of diffusion sampling steps. Higher values improve quality but slow generation | |
| guide_scale (`float`, *optional*, defaults 5.0): | |
| Classifier-free guidance scale. Controls prompt adherence vs. creativity. | |
| n_prompt (`str`, *optional*, defaults to ""): | |
| Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` | |
| seed (`int`, *optional*, defaults to -1): | |
| Random seed for noise generation. If -1, use random seed. | |
| offload_model (`bool`, *optional*, defaults to True): | |
| If True, offloads models to CPU during generation to save VRAM | |
| Returns: | |
| torch.Tensor: | |
| Generated video frames tensor. Dimensions: (C, N H, W) where: | |
| - C: Color channels (3 for RGB) | |
| - N: Number of frames (81) | |
| - H: Frame height (from size) | |
| - W: Frame width from size) | |
| """ | |
| # preprocess | |
| F = frame_num | |
| target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1, | |
| size[1] // self.vae_stride[1], | |
| size[0] // self.vae_stride[2]) | |
| seq_len = math.ceil((target_shape[2] * target_shape[3]) / | |
| (self.patch_size[1] * self.patch_size[2]) * | |
| target_shape[1] / self.sp_size) * self.sp_size | |
| if n_prompt == "": | |
| n_prompt = self.sample_neg_prompt | |
| seed = seed if seed >= 0 else random.randint(0, sys.maxsize) | |
| seed_g = torch.Generator(device=self.device) | |
| seed_g.manual_seed(seed) | |
| if not self.t5_cpu: | |
| self.text_encoder.model.to(self.device) | |
| context = self.text_encoder([input_prompt], self.device) | |
| context_null = self.text_encoder([n_prompt], self.device) | |
| if offload_model: | |
| self.text_encoder.model.cpu() | |
| else: | |
| context = self.text_encoder([input_prompt], torch.device('cpu')) | |
| context_null = self.text_encoder([n_prompt], torch.device('cpu')) | |
| context = [t.to(self.device) for t in context] | |
| context_null = [t.to(self.device) for t in context_null] | |
| noise = [ | |
| torch.randn( | |
| target_shape[0], | |
| target_shape[1], | |
| target_shape[2], | |
| target_shape[3], | |
| dtype=torch.float32, | |
| device=self.device, | |
| generator=seed_g) | |
| ] | |
| def noop_no_sync(): | |
| yield | |
| no_sync = getattr(self.model, 'no_sync', noop_no_sync) | |
| # evaluation mode | |
| with ( | |
| torch.amp.autocast('cuda', dtype=self.param_dtype), | |
| torch.no_grad(), | |
| no_sync(), | |
| ): | |
| if sample_solver == 'unipc': | |
| sample_scheduler = FlowUniPCMultistepScheduler( | |
| num_train_timesteps=self.num_train_timesteps, | |
| shift=1, | |
| use_dynamic_shifting=False) | |
| sample_scheduler.set_timesteps( | |
| sampling_steps, device=self.device, shift=shift) | |
| timesteps = sample_scheduler.timesteps | |
| elif sample_solver == 'dpm++': | |
| sample_scheduler = FlowDPMSolverMultistepScheduler( | |
| num_train_timesteps=self.num_train_timesteps, | |
| shift=1, | |
| use_dynamic_shifting=False) | |
| sampling_sigmas = get_sampling_sigmas(sampling_steps, shift) | |
| timesteps, _ = retrieve_timesteps( | |
| sample_scheduler, | |
| device=self.device, | |
| sigmas=sampling_sigmas) | |
| else: | |
| raise NotImplementedError("Unsupported solver.") | |
| # sample videos | |
| latents = noise | |
| mask1, mask2 = masks_like(noise, zero=False) | |
| arg_c = {'context': context, 'seq_len': seq_len} | |
| arg_null = {'context': context_null, 'seq_len': seq_len} | |
| if offload_model or self.init_on_cpu: | |
| self.model.to(self.device) | |
| torch.cuda.empty_cache() | |
| for _, t in enumerate(tqdm(timesteps)): | |
| latent_model_input = latents | |
| timestep = [t] | |
| timestep = torch.stack(timestep) | |
| temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten() | |
| temp_ts = torch.cat([ | |
| temp_ts, | |
| temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep | |
| ]) | |
| timestep = temp_ts.unsqueeze(0) | |
| noise_pred_cond = self.model( | |
| latent_model_input, t=timestep, **arg_c)[0] | |
| noise_pred_uncond = self.model( | |
| latent_model_input, t=timestep, **arg_null)[0] | |
| noise_pred = noise_pred_uncond + guide_scale * ( | |
| noise_pred_cond - noise_pred_uncond) | |
| temp_x0 = sample_scheduler.step( | |
| noise_pred.unsqueeze(0), | |
| t, | |
| latents[0].unsqueeze(0), | |
| return_dict=False, | |
| generator=seed_g)[0] | |
| latents = [temp_x0.squeeze(0)] | |
| x0 = latents | |
| if offload_model: | |
| self.model.cpu() | |
| torch.cuda.synchronize() | |
| torch.cuda.empty_cache() | |
| if self.rank == 0: | |
| videos = self.vae.decode(x0) | |
| del noise, latents | |
| del sample_scheduler | |
| if offload_model: | |
| gc.collect() | |
| torch.cuda.synchronize() | |
| if dist.is_initialized(): | |
| dist.barrier() | |
| return videos[0] if self.rank == 0 else None | |
| def i2v(self, | |
| input_prompt, | |
| img, | |
| max_area=704 * 1280, | |
| frame_num=121, | |
| shift=5.0, | |
| sample_solver='unipc', | |
| sampling_steps=40, | |
| guide_scale=5.0, | |
| n_prompt="", | |
| seed=-1, | |
| offload_model=True): | |
| r""" | |
| Generates video frames from input image and text prompt using diffusion process. | |
| Args: | |
| input_prompt (`str`): | |
| Text prompt for content generation. | |
| img (PIL.Image.Image): | |
| Input image tensor. Shape: [3, H, W] | |
| max_area (`int`, *optional*, defaults to 704*1280): | |
| Maximum pixel area for latent space calculation. Controls video resolution scaling | |
| frame_num (`int`, *optional*, defaults to 121): | |
| How many frames to sample from a video. The number should be 4n+1 | |
| shift (`float`, *optional*, defaults to 5.0): | |
| Noise schedule shift parameter. Affects temporal dynamics | |
| [NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0. | |
| sample_solver (`str`, *optional*, defaults to 'unipc'): | |
| Solver used to sample the video. | |
| sampling_steps (`int`, *optional*, defaults to 40): | |
| Number of diffusion sampling steps. Higher values improve quality but slow generation | |
| guide_scale (`float`, *optional*, defaults 5.0): | |
| Classifier-free guidance scale. Controls prompt adherence vs. creativity. | |
| n_prompt (`str`, *optional*, defaults to ""): | |
| Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` | |
| seed (`int`, *optional*, defaults to -1): | |
| Random seed for noise generation. If -1, use random seed | |
| offload_model (`bool`, *optional*, defaults to True): | |
| If True, offloads models to CPU during generation to save VRAM | |
| Returns: | |
| torch.Tensor: | |
| Generated video frames tensor. Dimensions: (C, N H, W) where: | |
| - C: Color channels (3 for RGB) | |
| - N: Number of frames (121) | |
| - H: Frame height (from max_area) | |
| - W: Frame width (from max_area) | |
| """ | |
| # preprocess | |
| ih, iw = img.height, img.width | |
| dh, dw = self.patch_size[1] * self.vae_stride[1], self.patch_size[ | |
| 2] * self.vae_stride[2] | |
| ow, oh = best_output_size(iw, ih, dw, dh, max_area) | |
| scale = max(ow / iw, oh / ih) | |
| img = img.resize((round(iw * scale), round(ih * scale)), Image.LANCZOS) | |
| # center-crop | |
| x1 = (img.width - ow) // 2 | |
| y1 = (img.height - oh) // 2 | |
| img = img.crop((x1, y1, x1 + ow, y1 + oh)) | |
| assert img.width == ow and img.height == oh | |
| # to tensor | |
| img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device).unsqueeze(1) | |
| F = frame_num | |
| seq_len = ((F - 1) // self.vae_stride[0] + 1) * ( | |
| oh // self.vae_stride[1]) * (ow // self.vae_stride[2]) // ( | |
| self.patch_size[1] * self.patch_size[2]) | |
| seq_len = int(math.ceil(seq_len / self.sp_size)) * self.sp_size | |
| seed = seed if seed >= 0 else random.randint(0, sys.maxsize) | |
| seed_g = torch.Generator(device=self.device) | |
| seed_g.manual_seed(seed) | |
| noise = torch.randn( | |
| self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1, | |
| oh // self.vae_stride[1], | |
| ow // self.vae_stride[2], | |
| dtype=torch.float32, | |
| generator=seed_g, | |
| device=self.device) | |
| if n_prompt == "": | |
| n_prompt = self.sample_neg_prompt | |
| # preprocess | |
| if not self.t5_cpu: | |
| self.text_encoder.model.to(self.device) | |
| context = self.text_encoder([input_prompt], self.device) | |
| context_null = self.text_encoder([n_prompt], self.device) | |
| if offload_model: | |
| self.text_encoder.model.cpu() | |
| else: | |
| context = self.text_encoder([input_prompt], torch.device('cpu')) | |
| context_null = self.text_encoder([n_prompt], torch.device('cpu')) | |
| context = [t.to(self.device) for t in context] | |
| context_null = [t.to(self.device) for t in context_null] | |
| z = self.vae.encode([img]) | |
| def noop_no_sync(): | |
| yield | |
| no_sync = getattr(self.model, 'no_sync', noop_no_sync) | |
| # evaluation mode | |
| with ( | |
| torch.amp.autocast('cuda', dtype=self.param_dtype), | |
| torch.no_grad(), | |
| no_sync(), | |
| ): | |
| if sample_solver == 'unipc': | |
| sample_scheduler = FlowUniPCMultistepScheduler( | |
| num_train_timesteps=self.num_train_timesteps, | |
| shift=1, | |
| use_dynamic_shifting=False) | |
| sample_scheduler.set_timesteps( | |
| sampling_steps, device=self.device, shift=shift) | |
| timesteps = sample_scheduler.timesteps | |
| elif sample_solver == 'dpm++': | |
| sample_scheduler = FlowDPMSolverMultistepScheduler( | |
| num_train_timesteps=self.num_train_timesteps, | |
| shift=1, | |
| use_dynamic_shifting=False) | |
| sampling_sigmas = get_sampling_sigmas(sampling_steps, shift) | |
| timesteps, _ = retrieve_timesteps( | |
| sample_scheduler, | |
| device=self.device, | |
| sigmas=sampling_sigmas) | |
| else: | |
| raise NotImplementedError("Unsupported solver.") | |
| # sample videos | |
| latent = noise | |
| mask1, mask2 = masks_like([noise], zero=True) | |
| latent = (1. - mask2[0]) * z[0] + mask2[0] * latent | |
| arg_c = { | |
| 'context': [context[0]], | |
| 'seq_len': seq_len, | |
| } | |
| arg_null = { | |
| 'context': context_null, | |
| 'seq_len': seq_len, | |
| } | |
| if offload_model or self.init_on_cpu: | |
| self.model.to(self.device) | |
| torch.cuda.empty_cache() | |
| for _, t in enumerate(tqdm(timesteps)): | |
| latent_model_input = [latent.to(self.device)] | |
| timestep = [t] | |
| timestep = torch.stack(timestep).to(self.device) | |
| temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten() | |
| temp_ts = torch.cat([ | |
| temp_ts, | |
| temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep | |
| ]) | |
| timestep = temp_ts.unsqueeze(0) | |
| noise_pred_cond = self.model( | |
| latent_model_input, t=timestep, **arg_c)[0] | |
| if offload_model: | |
| torch.cuda.empty_cache() | |
| noise_pred_uncond = self.model( | |
| latent_model_input, t=timestep, **arg_null)[0] | |
| if offload_model: | |
| torch.cuda.empty_cache() | |
| noise_pred = noise_pred_uncond + guide_scale * ( | |
| noise_pred_cond - noise_pred_uncond) | |
| temp_x0 = sample_scheduler.step( | |
| noise_pred.unsqueeze(0), | |
| t, | |
| latent.unsqueeze(0), | |
| return_dict=False, | |
| generator=seed_g)[0] | |
| latent = temp_x0.squeeze(0) | |
| latent = (1. - mask2[0]) * z[0] + mask2[0] * latent | |
| x0 = [latent] | |
| del latent_model_input, timestep | |
| if offload_model: | |
| self.model.cpu() | |
| torch.cuda.synchronize() | |
| torch.cuda.empty_cache() | |
| if self.rank == 0: | |
| videos = self.vae.decode(x0) | |
| del noise, latent, x0 | |
| del sample_scheduler | |
| if offload_model: | |
| gc.collect() | |
| torch.cuda.synchronize() | |
| if dist.is_initialized(): | |
| dist.barrier() | |
| return videos[0] if self.rank == 0 else None | |