diff --git a/tensors/comfyui.py b/tensors/comfyui.py
index 315e97d..c77400d 100644
--- a/tensors/comfyui.py
+++ b/tensors/comfyui.py
@@ -640,13 +640,9 @@ def run_workflow(
                 if on_progress:
                     on_progress(step, total, status)
 
-            return _wait_for_completion_ws(
-                prompt_id, base_url, client_id, timeout, on_progress=_console_progress
-            )
+            return _wait_for_completion_ws(prompt_id, base_url, client_id, timeout, on_progress=_console_progress)
     else:
-        return _wait_for_completion_ws(
-            prompt_id, base_url, client_id, timeout, on_progress=on_progress
-        )
+        return _wait_for_completion_ws(prompt_id, base_url, client_id, timeout, on_progress=on_progress)
 
 
 # ============================================================================