import gradio as gr
import requests
import random
import os
from PIL import Image

def process(prompt, seed, width, height):
    if seed == "":
        seed = random.randint(0, 999999999)
    else:
        seed = int(seed)

    filename = f"{random.randint(111111111, 999999999)}.jpg"
    file_path = os.path.join(os.path.dirname(__file__), filename)

    url = f"https://image.pollinations.ai/prompt/{prompt}?model=flux-realism&width={width}&height={height}&nologo=true&seed={seed}"
    response = requests.get(url)

    if response.status_code == 200:
        with open(file_path, "wb") as f:
            f.write(response.content)

        # Convert PNG to JPG using PIL
        try:
            img = Image.open(file_path)
            rgb_img = img.convert("RGB")
            rgb_img.save(file_path, "JPEG")
        except Exception as e:
            return f"Error converting image: {e}"

        return file_path
    else:
        return "Failed to generate image."

title = "Pollinations Image Generator"
description = "Generate high-resolution AI art using Pollinations API. You can customize the seed, resolution, and prompt."

iface = gr.Interface(
    fn=process,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Textbox(label="Seed (optional)", placeholder="Leave empty for random"),
        gr.Slider(512, 2048, step=64, label="Width", value=1024),
        gr.Slider(512, 2048, step=64, label="Height", value=1024),
    ],
    outputs=gr.File(label="Download Image"),
    title=title,
    description=description,
)

iface.launch()