diff --git a/comfy/patcher_extension.py b/comfy/patcher_extension.py index 189ee84ca..b869cdef9 100644 --- a/comfy/patcher_extension.py +++ b/comfy/patcher_extension.py @@ -13,6 +13,9 @@ class CallbacksMP: ON_REGISTER_ALL_HOOK_PATCHES = "on_register_all_hook_patches" ON_INJECT_MODEL = "on_inject_model" ON_EJECT_MODEL = "on_eject_model" + ON_SAMPLER_START = "on_sampler_start" + ON_SAMPLER_STEP = "on_sampler_step" + ON_SAMPLER_END = "on_sampler_end" # callbacks dict is in the format: # {"call_type": {"key": [Callable1, Callable2, ...]} } diff --git a/comfy/samplers.py b/comfy/samplers.py index 25c5a855f..206507aeb 100755 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -993,12 +993,65 @@ class KSAMPLER(Sampler): k_callback = None total_steps = len(sigmas) - 1 - if callback is not None: - k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps) + model_options = extra_args.get("model_options", {}) + callback_types = comfy.patcher_extension.CallbacksMP + get_callbacks = comfy.patcher_extension.get_all_callbacks + sampler_function_name = getattr(self.sampler_function, "__name__", self.sampler_function.__class__.__name__) + sampler_start_callbacks = get_callbacks(callback_types.ON_SAMPLER_START, model_options, is_model_options=True) + sampler_step_callbacks = get_callbacks(callback_types.ON_SAMPLER_STEP, model_options, is_model_options=True) + sampler_end_callbacks = get_callbacks(callback_types.ON_SAMPLER_END, model_options, is_model_options=True) - 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 + if len(sampler_start_callbacks) > 0: + sampler_info = { + "total_steps": total_steps, + "sample_sigmas": sigmas, + "noise_shape": tuple(noise.shape), + "latent_shape": tuple(latent_image.shape) if latent_image is not None else None, + "sampler_function": sampler_function_name, + } + for sampler_callback in sampler_start_callbacks: + sampler_callback(sampler_info) + + if callback is not None or len(sampler_step_callbacks) > 0: + def k_callback(x): + if callback is not None: + callback(x["i"], x["denoised"], x["x"], total_steps) + if len(sampler_step_callbacks) == 0: + return + + step = x["i"] + sigma_next = sigmas[step + 1] if step + 1 < len(sigmas) else None + sampler_info = { + "step": step, + "total_steps": total_steps, + "sigma": x.get("sigma", sigmas[step] if step < len(sigmas) else None), + "sigma_next": sigma_next, + "sigma_hat": x.get("sigma_hat", None), + "sample_sigmas": sigmas, + "x_shape": tuple(x["x"].shape) if "x" in x else None, + "denoised_shape": tuple(x["denoised"].shape) if "denoised" in x else None, + "sampler_function": sampler_function_name, + } + for sampler_callback in sampler_step_callbacks: + sampler_callback(sampler_info) + + samples = None + sampling_succeeded = False + try: + 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) + sampling_succeeded = True + return samples + finally: + if len(sampler_end_callbacks) > 0: + sampler_info = { + "total_steps": total_steps, + "sample_sigmas": sigmas, + "samples_shape": tuple(samples.shape) if sampling_succeeded else None, + "sampler_function": sampler_function_name, + } + for sampler_callback in sampler_end_callbacks: + sampler_callback(sampler_info) def ksampler(sampler_name, extra_options={}, inpaint_options={}):