gradio UI fix (#2)

- gradio UI fix (a2643c3ac77c686954c806729defdeb0a6f1231b)


Co-authored-by: Morpheus <6Morpheus6@users.noreply.huggingface.co>

app.py

@@ -3,6 +3,7 @@ sys.path.append('./')
 
 
 import os 
+import gc
 import cv2
 import torch
 import random
@@ -161,7 +162,9 @@ def create_image(image_pil,
                  seed,
                  target="Load only style blocks",
                  neg_content_prompt=None,
-                 neg_content_scale=0):
+                 neg_content_scale=0,
+                 progress=gr.Progress(track_tqdm=True)
+                ):
 
     if target =="Load original IP-Adapter":
         # target_blocks=["blocks"] for original IP-Adapter
@@ -211,13 +214,27 @@ def create_image(image_pil,
                                 image=canny_map,
                                 controlnet_conditioning_scale=float(control_scale),
                                 )
-    return images
+    
+    gradio_temp_dir = os.environ['GRADIO_TEMP_DIR']
+    temp_file_path = os.path.join(gradio_temp_dir, "image.png")
+    images[0].save(temp_file_path, format="PNG")
+    print(f"Image saved in: {temp_file_path}")
+
+    return images, temp_file_path
 
 def pil_to_cv2(image_pil):
     image_np = np.array(image_pil)
     image_cv2 = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
     return image_cv2
 
+def clear_cache(device="cuda"):
+    gc.collect()
+    if device == 'mps':
+        torch.mps.empty_cache()
+    elif device == 'cuda':
+        torch.cuda.empty_cache()
+    print(f"{device} cache cleared!")
+
 # Description
 title = r"""
 <h1 align="center">InstantStyle: Free Lunch towards Style-Preserving in Text-to-Image Generation</h1>
@@ -258,53 +275,58 @@ If our work is helpful for your research or applications, please cite us via:
 If you have any questions, please feel free to open an issue or directly reach us out at <b>haofanwang.ai@gmail.com</b>.
 """
 
-block = gr.Blocks(css="footer {visibility: hidden}").queue(max_size=10, api_open=False)
+css = """
+footer { visibility: hidden }
+#row-height { height: 65px !important }
+"""
+
+block = gr.Blocks(css=css).queue(max_size=10, api_open=False)
 with block:
     
     # description
     gr.Markdown(title)
     #gr.Markdown(description)
     
-    with gr.Tabs():
-        with gr.Row():
-            with gr.Column():
+    with gr.Row(equal_height=True):
+        with gr.Column():
+            
+            with gr.Row():
+                with gr.Column():
+                    image_pil = gr.Image(label="Style Image", type='pil')
+            
+            target = gr.Radio(["Load only style blocks", "Load style+layout block", "Load original IP-Adapter"], 
+                                value="Load only style blocks",
+                                label="Style mode")
+            
+            prompt = gr.Textbox(label="Prompt",
+                                value="a cat, masterpiece, best quality, high quality")
+            
+            scale = gr.Slider(minimum=0,maximum=2.0, step=0.01,value=1.0, label="Scale")
+            
+            with gr.Accordion(open=False, label="Advanced Options"):
                 
-                with gr.Row():
-                    with gr.Column():
-                        image_pil = gr.Image(label="Style Image", type='pil')
+                with gr.Column():
+                    src_image_pil = gr.Image(label="Source Image (optional)", type='pil')
+                control_scale = gr.Slider(minimum=0,maximum=1.0, step=0.01,value=0.5, label="Controlnet conditioning scale")
                 
-                target = gr.Radio(["Load only style blocks", "Load style+layout block", "Load original IP-Adapter"], 
-                                  value="Load only style blocks",
-                                  label="Style mode")
+                n_prompt = gr.Textbox(label="Neg Prompt", value="text, watermark, lowres, low quality, worst quality, deformed, glitch, low contrast, noisy, saturation, blurry")
                 
-                prompt = gr.Textbox(label="Prompt",
-                                    value="a cat, masterpiece, best quality, high quality")
-                
-                scale = gr.Slider(minimum=0,maximum=2.0, step=0.01,value=1.0, label="Scale")
-                
-                with gr.Accordion(open=False, label="Advanced Options"):
-                    
-                    with gr.Column():
-                        src_image_pil = gr.Image(label="Source Image (optional)", type='pil')
-                    control_scale = gr.Slider(minimum=0,maximum=1.0, step=0.01,value=0.5, label="Controlnet conditioning scale")
-                    
-                    n_prompt = gr.Textbox(label="Neg Prompt", value="text, watermark, lowres, low quality, worst quality, deformed, glitch, low contrast, noisy, saturation, blurry")
-                    
-                    neg_content_prompt = gr.Textbox(label="Neg Content Prompt", value="")
-                    neg_content_scale = gr.Slider(minimum=0, maximum=1.0, step=0.01,value=0.5, label="Neg Content Scale")
-
-                    guidance_scale = gr.Slider(minimum=1,maximum=15.0, step=0.01,value=5.0, label="guidance scale")
-                    num_samples= gr.Slider(minimum=1,maximum=4.0, step=1.0,value=1.0, label="num samples")
-                    num_inference_steps = gr.Slider(minimum=5,maximum=50.0, step=1.0,value=20, label="num inference steps")
-                    seed = gr.Slider(minimum=-1000000,maximum=1000000,value=1, step=1, label="Seed Value")
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                    
-                #generate_button = gr.Button("Generate Image")
+                neg_content_prompt = gr.Textbox(label="Neg Content Prompt", value="")
+                neg_content_scale = gr.Slider(minimum=0, maximum=1.0, step=0.01,value=0.5, label="Neg Content Scale")
+
+                guidance_scale = gr.Slider(minimum=1,maximum=15.0, step=0.01,value=5.0, label="guidance scale")
+                num_samples= gr.Slider(minimum=1,maximum=4.0, step=1.0,value=1.0, label="num samples")
+                num_inference_steps = gr.Slider(minimum=5,maximum=50.0, step=1.0,value=20, label="num inference steps")
+                seed = gr.Slider(minimum=-1000000,maximum=1000000,value=1, step=1, label="Seed Value")
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                 
-            with gr.Column():
-                generated_image = gr.Gallery(label="Generated Image")
-                generate_button = gr.Button("Generate Image")
-                gr.Markdown(description)
+            #generate_button = gr.Button("Generate Image")
+            
+        with gr.Column():
+            generated_image = gr.Gallery(label="Generated Image", scale=0.3)
+            download_image = gr.File(label="Download Image", elem_id="row-height", scale=0)
+            generate_button = gr.Button("Generate Image", min_width=2000, scale=0)
+            gr.Markdown(description)
 
         generate_button.click(
             fn=randomize_seed_fn,
@@ -327,7 +349,12 @@ with block:
                     target,
                     neg_content_prompt,
                     neg_content_scale], 
-            outputs=[generated_image])
+            outputs=[generated_image, download_image]
+        ).then(
+            fn=clear_cache,
+            inputs=[],
+            outputs=None
+        )
     
     gr.Examples(
         examples=get_example(),