Support Bounded Feedback Loops in the DAG Execution Engine#14550
Support Bounded Feedback Loops in the DAG Execution Engine#14550Prohect wants to merge 4 commits into
Conversation
Allow sampler nodes' internal iteration variables (e.g. step_index)
to flow back upstream through ComfyMathExpression nodes to control
per-step parameters (cfg, s_noise, eta, r) without triggering a
dependency cycle error.
Architecture: Two-level cycle handling
- Static validation: _is_bounded_feedback_cycle() allows cycles
where any node declares BOUNDED_FEEDBACK
- Graph building: _is_feedback_output() skips strong links for
declared feedback sockets, records them in feedback_links
Multi-hop chain walking via _build_feedback_fns() resolves
expression->CFGGuider/Sampler chains with simple_eval + MATH_FUNCTIONS,
composing per-step fn(step, total_steps) callables.
Sampler functions now re-read s_noise/eta/r each iteration via
_init_dynamic_options() / _refresh_dynamic_params() / _apply_dynamic_s_noise().
KSAMPLER.sample() conditionally injects mutable extra_options ref.
Safety: _dynamic_sampler_options popped at function top before model() calls.
One-line opt-in: BOUNDED_FEEDBACK = {'step_index'} on any node.
Prohect
requested review from
Kosinkadink,
alexisrolland,
comfyanonymous,
guill,
kijai and
rattus128
as code owners
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThis PR implements a "bounded feedback" iteration mechanism for the ComfyUI sampling pipeline. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
execution.py (1)
744-751: 💤 Low valueOutput socket 0 is hardcoded in
_find_consumers.The consumer search only matches links from
source_idat output socket0. This is correct forComfyMathExpression(single output), but could silently miss connections if the feature is ever extended to nodes with multiple outputs.Consider parameterizing the output socket
- def _find_consumers(source_id): + def _find_consumers(source_id, source_socket=0): consumers = [] for nid, n in prompt.items(): for iname, ival in n.get("inputs", {}).items(): if isinstance(ival, list) and len(ival) == 2 \ - and ival[0] == source_id and ival[1] == 0: + and ival[0] == source_id and ival[1] == source_socket: consumers.append((nid, n.get("class_type"), iname)) return consumers🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@execution.py` around lines 744 - 751, The _find_consumers function has the output socket hardcoded to 0 in the condition ival[1] == 0, which will only find consumers connected to output socket 0 and miss connections from other outputs. Parameterize the function by adding an output_socket parameter to _find_consumers alongside the existing source_id parameter, replace the hardcoded 0 with this new parameter in the comparison, and update all calls to _find_consumers throughout the code to pass the appropriate output socket index as an argument.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@execution.py`:
- Around line 744-751: The _find_consumers function has the output socket
hardcoded to 0 in the condition ival[1] == 0, which will only find consumers
connected to output socket 0 and miss connections from other outputs.
Parameterize the function by adding an output_socket parameter to
_find_consumers alongside the existing source_id parameter, replace the
hardcoded 0 with this new parameter in the comparison, and update all calls to
_find_consumers throughout the code to pass the appropriate output socket index
as an argument.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: dd5a4a3b-a9e5-45db-b35e-33fbf67738fd
📒 Files selected for processing (5)
comfy/k_diffusion/sampling.pycomfy/samplers.pycomfy_execution/graph.pycomfy_extras/nodes_custom_sampler.pyexecution.py
|
a demo workflow tested: it's tested on Windows 11 25h2 26200.8457, Nvidia RTX 3080 10G(610.62), result: |
- _FEEDBACK_DEFAULTS: add module-level docstring comment - _find_consumers: add source_socket parameter (default 0), add docstring - _is_sampler_target: add docstring
|
Updates pushed in 4fece5b:
|
1 participant
Closes #14551
Summary
Make ComfyUI's DAG execution engine support bounded feedback loops — where a sampler node's internal iteration variable (like
step_index) can flow back upstream through the graph viaComfyMathExpressionnodes to control its own parameters (CFG, s_noise, eta, r) on a per-step basis, without triggering a dependency cycle error.Example workflow
The core insight: the sampler loop is bounded (N steps, then terminates), so this isn't an infinite cycle. The DAG should recognize this and allow it.
Architecture
Two-Level Cycle Handling
Level 1 — Static Validation (
validate_inputsDFS)_is_bounded_feedback_cycle()checks whether at least one node in the detected cycle declaresBOUNDED_FEEDBACK. If so, the cycle is allowed — the repeated node is marked valid and traversal continues.Level 2 — Graph Building (
TopologicalSort.add_node)_is_feedback_output()checks whether a link's source socket is a declared bounded-iteration output. Feedback edges are:execution_list.feedback_linksMulti-Hop Chain Walking (
_build_feedback_fns)DFS walks from the feedback-receiving
ComfyMathExpressionnode through intermediate MathExpression nodes to find terminal CFGGuider / SamplerXXX targets. Each chain is composed into a callablefn(step, total_steps)evaluated withsimple_eval+ all 15 math functions.Per-Step Re-Read in Sampler Functions
_init_dynamic_options()pops_dynamic_sampler_optionsfromextra_args_refresh_dynamic_params()re-readss_noiseandetawith proper noise_scale handling_apply_dynamic_s_noise()re-readss_noisefrom the mutable dictSafety
KSAMPLER.sample()only injects_dynamic_sampler_optionswhen_feedback_param_fnsexists_dynamic_sampler_optionsat function top, preventing model call leaksSamplerCustomAdvancedfires after the model call for stepi, passingstep + 1Bootstrap Flow
feedback_linksget_input_data()supplies type-appropriate defaults for feedback sources not yet executedfeedback_links, calls_build_feedback_fns()Files Changed
execution.py_FEEDBACK_DEFAULTS,_is_bounded_feedback_cycle(),_is_feedback_link(),_get_feedback_default(),_build_feedback_fns()(multi-hop DFS chain builder with simple_eval),_resolve_input_value(),_collect_extra_names(), bootstrap + injection logic,validate_inputs()bounded-cycle bypasscomfy_execution/graph.pyfeedback_linksdict onTopologicalSort,_is_feedback_output(), strong-link skip + cache-link for feedback edgescomfy/k_diffusion/sampling.py_init_dynamic_options,_refresh_dynamic_params,_apply_dynamic_s_noise), per-loop s_noise/eta/r re-read in 23 stochastic sampler functionscomfy_extras/nodes_custom_sampler.pyBOUNDED_FEEDBACK = {"step_index"},io.Int.Output(display_name="step_index"), per-step callback wrapper updating guider.cfg and sampler.extra_optionscomfy/samplers.py_dynamic_sampler_optionsinjection inKSAMPLER.sample()Verification
Tested end-to-end with a 40-step sampler using expression
3 + a * 2wherea = step/40:pow()functions all verifiedDesign Principles
BOUNDED_FEEDBACK = {"step_index"}on any node enables feedbackBOUNDED_FEEDBACKdeclaration