UI improvement

- Fixed height
- Increase width
- Download button
- Clear cache
- Code fix for follow up generations

app.py

@@ -1,6 +1,8 @@
+import os
 import gradio as gr
 import torch
 import devicetorch
+import tempfile
 from diffusers import StableDiffusionXLPipeline, StableDiffusionPipeline, LCMScheduler
 from diffusers.schedulers import TCDScheduler
 
@@ -60,35 +62,41 @@ def generate_image(
 
     pipe = pipe_sdxl if mode == "sdxl" else pipe_sd15
 
-    if loaded != (ckpt + mode):
-        pipe.load_lora_weights(
+    pipe.load_lora_weights(
             "wangfuyun/PCM_Weights", weight_name=checkpoint, subfolder=mode
         )
-        loaded = ckpt + mode
-
-        if ckpt == "LCM-Like LoRA":
-            pipe.scheduler = LCMScheduler()
-        else:
-            pipe.scheduler = TCDScheduler(
-                num_train_timesteps=1000,
-                beta_start=0.00085,
-                beta_end=0.012,
-                beta_schedule="scaled_linear",
-                timestep_spacing="trailing",
-            )
+
+
+    if ckpt == "LCM-Like LoRA":
+        pipe.scheduler = LCMScheduler()
+    else:
+        pipe.scheduler = TCDScheduler(
+            num_train_timesteps=1000,
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            timestep_spacing="trailing",
+        )
 
     results = pipe(
         prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
     )
 
+    
+    gradio_temp_dir = os.environ.get('GRADIO_TEMP_DIR', tempfile.gettempdir())
+    temp_file_path = os.path.join(gradio_temp_dir, "image.png")
+    results.images[0].save(temp_file_path, format="PNG")
+
 #    if SAFETY_CHECKER:
 #        images, has_nsfw_concepts = check_nsfw_images(results.images)
 #        if any(has_nsfw_concepts):
 #            gr.Warning("NSFW content detected.")
 #            return Image.new("RGB", (512, 512))
 #        return images[0]
-    return results.images[0]
+    return results.images[0], temp_file_path
 
+def clear_cache():
+    devicetorch.empty_cache(torch)
 
 def update_steps(ckpt):
     num_inference_steps = checkpoints[ckpt][1]
@@ -99,7 +107,14 @@ def update_steps(ckpt):
 
 css = """
 .gradio-container {
-  max-width: 60rem !important;
+  max-width: 95vw !important;
+  margin: auto !important
+}
+.img img {
+    height: 70vh !important
+}
+#row-height {
+  height: 65px !important
 }
 """
 with gr.Blocks(css=css) as demo:
@@ -108,14 +123,12 @@ with gr.Blocks(css=css) as demo:
 # Phased Consistency Model
 
 Phased Consistency Model (PCM) is an image generation technique that addresses the limitations of the Latent Consistency Model (LCM) in high-resolution and text-conditioned image generation.
-PCM outperforms LCM across various generation settings and achieves state-of-the-art results in both image and video generation.
-
-[[paper](https://huggingface.co/papers/2405.18407)] [[arXiv](https://arxiv.org/abs/2405.18407)]  [[code](https://github.com/G-U-N/Phased-Consistency-Model)] [[project page](https://g-u-n.github.io/projects/pcm)]
+PCM outperforms LCM across various generation settings and achieves state-of-the-art results in both image and video generation.     [[paper](https://huggingface.co/papers/2405.18407)] [[arXiv](https://arxiv.org/abs/2405.18407)]  [[code](https://github.com/G-U-N/Phased-Consistency-Model)] [[project page](https://g-u-n.github.io/projects/pcm)]
 """
     )
     with gr.Group():
         with gr.Row():
-            prompt = gr.Textbox(label="Prompt", scale=8)
+            prompt = gr.Textbox(label="Prompt", scale=4)
             ckpt = gr.Dropdown(
                 label="Select inference steps",
                 choices=list(checkpoints.keys()),
@@ -137,10 +150,12 @@ PCM outperforms LCM across various generation settings and achieves state-of-the
                 show_progress=False,
             )
 
+        with gr.Row():
             submit_sdxl = gr.Button("Run on SDXL", scale=1)
             submit_sd15 = gr.Button("Run on SD15", scale=1)
 
-    img = gr.Image(label="PCM Image")
+    img = gr.Image(label="PCM Image", elem_classes="img")
+    download_image = gr.File(label="Download Image", file_count="single", interactive=False, elem_id="row-height")
     gr.Examples(
         examples=[
             [" astronaut walking on the moon", "4-Step", 4],
@@ -171,7 +186,7 @@ PCM outperforms LCM across various generation settings and achieves state-of-the
             ],
         ],
         inputs=[prompt, ckpt, steps],
-        outputs=[img],
+        outputs=[img, download_image],
         fn=generate_image,
         #cache_examples="lazy",
     )
@@ -180,14 +195,22 @@ PCM outperforms LCM across various generation settings and achieves state-of-the
         fn=generate_image,
         triggers=[ckpt.change, prompt.submit, submit_sdxl.click],
         inputs=[prompt, ckpt, steps],
-        outputs=[img],
-    )
+        outputs=[img, download_image],
+    ).then(
+            fn=clear_cache,
+            inputs=[],
+            outputs=None
+        )
     gr.on(
         fn=lambda *args: generate_image(*args, mode="sd15"),
         triggers=[submit_sd15.click],
         inputs=[prompt, ckpt, steps],
-        outputs=[img],
-    )
+        outputs=[img, download_image],
+    ).then(
+            fn=clear_cache,
+            inputs=[],
+            outputs=None
+        )
 
 
 demo.queue(api_open=False).launch(show_api=False)