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

# Load the model (make sure to use a model that exists on Hugging Face)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32).to(device)

def generate_image(caption):
    # Generate the image from the caption
    try:
        with torch.no_grad():
            image = model(caption).images[0]
        
        # Convert to PIL Image for drawing
        image = image.convert("RGBA")
        
        # Create a draw object
        draw = ImageDraw.Draw(image)
        
        # Define font size and color
        font_size = 40
        font_color = "white"
        
        # Load a font
        font = ImageFont.load_default()  # You can specify a TTF font file if needed
        
        # Calculate text size and position
        text_width, text_height = draw.textsize(caption, font=font)
        text_position = ((image.width - text_width) // 2, 10)  # Centered at the top
        
        # Draw the text on the image
        draw.text(text_position, caption, font=font, fill=font_color)
        
        return image
    except Exception as e:
        print(f"Error generating image: {e}")
        return None

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Text to Image Generation with Meme Caption")
    
    with gr.Row():
        caption_input = gr.Textbox(label="Enter your caption", placeholder="Type your caption here...")
        generate_button = gr.Button("Generate Image")
    
    output_image = gr.Image(label="Generated Image")
    
    generate_button.click(fn=generate_image, inputs=caption_input, outputs=output_image)

# Launch the app
demo.launch()