Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import DallEProcessor, DallEModel
+from transformers import DALLMiniProcessor, DALLMiniModel
 from PIL import Image
 import requests
 from flask import Flask, request, jsonify
@@ -8,9 +8,9 @@ import io
 app = Flask(__name__)
 
 # Initialize the DALL-E mini model and processor
-model_name = "dalle-mini/dalle-mini"
-processor = DallEProcessor.from_pretrained(model_name)
-model = DallEModel.from_pretrained(model_name)
+model_name = "flax-community/dalle-mini"
+processor = DALLMiniProcessor.from_pretrained(model_name)
+model = DALLMiniModel.from_pretrained(model_name)
 
 @app.route('/generate', methods=['POST'])
 def generate_image():
@@ -19,11 +19,11 @@ def generate_image():
 
     # Generate images from prompt
     inputs = processor(text=[prompt], return_tensors="pt")
-    outputs = model(**inputs)
+    outputs = model.generate(**inputs)
 
     # Post-process the generated image
-    generated_image = outputs.logits.argmax(-1)[0]
-    image = Image.fromarray(generated_image)
+    generated_image = outputs[0]
+    image = Image.fromarray(generated_image.numpy().astype('uint8'))
 
     # Save image to a BytesIO object
     img_byte_arr = io.BytesIO()