ALSv/self-forcing
0
1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import gc3import logging4import math5import os6import random7import sys8import types9from contextlib import contextmanager10from functools import partial11 12import torch13import torch.cuda.amp as amp14import torch.distributed as dist15from tqdm import tqdm16 17from .distributed.fsdp import shard_model18from .modules.model import WanModel19from .modules.t5 import T5EncoderModel20from .modules.vae import WanVAE21from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler,22 get_sampling_sigmas, retrieve_timesteps)23from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler24 25 26class WanT2V:27 28 def __init__(29 self,30 config,31 checkpoint_dir,32 device_id=0,33 rank=0,34 t5_fsdp=False,35 dit_fsdp=False,36 use_usp=False,37 t5_cpu=False,38 ):39 r"""40 Initializes the Wan text-to-video generation model components.41 42 Args:43 config (EasyDict):44 Object containing model parameters initialized from config.py45 checkpoint_dir (`str`):46 Path to directory containing model checkpoints47 device_id (`int`, *optional*, defaults to 0):48 Id of target GPU device49 rank (`int`, *optional*, defaults to 0):50 Process rank for distributed training51 t5_fsdp (`bool`, *optional*, defaults to False):52 Enable FSDP sharding for T5 model53 dit_fsdp (`bool`, *optional*, defaults to False):54 Enable FSDP sharding for DiT model55 use_usp (`bool`, *optional*, defaults to False):56 Enable distribution strategy of USP.57 t5_cpu (`bool`, *optional*, defaults to False):58 Whether to place T5 model on CPU. Only works without t5_fsdp.59 """60 self.device = torch.device(f"cuda:{device_id}")61 self.config = config62 self.rank = rank63 self.t5_cpu = t5_cpu64 65 self.num_train_timesteps = config.num_train_timesteps66 self.param_dtype = config.param_dtype67 68 shard_fn = partial(shard_model, device_id=device_id)69 self.text_encoder = T5EncoderModel(70 text_len=config.text_len,71 dtype=config.t5_dtype,72 device=torch.device('cpu'),73 checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),74 tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),75 shard_fn=shard_fn if t5_fsdp else None)76 77 self.vae_stride = config.vae_stride78 self.patch_size = config.patch_size79 self.vae = WanVAE(80 vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),81 device=self.device)82 83 logging.info(f"Creating WanModel from {checkpoint_dir}")84 self.model = WanModel.from_pretrained(checkpoint_dir)85 self.model.eval().requires_grad_(False)86 87 if use_usp:88 from xfuser.core.distributed import \89 get_sequence_parallel_world_size90 91 from .distributed.xdit_context_parallel import (usp_attn_forward,92 usp_dit_forward)93 for block in self.model.blocks:94 block.self_attn.forward = types.MethodType(95 usp_attn_forward, block.self_attn)96 self.model.forward = types.MethodType(usp_dit_forward, self.model)97 self.sp_size = get_sequence_parallel_world_size()98 else:99 self.sp_size = 1100 101 if dist.is_initialized():102 dist.barrier()103 if dit_fsdp:104 self.model = shard_fn(self.model)105 else:106 self.model.to(self.device)107 108 self.sample_neg_prompt = config.sample_neg_prompt109 110 def generate(self,111 input_prompt,112 size=(1280, 720),113 frame_num=81,114 shift=5.0,115 sample_solver='unipc',116 sampling_steps=50,117 guide_scale=5.0,118 n_prompt="",119 seed=-1,120 offload_model=True):121 r"""122 Generates video frames from text prompt using diffusion process.123 124 Args:125 input_prompt (`str`):126 Text prompt for content generation127 size (tupele[`int`], *optional*, defaults to (1280,720)):128 Controls video resolution, (width,height).129 frame_num (`int`, *optional*, defaults to 81):130 How many frames to sample from a video. The number should be 4n+1131 shift (`float`, *optional*, defaults to 5.0):132 Noise schedule shift parameter. Affects temporal dynamics133 sample_solver (`str`, *optional*, defaults to 'unipc'):134 Solver used to sample the video.135 sampling_steps (`int`, *optional*, defaults to 40):136 Number of diffusion sampling steps. Higher values improve quality but slow generation137 guide_scale (`float`, *optional*, defaults 5.0):138 Classifier-free guidance scale. Controls prompt adherence vs. creativity139 n_prompt (`str`, *optional*, defaults to ""):140 Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`141 seed (`int`, *optional*, defaults to -1):142 Random seed for noise generation. If -1, use random seed.143 offload_model (`bool`, *optional*, defaults to True):144 If True, offloads models to CPU during generation to save VRAM145 146 Returns:147 torch.Tensor:148 Generated video frames tensor. Dimensions: (C, N H, W) where:149 - C: Color channels (3 for RGB)150 - N: Number of frames (81)151 - H: Frame height (from size)152 - W: Frame width from size)153 """154 # preprocess155 F = frame_num156 target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1,157 size[1] // self.vae_stride[1],158 size[0] // self.vae_stride[2])159 160 seq_len = math.ceil((target_shape[2] * target_shape[3]) /161 (self.patch_size[1] * self.patch_size[2]) *162 target_shape[1] / self.sp_size) * self.sp_size163 164 if n_prompt == "":165 n_prompt = self.sample_neg_prompt166 seed = seed if seed >= 0 else random.randint(0, sys.maxsize)167 seed_g = torch.Generator(device=self.device)168 seed_g.manual_seed(seed)169 170 if not self.t5_cpu:171 self.text_encoder.model.to(self.device)172 context = self.text_encoder([input_prompt], self.device)173 context_null = self.text_encoder([n_prompt], self.device)174 if offload_model:175 self.text_encoder.model.cpu()176 else:177 context = self.text_encoder([input_prompt], torch.device('cpu'))178 context_null = self.text_encoder([n_prompt], torch.device('cpu'))179 context = [t.to(self.device) for t in context]180 context_null = [t.to(self.device) for t in context_null]181 182 noise = [183 torch.randn(184 target_shape[0],185 target_shape[1],186 target_shape[2],187 target_shape[3],188 dtype=torch.float32,189 device=self.device,190 generator=seed_g)191 ]192 193 @contextmanager194 def noop_no_sync():195 yield196 197 no_sync = getattr(self.model, 'no_sync', noop_no_sync)198 199 # evaluation mode200 with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():201 202 if sample_solver == 'unipc':203 sample_scheduler = FlowUniPCMultistepScheduler(204 num_train_timesteps=self.num_train_timesteps,205 shift=1,206 use_dynamic_shifting=False)207 sample_scheduler.set_timesteps(208 sampling_steps, device=self.device, shift=shift)209 timesteps = sample_scheduler.timesteps210 elif sample_solver == 'dpm++':211 sample_scheduler = FlowDPMSolverMultistepScheduler(212 num_train_timesteps=self.num_train_timesteps,213 shift=1,214 use_dynamic_shifting=False)215 sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)216 timesteps, _ = retrieve_timesteps(217 sample_scheduler,218 device=self.device,219 sigmas=sampling_sigmas)220 else:221 raise NotImplementedError("Unsupported solver.")222 223 # sample videos224 latents = noise225 226 arg_c = {'context': context, 'seq_len': seq_len}227 arg_null = {'context': context_null, 'seq_len': seq_len}228 229 for _, t in enumerate(tqdm(timesteps)):230 latent_model_input = latents231 timestep = [t]232 233 timestep = torch.stack(timestep)234 235 self.model.to(self.device)236 noise_pred_cond = self.model(237 latent_model_input, t=timestep, **arg_c)[0]238 noise_pred_uncond = self.model(239 latent_model_input, t=timestep, **arg_null)[0]240 241 noise_pred = noise_pred_uncond + guide_scale * (242 noise_pred_cond - noise_pred_uncond)243 244 temp_x0 = sample_scheduler.step(245 noise_pred.unsqueeze(0),246 t,247 latents[0].unsqueeze(0),248 return_dict=False,249 generator=seed_g)[0]250 latents = [temp_x0.squeeze(0)]251 252 x0 = latents253 if offload_model:254 self.model.cpu()255 if self.rank == 0:256 videos = self.vae.decode(x0)257 258 del noise, latents259 del sample_scheduler260 if offload_model:261 gc.collect()262 torch.cuda.synchronize()263 if dist.is_initialized():264 dist.barrier()265 266 return videos[0] if self.rank == 0 else None267