From e767ef3a2e5bc3d322458c2fc8396ab9e2f97a1f Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Mon, 22 Apr 2024 11:10:56 +0800 Subject: [PATCH] update dataset --- opensora/dataset/t2v_datasets.py | 14 ++++++++++---- opensora/models/ae/videobase/modules/conv.py | 5 ++--- opensora/models/diffusion/latte/modeling_latte.py | 2 +- opensora/sample/sample_t2v.py | 2 +- opensora/train/train_t2v.py | 2 +- 5 files changed, 15 insertions(+), 10 deletions(-) diff --git a/opensora/dataset/t2v_datasets.py b/opensora/dataset/t2v_datasets.py index 24e197e69..c548a453c 100644 --- a/opensora/dataset/t2v_datasets.py +++ b/opensora/dataset/t2v_datasets.py @@ -4,6 +4,7 @@ import numpy as np import torchvision from einops import rearrange from decord import VideoReader +from os.path import join as opj import torch import torchvision.transforms as transforms @@ -28,6 +29,7 @@ class T2V_dataset(Dataset): self.temporal_sample = temporal_sample self.tokenizer = tokenizer self.model_max_length = args.model_max_length + self.video_folder = args.video_folder self.v_decoder = DecordInit() with open(args.video_data_path, 'r') as f: @@ -63,8 +65,8 @@ class T2V_dataset(Dataset): # input_ids = torch.ones(1, 120).to(torch.long).squeeze(0) # cond_mask = torch.cat([torch.ones(1, 60).to(torch.long), torch.ones(1, 60).to(torch.long)], dim=1).squeeze(0) - video_path = self.vid_cap_list[idx]['path'].replace('/remote-home1/dataset/data_split_tt', '/yzwldata01/dataset/split_1024') - video = self.decord_read(video_path) + video_path = opj(self.video_folder, self.vid_cap_list[idx]['path']) + video = self.decord_read(video_path, self.vid_cap_list[idx]['frame_idx']) video = self.transform(video) # T C H W -> T C H W video = video.transpose(0, 1) # T C H W -> C T H W text = self.vid_cap_list[idx]['cap'][0] @@ -133,11 +135,15 @@ class T2V_dataset(Dataset): return video - def decord_read(self, path): + def decord_read(self, path, frame_idx=None): decord_vr = self.v_decoder(path) total_frames = len(decord_vr) # Sampling video frames - start_frame_ind, end_frame_ind = self.temporal_sample(total_frames) + if frame_idx is None: + start_frame_ind, end_frame_ind = self.temporal_sample(total_frames) + else: + start_frame_ind, end_frame_ind = frame_idx.split(':') + start_frame_ind, end_frame_ind = int(start_frame_ind), int(end_frame_ind) # assert end_frame_ind - start_frame_ind >= self.num_frames frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int) diff --git a/opensora/models/ae/videobase/modules/conv.py b/opensora/models/ae/videobase/modules/conv.py index 4159fdb0b..5a4c8ae27 100644 --- a/opensora/models/ae/videobase/modules/conv.py +++ b/opensora/models/ae/videobase/modules/conv.py @@ -94,6 +94,5 @@ class CausalConv3d(nn.Module): first_frame_pad = x[:, :, :1, :, :].repeat( (1, 1, self.time_kernel_size - 1, 1, 1) ) # b c t h w - x = torch.concatenate((first_frame_pad, x), dim=2) - return self.conv(x) - \ No newline at end of file + x = torch.concatenate((first_frame_pad, x), dim=2) # 3 + 16 + return self.conv(x) \ No newline at end of file diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py index 008ec7982..263cc1a50 100644 --- a/opensora/models/diffusion/latte/modeling_latte.py +++ b/opensora/models/diffusion/latte/modeling_latte.py @@ -145,7 +145,7 @@ class LatteT2V(ModelMixin, ConfigMixin): interpolation_scale_1d = (self.config.video_length - 1) // 16 # => 16 (= 16 Latte) has interpolation scale 1 else: interpolation_scale_1d = self.config.video_length // 16 # => 16 (= 16 Latte) has interpolation scale 1 - interpolation_scale_1d = self.config.video_length // 5 # + # interpolation_scale_1d = self.config.video_length // 5 # interpolation_scale_1d = max(interpolation_scale_1d, 1) temp_pos_embed = get_1d_sincos_pos_embed(inner_dim, video_length, interpolation_scale=interpolation_scale_1d) # 1152 hidden size self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False) diff --git a/opensora/sample/sample_t2v.py b/opensora/sample/sample_t2v.py index 270b8dd61..bfb743b65 100644 --- a/opensora/sample/sample_t2v.py +++ b/opensora/sample/sample_t2v.py @@ -41,7 +41,7 @@ def main(args): # Load model: # transformer_model = LatteT2V.from_pretrained(args.model_path, subfolder=args.version, cache_dir=args.cache_dir, torch_dtype=torch.float16).to(device) - transformer_model = LatteT2V.from_pretrained(args.model_path, torch_dtype=torch.float16).to(device) + transformer_model = LatteT2V.from_pretrained(args.model_path, low_cpu_mem_usage=False, device_map=None, torch_dtype=torch.float16).to(device) print(transformer_model.config) transformer_model.force_images = args.force_images tokenizer = T5Tokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir) diff --git a/opensora/train/train_t2v.py b/opensora/train/train_t2v.py index 108f54d92..a74c45bc4 100644 --- a/opensora/train/train_t2v.py +++ b/opensora/train/train_t2v.py @@ -204,7 +204,7 @@ def main(args): attention_type='default', video_length=video_length, attention_mode=args.attention_mode, - compress_kv_factor=args.compress_kv_factor + compress_kv_factor=args.compress_kv_factor, use_rope=args.use_rope, model_max_length=args.model_max_length, )