Update app.py

app.py

@@ -1,15 +1,31 @@
 import gradio as gr
+import torch
 from transformers import DalleBartTokenizer, DalleBartForConditionalGeneration
+from PIL import Image
+import numpy as np
 
 # Load the tokenizer and model
 tokenizer = DalleBartTokenizer.from_pretrained("dalle-mini/dalle-mini")
 model = DalleBartForConditionalGeneration.from_pretrained("dalle-mini/dalle-mini")
 
+# Ensure the model is in evaluation mode and use GPU if available
+model.eval()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
 def generate_image(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs)
-    image = outputs[0]  # This will depend on your model's output format
-    return image
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+        outputs = model.generate(**inputs)
+        
+        # Convert the output tensor to an image
+        # Adjust this line according to the model output format
+        image = outputs[0]  # Assuming this is an image tensor
+        image = Image.fromarray(np.array(image))  # Convert tensor to numpy array then to image
+
+        return image
+    except Exception as e:
+        return str(e)
 
 # Define Gradio interface
 iface = gr.Interface(fn=generate_image,
@@ -20,7 +36,3 @@ iface = gr.Interface(fn=generate_image,
 
 if __name__ == "__main__":
     iface.launch()
-
-
-
-