ALSv/self-forcing

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 numpy as np13import torch14import torch.cuda.amp as amp15import torch.distributed as dist16import torchvision.transforms.functional as TF17from tqdm import tqdm18 19from .distributed.fsdp import shard_model20from .modules.clip import CLIPModel21from .modules.model import WanModel22from .modules.t5 import T5EncoderModel23from .modules.vae import WanVAE24from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler,25                               get_sampling_sigmas, retrieve_timesteps)26from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler27 28 29class WanI2V:30 31    def __init__(32        self,33        config,34        checkpoint_dir,35        device_id=0,36        rank=0,37        t5_fsdp=False,38        dit_fsdp=False,39        use_usp=False,40        t5_cpu=False,41        init_on_cpu=True,42    ):43        r"""44        Initializes the image-to-video generation model components.45 46        Args:47            config (EasyDict):48                Object containing model parameters initialized from config.py49            checkpoint_dir (`str`):50                Path to directory containing model checkpoints51            device_id (`int`,  *optional*, defaults to 0):52                Id of target GPU device53            rank (`int`,  *optional*, defaults to 0):54                Process rank for distributed training55            t5_fsdp (`bool`, *optional*, defaults to False):56                Enable FSDP sharding for T5 model57            dit_fsdp (`bool`, *optional*, defaults to False):58                Enable FSDP sharding for DiT model59            use_usp (`bool`, *optional*, defaults to False):60                Enable distribution strategy of USP.61            t5_cpu (`bool`, *optional*, defaults to False):62                Whether to place T5 model on CPU. Only works without t5_fsdp.63            init_on_cpu (`bool`, *optional*, defaults to True):64                Enable initializing Transformer Model on CPU. Only works without FSDP or USP.65        """66        self.device = torch.device(f"cuda:{device_id}")67        self.config = config68        self.rank = rank69        self.use_usp = use_usp70        self.t5_cpu = t5_cpu71 72        self.num_train_timesteps = config.num_train_timesteps73        self.param_dtype = config.param_dtype74 75        shard_fn = partial(shard_model, device_id=device_id)76        self.text_encoder = T5EncoderModel(77            text_len=config.text_len,78            dtype=config.t5_dtype,79            device=torch.device('cpu'),80            checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),81            tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),82            shard_fn=shard_fn if t5_fsdp else None,83        )84 85        self.vae_stride = config.vae_stride86        self.patch_size = config.patch_size87        self.vae = WanVAE(88            vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),89            device=self.device)90 91        self.clip = CLIPModel(92            dtype=config.clip_dtype,93            device=self.device,94            checkpoint_path=os.path.join(checkpoint_dir,95                                         config.clip_checkpoint),96            tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))97 98        logging.info(f"Creating WanModel from {checkpoint_dir}")99        self.model = WanModel.from_pretrained(checkpoint_dir)100        self.model.eval().requires_grad_(False)101 102        if t5_fsdp or dit_fsdp or use_usp:103            init_on_cpu = False104 105        if use_usp:106            from xfuser.core.distributed import \107                get_sequence_parallel_world_size108 109            from .distributed.xdit_context_parallel import (usp_attn_forward,110                                                            usp_dit_forward)111            for block in self.model.blocks:112                block.self_attn.forward = types.MethodType(113                    usp_attn_forward, block.self_attn)114            self.model.forward = types.MethodType(usp_dit_forward, self.model)115            self.sp_size = get_sequence_parallel_world_size()116        else:117            self.sp_size = 1118 119        if dist.is_initialized():120            dist.barrier()121        if dit_fsdp:122            self.model = shard_fn(self.model)123        else:124            if not init_on_cpu:125                self.model.to(self.device)126 127        self.sample_neg_prompt = config.sample_neg_prompt128 129    def generate(self,130                 input_prompt,131                 img,132                 max_area=720 * 1280,133                 frame_num=81,134                 shift=5.0,135                 sample_solver='unipc',136                 sampling_steps=40,137                 guide_scale=5.0,138                 n_prompt="",139                 seed=-1,140                 offload_model=True):141        r"""142        Generates video frames from input image and text prompt using diffusion process.143 144        Args:145            input_prompt (`str`):146                Text prompt for content generation.147            img (PIL.Image.Image):148                Input image tensor. Shape: [3, H, W]149            max_area (`int`, *optional*, defaults to 720*1280):150                Maximum pixel area for latent space calculation. Controls video resolution scaling151            frame_num (`int`, *optional*, defaults to 81):152                How many frames to sample from a video. The number should be 4n+1153            shift (`float`, *optional*, defaults to 5.0):154                Noise schedule shift parameter. Affects temporal dynamics155                [NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0.156            sample_solver (`str`, *optional*, defaults to 'unipc'):157                Solver used to sample the video.158            sampling_steps (`int`, *optional*, defaults to 40):159                Number of diffusion sampling steps. Higher values improve quality but slow generation160            guide_scale (`float`, *optional*, defaults 5.0):161                Classifier-free guidance scale. Controls prompt adherence vs. creativity162            n_prompt (`str`, *optional*, defaults to ""):163                Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`164            seed (`int`, *optional*, defaults to -1):165                Random seed for noise generation. If -1, use random seed166            offload_model (`bool`, *optional*, defaults to True):167                If True, offloads models to CPU during generation to save VRAM168 169        Returns:170            torch.Tensor:171                Generated video frames tensor. Dimensions: (C, N H, W) where:172                - C: Color channels (3 for RGB)173                - N: Number of frames (81)174                - H: Frame height (from max_area)175                - W: Frame width from max_area)176        """177        img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device)178 179        F = frame_num180        h, w = img.shape[1:]181        aspect_ratio = h / w182        lat_h = round(183            np.sqrt(max_area * aspect_ratio) // self.vae_stride[1] //184            self.patch_size[1] * self.patch_size[1])185        lat_w = round(186            np.sqrt(max_area / aspect_ratio) // self.vae_stride[2] //187            self.patch_size[2] * self.patch_size[2])188        h = lat_h * self.vae_stride[1]189        w = lat_w * self.vae_stride[2]190 191        max_seq_len = ((F - 1) // self.vae_stride[0] + 1) * lat_h * lat_w // (192            self.patch_size[1] * self.patch_size[2])193        max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size194 195        seed = seed if seed >= 0 else random.randint(0, sys.maxsize)196        seed_g = torch.Generator(device=self.device)197        seed_g.manual_seed(seed)198        noise = torch.randn(199            16,200            21,201            lat_h,202            lat_w,203            dtype=torch.float32,204            generator=seed_g,205            device=self.device)206 207        msk = torch.ones(1, 81, lat_h, lat_w, device=self.device)208        msk[:, 1:] = 0209        msk = torch.concat([210            torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]211        ],212            dim=1)213        msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)214        msk = msk.transpose(1, 2)[0]215 216        if n_prompt == "":217            n_prompt = self.sample_neg_prompt218 219        # preprocess220        if not self.t5_cpu:221            self.text_encoder.model.to(self.device)222            context = self.text_encoder([input_prompt], self.device)223            context_null = self.text_encoder([n_prompt], self.device)224            if offload_model:225                self.text_encoder.model.cpu()226        else:227            context = self.text_encoder([input_prompt], torch.device('cpu'))228            context_null = self.text_encoder([n_prompt], torch.device('cpu'))229            context = [t.to(self.device) for t in context]230            context_null = [t.to(self.device) for t in context_null]231 232        self.clip.model.to(self.device)233        clip_context = self.clip.visual([img[:, None, :, :]])234        if offload_model:235            self.clip.model.cpu()236 237        y = self.vae.encode([238            torch.concat([239                torch.nn.functional.interpolate(240                    img[None].cpu(), size=(h, w), mode='bicubic').transpose(241                        0, 1),242                torch.zeros(3, 80, h, w)243            ],244                dim=1).to(self.device)245        ])[0]246        y = torch.concat([msk, y])247 248        @contextmanager249        def noop_no_sync():250            yield251 252        no_sync = getattr(self.model, 'no_sync', noop_no_sync)253 254        # evaluation mode255        with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():256 257            if sample_solver == 'unipc':258                sample_scheduler = FlowUniPCMultistepScheduler(259                    num_train_timesteps=self.num_train_timesteps,260                    shift=1,261                    use_dynamic_shifting=False)262                sample_scheduler.set_timesteps(263                    sampling_steps, device=self.device, shift=shift)264                timesteps = sample_scheduler.timesteps265            elif sample_solver == 'dpm++':266                sample_scheduler = FlowDPMSolverMultistepScheduler(267                    num_train_timesteps=self.num_train_timesteps,268                    shift=1,269                    use_dynamic_shifting=False)270                sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)271                timesteps, _ = retrieve_timesteps(272                    sample_scheduler,273                    device=self.device,274                    sigmas=sampling_sigmas)275            else:276                raise NotImplementedError("Unsupported solver.")277 278            # sample videos279            latent = noise280 281            arg_c = {282                'context': [context[0]],283                'clip_fea': clip_context,284                'seq_len': max_seq_len,285                'y': [y],286            }287 288            arg_null = {289                'context': context_null,290                'clip_fea': clip_context,291                'seq_len': max_seq_len,292                'y': [y],293            }294 295            if offload_model:296                torch.cuda.empty_cache()297 298            self.model.to(self.device)299            for _, t in enumerate(tqdm(timesteps)):300                latent_model_input = [latent.to(self.device)]301                timestep = [t]302 303                timestep = torch.stack(timestep).to(self.device)304 305                noise_pred_cond = self.model(306                    latent_model_input, t=timestep, **arg_c)[0].to(307                        torch.device('cpu') if offload_model else self.device)308                if offload_model:309                    torch.cuda.empty_cache()310                noise_pred_uncond = self.model(311                    latent_model_input, t=timestep, **arg_null)[0].to(312                        torch.device('cpu') if offload_model else self.device)313                if offload_model:314                    torch.cuda.empty_cache()315                noise_pred = noise_pred_uncond + guide_scale * (316                    noise_pred_cond - noise_pred_uncond)317 318                latent = latent.to(319                    torch.device('cpu') if offload_model else self.device)320 321                temp_x0 = sample_scheduler.step(322                    noise_pred.unsqueeze(0),323                    t,324                    latent.unsqueeze(0),325                    return_dict=False,326                    generator=seed_g)[0]327                latent = temp_x0.squeeze(0)328 329                x0 = [latent.to(self.device)]330                del latent_model_input, timestep331 332            if offload_model:333                self.model.cpu()334                torch.cuda.empty_cache()335 336            if self.rank == 0:337                videos = self.vae.decode(x0)338 339        del noise, latent340        del sample_scheduler341        if offload_model:342            gc.collect()343            torch.cuda.synchronize()344        if dist.is_initialized():345            dist.barrier()346 347        return videos[0] if self.rank == 0 else None348