handle single PIL Image

app.py

@@ -35,7 +35,7 @@ def predict(cg: ClimateGAN, api_key):
             "Both": "both",
         }
 
-        output_dict = cg.infer_single(img_np, painters[painter])
+        output_dict = cg.infer_single(img_np, painters[painter], as_pil_image=True)
 
         input_image = output_dict["input"]
         masked_input = output_dict["masked_input"]

climategan_wrapper.py

@@ -192,6 +192,7 @@ class ClimateGAN:
             "stable_flood",
             "stable_copy_flood",
         ],
+        as_pil_image=False,
     ):
         """
         Infers the image with the ClimateGAN model.
@@ -245,9 +246,18 @@ class ClimateGAN:
             if isinstance(orig_image, str)
             else orig_image
         )
+
+        pil_image = None
+        if as_pil_image:
+            pil_image = Image.fromarray(image_array)
+
         image = self._preprocess_image(image_array)
         output_dict = self.infer_preprocessed_batch(
-            image[None, ...], painter, prompt, concats
+            images=image[None, ...],
+            painter=painter,
+            prompt=prompt,
+            concats=concats,
+            pil_image=pil_image,
         )
         return {k: v[0] for k, v in output_dict.items()}
 
@@ -263,6 +273,7 @@ class ClimateGAN:
             "stable_flood",
             "stable_copy_flood",
         ],
+        pil_image=None,
     ):
         """
         Infers ClimateGAN predictions on a batch of preprocessed images.
@@ -293,6 +304,8 @@ class ClimateGAN:
                 in a new `{original_stem}_concat` image written. Defaults to:
                 ["input", "masked_input", "climategan_flood", "stable_flood",
                 "stable_copy_flood"].
+            pil_image (PIL.Image, optional): The original PIL image. If provided,
+                will be used for a single inference (batch_size=1)
 
         Returns:
             dict: a dictionary containing the output images
@@ -307,6 +320,10 @@ class ClimateGAN:
         if painter == "stable_diffusion":
             ignore_event.add("flood")
 
+        if pil_image is not None:
+            print("Warning: `pil_image` has been provided, it will override `images`")
+            images = self._preprocess_image(np.array(pil_image)[None, ...])
+
         # Retrieve numpy events as a dict {event: array[BxHxWxC]}
         outputs = self.trainer.infer_all(
             images,
@@ -336,8 +353,14 @@ class ClimateGAN:
             mask = outputs["mask"].squeeze(1)
             input_images = (
                 torch.tensor(images).permute(0, 3, 1, 2).to(self.trainer.device)
+                if pil_image is None
+                else pil_image
+            )
+            input_mask = (
+                torch.tensor(mask[:, None, ...] > 0).to(self.trainer.device)
+                if pil_image is None
+                else Image.fromarray(mask[0])
             )
-            input_mask = torch.tensor(mask[:, None, ...] > 0).to(self.trainer.device)
             floods = self.sdip_pipeline(
                 prompt=[prompt] * images.shape[0],
                 image=input_images,