Update app.py

app.py

@@ -1,24 +1,24 @@
 import gradio as gr
 import torch
-from transformers import DalleMini, DalleMiniProcessor
+from transformers import DalleBartTokenizer, DalleBartForConditionalGeneration
 from PIL import Image
+import io
 
-# Load model and processor
-model_id = "dalle-mini/dalle-mega"
-model = DalleMini.from_pretrained(model_id)
-processor = DalleMiniProcessor.from_pretrained(model_id)
+# 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)
 
 # Function to generate image
 def generate_image(prompt, num_inference_steps=50):
-    inputs = processor(prompt, return_tensors="pt")
+    inputs = tokenizer(prompt, return_tensors="pt")
     
     # Generate images
     with torch.no_grad():
-        outputs = model.generate(**inputs, num_inference_steps=num_inference_steps)
+        outputs = model.generate(**inputs, num_beams=num_inference_steps)
     
-    # Convert to PIL image
-    image = processor.decode(outputs[0], skip_special_tokens=True)
-    image = Image.open(io.BytesIO(image))
+    # Convert tensor to PIL image
+    image = Image.fromarray(outputs[0].cpu().numpy().astype('uint8'))
     
     return image