durgappc/infinitetalk2
0
1# Modified from transformers.models.t5.modeling_t52# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.3import logging4import math5import json6import os7 8import torch9import torch.nn as nn10import torch.nn.functional as F11 12from safetensors.torch import load_file13from optimum.quanto import quantize, freeze, qint8,requantize14 15from .tokenizers import HuggingfaceTokenizer16 17__all__ = [18 'T5Model',19 'T5Encoder',20 'T5Decoder',21 'T5EncoderModel',22]23 24 25def fp16_clamp(x):26 if x.dtype == torch.float16 and torch.isinf(x).any():27 clamp = torch.finfo(x.dtype).max - 100028 x = torch.clamp(x, min=-clamp, max=clamp)29 return x30 31 32def init_weights(m):33 if isinstance(m, T5LayerNorm):34 nn.init.ones_(m.weight)35 elif isinstance(m, T5Model):36 nn.init.normal_(m.token_embedding.weight, std=1.0)37 elif isinstance(m, T5FeedForward):38 nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)39 nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)40 nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)41 elif isinstance(m, T5Attention):42 nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5)43 nn.init.normal_(m.k.weight, std=m.dim**-0.5)44 nn.init.normal_(m.v.weight, std=m.dim**-0.5)45 nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5)46 elif isinstance(m, T5RelativeEmbedding):47 nn.init.normal_(48 m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5)49 50 51class GELU(nn.Module):52 53 def forward(self, x):54 return 0.5 * x * (1.0 + torch.tanh(55 math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))56 57 58class T5LayerNorm(nn.Module):59 60 def __init__(self, dim, eps=1e-6):61 super(T5LayerNorm, self).__init__()62 self.dim = dim63 self.eps = eps64 self.weight = nn.Parameter(torch.ones(dim))65 66 def forward(self, x):67 x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) +68 self.eps)69 if self.weight.dtype in [torch.float16, torch.bfloat16]:70 x = x.type_as(self.weight)71 return self.weight * x72 73 74class T5Attention(nn.Module):75 76 def __init__(self, dim, dim_attn, num_heads, dropout=0.1):77 assert dim_attn % num_heads == 078 super(T5Attention, self).__init__()79 self.dim = dim80 self.dim_attn = dim_attn81 self.num_heads = num_heads82 self.head_dim = dim_attn // num_heads83 84 # layers85 self.q = nn.Linear(dim, dim_attn, bias=False)86 self.k = nn.Linear(dim, dim_attn, bias=False)87 self.v = nn.Linear(dim, dim_attn, bias=False)88 self.o = nn.Linear(dim_attn, dim, bias=False)89 self.dropout = nn.Dropout(dropout)90 91 def forward(self, x, context=None, mask=None, pos_bias=None):92 """93 x: [B, L1, C].94 context: [B, L2, C] or None.95 mask: [B, L2] or [B, L1, L2] or None.96 """97 # check inputs98 context = x if context is None else context99 b, n, c = x.size(0), self.num_heads, self.head_dim100 101 # compute query, key, value102 q = self.q(x).view(b, -1, n, c)103 k = self.k(context).view(b, -1, n, c)104 v = self.v(context).view(b, -1, n, c)105 106 # attention bias107 attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))108 if pos_bias is not None:109 attn_bias += pos_bias110 if mask is not None:111 assert mask.ndim in [2, 3]112 mask = mask.view(b, 1, 1,113 -1) if mask.ndim == 2 else mask.unsqueeze(1)114 attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)115 116 # compute attention (T5 does not use scaling)117 attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias118 attn = F.softmax(attn.float(), dim=-1).type_as(attn)119 x = torch.einsum('bnij,bjnc->binc', attn, v)120 121 # output122 x = x.reshape(b, -1, n * c)123 x = self.o(x)124 x = self.dropout(x)125 return x126 127 128class T5FeedForward(nn.Module):129 130 def __init__(self, dim, dim_ffn, dropout=0.1):131 super(T5FeedForward, self).__init__()132 self.dim = dim133 self.dim_ffn = dim_ffn134 135 # layers136 self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())137 self.fc1 = nn.Linear(dim, dim_ffn, bias=False)138 self.fc2 = nn.Linear(dim_ffn, dim, bias=False)139 self.dropout = nn.Dropout(dropout)140 141 def forward(self, x):142 x = self.fc1(x) * self.gate(x)143 x = self.dropout(x)144 x = self.fc2(x)145 x = self.dropout(x)146 return x147 148 149class T5SelfAttention(nn.Module):150 151 def __init__(self,152 dim,153 dim_attn,154 dim_ffn,155 num_heads,156 num_buckets,157 shared_pos=True,158 dropout=0.1):159 super(T5SelfAttention, self).__init__()160 self.dim = dim161 self.dim_attn = dim_attn162 self.dim_ffn = dim_ffn163 self.num_heads = num_heads164 self.num_buckets = num_buckets165 self.shared_pos = shared_pos166 167 # layers168 self.norm1 = T5LayerNorm(dim)169 self.attn = T5Attention(dim, dim_attn, num_heads, dropout)170 self.norm2 = T5LayerNorm(dim)171 self.ffn = T5FeedForward(dim, dim_ffn, dropout)172 self.pos_embedding = None if shared_pos else T5RelativeEmbedding(173 num_buckets, num_heads, bidirectional=True)174 175 def forward(self, x, mask=None, pos_bias=None):176 e = pos_bias if self.shared_pos else self.pos_embedding(177 x.size(1), x.size(1))178 x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))179 x = fp16_clamp(x + self.ffn(self.norm2(x)))180 return x181 182 183class T5CrossAttention(nn.Module):184 185 def __init__(self,186 dim,187 dim_attn,188 dim_ffn,189 num_heads,190 num_buckets,191 shared_pos=True,192 dropout=0.1):193 super(T5CrossAttention, self).__init__()194 self.dim = dim195 self.dim_attn = dim_attn196 self.dim_ffn = dim_ffn197 self.num_heads = num_heads198 self.num_buckets = num_buckets199 self.shared_pos = shared_pos200 201 # layers202 self.norm1 = T5LayerNorm(dim)203 self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout)204 self.norm2 = T5LayerNorm(dim)205 self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout)206 self.norm3 = T5LayerNorm(dim)207 self.ffn = T5FeedForward(dim, dim_ffn, dropout)208 self.pos_embedding = None if shared_pos else T5RelativeEmbedding(209 num_buckets, num_heads, bidirectional=False)210 211 def forward(self,212 x,213 mask=None,214 encoder_states=None,215 encoder_mask=None,216 pos_bias=None):217 e = pos_bias if self.shared_pos else self.pos_embedding(218 x.size(1), x.size(1))219 x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e))220 x = fp16_clamp(x + self.cross_attn(221 self.norm2(x), context=encoder_states, mask=encoder_mask))222 x = fp16_clamp(x + self.ffn(self.norm3(x)))223 return x224 225 226class T5RelativeEmbedding(nn.Module):227 228 def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):229 super(T5RelativeEmbedding, self).__init__()230 self.num_buckets = num_buckets231 self.num_heads = num_heads232 self.bidirectional = bidirectional233 self.max_dist = max_dist234 235 # layers236 self.embedding = nn.Embedding(num_buckets, num_heads)237 238 def forward(self, lq, lk):239 device = self.embedding.weight.device240 # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \241 # torch.arange(lq).unsqueeze(1).to(device)242 rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \243 torch.arange(lq, device=device).unsqueeze(1)244 rel_pos = self._relative_position_bucket(rel_pos)245 rel_pos_embeds = self.embedding(rel_pos)246 rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(247 0) # [1, N, Lq, Lk]248 return rel_pos_embeds.contiguous()249 250 def _relative_position_bucket(self, rel_pos):251 # preprocess252 if self.bidirectional:253 num_buckets = self.num_buckets // 2254 rel_buckets = (rel_pos > 0).long() * num_buckets255 rel_pos = torch.abs(rel_pos)256 else:257 num_buckets = self.num_buckets258 rel_buckets = 0259 rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))260 261 # embeddings for small and large positions262 max_exact = num_buckets // 2263 rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) /264 math.log(self.max_dist / max_exact) *265 (num_buckets - max_exact)).long()266 rel_pos_large = torch.min(267 rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))268 rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)269 return rel_buckets270 271 272class T5Encoder(nn.Module):273 274 def __init__(self,275 vocab,276 dim,277 dim_attn,278 dim_ffn,279 num_heads,280 num_layers,281 num_buckets,282 shared_pos=True,283 dropout=0.1):284 super(T5Encoder, self).__init__()285 self.dim = dim286 self.dim_attn = dim_attn287 self.dim_ffn = dim_ffn288 self.num_heads = num_heads289 self.num_layers = num_layers290 self.num_buckets = num_buckets291 self.shared_pos = shared_pos292 293 # layers294 self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \295 else nn.Embedding(vocab, dim)296 self.pos_embedding = T5RelativeEmbedding(297 num_buckets, num_heads, bidirectional=True) if shared_pos else None298 self.dropout = nn.Dropout(dropout)299 self.blocks = nn.ModuleList([300 T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,301 shared_pos, dropout) for _ in range(num_layers)302 ])303 self.norm = T5LayerNorm(dim)304 305 # initialize weights306 self.apply(init_weights)307 308 def forward(self, ids, mask=None):309 x = self.token_embedding(ids)310 x = self.dropout(x)311 e = self.pos_embedding(x.size(1),312 x.size(1)) if self.shared_pos else None313 for block in self.blocks:314 x = block(x, mask, pos_bias=e)315 x = self.norm(x)316 x = self.dropout(x)317 return x318 319 320class T5Decoder(nn.Module):321 322 def __init__(self,323 vocab,324 dim,325 dim_attn,326 dim_ffn,327 num_heads,328 num_layers,329 num_buckets,330 shared_pos=True,331 dropout=0.1):332 super(T5Decoder, self).__init__()333 self.dim = dim334 self.dim_attn = dim_attn335 self.dim_ffn = dim_ffn336 self.num_heads = num_heads337 self.num_layers = num_layers338 self.num_buckets = num_buckets339 self.shared_pos = shared_pos340 341 # layers342 self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \343 else nn.Embedding(vocab, dim)344 self.pos_embedding = T5RelativeEmbedding(345 num_buckets, num_heads, bidirectional=False) if shared_pos else None346 self.dropout = nn.Dropout(dropout)347 self.blocks = nn.ModuleList([348 T5CrossAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,349 shared_pos, dropout) for _ in range(num_layers)350 ])351 self.norm = T5LayerNorm(dim)352 353 # initialize weights354 self.apply(init_weights)355 356 def forward(self, ids, mask=None, encoder_states=None, encoder_mask=None):357 b, s = ids.size()358 359 # causal mask360 if mask is None:361 mask = torch.tril(torch.ones(1, s, s).to(ids.device))362 elif mask.ndim == 2:363 mask = torch.tril(mask.unsqueeze(1).expand(-1, s, -1))364 365 # layers366 x = self.token_embedding(ids)367 x = self.dropout(x)368 e = self.pos_embedding(x.size(1),369 x.size(1)) if self.shared_pos else None370 for block in self.blocks:371 x = block(x, mask, encoder_states, encoder_mask, pos_bias=e)372 x = self.norm(x)373 x = self.dropout(x)374 return x375 376 377class T5Model(nn.Module):378 379 def __init__(self,380 vocab_size,381 dim,382 dim_attn,383 dim_ffn,384 num_heads,385 encoder_layers,386 decoder_layers,387 num_buckets,388 shared_pos=True,389 dropout=0.1):390 super(T5Model, self).__init__()391 self.vocab_size = vocab_size392 self.dim = dim393 self.dim_attn = dim_attn394 self.dim_ffn = dim_ffn395 self.num_heads = num_heads396 self.encoder_layers = encoder_layers397 self.decoder_layers = decoder_layers398 self.num_buckets = num_buckets399 400 # layers401 self.token_embedding = nn.Embedding(vocab_size, dim)402 self.encoder = T5Encoder(self.token_embedding, dim, dim_attn, dim_ffn,403 num_heads, encoder_layers, num_buckets,404 shared_pos, dropout)405 self.decoder = T5Decoder(self.token_embedding, dim, dim_attn, dim_ffn,406 num_heads, decoder_layers, num_buckets,407 shared_pos, dropout)408 self.head = nn.Linear(dim, vocab_size, bias=False)409 410 # initialize weights411 self.apply(init_weights)412 413 def forward(self, encoder_ids, encoder_mask, decoder_ids, decoder_mask):414 x = self.encoder(encoder_ids, encoder_mask)415 x = self.decoder(decoder_ids, decoder_mask, x, encoder_mask)416 x = self.head(x)417 return x418 419 420def _t5(name,421 encoder_only=False,422 decoder_only=False,423 return_tokenizer=False,424 tokenizer_kwargs={},425 dtype=torch.float32,426 device='cpu',427 **kwargs):428 # sanity check429 assert not (encoder_only and decoder_only)430 431 # params432 if encoder_only:433 model_cls = T5Encoder434 kwargs['vocab'] = kwargs.pop('vocab_size')435 kwargs['num_layers'] = kwargs.pop('encoder_layers')436 _ = kwargs.pop('decoder_layers')437 elif decoder_only:438 model_cls = T5Decoder439 kwargs['vocab'] = kwargs.pop('vocab_size')440 kwargs['num_layers'] = kwargs.pop('decoder_layers')441 _ = kwargs.pop('encoder_layers')442 else:443 model_cls = T5Model444 445 # init model446 with torch.device(device):447 model = model_cls(**kwargs)448 449 # set device450 model = model.to(dtype=dtype, device=device)451 452 # init tokenizer453 if return_tokenizer:454 from .tokenizers import HuggingfaceTokenizer455 tokenizer = HuggingfaceTokenizer(f'google/{name}', **tokenizer_kwargs)456 return model, tokenizer457 else:458 return model459 460 461def umt5_xxl(**kwargs):462 cfg = dict(463 vocab_size=256384,464 dim=4096,465 dim_attn=4096,466 dim_ffn=10240,467 num_heads=64,468 encoder_layers=24,469 decoder_layers=24,470 num_buckets=32,471 shared_pos=False,472 dropout=0.1)473 cfg.update(**kwargs)474 return _t5('umt5-xxl', **cfg)475 476 477class T5EncoderModel:478 479 def __init__(480 self,481 text_len,482 dtype=torch.bfloat16,483 device=None,484 checkpoint_path=None,485 tokenizer_path=None,486 shard_fn=None,487 quant=None,488 quant_dir=None489 ):490 assert quant is None or quant in ("int8", "fp8")491 self.text_len = text_len492 self.dtype = dtype493 # Defer CUDA device lookup to runtime (for ZeroGPU compatibility)494 self.device = device if device is not None else torch.cuda.current_device()495 self.checkpoint_path = checkpoint_path496 self.tokenizer_path = tokenizer_path497 498 # init model499 logging.info(f'loading {checkpoint_path}')500 if quant is not None:501 with torch.device('meta'):502 model = umt5_xxl(503 encoder_only=True,504 return_tokenizer=False,505 dtype=dtype,506 device=torch.device('meta'))507 logging.info(f'Loading quantized T5 from {os.path.join(quant_dir, f"t5_{quant}.safetensors")}')508 model_state_dict = load_file(os.path.join(quant_dir, f"t5_{quant}.safetensors"))509 with open(os.path.join(quant_dir, f"t5_map_{quant}.json"), "r") as f:510 quantization_map = json.load(f)511 requantize(model, model_state_dict, quantization_map, device='cpu')512 else:513 model = umt5_xxl(514 encoder_only=True,515 return_tokenizer=False,516 dtype=dtype,517 device=device).eval().requires_grad_(False)518 model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))519 self.model = model520 self.model.eval().requires_grad_(False)521 if shard_fn is not None:522 self.model = shard_fn(self.model, sync_module_states=False)523 else:524 self.model.to(self.device)525 # init tokenizer526 self.tokenizer = HuggingfaceTokenizer(527 name=tokenizer_path, seq_len=text_len, clean='whitespace')528 529 def __call__(self, texts, device):530 ids, mask = self.tokenizer(531 texts, return_mask=True, add_special_tokens=True)532 ids = ids.to(device)533 mask = mask.to(device)534 seq_lens = mask.gt(0).sum(dim=1).long()535 context = self.model(ids, mask)536 return [u[:v] for u, v in zip(context, seq_lens)]537