timesteps, set_timesteps(sigmas=..)

Author: hlky

src/diffusers/schedulers/scheduling_euler_discrete.py

@@ -403,10 +403,35 @@ def set_timesteps(
             num_inference_steps = len(timesteps) if timesteps is not None else len(sigmas) - 1
         self.num_inference_steps = num_inference_steps
 
+        if self.config.final_sigmas_type == "sigma_min" and not self.config.use_flow_match:
+            sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
+        elif self.config.final_sigmas_type == "zero":
+            sigma_last = 0
+        elif self.config.invert_sigmas:
+            sigma_last = 1
+        else:
+            raise ValueError(
+                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+            )
+
         if sigmas is not None and not self.config.use_flow_match:
             log_sigmas = np.log(np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5))
             sigmas = np.array(sigmas).astype(np.float32)
             timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas[:-1]])
+        elif sigmas is not None and self.config.use_flow_match:
+            sigmas = np.array(sigmas).astype(np.float32)
+            timesteps = sigmas * self.config.num_train_timesteps
+
+            if self.config.use_flow_match:
+                if self.config.use_dynamic_shifting:
+                    sigmas = self.time_shift(mu, 1.0, sigmas)
+                else:
+                    sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
+
+            if self.config.invert_sigmas:
+                sigmas = 1.0 - sigmas
+            timesteps = sigmas * self.config.num_train_timesteps
+            sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)
         elif sigmas is None:
             if timesteps is not None:
                 timesteps = np.array(timesteps).astype(np.float32)
@@ -482,18 +507,7 @@ def set_timesteps(
 
             if self.config.invert_sigmas:
                 sigmas = 1.0 - sigmas
-                timesteps = sigmas * self.config.num_train_timesteps
-
-            if self.config.final_sigmas_type == "sigma_min":
-                sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
-            elif self.config.final_sigmas_type == "zero":
-                sigma_last = 0
-            elif self.config.invert_sigmas:
-                sigma_last = 1
-            else:
-                raise ValueError(
-                    f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
-                )
+            timesteps = sigmas * self.config.num_train_timesteps
 
             sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)
 
@@ -502,8 +516,6 @@ def set_timesteps(
         # TODO: Support the full EDM scalings for all prediction types and timestep types
         if self.config.timestep_type == "continuous" and self.config.prediction_type == "v_prediction":
             self.timesteps = torch.Tensor([0.25 * sigma.log() for sigma in sigmas[:-1]]).to(device=device)
-        elif self.config.use_flow_match:
-            self.timesteps = sigmas[:-1] * self.config.num_train_timesteps
         else:
             self.timesteps = torch.from_numpy(timesteps.astype(np.float32)).to(device=device)