| 151 | |
| 152 | |
| 153 | class FlowMatchScheduler(): |
| 154 | |
| 155 | def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=5.0, sigma_max=1.0, sigma_min=0.0, inverse_timesteps=False, extra_one_step=True, reverse_sigmas=False): |
| 156 | self.num_train_timesteps = num_train_timesteps |
| 157 | self.shift = shift |
| 158 | self.sigma_max = sigma_max |
| 159 | self.sigma_min = sigma_min |
| 160 | self.inverse_timesteps = inverse_timesteps |
| 161 | self.extra_one_step = extra_one_step |
| 162 | self.reverse_sigmas = reverse_sigmas |
| 163 | self.num_inference_steps = num_inference_steps |
| 164 | self.set_timesteps(num_inference_steps) |
| 165 | |
| 166 | def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False): |
| 167 | sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength |
| 168 | if self.extra_one_step: |
| 169 | self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] |
| 170 | else: |
| 171 | self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) |
| 172 | if self.inverse_timesteps: |
| 173 | self.sigmas = torch.flip(self.sigmas, dims=[0]) |
| 174 | self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) |
| 175 | if self.reverse_sigmas: |
| 176 | self.sigmas = 1 - self.sigmas |
| 177 | self.timesteps = self.sigmas * self.num_train_timesteps |
| 178 | if training: |
| 179 | x = self.timesteps |
| 180 | y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2) |
| 181 | y_shifted = y - y.min() |
| 182 | bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) |
| 183 | self.linear_timesteps_weights = bsmntw_weighing |
| 184 | |
| 185 | |
| 186 | def step(self, model_output, timestep, sample, to_final=False): |
| 187 | if isinstance(timestep, torch.Tensor): |
| 188 | timestep = timestep.cpu() |
| 189 | timestep_id = torch.argmin((self.timesteps - timestep).abs()) |
| 190 | sigma = self.sigmas[timestep_id] |
| 191 | if to_final or timestep_id + 1 >= len(self.timesteps): |
| 192 | sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 |
| 193 | else: |
| 194 | sigma_ = self.sigmas[timestep_id + 1] |
| 195 | prev_sample = sample + model_output * (sigma_ - sigma) |
| 196 | return prev_sample |
| 197 | |
| 198 | |
| 199 | def return_to_timestep(self, timestep, sample, sample_stablized): |
| 200 | if isinstance(timestep, torch.Tensor): |
| 201 | timestep = timestep.cpu() |
| 202 | timestep_id = torch.argmin((self.timesteps - timestep).abs()) |
| 203 | sigma = self.sigmas[timestep_id] |
| 204 | model_output = (sample - sample_stablized) / sigma |
| 205 | return model_output |
| 206 | |
| 207 | |
| 208 | def add_noise(self, original_samples, noise, timestep): |
| 209 | if isinstance(timestep, torch.Tensor): |
| 210 | timestep = timestep.cpu() |