Refactor swizzle_execute for async and sync handling

- Enhanced the swizzle_execute function to differentiate between asynchronous and synchronous execution paths based on the origin_execute function's nature. - Improved the structure of the swizzle_execute function to ensure proper handling of parameters and execution flow for both async and sync scenarios. - Maintained existing error handling while ensuring consistent behavior across execution types. These changes improve the flexibility and robustness of the execution handling in the application.
tweak: optional
2025-07-23 22:41:18 +08:00 · 2025-07-23 21:59:13 +08:00 · 2025-07-23 21:52:51 +08:00 · 2025-07-23 21:46:55 +08:00 · 2025-07-23 12:41:15 +08:00 · 2025-07-23 11:58:47 +08:00
1 changed files with 76 additions and 33 deletions
@@ -265,7 +265,7 @@ def clear_current_prompt(sid):
    streaming_prompt_metadata[sid].running_prompt_ids.clear()


-def post_prompt(json_data):
+async def post_prompt(json_data):
    prompt_server = server.PromptServer.instance
    json_data = prompt_server.trigger_on_prompt(json_data)

@@ -281,7 +281,14 @@ def post_prompt(json_data):

    if "prompt" in json_data:
        prompt = json_data["prompt"]
-        valid = execution.validate_prompt(prompt)
+        prompt_id = json_data.get("prompt_id") or str(uuid.uuid4())
+
+        try:
+            valid = await execution.validate_prompt(prompt_id, prompt)
+        except TypeError as e:
+            logger.warning(f"Trying old validate_prompt signature: {e}")
+            valid = execution.validate_prompt(prompt)
+
        extra_data = {}
        if "extra_data" in json_data:
            extra_data = json_data["extra_data"]
@@ -292,8 +299,6 @@ def post_prompt(json_data):
        if "client_id" in json_data:
            extra_data["client_id"] = json_data["client_id"]
        if valid[0]:
