FrozenBurning
single view to 3D init release
81ecb2b
import torch
import torch.nn as nn
import torch.nn.functional as F
import open_clip
from dva.io import load_from_config
class TextConditioner(nn.Module):
def __init__(
self,
encoder_config,
):
super().__init__()
self.encoder = load_from_config(encoder_config)
@torch.no_grad()
def forward(self, batch, rm, amp=False, precision_dtype=torch.float32):
assert 'caption_token' in batch, "No tokenized caption in current batch for text conditions"
caption_token = batch['caption_token']
with torch.autocast(device_type='cuda', dtype=precision_dtype, enabled=amp):
results = self.encoder(caption_token)
return results
class CLIPTextEncoder(nn.Module):
def __init__(
self,
pretrained_path: str,
model_spec: str = 'ViT-L-14',
):
super().__init__()
self.model, _, _ = open_clip.create_model_and_transforms(model_spec, pretrained=pretrained_path)
self.model.eval()
@torch.no_grad()
def forward(self, text):
text_features = self.model.encode_text(text)
text_features /= text_features.norm(dim=-1, keepdim=True)
return text_features[:, None, :]