97 lines
2.7 KiB
Python
97 lines
2.7 KiB
Python
import requests
|
|
import json
|
|
import sys
|
|
import time
|
|
import websocket
|
|
import os
|
|
|
|
# additional Python packages required to run this script:
|
|
# websocket-client==1.8.0
|
|
|
|
def on_message(ws, message):
|
|
print("Received message:", message)
|
|
try:
|
|
data = json.loads(message)
|
|
# Look for outputs in the message
|
|
outputs = None
|
|
if "data" in data and "outputs" in data["data"]:
|
|
outputs = data["data"]["outputs"]
|
|
elif "outputs" in data:
|
|
outputs = data["outputs"]
|
|
if outputs:
|
|
# Find the STRING output (from S3SaveNode)
|
|
for node_id, node_outputs in outputs.items():
|
|
for output in node_outputs:
|
|
if "STRING" in output:
|
|
print("S3SaveNode output STRING:", output["STRING"])
|
|
ws.close()
|
|
return
|
|
# Optionally, close on execution_end
|
|
if data.get("type") == "execution_end":
|
|
print("Execution ended, but no output found.")
|
|
ws.close()
|
|
except Exception as e:
|
|
print("Error parsing message as JSON:", e)
|
|
|
|
def on_error(ws, error):
|
|
print("WebSocket error:", error)
|
|
|
|
def on_close(ws, close_status_code, close_msg):
|
|
print("WebSocket closed")
|
|
|
|
def on_open(ws):
|
|
print("WebSocket connection opened, waiting for output...")
|
|
|
|
if len(sys.argv) != 4:
|
|
print("Usage: python test_comfyui_node.py <workflow.json> <bearer_token> <comfyui_url>")
|
|
sys.exit(1)
|
|
|
|
workflow_path = sys.argv[1]
|
|
bearer_token = sys.argv[2]
|
|
comfyui_url = sys.argv[3]
|
|
if comfyui_url.endswith("/"):
|
|
comfyui_url = comfyui_url[:-1]
|
|
|
|
# Read workflow.json
|
|
with open(workflow_path, "r") as f:
|
|
workflow = json.load(f)
|
|
|
|
# Wrap workflow in a 'prompt' key if not already wrapped
|
|
if "prompt" not in workflow:
|
|
payload = {"prompt": workflow}
|
|
else:
|
|
payload = workflow
|
|
|
|
# Submit workflow to /prompt
|
|
url = f"{comfyui_url}/prompt"
|
|
headers = {
|
|
"Authorization": f"Bearer {bearer_token}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
response = requests.post(url, headers=headers, json=payload)
|
|
if response.status_code != 200:
|
|
print("Failed to submit prompt:", response.text)
|
|
sys.exit(1)
|
|
|
|
resp_json = response.json()
|
|
prompt_id = resp_json.get("prompt_id")
|
|
if not prompt_id:
|
|
print("No prompt_id in response:", resp_json)
|
|
sys.exit(1)
|
|
|
|
print("Submitted prompt, prompt_id:", prompt_id)
|
|
|
|
# Connect to /ws and listen for output
|
|
ws_url = comfyui_url.replace("http://", "ws://").replace("https://", "wss://") + "/ws"
|
|
ws = websocket.WebSocketApp(
|
|
ws_url,
|
|
header=[f"Authorization: Bearer {bearer_token}"],
|
|
on_open=on_open,
|
|
on_message=on_message,
|
|
on_error=on_error,
|
|
on_close=on_close
|
|
)
|
|
|
|
globals()['prompt_id'] = prompt_id
|
|
ws.run_forever()
|