Update app.py

app.py

@@ -21,7 +21,9 @@ 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)))
+    return image
     if isinstance(response, list):
         image_data = response[0]['image']
         image = Image.open(BytesIO(base64.b64decode(image_data)))