Update app.py

app.py

@@ -79,10 +79,20 @@ def create_gradio_app(api_spec, api_url):
     print(names)
     data_field = gr.State(value=names)
     inputs.append(data_field)
+    print(output_schema)
     if output_schema["type"] == "string":
-        output_component = gr.Gallery(label="Output Image")
+        output_component = gr.Textbox(label="Output")
         outputs.append(output_component)
-        outputs.append(data_field)
+    elif output_schema["type"] == "array":
+        if "format" in output_schema["items"]:
+            if(output_schema["items"]["format"] == "uri"):
+                output_component = gr.Gallery(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(data_field)
     #else if there's multiple outputs
     
     def predict(*args):
@@ -96,22 +106,24 @@ def create_gradio_app(api_spec, api_url):
             payload["input"][key] = value
         print(payload)
         response = requests.post(api_url, headers={"Content-Type": "application/json"}, json=payload)
-        json_response = response.json()
-
-        if "status" in json_response and json_response["status"] == "failed":
-            raise gr.Error("Failed to generate image")
+        print(response)
+        if response.status_code == 200:
+            json_response = response.json()
+            print(json_response)
+            if "status" in json_response and json_response["status"] == "failed":
+                raise gr.Error("Failed to generate image")
 
-        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
-
-    return gr.Interface(fn=predict, inputs=inputs, outputs=outputs if outputs else "label")
+            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
+        else:
+            raise gr.Error("The submission failed!")
+    return gr.Interface(fn=predict, inputs=inputs, outputs=outputs if outputs else "textbox")
 
 API_URL = "http://localhost:5000/predictions"
 app = create_gradio_app(api_spec, API_URL)