fix audio2audio

2025-05-15 08:15:15 +00:00
parent 05c58339f0
commit ded11fd02f
6 changed files with 78 additions and 29 deletions
@@ -71,9 +71,10 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
        max_shift: Optional[float] = 1.15,
        base_image_seq_len: Optional[int] = 256,
        max_image_seq_len: Optional[int] = 4096,
+        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)