Make the step index of sampler function accessible in model_options/transformer_options
- 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
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