import gradio as gr
import requests
from PIL import Image
from io import BytesIO

# Define style tokens or modifiers
styles = {
    "None": " ",
    "Hyper-Realistic": " hyper-realistic",
    "Realistic": " realistic",
    "Flux": " flux",
    "Abstract": " abstract",
    "Cartoon": " cartoon",
    "Ghibli": " ghibli style",
    "Anime": " anime style",
    "Watercolor": " watercolor painting",
    "Pixel Art": " pixel art",
    "Sketch": " sketch",
    "Digital Art": " digital art",
    "Oil Painting": " oil painting",
    "Pastel": " pastel painting",
    "Photorealistic": " photorealistic",
    "Sci-Fi": " sci-fi",
    "Fantasy": " fantasy",
}

def generate_image(prompt, style):
    # Append the style token to the prompt
    style_token = styles.get(style, "")
    full_prompt = prompt + style_token
    
    # Set image parameters
    width = 1024
    height = 1024
    seed = 42
    model = "flux"  # You can adjust the model if needed

    # Construct the URL for pollinations.ai
    image_url = f"https://pollinations.ai/p/{full_prompt}?width={width}&height={height}&seed={seed}&model={model}"
    
    # Download the image using requests
    response = requests.get(image_url)
    if response.status_code == 200:
        image = Image.open(BytesIO(response.content))
        return image
    else:
        return None

def crop_watermark(pil_img):
    """
    Crop the image from the bottom to remove the watermark.
    This function assumes the generated image is 1024x1024,
    and crops it to a height of 960 pixels (removing the bottom 64 pixels).
    """
    width, height = pil_img.size  # Should be 1024 x 1024
    # Crop from the top left (0,0) to (width, 960)
    cropped_img = pil_img.crop((0, 0, width, 960))
    return cropped_img

# Gradio interface function
def interface(prompt, style):
    image = generate_image(prompt, style)
    if image is None:
        return "Error generating image."
    # Crop out the bottom portion with the watermark
    image_cropped = crop_watermark(image)
    return image_cropped

# Define style options for the radio buttons
style_options = list(styles.keys())

# Build Gradio interface using Blocks
with gr.Blocks() as demo:
    gr.Markdown("## Canvasio: AI Image Generator")
    with gr.Row():
        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="A beautiful landscape")
    with gr.Row():
        style_choice = gr.Radio(choices=style_options, label="Select image style", value="Flux")
    generate_btn = gr.Button("Generate Image")
    output_image = gr.Image(label="Generated Image")
    
    generate_btn.click(fn=interface, inputs=[prompt_input, style_choice], outputs=output_image)

demo.launch()