Comfy-Org / Comfy-Org/ComfyUI

Make the step index of sampler function accessible in model_options/transformer_options

Open
#11,075 0 comments 0 reactions 0 assignees View on GitHub
Feature
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 6h
Merged PRs (30d)
155

Description

### Feature Idea

Although input sigma and `sample_sigmas` can indicate the step in some cases, high-order single-step samplers can pass a sigma that does not appear in `sample_sigmas`. For instance, the current Context Windows cannot handle mid-stage sigma (https://github.com/comfyanonymous/ComfyUI/issues/9659).

Another example is the Heun-style samplers, which can feed the next sigma (`sigmas[i + 1]`) to the model within the current timestep (`sigmas[i]`), but relying on input sigma and `sample_sigmas` makes it difficult to determine whether it comes from a mid-stage or the next timestep.

Therefore, explicitly passing the current step index `i` can help in those situations, and it also makes getting the current step simpler.

### Existing Solutions

A possible solution might be:
```
def add_sampler_step_to_extra_args(extra_args: dict, step: int = -1) -> dict:
if step < 0:
return extra_args
model_options = extra_args.get("model_options", {})
step_model_options = model_options | {"sampler_function_step": step}
return extra_args | {"model_options": step_model_options}

for i in trange(len(sigmas) - 1, disable=disable):
step_extra_args = add_sampler_step_to_extra_args(extra_args, step=i)
denoised = model(x, sigmas[i] * s_in, **step_extra_args)
```

Since this needs to be injected into all sampler functions, I'm not sure it's a good solution.

### Other

_No response_

Contributor guide

Open the contributing guide

Research direction

Start by tracing the sampler functions and their extra_args flow into model_options and transformer_options. Compare how current sigma and sample_sigmas are exposed, then determine where the current step index can be made available across sampler functions; done means high-order and Heun-style steps can distinguish their current timestep from mid-stage or next-timestep calls.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
ai
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.