Download test

app.py

@@ -23,6 +23,7 @@ else:
 
 import gradio as gr
 from diffusers.utils import load_image
+import tempfile
 from PIL import Image
 from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
 from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref
@@ -422,7 +423,7 @@ background: linear-gradient(135deg, #CCFED9, #95BAFF);
 html, body {
     height: 100%;
     margin: 0;
-    background: linear-gradient(135deg, #CCFED9, #95BAFF);
+    background: linear-gradient(135deg, #CCFED9, #95BAFF) !important;
 }
 
 .gradio-container {
@@ -490,6 +491,17 @@ input[type="radio"]:checked + span {
         
 '''
 
+def start_tryon_with_download(img_dict, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed):
+    # Call your existing start_tryon
+    output_img, mask_img = start_tryon(img_dict, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed)
+
+    # Save output image to temp file
+    tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    output_img.save(tmp_file.name)
+
+    # Return images for display plus file path for download
+    return output_img, mask_img, tmp_file.name
+
 image_blocks = gr.Blocks(css=css).queue()
 with image_blocks as demo:
     gr.Markdown("## <span id='app-title'>Virtual Try-on</span>")
@@ -502,7 +514,6 @@ with image_blocks as demo:
                 is_checked_crop = gr.Checkbox(label="Yes", info="Use auto-crop & resizing",value=False)
 
 
-
         with gr.Column():
             garm_img = gr.Image(label="Garment", sources='upload', type="pil")
             with gr.Row(elem_id="prompt-container"):
@@ -516,8 +527,7 @@ with image_blocks as demo:
         with gr.Column():
             # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
             image_out = gr.Image(label="Output", elem_id="output-img",show_share_button=False)
-
-
+            download_btn = gr.DownloadButton(label="Download Generated Image")
 
 
     with gr.Column():
@@ -529,7 +539,13 @@ with image_blocks as demo:
 
 
 
-    try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')
+    # try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')
+    try_button.click(
+        fn=start_tryon_with_download,
+        inputs=[imgs, garm_img, prompt, is_checked, is_checked_crop, denoise_steps, seed],
+        outputs=[image_out, masked_img, download_btn],
+        api_name='tryon'
+    )
 
 image_blocks.launch(auth=(HF_TEST_USER_NAME, HF_PASS), show_error=True)
 # image_blocks.launch(show_error=True)