Merge 6fc009d69e into 7f287b705e

comfy/samplers.py

@@ -996,6 +996,12 @@ class KSAMPLER(Sampler):
         if callback is not None:
             k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
 
+        # Expose mutable extra_options so sampler functions can re-read
+        # updated values at each step (e.g. s_noise varied by feedback).
+        # Only inject when the sampler has per-step feedback param functions,
+        # otherwise _dynamic_sampler_options would leak to the model call.
+        if hasattr(self, '_feedback_param_fns') and self._feedback_param_fns:
+            extra_args["_dynamic_sampler_options"] = self.extra_options
         samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
         samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
         return samples

comfy_execution/graph.py

@@ -111,6 +111,32 @@ class TopologicalSort:
         self.blocking = {} # Which nodes are blocked by this node
         self.externalBlocks = 0
         self.unblockedEvent = asyncio.Event()
+        # Tracks bounded-feedback edges that were intentionally excluded from
+        # strong (blocking) links.  Maps to_node_id -> list of (from_node_id,
+        # from_socket) so the execution layer can inject initial values for the
+        # iteration output that closes the cycle.
+        self.feedback_links = {}
+
+    def _is_feedback_output(self, from_node_id, from_socket):
+        """Return True when *from_socket* of *from_node_id* is a declared
+        bounded-iteration output (``BOUNDED_FEEDBACK``)."""
+        try:
+            class_type = self.dynprompt.get_node(from_node_id)["class_type"]
+            class_def = nodes.NODE_CLASS_MAPPINGS.get(class_type)
+        except (NodeNotFoundError, KeyError):
+            return False
+        if class_def is None:
+            return False
+        bounded = getattr(class_def, 'BOUNDED_FEEDBACK', None)
+        if not bounded:
+            return False
+        # Map socket index to name via RETURN_NAMES, falling back to the raw index.
+        return_names = getattr(class_def, 'RETURN_NAMES', None)
+        idx = int(from_socket)
+        if return_names is not None and 0 <= idx < len(return_names):
+            return return_names[idx] in bounded
+        # If the socket is already a string (uncommon), check directly.
+        return str(from_socket) in bounded
 
     def get_input_info(self, unique_id, input_name):
         class_type = self.dynprompt.get_node(unique_id)["class_type"]
@@ -163,6 +189,24 @@ class TopologicalSort:
                         links.append((from_node_id, from_socket, unique_id))
 
         for link in links:
+            from_node_id, from_socket, to_node_id = link
+            if self._is_feedback_output(from_node_id, from_socket):
+                # This edge carries an iteration variable (e.g. step_index)
+                # back upstream to close a bounded feedback cycle.  Don't
+                # create a strong (blocking) link — that would deadlock the
+                # topological dissolve.  Instead record it so the execution
+                # layer can seed the iteration output with an initial value.
+                if to_node_id not in self.feedback_links:
+                    self.feedback_links[to_node_id] = []
+                self.feedback_links[to_node_id].append((from_node_id, from_socket))
+                # Still ensure the source node is in the graph.
+                self.add_node(from_node_id)
+                # Create a cache link so the downstream node can read the
+                # placeholder value injected into the output cache by the
+                # execution bootstrap (only available on ExecutionList).
+                if hasattr(self, 'cache_link'):
+                    self.cache_link(from_node_id, to_node_id)
+                continue
             self.add_strong_link(*link)
 
     def add_external_block(self, node_id):

comfy_extras/nodes_custom_sampler.py

@@ -1011,6 +1011,10 @@ class RandomNoise(io.ComfyNode):
 
 
 class SamplerCustomAdvanced(io.ComfyNode):
+    # Declare which outputs are bounded iteration variables that may feed back
+    # through the graph to control upstream parameters (e.g. step_index -> cfg).
+    BOUNDED_FEEDBACK = {"step_index"}
+
     @classmethod
     def define_schema(cls):
         return io.Schema(
@@ -1026,6 +1030,7 @@ class SamplerCustomAdvanced(io.ComfyNode):
             outputs=[
                 io.Latent.Output(display_name="output"),
                 io.Latent.Output(display_name="denoised_output"),
+                io.Int.Output(display_name="step_index"),
             ]
         )
 
@@ -1041,8 +1046,30 @@ class SamplerCustomAdvanced(io.ComfyNode):
         if "noise_mask" in latent:
             noise_mask = latent["noise_mask"]
 
+        total_steps = sigmas.shape[-1] - 1
         x0_output = {}
-        callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output)
+        callback = latent_preview.prepare_callback(guider.model_patcher, total_steps, x0_output)
+
+        # ---- bounded-feedback per-step updates ----
+        # The execution engine may have injected per-step update functions
+        # onto the guider and/or sampler objects.  Wrap the callback to
+        # apply them before the *next* sampling step.  The k-diffusion
+        # callback fires *after* the model call for step i, so we pass
+        # i+1 so that step N uses parameters computed with a=N.
+        cfg_fn = getattr(guider, '_feedback_cfg_fn', None)
+        param_fns = getattr(sampler, '_feedback_param_fns', None)
+        _has_feedback = cfg_fn is not None or param_fns
+        if _has_feedback:
+            _orig_callback = callback
+            def _feedback_callback(step, x0, x, total_steps):
+                if cfg_fn is not None:
+                    guider.cfg = cfg_fn(step + 1, total_steps)
+                if param_fns is not None:
+                    for key, fn in param_fns.items():
+                        sampler.extra_options[key] = fn(step + 1, total_steps)
+                _orig_callback(step, x0, x, total_steps)
+            callback = _feedback_callback
+        # ----------------------------------------------------
 
         disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
         samples = guider.sample(noise.generate_noise(latent), latent_image, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed)
@@ -1061,7 +1088,7 @@ class SamplerCustomAdvanced(io.ComfyNode):
             out_denoised["samples"] = x0_out
         else:
             out_denoised = out
-        return io.NodeOutput(out, out_denoised)
+        return io.NodeOutput(out, out_denoised, total_steps)
 
     sample = execute