Update app.py

app.py

@@ -11,6 +11,8 @@ import hash_code_for_cached_output
 
 
 API_URL = os.environ.get("API_URL")
+TOKEN = os.environ.get("TOKEN")
+RESULT_URL = os.environ.get("RESULT_URL")
 supported_languages = ['zh', 'en', 'ja', 'ko', 'es', 'fr']
 supported_styles = {
     'zh': "zh_default",
@@ -139,42 +141,43 @@ def predict(prompt, style, audio_file_pth, speed, agree):
     }
     
     start = time.time()
-    # Send the data as a POST request
-    response = requests.post(API_URL, json=data, timeout=60)
+
+    headers = {
+        "Authorization": f"Bearer {TOKEN}"
+    }
+    
+    response = requests.post(API_URL, json=data, headers=headers, timeout=60)
     print(f'Get response successfully within {time.time() - start}')
 
-    # Check the response
-    if response.status_code == 200:
-        try:
-            json_data = json.loads(response.content)
-            text_hint += f"[ERROR] {json_data['error']} \n"
+    task_id = response.json()['task_id']
+    while True:
+        response = requests.post(RESULT_URL, json={'task_id': task_id}, headers=headers)
+        json_data = response.json()
+        status = json_data['status']
+        if status in ["CREATED", "RUNNING"]:
+            time.sleep(1)
+            continue
+        if status == 'FAILED':
+            text_hint += f"[HTTP ERROR] {json_data['error']} \n"
             gr.Warning(
-                f"[ERROR] {json_data['error']} \n"
+                f"[HTTP ERROR] {json_data['error']} \n"
             )
             return (
                 text_hint,
                 None,
                 None,
             )
-        except:
+        else:
+            decoded_bytes = base64.b64decode(json_data['result']['base64'].encode('utf-8'))
             with open(save_path, 'wb') as f:
-                f.write(response.content)
-    else:
-        text_hint += f"[HTTP ERROR] {response.status_code} - {response.text} \n"
-        gr.Warning(
-            f"[HTTP ERROR] {response.status_code} - {response.text} \n"
-        )
-        return (
-            text_hint,
-            None,
-            None,
-        )
-    text_hint += f'''Get response successfully \n'''
-    return (
-        text_hint,
-        save_path,
-        speaker_wav,
-    )
+                f.write(decoded_bytes)
+            
+            text_hint += f'''Get response successfully \n'''
+            return (
+                text_hint,
+                save_path,
+                speaker_wav,
+            )
 
 
 title = "MyShell OpenVoice V2"