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
91 lines
3.6 KiB
Python
91 lines
3.6 KiB
Python
from torchvision.transforms import Compose
|
|
from transformers import AutoTokenizer, AutoImageProcessor
|
|
|
|
from torchvision import transforms
|
|
from torchvision.transforms import Lambda
|
|
|
|
from opensora.dataset.t2v_datasets import T2V_dataset
|
|
from opensora.dataset.inpaint_datasets import Inpaint_dataset
|
|
from opensora.dataset.transform import ToTensorVideo, TemporalRandomCrop, RandomHorizontalFlipVideo, CenterCropResizeVideo, LongSideResizeVideo, SpatialStrideCropVideo
|
|
from opensora.models.causalvideovae import ae_norm, ae_denorm
|
|
|
|
def getdataset(args):
|
|
temporal_sample = TemporalRandomCrop(args.num_frames) # 16 x
|
|
norm_fun = ae_norm[args.ae]
|
|
resize_topcrop = [CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True), ]
|
|
resize = [CenterCropResizeVideo((args.max_height, args.max_width)), ]
|
|
transform = transforms.Compose([
|
|
ToTensorVideo(),
|
|
*resize,
|
|
# RandomHorizontalFlipVideo(p=0.5), # in case their caption have position decription
|
|
norm_fun
|
|
])
|
|
transform_topcrop = transforms.Compose([
|
|
ToTensorVideo(),
|
|
*resize_topcrop,
|
|
# RandomHorizontalFlipVideo(p=0.5), # in case their caption have position decription
|
|
norm_fun
|
|
])
|
|
# tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
|
|
tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
|
|
if args.dataset == 't2v':
|
|
return T2V_dataset(args, transform=transform, temporal_sample=temporal_sample, tokenizer=tokenizer,
|
|
transform_topcrop=transform_topcrop, video_length_tolerance_range=args.video_length_tolerance_range)
|
|
elif args.dataset == 'inpaint' or args.dataset == 'i2v':
|
|
return Inpaint_dataset(args, transform=transform, temporal_sample=temporal_sample, tokenizer=tokenizer,
|
|
transform_topcrop=transform_topcrop)
|
|
|
|
raise NotImplementedError(args.dataset)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from accelerate import Accelerator
|
|
from opensora.dataset.t2v_datasets import dataset_prog
|
|
import random
|
|
from tqdm import tqdm
|
|
args = type('args', (),
|
|
{
|
|
'ae': 'CausalVAEModel_4x8x8',
|
|
'dataset': 't2v',
|
|
'attention_mode': 'xformers',
|
|
'use_rope': True,
|
|
'model_max_length': 300,
|
|
'max_height': 320,
|
|
'max_width': 240,
|
|
'num_frames': 1,
|
|
'use_image_num': 0,
|
|
'interpolation_scale_t': 1,
|
|
'interpolation_scale_h': 1,
|
|
'interpolation_scale_w': 1,
|
|
'cache_dir': '../cache_dir',
|
|
'image_data': '/storage/ongoing/new/Open-Sora-Plan-bak/7.14bak/scripts/train_data/image_data.txt',
|
|
'video_data': '1',
|
|
'train_fps': 24,
|
|
'drop_short_ratio': 1.0,
|
|
'use_img_from_vid': False,
|
|
'speed_factor': 1.0,
|
|
'cfg': 0.1,
|
|
'text_encoder_name': 'google/mt5-xxl',
|
|
'dataloader_num_workers': 10,
|
|
|
|
}
|
|
)
|
|
accelerator = Accelerator()
|
|
dataset = getdataset(args)
|
|
num = len(dataset_prog.img_cap_list)
|
|
zero = 0
|
|
for idx in tqdm(range(num)):
|
|
image_data = dataset_prog.img_cap_list[idx]
|
|
caps = [i['cap'] if isinstance(i['cap'], list) else [i['cap']] for i in image_data]
|
|
try:
|
|
caps = [[random.choice(i)] for i in caps]
|
|
except Exception as e:
|
|
print(e)
|
|
# import ipdb;ipdb.set_trace()
|
|
print(image_data)
|
|
zero += 1
|
|
continue
|
|
assert caps[0] is not None and len(caps[0]) > 0
|
|
print(num, zero)
|
|
import ipdb;ipdb.set_trace()
|
|
print('end') |