scheduler cleanup

Author: miguelmartin75

src/diffusers/pipelines/cosmos/pipeline_cosmos25_predict.py

@@ -506,6 +506,7 @@ def __call__(
         callback_on_step_end_tensor_inputs: List[str] = ["latents"],
         max_sequence_length: int = 512,
         shift: float = 5.0,
+        timestep_scale: float = 0.001,
         conditional_frame_timestep: float = 0.1,
     ):
         r"""
@@ -635,14 +636,13 @@ def __call__(
             max_sequence_length=max_sequence_length,
         )
 
-        # 4. Prepare timesteps
-        self.scheduler.set_timesteps(num_inference_steps, shift=shift, device=device)
-        timesteps = self.scheduler.timesteps
-
-        # 5. Prepare latent variables
         vae_dtype = self.vae.dtype
         transformer_dtype = self.transformer.dtype
 
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, shift=shift, device=device, scale=timestep_scale)
+        timesteps = torch.tensor(self.scheduler.timesteps).to(transformer_dtype)
+
         num_frames_in = None
         if image is not None:
             # TODO: handle batch_size > 1
@@ -700,12 +700,8 @@ def __call__(
                 if self.interrupt:
                     continue
 
-                self._current_timestep = t
-
-                timestep = torch.stack([t]).to(torch.float32)
-                # TODO: make scheduler scale this instead
-                timestep *= 0.001  # NOTE: timestep scale
-                timestep = timestep.to(transformer_dtype)
+                self._current_timestep = t.cpu().item()
+                timestep = t.unsqueeze(0)
 
                 in_latents = cond_mask * cond_latent + (1 - cond_mask) * latents  # TODO: could use cond_indicator
                 in_latents = in_latents.to(transformer_dtype)

src/diffusers/schedulers/scheduling_flow_unipc_multistep.py

@@ -157,6 +157,7 @@ def set_timesteps(
         sigmas: Optional[List[float]] = None,
         mu: Optional[Union[float, None]] = None,
         shift: Optional[Union[float, None]] = None,
+        scale: float = 0.001,
     ):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -191,10 +192,7 @@ def set_timesteps(
                     shift = self.config.shift
                 sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore
 
-        if self.config.final_sigmas_type == "sigma_min":
-            # TODO(migmartin): this raises an error, rewrite this class
-            sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
-        elif self.config.final_sigmas_type == "zero":
+        if self.config.final_sigmas_type == "zero":
             sigma_last = 0
         else:
             raise ValueError(
@@ -205,7 +203,8 @@ def set_timesteps(
         sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)  # pyright: ignore
 
         self.sigmas = torch.from_numpy(sigmas)
-        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)
+        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.float32)
+        self.timesteps *= scale
 
         self.num_inference_steps = len(timesteps)