Algorithm 2, lines 3-4 — Single reverse step: sample x_{t-1} from p_θ(x_{t-1} | x_t). "x_{t-1} = 1/√α_t * (x_t - β_t/√(1-α̅_t) * ε_θ(x_t, t)) + σ_t * z" where z ~ N(0, I) if t > 1, else z = 0. Args: model: noise prediction network ε_θ x_t: (batch, C, H, W) — current no
(
model: nn.Module,
x_t: torch.Tensor,
t: torch.Tensor,
t_index: int,
schedule: Dict[str, torch.Tensor],
)
| 113 | |
| 114 | @torch.no_grad() |
| 115 | def p_sample( |
| 116 | model: nn.Module, |
| 117 | x_t: torch.Tensor, |
| 118 | t: torch.Tensor, |
| 119 | t_index: int, |
| 120 | schedule: Dict[str, torch.Tensor], |
| 121 | ) -> torch.Tensor: |
| 122 | """Algorithm 2, lines 3-4 — Single reverse step: sample x_{t-1} from p_θ(x_{t-1} | x_t). |
| 123 | |
| 124 | "x_{t-1} = 1/√α_t * (x_t - β_t/√(1-α̅_t) * ε_θ(x_t, t)) + σ_t * z" |
| 125 | |
| 126 | where z ~ N(0, I) if t > 1, else z = 0. |
| 127 | |
| 128 | Args: |
| 129 | model: noise prediction network ε_θ |
| 130 | x_t: (batch, C, H, W) — current noisy sample |
| 131 | t: (batch,) — current timestep (as tensor for model input) |
| 132 | t_index: integer timestep (for indexing schedule) |
| 133 | schedule: precomputed noise schedule |
| 134 | |
| 135 | Returns: |
| 136 | x_{t-1}: (batch, C, H, W) — denoised sample one step |
| 137 | """ |
| 138 | # Predict noise ε_θ(x_t, t) |
| 139 | predicted_noise = model(x_t, t) # (batch, C, H, W) |
| 140 | |
| 141 | # Extract schedule values |
| 142 | beta_t = schedule["betas"][t_index] |
| 143 | sqrt_recip_alpha_t = schedule["sqrt_recip_alphas"][t_index] |
| 144 | sqrt_one_minus_alpha_cumprod_t = schedule["sqrt_one_minus_alphas_cumprod"][t_index] |
| 145 | |
| 146 | # Algorithm 2, line 4 — Compute mean of p_θ(x_{t-1} | x_t) |
| 147 | # μ_θ = 1/√α_t * (x_t - β_t/√(1-α̅_t) * ε_θ(x_t, t)) |
| 148 | mean = sqrt_recip_alpha_t * ( |
| 149 | x_t - beta_t / sqrt_one_minus_alpha_cumprod_t * predicted_noise |
| 150 | ) |
| 151 | |
| 152 | if t_index == 0: |
| 153 | # Algorithm 2, line 3 — z = 0 when t = 1 (final step) |
| 154 | return mean |
| 155 | else: |
| 156 | # Algorithm 2, line 3 — z ~ N(0, I) when t > 1 |
| 157 | # §3.4 — σ²_t = β_t (fixed small variance) |
| 158 | sigma_t = torch.sqrt(schedule["betas"][t_index]) |
| 159 | noise = torch.randn_like(x_t) |
| 160 | return mean + sigma_t * noise |
| 161 | |
| 162 | |
| 163 | @torch.no_grad() |