Update app.py

app.py

@@ -21,12 +21,12 @@ models = [
 def generate_image(prompt, model_name):
     client = InferenceClient(model_name)
     response = client.text_to_image(prompt)
-    image_data = response[0]['image']
-    image = Image.open(BytesIO(base64.b64decode(image_data)))
+    #image_data = response[0]['image']
+    image = Image.open(BytesIO(base64.b64decode(response)))
     return image
     if isinstance(response, list):
         image_data = response[0]['image']
-        image = Image.open(BytesIO(base64.b64decode(image_data)))
+        image = Image.open(BytesIO(base64.b64decode(response)))
         return image
     else:
         return "Failed to generate image."