Merge ffdc23c6dd into b08debceca

Addresses review feedback: the scheduler error path blamed available[0]
even when picking failed while inspecting a later ready node, misreporting
the node to the frontend.

Instead of threading the node id through the exception, make is_output and
is_async fully defensive. They are pure ordering heuristics, so a malformed
node (a FUNCTION typo, or schema-derived attributes that raise) just means
"not prioritized"; the node then runs through normal execution, where the
error is reported against the correct node. The stage_node_execution
try/except remains as a backstop only.

Add a test for a node whose attribute access raises during the heuristics.

comfy/text_encoders/llama.py

@@ -937,22 +937,41 @@ class BaseGenerate:
             return torch.argmax(logits, dim=-1, keepdim=True)
 
         # Sampling mode
-        if repetition_penalty != 1.0:
-            for i in range(logits.shape[0]):
-                for token_id in set(token_history):
-                    logits[i, token_id] *= repetition_penalty if logits[i, token_id] < 0 else 1/repetition_penalty
-
-        if presence_penalty is not None and presence_penalty != 0.0:
-            for i in range(logits.shape[0]):
-                for token_id in set(token_history):
-                    logits[i, token_id] -= presence_penalty
+        if len(token_history) > 0 and (repetition_penalty != 1.0 or (presence_penalty is not None and presence_penalty != 0.0)):
+            token_ids = torch.tensor(list(set(token_history)), device=logits.device)
+            token_logits = logits[:, token_ids]
+            if repetition_penalty != 1.0:
+                token_logits = torch.where(token_logits < 0, token_logits * repetition_penalty, token_logits / repetition_penalty)
+            if presence_penalty is not None and presence_penalty != 0.0:
+                token_logits = token_logits - presence_penalty
+            logits[:, token_ids] = token_logits
 
         if temperature != 1.0:
             logits = logits / temperature
 
         if top_k > 0:
-            indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
-            logits[indices_to_remove] = torch.finfo(logits.dtype).min
+            top_k = min(top_k, logits.shape[-1])
+            logits, top_indices = torch.topk(logits, top_k)
+
+            if min_p > 0.0:
+                probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)
+                top_probs, _ = probs_before_filter.max(dim=-1, keepdim=True)
+                min_threshold = min_p * top_probs
+                indices_to_remove = probs_before_filter < min_threshold
+                logits[indices_to_remove] = torch.finfo(logits.dtype).min
+
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 0] = False
+                indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
+                indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove)
+                logits[indices_to_remove] = torch.finfo(logits.dtype).min
+
+            probs = torch.nn.functional.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1, generator=generator)
+            return top_indices.gather(1, next_token)
 
         if min_p > 0.0:
             probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)

comfy_execution/graph.py

@@ -3,6 +3,7 @@ from typing import Type, Literal
 import nodes
 import asyncio
 import inspect
+import traceback
 from comfy_execution.graph_utils import is_link, ExecutionBlocker
 from comfy.comfy_types.node_typing import ComfyNodeABC, InputTypeDict, InputTypeOptions
 
@@ -263,7 +264,25 @@ class ExecutionList(TopologicalSort):
             }
             return None, error_details, ex
 
-        self.staged_node_id = self.ux_friendly_pick_node(available)
+        try:
+            self.staged_node_id = self.ux_friendly_pick_node(available)
+        except Exception as ex:
+            # Backstop: the ordering heuristics in ux_friendly_pick_node are
+            # defensive, but should anything else there fail, surface it as an
+            # execution error instead of letting it kill the prompt worker
+            # thread. Blame an available node (best effort).
+            blamed_node = self.dynprompt.get_display_node_id(available[0])
+            exception_type = type(ex).__qualname__
+            if type(ex).__module__ != "builtins":
+                exception_type = type(ex).__module__ + "." + exception_type
+            error_details = {
+                "node_id": blamed_node,
+                "exception_message": str(ex),
+                "exception_type": exception_type,
+                "traceback": traceback.format_tb(ex.__traceback__),
+                "current_inputs": []
+            }
+            return None, error_details, ex
         return self.staged_node_id, None, None
 
     def ux_friendly_pick_node(self, node_list):
@@ -271,19 +290,28 @@ class ExecutionList(TopologicalSort):
         # Technically this has no effect on the overall length of execution, but it feels better as a user
         # for a PreviewImage to display a result as soon as it can
         # Some other heuristics could probably be used here to improve the UX further.
