Skip to content

Commit

Permalink
fix: use better assumptions for images_received
Browse files Browse the repository at this point in the history
The presence of the key `output` in `data` during `send_sync` calls was enough to assert `data["output"]["images"]` was present and populated (i.e., we are post-inference and receiving the final images from comfyui). However, comfyui internals have changed so that `data["output"]["images"]` can be equal to `None` in ordinary circumstances during other stages of generation. As such, I have changed the logic to check for the key and for non-`None` values.
  • Loading branch information
tazlin committed Aug 23, 2024
1 parent 67f44c2 commit 2c8200e
Showing 1 changed file with 9 additions and 2 deletions.
11 changes: 9 additions & 2 deletions hordelib/comfy_horde.py
Original file line number Diff line number Diff line change
Expand Up @@ -650,8 +650,15 @@ def reconnect_input(cls, dct, input, output):
# This is the callback handler for comfy async events.
def send_sync(self, label: str, data: dict, _id: str) -> None:
# Get receive image outputs via this async mechanism
if "output" in data and "images" in data["output"]:
images_received = data["output"]["images"]
output = data.get("output", None)
images_received = None
if output is not None and "images" in output:
images_received = output.get("images", None)

if images_received is not None:
if len(images_received) == 0:
logger.warning("Received no output images from comfyui")

for image_info in images_received:
if not isinstance(image_info, dict):
logger.error(f"Received non dict output from comfyui: {image_info}")
Expand Down

0 comments on commit 2c8200e

Please sign in to comment.