Update app.py

app.py

@@ -1,46 +1,30 @@
-import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-from transformers import DalleBartTokenizer, DalleBartForConditionalGeneration
-from PIL import Image
-import io
+import gradio as gr
 
-# Load model and tokenizer
-model_id = "dalle-mini/dalle-mini"  # Example model id; adjust if needed
-model = DalleBartForConditionalGeneration.from_pretrained(model_id)
-tokenizer = DalleBartTokenizer.from_pretrained(model_id)
+# Load the model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("dalle-mini/dalle-mega")
+model = AutoModelForCausalLM.from_pretrained("dalle-mini/dalle-mega")
 
-# Function to generate image
-def generate_image(prompt, num_inference_steps=50):
+# Define the function for Gradio interface
+def generate_image(prompt):
     inputs = tokenizer(prompt, return_tensors="pt")
     
-    # Generate images
+    # Generate image (or output) using the model
     with torch.no_grad():
-        outputs = model.generate(**inputs, num_beams=num_inference_steps)
-    
-    # Convert tensor to PIL image
-    image = Image.fromarray(outputs[0].cpu().numpy().astype('uint8'))
+        outputs = model.generate(**inputs)
     
-    return image
-
-# Define the Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Text to Image Generation")
-
-    with gr.Row():
-        prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here...")
-        num_inference_steps = gr.Slider(minimum=1, maximum=50, step=1, value=28, label="Number of Inference Steps")
-
-    with gr.Row():
-        generate_button = gr.Button("Generate Image")
-
-    result = gr.Image(label="Generated Image")
-
-    # Connect the function to the button
-    generate_button.click(
-        fn=generate_image,
-        inputs=[prompt, num_inference_steps],
-        outputs=result
-    )
-
-# Launch the Gradio app
-demo.launch()
+    # Convert output to a format suitable for Gradio
+    # This part may need to be adapted based on actual output format
+    return outputs
+
+# Set up Gradio interface
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=gr.Textbox(label="Enter prompt"),
+    outputs=gr.Image(type="pil", label="Generated Image"),
+    live=True
+)
+
+# Launch the app
+iface.launch()