runtime/orbitquant_wan_a2/diffusers_components.py

from __future__ import annotations

import gc
from pathlib import Path
import torch
import torch.nn.functional as F

CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
CLIP_STD = [0.26862954, 0.26130258, 0.27577711]


class CachedT5Adapter:
    def __init__(self, model_root: str | Path, dtype=torch.bfloat16):
        from transformers import T5TokenizerFast, UMT5EncoderModel
        root = Path(model_root)
        self.tokenizer = T5TokenizerFast.from_pretrained(str(root / 'tokenizer'), local_files_only=True)
        try:
            self.model = UMT5EncoderModel.from_pretrained(
                str(root / 'text_encoder'), dtype=dtype, local_files_only=True, low_cpu_mem_usage=True
            )
        except TypeError:
            self.model = UMT5EncoderModel.from_pretrained(
                str(root / 'text_encoder'), torch_dtype=dtype, local_files_only=True, low_cpu_mem_usage=True
            )
        self.cache = {}
        self.dtype = dtype

    @torch.no_grad()
    def precompute(self, prompts: list[str], device='cuda', max_length=512):
        self.model.to(device).eval()
        for prompt in dict.fromkeys(prompts):
            enc = self.tokenizer(
                prompt, padding='max_length', max_length=max_length, truncation=True,
                add_special_tokens=True, return_attention_mask=True, return_tensors='pt'
            )
            ids = enc.input_ids.to(device)
            mask = enc.attention_mask.to(device)
            hidden = self.model(ids, attention_mask=mask).last_hidden_state[0]
            length = int(mask[0].sum().item())
            # Official T5 wrapper returns one variable-length embedding tensor.
            self.cache[prompt] = hidden[:length].to(dtype=self.dtype).contiguous()
        self.model.to('cpu')
        torch.cuda.empty_cache()
        return self

    def release_model(self):
        if hasattr(self, 'model'):
            del self.model
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def __call__(self, prompt):
        if isinstance(prompt, (list, tuple)):
            if len(prompt) != 1:
                return [self.cache[p].clone() for p in prompt]
            prompt = prompt[0]
        if prompt not in self.cache:
            raise RuntimeError(f'T5 prompt not precomputed: {prompt!r}')
        return [self.cache[prompt]]


class DiffusersCLIPAdapter:
    def __init__(self, model_root: str | Path, dtype=torch.float16):
        from transformers import CLIPVisionModel
        root = Path(model_root)
        try:
            self.model = CLIPVisionModel.from_pretrained(
                str(root / 'image_encoder'), dtype=dtype, local_files_only=True, low_cpu_mem_usage=True
            )
        except TypeError:
            self.model = CLIPVisionModel.from_pretrained(
                str(root / 'image_encoder'), torch_dtype=dtype, local_files_only=True, low_cpu_mem_usage=True
            )
        self.dtype = dtype

    def to(self, device):
        self.model.to(device).eval()
        return self

    @torch.no_grad()
    def visual(self, videos):
        outputs = []
        device = next(self.model.parameters()).device
        for tensor in videos:
            # Official call provides C,T,H,W and uses T=1 in Animate-2.
            if tensor.ndim != 4:
                raise ValueError(f'CLIP adapter expects C,T,H,W, got {tensor.shape}')
            frames = tensor.permute(1,0,2,3).float()
            frames = F.interpolate(frames, size=(224,224), mode='bicubic', align_corners=False)
            frames = frames.mul(0.5).add(0.5)
            mean = torch.tensor(CLIP_MEAN, device=device).view(1,3,1,1)
            std = torch.tensor(CLIP_STD, device=device).view(1,3,1,1)
            frames = ((frames.to(device) - mean) / std).to(self.dtype)
            h = self.model(pixel_values=frames, output_hidden_states=True).hidden_states[-2]
            # Animate-2 calls this adapter on a single frame; preserve [B,257,1280].
            if h.shape[0] != 1:
                raise RuntimeError('Animate-2 CLIP adapter currently expects a one-frame visual input')
            outputs.append(h[0])
        return torch.stack(outputs, dim=0)


class DiffusersWanVAEAdapter:
    def __init__(self, model_root: str | Path, dtype=torch.bfloat16):
        from diffusers import AutoencoderKLWan
        root = Path(model_root)
        try:
            self.model = AutoencoderKLWan.from_pretrained(
                str(root / 'vae'), dtype=dtype, local_files_only=True, low_cpu_mem_usage=True
            )
        except TypeError:
            self.model = AutoencoderKLWan.from_pretrained(
                str(root / 'vae'), torch_dtype=dtype, local_files_only=True, low_cpu_mem_usage=True
            )
        self.dtype = dtype

    def to(self, device):
        self.model.to(device).eval()
        return self

    def _latent_stats(self, z):
        cfg = self.model.config
        channels = int(z.shape[1])
        mean_values = getattr(cfg, 'latents_mean', None)
        std_values = getattr(cfg, 'latents_std', None)
        if mean_values is None or std_values is None:
            raise RuntimeError('AutoencoderKLWan config is missing latents_mean/latents_std')
        if len(mean_values) != channels or len(std_values) != channels:
            raise RuntimeError(
                f'AutoencoderKLWan latent-stat size mismatch: channels={channels}, '
                f'mean={len(mean_values)}, std={len(std_values)}'
            )
        mean = torch.tensor(mean_values, device=z.device, dtype=z.dtype).view(1, channels, 1, 1, 1)
        std = torch.tensor(std_values, device=z.device, dtype=z.dtype).view(1, channels, 1, 1, 1)
        return mean, std

    def _standardize(self, z):
        mean, std = self._latent_stats(z)
        return (z - mean) / std

    def _destandardize(self, z):
        mean, std = self._latent_stats(z)
        return z * std + mean

    @torch.no_grad()
    def encode(self, x):
        if isinstance(x, (list, tuple)):
            x = torch.stack(list(x), dim=0)
        if x.ndim == 4:
            x = x.unsqueeze(0)
        out = self.model.encode(x.to(device=next(self.model.parameters()).device, dtype=self.dtype))
        if hasattr(out, 'latent_dist'):
            z = out.latent_dist.mode()
        elif hasattr(out, 'latents'):
            z = out.latents
        else:
            z = out[0] if isinstance(out, (list, tuple)) else out
        z = self._standardize(z)
        return [z[i] for i in range(z.shape[0])]

    @torch.no_grad()
    def decode(self, z):
        if isinstance(z, (list, tuple)):
            z = torch.stack(list(z), dim=0)
        if z.ndim == 4:
            z = z.unsqueeze(0)
        z = self._destandardize(z.to(device=next(self.model.parameters()).device, dtype=self.dtype))
        out = self.model.decode(z, return_dict=False)[0]
        return [out[i] for i in range(out.shape[0])]