Update app.py

app.py

@@ -1,23 +1,30 @@
 import gradio as gr
-from transformers import eBart
-import torch
+import jax
+import jax.numpy as jnp
+from flax.jax_utils import replicate
+from transformers import DalleBartProcessor, FlaxDalleBartForConditionalGeneration
 
-# Load the eBart model
-model = eBart.from_pretrained("dalle-mini/dalle-mega")
+# Load the model and processor
+processor = DalleBartProcessor.from_pretrained("dalle-mini/dalle-mega")
+model = FlaxDalleBartForConditionalGeneration.from_pretrained("dalle-mini/dalle-mega")
 
-# Define a function to generate an image from text
+# Function to generate an image from a text prompt
 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
+    inputs = processor([prompt], return_tensors="jax", padding="max_length", truncation=True, max_length=128)
+    # Generate images
+    images = model.generate(**replicate(inputs.data), do_sample=True, num_beams=1, num_return_sequences=1)
+    
+    # Post-process image for display (convert to PIL image format)
+    image = images[0]  # assuming single image output
+    return image  # return the generated image
 
 # 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."
+    fn=generate_image,               # Function to generate image
+    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your prompt"),   # Textbox input
+    outputs="image",                 # Output as an image
+    title="DALL-E Mini Image Generator",
+    description="Generate images from text prompts using DALL-E Mini model."
 )
 
 # Launch the app