Bước Lấy Mẫu Tất Định Tốc Độ Cao DDIM (Song et al. ICLR 2021)

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

• DDPM đòi hỏi chạy tuần tự đủ T = 1000 bước gây tốn kém thời gian suy luận (10-30 giây cho một bức ảnh).
• Song et al. (2020) chứng minh rằng có một họ các quá trình phi-Markovian (Non-Markovian) chia sẻ cùng phân phối biên q(x_t | x₀) nhưng cho phép suy luận với bước nhảy bất kỳ (chỉ cần 20 - 50 bước).
• Công thức cập nhật DDIM:
• Ước lượng ảnh gốc x₀:

x₀ = x_t - √(1 - alpha_t) · epsilon_theta(x_t, t)√(alpha_t)

• Phương sai nhiễu:

sigma_t = eta · √(1 - alpha_t-11 - alpha_t) √(1 - alpha_talpha_t-1)

• Khi eta = 0 (DDIM Tất Định / ODE Trajectory): sigma_t = 0.
• Cập nhật bước tiếp theo x_t-1:

x_t-1 = √(alpha_t-1) · x₀ + √(1 - alpha_t-1 - sigma_t²) · epsilon_theta(x_t, t) + sigma_t · z

Yêu cầu

Viết hàm:

def ddim_step(
    x_t: torch.Tensor,
    predicted_noise: torch.Tensor,
    alpha_bar_t: float,
    alpha_bar_prev: float,
    eta: float = 0.0,
    z: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
    pass

• Trả về tuple (x_prev, pred_x0):
• x_prev: Tensor x_t-1.
• pred_x0: Tensor ảnh gốc ước lượng x₀.

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ả.