Bước Lấy Mẫu Khuếch Tán Ngược DDPM (Denoising Step - Ho et al. 2020)

Mục tiêu kiến thức

• Quá trình khuếch tán ngược khử dần nhiễu từ x_T sim N(0, I) về ảnh rõ nét x₀.
• Xác suất hậu nghiệm từng bước (Algorithm 2 - Ho et al. 2020):

p_theta(x_t-1 | x_t) = N(x_t-1; mu_theta(x_t, t), tilde{beta}_t I)

• Kỳ vọng mu_theta(x_t, t) được tính từ ước lượng nhiễu epsilon_theta(x_t, t):

mu_theta(x_t, t) = 1√(alpha_t) ( x_t - beta_t√(1 - alpha_t) · epsilon_theta(x_t, t) )

• Phương sai hậu nghiệm tilde{beta}_t:

tilde{beta}_t = 1 - alpha_t-11 - alpha_t · beta_t

• Công thức cập nhật điểm ảnh:

x_t-1 = mu_theta(x_t, t) + sigma_t · z, với sigma_t = √(tilde{beta}_t) (Đặc biệt: nếu t = 0, không cộng thêm nhiễu, tức sigma₀ = 0).

Yêu cầu

Viết hàm:

def p_sample_step(
    x_t: torch.Tensor,
    t: int,
    predicted_noise: torch.Tensor,
    betas: torch.Tensor,
    alphas: torch.Tensor,
    alphas_cumprod: torch.Tensor,
    z: torch.Tensor | None = None
) -> torch.Tensor:
    pass

• x_t: Tensor (B, C, H, W) hoặc (B, D).
• t: Số nguyên chỉ số bước thời gian t ∈ [0, T-1].
• predicted_noise: Tensor cùng shape với x_t do mô hình dự đoán.
• z: Tensor Gauss ngẫu nhiên N(0, I) hoặc None.
• Trả về x_t-1.

Ràng buộc

• Thời gian chạy tối đa: 2000ms.
• Giới hạn bộ nhớ: 256MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.