Ensure sampler end callbacks run on errors

This commit is contained in:
Masahiro 2026-05-27 12:39:10 +12:00
parent 5d614c84cb
commit e435cc668f

View File

@ -1038,18 +1038,23 @@ class KSAMPLER(Sampler):
for sampler_callback in sampler_step_callbacks: for sampler_callback in sampler_step_callbacks:
sampler_callback(sampler_info) sampler_callback(sampler_info)
samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options) samples = None
samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples) sampling_succeeded = False
if len(sampler_end_callbacks) > 0: try:
sampler_info = { samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
"total_steps": total_steps, samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
"sample_sigmas": sigmas, sampling_succeeded = True
"samples_shape": tuple(samples.shape), return samples
"sampler_function": sampler_function_name, finally:
} if len(sampler_end_callbacks) > 0:
for sampler_callback in sampler_end_callbacks: sampler_info = {
sampler_callback(sampler_info) "total_steps": total_steps,
return samples "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={}): def ksampler(sampler_name, extra_options={}, inpaint_options={}):