Peiyuan Zhang
5680039dbe
Update dependencies; Setup code for debugging Delete unused files and code Remove all npu code Update torchvision imports Update EMA model and t2v_debug.sh script Delete npu related stuff and remove inpaint module Remove compress kv Fix warning with dataset handling and model loading Update PyTorch index URLs and video length tolerance range Remove UDIT and inpaint fix typo for dataset download include pretrained open-sora Add pretrained model for OpenSoraT2V-ROPE-L Update max height and width for video processing
76 lines
2.8 KiB
Python
76 lines
2.8 KiB
Python
import torch
|
|
from torch import nn
|
|
from transformers import CLIPModel, CLIPProcessor
|
|
|
|
from opensora.utils.utils import get_precision
|
|
|
|
|
|
class T5Wrapper(nn.Module):
|
|
def __init__(self, args, **kwargs):
|
|
super(T5Wrapper, self).__init__()
|
|
self.model_name = args.text_encoder_name
|
|
if 'mt5' in self.model_name:
|
|
from transformers import MT5EncoderModel
|
|
self.text_enc = MT5EncoderModel.from_pretrained(self.model_name, cache_dir=args.cache_dir, **kwargs).eval()
|
|
|
|
elif 't5' in self.model_name:
|
|
from transformers import T5EncoderModel
|
|
self.text_enc = T5EncoderModel.from_pretrained(self.model_name, cache_dir=args.cache_dir, **kwargs).eval()
|
|
# self.text_enc = T5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/models--DeepFloyd--t5-v1_1-xxl/snapshots/c9c625d2ec93667ec579ede125fd3811d1f81d37", cache_dir=args.cache_dir, **kwargs).eval()
|
|
|
|
def forward(self, input_ids, attention_mask):
|
|
text_encoder_embs = self.text_enc(input_ids=input_ids, attention_mask=attention_mask)['last_hidden_state']
|
|
return text_encoder_embs.detach()
|
|
|
|
class CLIPWrapper(nn.Module):
|
|
def __init__(self, args):
|
|
super(CLIPWrapper, self).__init__()
|
|
self.model_name = args.text_encoder_name
|
|
dtype = get_precision(args)
|
|
model_kwargs = {'cache_dir': args.cache_dir, 'low_cpu_mem_usage': True, 'torch_dtype': dtype}
|
|
self.text_enc = CLIPModel.from_pretrained(self.model_name, **model_kwargs).eval()
|
|
|
|
def forward(self, prompts): # prompts = ['I am a test caption', 'Test twice']
|
|
# Process the prompts and move to the correct device
|
|
inputs = self.clip_processor(text=prompts, return_tensors="pt", padding=True, truncation=True).to(self.device)
|
|
|
|
# Generate text features
|
|
with torch.no_grad():
|
|
text_features = self.clip_model.get_text_features(**inputs)
|
|
|
|
return text_features.detach()
|
|
|
|
|
|
|
|
text_encoder = {
|
|
'google/mt5-xl': T5Wrapper,
|
|
'google/mt5-xxl': T5Wrapper,
|
|
'google/umt5-xl': T5Wrapper,
|
|
'google/umt5-xxl': T5Wrapper,
|
|
'DeepFloyd/t5-v1_1-xxl': T5Wrapper,
|
|
'clip': CLIPWrapper
|
|
}
|
|
|
|
|
|
def get_text_enc(args):
|
|
"""deprecation"""
|
|
encoder_key = None
|
|
for key in text_encoder.keys():
|
|
if key in args.text_encoder_name:
|
|
encoder_key = key
|
|
break
|
|
text_enc = text_encoder.get(encoder_key, None)
|
|
assert text_enc is not None
|
|
return text_enc(args)
|
|
|
|
def get_text_warpper(text_encoder_name):
|
|
"""deprecation"""
|
|
encoder_key = None
|
|
for key in text_encoder.keys():
|
|
if key in text_encoder_name:
|
|
encoder_key = key
|
|
break
|
|
text_enc = text_encoder.get(encoder_key, None)
|
|
assert text_enc is not None
|
|
return text_enc
|