Update app.py

app.py

@@ -1,24 +1,33 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import base64
+from io import BytesIO
+from PIL import Image
 
-# List of models and their corresponding names
+# Define the list of models
 models = [
     "CompVis/stable-diffusion-v1-4",
     "runwayml/stable-diffusion-v1-5",
-    "hakurei/waifu-diffusion",
-    "prompthero/openjourney",
-    "jacek-huszar/$(IMG2TEXT_MODEL_NAME)",
-    "cjwbw/canvae-text-to-image",
-    "tdbooth/hollow-kingdom-ddim",
-    "kuprel/min-dalle",
-    "osanseviero/stable-diffusion-webui",
-    "lambdal/text-to-image"
+    "stabilityai/stable-diffusion-2-1-base",
+    "stabilityai/stable-diffusion-2-1",
+    "CompVis/ldm-text2im-large-256",
+    "lambdalabs/sd-text2img-base-2-0",
+    "ZB-Tech/Text-to-Image",
+    "cloudqi/cqi_text_to_image_pt_v0",
+    "kothariyashhh/GenAi-Texttoimage",
+    "sairajg/Text_To_Image"
 ]
 
 def generate_image(prompt, model_name):
     client = InferenceClient(model_name)
     response = client.text_to_image(prompt)
-    return response.image
+    
+    if isinstance(response, list):
+        image_data = response[0]['image']
+        image = Image.open(BytesIO(base64.b64decode(image_data)))
+        return image
+    else:
+        return "Failed to generate image."
 
 # Create Gradio Interface
 with gr.Blocks() as demo: