add prints

app.py

@@ -20,6 +20,7 @@ from gradio.components import (
     Textbox,
 )
 from skimage import io
+from datetime import datetime
 
 from climategan_wrapper import ClimateGAN
 
@@ -122,6 +123,7 @@ def toggle(radio):
 
 def predict(cg: ClimateGAN, api_key):
     def _predict(*args):
+        print(f"Starting inference ({str(datetime.now())})")
         image = place = painter = radio = None
         if api_key:
             radio, image, place, painter = args
@@ -134,7 +136,9 @@ def predict(cg: ClimateGAN, api_key):
             address = geocode_result[0]["formatted_address"]
             static_map_url = f"https://maps.googleapis.com/maps/api/streetview?size=640x640&location={parse.quote(address)}&source=outdoor&key={api_key}"
             img_np = io.imread(static_map_url)
+            print("Using GSV image")
         else:
+            print("Using user image")
             img_np = image
 
         painters = {
@@ -142,7 +146,7 @@ def predict(cg: ClimateGAN, api_key):
             "Stable Diffusion Painter": "stable_diffusion",
             "Both": "both",
         }
-
+        print("Using painter", painters[painter])
         output_dict = cg.infer_single(img_np, painters[painter], as_pil_image=True)
 
         input_image = output_dict["input"]

climategan_wrapper.py

@@ -264,7 +264,7 @@ class ClimateGAN:
         pil_image = None
         if as_pil_image:
             pil_image = Image.fromarray(image_array)
-
+        print("Preprocessing image")
         image = self._preprocess_image(image_array)
         output_dict = self.infer_preprocessed_batch(
             images=image[None, ...],
@@ -273,6 +273,7 @@ class ClimateGAN:
             concats=concats,
             pil_image=pil_image,
         )
+        print("Inference done")
         return {k: v[0] for k, v in output_dict.items()}
 
     def infer_preprocessed_batch(
@@ -340,6 +341,7 @@ class ClimateGAN:
             pil_image = Image.fromarray(((images[0] + 1) / 2 * 255).astype(np.uint8))
 
         # Retrieve numpy events as a dict {event: array[BxHxWxC]}
+        print("Inferring ClimateGAN events")
         outputs = self.trainer.infer_all(
             images,
             numpy=True,
@@ -376,8 +378,7 @@ class ClimateGAN:
                 if pil_image is None
                 else Image.fromarray(mask[0])
             )
-            print("input_mask size: ", input_mask.size)
-            print("input_images size: ", input_images.size)
+            print("Inferring stable diffusion in-painting for 50 steps")
             floods = self.sdip_pipeline(
                 prompt=[prompt] * images.shape[0],
                 image=input_images,
@@ -386,6 +387,7 @@ class ClimateGAN:
                 width=640,
                 num_inference_steps=50,
             )
+            print("Stable diffusion in-painting done")
 
             bin_mask = mask[..., None] > 0
             flood = np.stack([np.array(i) for i in floods.images])
@@ -394,6 +396,7 @@ class ClimateGAN:
             outputs["stable_copy_flood"] = copy_flood
 
         if concats:
+            print("Concatenating flood images")
             outputs["concat"] = concat_events(outputs, concats, axis=2)
 
         return {k: v.squeeze(1) if v.shape[1] == 1 else v for k, v in outputs.items()}