From 55139fc8e06265c647df14e6a97a8bce90524b34 Mon Sep 17 00:00:00 2001 From: tazlin Date: Fri, 23 Aug 2024 16:53:02 -0400 Subject: [PATCH] fix: use better assumptions for `images_received` 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. --- hordelib/comfy_horde.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/hordelib/comfy_horde.py b/hordelib/comfy_horde.py index e30c275e..732ea8fe 100644 --- a/hordelib/comfy_horde.py +++ b/hordelib/comfy_horde.py @@ -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}")