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 ∼ 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_θ(x_t-1 | x_t) = N(x_t-1; ; μ_θ(x_t, t), ; β_t I)

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

μ_θ(x_t, t) = 1√(α_t) ( x_t - β_t√(1 - α_t) · ε_θ(x_t, t) )

• Phương sai hậu nghiệm β_t:

β_t = 1 - α_t-11 - α_t · β_t

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

x_t-1 = μ_θ(x_t, t) + σ_t · z, với σ_t = √(β_t) (Đặc biệt: nếu t = 0, không cộng thêm nhiễu, tức σ₀ = 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ả.