save code

2025-10-04 12:53:01 +08:00
parent 20e6b294f1
commit 1fbb91eadb
2 changed files with 185 additions and 120 deletions
@@ -1,6 +1,73 @@
 import cv2
 import mediapipe as mp

+async def listen_node_output(prompt_id, target_node_id):
+    uri = "ws://localhost:8188/ws"
+    async with websockets.connect(uri) as websocket:
+        await websocket.send(json.dumps({"prompt_id": prompt_id}))
+        
+        while True:
+            message = await websocket.recv()
+            
+            # 1. 处理二进制数据
+            if isinstance(message, bytes):
+                try:
+                    message = message.decode("utf-8")
+                except UnicodeDecodeError:
+                    print(f"[{datetime.now().strftime('%H:%M:%S')}] 收到二进制数据 (长度: {len(message)} bytes)")
+                    continue
+            
+            # 2. 解析JSON
+            try:
+                data = json.loads(message)
+                print(f"cur data {data}")
+            except json.JSONDecodeError:
+                print(f"[{datetime.now().strftime('%H:%M:%S')}] 非JSON数据: {message[:100]}...")
+                continue
+
+            if data.get("type") == "progress":
+                pdata = data.get('data', {})
+                print(f"cur progress {pdata}")
+                # if value >= max:
+                #     return f"finished value:{value} max:{max}"
+
+            if data.get("type") == "progress_state":
+                nodes = data.get('data', {}).get('nodes', {})
+                for node in nodes:
+                    if node.get("state", None) != 'finished':
+                        print(f"progress_state node:{node}")
+            
+            # 3. 输出运行状态
+            if data.get("type") == "status":
+                print(f"[{datetime.now().strftime('%H:%M:%S')}] 系统状态: {data.get('data', {}).get('status', {})}")
+                status = data.get('data', {}).get('status', None)
+                print(f"statue:{status}")
+                if status:
+                    exec_info = status.get('exec_info', None)
+                    print(f"exec_info:{exec_info}")
+                    if exec_info:
+                        queue_remaining = exec_info.get('queue_remaining', None)
+                        print(f"queue_remaining:{queue_remaining}")
+                        if queue_remaining == 0:
+                            return "process end"
+                continue
+                
+            if data.get("type") == "executing":
+                node_id = data.get("data", {}).get("node")
+                progress = data.get("data", {}).get("progress", 0)
+                print(f"[{datetime.now().strftime('%H:%M:%S')}] 正在执行节点 {node_id} (进度: {progress:.0%})")
+            
+            # 4. 目标节点完成
+            if data.get("type") == "executed" and data.get("data", {}).get("node") == target_node_id:
+                output = data.get("data", {}).get("output")
+                print(f"[{datetime.now().strftime('%H:%M:%S')}] 节点 {target_node_id} 完成!")
+                return output
+
+
+
+    # output = asyncio.run(listen_node_output(prompt_id, "104"))
+    # print("最终输出:", output)
+
 # def is_thigh_visible(image_path):
 #     # 初始化MediaPipe Pose模型
 #     mp_pose = mp.solutions.pose