File size: 6,930 Bytes
f2c0505 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | 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])]
|