diff --git a/CHANGELOG.md b/CHANGELOG.md index 13401a956..8399227b5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,7 @@ - further work on type consistency and type checking, thanks @awsr - log captured exceptions - improve temp folder handling and cleanup + - remove torch errors/warings on fast server shutdown - add ui placeholders for future agent-scheduler work, thanks @ryanmeador - implement abort system on repeated errors, thanks @awsr currently used by lora and textual-inversion loaders diff --git a/modules/loader.py b/modules/loader.py index c6e25a1d3..951728682 100644 --- a/modules/loader.py +++ b/modules/loader.py @@ -46,6 +46,13 @@ except Exception as e: sys.exit(1) timer.startup.record("scipy") +try: + import atexit + import torch._inductor.async_compile as ac + atexit.unregister(ac.shutdown_compile_workers) +except Exception: + pass + import torch # pylint: disable=C0411 if torch.__version__.startswith('2.5.0'): errors.log.warning(f'Disabling cuDNN for SDP on torch={torch.__version__}') diff --git a/modules/res4lyf/TASK.md b/modules/res4lyf/TASK.md index e3c376d7d..4fa0d1e82 100644 --- a/modules/res4lyf/TASK.md +++ b/modules/res4lyf/TASK.md @@ -1,20 +1,23 @@ -# RES4LYF DIFFUSION SCHEDULERS +# TASK: Schedulers -- Schedulers codebase is in `modules/res4lyf` - do not modify any other files -- Testing notes: - - using `epsilon` prediction type - - using `StableDiffusionXLPipeline` pipeline for text2image - - using `StableDiffusionXLInpaintPipeline` for inpainting and image2image -- *ETDRKScheduler, LawsonScheduler, ABNorsettScheduler, RESSinglestepScheduler, RESSinglestepSDEScheduler, PECScheduler*: +## Notes + +This is a codebase for diffusion schedulers implemented for `diffusers` library and ported from `res4lyf` repository at + +Ported schedulers codebase is in `modules/res4lyf`, do not modify any other files + +## Testing + +Current focus is on following code-paths: +- using `epsilon` prediction type +- using `StableDiffusionXLPipeline` pipeline for *text2image* + +## Results + +- *ETDRKScheduler, LawsonScheduler, ABNorsettScheduler, RESSinglestepScheduler, RESSinglestepSDEScheduler, PECScheduler, etc.*: do NOT modify behavior and codebase for these schedulers as they produce good outputs under all circumstances if needed, you can use them as gold-standard references to compare other schedulers against -- *LinearRKScheduler, LobattoScheduler, RadauIIAScheduler, GaussLegendreScheduler, SpecializedRKScheduler, RungeKuttaScheduler*: - work well for text2image, but then in image2image it produces pure black image -- *RESUnifiedScheduler*: - works fine with `rk_type=res_2s` and similar single-step params, - but produces too much noise with `rk_type=res_2m` and similar multi-step params -- *DEISMultistepScheduler, RESMultistepScheduler* have the same problem - while *RESSinglestepScheduler* works fine -- *CommonSigmaScheduler, LangevinDynamicsScheduler*: - do not work with `epsilon` prediction type, results in pure noise +- *RESUnifiedScheduler*, *DEISMultistepScheduler, RESMultistepScheduler* + work fine with `rk_type=res_2s`, `rk_type=deis_1s` and similar single-step params, + but with `rk_type=res_2m`, `rk_type=deis_2m` and similar multi-step params + image looks fine in early steps, but then degrages at the final steps with what looks like too much noise diff --git a/modules/res4lyf/common_sigma_scheduler.py b/modules/res4lyf/common_sigma_scheduler.py index 46be9207d..61f47c015 100644 --- a/modules/res4lyf/common_sigma_scheduler.py +++ b/modules/res4lyf/common_sigma_scheduler.py @@ -125,8 +125,8 @@ class CommonSigmaScheduler(SchedulerMixin, ConfigMixin): # Derived sigma range from alphas_cumprod base_sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5) - sigma_max = base_sigmas[0] - sigma_min = base_sigmas[-1] + sigma_max = base_sigmas[-1] + sigma_min = base_sigmas[0] t = torch.linspace(0, 1, num_inference_steps) profile = self.config.profile diff --git a/modules/res4lyf/deis_scheduler_alt.py b/modules/res4lyf/deis_scheduler_alt.py index d28b27fe0..689c710da 100644 --- a/modules/res4lyf/deis_scheduler_alt.py +++ b/modules/res4lyf/deis_scheduler_alt.py @@ -212,19 +212,10 @@ class DEISMultistepScheduler(SchedulerMixin, ConfigMixin): return self._step_index def index_for_timestep(self, timestep, schedule_timesteps=None): - if self._step_index is not None: - return self._step_index - + from .scheduler_utils import index_for_timestep if schedule_timesteps is None: schedule_timesteps = self.timesteps - - if isinstance(schedule_timesteps, torch.Tensor): - schedule_timesteps = schedule_timesteps.detach().cpu().numpy() - - if isinstance(timestep, torch.Tensor): - timestep = timestep.detach().cpu().numpy() - - return np.abs(schedule_timesteps - timestep).argmin().item() + return index_for_timestep(timestep, schedule_timesteps) def _init_step_index(self, timestep): if self._step_index is None: @@ -286,30 +277,51 @@ class DEISMultistepScheduler(SchedulerMixin, ConfigMixin): x0s = [denoised] + self.model_outputs[::-1] orders = min(len(x0s), self.config.solver_order) - if orders == 1: + # Force Order 1 at the end of schedule + if self.num_inference_steps is not None and step_index >= self.num_inference_steps - 3: + res = phi_1 * denoised + elif orders == 1: res = phi_1 * denoised elif orders == 2: # Use phi(2) for 2nd order interpolation h_prev = -np.log(self._sigmas_cpu[step_index] / (self._sigmas_cpu[step_index - 1] + 1e-9)) h_prev_t = torch.tensor(h_prev, device=sample.device, dtype=sample.dtype) r = h_prev_t / (h + 1e-9) - phi_2 = phi(2) - # Correct Adams-Bashforth-like coefficients: b2 = -phi_2 / r - b2 = -phi_2 / (r + 1e-9) - b1 = phi_1 - b2 - res = b1 * x0s[0] + b2 * x0s[1] + h_prev = -np.log(self._sigmas_cpu[step_index] / (self._sigmas_cpu[step_index - 1] + 1e-9)) + h_prev_t = torch.tensor(h_prev, device=sample.device, dtype=sample.dtype) + r = h_prev_t / (h + 1e-9) + + # Hard Restart + if r < 0.5 or r > 2.0: + res = phi_1 * denoised + else: + phi_2 = phi(2) + # Correct Adams-Bashforth-like coefficients: b2 = -phi_2 / r + b2 = -phi_2 / (r + 1e-9) + b1 = phi_1 - b2 + res = b1 * x0s[0] + b2 * x0s[1] elif orders == 3: + # 3rd order with varying step sizes # 3rd order with varying step sizes h_p1 = -np.log(self._sigmas_cpu[step_index] / (self._sigmas_cpu[step_index - 1] + 1e-9)) h_p2 = -np.log(self._sigmas_cpu[step_index] / (self._sigmas_cpu[step_index - 2] + 1e-9)) r1 = torch.tensor(h_p1, device=sample.device, dtype=sample.dtype) / (h + 1e-9) r2 = torch.tensor(h_p2, device=sample.device, dtype=sample.dtype) / (h + 1e-9) - phi_2, phi_3 = phi(2), phi(3) - denom = r2 - r1 + 1e-9 - b3 = (phi_3 + r1 * phi_2) / (r2 * denom) - b2 = -(phi_3 + r2 * phi_2) / (r1 * denom) - b1 = phi_1 - b2 - b3 - res = b1 * x0s[0] + b2 * x0s[1] + b3 * x0s[2] + h_p1 = -np.log(self._sigmas_cpu[step_index] / (self._sigmas_cpu[step_index - 1] + 1e-9)) + h_p2 = -np.log(self._sigmas_cpu[step_index] / (self._sigmas_cpu[step_index - 2] + 1e-9)) + r1 = torch.tensor(h_p1, device=sample.device, dtype=sample.dtype) / (h + 1e-9) + r2 = torch.tensor(h_p2, device=sample.device, dtype=sample.dtype) / (h + 1e-9) + + # Hard Restart + if r1 < 0.5 or r1 > 2.0 or r2 < 0.5 or r2 > 2.0: + res = phi_1 * denoised + else: + phi_2, phi_3 = phi(2), phi(3) + denom = r2 - r1 + 1e-9 + b3 = (phi_3 + r1 * phi_2) / (r2 * denom) + b2 = -(phi_3 + r2 * phi_2) / (r1 * denom) + b1 = phi_1 - b2 - b3 + res = b1 * x0s[0] + b2 * x0s[1] + b3 * x0s[2] else: # Fallback to Euler or lower order res = phi_1 * denoised diff --git a/modules/res4lyf/gauss_legendre_scheduler.py b/modules/res4lyf/gauss_legendre_scheduler.py index 0a7c87dc1..4740b7565 100644 --- a/modules/res4lyf/gauss_legendre_scheduler.py +++ b/modules/res4lyf/gauss_legendre_scheduler.py @@ -227,7 +227,7 @@ class GaussLegendreScheduler(SchedulerMixin, ConfigMixin): self.model_outputs = [] self.sample_at_start_of_step = None - self._step_index = 0 + self._step_index = None @property def step_index(self): @@ -322,6 +322,16 @@ class GaussLegendreScheduler(SchedulerMixin, ConfigMixin): derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample) if self.sample_at_start_of_step is None: + if stage_index > 0: + # Mid-step fallback for Img2Img/Inpainting + sigma_next_t = self.sigmas[self._step_index + 1] + dt = sigma_next_t - sigma_t + prev_sample = sample + dt * derivative + self._step_index += 1 + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) + self.sample_at_start_of_step = sample self.model_outputs = [derivative] * stage_index diff --git a/modules/res4lyf/langevin_dynamics_scheduler.py b/modules/res4lyf/langevin_dynamics_scheduler.py index 6aeea1a21..fc2a83846 100644 --- a/modules/res4lyf/langevin_dynamics_scheduler.py +++ b/modules/res4lyf/langevin_dynamics_scheduler.py @@ -140,6 +140,8 @@ class LangevinDynamicsScheduler(SchedulerMixin, ConfigMixin): trajectory.append(x.item()) sigmas = np.array(trajectory) + # Force monotonicity to prevent negative h in step() + sigmas = np.sort(sigmas)[::-1] sigmas[-1] = end_sigma if self.config.use_karras_sigmas: diff --git a/modules/res4lyf/linear_rk_scheduler.py b/modules/res4lyf/linear_rk_scheduler.py index 6048f5749..8eb26c446 100644 --- a/modules/res4lyf/linear_rk_scheduler.py +++ b/modules/res4lyf/linear_rk_scheduler.py @@ -67,7 +67,7 @@ class LinearRKScheduler(SchedulerMixin, ConfigMixin): # Internal state self.model_outputs = [] self.sample_at_start_of_step = None - self._step_index = 0 + self._step_index = None def _get_tableau(self): v = str(self.config.variant).lower().strip() @@ -183,7 +183,7 @@ class LinearRKScheduler(SchedulerMixin, ConfigMixin): self.model_outputs = [] self.sample_at_start_of_step = None - self._step_index = 0 + self._step_index = None @property def step_index(self): @@ -261,6 +261,16 @@ class LinearRKScheduler(SchedulerMixin, ConfigMixin): derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample) if self.sample_at_start_of_step is None: + if stage_index > 0: + # Mid-step fallback for Img2Img/Inpainting + sigma_next_t = self.sigmas[self._step_index + 1] + dt = sigma_next_t - sigma_t + prev_sample = sample + dt * derivative + self._step_index += 1 + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) + self.sample_at_start_of_step = sample self.model_outputs = [derivative] * stage_index diff --git a/modules/res4lyf/lobatto_scheduler.py b/modules/res4lyf/lobatto_scheduler.py index 94282109e..bba39e07a 100644 --- a/modules/res4lyf/lobatto_scheduler.py +++ b/modules/res4lyf/lobatto_scheduler.py @@ -68,7 +68,7 @@ class LobattoScheduler(SchedulerMixin, ConfigMixin): # Internal state self.model_outputs = [] self.sample_at_start_of_step = None - self._step_index = 0 + self._step_index = None def _get_tableau(self): v = self.config.variant @@ -183,7 +183,7 @@ class LobattoScheduler(SchedulerMixin, ConfigMixin): self.model_outputs = [] self.sample_at_start_of_step = None - self._step_index = 0 + self._step_index = None @property def step_index(self): @@ -261,6 +261,16 @@ class LobattoScheduler(SchedulerMixin, ConfigMixin): derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample) if self.sample_at_start_of_step is None: + if stage_index > 0: + # Mid-step fallback for Img2Img/Inpainting + sigma_next_t = self.sigmas[self._step_index + 1] + dt = sigma_next_t - sigma_t + prev_sample = sample + dt * derivative + self._step_index += 1 + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) + self.sample_at_start_of_step = sample self.model_outputs = [derivative] * stage_index diff --git a/modules/res4lyf/radau_iia_scheduler.py b/modules/res4lyf/radau_iia_scheduler.py index 741646305..f88474fad 100644 --- a/modules/res4lyf/radau_iia_scheduler.py +++ b/modules/res4lyf/radau_iia_scheduler.py @@ -68,7 +68,7 @@ class RadauIIAScheduler(SchedulerMixin, ConfigMixin): # Internal state self.model_outputs = [] self.sample_at_start_of_step = None - self._step_index = 0 + self._step_index = None def _get_tableau(self): v = self.config.variant @@ -217,7 +217,7 @@ class RadauIIAScheduler(SchedulerMixin, ConfigMixin): self.model_outputs = [] self.sample_at_start_of_step = None - self._step_index = 0 + self._step_index = None @property def step_index(self): @@ -308,6 +308,16 @@ class RadauIIAScheduler(SchedulerMixin, ConfigMixin): derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample) if self.sample_at_start_of_step is None: + if stage_index > 0: + # Mid-step fallback for Img2Img/Inpainting + sigma_next_t = self.sigmas[self._step_index + 1] + dt = sigma_next_t - sigma_t + prev_sample = sample + dt * derivative + self._step_index += 1 + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) + self.sample_at_start_of_step = sample self.model_outputs = [derivative] * stage_index diff --git a/modules/res4lyf/res_multistep_scheduler.py b/modules/res4lyf/res_multistep_scheduler.py index e4f3d5ab4..bd85876cf 100644 --- a/modules/res4lyf/res_multistep_scheduler.py +++ b/modules/res4lyf/res_multistep_scheduler.py @@ -242,6 +242,11 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin): if variant.startswith("res"): # REiS Multistep logic c2, c3 = 0.5, 1.0 + + # Force Order 1 at the end of schedule + if self.num_inference_steps is not None and self._step_index >= self.num_inference_steps - 3: + curr_order = 1 + if curr_order == 2: h_prev = -torch.log(self.prev_sigmas[-1] / self.prev_sigmas[-2]) c2 = (-h_prev / h).item() if h > 0 else 0.5 @@ -260,21 +265,43 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin): res = phi_1 * x0 elif curr_order == 2: # b2 = -phi_2 / r - b2 = -phi(2) / ((-h_prev / h) + 1e-9) - b1 = phi_1 - b2 - res = b1 * self.x0_outputs[-1] + b2 * self.x0_outputs[-2] + # b2 = -phi_2 / r = -phi(2) / (h_prev/h) + # Here we use: b2 = phi(2) / ((-h_prev / h) + 1e-9) + # Since (-h_prev/h) is negative (-r), this gives correct negative sign for b2. + + # Stability check + r_check = h_prev / (h + 1e-9) # This is effectively -r if using h_prev definition above? + # Wait, h_prev above is -log(). Positive. + # h is positive. + # So h_prev/h is positive. defined as r in other files. + # But here code uses -h_prev / h in denominator. + + # Stability check + r_check = h_prev / (h + 1e-9) + + # Hard Restart + if r_check < 0.5 or r_check > 2.0: + res = phi_1 * x0 + else: + b2 = phi(2) / ((-h_prev / h) + 1e-9) + b1 = phi_1 - b2 + res = b1 * self.x0_outputs[-1] + b2 * self.x0_outputs[-2] elif curr_order == 3: # Generalized AB3 for Exponential Integrators h_p1 = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-2] + 1e-9)) h_p2 = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-3] + 1e-9)) r1 = h_p1 / (h + 1e-9) r2 = h_p2 / (h + 1e-9) - phi_2, phi_3 = phi(2), phi(3) - denom = r2 - r1 + 1e-9 - b3 = (phi_3 + r1 * phi_2) / (r2 * denom) - b2 = -(phi_3 + r2 * phi_2) / (r1 * denom) - b1 = phi_1 - b2 - b3 - res = b1 * self.x0_outputs[-1] + b2 * self.x0_outputs[-2] + b3 * self.x0_outputs[-3] + + if r1 < 0.5 or r1 > 2.0 or r2 < 0.5 or r2 > 2.0: + res = phi_1 * x0 + else: + phi_2, phi_3 = phi(2), phi(3) + denom = r2 - r1 + 1e-9 + b3 = (phi_3 + r1 * phi_2) / (r2 * denom) + b2 = -(phi_3 + r2 * phi_2) / (r1 * denom) + b1 = phi_1 - b2 - b3 + res = b1 * self.x0_outputs[-1] + b2 * self.x0_outputs[-2] + b3 * self.x0_outputs[-3] else: res = phi_1 * x0 @@ -341,8 +368,13 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin): elif order == 2: h_prev = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-2] + 1e-9)) r = h_prev / (h + 1e-9) - phi_2 = phi(2) + # Correct Adams-Bashforth-like coefficients for Exponential Integrators + + # Hard Restart for stability + if r < 0.5 or r > 2.0: + return [[phi_1]] + b2 = -phi_2 / (r + 1e-9) b1 = phi_1 - b2 return [[b1, b2]] @@ -352,6 +384,9 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin): r1 = h_prev1 / (h + 1e-9) r2 = h_prev2 / (h + 1e-9) + if r1 < 0.5 or r1 > 2.0 or r2 < 0.5 or r2 > 2.0: + return [[phi_1]] + phi_2 = phi(2) phi_3 = phi(3) diff --git a/modules/res4lyf/res_unified_scheduler.py b/modules/res4lyf/res_unified_scheduler.py index f1d3fdb02..3a48744bb 100644 --- a/modules/res4lyf/res_unified_scheduler.py +++ b/modules/res4lyf/res_unified_scheduler.py @@ -177,10 +177,22 @@ class RESUnifiedScheduler(SchedulerMixin, ConfigMixin): # phi_2 = phi(2) # Moved inside conditional blocks as needed history_len = len(self.x0_outputs) + + # Stability: Force Order 1 for final few steps to prevent degradation at low noise levels + if self.num_inference_steps is not None and self._step_index >= self.num_inference_steps - 3: + return [phi_1], h if self.config.rk_type in ["res_2m", "deis_2m"] and history_len >= 2: h_prev = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-2] + 1e-9)) r = h_prev / (h + 1e-9) + + h_prev = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-2] + 1e-9)) + r = h_prev / (h + 1e-9) + + # Hard Restart: if step sizes vary too wildly, fallback to order 1 + if r < 0.5 or r > 2.0: + return [phi_1], h + phi_2 = phi(2) # Correct Adams-Bashforth-like coefficients for Exponential Integrators b2 = -phi_2 / (r + 1e-9) @@ -191,6 +203,15 @@ class RESUnifiedScheduler(SchedulerMixin, ConfigMixin): h_prev2 = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-3] + 1e-9)) r1 = h_prev1 / (h + 1e-9) r2 = h_prev2 / (h + 1e-9) + + h_prev1 = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-2] + 1e-9)) + h_prev2 = -torch.log(self.prev_sigmas[-1] / (self.prev_sigmas[-3] + 1e-9)) + r1 = h_prev1 / (h + 1e-9) + r2 = h_prev2 / (h + 1e-9) + + # Hard Restart check + if r1 < 0.5 or r1 > 2.0 or r2 < 0.5 or r2 > 2.0: + return [phi_1], h phi_2 = phi(2) phi_3 = phi(3) diff --git a/modules/res4lyf/rungekutta_44s_scheduler.py b/modules/res4lyf/rungekutta_44s_scheduler.py index 1a9666e76..444a5e539 100644 --- a/modules/res4lyf/rungekutta_44s_scheduler.py +++ b/modules/res4lyf/rungekutta_44s_scheduler.py @@ -195,6 +195,16 @@ class RungeKutta44Scheduler(SchedulerMixin, ConfigMixin): derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample) if self.sample_at_start_of_step is None: + if stage_index > 0: + # Mid-step fallback for Img2Img/Inpainting + sigma_next_t = self._sigmas_cpu[self._step_index + 1] + dt = sigma_next_t - sigma_t + prev_sample = sample + dt * derivative + self._step_index += 1 + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) + self.sample_at_start_of_step = sample self.model_outputs = [derivative] * stage_index diff --git a/modules/res4lyf/rungekutta_57s_scheduler.py b/modules/res4lyf/rungekutta_57s_scheduler.py index 8bf3ede97..b8294a091 100644 --- a/modules/res4lyf/rungekutta_57s_scheduler.py +++ b/modules/res4lyf/rungekutta_57s_scheduler.py @@ -168,19 +168,10 @@ class RungeKutta57Scheduler(SchedulerMixin, ConfigMixin): return self._step_index def index_for_timestep(self, timestep, schedule_timesteps=None): - if self._step_index is not None: - return self._step_index - + from .scheduler_utils import index_for_timestep if schedule_timesteps is None: - schedule_timesteps = self._timesteps_cpu - else: - if isinstance(schedule_timesteps, torch.Tensor): - schedule_timesteps = schedule_timesteps.detach().cpu().numpy() - - if isinstance(timestep, torch.Tensor): - timestep = timestep.detach().cpu().numpy() - - return np.abs(schedule_timesteps - timestep).argmin().item() + schedule_timesteps = self.timesteps + return index_for_timestep(timestep, schedule_timesteps) def _init_step_index(self, timestep): if self._step_index is None: @@ -237,6 +228,16 @@ class RungeKutta57Scheduler(SchedulerMixin, ConfigMixin): derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample) if self.sample_at_start_of_step is None: + if stage_index > 0: + # Mid-step fallback for Img2Img/Inpainting + sigma_next_t = self._sigmas_cpu[self._step_index + 1] + dt = sigma_next_t - sigma_t + prev_sample = sample + dt * derivative + self._step_index += 1 + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) + self.sample_at_start_of_step = sample self.model_outputs = [derivative] * stage_index diff --git a/modules/res4lyf/rungekutta_67s_scheduler.py b/modules/res4lyf/rungekutta_67s_scheduler.py index 5e3b9af86..96a8d8a0d 100644 --- a/modules/res4lyf/rungekutta_67s_scheduler.py +++ b/modules/res4lyf/rungekutta_67s_scheduler.py @@ -167,19 +167,10 @@ class RungeKutta67Scheduler(SchedulerMixin, ConfigMixin): return self._step_index def index_for_timestep(self, timestep, schedule_timesteps=None): - if self._step_index is not None: - return self._step_index - + from .scheduler_utils import index_for_timestep if schedule_timesteps is None: - schedule_timesteps = self._timesteps_cpu - else: - if isinstance(schedule_timesteps, torch.Tensor): - schedule_timesteps = schedule_timesteps.detach().cpu().numpy() - - if isinstance(timestep, torch.Tensor): - timestep = timestep.detach().cpu().numpy() - - return np.abs(schedule_timesteps - timestep).argmin().item() + schedule_timesteps = self.timesteps + return index_for_timestep(timestep, schedule_timesteps) def _init_step_index(self, timestep): if self._step_index is None: @@ -237,6 +228,16 @@ class RungeKutta67Scheduler(SchedulerMixin, ConfigMixin): derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample) if self.sample_at_start_of_step is None: + if stage_index > 0: + # Mid-step fallback for Img2Img/Inpainting + sigma_next_t = self._sigmas_cpu[self._step_index + 1] + dt = sigma_next_t - sigma_t + prev_sample = sample + dt * derivative + self._step_index += 1 + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) + self.sample_at_start_of_step = sample self.model_outputs = [derivative] * stage_index diff --git a/modules/res4lyf/scheduler_utils.py b/modules/res4lyf/scheduler_utils.py index f18f38b64..99dab451a 100644 --- a/modules/res4lyf/scheduler_utils.py +++ b/modules/res4lyf/scheduler_utils.py @@ -90,13 +90,22 @@ def get_dynamic_shift(mu, base_shift, max_shift, base_seq_len, max_seq_len): return m * mu + b def index_for_timestep(timestep, timesteps): + import numpy as _np + + # Normalize inputs to numpy arrays for a robust, device-agnostic argmin if isinstance(timestep, torch.Tensor): - timestep = timestep.to(timesteps.device) - - # Use argmin for robustness against float precision issues - # and to handle timesteps that might be slightly outside the schedule - dists = torch.abs(timesteps - timestep) - return torch.argmin(dists).item() + timestep_np = timestep.detach().cpu().numpy() + else: + timestep_np = _np.array(timestep) + + if isinstance(timesteps, torch.Tensor): + timesteps_np = timesteps.detach().cpu().numpy() + else: + timesteps_np = _np.array(timesteps) + + # Use numpy argmin on absolute difference for stability + idx = _np.abs(timesteps_np - timestep_np).argmin() + return int(idx) def add_noise_to_sample( original_samples: torch.Tensor, diff --git a/modules/res4lyf/specialized_rk_scheduler.py b/modules/res4lyf/specialized_rk_scheduler.py index 91f95da18..4e04cc56a 100644 --- a/modules/res4lyf/specialized_rk_scheduler.py +++ b/modules/res4lyf/specialized_rk_scheduler.py @@ -68,7 +68,7 @@ class SpecializedRKScheduler(SchedulerMixin, ConfigMixin): # Internal state self.model_outputs = [] self.sample_at_start_of_step = None - self._step_index = 0 + self._step_index = None def _get_tableau(self): v = self.config.variant @@ -200,19 +200,10 @@ class SpecializedRKScheduler(SchedulerMixin, ConfigMixin): return self._step_index def index_for_timestep(self, timestep, schedule_timesteps=None): - if self._step_index is not None: - return self._step_index - + from .scheduler_utils import index_for_timestep if schedule_timesteps is None: - schedule_timesteps = self._timesteps_cpu - else: - if isinstance(schedule_timesteps, torch.Tensor): - schedule_timesteps = schedule_timesteps.detach().cpu().numpy() - - if isinstance(timestep, torch.Tensor): - timestep = timestep.detach().cpu().numpy() - - return np.abs(schedule_timesteps - timestep).argmin().item() + schedule_timesteps = self.timesteps + return index_for_timestep(timestep, schedule_timesteps) def _init_step_index(self, timestep): if self._step_index is None: @@ -294,6 +285,16 @@ class SpecializedRKScheduler(SchedulerMixin, ConfigMixin): derivative = (sample - denoised) / sigma_t if sigma_t > 1e-6 else torch.zeros_like(sample) if self.sample_at_start_of_step is None: + if stage_index > 0: + # Mid-step fallback for Img2Img/Inpainting + sigma_next_t = self._sigmas_cpu[self._step_index + 1] + dt = sigma_next_t - sigma_t + prev_sample = sample + dt * derivative + self._step_index += 1 + if not return_dict: + return (prev_sample,) + return SchedulerOutput(prev_sample=prev_sample) + self.sample_at_start_of_step = sample self.model_outputs = [derivative] * stage_index diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index 678dbc248..dee2ed6ae 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -332,11 +332,11 @@ samplers_data_diffusers = [ SamplerData('Gauss-Legendre 2S', lambda model: DiffusionSampler('Gauss-Legendre 2S', GaussLegendreScheduler, model), [], {}), SamplerData('Gauss-Legendre 3S', lambda model: DiffusionSampler('Gauss-Legendre 3S', GaussLegendreScheduler, model), [], {}), SamplerData('Gauss-Legendre 4S', lambda model: DiffusionSampler('Gauss-Legendre 4S', GaussLegendreScheduler, model), [], {}), + SamplerData('Specialized-RK 3S', lambda model: DiffusionSampler('Specialized-RK 3S', SpecializedRKScheduler, model), [], {}), + SamplerData('Specialized-RK 4S', lambda model: DiffusionSampler('Specialized-RK 4S', SpecializedRKScheduler, model), [], {}), SamplerData('Runge-Kutta 4/4', lambda model: DiffusionSampler('Runge-Kutta 4/4', RungeKutta44Scheduler, model), [], {}), SamplerData('Runge-Kutta 5/7', lambda model: DiffusionSampler('Runge-Kutta 5/7', RungeKutta57Scheduler, model), [], {}), SamplerData('Runge-Kutta 6/7', lambda model: DiffusionSampler('Runge-Kutta 6/7', RungeKutta67Scheduler, model), [], {}), - SamplerData('Specialized-RK 3S', lambda model: DiffusionSampler('Specialized-RK 3S', SpecializedRKScheduler, model), [], {}), - SamplerData('Specialized-RK 4S', lambda model: DiffusionSampler('Specialized-RK 4S', SpecializedRKScheduler, model), [], {}), SamplerData('Same as primary', None, [], {}), ]