Merge pull request #25 from JCBrouwer/video-folder-dataset

Add video folder dataset

Author: ExponentialML

.gitignore

@@ -1,3 +1,5 @@
+text-to-video-ms-1.7b/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

configs/video_folder.yaml

@@ -0,0 +1,38 @@
+pretrained_model_path: "./text-to-video-ms-1.7b/"
+output_dir: "./output"
+train_text_encoder: False
+
+train_data:
+  type: folder
+  path: "path/to/folder/of/videos/"
+  n_sample_frames: 16
+  width: 256
+  height: 256
+  fps: 24
+  fallback_prompt: ""  # used when a video doesn't have a corresponding .txt file with a prompt
+
+validation_data:
+  prompt: ""
+  sample_preview: True
+  num_frames: 48
+  width: 256
+  height: 256
+  num_inference_steps: 50
+  guidance_scale: 9
+
+learning_rate: 1e-5
+adam_weight_decay: 1e-2
+train_batch_size: 1
+max_train_steps: 50000
+checkpointing_steps: 5000
+validation_steps: 500
+trainable_modules:
+  - "attn1"
+  - "attn2"
+  - "attn3"
+seed: 1234
+mixed_precision: "fp16"
+use_8bit_adam: False # This seems to be incompatible at the moment. 
+gradient_checkpointing: True
+enable_xformers_memory_efficient_attention: False
+enable_torch_2_attn: True

train.py

@@ -10,6 +10,7 @@
 from typing import Dict, Optional, Tuple
 from omegaconf import OmegaConf
 
+import cv2
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
@@ -24,7 +25,7 @@
 from accelerate.logging import get_logger
 from accelerate.utils import set_seed
 
-from .models.unet_3d_condition import UNet3DConditionModel
+from models.unet_3d_condition import UNet3DConditionModel
 from diffusers.models import AutoencoderKL
 from diffusers import DPMSolverMultistepScheduler, DDPMScheduler, TextToVideoSDPipeline
 from diffusers.optimization import get_scheduler
@@ -34,7 +35,7 @@
 from diffusers.models.attention import BasicTransformerBlock
 
 from transformers import CLIPTextModel, CLIPTokenizer
-from utils.dataset import VideoDataset
+from utils.dataset import VideoDataset, VideoFolderDataset
 from einops import rearrange, repeat
 
 already_printed_unet = False
@@ -60,6 +61,14 @@ def accelerate_set_verbose(accelerator):
         transformers.utils.logging.set_verbosity_error()
         diffusers.utils.logging.set_verbosity_error()
 
+def export_to_video(video_frames, output_video_path, fps):
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    h, w, _ = video_frames[0].shape
+    video_writer = cv2.VideoWriter(output_video_path, fourcc, fps=fps, frameSize=(w, h))
+    for i in range(len(video_frames)):
+        img = cv2.cvtColor(video_frames[i], cv2.COLOR_RGB2BGR)
+        video_writer.write(img)
+
 def create_output_folders(output_dir, config):
     now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
     out_dir = os.path.join(output_dir, f"train_{now}")
@@ -306,7 +315,10 @@ def main(
     )
 
     # Get the training dataset
-    train_dataset = VideoDataset(**train_data, tokenizer=tokenizer)
+    if train_data.pop("type", "regular") == "folder":
+        train_dataset = VideoFolderDataset(**train_data, tokenizer=tokenizer)
+    else:
+        train_dataset = VideoDataset(**train_data, tokenizer=tokenizer)
 
     # DataLoaders creation:
     train_dataloader = torch.utils.data.DataLoader(
@@ -513,7 +525,7 @@ def finetune_unet(batch, train_encoder=False):
                                     num_inference_steps=validation_data.num_inference_steps,
                                     guidance_scale=validation_data.guidance_scale
                                 ).frames
-                            video_path = export_to_video(video_frames, out_file)
+                            export_to_video(video_frames, out_file, train_data.get('fps', 8))
 
                             del pipeline
                             gc.collect()

utils/dataset.py

@@ -9,6 +9,7 @@
 
 from torch.utils.data import Dataset
 from einops import rearrange
+from glob import glob
 
 class VideoDataset(Dataset):
     def __init__(
@@ -199,3 +200,66 @@ def __getitem__(self, index):
         }
 
         return example
+
+class VideoFolderDataset(Dataset):
+    def __init__(
+        self,
+        tokenizer=None,
+        width: int = 256,
+        height: int = 256,
+        n_sample_frames: int = 16,
+        fps: int = 8,
+        path: str = "./data",
+        fallback_prompt: str = "",
+        **kwargs
+    ):
+        self.tokenizer = tokenizer
+
+        self.fallback_prompt = fallback_prompt
+
+        self.video_files = glob(f"{path}/*.mp4")
+
+        self.width = width
+        self.height = height
+
+        self.n_sample_frames = n_sample_frames
+        self.fps = fps
+
+    def get_prompt_ids(self, prompt):
+        return self.tokenizer(
+            prompt,
+            truncation=True,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        ).input_ids
+
+    def __len__(self):
+        return len(self.video_files)
+
+    def __getitem__(self, index):
+        vr = decord.VideoReader(self.video_files[index], width=self.width, height=self.height)
+        native_fps = vr.get_avg_fps()
+        every_nth_frame = round(native_fps / self.fps)
+
+        effective_length = len(vr) // every_nth_frame
+
+        if effective_length < self.n_sample_frames:
+            return self.__getitem__(random.randint(0, len(self.video_files) - 1))
+
+        effective_idx = random.randint(0, effective_length - self.n_sample_frames)
+
+        idxs = every_nth_frame * np.arange(effective_idx, effective_idx + self.n_sample_frames)
+
+        video = vr.get_batch(idxs)
+        video = rearrange(video, "f h w c -> f c h w")
+
+        if os.path.exists(self.video_files[index].replace(".mp4", ".txt")):
+            with open(self.video_files[index].replace(".mp4", ".txt"), "r") as f:
+                prompt = f.read()
+        else:
+            prompt = self.fallback_prompt
+
+        prompt_ids = self.get_prompt_ids(prompt)
+
+        return {"pixel_values": (video / 127.5 - 1.0), "prompt_ids": prompt_ids[0], "text_prompt": prompt}