+        # These node-ordering heuristics only affect *order*, never correctness.
+        # A malformed node (e.g. a FUNCTION typo, or a node whose schema-derived
+        # attributes raise) must not crash scheduling: failing a heuristic just
+        # means "not prioritized". The node then proceeds to normal execution,
+        # where the real error is raised and reported against the correct node.
         def is_output(node_id):
             class_type = self.dynprompt.get_node(node_id)["class_type"]
             class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
-            if hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE == True:
-                return True
-            return False
+            try:
+                return hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE == True
+            except Exception:
+                return False
 
         # If an available node is async, do that first.
         # This will execute the asynchronous function earlier, reducing the overall time.
         def is_async(node_id):
             class_type = self.dynprompt.get_node(node_id)["class_type"]
             class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
-            return inspect.iscoroutinefunction(getattr(class_def, class_def.FUNCTION))
+            try:
+                return inspect.iscoroutinefunction(getattr(class_def, class_def.FUNCTION))
+            except Exception:
+                return False
 
         for node_id in node_list:
             if is_output(node_id) or is_async(node_id):

requirements.txt

@@ -1,6 +1,6 @@
 comfyui-frontend-package==1.45.20
 comfyui-workflow-templates==0.11.2
-comfyui-embedded-docs==0.5.6
+comfyui-embedded-docs==0.5.7
 torch
 torchsde
 torchvision

tests-unit/execution_test/scheduler_malformed_node_test.py

@@ -0,0 +1,97 @@
+"""Regression tests for scheduler resilience to malformed nodes.
+
+A node whose FUNCTION points at a method that does not exist (e.g. a typo in a
+custom node) used to raise inside the scheduling heuristic, escaping the prompt
+worker's error handling and silently killing the worker thread. Scheduling must
+instead either proceed (so the error surfaces through normal execution) or report
+the failure as an execution error.
+"""
+import asyncio
+
+import nodes
+from comfy_execution.graph import DynamicPrompt, ExecutionList
+
+
+class _MalformedV1Node:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {}}
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "invert"  # the actual method below is misspelled
+    OUTPUT_NODE = True
+    CATEGORY = "Test"
+
+    def invvert(self):
+        return (None,)
+
+
+class _RaisingDescriptor:
+    def __get__(self, obj, owner):
+        raise RuntimeError("schema error")
+
+
+class _SchemaRaisesNode:
+    """A node whose schema-derived attribute access raises, as a broken V3 node would."""
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {}}
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "run"
+    OUTPUT_NODE = _RaisingDescriptor()
+    CATEGORY = "Test"
+
+    def run(self):
+        return (None,)
+
+
+class _FakeOutputCache:
+    def all_node_ids(self):
+        return set()
+
+    async def get(self, node_id):
+        return None
+
+
+def _make_execution_list(class_type, class_def):
+    nodes.NODE_CLASS_MAPPINGS[class_type] = class_def
+    prompt = {"1": {"class_type": class_type, "inputs": {}}}
+    execution_list = ExecutionList(DynamicPrompt(prompt), _FakeOutputCache())
+    execution_list.add_node("1")
+    return execution_list
+
+
+def test_malformed_function_does_not_crash_scheduler():
+    """A FUNCTION-typo node schedules without raising; the error surfaces later."""
+    execution_list = _make_execution_list("MalformedV1Node", _MalformedV1Node)
+    node_id, error, ex = asyncio.run(execution_list.stage_node_execution())
+    assert ex is None
+    assert error is None
+    assert node_id == "1"
+
+
+def test_schema_attribute_error_does_not_crash_scheduler():
+    """A node whose attribute access raises during heuristics still schedules."""
+    execution_list = _make_execution_list("SchemaRaisesNode", _SchemaRaisesNode)
+    node_id, error, ex = asyncio.run(execution_list.stage_node_execution())
+    assert ex is None
+    assert error is None
+    assert node_id == "1"
+
+
+def test_pick_node_failure_is_reported_not_raised():
+    """An unexpected scheduling error is returned as an error, not raised."""
+    execution_list = _make_execution_list("MalformedV1Node", _MalformedV1Node)
+
+    def raise_on_pick(_available):
+        raise RuntimeError("boom")
+
+    execution_list.ux_friendly_pick_node = raise_on_pick
+    node_id, error, ex = asyncio.run(execution_list.stage_node_execution())
+    assert node_id is None
+    assert isinstance(ex, RuntimeError)
+    assert error["node_id"] == "1"
+    assert error["exception_type"] == "RuntimeError"
+    assert error["exception_message"] == "boom"
+    assert error["traceback"]