refac/fix: comfyui filter output node type

Co-Authored-By: Paul <239564541+mirrordna-reflection-protocol@users.noreply.github.com>
This commit is contained in:
Timothy Jaeryang Baek 2025-12-30 12:24:48 +04:00
parent c02451178c
commit aaea9a5956

View file

@ -64,8 +64,8 @@ def get_history(prompt_id, base_url, api_key):
return json.loads(response.read()) return json.loads(response.read())
def get_images(ws, prompt, client_id, base_url, api_key): def get_images(ws, workflow, client_id, base_url, api_key):
prompt_id = queue_prompt(prompt, client_id, base_url, api_key)["prompt_id"] prompt_id = queue_prompt(workflow, client_id, base_url, api_key)["prompt_id"]
output_images = [] output_images = []
while True: while True:
out = ws.recv() out = ws.recv()
@ -81,12 +81,16 @@ def get_images(ws, prompt, client_id, base_url, api_key):
history = get_history(prompt_id, base_url, api_key)[prompt_id] history = get_history(prompt_id, base_url, api_key)[prompt_id]
for node_id in history["outputs"]: for node_id in history["outputs"]:
node_output = history["outputs"][node_id] node_output = history["outputs"][node_id]
if "images" in node_output: if node_id in workflow and workflow[node_id].get("class_type") in [
for image in node_output["images"]: "SaveImage",
url = get_image_url( "PreviewImage",
image["filename"], image["subfolder"], image["type"], base_url ]:
) if "images" in node_output:
output_images.append({"url": url}) for image in node_output["images"]:
url = get_image_url(
image["filename"], image["subfolder"], image["type"], base_url
)
output_images.append({"url": url})
return {"data": output_images} return {"data": output_images}