Consistency Models: Single-Step Generation

Module: Generative AI | Difficulty: Advanced

Consistency Function

for all along the ODE trajectory.

Boundary Condition

Self-Consistency Property

for any on the same trajectory.

Training

class ConsistencyModel(nn.Module):
    def __init__(self, model, sigma_min=0.002, sigma_max=80):
        super().__init__()
        self.model = model
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max
    def forward(self, x, sigma):
        c_in = 1 / (sigma**2 + 1)**0.5
        c_out = sigma / (sigma**2 + 1)**0.5
        return x + c_out * self.model(c_in * x, sigma)