update reqs

app.py

@@ -42,7 +42,7 @@ def predict(cg: ClimateGAN, api_key):
         masked_input = output_dict["masked_input"]
         wildfire = output_dict["wildfire"]
         smog = output_dict["smog"]
-        depth = np.repeat(output_dict["depth"][..., None], 3, axis=-1)
+        depth = np.repeat(output_dict["depth"], 3, axis=-1)
         segmentation = output_dict["segmentation"]
 
         climategan_flood = output_dict.get(
@@ -145,19 +145,27 @@ if __name__ == "__main__":
                 gr.HTML(
                     dedent(
                         """
+
+                    <br>
+                    <br>
+                    <br>
+
+
                     <p style='text-align: center'>
                         Visit
                         <a href='https://thisclimatedoesnotexist.com/'>
                             ThisClimateDoesNotExist.com
                         </a>
-                        for more information.
-                        |
+                        for more information
+                        &nbsp;&nbsp;|&nbsp;&nbsp;
                         Original
                         <a href='https://github.com/cc-ai/climategan'>
                             ClimateGAN GitHub Repo
                         </a>
                     </p>
 
+                    <br>
+
                     <p>
                         After you have selected an image and started the inference you
                         will see all the outputs of ClimateGAN, including intermediate
@@ -165,12 +173,18 @@ if __name__ == "__main__":
                         depth maps used to produce the 3 events.
                     </p>
 
+                    <br>
+
                     <p>
                         This Space makes use of recent Stable Diffusion in-painting
                         pipelines to replace ClimateGAN's original Painter. If you
                         select 'Both' painters, you will see a comparison
                     </p>
 
+                    <br>
+
+                    <br>
+
                     <p>
                         Read the original
                         <a
@@ -201,7 +215,7 @@ if __name__ == "__main__":
                         default="Both",
                     )
                 ]
-                btn = gr.Button("See for yourself!", label="Run")
+                btn = gr.Button("See for yourself!", label="Run", variant="primary")
         with gr.Row():
             gr.Markdown("## Outputs")
         with gr.Row():

climategan_wrapper.py

@@ -250,7 +250,10 @@ class ClimateGAN:
                 "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)),
+                "depth": np.random.randint(0, 255, (640, 640, 1)),
+                "segmentation": np.random.randint(0, 255, (640, 640, 3)),
             }
+            return
 
         image_array = (
             np.array(Image.open(orig_image))

requirements.txt

@@ -146,7 +146,8 @@ threadpoolctl==2.2.0
 tifffile==2021.8.30
 tokenizers==0.13.1
 tomli==2.0.1
-torch==1.7.1
+--extra-index-url https://download.pytorch.org/whl/cu113
+torch
 torch-optimizer==0.1.0
 torchvision==0.8.2
 tqdm==4.62.2