| from abc import ABC |
| from typing import Sequence, Union |
| import torch |
| from torch.distributions import LogisticNormal |
|
|
| class Timesteps(ABC): |
| """ |
| Timesteps base class. |
| """ |
|
|
| def __init__(self, T: Union[int, float]): |
| assert T > 0 |
| self._T = T |
|
|
| @property |
| def T(self) -> Union[int, float]: |
| """ |
| Maximum timestep inclusive. |
| int if discrete, float if continuous. |
| """ |
| return self._T |
|
|
| def is_continuous(self) -> bool: |
| """ |
| Whether the schedule is continuous. |
| """ |
| return isinstance(self.T, float) |
|
|
| class LogitNormalTrainingTimesteps(Timesteps): |
| """ |
| Logit-Normal sampling of timesteps in [0, T]. |
| """ |
|
|
| def __init__(self, T: Union[int, float], loc: float, scale: float): |
| super().__init__(T) |
| self.dist = LogisticNormal(loc, scale) |
|
|
| def sample( |
| self, |
| size: Sequence[int], |
| device: torch.device = "cpu", |
| ) -> torch.Tensor: |
| t = self.dist.sample(size)[..., 0].to(device).mul_(self.T) |
| return t if self.is_continuous() else t.round().int() |
|
|
