Update app.py

app.py

@@ -1,70 +1,110 @@
+import gradio as gr
 import requests
-import base64
+import json
 import io
 from PIL import Image
-import gradio as gr
-import json
+import base64
+from prance import ResolvingParser
+
 
-# API and Schema URLs
-API_URL = "http://localhost:5000/predictions"
 SCHEMA_URL = "http://localhost:5000/openapi.json"
+FILENAME = "openapi.json"
+schema_response = requests.get(SCHEMA_URL)
+openapi_spec = schema_response
+
+r = requests.get(SCHEMA_URL)
+print(r.content)
+with open(FILENAME, "wb") as f:
+    f.write(r.content)
+
+parser = ResolvingParser(FILENAME)
+print(parser.specification)
+
 
-def fetch_api_spec(url):
-    response = requests.get(url)
-    return response.json()
+def extract_property_info(prop):
+    # Handle 'allOf' by merging all contained properties (assuming simple case of enum merging)
+    if "allOf" in prop:
+        combined_prop = {}
+        for subprop in prop["allOf"]:
+            combined_prop.update(subprop)
+        prop = combined_prop
+    return prop
 
-def create_gradio_app_from_api_spec(api_spec):
-    input_properties = api_spec['components']['schemas']['Input']['properties']
+
+def create_gradio_app(api_spec, api_url):
     inputs = []
-    for prop, details in input_properties.items():
-        if 'enum' in details:
-            choices = details['enum']
-            inputs.append(gr.Dropdown(choices=choices, label=prop, value=details.get('default')))
-        elif details['type'] == 'integer':
-            inputs.append(gr.Number(label=prop, value=details.get('default'), minimum=details.get('minimum'), maximum=details.get('maximum')))
-        elif details['type'] == 'number':
-            inputs.append(gr.Slider(minimum=details.get('minimum'), maximum=details.get('maximum'), value=details.get('default'), label=prop))
-        elif details['type'] == 'string' and 'format' in details and details['format'] == 'uri':
-            inputs.append(gr.Image(label=prop))
-        elif details['type'] == 'string':
-            inputs.append(gr.Textbox(label=prop, value=details.get('default')))
-        elif details['type'] == 'boolean':
-            inputs.append(gr.Checkbox(label=prop, value=details.get('default')))
-    
-    def predict_function(**kwargs):
-        # Adjust the input kwargs for image inputs to convert them to the expected format by the API if needed
-        payload = {
-            "input": kwargs
-        }
-        print(payload)
-        response = requests.post(API_URL, headers={"Content-Type": "application/json"}, json=payload)
+    input_schema = api_spec["components"]["schemas"]["PredictionRequest"]["properties"][
+        "input"
+    ]["properties"]
+
+    for name, prop in input_schema.items():
+        prop = extract_property_info(
+            prop
+        )  # Extract property info correctly for 'allOf'
+        print(prop)
+        if "enum" in prop:
+            input_field = gr.Dropdown(
+                choices=prop["enum"], label=name, value=prop.get("default")
+            )
+        elif prop["type"] == "integer":
+            input_field = gr.Number(
+                label=name,
+                value=prop.get("default"),
+                minimum=prop.get("minimum"),
+                maximum=prop.get("maximum"),
+                step=1,
+            )
+        elif prop["type"] == "number":
+            input_field = gr.Number(
+                label=name,
+                value=prop.get("default"),
+                minimum=prop.get("minimum"),
+                maximum=prop.get("maximum"),
+            )
+        elif prop["type"] == "boolean":
+            input_field = gr.Checkbox(label=name, value=prop.get("default"))
+        elif prop["type"] == "string" and prop.get("format") == "uri":
+            input_field = gr.File(label=name)
+        else:  # Assuming string type for simplicity, can add more types as needed
+            input_field = gr.Textbox(label=name, value=prop.get("default"))
+        inputs.append(input_field)
+
+    def predict(**kwargs):
+        payload = {"input": {}}
+        for key, value in kwargs.items():
+            if isinstance(
+                value, io.BytesIO
+            ):  # For image inputs, convert to the desired format
+                value.seek(0)
+                value = (
+                    "data:image/jpeg;base64," + base64.b64encode(value.read()).decode()
+                )
+            payload["input"][key] = value
+
+        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":
+        if "status" in json_response and json_response["status"] == "failed":
             raise gr.Error("Failed to generate image")
 
-        output_spec = api_spec['components']['schemas']['Output']
-        if output_spec['items']['type'] == 'string' and output_spec['items']['format'] == 'uri':
-            outputs = []
-            for uri in json_response["output"]:
-                if uri.startswith("data:image"):
-                    base64_image = uri.split(",")[1]  # Strip the prefix part
-                    image_data = base64.b64decode(base64_image)
-                    image_stream = io.BytesIO(image_data)
-                    image = Image.open(image_stream)
-                    outputs.append(image)
-                else:
-                    outputs.append(uri)
-            return outputs
-        else:
-            return json_response["output"]
-    
-    iface = gr.Interface(fn=predict_function, inputs=inputs, outputs=gr.outputs.Image(type="pil"), title=api_spec['info']['title'])
-    return iface
-
-# Fetch API Specification
-api_spec = fetch_api_spec(SCHEMA_URL)
-
-# Create and Launch Gradio App
-gradio_app = create_gradio_app_from_api_spec(api_spec)
-gradio_app.launch()
+        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
+
+    output_component = gr.Gallery(label="Output Images")
+    return gr.Interface(fn=predict, inputs=inputs, outputs=output_component)
+
+
+# Use the modified function with the API URL
+api_spec = parser.specification
+
+API_URL = "http://localhost:5000/predictions"
+app = create_gradio_app(api_spec, API_URL)
+app.launch()