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])]