Added ability to set timesteps to linear for flowmatching schedule

2026-04-26 09:19:20 +00:00 · 2024-08-11 13:06:08 -06:00
parent f8f0657b68
commit a6aa4b2c7d
3 changed files with 19 additions and 11 deletions
--- a/jobs/process/BaseSDTrainProcess.py
+++ b/jobs/process/BaseSDTrainProcess.py
@@ -908,7 +908,9 @@ class BaseSDTrainProcess(BaseTrainProcess):
                    )
                elif self.train_config.noise_scheduler == 'flowmatch':
                    self.sd.noise_scheduler.set_train_timesteps(
-                        num_train_timesteps, device=self.device_torch
+                        num_train_timesteps,
+                        device=self.device_torch,
+                        linear=self.train_config.linear_timesteps
                    )
                else:
                    self.sd.noise_scheduler.set_timesteps(