Merge pull request #1420 from bghira/bugfix/video-model-image-training

When running sample transforms, the dataset_type should be considered so that we do not run video transforms on image

Author: bghira

helpers/caching/vae.py

@@ -116,7 +116,7 @@ def __init__(
         self.vae_batch_size = vae_batch_size
         self.instance_data_dir = instance_data_dir
         self.model = model
-        self.transform_sample = model.get_transforms()
+        self.transform_sample = model.get_transforms(dataset_type=dataset_type)
         self.num_video_frames = None
         if self.dataset_type == "video":
             self.num_video_frames = num_video_frames
@@ -894,13 +894,19 @@ def _encode_images_in_batch(
                 count_to_process = min(qlen, self.vae_batch_size)
                 for idx in range(0, count_to_process):
                     if image_pixel_values:
-                        pixel_values, filepath, aspect_bucket, is_final_sample = (
-                            image_pixel_values.pop()
-                        )
+                        (
+                            pixel_values,
+                            filepath,
+                            aspect_bucket,
+                            is_final_sample,
+                        ) = image_pixel_values.pop()
                     else:
-                        pixel_values, filepath, aspect_bucket, is_final_sample = (
-                            self.vae_input_queue.get()
-                        )
+                        (
+                            pixel_values,
+                            filepath,
+                            aspect_bucket,
+                            is_final_sample,
+                        ) = self.vae_input_queue.get()
 
                     if batch_aspect_bucket is None:
                         batch_aspect_bucket = aspect_bucket

helpers/models/common.py

@@ -5,6 +5,7 @@
 import logging
 import inspect
 import os
+from torchvision import transforms
 from diffusers import DiffusionPipeline
 from torch.distributions import Beta
 from helpers.training.wrappers import unwrap_model
@@ -250,10 +251,16 @@ def get_flavour_choices(cls):
         """
         return list(cls.HUGGINGFACE_PATHS.keys())
 
-    def get_transforms(self):
+    def get_transforms(self, dataset_type: str = "image"):
         """
         Returns nothing, but subclasses can implement different torchvision transforms as needed.
+
+        dataset_type is passed in for models that support transforming videos or images etc.
         """
+        if dataset_type in ["video"]:
+            raise ValueError(
+                f"{dataset_type} transforms are not supported by {self.NAME}."
+            )
         return transforms.Compose(
             [
                 transforms.ToTensor(),
@@ -1328,12 +1335,10 @@ def __init__(self, config, accelerator):
         # }
         # The trainer or child class might call self._init_text_encoders() at the right time.
 
-    def get_transforms(self):
-        from torchvision import transforms
-
+    def get_transforms(self, dataset_type: str = "image"):
         return transforms.Compose(
             [
-                VideoToTensor(),
+                VideoToTensor() if dataset_type == "video" else transforms.ToTensor(),
             ]
         )