with Download Buttons

app.py

@@ -4,7 +4,9 @@ import numpy as np
 import random
 import os
 import sys
+import tempfile
 
+from PIL import Image
 from diffusers.utils import load_image
 from diffusers import EulerDiscreteScheduler, T2IAdapter
 
@@ -174,7 +176,17 @@ def generate_image(
         adapter_conditioning_scale=adapter_conditioning_scale,
         adapter_conditioning_factor=adapter_conditioning_factor,
     ).images
-    return images, gr.update(visible=True)
+
+    # Save images to temp files for download
+    temp_paths = []
+    for i, img in enumerate(images):
+        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+        img.save(tmp.name)
+        temp_paths.append(tmp.name)
+
+    return images, gr.update(visible=True), temp_paths[0], temp_paths[1]
+
+    # return images, gr.update(visible=True)
 
 def swap_to_gallery(images):
     return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
@@ -391,6 +403,14 @@ background: linear-gradient(135deg, #CCFED9, #95BAFF);
     background: #553566 !important;
     color: white !important;
     }
+#download_btn1 {
+background: #553566 !important;
+color: white !important;
+}
+#download_btn2 {
+background: #553566 !important;
+color: white !important;
+}      
 html body h1#app-title {
     text-align: center;
     font-size: 40px;
@@ -471,7 +491,7 @@ with gr.Blocks(css=css) as demo:
                         label="Drag (Select) 1 or more photos of your face",
                         file_types=["image"]
                     )
-            uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200)
+            uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200, show_share_button=False, show_download_button=False)
             with gr.Column(visible=False, elem_id="reupload_wrapper") as clear_button:
                 remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
 
@@ -593,7 +613,11 @@ with gr.Blocks(css=css) as demo:
                 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
         with gr.Column():
             # gallery = gr.Gallery(label="Generated Images")
-            gallery = gr.Gallery(label="Generated Images", show_share_button=False)
+            gallery = gr.Gallery(label="Generated Images", show_share_button=False, show_download_button=False)
+            download_btn1 = gr.DownloadButton(label="Download Image 1", elem_id="download_btn1")
+            download_btn2 = gr.DownloadButton(label="Download Image 2", elem_id="download_btn2")
+
+            
             usage_tips = gr.Markdown(label="Usage tips of PhotoMaker", value=tips ,visible=False, elem_id="usage_tips_section")
 
             # usage_tips = gr.Markdown(label="Usage tips of PhotoMaker", value=tips ,visible=False)
@@ -631,7 +655,7 @@ with gr.Blocks(css=css) as demo:
         ).then(
             fn=generate_image,
             inputs=input_list,
-            outputs=[gallery, usage_tips]
+            outputs=[gallery, usage_tips, download_btn1, download_btn2]
         )
         
         # When user clicks on a suggestion, fill the prompt textbox and close popup