Update app.py

app.py

@@ -2,36 +2,53 @@ import gradio as gr
 import requests
 from PIL import Image
 import io
-import os
 
-def generate_kontext_image(input_image, prompt, width=1024, height=1024, seed=-1, model="kontext", nologo=True, enhance=False):
+def generate_kontext_image(input_image, prompt, width=1024, height=1024, seed=-1, model="dreamshaper", nologo=True, enhance=False):
     """
-    Generate a transformed image using the Kontext model from Pollinations API.
+    Generate a transformed image using the Pollinations API.
     
     Args:
         input_image (PIL.Image): Input image to transform.
         prompt (str): Prompt for the transformation.
-        width (int): Width of the output image (default: 1024).
-        height (int): Height of the output image (default: 1024).
-        seed (int): Random seed for generation (default: -1 for random).
-        model (str): Model to use (default: 'kontext').
-        nologo (bool): Whether to exclude logo (default: True).
-        enhance (bool): Whether to enhance the image (default: False).
+        width (int): Width of the output image.
+        height (int): Height of the output image.
+        seed (int): Random seed for generation (-1 for random).
+        model (str): Model to use (default: 'dreamshaper').
+        nologo (bool): Whether to exclude logo.
+        enhance (bool): Whether to enhance the image.
     
     Returns:
         PIL.Image or str: Generated image or error message.
     """
-    # Save the uploaded image temporarily to get a public URL
-    temp_image_path = "temp_input_image.jpg"
-    input_image.save(temp_image_path)
-
-    # Note: In a production environment, you should upload the image to a public URL
-    # For this example, we'll assume the image is accessible locally or via a public URL
-    # Replace with actual public URL if deploying (e.g., upload to a cloud service)
-    input_image_url = "https://example.com/input-image.jpg"  # Placeholder; replace with actual URL
+    # Step 1: Convert the input PIL Image to bytes for uploading.
+    image_bytes = io.BytesIO()
+    input_image.save(image_bytes, format='JPEG')
+    image_bytes.seek(0)
+    
+    input_image_url = ""
+    # Step 2: Upload the image bytes to get a public URL.
+    try:
+        upload_response = requests.post(
+            'https://image.pollinations.ai/upload',
+            files={'file': ('input_image.jpg', image_bytes, 'image/jpeg')}
+        )
+        upload_response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)
+        upload_result = upload_response.json()
+        input_image_url = upload_result.get('ipfs')
+        
+        if not input_image_url:
+            return "Error: Could not retrieve a public URL after uploading the image."
+            
+    except requests.RequestException as e:
+        return f"Error: Failed to upload the image to the server - {e}"
 
-    # Construct the API URL with query parameters
+    # Step 3: Use the public URL to make the final image generation request.
     base_url = "https://image.pollinations.ai/prompt"
+    
+    # URL-encode the prompt to handle spaces and special characters
+    encoded_prompt = requests.utils.quote(prompt)
+    api_url = f"{base_url}/{encoded_prompt}"
+
     query_params = {
         "model": model,
         "image": input_image_url,
@@ -41,65 +58,51 @@ def generate_kontext_image(input_image, prompt, width=1024, height=1024, seed=-1
         "nologo": str(nologo).lower(),
         "enhance": str(enhance).lower()
     }
-    api_url = f"{base_url}/{prompt.replace(' ', '%20')}"
 
     try:
         # Make the API request
         response = requests.get(api_url, params=query_params, stream=True)
+        response.raise_for_status()
+        
+        # Convert response content to a PIL Image and return it
+        output_image = Image.open(io.BytesIO(response.content))
+        return output_image
         
-        # Check if the request was successful
-        if response.status_code == 200:
-            # Convert response content to PIL Image
-            output_image = Image.open(io.BytesIO(response.content))
-            # Clean up temporary file
-            if os.path.exists(temp_image_path):
-                os.remove(temp_image_path)
-            return output_image
-        else:
-            return f"Error: API request failed with status code {response.status_code}"
     except requests.RequestException as e:
-        return f"Error: Failed to connect to the API - {e}"
-    finally:
-        # Ensure temporary file is removed even if an error occurs
-        if os.path.exists(temp_image_path):
-            os.remove(temp_image_path)
+        error_details = str(e)
+        try:
+            # Try to get a more specific error message from the API response
+            error_details = e.response.json().get("message", e.response.text)
+        except:
+            pass # Keep the original exception text if parsing fails
+        return f"Error: API request failed. Details: {error_details}"
+
 
 def app_interface(input_image, prompt, width, height, seed, nologo, enhance):
     """
     Gradio interface function to handle user inputs and display results.
-    
-    Args:
-        input_image (PIL.Image): Uploaded image.
-        prompt (str): Transformation prompt.
-        width (int): Output image width.
-        height (int): Output image height.
-        seed (int): Random seed.
-        nologo (bool): Exclude logo.
-        enhance (bool): Enhance image.
-    
-    Returns:
-        PIL.Image or str: Generated image or error message.
     """
     if input_image is None:
         return "Please upload an image."
     if not prompt:
         return "Please provide a prompt."
     
+    # Using 'dreamshaper' as it's a free and effective model for this task
     return generate_kontext_image(
         input_image=input_image,
         prompt=prompt,
         width=width,
         height=height,
         seed=seed,
-        model="kontext",
+        model="dreamshaper",
         nologo=nologo,
         enhance=enhance
     )
 
 # Define the Gradio interface
-with gr.Blocks(title="Kontext Image Transformation") as demo:
-    gr.Markdown("# Kontext Image Transformation App")
-    gr.Markdown("Upload an image, provide a transformation prompt, and generate a new image using the Kontext model.")
+with gr.Blocks(title="Image Transformation") as demo:
+    gr.Markdown("# Image Transformation App")
+    gr.Markdown("Upload an image, provide a transformation prompt, and generate a new image.")
     
     with gr.Row():
         with gr.Column():