Update app.py

app.py

@@ -83,7 +83,6 @@ def predict(prompt, guidance, steps, seed=1231231):
     if nsfw_content_detected:
         nsfw=gr.Button("🕹️NSFW🎨", scale=1)
 
-
     # Generate file name
     date_str = datetime.datetime.now().strftime("%Y%m%d")
     safe_prompt = prompt.replace(" ", "_")[:50]  # Truncate long prompts
@@ -94,11 +93,17 @@ def predict(prompt, guidance, steps, seed=1231231):
         image_path = os.path.join("", filename)  # Specify your directory
         results.images[0].save(image_path)
         print(f"#Image saved as {image_path}")
-    
-    return results.images[0] if len(results.images) > 0 else None
 
-        
-    #return results.images[0]
+        # Create a download link
+        with open(image_path, "rb") as file:
+            btn = st.download_button(
+                label="Download 📥",
+                data=file,
+                file_name=filename,
+                mime="image/png"
+            )
+
+    return results.images[0] if len(results.images) > 0 else None
 
 
 css = """