app.py

import gradio as gr
import replicate
import os
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')

def process_images(prompt, image1, image2=None):
    """
    Process uploaded images with Replicate API
    """
    if not image1:
        return None, "Please upload at least one image"
    
    # Check if API token is set
    if not os.getenv('REPLICATE_API_TOKEN'):
        return None, "⚠️ Please set REPLICATE_API_TOKEN environment variable"
    
    try:
        import tempfile
        import base64
        
        # Save images temporarily and create data URIs
        image_inputs = []
        
        # 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
        
        # 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..."
        
        # 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
                }
            )
        
        # Handle different output formats
        output_url = None
        
        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()
        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)
        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)}"

# Create Gradio interface with gradient theme
css = """
.gradio-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    font-family: 'Inter', sans-serif;
}
.gr-button {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    border: none;
    color: white;
    font-weight: bold;
    transition: transform 0.2s;
}
.gr-button:hover {
    transform: scale(1.05);
    box-shadow: 0 10px 20px rgba(0,0,0,0.2);
}
.gr-input {
    border-radius: 10px;
    border: 2px solid rgba(255,255,255,0.3);
    background: rgba(255,255,255,0.9);
}
.header-text {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-size: 2.5em;
    font-weight: bold;
    text-align: center;
    margin-bottom: 20px;
}
.description-text {
    color: white;
    text-align: center;
    font-size: 1.1em;
    margin-bottom: 30px;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
}
"""

# Build the Gradio interface
with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.HTML("""
        <div class="header-text">🎨 AI Image Style Transfer Studio</div>
        <div class="description-text">
            Upload 1-2 images and describe how you want them styled.
            The AI will create a beautiful transformation!
        </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Input Section")
            
            prompt = gr.Textbox(
                label="✏️ Style Prompt",
                placeholder="Describe how you want to style your images...",
                lines=3,
                value="Make the sheets in the style of the logo. Make the scene natural."
            )
            
            with gr.Row():
                image1 = gr.Image(
                    label="Image 1 (Required)",
                    type="pil",
                    height=200
                )
                image2 = gr.Image(
                    label="Image 2 (Optional)",
                    type="pil",
                    height=200
                )
            
            generate_btn = gr.Button(
                "🚀 Generate Styled Image",
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("""
            #### 💡 Tips:
            - Upload high-quality images for best results
            - Be specific in your style description
            - Experiment with different prompts!
            """)
        
        with gr.Column(scale=1):
            gr.Markdown("### 🎯 Output Section")
            
            output_image = gr.Image(
                label="Generated Image",
                type="pil",
                height=400
            )
            
            status = gr.Textbox(
                label="Status",
                interactive=False,
                lines=2
            )
    
    # Examples section
    with gr.Row():
        gr.Examples(
            examples=[
                ["Transform into watercolor painting style", None, None],
                ["Make it look like a vintage photograph", None, None],
                ["Apply cyberpunk neon style", None, None],
                ["Convert to minimalist line art", None, None],
            ],
            inputs=[prompt, image1, image2],
            label="Example Prompts"
        )
    
    # Event handlers
    generate_btn.click(
        fn=process_images,
        inputs=[prompt, image1, image2],
        outputs=[output_image, status],
        api_name="generate"
    )
    
    # Additional information
    gr.Markdown("""
    ---
    ### ⚙️ Setup Instructions:
    
    1. **Set Environment Variable:**
       ```bash
       export REPLICATE_API_TOKEN="your_api_token_here"
       ```
    
    2. **Install Required Packages:**
       ```bash
       pip install gradio replicate pillow requests
       ```
    
    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.)
       - 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
    - Never commit API keys to version control
    - Consider implementing user authentication for production
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )