Fixed the dataset not converting properly the images to RGBA when using 4 channels for training on the maskgit.

Author: ZeroCool940711

muse_maskgit_pytorch/dataset.py

@@ -243,13 +243,6 @@ def __init__(
             self.caption_pair.append(captions)
 
         transform_list = [
-            T.Lambda(
-                lambda img: img.convert("RGBA")
-                if img.mode != "RGBA" and alpha_channel
-                else img
-                if img.mode == "RGB" and not alpha_channel
-                else img.convert("RGB")
-            ),
             T.Resize(image_size),
         ]
         if flip:
@@ -258,6 +251,10 @@ def __init__(
             transform_list.append(T.CenterCrop(image_size))
         if random_crop:
             transform_list.append(T.RandomCrop(image_size, pad_if_needed=True))
+        if alpha_channel:
+            transform_list.append(T.Lambda(lambda img: img.convert("RGBA") if img.mode != "RGBA" else img))
+        else:
+            transform_list.append(T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img))
         transform_list.append(T.ToTensor())
         self.transform = T.Compose(transform_list)