fix sigmas always populated after set_timesteps

This commit is contained in:
CalamitousFelicitousness
2026-04-09 16:16:23 +01:00
parent 4d347049eb
commit 153a5febde
+14 -2
View File
@@ -234,7 +234,13 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin):
self.num_inference_steps = len(timesteps)
timesteps = np.array(timesteps, dtype=np.int64)
self.timesteps = torch.from_numpy(timesteps).to(device)
self.sigmas = None
# Compute sigmas from alphas_cumprod so pipelines can access self.sigmas
sigmas_arr = []
for t in timesteps:
acp = self.alphas_cumprod[t].item()
sigmas_arr.append(((1 - acp) / max(acp, 1e-8)) ** 0.5)
sigmas_arr.append(0.0)
self.sigmas = torch.tensor(sigmas_arr, dtype=torch.float64, device=device)
self._flow_alphas = None
self._flow_sigmas = None
self._flow_lambdas = None
@@ -257,7 +263,13 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin):
raise ValueError(f"{self.config.timestep_spacing} is not supported")
self.timesteps = torch.from_numpy(timesteps).to(device)
self.sigmas = None
# Compute sigmas from alphas_cumprod so pipelines can access self.sigmas
sigmas_arr = []
for t in timesteps:
acp = self.alphas_cumprod[t].item()
sigmas_arr.append(((1 - acp) / max(acp, 1e-8)) ** 0.5)
sigmas_arr.append(0.0)
self.sigmas = torch.tensor(sigmas_arr, dtype=torch.float64, device=device)
self._flow_alphas = None
self._flow_sigmas = None
self._flow_lambdas = None