Update app.py

app.py

@@ -65,45 +65,50 @@ def process_image_workflow(image_path, positive_prompt, seed=-1):
     except requests.exceptions.RequestException as e:
         return f"任務建立連線錯誤: {e}", None
 
-    # --- 步驟 3: 輪詢任務狀態 ---
-    max_retries = 30
-    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_INFO_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
-
-            data = status_data.get("data")
-            if not data:
-                yield "查詢回傳 data 為空，請稍後重試。", None
-                return "查詢回傳 data 為空，請稍後重試。", None
-
-            current_status = data.get("taskStatus")
-            if current_status == "SUCCESS":
-                image_url = data["imageInfoList"][0]["fileUrl"]
-                yield "任務完成！", image_url
-                return "任務完成！", image_url
-            elif current_status == "FAILED":
-                error_msg = data.get("errorMsg", "Unknown error")
-                yield f"任務失敗！錯誤訊息: {error_msg}", None
-                return f"任務失敗！錯誤訊息: {error_msg}", None
-
-        except requests.exceptions.RequestException as e:
-            yield f"查詢連線錯誤: {e}", None
-            return f"查詢連線錯誤: {e}", None
-
-    yield "任務查詢超時，請使用 TaskID 手動查詢。", None
-    return "任務查詢超時，請使用 TaskID 手動查詢。", None
+ # --- 步驟 3: 輪詢任務狀態 (修改版) ---
+max_retries = 30
+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_INFO_URL, headers={"Content-Type": "application/json"}, data=json.dumps(status_payload))
+        
+        # 檢查 HTTP 狀態碼
+        if status_response.status_code == 404:
+            print("Task not found yet, retrying...")
+            continue  # 繼續下一次迴圈
+        
+        status_response.raise_for_status() # 如果不是 200 或 404，則拋出錯誤
+        status_data = status_response.json()
+        print("查詢回傳:", status_data)
+
+        if status_data.get("code") != 0:
+            return f"查詢任務失敗！訊息: {status_data.get('msg')}", None
+
+        data = status_data.get("data")
+        if not data:
+            return "查詢回傳 data 為空，請稍後重試。", None
+
+        current_status = data.get("taskStatus")
+        if current_status == "SUCCESS":
+            image_url = data["imageInfoList"][0]["fileUrl"]
+            yield "任務完成！", image_url
+            return "任務完成！", image_url
+        elif current_status == "FAILED":
+            error_msg = data.get("errorMsg", "Unknown error")
+            yield f"任務失敗！錯誤訊息: {error_msg}", None
+            return f"任務失敗！錯誤訊息: {error_msg}", None
+
+    except requests.exceptions.RequestException as e:
+        # 如果是連線錯誤，直接返回
+        yield f"查詢連線錯誤: {e}", None
+        return f"查詢連線錯誤: {e}", None
+
+yield "任務查詢超時，請使用 TaskID 手動查詢。", None
+return "任務查詢超時，請使用 TaskID 手動查詢。", None
 
 # --- Gradio 介面 ---
 with gr.Blocks() as demo: