Make start/end_percent for timestep scheduling choose steps more predictably across nodes

comfy/ldm/sensenova/sampling.py

@@ -57,7 +57,7 @@ def percent_to_sigma(self, percent):
             return 1.0
         if percent >= 1.0:
             return 0.0
-        return float(time_snr_shift(self.shift, 1.0 - percent))
+        return self.sigma(torch.tensor(percent * self.multiplier, dtype=torch.float32)).item()
 
     def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
         sigma = comfy.model_sampling.reshape_sigma(sigma, noise.ndim)

comfy/model_sampling.py

@@ -265,7 +265,7 @@ def percent_to_sigma(self, percent):
         percent = 1.0 - percent
 
         log_sigma_min = math.log(self.sigma_min)
-        return math.exp((math.log(self.sigma_max) - log_sigma_min) * percent + log_sigma_min)
+        return torch.tensor(math.exp((math.log(self.sigma_max) - log_sigma_min) * percent + log_sigma_min), dtype=torch.float32).item()
 
 
 class ModelSamplingContinuousV(ModelSamplingContinuousEDM):
@@ -323,7 +323,7 @@ def percent_to_sigma(self, percent):
             return 1.0
         if percent >= 1.0:
             return 0.0
-        return time_snr_shift(self.shift, 1.0 - percent)
+        return self.sigma(torch.tensor((1.0 - percent) * self.multiplier, dtype=torch.float32)).item()
 
 class ModelSamplingAV(ModelSamplingDiscreteFlow):
     """Flow sampling for packed audio-video latents whose audio stream has its own flow shift.
@@ -437,7 +437,7 @@ def percent_to_sigma(self, percent):
             return 1.0
         if percent >= 1.0:
             return 0.0
-        return flux_time_shift(self.shift, 1.0, 1.0 - percent)
+        return self.sigma(torch.tensor(1.0 - percent, dtype=torch.float32)).item()
 
 
 class ModelSamplingCosmosRFlow(ModelSamplingContinuousEDM):

comfy_extras/nodes_easycache.py

@@ -227,7 +227,7 @@ def __init__(self, reuse_threshold: float, start_percent: float, end_percent: fl
         self.output_channels = output_channels
 
     def is_past_end_timestep(self, timestep: float) -> bool:
-        return not (timestep[0] > self.end_t).item()
+        return not (timestep[0] >= self.end_t).item()
 
     def should_do_easycache(self, timestep: float) -> bool:
         return (timestep[0] <= self.start_t).item()
@@ -418,7 +418,7 @@ def has_cache_diff(self) -> bool:
         return self.cache_diff is not None
 
     def is_past_end_timestep(self, timestep: float) -> bool:
-        return not (timestep[0] > self.end_t).item()
+        return not (timestep[0] >= self.end_t).item()
 
     def should_do_easycache(self, timestep: float) -> bool:
         return (timestep[0] <= self.start_t).item()

comfy_extras/nodes_hidream_o1.py

@@ -200,9 +200,8 @@ def execute(cls, *, model, start_percent: float, end_percent: float, pattern: st
 
         m = model.clone()
         model_sampling = m.get_model_object("model_sampling")
-        multiplier = float(model_sampling.multiplier)
-        start_t = float(model_sampling.percent_to_sigma(start_percent)) * multiplier
-        end_t = float(model_sampling.percent_to_sigma(end_percent)) * multiplier
+        start_sigma = model_sampling.percent_to_sigma(start_percent)
+        end_sigma = model_sampling.percent_to_sigma(end_percent)
 
         edge_ramp_cache: dict = {}
 
@@ -224,15 +223,15 @@ def get_edge_ramp(H: int, W: int, device, dtype) -> torch.Tensor:
 
         def smoothing_wrapper(executor, *args, **kwargs):
             x = args[0]
-            t = float(args[1][0])
+            sigma = float(args[3]["sigmas"][0])
             pred = executor(*args, **kwargs)
-            if not (end_t <= t <= start_t):
+            if not (end_sigma <= sigma <= start_sigma):
                 return pred
             # Pick shift-level by sigma phase across the gated range.
             if len(shift_levels) == 1:
                 level_idx = 0
             else:
-                phase = (start_t - t) / max(start_t - end_t, 1e-8)
+                phase = (start_sigma - sigma) / max(start_sigma - end_sigma, 1e-8)
                 level_idx = min(int(phase * len(shift_levels)), len(shift_levels) - 1)
             shifts = shift_levels[level_idx]
             window_tiles = window_tile_levels[level_idx]