diff --git a/script_examples/websockets_api_example.py b/script_examples/websockets_api_example.py index 57a6cbd9..04c9fa21 100644 --- a/script_examples/websockets_api_example.py +++ b/script_examples/websockets_api_example.py @@ -41,15 +41,14 @@ def get_images(ws, prompt): continue #previews are binary data history = get_history(prompt_id)[prompt_id] - for o in history['outputs']: - for node_id in history['outputs']: - node_output = history['outputs'][node_id] - if 'images' in node_output: - images_output = [] - for image in node_output['images']: - image_data = get_image(image['filename'], image['subfolder'], image['type']) - images_output.append(image_data) - output_images[node_id] = images_output + for node_id in history['outputs']: + node_output = history['outputs'][node_id] + images_output = [] + if 'images' in node_output: + for image in node_output['images']: + image_data = get_image(image['filename'], image['subfolder'], image['type']) + images_output.append(image_data) + output_images[node_id] = images_output return output_images diff --git a/tests/inference/test_execution.py b/tests/inference/test_execution.py index ffc0c482..c7daddeb 100644 --- a/tests/inference/test_execution.py +++ b/tests/inference/test_execution.py @@ -95,17 +95,16 @@ class ComfyClient: pass # Probably want to store this off for testing history = self.get_history(prompt_id)[prompt_id] - for o in history['outputs']: - for node_id in history['outputs']: - node_output = history['outputs'][node_id] - result.outputs[node_id] = node_output - if 'images' in node_output: - images_output = [] - for image in node_output['images']: - image_data = self.get_image(image['filename'], image['subfolder'], image['type']) - image_obj = Image.open(BytesIO(image_data)) - images_output.append(image_obj) - node_output['image_objects'] = images_output + for node_id in history['outputs']: + node_output = history['outputs'][node_id] + result.outputs[node_id] = node_output + images_output = [] + if 'images' in node_output: + for image in node_output['images']: + image_data = self.get_image(image['filename'], image['subfolder'], image['type']) + image_obj = Image.open(BytesIO(image_data)) + images_output.append(image_obj) + node_output['image_objects'] = images_output return result diff --git a/tests/inference/test_inference.py b/tests/inference/test_inference.py index 141cc5c7..2e11778f 100644 --- a/tests/inference/test_inference.py +++ b/tests/inference/test_inference.py @@ -109,15 +109,14 @@ class ComfyClient: continue #previews are binary data history = self.get_history(prompt_id)[prompt_id] - for o in history['outputs']: - for node_id in history['outputs']: - node_output = history['outputs'][node_id] - if 'images' in node_output: - images_output = [] - for image in node_output['images']: - image_data = self.get_image(image['filename'], image['subfolder'], image['type']) - images_output.append(image_data) - output_images[node_id] = images_output + for node_id in history['outputs']: + node_output = history['outputs'][node_id] + images_output = [] + if 'images' in node_output: + for image in node_output['images']: + image_data = self.get_image(image['filename'], image['subfolder'], image['type']) + images_output.append(image_data) + output_images[node_id] = images_output return output_images