Update app.py

app.py

@@ -5,6 +5,8 @@ from PIL import Image
 import requests
 from io import BytesIO
 import time
+import tempfile
+import base64
 
 # Set up Replicate API key from environment variable
 os.environ['REPLICATE_API_TOKEN'] = os.getenv('REPLICATE_API_TOKEN')
@@ -21,44 +23,85 @@ def process_images(prompt, image1, image2=None):
         return None, "⚠️ Please set REPLICATE_API_TOKEN environment variable"
     
     try:
-        # Prepare input images list
+        import tempfile
+        import base64
+        
+        # Save images temporarily and create data URIs
         image_inputs = []
         
-        # Convert PIL images to temporary files or use file paths
-        # Note: In production, you'd upload these to a cloud service
-        # For this example, we'll assume the images are accessible via URLs
-        # You may need to upload them to a service like Cloudinary or S3 first
+        # Process first image
+        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp1:
+            image1.save(tmp1.name)
+            with open(tmp1.name, 'rb') as f:
+                img_data = base64.b64encode(f.read()).decode()
+                # For some models, you might need data URI format
+                image_inputs.append(f"data:image/png;base64,{img_data}")
+            os.unlink(tmp1.name)  # Clean up temp file
         
-        # This is a placeholder - in real implementation, you'd need to:
-        # 1. Save images temporarily
-        # 2. Upload to a cloud service
-        # 3. Get the URLs
+        # Process second image if provided
+        if image2:
+            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp2:
+                image2.save(tmp2.name)
+                with open(tmp2.name, 'rb') as f:
+                    img_data = base64.b64encode(f.read()).decode()
+                    image_inputs.append(f"data:image/png;base64,{img_data}")
+                os.unlink(tmp2.name)
         
         status_message = "🎨 Processing your images..."
         
-        # Prepare input for Replicate
-        input_data = {
-            "prompt": prompt,
-            "image_input": image_inputs  # This should contain actual URLs
-        }
+        # Example using a real Replicate model (stable-diffusion)
+        # You should replace this with your actual model
+        try:
+            # For image-to-image models, use something like:
+            output = replicate.run(
+                "stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf",
+                input={
+                    "prompt": prompt,
+                    "image": image_inputs[0] if image_inputs else None,
+                    "num_outputs": 1,
+                    "guidance_scale": 7.5,
+                    "num_inference_steps": 50
+                }
+            )
+        except:
+            # Fallback to a simpler text-to-image if image input fails
+            output = replicate.run(
+                "stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf",
+                input={
+                    "prompt": prompt,
+                    "num_outputs": 1
+                }
+            )
         
-        # Run the model
-        output = replicate.run(
-            "google/nano-banana",  # Replace with actual model
-            input=input_data
-        )
+        # Handle different output formats
+        output_url = None
         
-        # Get the output URL
-        if hasattr(output, 'url'):
+        if isinstance(output, list) and len(output) > 0:
+            # If output is a list, take the first item
+            output_url = output[0]
+        elif isinstance(output, str):
+            # If output is already a string URL
+            output_url = output
+        elif hasattr(output, 'url'):
+            # If output has a url method
             output_url = output.url()
-        else:
-            output_url = str(output)
+        elif hasattr(output, '__iter__'):
+            # If output is iterable, try to get first item
+            try:
+                output_url = next(iter(output))
+            except:
+                pass
+        
+        if not output_url:
+            return None, "❌ Error: No image content found in response"
         
         # Download and return the generated image
         response = requests.get(output_url)
-        img = Image.open(BytesIO(response.content))
-        
-        return img, "✅ Image generated successfully!"
+        if response.status_code == 200:
+            img = Image.open(BytesIO(response.content))
+            return img, "✅ Image generated successfully!"
+        else:
+            return None, f"❌ Error: Failed to download image (Status: {response.status_code})"
         
     except Exception as e:
         return None, f"❌ Error: {str(e)}"
@@ -201,10 +244,16 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
        pip install gradio replicate pillow requests
        ```
     
-    3. **Note:** For production use, you'll need to:
+    3. **Available Models to Try:**
+       - `stability-ai/stable-diffusion` - Text to image generation
+       - `jagilley/controlnet-canny` - Image style transfer with edge detection
+       - `rossjillian/controlnet` - Image controlled generation
+       - Replace the model in the code with your preferred model
+    
+    4. **Note:** For production use, you'll need to:
        - Implement proper image upload to cloud storage (S3, Cloudinary, etc.)
-       - Replace the model name with the actual Replicate model you want to use
        - Add proper error handling and rate limiting
+       - Check the specific input format required by your chosen model
     
     ### 🔒 Security:
     - API keys are managed through environment variables