correct ldm setups

caradryanl · May 13, 2024 · 8b286d8 · 8b286d8
1 parent a1ef156
commit 8b286d8
Showing 1 changed file with 19 additions and 9 deletions.
diff --git a/diffusers/stable_copyright/data_utils.py b/diffusers/stable_copyright/data_utils.py
@@ -177,20 +177,30 @@ def __getitem__(self, index: int):
 
 
 def load_dataset(dataset_root, ckpt_path, dataset: str='laion-aesthetic-2-5k', batch_size: int=6, model_type='sd'):
-    resolution = 512
-    transform = transforms.Compose(
-        [
-            transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR),
-            transforms.CenterCrop(resolution),
-            transforms.ToTensor(),
-            transforms.Normalize([0.5], [0.5]),
-        ]
-    )
+
     if model_type != 'ldm':
+        resolution = 512
+        transform = transforms.Compose(
+            [
+                transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR),
+                transforms.CenterCrop(resolution),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+        )
         tokenizer = CLIPTokenizer.from_pretrained(
             ckpt_path, subfolder="tokenizer", revision=None
         )
     else:
+        resolution = 256
+        transform = transforms.Compose(
+            [
+                transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR),
+                # transforms.CenterCrop(resolution),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+        )
         tokenizer = None
 
     train_dataset = Dataset(