Update app.py

app.py

@@ -3,6 +3,7 @@ import requests
 import json
 import io
 import os
+import uuid
 from PIL import Image
 import base64
 from prance import ResolvingParser
@@ -80,24 +81,9 @@ def create_gradio_app(api_spec, api_url):
     data_field = gr.State(value=names)
     inputs.append(data_field)
     print(output_schema)
-    if output_schema["type"] == "string":
-        if "format" in output_schema:
-            if(output_schema["format"] == "uri"):
-                output_component = gr.Image(label=output_schema["title"])
-            else:
-                output_component = gr.Textbox(label="Output")
-        else:
-            output_component = gr.Textbox(label="Output")
-        outputs.append(output_component)
-    elif output_schema["type"] == "array":
-        if "format" in output_schema["items"]:
-            if(output_schema["items"]["format"] == "uri"):
-                output_component = gr.Image(label=output_schema["title"])
-            else:
-                output_component = gr.Textbox(label=output_schema["title"])
-        else:
-            output_component = gr.Textbox(label=output_schema["title"])
-        outputs.append(output_component)
+    outputs.append(gr.Image(label=output_schema["title"], visible=True))
+    outputs.append(gr.Audio(label=output_schema["title"], visible=False))
+    outputs.append(gr.Textbox(label=output_schema["title"], visible=False))
     outputs.append(data_field)
     #else if there's multiple outputs
     
@@ -117,19 +103,31 @@ def create_gradio_app(api_spec, api_url):
             json_response = response.json()
             print(json_response)
             if "status" in json_response and json_response["status"] == "failed":
-                raise gr.Error("Failed to generate image")
+                raise gr.Error("Failed to generate output")
 
             output_images = []
             for output_uri in json_response["output"]:
-                base64_image = output_uri.replace("data:image/png;base64,", "")
-                image_data = base64.b64decode(base64_image)
-                image_stream = io.BytesIO(image_data)
-                output_images.append(Image.open(image_stream))
-
-            return output_images[0], keys
+                if output_uri.startswith("data:image"):
+                    # Process as image
+                    base64_data = output_uri.split(",", 1)[1]
+                    image_data = base64.b64decode(base64_data)
+                    image_stream = io.BytesIO(image_data)
+                    image = Image.open(image_stream)
+                    return gr.update(visible=True, value=image), gr.update(visible=False), gr.update(visible=False), keys
+                elif output_uri.startswith("data:audio"):
+                    # Process as audio
+                    base64_data = output_uri.split(",", 1)[1]
+                    audio_data = base64.b64decode(base64_data)
+                    audio_stream = io.BytesIO(audio_data)
+                    # Here you can save the audio or return the stream for further processing
+                    filename = f"{uuid.uuid4()}.wav"  # Change format as needed
+                    with open(filename, "wb") as audio_file:
+                        audio_file.write(audio_stream.getbuffer())
+                    return gr.update(visible=False), gr.update(visible=True, value=filename), gr.update(visible=False), keys
+            
         else:
             raise gr.Error("The submission failed!")
-    return gr.Interface(fn=predict, inputs=inputs, outputs=outputs if outputs else "textbox")
+    return gr.Interface(fn=predict, inputs=inputs, outputs=outputs)
 
 API_URL = "http://localhost:5000/predictions"
 app = create_gradio_app(api_spec, API_URL)