Update app.py

app.py

@@ -9,7 +9,40 @@ from utils.gradio_helpers import parse_outputs, process_outputs
 names = ['image', 'rotate_pitch', 'rotate_yaw', 'rotate_roll', 'blink', 'eyebrow', 'wink', 'pupil_x', 'pupil_y', 'aaa', 'eee', 'woo', 'smile', 'src_ratio', 'sample_ratio', 'crop_factor', 'output_format', 'output_quality']
 
 def predict(request: gr.Request, *args, progress=gr.Progress(track_tqdm=True)):
-    # ... (keep the predict function as is)
+    headers = {'Content-Type': 'application/json'}
+
+    payload = {"input": {}}
+    
+    base_url = "http://0.0.0.0:7860"
+    for i, key in enumerate(names):
+        value = args[i]
+        if value and (os.path.exists(str(value))):
+            value = f"{base_url}/file=" + value
+        if value is not None and value != "":
+            payload["input"][key] = value
+
+    response = requests.post("http://0.0.0.0:5000/predictions", headers=headers, json=payload)
+
+    if response.status_code == 201:
+        follow_up_url = response.json()["urls"]["get"]
+        response = requests.get(follow_up_url, headers=headers)
+        while response.json()["status"] != "succeeded":
+            if response.json()["status"] == "failed":
+                raise gr.Error("The submission failed!")
+            response = requests.get(follow_up_url, headers=headers)
+            time.sleep(1)
+    if response.status_code == 200:
+        json_response = response.json()
+        if(outputs[0].get_config()["name"] == "json"):
+            return json_response["output"]
+        predict_outputs = parse_outputs(json_response["output"])
+        processed_outputs = process_outputs(predict_outputs)
+        
+        return tuple(processed_outputs) if len(processed_outputs) > 1 else processed_outputs[0]
+    else:
+        if(response.status_code == 409):
+            raise gr.Error(f"Sorry, the Cog image is still processing. Try again in a bit.")
+        raise gr.Error(f"The submission failed! Error: {response.status_code}")
 
 def check_password(password):
     return password == "pixo"