app.py

import gradio as gr
import requests
import json
import os
import time
from collections import defaultdict
from PIL import Image
import io

BASE_URL = "https://api.jigsawstack.com/v1"
headers = {
    "x-api-key": os.getenv("JIGSAWSTACK_API_KEY")
}


# Rate limiting configuration
request_times = defaultdict(list)
MAX_REQUESTS = 20  # Maximum requests per time window
TIME_WINDOW = 3600   # Time window in seconds (1 hour)

def get_real_ip(request: gr.Request):
    """Extract real IP address using x-forwarded-for header or fallback"""
    if not request:
        return "unknown"
    
    forwarded = request.headers.get("x-forwarded-for")
    if forwarded:
        ip = forwarded.split(",")[0].strip()  # First IP in the list is the client's
    else:
        ip = request.client.host  # fallback
    return ip

def check_rate_limit(request: gr.Request):
    """Check if the current request exceeds rate limits"""
    if not request:
        return True, "Rate limit check failed - no request info"
    
    ip = get_real_ip(request)
    now = time.time()

    # Clean up old timestamps outside the time window
    request_times[ip] = [t for t in request_times[ip] if now - t < TIME_WINDOW]
    

    # Check if rate limit exceeded
    if len(request_times[ip]) >= MAX_REQUESTS:
        time_remaining = int(TIME_WINDOW - (now - request_times[ip][0]))
        time_remaining_minutes = round(time_remaining / 60, 1)      
        time_window_minutes = round(TIME_WINDOW / 60, 1)
        
        return False, f"Rate limit exceeded. You can make {MAX_REQUESTS} requests per {time_window_minutes} minutes. Try again in {time_remaining_minutes} minutes."
    
    # Add current request timestamp
    request_times[ip].append(now)
    return True, ""

def generate_image(prompt, aspect_ratio, width, height, steps, negative_prompt, guidance, seed, request: gr.Request):
    rate_limit_ok, rate_limit_msg = check_rate_limit(request)
    if not rate_limit_ok:
        return None, {"error": "Rate limit exceeded", "message": rate_limit_msg}
    
    # Validate required inputs
    if not prompt or not prompt.strip():
        return None, {"error": "Prompt is required"}
    
    if len(prompt.strip()) > 5000:
        return None, {"error": "Prompt must be between 1-5000 characters"}
    
    # Build payload with required and optional parameters
    payload = {
        "prompt": prompt.strip()
    }
    
    # Add optional parameters if provided
    if aspect_ratio and aspect_ratio.strip():
        payload["aspect_ratio"] = aspect_ratio.strip()
    
    if width:
        try:
            width_int = int(width)
            if 256 <= width_int <= 1920:
                payload["width"] = width_int
            else:
                return None, {"error": "Width must be between 256-1920 pixels"}
        except ValueError:
            return None, {"error": "Width must be a valid number"}
    
    if height:
        try:
            height_int = int(height)
            if 256 <= height_int <= 1920:
                payload["height"] = height_int
            else:
                return None, {"error": "Height must be between 256-1920 pixels"}
        except ValueError:
            return None, {"error": "Height must be a valid number"}
    
    if steps:
        try:
            steps_int = int(steps)
            if 1 <= steps_int <= 90:
                payload["steps"] = steps_int
            else:
                return None, {"error": "Steps must be between 1-90"}
        except ValueError:
            return None, {"error": "Steps must be a valid number"}
    
    if negative_prompt and negative_prompt.strip():
        payload["negative_prompt"] = negative_prompt.strip()
    
    if guidance:
        try:
            guidance_float = float(guidance)
            if 1 <= guidance_float <= 28:
                payload["guidance"] = guidance_float
            else:
                return None, {"error": "Guidance must be between 1-28"}
        except ValueError:
            return None, {"error": "Guidance must be a valid number"}
    
    if seed:
        try:
            seed_int = int(seed)
            payload["seed"] = seed_int
        except ValueError:
            return None, {"error": "Seed must be a valid number"}
    
    try:
        r = requests.post(f"{BASE_URL}/ai/image_generation", headers=headers, json=payload)
        r.raise_for_status()
        
        # The API returns the image directly as binary data
        if r.headers.get('content-type', '').startswith('image/'):
            # Convert bytes to PIL Image for Gradio
            
            # Create PIL Image from bytes
            image = Image.open(io.BytesIO(r.content))
            
            return image, {"success": True, "message": "Image generated successfully!"}
        else:
            # If not an image, try to parse as JSON for error
            try:
                error_data = r.json()
                return None, {"error": "Image generation failed", "details": error_data}
            except:
                return None, {"error": "Unexpected response format"}
                
    except requests.exceptions.RequestException as req_err:
        return None, {"error": "Request failed", "message": str(req_err)}
    except Exception as e:
        return None, {"error": "Unexpected error", "message": str(e)}

# ----------------- Gradio UI ------------------

