From 34eda0f853daffdfcab04e1b3187de40f1d30bbf Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <128333288+ltdrdata@users.noreply.github.com>
Date: Thu, 29 Aug 2024 06:46:30 +0900
Subject: [PATCH] fix: remove redundant useless loop (#4656)

fix: potential error of undefined variable

https://github.com/comfyanonymous/ComfyUI/discussions/4650
---
 script_examples/websockets_api_example.py | 17 ++++++++---------
 tests/inference/test_execution.py         | 21 ++++++++++-----------
 tests/inference/test_inference.py         | 17 ++++++++---------
 3 files changed, 26 insertions(+), 29 deletions(-)

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