wan/scail.py

# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import logging
import math
import os
import random
import sys
import types
from contextlib import contextmanager
from functools import partial

import numpy as np
import torch
import torch.cuda.amp as amp
import torch.distributed as dist
import torchvision.transforms.functional as TF
import torch.nn.functional as F
from tqdm import tqdm
from einops import rearrange
from safetensors.torch import load_file
import gc

from .distributed.fsdp import shard_model
from .modules.clip import CLIPModel
from .modules.model_scail import SCAILModel
from .modules.model_scail2 import SCAIL2Model
from .modules.t5 import T5EncoderModel
from .modules.vae import WanVAE
from .utils.fm_solvers import (
    FlowDPMSolverMultistepScheduler,
    get_sampling_sigmas,
    retrieve_timesteps,
)
from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
from .utils.lora import fuse_lora_with_diff_b
from .utils.scail_utils import extract_and_compress_mask_to_latent

class SCAIL2Pipeline:

    def __init__(
        self,
        config,
        checkpoint_dir,
        scail_safetensors_path, 
        scail_config_path="./config.json",
        device_id=0,
        rank=0,
        t5_fsdp=False,
        dit_fsdp=False,
        use_usp=False,
        t5_cpu=False,
        init_on_cpu=True,
        lora_path=None,
        lora_alpha=None,
    ):
        r"""
        Initializes the image-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_usp (`bool`, *optional*, defaults to False):
                Enable distribution strategy of USP.
            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.
        """
        self.device = torch.device(f"cuda:{device_id}")
        self.config = config
        self.rank = rank
        self.use_usp = use_usp
        self.t5_cpu = t5_cpu
        self.lora_path = lora_path
        self.lora_alpha = lora_alpha

        self.num_train_timesteps = config.num_train_timesteps
        self.param_dtype = config.param_dtype

        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 = WanVAE(
            vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
            device=self.device)

        self.clip = CLIPModel(
            dtype=config.clip_dtype,
            device=self.device,
            checkpoint_path=os.path.join(checkpoint_dir,
                                         config.clip_checkpoint),
            tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))

        logging.info(f"Creating WanSCAILModel from {scail_safetensors_path}")
        self.model = SCAIL2Model.from_config(scail_config_path)
        state_dict = load_file(scail_safetensors_path)
        self.model.load_state_dict(state_dict)
        if self.lora_path is not None:
            if self.lora_alpha is None:
                self.lora_alpha = 1.0
            self.fuse_lora(self.lora_path, self.lora_alpha)
        self.model.eval().requires_grad_(False)

        if t5_fsdp or dit_fsdp or use_usp:
            init_on_cpu = False

        if use_usp:
            from xfuser.core.distributed import get_sequence_parallel_world_size

            from .distributed.xdit_context_parallel import (
                usp_attn_forward,
                usp_dit_forward,
            )
            for block in self.model.blocks:
                block.self_attn.forward = types.MethodType(
                    usp_attn_forward, block.self_attn)
            self.model.forward = types.MethodType(usp_dit_forward, self.model)
            self.sp_size = get_sequence_parallel_world_size()
        else:
            self.sp_size = 1

        if dist.is_initialized():
            dist.barrier()
        if dit_fsdp:
            self.model = shard_fn(self.model)
        else:
            if not init_on_cpu:
                self.model.to(self.device)

        self.sample_neg_prompt = config.sample_neg_prompt

    def fuse_lora(self, lora_path, alpha=1.0):
        logging.info(f"Fusing LoRA from {lora_path}, strength = {alpha}.")
        lora_state_dict = load_file(lora_path)
        fuse_lora_with_diff_b(self.model, lora_state_dict, alpha=alpha)

    def generate(self,
                 input_prompt,
                 img,
                 ref_mask_img: torch.Tensor,
                 pose_video: torch.Tensor,
                 driving_mask_video: torch.Tensor,
                 replace_flag: bool,
                 segment_len=81,
                 segment_overlap=5,
                 shift=5.0,
                 sample_solver='unipc',
                 sampling_steps=40,
                 guide_scale=5.0,
                 n_prompt=None,
                 seed=-1,
                 offload_model=True,
                 additional_ref_imgs: list[torch.Tensor] = None,
                 additional_ref_mask_imgs: list[torch.Tensor] = None,
                 **kwargs):
        r"""
        Generates video frames from input image and text prompt using diffusion process.

        Args:
            input_prompt (`str`):
                Text prompt for content generation.
            img (torch.Tensor):
                Input image tensor. Shape: [3, H, W], Range: (-1, 1)
            ref_mask_img (torch.Tensor):
                Input image mask tensor. Shape: [3, H, W], Range: (-1, 1)
            pose_video (torch.Tensor):
                Input pose video. Shape: [T, C, H, W]
            driving_mask_video (torch.Tensor):
                Input driving mask tensor. Shape: [3, T, H, W], Range: (-1, 1)
            replace_flag (bool):
                True for replacement mode, False for animation mode
            segment_len (`int`, *optional*, defaults to 81):
                Number of pixel frames sampled in each segment.
            segment_overlap (`int`, *optional*, defaults to 5):
                Number of pixel frames shared with the previous segment as clean history.
            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 None):
                Negative prompt for content exclusion. If not given, use ""
            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, T, H, W).
        """
        if segment_len <= 0:
            raise ValueError("segment_len must be positive")
        if segment_overlap <= 0 or segment_overlap >= segment_len:
            raise ValueError("segment_overlap must be in (0, segment_len)")

        pose_video = pose_video.to(self.device)
        driving_mask_video = driving_mask_video.to(self.device)
        if not isinstance(img, torch.Tensor):
            img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device) # 3 H W
        else:
            img = img.to(self.device) # 3 H W, -1 ~ 1
        ori_img = img.unsqueeze(0).to(self.device) # 1, 3, H, W

        if not isinstance(ref_mask_img, torch.Tensor):
            ref_mask_img = TF.to_tensor(ref_mask_img).sub_(0.5).div_(0.5).to(self.device) # 3 H W
        else:
            ref_mask_img = ref_mask_img.to(self.device) # 3 H W, -1 ~ 1

        if additional_ref_imgs is not None:
            if additional_ref_mask_imgs is None:
                raise ValueError('additional_ref_mask_imgs is required when additional_ref_imgs is provided.')
            if isinstance(additional_ref_imgs, torch.Tensor):
                additional_ref_imgs = [additional_ref_imgs]
            if isinstance(additional_ref_mask_imgs, torch.Tensor):
                additional_ref_mask_imgs = [additional_ref_mask_imgs]
            if len(additional_ref_imgs) != len(additional_ref_mask_imgs):
                raise ValueError(
                    'additional_ref_imgs and additional_ref_mask_imgs must have the same length, '
                    'got %d and %d.' % (len(additional_ref_imgs), len(additional_ref_mask_imgs)))
            additional_ref_imgs = [
                TF.to_tensor(u).sub_(0.5).div_(0.5).to(self.device)
                if not isinstance(u, torch.Tensor) else u.to(self.device)
                for u in additional_ref_imgs
            ]
            additional_ref_mask_imgs = [
                TF.to_tensor(u).sub_(0.5).div_(0.5).to(self.device)
                if not isinstance(u, torch.Tensor) else u.to(self.device)
                for u in additional_ref_mask_imgs
            ]
        elif additional_ref_mask_imgs is not None:
            raise ValueError('additional_ref_mask_imgs requires additional_ref_imgs.')
        num_frames = pose_video.shape[0]
        if driving_mask_video.shape[1] != num_frames:
            raise ValueError(
                f"pose_video and driving_mask_video must have the same frame count, "
                f"got {num_frames} and {driving_mask_video.shape[1]}")

        def build_segments(total_frames):
            if total_frames <= segment_len:
                keep = ((total_frames - 1) // self.vae_stride[0]) * self.vae_stride[0] + 1
                return [(0, keep)]
            segments = []
            start = 0
            stride = segment_len - segment_overlap
            while start < total_frames:
                end = start + segment_len
                if end > total_frames:
                    break
                segments.append((start, end))
                start += stride
            return segments

        segments = build_segments(num_frames)
        if len(segments) == 0:
            raise ValueError(
                f"No valid segment was produced for {num_frames} frames. "
                f"Use a longer driving video or reduce segment_len.")
        if len(segments) > 1:
            logging.info(
                f"Sampling {len(segments)} segments with segment_len={segment_len}, "
                f"segment_overlap={segment_overlap}.")

        ref_latent = self.vae.encode([rearrange(ori_img, 't c h w -> c t h w')])[0]
        
        additional_ref_latent = None
        additional_ref_mask_latent_28ch = None
        if additional_ref_imgs is not None:
            additional_ref_latents = []
            additional_ref_mask_latents = []
            for additional_ref_img, additional_ref_mask_img in zip(additional_ref_imgs, additional_ref_mask_imgs):
                ori_additional_ref_img = additional_ref_img.unsqueeze(0).to(self.device)
                additional_ref_latents.append(
                    self.vae.encode([rearrange(ori_additional_ref_img, 't c h w -> c t h w')])[0]
                )
                additional_ref_mask_latents.append(
                    extract_and_compress_mask_to_latent(
                        additional_ref_mask_img.unsqueeze(1), additional_spatial_downsample=1
                    )
                )
            additional_ref_latent = torch.cat(additional_ref_latents, dim=1)
            additional_ref_mask_latent_28ch = torch.cat(additional_ref_mask_latents, dim=1)
        ref_mask_latent_28ch = extract_and_compress_mask_to_latent(
            ref_mask_img.unsqueeze(1), additional_spatial_downsample=1
        )  # (28, 1, H_lat, W_lat)
        lat_c = ref_latent.shape[0]

        # TODO: support sequence_parallel
        max_seq_len = 1e10
        # max_seq_len = ((F - 1) // self.vae_stride[0] + 1) * lat_h * lat_w // (
        #     self.patch_size[1] * self.patch_size[2])
        # max_seq_len = int(math.ceil(max_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)

        if n_prompt is None:
            n_prompt = ""

        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]

        self.clip.model.to(self.device)
        clip_context = self.clip.visual([img[:, None, :, :]])
        if offload_model:
            self.clip.model.cpu()

        @contextmanager
        def noop_no_sync():
            yield

        no_sync = getattr(self.model, 'no_sync', noop_no_sync)

        def apply_clean_history(latent, history_latent):
            if history_latent is None:
                return latent
            history_t = history_latent.shape[1]
            latent[:, :history_t] = history_latent.to(device=latent.device, dtype=latent.dtype)
            return latent

        output_segments = []
        prev_history_pixel = None

        with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():

            def build_sample_scheduler():
                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.")
                return sample_scheduler, timesteps

            def sample_func(latent, arg_c, arg_null, history_latent):
                if offload_model:
                    self.model.to(self.device)
                latent = apply_clean_history(latent, history_latent)
                for _, t in enumerate(tqdm(timesteps)):
                    latent_model_input = [apply_clean_history(latent.to(self.device), history_latent)]
                    timestep = [t]

                    timestep = torch.stack(timestep).to(self.device)

                    noise_pred_cond = self.model(
                        latent_model_input, t=timestep, **arg_c)[0].to(
                            torch.device('cpu') if offload_model else self.device)
                    if offload_model:
                        torch.cuda.empty_cache()
                    if guide_scale <= 1.0:
                        noise_pred = noise_pred_cond
                    else:
                        noise_pred_uncond = self.model(
                            latent_model_input, t=timestep, **arg_null)[0].to(
                                torch.device('cpu') if offload_model else self.device)
                        if offload_model:
                            torch.cuda.empty_cache()
                        noise_pred = noise_pred_uncond + guide_scale * (
                            noise_pred_cond - noise_pred_uncond)

                    latent = latent.to(
                        torch.device('cpu') if offload_model else self.device)

                    temp_x0 = sample_scheduler.step(
                        noise_pred.unsqueeze(0),
                        t,
                        latent.unsqueeze(0),
                        return_dict=False,
                        generator=seed_g)[0]
                    latent = apply_clean_history(temp_x0.squeeze(0), history_latent)

                    x0 = [latent.to(self.device)]
                    del latent_model_input, timestep

                if offload_model:
                    self.model.cpu()
                    torch.cuda.empty_cache()

                if self.rank == 0:
                    videos = self.vae.decode(x0)
                return videos

            for seg_idx, (seg_start, seg_end) in enumerate(segments):
                logging.info(
                    f"Processing segment {seg_idx + 1}/{len(segments)}: "
                    f"frames [{seg_start}, {seg_end})")
                sample_scheduler, timesteps = build_sample_scheduler()

                pose_segment = pose_video[seg_start:seg_end]
                smpl_render_video = F.interpolate(
                    pose_segment, scale_factor=0.5, mode='bilinear', align_corners=False)
                pose_latent = self.vae.encode([rearrange(smpl_render_video, 't c h w -> c t h w')])[0]

                lat_t = pose_latent.shape[1]
                _, lat_h, lat_w = ref_latent.shape[1:]

                null_noisy_mask = torch.zeros(
                    ref_mask_latent_28ch.shape[0], lat_t, lat_h, lat_w,
                    device=self.device, dtype=ref_mask_latent_28ch.dtype)
                ref_masks = torch.cat([ref_mask_latent_28ch, null_noisy_mask], dim=1)

                driving_mask_segment = driving_mask_video[:, seg_start:seg_end]
                driving_mask_segment = F.interpolate(
                    driving_mask_segment, scale_factor=0.5, mode='bilinear', align_corners=False)
                driving_masks = extract_and_compress_mask_to_latent(
                    driving_mask_segment, additional_spatial_downsample=1
                )

                history_latent = None
                history_mask = None
                if seg_idx > 0:
                    if prev_history_pixel is None:
                        raise RuntimeError("Missing previous segment history frames.")
                    history_latent = self.vae.encode([
                        prev_history_pixel.to(self.device, dtype=self.param_dtype)
                    ])[0]
                    history_t = min(history_latent.shape[1], lat_t)
                    history_mask = torch.zeros(
                        4, lat_t, lat_h, lat_w, device=self.device, dtype=torch.float32)
                    history_mask[:, :history_t] = 1
                    logging.info(
                        f"Using {prev_history_pixel.shape[1]} clean history frames "
                        f"({history_t} latent frames).")

                noise = torch.randn(
                    lat_c,
                    lat_t,
                    lat_h,
                    lat_w,
                    dtype=torch.float32,
                    generator=seed_g,
                    device=self.device)

                arg_c = {
                    'context': [context[0]],
                    'clip_fea': clip_context,
                    'seq_len': max_seq_len,
                    'ref_latents': [ref_latent],
                    'ref_masks': [ref_masks],
                    'pose_latents': [pose_latent],
                    'driving_masks': [driving_masks],
                    'history_mask': [history_mask] if history_mask is not None else None,
                    'replace_flag': replace_flag,
                    'additional_ref_latents': None if additional_ref_latent is None else [additional_ref_latent],
                    'additional_ref_masks': None if additional_ref_mask_latent_28ch is None else [additional_ref_mask_latent_28ch],
                }

                arg_null = {
                    'context': context_null,
                    'clip_fea': clip_context,
                    'seq_len': max_seq_len,
                    'ref_latents': [ref_latent],
                    'ref_masks': [ref_masks],
                    'pose_latents': [pose_latent],
                    'driving_masks': [driving_masks],
                    'history_mask': [history_mask] if history_mask is not None else None,
                    'replace_flag': replace_flag,
                    'additional_ref_latents': None if additional_ref_latent is None else [additional_ref_latent],
                    'additional_ref_masks': None if additional_ref_mask_latent_28ch is None else [additional_ref_mask_latent_28ch],
                }

                if offload_model:
                    torch.cuda.empty_cache()

                videos = sample_func(noise, arg_c, arg_null, history_latent)
                segment_video = videos[0] if self.rank == 0 else None
                if self.rank == 0:
                    if seg_idx == 0:
                        output_segments.append(segment_video.cpu())
                    else:
                        output_segments.append(segment_video[:, segment_overlap:].cpu())
                    if seg_idx < len(segments) - 1:
                        prev_history_pixel = segment_video[:, -segment_overlap:].contiguous()

                del noise, pose_latent, ref_masks, driving_masks, sample_scheduler
                if history_latent is not None:
                    del history_latent, history_mask
                if offload_model:
                    torch.cuda.empty_cache()

        if offload_model:
            gc.collect()
            torch.cuda.synchronize()
        if dist.is_initialized():
            dist.barrier()

        if self.rank == 0:
            return torch.cat(output_segments, dim=1).to(self.device)
        return None