Update app.py

app.py

@@ -1,23 +1,24 @@
 import gradio as gr
 from transformers import eBart
+import torch
 
-# Load the DALL-E Mega model
+# Load the eBart model
 model = eBart.from_pretrained("dalle-mini/dalle-mega")
 
-# Define the function to generate images
-def generate_image(text):
-    inputs = model.tokenizer(text, return_tensors="pt")
-    outputs = model.generate(**inputs)
-    image_url = model.tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return image_url
+# 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 the Gradio interface
-ui = gr.Interface(
-    fn=generate_image,
-    inputs="text",
-    outputs="image",
-    title="DALL-E Mega Image Generator"
+# 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 Gradio app
-ui.launch()
+# Launch the app
+iface.launch()