Update app.py

app.py

@@ -2,10 +2,13 @@ import gradio as gr
 import threading
 import os
 import torch
+from PIL import Image
 
+# Set environment variables and threading
 os.environ["OMP_NUM_THREADS"] = str(os.cpu_count())
 torch.set_num_threads(os.cpu_count())
 
+# Load models
 model1 = gr.load("models/prithivMLmods/SD3.5-Turbo-Realism-2.0-LoRA")
 model2 = gr.load("models/Purz/face-projection")
 
@@ -19,46 +22,67 @@ def generate_images(text, selected_model):
     elif selected_model == "Model 2 (Face Projection)":
         model = model2
     else:
-        return ["Invalid model selection."] * 3
+        return ["Invalid model selection."] * 3, [""] * 3
 
     results = []
+    download_links = []
+
     for i in range(3):
         if stop_event.is_set():
-            return ["Image generation stopped by user."] * 3
+            return ["Image generation stopped by user."] * 3, [""] * 3
 
         modified_text = f"{text} variation {i+1}"
         result = model(modified_text)
+
+        # Save image locally
+        filename = f"generated_image_{i+1}.png"
+        result.save(filename)
+
         results.append(result)
+        download_links.append(filename)  # Store file path for download
 
-    return results
+    return results, download_links
 
 def stop_generation():
     """Stops the ongoing image generation by setting the stop_event flag."""
     stop_event.set()
-    return ["Generation stopped."] * 3
+    return ["Generation stopped."] * 3, [""] * 3
 
-with gr.Blocks() as interface:#...
-    gr.Markdown(
-        "### ⚠ Sorry for the inconvenience. The Space is currently running on the CPU, which might affect performance. We appreciate your understanding."
-    )
-    
-    text_input = gr.Textbox(label="Type here your imagination:", placeholder="Type your prompt...")
+# Gradio UI
+with gr.Blocks() as interface:
+    gr.Markdown("### Text-to-Image Generator with Download Option")
+
+    text_input = gr.Textbox(label="Enter Your Prompt:", placeholder="Describe an image...")
     model_selector = gr.Radio(
         ["Model 1 (Turbo Realism)", "Model 2 (Face Projection)"],
         label="Select Model",
         value="Model 1 (Turbo Realism)"
     )
-    
+
     with gr.Row():
         generate_button = gr.Button("Generate 3 Images 🎨")
-        stop_button = gr.Button("Stop Image Generation")
-    
+        stop_button = gr.Button("Stop Generation")
+
     with gr.Row():
         output1 = gr.Image(label="Generated Image 1")
         output2 = gr.Image(label="Generated Image 2")
         output3 = gr.Image(label="Generated Image 3")
+
+    with gr.Row():
+        download1 = gr.File(label="Download Image 1")
+        download2 = gr.File(label="Download Image 2")
+        download3 = gr.File(label="Download Image 3")
+
+    generate_button.click(
+        generate_images, 
+        inputs=[text_input, model_selector], 
+        outputs=[output1, output2, output3, download1, download2, download3]
+    )
     
-    generate_button.click(generate_images, inputs=[text_input, model_selector], outputs=[output1, output2, output3])
-    stop_button.click(stop_generation, inputs=[], outputs=[output1, output2, output3])
+    stop_button.click(
+        stop_generation, 
+        inputs=[], 
+        outputs=[output1, output2, output3, download1, download2, download3]
+    )
 
-interface.launch()
+interface.launch()