Fix res_multistep steps

2025-12-29 15:39:27 +02:00
parent 351829a2b6
commit be41f67fae
1 changed files with 4 additions and 6 deletions
@@ -35,9 +35,7 @@ class FlowMatchSchedulerResMultistep():
        self.sigmas = torch.FloatTensor(sigmas)
        self.sigmas = self.shift * self.sigmas / \
             (1 + (self.shift - 1) * self.sigmas)
-        self.timesteps = self.sigmas * self.num_train_timesteps
-        #print(f"Timesteps: {self.timesteps}, Sigmas: {self.sigmas}")
-
+        self.timesteps = self.sigmas[:-1] * self.num_train_timesteps

    def step(self, model_output, timestep, sample):
        if timestep.ndim == 2: