Function sample

skills/paper2code/worked/ddpm/src/utils.py:164–200

Algorithm 2 — Full reverse sampling process. "Algorithm 2 Sampling 1: x_T ~ N(0, I) 2: for t = T, ..., 1 do 3: z ~ N(0, I) if t > 1, else z = 0 4: x_{t-1} = 1/√α_t (x_t - β_t/√(1-α̅_t) ε_θ(x_t, t)) + σ_t z 5: end for 6: return x_0" Args: model: noi

(
    model: nn.Module,
    schedule: Dict[str, torch.Tensor],
    image_shape: tuple,
    device: torch.device,
)

@torch.no_grad()
def sample(
    model: nn.Module,
    schedule: Dict[str, torch.Tensor],
    image_shape: tuple,
    device: torch.device,
) -> torch.Tensor:
    """Algorithm 2 — Full reverse sampling process.

    "Algorithm 2 Sampling
     1: x_T ~ N(0, I)
     2: for t = T, ..., 1 do
     3:   z ~ N(0, I) if t > 1, else z = 0
     4:   x_{t-1} = 1/√α_t (x_t - β_t/√(1-α̅_t) ε_θ(x_t, t)) + σ_t z
     5: end for
     6: return x_0"

    Args:
        model: noise prediction network ε_θ (should be in eval mode, ideally EMA weights)
        schedule: precomputed noise schedule
        image_shape: (batch, C, H, W) — shape of images to generate
        device: torch device

    Returns:
        x_0: (batch, C, H, W) — generated images
    """
    model.eval()
    timesteps = len(schedule["betas"])

    # Algorithm 2, line 1: x_T ~ N(0, I)
    x = torch.randn(image_shape, device=device)

    # Algorithm 2, lines 2-5: reverse iterate from t=T to t=1
    for t_index in reversed(range(timesteps)):
        t = torch.full((image_shape[0],), t_index, device=device, dtype=torch.long)
        x = p_sample(model, x, t, t_index, schedule)

    return x


class EMA: