Update app.py

app.py

@@ -1,7 +1,25 @@
 import gradio as gr
+import torch
 
 # Load the new model from Shakker-Labs
 demo = gr.load("models/Shakker-Labs/AWPortrait-FL")
 
-# Launch the Gradio interface
-demo.launch()
+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()