with gr.Blocks() as demo:
    gr.Markdown("""
    <div style='text-align: center; margin-bottom: 24px;'>
        <h1 style='font-size:2.2em; margin-bottom: 0.2em;'>🧩 Image Generation</h1>
        <p style='font-size:1.2em; margin-top: 0;'>Generate high-quality images from text prompts using advanced AI models.</p>
        <p style='font-size:1em; margin-top: 0.5em;'>For more details and API usage, see the <a href='https://jigsawstack.com/docs/api-reference/ai/image-generation' target='_blank'>documentation</a>.</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("#### Input")
            prompt = gr.Textbox(
                label="Text Prompt", 
                lines=4, 
                placeholder="Enter text to describe the image"
            )
            
            gr.Markdown("#### Quick Examples")
            gr.Markdown("Click any example below to auto-fill the prompt:")
            
            with gr.Row():
                img_example_btn1 = gr.Button("🏔️ Mountain Landscape", size="sm")
                img_example_btn2 = gr.Button("🐱 Cute Cat Portrait", size="sm")
                img_example_btn3 = gr.Button("🏙️ Cyberpunk City", size="sm")
            
            with gr.Row():
                img_example_btn4 = gr.Button("🌊 Ocean Sunset", size="sm")
                img_example_btn5 = gr.Button("🚀 Space Station", size="sm")
                img_example_btn6 = gr.Button("🏰 Medieval Castle", size="sm")
            
            with gr.Row():
                img_example_btn7 = gr.Button("🎨 Abstract Art", size="sm")
                img_example_btn8 = gr.Button("🍕 Delicious Pizza", size="sm")
            
            gr.Markdown("#### Image Options")
            aspect_ratio = gr.Textbox(
                label="Aspect Ratio (optional, e.g., 16:9, 4:3)", 
                placeholder="16:9, 4:3, etc.",
                info="Common ratios: 16:9, 4:3, 1:1"
            )
            width = gr.Textbox(
                label="Width (optional, pixels)", 
                placeholder="1024",
                info="Minimum: 256, Maximum: 1920"
            )
            height = gr.Textbox(
                label="Height (optional, pixels)", 
                placeholder="768",
                info="Minimum: 256, Maximum: 1920"
            )
            steps = gr.Textbox(
                label="Steps (optional, 1-90)", 
                placeholder="50",
                info="Minimum: 1, Maximum: 90"
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt (optional)", 
                placeholder="A bad image",
                info="Optional: describe what you don't want in the image"
            )
            guidance = gr.Textbox(
                label="Guidance (optional, 1-28)", 
                placeholder="7.5",
                info="Minimum: 1, Maximum: 28"
            )
            seed = gr.Textbox(
                label="Seed (optional)", 
                placeholder="Random",
                info="Optional: randomize image generation"
            )

        with gr.Column():
            gr.Markdown("#### Preview")
            image_preview = gr.Image(label="Generated Image")
            generate_image_btn = gr.Button("Generate Image", variant="primary")
            generate_image_result = gr.JSON(label="Image Generation Result")

    # Example functions to auto-fill prompt field
    def fill_img_example_1():
        return "A majestic mountain landscape with snow-capped peaks, crystal clear lake in the foreground, golden hour lighting, photorealistic, 8K resolution"
    
    def fill_img_example_2():
        return "A cute fluffy cat sitting on a windowsill, looking directly at camera with big green eyes, soft natural lighting, high quality, detailed fur"
    
    def fill_img_example_3():
        return "A futuristic cyberpunk cityscape at night, neon lights, flying cars, tall skyscrapers, rain-slicked streets, cinematic lighting"
    
    def fill_img_example_4():
        return "A breathtaking ocean sunset over calm waters, orange and pink sky, silhouettes of palm trees, peaceful atmosphere, high resolution"
    
    def fill_img_example_5():
        return "A massive space station orbiting Earth, detailed engineering, solar panels, astronauts in space suits, stars and nebula background"
    
    def fill_img_example_6():
        return "An ancient medieval castle on a hilltop, stone walls and towers, flags waving, dramatic clouds, fantasy atmosphere, detailed architecture"
    
    def fill_img_example_7():
        return "Abstract digital art with flowing colors, geometric shapes, vibrant gradients, modern aesthetic, artistic composition"
    
    def fill_img_example_8():
        return "A delicious pizza with melted cheese, pepperoni, fresh basil, wood-fired oven, appetizing, high quality food photography"

    # Connect example buttons to auto-fill functions
    img_example_btn1.click(fill_img_example_1, outputs=[prompt])
    img_example_btn2.click(fill_img_example_2, outputs=[prompt])
    img_example_btn3.click(fill_img_example_3, outputs=[prompt])
    img_example_btn4.click(fill_img_example_4, outputs=[prompt])
    img_example_btn5.click(fill_img_example_5, outputs=[prompt])
    img_example_btn6.click(fill_img_example_6, outputs=[prompt])
    img_example_btn7.click(fill_img_example_7, outputs=[prompt])
    img_example_btn8.click(fill_img_example_8, outputs=[prompt])

    generate_image_btn.click(
        generate_image,
        inputs=[prompt, aspect_ratio, width, height, steps, negative_prompt, guidance, seed],
        outputs=[image_preview, generate_image_result]
    )

demo.launch()