import gradio as gr
import torch
from transformers import DalleBartTokenizer, DalleBartForConditionalGeneration
from PIL import Image
import io

# Load model and tokenizer
model_id = "dalle-mini/dalle-mini"  # Example model id; adjust if needed
model = DalleBartForConditionalGeneration.from_pretrained(model_id)
tokenizer = DalleBartTokenizer.from_pretrained(model_id)

# Function to generate image
def generate_image(prompt, num_inference_steps=50):
    inputs = tokenizer(prompt, return_tensors="pt")
    
    # Generate images
    with torch.no_grad():
        outputs = model.generate(**inputs, num_beams=num_inference_steps)
    
    # Convert tensor to PIL image
    image = Image.fromarray(outputs[0].cpu().numpy().astype('uint8'))
    
    return image

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Text to Image Generation")

    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here...")
        num_inference_steps = gr.Slider(minimum=1, maximum=50, step=1, value=28, label="Number of Inference Steps")

    with gr.Row():
        generate_button = gr.Button("Generate Image")

    result = gr.Image(label="Generated Image")

    # Connect the function to the button
    generate_button.click(
        fn=generate_image,
        inputs=[prompt, num_inference_steps],
        outputs=result
    )

# Launch the Gradio app
demo.launch()