Fix division-by-zero NaN in shift_terminal with a single denoising step

src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

@@ -350,8 +350,10 @@ def set_timesteps(
         else:
             sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)
 
-        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
-        if self.config.shift_terminal:
+        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value. This is
+        #    skipped when there is only a single step, since there is nothing to stretch and the terminal rescaling
+        #    otherwise divides by zero (the single sigma is always 1.0, i.e. `one_minus_z[-1]` is always 0).
+        if self.config.shift_terminal and len(sigmas) > 1:
             sigmas = self.stretch_shift_to_terminal(sigmas)
 
         # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules

src/diffusers/schedulers/scheduling_flow_match_lcm.py

@@ -359,8 +359,10 @@ def set_timesteps(
         else:
             sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)  # type: ignore
 
-        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
-        if self.config.shift_terminal:
+        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value. This is
+        #    skipped when there is only a single step, since there is nothing to stretch and the terminal rescaling
+        #    otherwise divides by zero (the single sigma is always 1.0, i.e. `one_minus_z[-1]` is always 0).
+        if self.config.shift_terminal and len(sigmas) > 1:
             sigmas = self.stretch_shift_to_terminal(sigmas)  # type: ignore
 
         # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules

src/diffusers/schedulers/scheduling_unipc_multistep.py

@@ -432,7 +432,7 @@ def set_timesteps(
                 sigmas = self.time_shift(mu, 1.0, sigmas)
             else:
                 sigmas = self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas)
-            if self.config.shift_terminal:
+            if self.config.shift_terminal and len(sigmas) > 1:
                 sigmas = self.stretch_shift_to_terminal(sigmas)
             eps = 1e-6
             if np.fabs(sigmas[0] - 1) < eps: