fix: remove redundant useless loop (comfyanonymous#4656)

fix: potential error of undefined variable comfyanonymous#4650
jeremiahyan · Aug 28, 2024 · 34eda0f · 34eda0f
1 parent d31e226
commit 34eda0f
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 29 deletions.
diff --git 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
@@ -95,17 +95,16 @@ def run(self, graph):
  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
@@ -109,15 +109,14 @@ def get_images(self, graph, save=True):
  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