import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import io

# Load the model globally once
pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/sd-turbo"
)
pipe.to("cpu")  # Free hardware: CPU

# Style presets
STYLE_PRESETS = {
    "Minimal": "minimal, clean lines, flat design",
    "Modern": "modern, bold, gradient",
    "Retro": "retro, vintage, 90s style"
}

# Resolutions
RESOLUTIONS = ["256x256", "384x384", "512x512", "768x768"]

def generate_logo(prompt, style, resolution, transparent):
    # Apply style
    style_prompt = STYLE_PRESETS.get(style, "")
    full_prompt = f"{style_prompt}, {prompt}".strip()

    # Resolution
    width, height = map(int, resolution.split("x"))

    # Run pipeline
    image = pipe(full_prompt, height=height, width=width, num_inference_steps=15).images[0]

    # Convert to RGBA if transparent
    if transparent:
        image = image.convert("RGBA")

    return image

def get_image_bytes(img):
    buffer = io.BytesIO()
    img.save(buffer, format="PNG")
    buffer.seek(0)
    return buffer

with gr.Blocks(theme=gr.themes.Default(), title="🎨 AI Logo Generator") as demo:
    gr.Markdown("## 🎨 AI Logo Generator\nEnter a description to generate a high-quality logo")

    with gr.Row():
        prompt = gr.Textbox(label="Logo Description", placeholder="e.g., A minimalist logo of owl for tech startup")
        style = gr.Dropdown(label="Style", choices=list(STYLE_PRESETS.keys()), value="Minimal")
        resolution = gr.Dropdown(label="Resolution", choices=RESOLUTIONS, value="512x512")

    with gr.Row():
        transparent = gr.Checkbox(label="Make background transparent", value=False)
        dark_mode = gr.Checkbox(label="Dark Mode UI", value=False)
        share_link = gr.Checkbox(label="Make app public (Gradio share link)", value=True)

    btn = gr.Button("Generate")

    output_img = gr.Image(label="Generated Logo", type="pil")
    download_btn = gr.File(label="Download Logo (PNG)")

    def process_all(prompt, style, resolution, transparent, dark_mode, share_link):
        image = generate_logo(prompt, style, resolution, transparent)
        file = get_image_bytes(image)
        return image, (file, "logo.png")

    btn.click(
        fn=process_all,
        inputs=[prompt, style, resolution, transparent, dark_mode, share_link],
        outputs=[output_img, download_btn]
    )

if __name__ == "__main__":
    demo.launch(share=True)