diff --git a/restart_sampling.py b/restart_sampling.py index 02a9381..5845b04 100644 --- a/restart_sampling.py +++ b/restart_sampling.py @@ -17,7 +17,23 @@ def add_restart_segment(restart_segments, n_restart, k, t_min, t_max): return restart_segments -def prepare_restart_segments(restart_info): +def resolve_t_value(val, ms): + if isinstance(val, (float, int)): + if val >= 0.0: + return val + if val >= -1000: + return ms.sigma(torch.FloatTensor([abs(int(val))], device="cpu")).item() + if isinstance(val, str) and val.endswith("%"): + try: + val = float(val[:-1]) + if val >= 0 and val <= 100: + return ms.percent_to_sigma(1.0 - val / 100.0) + except ValueError: + pass + raise ValueError("bad t_min or t_max value") + + +def prepare_restart_segments(restart_info, ms): try: restart_arrays = ast.literal_eval(f"[{restart_info}]") except SyntaxError as e: @@ -27,8 +43,10 @@ def prepare_restart_segments(restart_info): for arr in restart_arrays: if len(arr) != 4: raise ValueError("Restart segment must have 4 values") - n_restart, k, t_min, t_max = arr + n_restart, k, val_min, val_max = arr n_restart, k = int(n_restart), int(k) + t_min = resolve_t_value(val_min, ms) + t_max = resolve_t_value(val_max, ms) restart_segments = add_restart_segment(restart_segments, n_restart, k, t_min, t_max) return restart_segments @@ -73,13 +91,14 @@ def restart_sampling(model, seed, steps, cfg, sampler, scheduler, positive, nega if isinstance(sampler, str): sampler = sampler_object(sampler) - restart_segments = prepare_restart_segments(restart_info) comfy.model_management.load_models_gpu([model]) real_model = model while hasattr(real_model, "model"): real_model = real_model.model + restart_segments = prepare_restart_segments(restart_info, real_model.model_sampling) + effective_steps = steps if step_range is not None or denoise > 0.9999 else int(steps / denoise) sigmas = calc_sigmas(scheduler, effective_steps, float(real_model.model_sampling.sigma_min), float(real_model.model_sampling.sigma_max), @@ -185,17 +204,20 @@ class KSamplerRestartWrapper: **kwargs) pbar.update(1) step += 1 - if sigmas[i + 1].item() in segments: - seg = segments[sigmas[i + 1].item()] - s_min, s_max, k, n_restart = sigmas[i + 1], seg['t_max'], seg['k'], seg['n'] - seg_sigmas = calc_sigmas(self.restart_scheduler, n_restart, s_min, - s_max, self.real_model, device=x.device) - for _ in range(k): - x += torch.randn_like(x) * (s_max ** 2 - s_min ** 2) ** 0.5 - for j in range(n_restart - 1): - x = ksampler.sampler_function(model, x, torch.tensor( - [seg_sigmas[j], seg_sigmas[j + 1]], device=x.device), *args, extra_args=extra_args, - callback=callback_wrapper, disable=True, **kwargs) - pbar.update(1) - step += 1 + s_min = sigmas[i + 1].item() + seg = segments.get(s_min) + if seg is None: + continue + seg = segments[s_min] + s_max, k, n_restart = seg['t_max'], seg['k'], seg['n'] + seg_sigmas = calc_sigmas(self.restart_scheduler, n_restart, s_min, + s_max, self.real_model, device=x.device) + for _ in range(k): + x += torch.randn_like(x) * (s_max ** 2 - s_min ** 2) ** 0.5 + for j in range(n_restart - 1): + x = ksampler.sampler_function(model, x, torch.tensor( + [seg_sigmas[j], seg_sigmas[j + 1]], device=x.device), *args, extra_args=extra_args, + callback=callback_wrapper, disable=True, **kwargs) + pbar.update(1) + step += 1 return x