Blame for 1.3B/code/utils/wan_wrapper.py

import types
from typing import List, Optional
import torch
from torch import nn

from utils.paths import model_path
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

_DEFAULT_TEXT_ENCODER_NAME = "Wan2.1-T2V-1.3B"


class WanTextEncoder(torch.nn.Module):
    def __init__(self, dtype: torch.dtype = torch.float32) -> None:
        super().__init__()

        self.text_encoder = umt5_xxl(
            encoder_only=True,
            return_tokenizer=False,
            dtype=dtype,
            device=torch.device('cpu')
        ).eval().requires_grad_(False)
        encoder_root = model_path(_DEFAULT_TEXT_ENCODER_NAME)
        self.text_encoder.load_state_dict(
            torch.load(str(encoder_root / "models_t5_umt5-xxl-enc-bf16.pth"),
                       map_location='cpu', weights_only=False)
        )

        self.tokenizer = HuggingfaceTokenizer(
            name=str(encoder_root / "google" / "umt5-xxl") + "/",
            seq_len=512, clean='whitespace',
        )

    @property
    def device(self):
        # Assume we are always on GPU
        return torch.cuda.current_device()

    def forward(
        self,
        text_prompts: List[str],
        return_attention_mask: bool = False,
    ) -> 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  # set padding to 0.0

        result = {
            "prompt_embeds": context
        }
        if return_attention_mask:
            result["prompt_attention_mask"] = mask.gt(0)
        return result


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.9160
        ]
        self.mean = torch.tensor(mean, dtype=torch.float32)
        self.std = torch.tensor(std, dtype=torch.float32)

        # init model
        self.model = _video_vae(
            pretrained_path=str(model_path(_DEFAULT_TEXT_ENCODER_NAME) / "Wan2.1_VAE.pth"),
            z_dim=16,
        ).eval().requires_grad_(False)

    def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
        # pixel: [batch_size, num_channels, num_frames, height, width]
        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)
        # from [batch_size, num_channels, num_frames, height, width]
        # to [batch_size, num_frames, num_channels, height, width]
        output = output.permute(0, 2, 1, 3, 4)
        return output

    def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torch.Tensor:
        # from [batch_size, num_frames, num_channels, height, width]
        # to [batch_size, num_channels, num_frames, height, width]
        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)
        # from [batch_size, num_channels, num_frames, height, width]
        # to [batch_size, num_frames, num_channels, height, width]
        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,
            attention_type="softmax",       # one of {"softmax", "linear", "hybrid"}; only used when is_causal=True
            num_frame_per_block=1,          # only used when attention_type in {"linear", "hybrid"}
            hybrid_gate=False,              # legacy bool: equivalent to hybrid_gate_type="lambda".  Kept for back-compat.
            hybrid_gate_init_lambda=0.12,   # initial value of λ before any training: σ(bias) where bias = logit(init_lambda).  ``lambda`` mode only.
            hybrid_gate_type="none",        # one of {"none", "lambda", "alpha_a0", "alpha_a1", "alpha_a2", "alpha_a3"}; alpha modes use the revised normalized-branch mixture.
            alpha_init_bias=-3.0,           # per-head bias init for alpha modes; -3.0 → α≈0.047 at init (local-dominant).
            alpha_detach_denom_delta=True,  # detach log(Z_c)-log(Z_i) in the alpha logit so the gate's gradient flows only through the explicit base.
            alpha_detach_state_summary=True,  # detach the prior-state summary used by alpha_a2/a3 so the gate doesn't differentiate through prior blocks' k_phi.
            # ----- fast alpha hybrid path (only consulted when hybrid_gate_type startswith "alpha_") -----
            use_fast_alpha_hybrid=False,    # master switch; off by default for full backward compat.
            hybrid_local_impl="torch",      # "torch" | "flash"
            hybrid_cross_impl="torch",      # "torch" | "torch_cublas" (alias) | "triton" (reserved, raises)
            alpha_combine_impl="torch",     # "torch" | "triton" (fused Triton combine for the largest single speedup)
            fallback_to_torch_on_error=True,
            match_torch_centering=True,     # only consulted when hybrid_local_impl="flash"
            # ``fast_alpha_maxperf=True`` LOCKS the maxperf config (Flash + Triton-combine,
            # no fallback, no centering, no alpha return).  Overrides the per-knob flags
            # above; pre-flights FA + Triton availability and raises on any mismatch.
            fast_alpha_maxperf=False,
            # ----- mixed-layer attention (some softmax + some linear/hybrid layers) -----
            softmax_layer_indices=None,         # explicit list of 0-indexed softmax layers
            softmax_layer_interval=None,        # shorthand: every Nth layer is softmax
            softmax_layer_local_attn_size=-1,   # forwarded to CausalWanSelfAttention
            softmax_layer_sink_size=0,          # forwarded to CausalWanSelfAttention
    ):
        super().__init__()

        model_dir = str(model_path(model_name)) + "/"
        if is_causal:
            if attention_type == "softmax":
                self.model = CausalWanModel.from_pretrained(
                    model_dir, local_attn_size=local_attn_size, sink_size=sink_size)
            else:
                raise ValueError(f"unknown attention_type {attention_type!r}")
        else:
            self.model = WanModel.from_pretrained(model_dir)
        self.attention_type = attention_type
        self.is_causal = bool(is_causal)
        self.model.eval()

        # For non-causal diffusion, all frames share the same timestep
        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 = 32760  # [1, 21, 16, 60, 104]
        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:
        # NOTE: This is hard coded for WAN2.1-T2V-1.3B for now!!!!!!!!!!!!!!!!!!!!
        self._cls_pred_branch = nn.Sequential(
            # Input: [B, 384, 21, 60, 104]
            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)
        # self.has_cls_branch = True

    def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
        """
        Convert flow matching's prediction to x0 prediction.
        flow_pred: the prediction with shape [B, C, H, W]
        xt: the input noisy data with shape [B, C, H, W]
        timestep: the timestep with shape [B]

        pred = noise - x0
        x_t = (1-sigma_t) * x0 + sigma_t * noise
        we have x0 = x_t - sigma_t * pred

        """
        # use higher precision for calculations
        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:
        """
        Convert x0 prediction to flow matching's prediction.
        x0_pred: the x0 prediction with shape [B, C, H, W]
        xt: the input noisy data with shape [B, C, H, W]
        timestep: the timestep with shape [B]

        pred = (x_t - x_0) / sigma_t
        """
        # use higher precision for calculations
        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,
        return_kv_features: bool = False,
        kv_feature_layers: Optional[List[int]] = None,
        kv_token_indices: Optional[torch.Tensor] = None,
        return_latent_features: bool = False,
        latent_feature_layer: Optional[int] = None,
        latent_token_indices: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        prompt_embeds = conditional_dict["prompt_embeds"]

        # [B, F] -> [B]
        if self.uniform_timestep:
            input_timestep = timestep[:, 0]
        else:
            input_timestep = timestep

        logits = None
        if return_kv_features or return_latent_features:
            if return_kv_features and return_latent_features:
                raise ValueError(
                    "KV and post-block latent feature extraction are mutually exclusive")
            if return_kv_features:
                if kv_feature_layers is None:
                    raise ValueError(
                        "return_kv_features requires kv_feature_layers")
                if latent_feature_layer is not None or latent_token_indices is not None:
                    raise ValueError(
                        "latent feature arguments cannot be used in KV feature mode")
            else:
                if latent_feature_layer is None:
                    raise ValueError(
                        "return_latent_features requires latent_feature_layer")
                if kv_feature_layers is not None or kv_token_indices is not None:
                    raise ValueError(
                        "KV feature arguments cannot be used in latent feature mode")
            if self.is_causal:
                raise ValueError(
                    "Wan feature extraction requires the frozen non-causal Wan encoder")
            if getattr(self.model, "model_type", "t2v") != "t2v":
                raise ValueError("Wan feature extraction currently supports T2V only")
            if kv_cache is not None or crossattn_cache is not None or clean_x is not None:
                raise ValueError(
                    "Wan feature extraction is mutually exclusive with "
                    "generation caches/teacher forcing")
            if noisy_image_or_video.ndim != 5:
                raise ValueError(
                    "Wan feature input must have shape [batch, frames, channels, height, width]")
            patch = tuple(int(v) for v in self.model.patch_size)
            _, frames, _, height, width = noisy_image_or_video.shape
            feature_seq_len = (
                ((frames - patch[0]) // patch[0] + 1)
                * ((height - patch[1]) // patch[1] + 1)
                * ((width - patch[2]) // patch[2] + 1)
            )
            return self.model(
                noisy_image_or_video.permute(0, 2, 1, 3, 4),
                t=input_timestep,
                context=prompt_embeds,
                seq_len=feature_seq_len,
                kv_feature_layers=kv_feature_layers,
                kv_token_indices=kv_token_indices,
                latent_feature_layer=latent_feature_layer,
                latent_token_indices=latent_token_indices,
            )
        if latent_feature_layer is not None or latent_token_indices is not None:
            raise ValueError(
                "latent feature arguments require return_latent_features=True")

        # X0 prediction
        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
            ).permute(0, 2, 1, 3, 4)
        else:
            if clean_x is not None:
                # teacher forcing
                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,
                ).permute(0, 2, 1, 3, 4)
            else:
                if 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
                    )
                    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
                    ).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:
        """
        Update the current scheduler with the interface's static method
        """
        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):
        """
        A few custom initialization steps that should be called after the object is created.
        Currently, the only one we have is to bind a few methods to scheduler.
        We can gradually add more methods here if needed.
        """
        self.get_scheduler()