utils/wan_wrapper.py

import inspect
import json
import os
import types
from typing import List, Optional

import torch
from einops import rearrange
from safetensors.torch import load_file as load_safetensors
from torch import nn

from utils.scheduler import FlowMatchScheduler, SchedulerInterface
from wan.modules.causal_model import CausalWanModel
from wan.modules.clip import CLIPModel
from wan.modules.model import WanModel
from wan.modules.vae import _video_vae


class WanTextEncoder(nn.Module):
    """Stage1-only stub kept for import compatibility."""

    def __init__(self) -> None:
        super().__init__()

    def forward(self, text_prompts: List[str]) -> dict:
        return {}


class WanVAEWrapper(nn.Module):
    def __init__(
        self,
        vae_path: str = "ckpts/MG2-base/Wan2.1_VAE.pth",
        clip_checkpoint_path: str = "ckpts/MG2-base/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",
        clip_tokenizer_path: str = "ckpts/MG2-base/xlm-roberta-large",
    ):
        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)

        self.model = _video_vae(
            pretrained_path=vae_path,
            z_dim=16,
        ).eval().requires_grad_(False)

        self.clip = CLIPModel(
            dtype=torch.float32,
            device=torch.device("cpu"),
            checkpoint_path=clip_checkpoint_path,
            tokenizer_path=clip_tokenizer_path,
        )

    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:
        target_dtype = next(self.model.parameters()).dtype
        latent = latent.to(dtype=target_dtype)
        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),
        ]

        decode_function = self.model.cached_decode if use_cache else 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 encode_visual_context_from_pixels(self, pixel: torch.Tensor) -> torch.Tensor:
        if pixel.ndim != 5:
            raise ValueError(f"Expected pixel shape [B, F, C, H, W], got {tuple(pixel.shape)}")

        first_frame = pixel[:, 0].float()
        target_device = first_frame.device
        self.clip.model.to(target_device)
        videos = [img[:, None, :, :] for img in first_frame]
        with torch.no_grad():
            visual_context = self.clip.visual(videos)
        return visual_context

    def encode_visual_context_from_latent(self, latent: torch.Tensor) -> torch.Tensor:
        pixels = self.decode_to_pixel(latent)
        return self.encode_visual_context_from_pixels(pixels)


