Add sampler lifecycle callbacks

[masahiroteraoka]

Summary

Adds three sampler lifecycle callback identifiers to CallbacksMP:

• ON_SAMPLER_START
• ON_SAMPLER_STEP
• ON_SAMPLER_END

KSAMPLER.sample() now emits these callbacks around the existing k-diffusion callback path. The callbacks receive sampling metadata such as step index, sigmas, total step count, sampler function name, and tensor shapes.

Motivation

Some custom nodes need to observe sampler lifecycle information for debugging, profiling, progress tracking, or step-aware extension behavior. Today those extensions often need to wrap or replace sampler internals even when they only need metadata.

Comfy already exposes wrappers and callbacks through patcher_extension; this change extends that existing mechanism to sampler lifecycle events without adding node UI or changing default sampling behavior.

Scope

This intentionally does not expose mutable latent tensors to callbacks. The step callback includes x_shape and denoised_shape, not the tensors themselves.

If there is interest in a latent observer or transform callback later, that should be handled separately because it needs a more explicit contract for ordering, mutation, dtype/device/shape validation, batching, masks, determinism, and multiple callbacks.

Compatibility

• No callbacks registered: no intended behavior change.
• Existing preview/progress callback path is preserved.
• No new node inputs or UI changes.

Testing

• python3 -m py_compile comfy/patcher_extension.py comfy/samplers.py
• git diff --check
• Local smoke test with a stub KSAMPLER sampler function verified that ON_SAMPLER_START, ON_SAMPLER_STEP, and ON_SAMPLER_END callbacks are emitted in order while preserving the existing callback path.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: bb7d3282-c242-40f4-8bf0-7bd1323edbd0

📥 Commits

Reviewing files that changed from the base of the PR and between 5d614c8 and e435cc6.

📒 Files selected for processing (1)

• comfy/samplers.py

🚧 Files skipped from review as they are similar to previous changes (1)

• comfy/samplers.py

---

📝 Walkthrough

Walkthrough

This PR adds sampler lifecycle callback support to ComfyUI's sampling pipeline. It introduces three new callback event identifiers (ON_SAMPLER_START, ON_SAMPLER_STEP, ON_SAMPLER_END) in CallbacksMP and integrates callback dispatch into KSAMPLER.sample. KSAMPLER.sample now reads callbacks from extra_args["model_options"], invokes start callbacks with sampler metadata, calls step callbacks (and the original callback) each step with sigma and tensor-shape info, and invokes end callbacks with the final samples shape.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Add sampler lifecycle callbacks' directly and clearly summarizes the main change: introducing three new callback identifiers for sampler lifecycle events.
Description check	✅ Passed	The description is comprehensive and directly related to the changeset, explaining the motivation, scope, compatibility, and testing for the sampler lifecycle callbacks feature.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ 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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 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.

Inline comments:
In `@comfy/samplers.py`:
- Around line 1041-1051: The sampling call using self.sampler_function(...)
should be wrapped in a try/finally so sampler_end_callbacks are always invoked
even on exceptions; move the inverse_noise_scaling and the construction of
sampler_info into the try block for the success path, and in the finally iterate
sampler_end_callbacks to emit the end event with sampler_info where
"samples_shape": tuple(samples.shape) on success but "samples_shape": None on
failure, referencing self.sampler_function,
model_wrap.inner_model.model_sampling.inverse_noise_scaling,
sampler_end_callbacks, sampler_info and k_callback so listeners aren't left
stuck when sampling throws.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: bd84356e-7c9f-4667-b968-688d78590967

📥 Commits

Reviewing files that changed from the base of the PR and between f9f54ca and 5d614c8.

📒 Files selected for processing (2)

• comfy/patcher_extension.py
• comfy/samplers.py