update layout and add dev mode

app.py

@@ -6,18 +6,19 @@ import gradio as gr
 import googlemaps
 from skimage import io
 from urllib import parse
+import numpy as np
 from climategan_wrapper import ClimateGAN
 
 
-def predict(api_key):
+def predict(cg: ClimateGAN, api_key):
     def _predict(*args):
-        print("args: ", args)
-        image = place = None
-        if len(args) == 1:
+        image = place = painter = None
+        if len(args) == 2:
             image = args[0]
+            painter = args[1]
         else:
-            assert len(args) == 2, "Unknown number of inputs {}".format(len(args))
-            image, place = args
+            assert len(args) == 3, "Unknown number of inputs {}".format(len(args))
+            image, place, painter = args
 
         if api_key and place:
             geocode_result = gmaps.geocode(place)
@@ -27,8 +28,40 @@ def predict(api_key):
             img_np = io.imread(static_map_url)
         else:
             img_np = image
-        flood, wildfire, smog = model.inference(img_np)
-        return img_np, flood, wildfire, smog
+        output_dict = cg.infer_single(img_np, painter)
+
+        input_image = output_dict["input"]
+        masked_input = output_dict["masked_input"]
+        wildfire = output_dict["wildfire"]
+        smog = output_dict["smog"]
+
+        climategan_flood = output_dict.get(
+            "climategan_flood",
+            np.ones(input_image.shape) * 255,
+        )
+        stable_flood = output_dict.get(
+            "stable_flood",
+            np.ones(input_image.shape) * 255,
+        )
+        stable_copy_flood = output_dict.get(
+            "stable_copy_flood",
+            np.ones(input_image.shape) * 255,
+        )
+        concat = output_dict.get(
+            "concat",
+            np.ones(input_image.shape) * 255,
+        )
+
+        return (
+            input_image,
+            masked_input,
+            climategan_flood,
+            stable_flood,
+            stable_copy_flood,
+            concat,
+            wildfire,
+            smog,
+        )
 
     return _predict
 
@@ -40,28 +73,73 @@ if __name__ == "__main__":
     if api_key is not None:
         gmaps = googlemaps.Client(key=api_key)
 
-    model = ClimateGAN(model_path="config/model/masker")
-
-    inputs = inputs = [gr.inputs.Image(label="Input Image")]
-    if api_key:
-        inputs += [gr.inputs.Textbox(label="Address or place name")]
+    cg = ClimateGAN(model_path="config/model/masker", dev_mode=True)
+    cg._setup_stable_diffusion()
 
