import gradio as gr
# Load the new model from Shakker-Labs
demo = gr.load("models/Shakker-Labs/AWPortrait-FL")

def process_output(*outputs):
    # Assuming the model returns a tuple, and the first element is the image
    if isinstance(outputs, tuple):
        # Extract the image from the tuple
        image = outputs[0]
        return image
    return outputs

# Create a wrapper to process the output correctly
def wrapped_inference(*inputs):
    outputs = demo(*inputs)
    return process_output(*outputs)

# Launch the Gradio interface with the corrected output processing
gr.Interface(
    wrapped_inference,
    demo.inputs,
    demo.outputs,
).launch()