import gradio as gr
from diffusers import AutoPipelineForText2Image
from PIL import Image, ImageDraw, ImageFont
import torch
import random

device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/sdxl-turbo",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    variant="fp16" if device == "cuda" else None
)
pipe = pipe.to(device)

MAX_SEED = 2**32 - 1

def add_watermark(image):
    draw = ImageDraw.Draw(image)
    font = ImageFont.load_default()
    text = "SelamGPT"
    margin = 10
    x = image.width - draw.textlength(text, font=font) - margin
    y = image.height - 20
    draw.text((x, y), text, font=font, fill=(255, 255, 255))
    return image

def generate(prompt, seed, randomize_seed):
    if randomize_seed or seed == 0:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)

    image = pipe(
        prompt=prompt,
        num_inference_steps=2,
        guidance_scale=0.0,
        generator=generator,
    ).images[0]

    image = add_watermark(image)
    return image, seed

examples = [
    "Futuristic Ethiopian city at sunset, detailed, cinematic",
    "አንድ ጫካ ውስጥ የተሰወረ የሳይባር ከተማ በመስሎ የታሰረ ቀይ ብርሃን",
]

with gr.Blocks() as demo:
    gr.Markdown("## 🎨 SelamGPT - Super Fast Text-to-Image Generator")

    prompt = gr.Textbox(label="Prompt", placeholder="Type your idea in English or Amharic")
    run = gr.Button("Generate")

    result = gr.Image(label="Generated Image")
    with gr.Accordion("Advanced Settings", open=False):
        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

    gr.Examples(examples=examples, inputs=[prompt])

    run.click(fn=generate, inputs=[prompt, seed, randomize_seed], outputs=[result, seed])

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