update dataset

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)
 

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)
-    
+        x = torch.concatenate((first_frame_pad, x), dim=2)  # 3 + 16
+        return self.conv(x)

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)

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)

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, 
     )