add download element

app.py

@@ -1,16 +1,27 @@
+import os
+import gc
+import random
+import tempfile
+import torch
+import devicetorch
 import gradio as gr
 import numpy as np
 # import spaces
-import torch
-import random
 from PIL import Image
 
 from diffusers import FluxKontextPipeline
 from diffusers.utils import load_image
+from dfloat11 import DFloat11Model
 
 MAX_SEED = np.iinfo(np.int32).max
 
-pipe = FluxKontextPipeline.from_pretrained("fuliucansheng/FLUX.1-Kontext-dev-diffusers", torch_dtype=torch.bfloat16).to("cuda")
+pipe = FluxKontextPipeline.from_pretrained("fuliucansheng/FLUX.1-Kontext-dev-diffusers", torch_dtype=torch.bfloat16)
+DFloat11Model.from_pretrained(
+    "DFloat11/FLUX.1-Kontext-dev-DF11",
+    device="cpu",
+    bfloat16_model=pipe.transformer,
+)
+pipe.enable_model_cpu_offload()
 
 # @spaces.GPU
 def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress(track_tqdm=True)):
@@ -74,11 +85,22 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
             num_inference_steps=steps,
             generator=torch.Generator().manual_seed(seed),
         ).images[0]
-    return image, seed, gr.Button(visible=True)
+
+    gradio_temp_dir = os.environ.get('GRADIO_TEMP_DIR', tempfile.gettempdir())
+    temp_file_path = os.path.join(gradio_temp_dir, "image.png")
+    image.save(temp_file_path, format="PNG")
+    print(f"Image saved in: {temp_file_path}")
+
+    gc.collect()
+    devicetorch.empty_cache(torch)
+
+    return image, temp_file_path, seed, gr.Button(visible=True)
 
 # @spaces.GPU
 def infer_example(input_image, prompt):
     image, seed, _ = infer(input_image, prompt)
+    gc.collect()
+    devicetorch.empty_cache(torch)
     return image, seed
 
 css="""
@@ -90,7 +112,11 @@ css="""
     height: 70vh; !Important
 }
 #row {
-    min-height: 35vh; !Important
+    min-height: 40vh; !Important
+}
+
+#row-height {
+    height: 65px !important
 }
 """
 
@@ -108,17 +134,22 @@ Image editing and manipulation model guidance-distilled from FLUX.1 Kontext [pro
                 result = gr.Image(label="Result", show_label=False, interactive=False, elem_classes="input-image", elem_id="row")
                 reuse_button = gr.Button("Reuse this image", visible=False)
         
-        with gr.Row():
+        with gr.Row(equal_height=True):
+            with gr.Column():
                 prompt = gr.Text(
                     label="Prompt",
-                    show_label=False,
-                    max_lines=1,
+                    show_label=True,
+                    lines=3,
+                    max_lines=3,
                     placeholder="Enter your prompt for editing (e.g., 'Remove glasses', 'Add a hat')",
-                    container=False,
-                    scale=2
+                    container=True,
+                    scale=1
                 )
+
+            with gr.Column():
+                download_image = gr.File(label="Download Image", elem_id="row-height", scale=0)
                 run_button = gr.Button("Run", scale=1)
-        
+
         with gr.Row():
             with gr.Accordion("Advanced Settings", open=False):
                 
@@ -155,7 +186,7 @@ Image editing and manipulation model guidance-distilled from FLUX.1 Kontext [pro
                 ["cat.png", "make this cat happy"]
             ],
             inputs=[input_image, prompt],
-            outputs=[result, seed],
+            outputs=[result, download_image, seed],
             fn=infer_example,
             cache_examples=False
         )
@@ -164,7 +195,7 @@ Image editing and manipulation model guidance-distilled from FLUX.1 Kontext [pro
         triggers=[run_button.click, prompt.submit],
         fn = infer,
         inputs = [input_image, prompt, seed, randomize_seed, guidance_scale, steps],
-        outputs = [result, seed, reuse_button]
+        outputs = [result, download_image, seed, reuse_button]
     )
     reuse_button.click(
         fn = lambda image: image,