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

# Initialize the model globally for faster inference
pipe = None

def load_model():
    global pipe
    if pipe is None:
        pipe = DiffusionPipeline.from_pretrained("segmind/tiny-sd", torch_dtype=torch.float32)
        pipe = pipe.to("cpu")
    return pipe

def generate_image_with_text(prompt, caption_text):
    """
    Generate an image based on a text prompt and add simple text with border at the bottom
    """
    # Ensure model is loaded
    pipe = load_model()
    
    # Generate image
    image = pipe(prompt, num_inference_steps=20, height=384, width=384).images[0]
    
    # Add space for text at bottom
    width, height = image.size
    new_img = Image.new("RGB", (width, height + 40), (0, 0, 0))
    new_img.paste(image, (0, 0))
    
    # Add text with border
    draw = ImageDraw.Draw(new_img)
    
    # Use default font
    font = ImageFont.load_default()
    
    # Calculate text position to center it
    text_width = draw.textlength(caption_text, font=font)
    text_x = (width - text_width) / 2
    text_y = height + 10
    
    # Draw text border (offset in 4 directions)
    offset = 1  # Border thickness
    for dx, dy in [(-offset, -offset), (-offset, offset), (offset, -offset), (offset, offset),
                  (0, -offset), (0, offset), (-offset, 0), (offset, 0)]:
        draw.text((text_x + dx, text_y + dy), caption_text, fill=(0, 0, 0), font=font)
    
    # Draw main text in white
    draw.text((text_x, text_y), caption_text, fill=(255, 255, 255), font=font)
    
    return new_img

# Create Gradio interface
title = "Text-to-Image Generator with Caption"
description = """
Generate an image from a text prompt and add a caption at the bottom.
The model used is lightweight and runs on CPU.
"""

demo = gr.Interface(
    fn=generate_image_with_text,
    inputs=[
        gr.Textbox(label="Image Prompt", placeholder="Describe the image you want to generate..."),
        gr.Textbox(label="Caption Text", placeholder="Text to display at the bottom of the image")
    ],
    outputs=gr.Image(type="pil", label="Generated Image"),
    title=title,
    description=description,
    examples=[
        ["A serene mountain landscape at sunset", "Beautiful Sunset View"],
        ["A cute cat playing with yarn", "Playful Kitten"],
        ["Abstract colorful shapes", "Modern Art"]
    ],
    cache_examples=False,
)

# Load the model when the app starts
load_model()

# Launch the app
if __name__ == "__main__":
    demo.launch()