|
|
from abc import ABC, abstractmethod |
|
|
from typing import Sequence, Union |
|
|
import torch |
|
|
|
|
|
from ..types import SamplingDirection |
|
|
|
|
|
|
|
|
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 SamplingTimesteps(Timesteps): |
|
|
""" |
|
|
Sampling timesteps. |
|
|
It defines the discretization of sampling steps. |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
T: Union[int, float], |
|
|
timesteps: torch.Tensor, |
|
|
direction: SamplingDirection, |
|
|
): |
|
|
assert timesteps.ndim == 1 |
|
|
super().__init__(T) |
|
|
self.timesteps = timesteps |
|
|
self.direction = direction |
|
|
|
|
|
def __len__(self) -> int: |
|
|
""" |
|
|
Number of sampling steps. |
|
|
""" |
|
|
return len(self.timesteps) |
|
|
|
|
|
def __getitem__(self, idx: Union[int, torch.IntTensor]) -> torch.Tensor: |
|
|
""" |
|
|
The timestep at the sampling step. |
|
|
Returns a scalar tensor if idx is int, |
|
|
or tensor of the same size if idx is a tensor. |
|
|
""" |
|
|
return self.timesteps[idx] |
|
|
|
|
|
def index(self, t: torch.Tensor) -> torch.Tensor: |
|
|
""" |
|
|
Find index by t. |
|
|
Return index of the same shape as t. |
|
|
Index is -1 if t not found in timesteps. |
|
|
""" |
|
|
i, j = t.reshape(-1, 1).eq(self.timesteps).nonzero(as_tuple=True) |
|
|
idx = torch.full_like(t, fill_value=-1, dtype=torch.int) |
|
|
idx.view(-1)[i] = j.int() |
|
|
return idx |
|
|
|