-    gr.Interface(
-        predict(api_key),
-        inputs=inputs,
-        outputs=[
-            gr.outputs.Image(type="numpy", label="Original image"),
-            gr.outputs.Image(type="numpy", label="Flooding"),
-            gr.outputs.Image(type="numpy", label="Wildfire"),
-            gr.outputs.Image(type="numpy", label="Smog"),
-        ],
-        title="ClimateGAN: Visualize Climate Change",
-        description='Climate change does not impact everyone equally. This Space shows the effects of the climate emergency, "one address at a time". Visit the original experience at <a href="https://thisclimatedoesnotexist.com/">ThisClimateDoesNotExist.com</a>.<br>Enter an address or place name, and ClimateGAN will generate images showing how the location could be impacted by flooding, wildfires, or smog.',  # noqa: E501
-        article="<p style='text-align: center'>This project is an unofficial clone of <a href='https://thisclimatedoesnotexist.com/'>ThisClimateDoesNotExist</a> | <a href='https://github.com/cc-ai/climategan'>ClimateGAN GitHub Repo</a></p>",  # noqa: E501
-        # examples=[
-        #     "Vancouver Art Gallery",
-        #     "Chicago Bean",
-        #     "Duomo Siracusa",
-        # ],
-        css=".footer{display:none !important}",
-    ).launch()
+    with gr.Blocks() as blocks:
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("# ClimateGAN: Visualize Climate Change")
+                gr.HTML(
+                    'Climate change does not impact everyone equally. This Space shows the effects of the climate emergency, "one address at a time". Visit the original experience at <a href="https://thisclimatedoesnotexist.com/">ThisClimateDoesNotExist.com</a>.<br>Enter an address or place name, and ClimateGAN will generate images showing how the location could be impacted by flooding, wildfires, or smog.'  # noqa: E501
+                )
+            with gr.Column():
+                gr.HTML(
+                    "<p style='text-align: center'>This project is an unofficial clone of <a href='https://thisclimatedoesnotexist.com/'>ThisClimateDoesNotExist</a> | <a href='https://github.com/cc-ai/climategan'>ClimateGAN GitHub Repo</a></p>"  # noqa: E501
+                )
+        with gr.Row():
+            gr.Markdown("## Inputs")
+        with gr.Row():
+            with gr.Column():
+                inputs = [gr.inputs.Image(label="Input Image")]
+            with gr.Column():
+                if api_key:
+                    inputs += [gr.inputs.Textbox(label="Address or place name")]
+                inputs += [
+                    gr.inputs.Dropdown(
+                        choices=[
+                            "ClimateGAN Painter",
+                            "Stable Diffusion Painter",
+                            "Both",
+                        ],
+                        label="Choose Flood Painter",
+                        default="Both",
+                    )
+                ]
+                btn = gr.Button("See for yourself!", label="Run")
+        with gr.Row():
+            gr.Markdown("## Outputs")
+        with gr.Row():
+            outputs = []
+            outputs.append(
+                gr.outputs.Image(type="numpy", label="Original image"),
+            )
+            outputs.append(
+                gr.outputs.Image(type="numpy", label="Masked input image"),
+            )
+        with gr.Row():
+            outputs.append(
+                gr.outputs.Image(type="numpy", label="ClimateGAN-Flooded image"),
+            )
+            outputs.append(
+                gr.outputs.Image(type="numpy", label="Stable Diffusion-Flooded image"),
+            )
+            outputs.append(
+                gr.outputs.Image(
+                    type="numpy",
+                    label="Stable Diffusion-Flooded image (restricted to masked area)",
+                )
+            ),
+        with gr.Row():
+            outputs.append(
+                gr.outputs.Image(type="numpy", label="Comparison of previous images"),
+            )
+        with gr.Row():
+            outputs.append(
+                gr.outputs.Image(type="numpy", label="Wildfire"),
+            )
+            outputs.append(
+                gr.outputs.Image(type="numpy", label="Smog"),
+            )
+        btn.click(predict(cg, api_key), inputs=inputs, outputs=outputs)
+    blocks.launch()

climategan_wrapper.py

@@ -115,7 +115,7 @@ def to_m1_p1(img):
 
 # No need to do any timing in this, since it's just for the HF Space
 class ClimateGAN:
-    def __init__(self, model_path) -> None:
+    def __init__(self, model_path, dev_mode=False) -> None:
         """
         A wrapper for the ClimateGAN model that you can use to generate
         events from images or folders containing images.
@@ -125,6 +125,10 @@ class ClimateGAN:
         """
         torch.set_grad_enabled(False)
         self.target_size = 640
+        self._stable_diffusion_is_setup = False
+        self.dev_mode = dev_mode
+        if self.dev_mode:
+            return
         self.trainer = Trainer.resume_from_path(
             model_path,
             setup=True,
@@ -132,7 +136,6 @@ class ClimateGAN:
             new_exp=None,
         )
         self.trainer.G.half()
-        self._stable_diffusion_is_setup = False
 
     def _setup_stable_diffusion(self):
         """
@@ -140,6 +143,9 @@ class ClimateGAN:
         Make sure you have accepted the license on the model's card
         https://huggingface.co/CompVis/stable-diffusion-v1-4
         """
+        if self.dev_mode:
+            return
+
         try:
             self.sdip_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
                 "runwayml/stable-diffusion-inpainting",
@@ -216,6 +222,18 @@ class ClimateGAN:
             dict: a dictionary containing the output images {k: HxWxC}. C is omitted
                 for masks (HxW).
         """
+        if self.dev_mode:
+            return {
+                "input": np.random.randint(0, 255, (640, 640, 3)),
+                "mask": np.random.randint(0, 255, (640, 640)),
+                "masked_input": np.random.randint(0, 255, (640, 640, 3)),
+                "climategan_flood": np.random.randint(0, 255, (640, 640, 3)),
+                "stable_flood": np.random.randint(0, 255, (640, 640, 3)),
+                "stable_copy_flood": np.random.randint(0, 255, (640, 640, 3)),
+                "concat": np.random.randint(0, 255, (640, 640 * 5, 3)),
+                "smog": np.random.randint(0, 255, (640, 640, 3)),
+                "wildfire": np.random.randint(0, 255, (640, 640, 3)),
+            }
         image_array = (
             np.array(Image.open(orig_image))
             if isinstance(orig_image, str)