Update app.py

app.py

@@ -1,8 +1,6 @@
 import gradio as gr
 import requests
 import json
-import asyncio
-import websockets
 import time
 
 # --- 配置資訊 ---
@@ -10,6 +8,8 @@ API_KEY = "d72ef10876b9403a9479c6e31f27a4be".strip()
 WORKFLOW_ID = "1958842411961233410"
 UPLOAD_URL = "https://www.runninghub.cn/task/openapi/upload"
 CREATE_TASK_URL = "https://www.runninghub.cn/task/openapi/create"
+GET_TASK_STATUS_URL = "https://www.runninghub.cn/task/openapi/status"
+GET_TASK_OUTPUTS_URL = "https://www.runninghub.cn/task/openapi/outputs"
 HEADERS = {"Host": "www.runninghub.cn"}
 
 # 工作流節點 ID
@@ -18,33 +18,33 @@ POSITIVE_PROMPT_NODE_ID = "144"
 SEED_NODE_ID = "132"
 
 # --- 核心函式 ---
-async def process_image_workflow_async(image_path, seed=-1):
+def process_image_workflow(image_path, seed=-1):
     positive_prompt = "a charming young woman with bright eyes and a subtle smile"
 
     if image_path is None:
         return "請先上傳圖片！", None
 
-    # --- 步驟 1: 上傳圖片 (使用非同步) ---
+    # --- 步驟 1: 上傳圖片 ---
     yield "正在上傳圖片...", None
     try:
         with open(image_path, "rb") as f:
             files = {"file": f}
             data = {"apiKey": API_KEY, "fileType": "image"}
-            # 使用 requests 函式庫發送同步請求
             upload_response = requests.post(UPLOAD_URL, headers=HEADERS, data=data, files=files)
             upload_response.raise_for_status()
             upload_data = upload_response.json()
+            print("上傳回傳:", upload_data)
 
             if upload_data.get("code") != 0:
                 return f"圖片上傳失敗！訊息: {upload_data.get('msg')}", None
-            
+
             rh_image_path = upload_data["data"]["fileName"]
             yield f"圖片上傳成功！fileName: {rh_image_path}", None
 
     except requests.exceptions.RequestException as e:
         return f"上傳連線錯誤: {e}", None
 
-    # --- 步驟 2: 建立工作流任務 (使用非同步) ---
+    # --- 步驟 2: 建立工作流任務 ---
     yield "正在建立任務...", None
     node_info_list = [
         {"nodeId": LOAD_IMAGE_NODE_ID, "fieldName": "image", "fieldValue": rh_image_path},
@@ -57,50 +57,69 @@ async def process_image_workflow_async(image_path, seed=-1):
         create_response = requests.post(CREATE_TASK_URL, headers={"Content-Type": "application/json"}, data=json.dumps(payload))
         create_response.raise_for_status()
         create_data = create_response.json()
+        print("建立任務回傳:", create_data)
 
         if create_data.get("code") != 0:
             return f"任務建立失敗！訊息: {create_data.get('msg')}", None
-        
+
         task_id = create_data["data"]["taskId"]
-        net_wss_url = create_data["data"]["netWssUrl"]
         yield f"任務建立成功！TaskID: {task_id}", None
 
     except requests.exceptions.RequestException as e:
         return f"任務建立連線錯誤: {e}", None
 
-    # --- 步驟 3: 使用 WebSocket 監聽任務狀態 ---
-    yield "正在等待任務完成...", None
+    # --- 步驟 3: 輪詢任務狀態 ---
+    max_retries = 40
+    retry_interval = 5
+
+    for i in range(max_retries):
+        yield f"查詢任務狀態... (第 {i+1} 次)", None
+        time.sleep(retry_interval)
+        try:
+            status_payload = {"apiKey": API_KEY, "taskId": task_id}
+            status_response = requests.post(GET_TASK_STATUS_URL, headers={"Content-Type": "application/json"}, data=json.dumps(status_payload))
+            status_response.raise_for_status()
+            status_data = status_response.json()
+            print("查詢狀態回傳:", status_data)
+
+            if status_data.get("code") != 0:
+                yield f"查詢任務失敗！訊息: {status_data.get('msg')}", None
+                return f"查詢任務失敗！訊息: {status_data.get('msg')}", None
+
+            current_status = status_data.get("data")
+            if current_status == "SUCCESS":
+                yield "任務狀態為成功，正在取得圖片結果...", None
+                break  # 跳出輪詢迴圈
+            elif current_status == "FAILED":
+                yield f"任務失敗！", None
+                return f"任務失敗！", None
+        
+        except requests.exceptions.RequestException as e:
+            yield f"查詢連線錯誤: {e}", None
+            return f"查詢連線錯誤: {e}", None
+    else:
+        yield "任務查詢超時，請使用 TaskID 手動查詢。", None
+        return "任務查詢超時，請使用 TaskID 手動查詢。", None
+
+    # --- 步驟 4: 取得任務結果 ---
     try:
-        async with websockets.connect(net_wss_url, extra_headers={"Host": "www.runninghub.ai"}) as ws:
-            while True:
-                message = await asyncio.wait_for(ws.recv(), timeout=1200) # 設定超時時間，例如20分鐘
-                data = json.loads(message)
-
-                if data.get("type") == "execution_status":
-                    status = data.get("data", {}).get("taskStatus")
-                    if status == "SUCCESS":
-                        yield "任務完成！", data["data"]["imageInfoList"][0]["fileUrl"]
-                        return "任務完成！", data["data"]["imageInfoList"][0]["fileUrl"]
-                    elif status == "FAILED":
-                        error_msg = data["data"].get("errorMsg", "Unknown error")
-                        return f"任務失敗！錯誤訊息: {error_msg}", None
-    
-    except asyncio.TimeoutError:
-        return "WebSocket 連線超時，請使用 TaskID 手動查詢。", None
-    except websockets.exceptions.ConnectionClosedError as e:
-        return f"WebSocket 連線關閉，錯誤: {e}", None
+        outputs_payload = {"apiKey": API_KEY, "taskId": task_id}
+        outputs_response = requests.post(GET_TASK_OUTPUTS_URL, headers={"Content-Type": "application/json"}, data=json.dumps(outputs_payload))
+        outputs_response.raise_for_status()
+        outputs_data = outputs_response.json()
+        print("取得結果回傳:", outputs_data)
+
+        if outputs_data.get("code") != 0:
+            return f"取得結果失敗！訊息: {outputs_data.get('msg')}", None
+
+        image_url = outputs_data["data"][0]["fileUrl"]
+        yield "任務完成！", image_url
+        return "任務完成！", image_url
+
+    except requests.exceptions.RequestException as e:
+        return f"取得結果連線錯誤: {e}", None
     except Exception as e:
-        return f"WebSocket 監聽錯誤: {e}", None
-    
-# Gradio 需要一個包裝函式來處理非同步
-def process_image_workflow(image_path, seed=-1):
-    loop = asyncio.get_event_loop()
-    if loop.is_running():
-        # 如果事件迴圈已經在運行，直接在現有迴圈中執行
-        return asyncio.run(process_image_workflow_async(image_path, seed))
-    else:
-        # 如果事件迴圈沒有運行，則啟動一個
-        return loop.run_until_complete(process_image_workflow_async(image_path, seed))
+        return f"處理結果時發生錯誤: {e}", None
 
 # --- Gradio 介面 ---
 with gr.Blocks() as demo: