app.py

import gradio as gr
import requests
import json
import io
import os
from PIL import Image
import base64
from prance import ResolvingParser

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)
api_spec = parser.specification
print(parser.specification)

def extract_property_info(prop):
    combined_prop = {}
    merge_keywords = ["allOf", "anyOf", "oneOf"]
    
    for keyword in merge_keywords:
        if keyword in prop:
            for subprop in prop[keyword]:
                combined_prop.update(subprop)
            del prop[keyword]
    
    if not combined_prop:
        combined_prop = prop.copy()
    
    for key in ['description', 'default']:
        if key in prop:
            combined_prop[key] = prop[key]
    
    return combined_prop

def sort_properties_by_order(properties):
    ordered_properties = sorted(properties.items(), key=lambda x: x[1].get('x-order', float('inf')))
    return ordered_properties

def create_gradio_app(api_spec, api_url):
    inputs = []
    outputs = []
    input_schema = api_spec["components"]["schemas"]["Input"]["properties"]
    output_schema = api_spec["components"]["schemas"]["Output"]
    ordered_input_schema = sort_properties_by_order(input_schema)
    names = []
    for name, prop in ordered_input_schema:

        prop = extract_property_info(prop)
        if "enum" in prop:
            input_field = gr.Dropdown(
                choices=prop["enum"], label=prop.get("title"), info=prop.get("description"), value=prop.get("default")
            )
        elif prop["type"] == "integer":
            input_field = gr.Slider(
                label=prop.get("title"), info=prop.get("description"), value=prop.get("default"),
                minimum=prop.get("minimum"), maximum=prop.get("maximum"), step=1,
            )
        elif prop["type"] == "number":
            input_field = gr.Slider(
                label=prop.get("title"), info=prop.get("description"), value=prop.get("default"),
                minimum=prop.get("minimum"), maximum=prop.get("maximum"),
            )
        elif prop["type"] == "boolean":
            input_field = gr.Checkbox(label=prop.get("title"), info=prop.get("description"), value=prop.get("default"))
        elif prop["type"] == "string" and prop.get("format") == "uri":
            input_field = gr.File(label=prop.get("title"))
        else:
            input_field = gr.Textbox(label=prop.get("title"), info=prop.get("description"))
        inputs.append(input_field)
        names.append(name)
    print(names)
    data_field = gr.State(value=names)
    inputs.append(data_field)
    if output_schema["type"] == "string":
        output_component = gr.Gallery(label="Output Image")
        outputs.append(output_component)
        outputs.append(data_field)
    #else if there's multiple outputs
    
    def predict(*args):
        print(args)
        keys = args[-1]
        payload = {"input": {}}
        for i, key in enumerate(keys):
            value = args[i]
            if(os.path.exists(value)):
                value = "http://localhost:7860/file=" + value
            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")

        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")


API_URL = "http://localhost:5000/predictions"
app = create_gradio_app(api_spec, API_URL)
app.launch(share=True)