-            # if the prompt id is provided
-            prompt_id = json_data.get("prompt_id") or str(uuid.uuid4())
            outputs_to_execute = valid[2]
            prompt_server.prompt_queue.put(
                (number, prompt_id, prompt, extra_data, outputs_to_execute)
@@ -500,15 +505,15 @@ def send_prompt(sid: str, inputs: StreamingPrompt):

    prompt_id = str(uuid.uuid4())

-    prompt = {
-        "prompt": workflow_api,
-        "client_id": sid,  # "comfy_deploy_instance", #api.client_id
-        "prompt_id": prompt_id,
-        "extra_data": {"extra_pnginfo": {"workflow": workflow}},
-    }
+    # prompt = {
+    #     "prompt": workflow_api,
+    #     "client_id": sid,  # "comfy_deploy_instance", #api.client_id
+    #     "prompt_id": prompt_id,
+    #     "extra_data": {"extra_pnginfo": {"workflow": workflow}},
+    # }

    try:
-        res = post_prompt(prompt)
+        # res = post_prompt(prompt)
        inputs.running_prompt_ids.add(prompt_id)
        prompt_metadata[prompt_id] = SimplePrompt(
            status_endpoint=inputs.status_endpoint,
@@ -519,7 +524,7 @@ def send_prompt(sid: str, inputs: StreamingPrompt):
    except Exception as e:
        error_type = type(e).__name__
        stack_trace_short = traceback.format_exc().strip().split("\n")[-2]
-        stack_trace = traceback.format_exc().strip()
+        # stack_trace = traceback.format_exc().strip()
        logger.info(f"error: {error_type}, {e}")
        logger.info(f"stack trace: {stack_trace_short}")

@@ -595,7 +600,7 @@ async def comfy_deploy_run(request):
    )

    try:
-        res = post_prompt(prompt)
+        res = await post_prompt(prompt)
    except Exception as e:
        error_type = type(e).__name__
        stack_trace_short = traceback.format_exc().strip().split("\n")[-2]
@@ -664,7 +669,7 @@ async def stream_prompt(data, token):
    # log('info', "Begin prompt", prompt=prompt)

    try:
-        res = post_prompt(prompt)
+        res = await post_prompt(prompt)
    except Exception as e:
        error_type = type(e).__name__
        stack_trace_short = traceback.format_exc().strip().split("\n")[-2]
@@ -1266,22 +1271,11 @@ def handle_execute(class_type, last_node_id, prompt_id, server, unique_id):

 try:
    origin_execute = execution.execute
+    is_async = asyncio.iscoroutinefunction(origin_execute)

-    def swizzle_execute(
-        server,
-        dynprompt,
-        caches,
-        current_item,
-        extra_data,
-        executed,
-        prompt_id,
-        execution_list,
-        pending_subgraph_results,
-    ):
-        unique_id = current_item
-        class_type = dynprompt.get_node(unique_id)["class_type"]
-        last_node_id = server.last_node_id
-        result = origin_execute(
+    if is_async:
+
+        async def swizzle_execute(
            server,
            dynprompt,
            caches,
@@ -1291,12 +1285,61 @@ try:
            prompt_id,
            execution_list,
            pending_subgraph_results,
-        )
-        handle_execute(class_type, last_node_id, prompt_id, server, unique_id)
-        return result
+            pending_async_nodes,
+        ):
+            unique_id = current_item
+            class_type = dynprompt.get_node(unique_id)["class_type"]
+            last_node_id = server.last_node_id
+
+            result = await origin_execute(
+                server,
+                dynprompt,
+                caches,
+                current_item,
+                extra_data,
+                executed,
+                prompt_id,
+                execution_list,
+                pending_subgraph_results,
+                pending_async_nodes,
+            )
+
+            handle_execute(class_type, last_node_id, prompt_id, server, unique_id)
+            return result
+    else:
+
+        def swizzle_execute(
+            server,
+            dynprompt,
+            caches,
+            current_item,
+            extra_data,
+            executed,
+            prompt_id,
+            execution_list,
+            pending_subgraph_results,
+        ):
+            unique_id = current_item
+            class_type = dynprompt.get_node(unique_id)["class_type"]
+            last_node_id = server.last_node_id
+
+            result = origin_execute(
+                server,
+                dynprompt,
+                caches,
+                current_item,
+                extra_data,
+                executed,
+                prompt_id,
+                execution_list,
+                pending_subgraph_results,
+            )
+
+            handle_execute(class_type, last_node_id, prompt_id, server, unique_id)
+            return result

    execution.execute = swizzle_execute
-except Exception as e:
+except Exception:
    pass
Author	SHA1	Message	Date
KarrixLee	761e6db573	Refactor swizzle_execute for async and sync handling - Enhanced the swizzle_execute function to differentiate between asynchronous and synchronous execution paths based on the origin_execute function's nature. - Improved the structure of the swizzle_execute function to ensure proper handling of parameters and execution flow for both async and sync scenarios. - Maintained existing error handling while ensuring consistent behavior across execution types. These changes improve the flexibility and robustness of the execution handling in the application.	2025-07-23 22:41:18 +08:00
KarrixLee	2e527f74b5	tweak: optional	2025-07-23 21:59:13 +08:00
KarrixLee	7d8f4563af	Refactor post_prompt and origin_execute for synchronous handling - Modified the post_prompt function to call validate_prompt synchronously when a TypeError occurs, improving error handling. - Updated the origin_execute function to execute synchronously, ensuring consistent behavior during execution. These changes enhance the robustness of prompt validation and execution processes in the application.	2025-07-23 21:52:51 +08:00
KarrixLee	fba00a3c4d	Enhance error handling in post_prompt and origin_execute functions - Added try-except blocks in post_prompt to handle TypeErrors during prompt validation, allowing for fallback to an older signature. - Implemented similar error handling in the origin_execute function to manage potential TypeErrors, ensuring robust execution flow. - Improved logging to capture issues with function signatures, aiding in debugging. These changes improve the resilience of the application when dealing with prompt and execution validation.	2025-07-23 21:46:55 +08:00
KarrixLee	7f56d18599	Refactor handle_execute to support asynchronous execution - Updated the swizzle_execute function to be asynchronous, allowing for non-blocking execution of the origin_execute function. - Added missing parameters for pending_async_nodes in the swizzle_execute function call. - Ensured that the result from origin_execute is awaited, improving the handling of asynchronous operations. These changes enhance the performance and responsiveness of the execution handling in the application.	2025-07-23 12:41:15 +08:00
KarrixLee	bdb2c1b85d	Refactor post_prompt and send_prompt functions for async handling - Updated the post_prompt function to be asynchronous, allowing for non-blocking execution when validating prompts. - Adjusted calls to post_prompt in send_prompt and comfy_deploy_run to await the asynchronous execution, ensuring proper handling of prompt submissions. - Commented out unused prompt construction code to streamline the function. These changes enhance the performance and responsiveness of prompt handling in the application.	2025-07-23 11:58:47 +08:00
KarrixLee	ee658f90b2	Refactor post_prompt function to include prompt_id handling - Updated the post_prompt function in custom_routes.py to retrieve and validate a prompt_id from the incoming JSON data, defaulting to a new UUID if not provided. - Adjusted the validation call to use prompt_id alongside the prompt, ensuring proper identification and processing of prompts. These changes enhance the functionality of prompt handling within the application.	2025-07-23 11:42:16 +08:00