Update app.py

app.py

@@ -3,6 +3,7 @@ import requests
 from io import BytesIO
 from PIL import Image
 import os
+import tempfile
 
 TOKEN = os.getenv("TOKEN1")
 API_URL = os.getenv("API_URL")
@@ -19,21 +20,14 @@ def get_image_from_url(url):
         response = requests.get(url, stream=True)
         response.raise_for_status()
         image = Image.open(BytesIO(response.content))
-        
-        # Convert to PNG if the image is WebP
-        if image.format == "WEBP":
-            image = image.convert("RGBA")  # Convert to RGBA to handle transparency
-            output_buffer = BytesIO()
-            image.save(output_buffer, format="PNG")
-            output_buffer.seek(0)  # Reset buffer position for reading
-            image = Image.open(output_buffer) # Re-open image from buffer
-        
-        return image
+        return image, url # Return the image and the URL
+
 
     except requests.exceptions.RequestException as e:
-        return f"Error fetching image: {e}"
+        return f"Error fetching image: {e}", None
     except Exception as e:
-        return f"Error processing image: {e}"
+        return f"Error processing image: {e}", None
+
 
 def generate_image(prompt, aspect_ratio, realism):
     global token_id
@@ -58,17 +52,34 @@ def generate_image(prompt, aspect_ratio, realism):
                     response_data = generate_image(prompt, aspect_ratio, realism)
                     tokens_tried = 0
                     return response_data
-                return "No credits available"
-            return response_data
+                return "No credits available", None
+            return response_data, None
         elif "output" in response_data:
             url = response_data['output']
-            return get_image_from_url(url)  # Return the image directly
+            image, url = get_image_from_url(url)
+            return image, url  # Return the image and the URL
         else:
-            return "Error: Unexpected response from server"
+            return "Error: Unexpected response from server", None
     except Exception as e:
-        return f"Error"
-
+        return f"Error", None
 
+def download_image(image_url):
+    if not image_url:
+         return None # Return None if image_url is empty
+    try:
+        response = requests.get(image_url, stream=True)
+        response.raise_for_status()
+        
+        # Create a temporary file
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
+            for chunk in response.iter_content(chunk_size=8192):
+                tmp_file.write(chunk)
+            temp_file_path = tmp_file.name
+        return temp_file_path
+    except Exception as e:
+         return None
+    
+# Define the Gradio interface
 interface = gr.Interface(
     fn=generate_image,
     inputs=[
@@ -80,7 +91,10 @@ interface = gr.Interface(
         ),
         gr.Checkbox(label="Realism", value=False),  # Checkbox for realism (True/False)
     ],
-    outputs=gr.Image(type="pil"),  # Output image
+    outputs=[
+        gr.Image(type="pil"),  # Output image
+        gr.File(label="Download Image", file_types = [".png", ".jpg", ".jpeg"]), # Output File to be downloaded
+    ],
     title="Image Generator",
     description="Provide a prompt, select an aspect ratio, and set realism to generate an image.",
 )