Fix sigmas/timesteps

2026-04-18 01:30:02 +00:00 · 2025-06-03 20:55:29 -04:00
parent 415233993a
commit db05136480
4 changed files with 5 additions and 5 deletions
--- a/tests/library/test_flux_train_utils.py
+++ b/tests/library/test_flux_train_utils.py
@@ -13,6 +13,9 @@ class MockNoiseScheduler:
        self.config.num_train_timesteps = num_train_timesteps
        self.timesteps = torch.arange(num_train_timesteps, dtype=torch.long)

+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+

 # Create fixtures for commonly used objects
@pytest.fixture