utils/wan_wrapper.py

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.model import WanModel, RegisterTokens, GanAttentionBlock
from wan.modules.vae import _video_vae
from wan.modules.t5 import umt5_xxl
from wan.modules.causal_model import CausalWanModel
from wan.modules.causal_model_infinity import CausalWanModel as CausalWanModelInfinity
from wan.modules.causal_model_infinity_memory import CausalWanModelInfinityMemory

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('wan_models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth', map_location='cpu', weights_only=False))
        if torch.cuda.is_available():
            self.text_encoder = self.text_encoder.cuda()
        self.tokenizer = HuggingfaceTokenizer(name='wan_models/Wan2.1-T2V-1.3B/google/umt5-xxl/', seq_len=512, clean='whitespace')

    @property
    def device(self):
        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
        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.916]
        self.mean = torch.tensor(mean, dtype=torch.float32)
        self.std = torch.tensor(std, dtype=torch.float32)
        self.model = _video_vae(pretrained_path='wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth', z_dim=16).eval().requires_grad_(False)

    def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
        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)
        output = output.permute(0, 2, 1, 3, 4)
        return output

    def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool=False) -> torch.Tensor:
        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:
            output.append(decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0))
        output = torch.stack(output, dim=0)
        output = output.permute(0, 2, 1, 3, 4)
        return output

    def decode_to_pixel_chunk(self, latent: torch.Tensor, use_cache: bool=False, chunk_size: int=120) -> torch.Tensor:
        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:
            num_frames = u.shape[1]
            if num_frames <= chunk_size:
                decoded = decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)
                decoded = decoded.cpu()
            else:
                decoded_chunks = []
                for start_idx in range(0, num_frames, chunk_size):
                    end_idx = min(start_idx + chunk_size, num_frames)
                    chunk = u[:, start_idx:end_idx, :, :]
                    self.model.clear_cache()
                    decoded_chunk = decode_function(chunk.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)
                    decoded_chunks.append(decoded_chunk.cpu())
                    del decoded_chunk
                    torch.cuda.empty_cache()
                decoded = torch.cat(decoded_chunks, dim=1)
                self.model.clear_cache()
            output.append(decoded)
        output = torch.stack(output, dim=0)
        output = output.permute(0, 2, 1, 3, 4)
        return output

class WanDiffusionWrapper(torch.nn.Module):

    def __init__(self, model_name='Wan2.1-T2V-1.3B', timestep_shift=8.0, is_causal=False, local_attn_size=-1, sink_size=0, use_infinite_attention=False, dr_rope=False, tri_rope_cont=False, tri_rope_pmax=21, relative_rope=False, relative_rope_pmax=21):
        super().__init__()
        if is_causal:
            if use_infinite_attention:
                if relative_rope:
                    self.model = CausalWanModelInfinityMemory.from_pretrained(f'wan_models/{model_name}/', local_attn_size=local_attn_size, sink_size=sink_size)
                    self.model.enable_infmem(relative_rope=True, relative_rope_pmax=relative_rope_pmax, num_frame_per_block_attr=3)
                else:
                    self.model = CausalWanModelInfinity.from_pretrained(f'wan_models/{model_name}/', local_attn_size=local_attn_size, sink_size=sink_size)
            else:
                self.model = CausalWanModel.from_pretrained(f'wan_models/{model_name}/', local_attn_size=local_attn_size, sink_size=sink_size, dr_rope=dr_rope, tri_rope_cont=tri_rope_cont, tri_rope_pmax=tri_rope_pmax, relative_rope=relative_rope, relative_rope_pmax=relative_rope_pmax)
        else:
            self.model = WanModel.from_pretrained(f'wan_models/{model_name}/')
        self.model.eval()
        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)
        self.seq_len = 1560 * local_attn_size if local_attn_size > 21 else 32760
        self.post_init()

    def enable_gradient_checkpointing(self) -> None:
        self.model.enable_gradient_checkpointing()

    def adding_cls_branch(self, atten_dim=1536, num_class=4, time_embed_dim=0) -> None:
        self._cls_pred_branch = nn.Sequential(nn.LayerNorm(atten_dim * 3 + time_embed_dim), nn.Linear(atten_dim * 3 + time_embed_dim, 1536), nn.SiLU(), nn.Linear(atten_dim, num_class))
        self._cls_pred_branch.requires_grad_(True)
        num_registers = 3
        self._register_tokens = RegisterTokens(num_registers=num_registers, dim=atten_dim)
        self._register_tokens.requires_grad_(True)
        gan_ca_blocks = []
        for _ in range(num_registers):
            block = GanAttentionBlock()
            gan_ca_blocks.append(block)
        self._gan_ca_blocks = nn.ModuleList(gan_ca_blocks)
        self._gan_ca_blocks.requires_grad_(True)

    def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
        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)

    @staticmethod
    def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
        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, classify_mode: Optional[bool]=False, concat_time_embeddings: Optional[bool]=False, clean_x: Optional[torch.Tensor]=None, aug_t: Optional[torch.Tensor]=None, cache_start: Optional[int]=None, sink_recache_after_switch=False) -> torch.Tensor:
        prompt_embeds = conditional_dict['prompt_embeds']
        if self.uniform_timestep:
            input_timestep = timestep[:, 0]
        else:
            input_timestep = timestep
        logits = None
        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, sink_recache_after_switch=sink_recache_after_switch).permute(0, 2, 1, 3, 4)
        elif clean_x 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, clean_x=clean_x.permute(0, 2, 1, 3, 4), aug_t=aug_t, sink_recache_after_switch=sink_recache_after_switch).permute(0, 2, 1, 3, 4)
        elif classify_mode:
            flow_pred, logits = self.model(noisy_image_or_video.permute(0, 2, 1, 3, 4), t=input_timestep, context=prompt_embeds, seq_len=self.seq_len, classify_mode=True, register_tokens=self._register_tokens, cls_pred_branch=self._cls_pred_branch, gan_ca_blocks=self._gan_ca_blocks, concat_time_embeddings=concat_time_embeddings, sink_recache_after_switch=sink_recache_after_switch)
            flow_pred = flow_pred.permute(0, 2, 1, 3, 4)
        else:
            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, sink_recache_after_switch=sink_recache_after_switch).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])
        if logits is not None:
            return (flow_pred, pred_x0, logits)
        return (flow_pred, pred_x0)

    def get_scheduler(self) -> SchedulerInterface:
        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):
        self.get_scheduler()