res4lyf epsilon validated

Signed-off-by: Vladimir Mandic <mandic00@live.com>

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  

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__}')

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 <https://github.com/ClownsharkBatwing/RES4LYF>
+
+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

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

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

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
 

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:

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
 

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
 

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
 

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)
             

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)

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
 

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
 

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
 

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,

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
 

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, [], {}),
 ]