fix audio2audio

2025-05-15 08:15:15 +00:00
parent 05c58339f0
commit ded11fd02f
6 changed files with 78 additions and 29 deletions
@@ -66,9 +66,10 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
        self,
        num_train_timesteps: int = 1000,
        shift: float = 1.0,
+        sigma_max: Optional[float] = 1.0,
    ):
        timesteps = np.linspace(
-            1, num_train_timesteps, num_train_timesteps, dtype=np.float32
+            1.0, sigma_max*num_train_timesteps, num_train_timesteps, dtype=np.float32
        )[::-1].copy()
        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)