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

# Load model
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch_dtype,
    revision="fp16" if device == "cuda" else None
)
pipe = pipe.to(device)

MAX_SEED = 2**32 - 1

# Add "SelamGPT" watermark to image
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

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

    image = pipe(prompt=prompt, generator=generator).images[0]
    image = add_watermark(image)
    return image, seed

examples = [
    "አዲስ ዘመናዊ ከተማ በደመና መካከል", 
    "A futuristic Ethiopian skyline at night", 
    "አንድ የፍቅር ድራማ በተራራ በታች", 
]

with gr.Blocks() as demo:
    gr.Markdown("# SelamGPT ፡ Text-to-Image Generator 🖼️\nGenerate creative visuals from your imagination!")

    prompt = gr.Textbox(label="Image Prompt (in Amharic or English)", placeholder="e.g. አዲስ ከተማ በጨለማ ውስጥ")
    run_button = 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_button.click(fn=generate, inputs=[prompt, seed, randomize_seed], outputs=[result, seed])

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