Update modelling_magiv2.py

modelling_magiv2.py

@@ -29,13 +29,15 @@ class Magiv2Model(PreTrainedModel):
     def move_to_device(self, input):
         return move_to_device(input, self.device)
     
-    def forward(self, images, move_to_device_fn=None, mask_ratio=0.0, batch_size=256):
+    def forward(self, images, move_to_device_fn=None, mask_ratio=0.0, batch_size=256, convert_to_grayscale=True):
         move_to_device_fn = self.move_to_device if move_to_device_fn is None else move_to_device_fn
         if len(images) == 0:
             return move_to_device_fn(torch.zeros(len(images), self.config.crop_embedding_model_config.hidden_size))
 
         assert all(isinstance(image, PIL.Image.Image) for image in images), "please provide a list of PIL images"
-        images = [np.array(image.convert("L").convert("RGB")) for image in images]
+        if convert_to_grayscale:
+            images = [x.convert("L") for x in images]
+        images = [np.array(image.convert("RGB")) for image in images]
         images = self.processor(images, return_tensors="pt").pixel_values
         images = move_to_device_fn(images)