§3, Eq. 4 — Forward process: sample x_t from q(x_t | x_0). "A notable property is that we can sample x_t at any arbitrary time step t in closed form: q(x_t | x_0) = N(x_t; √α̅_t x_0, (1 - α̅_t)I)" x_t = √α̅_t * x_0 + √(1 - α̅_t) * ε, where ε ~ N(0, I) Args: x_0: (batch, C
(
x_0: torch.Tensor,
t: torch.Tensor,
schedule: Dict[str, torch.Tensor],
noise: Optional[torch.Tensor] = None,
)
| 75 | |
| 76 | |
| 77 | def q_sample( |
| 78 | x_0: torch.Tensor, |
| 79 | t: torch.Tensor, |
| 80 | schedule: Dict[str, torch.Tensor], |
| 81 | noise: Optional[torch.Tensor] = None, |
| 82 | ) -> torch.Tensor: |
| 83 | """§3, Eq. 4 — Forward process: sample x_t from q(x_t | x_0). |
| 84 | |
| 85 | "A notable property is that we can sample x_t at any arbitrary time step t |
| 86 | in closed form: q(x_t | x_0) = N(x_t; √α̅_t x_0, (1 - α̅_t)I)" |
| 87 | |
| 88 | x_t = √α̅_t * x_0 + √(1 - α̅_t) * ε, where ε ~ N(0, I) |
| 89 | |
| 90 | Args: |
| 91 | x_0: (batch, C, H, W) — clean images |
| 92 | t: (batch,) — timestep indices |
| 93 | schedule: precomputed noise schedule |
| 94 | noise: optional pre-sampled noise (for reproducibility) |
| 95 | |
| 96 | Returns: |
| 97 | x_t: (batch, C, H, W) — noisy images at timestep t |
| 98 | """ |
| 99 | if noise is None: |
| 100 | noise = torch.randn_like(x_0) |
| 101 | |
| 102 | # Extract schedule values for timestep t, reshape for broadcasting |
| 103 | sqrt_alpha_cumprod = schedule["sqrt_alphas_cumprod"][t] # (batch,) |
| 104 | sqrt_one_minus_alpha_cumprod = schedule["sqrt_one_minus_alphas_cumprod"][t] # (batch,) |
| 105 | |
| 106 | # Reshape for broadcasting with (batch, C, H, W) |
| 107 | sqrt_alpha_cumprod = sqrt_alpha_cumprod.view(-1, 1, 1, 1) |
| 108 | sqrt_one_minus_alpha_cumprod = sqrt_one_minus_alpha_cumprod.view(-1, 1, 1, 1) |
| 109 | |
| 110 | # §3, Eq. 4 — x_t = √α̅_t * x_0 + √(1 - α̅_t) * ε |
| 111 | return sqrt_alpha_cumprod * x_0 + sqrt_one_minus_alpha_cumprod * noise |
| 112 | |
| 113 | |
| 114 | @torch.no_grad() |