class WanDiffusionWrapper(nn.Module):
    @staticmethod
    def _resolve_model_path(model_name: str) -> str:
        # Accept absolute paths, relative paths (./ckpts/...), or bare names (legacy)
        if os.path.isabs(model_name) or os.path.exists(model_name):
            return model_name
        # Legacy fallback: bare model name → look in wan_models/
        fallback = f"wan_models/{model_name}/"
        if os.path.exists(fallback):
            return fallback
        # Return as-is and let downstream raise a clear error
        return model_name

    @staticmethod
    def _find_first_existing(directory: str, candidates: list[str]) -> Optional[str]:
        for candidate in candidates:
            path = os.path.join(directory, candidate)
            if os.path.exists(path):
                return path
        return None

    @staticmethod
    def _load_local_state_dict(weights_path: str) -> dict:
        if weights_path.endswith(".safetensors"):
            return load_safetensors(weights_path, device="cpu")

        state_dict = torch.load(weights_path, map_location="cpu")
        if isinstance(state_dict, dict):
            for key in ("state_dict", "model", "generator", "generator_ema"):
                nested = state_dict.get(key)
                if isinstance(nested, dict):
                    return nested
        if not isinstance(state_dict, dict):
            raise TypeError(f"Unsupported checkpoint format at {weights_path}")
        return state_dict

    @staticmethod
    def _build_model_from_config(model_cls, config_dict: dict, overrides: dict):
        valid_keys = set(inspect.signature(model_cls.__init__).parameters) - {"self"}
        init_kwargs = {
            key: value
            for key, value in config_dict.items()
            if not key.startswith("_") and key in valid_keys
        }
        init_kwargs.update({key: value for key, value in overrides.items() if key in valid_keys})
        return model_cls(**init_kwargs)

    @classmethod
    def _load_local_model(
        cls,
        model_path: str,
        is_causal: bool,
        local_attn_size: int,
        sink_size: int,
        action_config: Optional[dict],
    ):
        if not os.path.isdir(model_path):
            return None

        config_path = cls._find_first_existing(model_path, ["config.json", "base_config.json"])
        weights_path = cls._find_first_existing(
            model_path,
            [
                "diffusion_pytorch_model.safetensors",
                "base_distill.safetensors",
                "diffusion_pytorch_model.bin",
                "pytorch_model.bin",
            ],
        )
        if config_path is None or weights_path is None:
            return None

        with open(config_path, encoding="utf-8") as f:
            config_dict = json.load(f)

        model_cls = CausalWanModel if is_causal else WanModel
        overrides = {}
        if is_causal:
            overrides["local_attn_size"] = local_attn_size
            overrides["sink_size"] = sink_size
            overrides["action_config"] = {} if action_config is None else dict(action_config)
        else:
            # Bid `WanModel` also accepts action_config for I2V w/ MG2-style
            # mouse+keyboard injection (see wan/modules/model.py:283).  We
            # need to pass it so that ActionModule is constructed and the
            # base ckpt's .action_model.* weights load correctly.  Without
            # this, the bid model would be built without ActionModule and
            # all action_model.* keys would land in `unexpected`, causing
            # _load_local_model to fail.
            overrides["action_config"] = {} if action_config is None else dict(action_config)

        model = cls._build_model_from_config(model_cls, config_dict, overrides)
        state_dict = cls._load_local_state_dict(weights_path)
        candidates = [state_dict]
        if any(key.startswith("model.") for key in state_dict):
            candidates.append({
                key[len("model."):] if key.startswith("model.") else key: value
                for key, value in state_dict.items()
            })

        # When action is ablated (action_config is None / empty), the base ckpt
        # still contains .action_model.* weights. Those are legitimately unused
        # and should not be treated as a load failure.
        action_disabled = not overrides.get("action_config")
        # When action is enabled on a SUBSET of blocks (e.g. MG2 distilled
        # recipe: blocks=[0..14] on a 30-layer DiT), the ckpt may have
        # action_model.* weights for blocks that our model doesn't instantiate.
        # Those are also legitimately unused and should not fail the load.
        action_blocks_subset = None
        if not action_disabled:
            ac_blocks = overrides.get("action_config", {}).get("blocks", None)
            num_layers = config_dict.get("num_layers", 30)
            if ac_blocks is not None and len(ac_blocks) < num_layers:
                action_blocks_subset = set(int(b) for b in ac_blocks)

        best_missing = None
        best_unexpected = None
        model_state = model.state_dict()
        for candidate in candidates:
            # Shape-mismatch filter.  When the action interface is re-shaped
            # relative to the checkpoint (e.g. the CrossFPS adaptation widens
            # mouse_dim_in from 2 to 4, so every block's
            # action_model.mouse_mlp.0.weight goes 1560 -> 1584 input units),
            # the corresponding tensors cannot be copied.  We drop them here
            # and let them keep their freshly initialized values; the trainer
            # is responsible for seeding them properly (see
            # `graft_mouse_mlp_weight` in trainer/diffusion.py).  Without this
            # the load would raise and no re-shaped build could ever start.
            filtered = {}
            skipped_shape = []
            for k, v in candidate.items():
                if k in model_state and hasattr(v, "shape") and v.shape != model_state[k].shape:
                    skipped_shape.append((k, tuple(v.shape), tuple(model_state[k].shape)))
                    continue
                filtered[k] = v
            if skipped_shape:
                for k, ckpt_shape, model_shape in skipped_shape[:3]:
                    print(f"[wan_load] skipped shape-mismatch: {k} "
                          f"ckpt={ckpt_shape} model={model_shape}")
                if len(skipped_shape) > 3:
                    print(f"[wan_load] ... and {len(skipped_shape) - 3} more shape "
                          f"mismatches (all left at their initialized values)")

            missing, unexpected = model.load_state_dict(filtered, strict=False)
            if skipped_shape:
                # A deliberately skipped tensor shows up in `missing`; that is
                # expected, not a load failure.
                skipped_names = {k for k, _, _ in skipped_shape}
                missing = [k for k in missing if k not in skipped_names]
            if action_disabled:
                unexpected = [k for k in unexpected if ".action_model." not in k]
            elif action_blocks_subset is not None:
                import re
                def _in_subset(k):
                    m = re.match(r"blocks\.(\d+)\.action_model\.", k)
                    if m is None:
                        return True  # not an action_model key, keep as-is
                    return int(m.group(1)) in action_blocks_subset
                unexpected = [k for k in unexpected if _in_subset(k)]
            if not missing and not unexpected:
                return model
            if best_missing is None or len(missing) + len(unexpected) < len(best_missing) + len(best_unexpected):
                best_missing, best_unexpected = missing, unexpected

        raise RuntimeError(
            f"Failed to load local model from {model_path}. "
            f"Missing keys: {best_missing[:10]}, unexpected keys: {best_unexpected[:10]}"
        )

    def __init__(
        self,
        model_name="Wan2.1-T2V-1.3B",
        timestep_shift=5.0,
        is_causal=False,
        local_attn_size=-1,
        sink_size=0,
        action_config=None,
    ):
        super().__init__()

        model_path = self._resolve_model_path(model_name)
        self.model = self._load_local_model(
            model_path=model_path,
            is_causal=is_causal,
            local_attn_size=local_attn_size,
            sink_size=sink_size,
            action_config=action_config,
        )
        if self.model is None:
            if is_causal:
                self.model = CausalWanModel.from_pretrained(
                    model_path,
                    local_attn_size=local_attn_size,
                    sink_size=sink_size,
                    action_config={} if action_config is None else dict(action_config),
                    low_cpu_mem_usage=False,
                )
            else:
                self.model = WanModel.from_pretrained(
                    model_path,
                    action_config={} if action_config is None else dict(action_config),
                    low_cpu_mem_usage=False,
                )
        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 = None
        self.post_init()

    def enable_gradient_checkpointing(self) -> None:
        try:
            self.model.enable_gradient_checkpointing()
        except TypeError:
            self.model._set_gradient_checkpointing(None, 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 _prepare_model_kwargs(self, conditional_dict: dict, device: torch.device, dtype: torch.dtype) -> dict:
        visual_context = conditional_dict["visual_context"].to(device=device, dtype=dtype)
        cond_concat = conditional_dict["cond_concat"].permute(0, 2, 1, 3, 4).to(device=device, dtype=dtype)
        mouse_condition = conditional_dict.get("mouse_condition", conditional_dict.get("mouse_cond", None))
        keyboard_condition = conditional_dict.get("keyboard_condition", conditional_dict.get("keyboard_cond", None))

        kwargs = {
            "visual_context": visual_context,
            "cond_concat": cond_concat,
        }
        if mouse_condition is not None:
            kwargs["mouse_cond"] = mouse_condition.to(device=device, dtype=dtype)
        if keyboard_condition is not None:
            kwargs["keyboard_cond"] = keyboard_condition.to(device=device, dtype=dtype)
        return kwargs

    def forward(
        self,
        noisy_image_or_video: torch.Tensor,
        conditional_dict: dict,
        timestep: torch.Tensor,
        kv_cache: Optional[List[dict]] = None,
        kv_cache_mouse: Optional[List[dict]] = None,
        kv_cache_keyboard: Optional[List[dict]] = None,
        crossattn_cache: Optional[List[dict]] = None,
        current_start: Optional[int] = None,
        cache_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,
    ) -> torch.Tensor:
        del classify_mode, concat_time_embeddings

        def normalize_latent_layout(latent: torch.Tensor, name: str) -> torch.Tensor:
            if latent.ndim != 5:
                raise ValueError(f"Expected 5D latent tensor for {name}, got {tuple(latent.shape)}")
            if latent.shape[2] == 16:
                return latent
            if latent.shape[1] == 16:
                return latent.permute(0, 2, 1, 3, 4)
            raise ValueError(f"Cannot infer latent layout for {name} from shape {tuple(latent.shape)}")

        latent_bfchw = normalize_latent_layout(noisy_image_or_video, "noisy_image_or_video")
        clean_latent_bfchw = None if clean_x is None else normalize_latent_layout(clean_x, "clean_x")

        input_timestep = timestep[:, 0] if self.uniform_timestep else timestep
        # Use the latent's own dtype (bf16 under mixed_precision) as the
        # model's effective forward dtype.  Querying self.model.dtype can
        # return fp32 on FSDP-wrapped modules because the underlying module
        # still carries fp32 master params — but forward pass sees bf16
        # shards, and our inputs (x + cond_concat) must match.
        model_dtype = latent_bfchw.dtype

        model_kwargs = self._prepare_model_kwargs(
            conditional_dict,
            device=latent_bfchw.device,
            dtype=model_dtype,
        )

        latent_input = latent_bfchw.permute(0, 2, 1, 3, 4).to(model_dtype)
        if kv_cache is not None:
            flow_pred = self.model(
                latent_input,
                t=input_timestep,
                kv_cache=kv_cache,
                kv_cache_mouse=kv_cache_mouse,
                kv_cache_keyboard=kv_cache_keyboard,
                crossattn_cache=crossattn_cache,
                current_start=current_start,
                cache_start=cache_start,
                **model_kwargs,
            ).permute(0, 2, 1, 3, 4)
        else:
            forward_kwargs = dict(model_kwargs)
            if clean_latent_bfchw is not None:
                forward_kwargs["clean_x"] = clean_latent_bfchw.permute(0, 2, 1, 3, 4).to(model_dtype)
                forward_kwargs["aug_t"] = aug_t
            flow_pred = self.model(
                latent_input,
                t=input_timestep,
                **forward_kwargs,
            ).permute(0, 2, 1, 3, 4)

        pred_x0 = self._convert_flow_pred_to_x0(
            flow_pred=flow_pred.flatten(0, 1),
            xt=latent_bfchw.flatten(0, 1),
            timestep=timestep.flatten(0, 1),
        ).unflatten(0, flow_pred.shape[:2])
        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()