import json
import os
from pathlib import Path
import types
from typing import List, Optional

import numpy as np
import torch
import torch.distributed as dist
from safetensors import safe_open
from torch import amp
from torch import nn

from utils.scheduler import FlowMatchScheduler, SchedulerInterface
from wan.configs import WAN_CONFIGS
from wan.modules.causal_model import CausalWanModel
from wan.modules.clip import CLIPModel
from wan.modules.model import GanAttentionBlock, RegisterTokens, WanModel, rope_params
from wan.modules.t5 import umt5_xxl
from wan.modules.tokenizers import HuggingfaceTokenizer
from wan.modules.vae import _video_vae
from wan.modules.vae2_2 import _video_vae as _video_vae_2_2


class WanTextEncoder(torch.nn.Module):
    def __init__(
        self,
        model_name: str = "Wan2.1-T2V-1.3B",
        model_dir: str | os.PathLike[str] | None = None,
        official_model_dir: str | os.PathLike[str] | None = None,
        compute_dtype: torch.dtype | None = None,
        output_dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        self.model_name = model_name
        model_dir = self._resolve_shared_model_dir(model_name=model_name, model_dir=model_dir)

        if official_model_dir is None:
            checkpoint_dtype = WAN_CONFIGS[self.model_name].param_dtype
            param_dtype = checkpoint_dtype if compute_dtype is None else compute_dtype
            state_dict = torch.load(
                model_dir / "models_t5_umt5-xxl-enc-bf16.pth",
                map_location="cpu",
                mmap=True,
                weights_only=True,
            )
            tokenizer_path = model_dir / "google" / "umt5-xxl"
            default_output_dtype = checkpoint_dtype
        else:
            param_dtype = torch.float32 if compute_dtype is None else compute_dtype
            text_encoder_dir, tokenizer_path = self._resolve_official_text_paths(official_model_dir)
            state_dict = self._load_official_umt5_state_dict(text_encoder_dir, target_dtype=param_dtype)
            default_output_dtype = torch.bfloat16

        if official_model_dir is None and param_dtype != WAN_CONFIGS[self.model_name].param_dtype:
            # Upcast bf16-distributed weights for numerically safer text encoding without changing the on-disk asset.
            for key, tensor in state_dict.items():
                state_dict[key] = tensor.to(dtype=param_dtype)

        self.compute_dtype = param_dtype
        self.output_dtype = default_output_dtype if output_dtype is None else output_dtype

        self.text_encoder = (
            umt5_xxl(encoder_only=True, return_tokenizer=False, dtype=param_dtype, device=torch.device("meta"))
            .eval()
            .requires_grad_(False)
        )
        self.text_encoder.load_state_dict(state_dict, strict=True, assign=True)

        self.tokenizer = HuggingfaceTokenizer(name=str(tokenizer_path), seq_len=512, clean="whitespace")

    @staticmethod
    def _resolve_shared_model_dir(
        model_name: str,
        model_dir: str | os.PathLike[str] | None,
    ) -> Path:
        if model_dir is None:
            return Path("wan_models") / model_name
        return Path(model_dir).expanduser().resolve()

    @staticmethod
    def _resolve_official_text_paths(
        official_model_dir: str | os.PathLike[str],
    ) -> tuple[Path, Path]:
        root = Path(official_model_dir).expanduser().resolve()
        shared_text_encoder_dir = root / "text_encoder"
        shared_tokenizer_dir = root / "tokenizer"
        if shared_text_encoder_dir.is_dir() and shared_tokenizer_dir.is_dir():
            return shared_text_encoder_dir, shared_tokenizer_dir

        if root.is_dir() and root.name == "text_encoder":
            tokenizer_dir = root.parent / "tokenizer"
            if tokenizer_dir.is_dir():
                return root, tokenizer_dir

        raise FileNotFoundError(
            f"Could not resolve official text encoder/tokenizer paths from `{root}`. "
            "Expected either a shared model dir with `text_encoder/` and `tokenizer/`, "
            "or the `text_encoder/` directory itself."
        )

    @staticmethod
    def _build_official_umt5_key_map(num_layers: int) -> dict[str, str]:
        key_map = {
            "shared.weight": "token_embedding.weight",
            "encoder.final_layer_norm.weight": "norm.weight",
        }
        for index in range(num_layers):
            official_prefix = f"encoder.block.{index}"
            sf_prefix = f"blocks.{index}"
            key_map.update(
                {
                    f"{official_prefix}.layer.0.SelfAttention.q.weight": f"{sf_prefix}.attn.q.weight",
                    f"{official_prefix}.layer.0.SelfAttention.k.weight": f"{sf_prefix}.attn.k.weight",
                    f"{official_prefix}.layer.0.SelfAttention.v.weight": f"{sf_prefix}.attn.v.weight",
                    f"{official_prefix}.layer.0.SelfAttention.o.weight": f"{sf_prefix}.attn.o.weight",
                    f"{official_prefix}.layer.0.SelfAttention.relative_attention_bias.weight": (
                        f"{sf_prefix}.pos_embedding.embedding.weight"
                    ),
                    f"{official_prefix}.layer.0.layer_norm.weight": f"{sf_prefix}.norm1.weight",
                    f"{official_prefix}.layer.1.DenseReluDense.wi_0.weight": f"{sf_prefix}.ffn.gate.0.weight",
                    f"{official_prefix}.layer.1.DenseReluDense.wi_1.weight": f"{sf_prefix}.ffn.fc1.weight",
                    f"{official_prefix}.layer.1.DenseReluDense.wo.weight": f"{sf_prefix}.ffn.fc2.weight",
                    f"{official_prefix}.layer.1.layer_norm.weight": f"{sf_prefix}.norm2.weight",
                }
            )
        return key_map

    @classmethod
    def _load_official_umt5_state_dict(
        cls,
        text_encoder_dir: Path,
        target_dtype: torch.dtype,
    ) -> dict[str, torch.Tensor]:
        key_map = cls._build_official_umt5_key_map(num_layers=24)
        index_path = text_encoder_dir / "model.safetensors.index.json"

        if index_path.is_file():
            weight_map = json.loads(index_path.read_text())["weight_map"]
            keys_by_file: dict[str, list[str]] = {}
            for official_key in key_map:
                if official_key not in weight_map:
                    raise KeyError(f"Missing official text encoder weight `{official_key}` in `{index_path}`.")
                keys_by_file.setdefault(weight_map[official_key], []).append(official_key)

            state_dict: dict[str, torch.Tensor] = {}
            for relative_path, official_keys in keys_by_file.items():
                with safe_open(str(text_encoder_dir / relative_path), framework="pt", device="cpu") as tensors:
                    for official_key in official_keys:
                        tensor = tensors.get_tensor(official_key).to(dtype=target_dtype)
                        state_dict[key_map[official_key]] = tensor
            return state_dict

        safetensors_files = sorted(text_encoder_dir.glob("*.safetensors"))
        if len(safetensors_files) != 1:
            raise FileNotFoundError(
                f"Expected either `model.safetensors.index.json` or a single `.safetensors` file in `{text_encoder_dir}`."
            )

        state_dict = {}
        with safe_open(str(safetensors_files[0]), framework="pt", device="cpu") as tensors:
            for official_key, sf_key in key_map.items():
                tensor = tensors.get_tensor(official_key).to(dtype=target_dtype)
                state_dict[sf_key] = tensor
        return state_dict

    @property
    def device(self):
        return next(self.text_encoder.parameters()).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  # set padding to 0.0

        if self.output_dtype is not None and context.dtype != self.output_dtype:
            context = context.to(dtype=self.output_dtype)

        return {"prompt_embeds": context}


class WanVAEWrapper(torch.nn.Module):
    def __init__(self, model_name, model_dir: str | os.PathLike[str] | None = None):
        super().__init__()
        self.model_name = model_name
        self.dtype = torch.bfloat16

        shared_model_dir = WanTextEncoder._resolve_shared_model_dir(model_name=model_name, model_dir=model_dir)
        vae_path = shared_model_dir / WAN_CONFIGS[self.model_name].vae_checkpoint
        if "5B" in self.model_name:
            self.mean = torch.tensor(
                [
                    -0.2289,
                    -0.0052,
                    -0.1323,
                    -0.2339,
                    -0.2799,
                    0.0174,
                    0.1838,
                    0.1557,
                    -0.1382,
                    0.0542,
                    0.2813,
                    0.0891,
                    0.1570,
                    -0.0098,
                    0.0375,
                    -0.1825,
                    -0.2246,
                    -0.1207,
                    -0.0698,
                    0.5109,
                    0.2665,
                    -0.2108,
                    -0.2158,
                    0.2502,
                    -0.2055,
                    -0.0322,
                    0.1109,
                    0.1567,
                    -0.0729,
                    0.0899,
                    -0.2799,
                    -0.1230,
                    -0.0313,
                    -0.1649,
                    0.0117,
                    0.0723,
                    -0.2839,
                    -0.2083,
                    -0.0520,
                    0.3748,
                    0.0152,
                    0.1957,
                    0.1433,
                    -0.2944,
                    0.3573,
                    -0.0548,
                    -0.1681,
                    -0.0667,
                ],
                dtype=torch.float32,
            )
            self.std = torch.tensor(
                [
                    0.4765,
                    1.0364,
                    0.4514,
                    1.1677,
                    0.5313,
                    0.4990,
                    0.4818,
                    0.5013,
                    0.8158,
                    1.0344,
                    0.5894,
                    1.0901,
                    0.6885,
                    0.6165,
                    0.8454,
                    0.4978,
                    0.5759,
                    0.3523,
                    0.7135,
                    0.6804,
                    0.5833,
                    1.4146,
                    0.8986,
                    0.5659,
                    0.7069,
                    0.5338,
                    0.4889,
                    0.4917,
                    0.4069,
                    0.4999,
                    0.6866,
                    0.4093,
                    0.5709,
                    0.6065,
                    0.6415,
                    0.4944,
                    0.5726,
                    1.2042,
                    0.5458,
                    1.6887,
                    0.3971,
                    1.0600,
                    0.3943,
                    0.5537,
                    0.5444,
                    0.4089,
                    0.7468,
                    0.7744,
                ],
                dtype=torch.float32,
            )
            cfg = {
                "dim": 160,
                "z_dim": 48,
                "dim_mult": [1, 2, 4, 4],
                "num_res_blocks": 2,
                "attn_scales": [],
                "temperal_downsample": [False, True, True],
                "dropout": 0.0,
            }
            # initialize model
            self.model = _video_vae_2_2(cfg, pretrained_path=vae_path).eval().requires_grad_(False)
        else:
            self.mean = torch.tensor(
                [
                    -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,
                ],
                dtype=torch.float32,
            )
            self.std = torch.tensor(
                [
                    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,
                ],
                dtype=torch.float32,
            )
            cfg = {
                "dim": 96,
                "z_dim": 16,
                "dim_mult": [1, 2, 4, 4],
                "num_res_blocks": 2,
                "attn_scales": [],
                "temperal_downsample": [False, True, True],
                "dropout": 0.0,
            }

            # initialize model
            self.model = _video_vae(cfg, pretrained_path=vae_path).eval().requires_grad_(False)

    def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
        # pixel: [batch_size, num_channels, num_frames, height, width]
        model_dtype = next(self.model.parameters()).dtype
        device = pixel.device
        pixel = pixel.to(dtype=model_dtype)
        scale = [self.mean.to(device=device, dtype=model_dtype), 1.0 / self.std.to(device=device, dtype=model_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]
        model_dtype = next(self.model.parameters()).dtype
        zs = latent.to(dtype=model_dtype).permute(0, 2, 1, 3, 4)
        if use_cache:
            assert latent.shape[0] == 1, "Batch size must be 1 when using cache"

        device = latent.device
        scale = [self.mean.to(device=device, dtype=model_dtype), 1.0 / self.std.to(device=device, dtype=model_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,
        model_config,
        timestep_shift=8.0,
        is_causal=False,
        local_attn_size=-1,
        sink_size=0,
        use_sp=False,
        model_dir: str | os.PathLike[str] | None = None,
    ):
        super().__init__()

        # Wan specific hyperparameters
        self.wan_config = WAN_CONFIGS[model_name]
        self.vae_stride = self.wan_config.vae_stride
        self.patch_size = self.wan_config.patch_size
        self.model_name = model_name
        self.shared_model_dir = WanTextEncoder._resolve_shared_model_dir(model_name=model_name, model_dir=model_dir)
        self.use_sp = use_sp

        self.model_config = model_config
        self.sink_size = sink_size
        self.local_attn_size = local_attn_size
        _, _f, _c, _h, _w = self.model_config.image_or_video_shape
        self.frame_seq_len = _h * _w // np.prod(self.patch_size)
        self.seq_len = _f * _h * _w // np.prod(self.patch_size)

        if is_causal:
            self.max_attention_size = _f * _h * _w // np.prod(self.patch_size)
            self.model = self._build_causal_model(model_name=model_name)
        else:
            self.model = WanModel.from_pretrained(str(self.shared_model_dir))
        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.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
        see derivations https://chatgpt.com/share/67bf8589-3d04-8008-bc6e-4cf1a24e2d0e
        """
        # use higher precision for calculations
        original_dtype = flow_pred.dtype
        flow_pred, xt, sigmas, timesteps = [
            x.double().to(flow_pred.device) for x in [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 = [
            x.double().to(x0_pred.device) for x in [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,
    ) -> torch.Tensor:
        model_dtype = next(self.model.parameters()).dtype
        prompt_embeds = conditional_dict["prompt_embeds"].to(dtype=model_dtype)
        model_input = noisy_image_or_video.to(dtype=model_dtype)

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

        autocast_enabled = model_input.is_cuda and model_dtype in (torch.float16, torch.bfloat16)
        autocast_ctx = amp.autocast("cuda", dtype=model_dtype, enabled=autocast_enabled)

        # X0 prediction
        with autocast_ctx:
            if kv_cache is not None:
                flow_pred = self.model(
                    model_input.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=0 if cache_start is None else cache_start,
                ).permute(0, 2, 1, 3, 4)
            else:
                if clean_x is not None:
                    # teacher forcing
                    clean_x_model = clean_x.to(dtype=model_dtype)
                    flow_pred = self.model(
                        model_input.permute(0, 2, 1, 3, 4),
                        t=input_timestep,
                        context=prompt_embeds,
                        seq_len=self.seq_len,
                        clean_x=clean_x_model.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(
                            model_input.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(
                            model_input.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()

    def _build_causal_model(self, model_name: str) -> CausalWanModel:
        config_path = self.shared_model_dir / "config.json"
        if not config_path.is_file():
            raise FileNotFoundError(f"Missing shared Wan config: {config_path}")

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

        shared_config = WAN_CONFIGS[model_name]
        with torch.device("meta"):
            model = CausalWanModel(
                model_type=model_config.get("model_type", "t2v"),
                patch_size=tuple(getattr(shared_config, "patch_size", (1, 2, 2))),
                text_len=model_config.get("text_len", getattr(shared_config, "text_len", 512)),
                in_dim=model_config["in_dim"],
                dim=model_config["dim"],
                ffn_dim=model_config["ffn_dim"],
                freq_dim=model_config["freq_dim"],
                text_dim=model_config.get("text_dim", 4096),
                out_dim=model_config["out_dim"],
                num_heads=model_config["num_heads"],
                num_layers=model_config["num_layers"],
                sink_size=self.sink_size,
                local_attn_size=self.local_attn_size,
                max_attention_size=self.max_attention_size,
                qk_norm=getattr(shared_config, "qk_norm", True),
                cross_attn_norm=getattr(shared_config, "cross_attn_norm", True),
                eps=model_config.get("eps", getattr(shared_config, "eps", 1e-6)),
            )
        if getattr(model, "freqs", None) is not None and model.freqs.is_meta:
            d = model.dim // model.num_heads
            model.freqs = torch.cat(
                [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))],
                dim=1,
            )
        if self.use_sp:
            if not dist.is_initialized():
                raise RuntimeError("Sequence parallel requires an initialized torch.distributed process group.")
            world_size = dist.get_world_size()
            if world_size <= 1:
                raise ValueError("Sequence parallel requires WORLD_SIZE > 1.")
            if model.num_heads % world_size != 0:
                raise ValueError(
                    f"Sequence parallel requires `num_heads` ({model.num_heads}) to be divisible by WORLD_SIZE ({world_size})."
                )
            if self.frame_seq_len % world_size != 0:
                raise ValueError(
                    f"Sequence parallel requires per-frame token count ({self.frame_seq_len}) "
                    f"to be divisible by WORLD_SIZE ({world_size})."
                )
            model.use_sp = True
            for block in model.blocks:
                block.self_attn.use_sp = True
        return model


class WanCLIPEncoder(torch.nn.Module):
    def __init__(self, model_name="Wan2.1-T2V-14B", model_dir: str | os.PathLike[str] | None = None):
        super().__init__()
        self.model_name = model_name
        shared_model_dir = WanTextEncoder._resolve_shared_model_dir(model_name=model_name, model_dir=model_dir)
        self.image_encoder = CLIPModel(
            dtype=torch.float16,
            device=torch.device("cpu"),
            checkpoint_path=str(shared_model_dir / "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"),
        )

    @property
    def device(self):
        return self.image_encoder.device

    def forward(self, img):
        # img = TF.to_tensor(img).sub_(0.5).div_(0.5).cuda()
        img = img[:, None, :, :].to(self.device)
        clip_encoder_out = self.image_encoder.visual([img]).squeeze(0)
        return clip_encoder_out