64 lines
1.8 KiB
Python
64 lines
1.8 KiB
Python
import time
|
|
import requests
|
|
import json
|
|
import uuid
|
|
|
|
COMFY_URL = "http://localhost:8188"
|
|
client_id = str(uuid.uuid4())
|
|
|
|
with open("text_to_image.json", "r") as f:
|
|
prompt = json.load(f)
|
|
|
|
response = requests.post(
|
|
f"{COMFY_URL}/prompt", json={"prompt": prompt, "client_id": client_id}, timeout=10
|
|
)
|
|
|
|
result = response.json()
|
|
|
|
if "prompt_id" in result:
|
|
prompt_id = result["prompt_id"]
|
|
print(f"The task has been submit prompt ID: {prompt_id}")
|
|
|
|
history_url = f"{COMFY_URL}/history/{prompt_id}"
|
|
while True:
|
|
history_resp = requests.get(history_url)
|
|
history_data = history_resp.json()
|
|
|
|
if prompt_id in history_data:
|
|
outputs = history_data[prompt_id]["outputs"]
|
|
print(f"The task has finished, output message: {outputs}")
|
|
break
|
|
time.sleep(1)
|
|
|
|
stats_response = requests.post(
|
|
f"{COMFY_URL}/exec-timer/stat",
|
|
json={
|
|
"prompt_id": prompt_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
stats_response.raise_for_status()
|
|
stats_dict = stats_response.json()
|
|
|
|
assert stats_dict["prompt_id"] == prompt_id
|
|
resp_msg = stats_dict["msg"]
|
|
resp_code = stats_dict["code"]
|
|
node_stats = stats_dict["data"]
|
|
|
|
TABLE_HEADER = ["GLOBAL_STEP", "CLASS TYPE", "Execution Time(ms)", "Unique ID"]
|
|
print(
|
|
f"{TABLE_HEADER[0]:<30} {TABLE_HEADER[1]:<30} {TABLE_HEADER[2]:<30} {TABLE_HEADER[3]:<30}"
|
|
)
|
|
|
|
for key, value in node_stats.items():
|
|
if key == "peak_memory_usage" or key == "total_time":
|
|
print(f"{key:<30} {value:<30}")
|
|
else:
|
|
print(
|
|
f"{key:<30} {value['class_type']:<30} {value['exec_time_ms']:<30} {value['unique_id']:<30}"
|
|
)
|
|
|
|
else:
|
|
print(f"Failed to submit task, {result.get('error', 'Unknown error')}")
|