Update app.py

app.py

@@ -1,30 +1,23 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 import gradio as gr
+from transformers import eBart
 
-# Load the model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("dalle-mini/dalle-mega")
-model = AutoModelForCausalLM.from_pretrained("dalle-mini/dalle-mega")
+# Load the DALL-E Mega model
+model = eBart.from_pretrained("dalle-mini/dalle-mega")
 
-# Define the function for Gradio interface
-def generate_image(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    
-    # Generate image (or output) using the model
-    with torch.no_grad():
-        outputs = model.generate(**inputs)
-    
-    # Convert output to a format suitable for Gradio
-    # This part may need to be adapted based on actual output format
-    return outputs
+# 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
 
-# Set up Gradio interface
-iface = gr.Interface(
+# Create the Gradio interface
+ui = gr.Interface(
     fn=generate_image,
-    inputs=gr.Textbox(label="Enter prompt"),
-    outputs=gr.Image(type="pil", label="Generated Image"),
-    live=True
+    inputs="text",
+    outputs="image",
+    title="DALL-E Mega Image Generator"
 )
 
-# Launch the app
-iface.launch()
+# Launch the Gradio app
+ui.launch()