Update app.py

app.py

@@ -60,7 +60,18 @@ if TORCH_COMPILE:
 pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
 pipe.fuse_lora()
 
-
+def safe_filename(text):
+    """Generate a safe filename from a string."""
+    safe_text = re.sub(r'\W+', '_', text)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    return f"{safe_text}_{timestamp}.png"
+    
+def encode_image(image):
+    """Encode image to base64."""
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode()
+    
 def predict(prompt, guidance, steps, seed=1231231):
     generator = torch.manual_seed(seed)
     last_time = time.time()
@@ -93,15 +104,12 @@ 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}")
+        filename = safe_filename(prompt)
+        image.save(filename)
+        encoded_image = encode_image(image)
+        html_link = f'<a href="data:image/png;base64,{encoded_image}" download="{filename}">Download Image</a>'
+    
 
-        # 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