File size: 570 Bytes
f15d29e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import abc
from typing import Generic, TypeVar
import torch
from ...diffusion.data.batched_data import BatchedData
Diffusable = TypeVar("Diffusable", bound=BatchedData)
class ScoreModel(torch.nn.Module, Generic[Diffusable], abc.ABC):
"""Abstract base class for score models."""
@abc.abstractmethod
def forward(self, x: Diffusable, t: torch.Tensor) -> Diffusable:
"""Args:
x: batch of noisy data
t: timestep. Shape (batch_size, 1)
"""
...
|