File size: 716 Bytes
f15d29e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
from typing import Protocol
from ...diffusion.corruption.corruption import Corruption
from ...diffusion.corruption.sde_lib import ScoreFunction
from ...diffusion.sampling.predictors import Predictor
from ...diffusion.sampling.predictors_correctors import LangevinCorrector
class PredictorPartial(Protocol):
def __call__(self, *, corruption: Corruption, score_fn: ScoreFunction | None) -> Predictor:
raise NotImplementedError
class CorrectorPartial(Protocol):
def __call__(
self, *, corruption: Corruption, n_steps: int, score_fn: ScoreFunction | None
) -> LangevinCorrector:
raise NotImplementedError
|