Update app.py

app.py

@@ -1,30 +1,38 @@
 import gradio as gr
+from transformers import DalleBartProcessor, FlaxDalleBartForConditionalGeneration
+from PIL import Image
+import numpy as np
 import jax
 import jax.numpy as jnp
-from flax.jax_utils import replicate
-from transformers import DalleBartProcessor, FlaxDalleBartForConditionalGeneration
 
-# Load the model and processor
+# Load the DALL-E Mega model and processor
 processor = DalleBartProcessor.from_pretrained("dalle-mini/dalle-mega")
 model = FlaxDalleBartForConditionalGeneration.from_pretrained("dalle-mini/dalle-mega")
 
 # Function to generate an image from a text prompt
 def generate_image(prompt):
-    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)
+    # Process the prompt
+    inputs = processor([prompt], return_tensors="jax", padding="max_length", truncation=True, max_length=64)
+    
+    # Generate the images
+    outputs = model.generate(**inputs, do_sample=True, num_beams=4, num_return_sequences=1)
+    
+    # Decode the images and convert them to displayable format
+    images = model.decode(outputs.sequences)
+    images = jax.device_get(images)
+    
+    # Convert to a PIL image
+    pil_img = Image.fromarray(np.asarray(images[0]).astype(np.uint8))
     
-    # Post-process image for display (convert to PIL image format)
-    image = images[0]  # assuming single image output
-    return image  # return the generated image
+    return pil_img
 
 # Create Gradio interface
 iface = gr.Interface(
-    fn=generate_image,               # Function to generate image
-    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your prompt"),   # Textbox input
+    fn=generate_image,               # Function to generate the image
+    inputs=gr.Textbox(lines=2, placeholder="Enter your text prompt"),   # Input textbox for the prompt
     outputs="image",                 # Output as an image
-    title="DALL-E Mini Image Generator",
-    description="Generate images from text prompts using DALL-E Mini model."
+    title="DALL-E Mini/Mega Image Generator",
+    description="Generate images from text prompts using the DALL-E Mega model."
 )
 
 # Launch the app