import gradio as gr
from transformers import eBart
import torch

# Load the eBart model
model = eBart.from_pretrained("dalle-mini/dalle-mega")

# Define a function to generate an image from text
def generate_image(prompt):
    inputs = model.prepare_inputs_for_generation(prompt)
    outputs = model.generate(inputs)
    return outputs  # You may need to convert it to a displayable image format depending on the model output

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,               # Function that takes a prompt and returns an image
    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your prompt"),   # Textbox input for the prompt
    outputs="image",                 # Output is an image
    title="eBart DALL-E Mega Image Generator",
    description="Generate images from text prompts using the DALL-E Mega model."
)

# Launch the app
iface.launch()