Update app.py

app.py

@@ -1,12 +1,16 @@
+import gradio as gr
 import requests
+from PIL import Image
+import io
+import os
 
-def generate_kontext_image(input_image_url, prompt="transform this image", 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="kontext", nologo=True, enhance=False):
     """
     Generate a transformed image using the Kontext model from Pollinations API.
     
     Args:
-        input_image_url (str): URL of the input image to transform.
-        prompt (str): Prompt for the transformation (default: 'transform this image').
+        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).
@@ -15,8 +19,17 @@ def generate_kontext_image(input_image_url, prompt="transform this image", width
         enhance (bool): Whether to enhance the image (default: False).
     
     Returns:
-        bytes: Raw image data if successful, None otherwise.
+        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
+
     # Construct the API URL with query parameters
     base_url = "https://image.pollinations.ai/prompt"
     query_params = {
@@ -36,33 +49,78 @@ def generate_kontext_image(input_image_url, prompt="transform this image", width
         
         # Check if the request was successful
         if response.status_code == 200:
-            return response.content  # Return raw image data
+            # 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:
-            print(f"Error: API request failed with status code {response.status_code}")
-            return None
+            return f"Error: API request failed with status code {response.status_code}"
     except requests.RequestException as e:
-        print(f"Error: Failed to connect to the API - {e}")
-        return None
+        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)
 
-# Example usage
-if __name__ == "__main__":
-    # Replace with a valid input image URL
-    input_image = "https://example.com/input-image.jpg"
-    output_image_data = generate_kontext_image(
-        input_image_url=input_image,
-        prompt="transform this image into a surreal painting",
-        width=800,
-        height=600,
-        seed=42,
+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."
+    
+    return generate_kontext_image(
+        input_image=input_image,
+        prompt=prompt,
+        width=width,
+        height=height,
+        seed=seed,
         model="kontext",
-        nologo=True,
-        enhance=False
+        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.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Upload Image")
+            prompt = gr.Textbox(label="Prompt", placeholder="e.g., transform this image into a surreal painting")
+            width = gr.Slider(minimum=256, maximum=2048, value=1024, step=1, label="Width")
+            height = gr.Slider(minimum=256, maximum=2048, value=1024, step=1, label="Height")
+            seed = gr.Number(value=-1, label="Seed (-1 for random)", precision=0)
+            nologo = gr.Checkbox(value=True, label="No Logo")
+            enhance = gr.Checkbox(value=False, label="Enhance Image")
+            submit_button = gr.Button("Generate Image")
+        
+        with gr.Column():
+            output = gr.Image(label="Generated Image")
+    
+    submit_button.click(
+        fn=app_interface,
+        inputs=[input_image, prompt, width, height, seed, nologo, enhance],
+        outputs=output
     )
 
-    # Save the output image if successful
-    if output_image_data:
-        with open("output_image.jpg", "wb") as f:
-            f.write(output_image_data)
-        print("Image saved as output_image.jpg")
-    else:
-        print("Failed to generate image")
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()