wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-07 15:10:50 +08:00
Code Issues Actions 17 Packages Projects Releases Wiki Activity

Merge e435cc668f into b08debceca

Browse Source
This commit is contained in:
masahiroteraoka 2026-07-06 17:34:00 +08:00 committed by GitHub
commit 34fab2c709
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 61 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
comfy/patcher_extension.py
View File

@ -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, ...]} }

63
comfy/samplers.py
View File

@ -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={}):