Spaces:
Running on Zero
Running on Zero
File size: 2,207 Bytes
2fc2e78 619344d 2fc2e78 619344d 2fc2e78 619344d 2fc2e78 619344d 2fc2e78 619344d 2fc2e78 | 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 | import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from contextlib import nullcontext
from .quant import VectorQuantizerM
from .vqvae import AttnProjection
def is_rank0():
return (not dist.is_initialized()) or (dist.get_rank() == 0)
class UniTok(nn.Module):
def __init__(self, args):
super().__init__()
try:
embed_dim = args.unitok_embed_dim
vocab_width = args.unitok_vocab_width
vocab_size = args.unitok_vocab_size
e_temp =args.unitok_e_temp
num_codebooks = args.unitok_num_codebooks
le = args.unitok_le
except:
embed_dim = args['unitok_embed_dim']
vocab_width = args['unitok_vocab_width']
vocab_size = args['unitok_vocab_size']
e_temp =args['unitok_e_temp']
num_codebooks = args['unitok_num_codebooks']
le = args['unitok_le']
self.quant_proj = AttnProjection(embed_dim, vocab_width, embed_dim // vocab_width)
self.quantizer = VectorQuantizerM(
vocab_size=vocab_size,
vocab_width=vocab_width,
beta=0.25,
use_entropy_loss=le > 0,
entropy_temp=e_temp,
num_codebooks=num_codebooks,
)
self.post_quant_proj = AttnProjection(vocab_width, embed_dim, embed_dim // vocab_width)
self.fc_norm = nn.LayerNorm(embed_dim, eps=1e-6)
self.projection = nn.Linear(embed_dim, embed_dim)
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
self.vocab_size = vocab_size
def forward(self, input_feature):
with torch.amp.autocast(device_type="cuda", enabled=False):
img_tokens = self.quant_proj(input_feature)
img_tokens, vq_loss, entropy_loss, usages = self.quantizer(img_tokens)
img_tokens = self.post_quant_proj(img_tokens)
output_dict = {
"img_rec": img_tokens,
"vq_loss": vq_loss,
"entropy_loss": entropy_loss,
"codebook_usages": usages,
"logit_scale": self.logit_scale.exp()
}
return output_dict
|