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

# Load the tokenizer and model
tokenizer = DalleBartTokenizer.from_pretrained("dalle-mini/dalle-mini")
model = DalleBartForConditionalGeneration.from_pretrained("dalle-mini/dalle-mini")

# Ensure the model is in evaluation mode and use GPU if available
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

def generate_image(prompt):
    try:
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        outputs = model.generate(**inputs)
        
        # Convert the output tensor to an image
        # Adjust this line according to the model output format
        image = outputs[0]  # Assuming this is an image tensor
        image = Image.fromarray(np.array(image))  # Convert tensor to numpy array then to image

        return image
    except Exception as e:
        return str(e)

# Define Gradio interface
iface = gr.Interface(fn=generate_image,
                     inputs="text",
                     outputs="image",
                     title="DALL-E Mini Image Generator",
                     description="Enter a prompt to generate an image.")

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