| """FlowMatchEulerDiscreteScheduler implementation.""" |
|
|
| |
| from dataclasses import dataclass |
| import math |
| from typing import List, Optional, Tuple, Union |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| @dataclass |
| class SchedulerOutput: |
| prev_sample: torch.FloatTensor |
|
|
|
|
| class SchedulerConfig: |
| def __init__(self, **kwargs): |
| self.__dict__.update(kwargs) |
|
|
| def get(self, key, default=None): |
| return self.__dict__.get(key, default) |
|
|
| def __getattr__(self, name): |
| return self.__dict__.get(name) |
|
|
|
|
| class FlowMatchEulerDiscreteScheduler: |
| """Euler scheduler for flow matching.""" |
|
|
| def __init__( |
| self, |
| num_train_timesteps: int = 1000, |
| shift: float = 1.0, |
| use_dynamic_shifting: bool = False, |
| **kwargs, |
| ): |
| self.num_train_timesteps = num_train_timesteps |
| self.shift = shift |
| self.use_dynamic_shifting = use_dynamic_shifting |
| self.config = SchedulerConfig( |
| num_train_timesteps=num_train_timesteps, |
| shift=shift, |
| use_dynamic_shifting=use_dynamic_shifting, |
| ) |
|
|
| timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() |
| timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) |
| sigmas = timesteps / num_train_timesteps |
|
|
| if not use_dynamic_shifting: |
| sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) |
|
|
| self.timesteps = sigmas * num_train_timesteps |
| self.sigmas = sigmas.to("cpu") |
| self.sigma_min = self.sigmas[-1].item() |
| self.sigma_max = self.sigmas[0].item() |
|
|
| self._step_index = None |
| self._begin_index = None |
|
|
| def set_timesteps( |
| self, |
| num_inference_steps: Optional[int] = None, |
| device: Union[str, torch.device] = None, |
| sigmas: Optional[List[float]] = None, |
| mu: Optional[float] = None, |
| timesteps: Optional[List[float]] = None, |
| ): |
| passed_timesteps = timesteps |
| if num_inference_steps is None: |
| num_inference_steps = len(sigmas) if sigmas is not None else len(timesteps) |
|
|
| self.num_inference_steps = num_inference_steps |
|
|
| if sigmas is None: |
| if timesteps is None: |
| timesteps = np.linspace( |
| self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps + 1 |
| )[:-1] |
| sigmas = timesteps / self.num_train_timesteps |
| else: |
| sigmas = np.array(sigmas).astype(np.float32) |
|
|
| if self.use_dynamic_shifting: |
| sigmas = self.time_shift(mu, 1.0, sigmas) |
| else: |
| sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) |
|
|
| sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) |
|
|
| if passed_timesteps is None: |
| timesteps = sigmas * self.num_train_timesteps |
| else: |
| timesteps = torch.from_numpy(passed_timesteps).to(dtype=torch.float32, device=device) |
|
|
| sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) |
|
|
| self.timesteps = timesteps |
| self.sigmas = sigmas |
| self._step_index = None |
| self._begin_index = None |
|
|
| def index_for_timestep(self, timestep, schedule_timesteps=None): |
| if schedule_timesteps is None: |
| schedule_timesteps = self.timesteps |
|
|
| indices = (schedule_timesteps == timestep).nonzero() |
| pos = 1 if len(indices) > 1 else 0 |
| return indices[pos].item() |
|
|
| def _init_step_index(self, timestep): |
| if self._begin_index is None: |
| if isinstance(timestep, torch.Tensor): |
| timestep = timestep.to(self.timesteps.device) |
| self._step_index = self.index_for_timestep(timestep) |
| else: |
| self._step_index = self._begin_index |
|
|
| def step( |
| self, |
| model_output: torch.FloatTensor, |
| timestep: Union[float, torch.FloatTensor], |
| sample: torch.FloatTensor, |
| return_dict: bool = True, |
| **kwargs, |
| ) -> Union[SchedulerOutput, Tuple]: |
| """Predict the sample at the previous timestep.""" |
| if self._step_index is None: |
| self._init_step_index(timestep) |
|
|
| sample = sample.to(torch.float32) |
| sigma_idx = self._step_index |
| sigma = self.sigmas[sigma_idx] |
| sigma_next = self.sigmas[sigma_idx + 1] |
|
|
| dt = sigma_next - sigma |
| prev_sample = sample + dt * model_output |
| self._step_index += 1 |
| prev_sample = prev_sample.to(model_output.dtype) |
|
|
| if not return_dict: |
| return (prev_sample,) |
| return SchedulerOutput(prev_sample=prev_sample) |
|
|
| def _sigma_to_t(self, sigma): |
| return sigma * self.num_train_timesteps |
|
|
| def time_shift(self, mu: float, sigma: float, t: torch.Tensor): |
| return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) |
|
|
