Diffusion

# 디퓨전 핵심 개념 구현 (간소화 버전)
import torch
import torch.nn as nn
import numpy as np

class SimpleDiffusion:
    """간단한 디퓨전 프로세스 구현"""
    def __init__(self, timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.timesteps = timesteps
        # 선형 노이즈 스케줄
        self.betas = torch.linspace(beta_start, beta_end, timesteps)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)

    def forward_process(self, x0, t):
        """Forward diffusion: x0에 노이즈 추가하여 xt 생성"""
        sqrt_alpha_cumprod = torch.sqrt(self.alphas_cumprod[t])
        sqrt_one_minus = torch.sqrt(1.0 - self.alphas_cumprod[t])
        noise = torch.randn_like(x0)
        xt = sqrt_alpha_cumprod * x0 + sqrt_one_minus * noise
        return xt, noise

    def reverse_step(self, xt, predicted_noise, t):
        """Reverse diffusion: 예측 노이즈로 xt-1 복원"""
        alpha = self.alphas[t]
        alpha_cumprod = self.alphas_cumprod[t]
        beta = self.betas[t]

        # 평균 계산
        mean = (1 / torch.sqrt(alpha)) * (
            xt - (beta / torch.sqrt(1 - alpha_cumprod)) * predicted_noise
        )
        # 노이즈 추가 (t > 0일 때)
        if t > 0:
            noise = torch.randn_like(xt)
            return mean + torch.sqrt(beta) * noise
        return mean

# 학습 루프 예시
diffusion = SimpleDiffusion()
model = nn.Sequential(nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 784))

for epoch in range(100):
    x0 = torch.randn(32, 784)  # 배치 데이터
    t = torch.randint(0, 1000, (32,))  # 랜덤 타임스텝
    xt, noise = diffusion.forward_process(x0, t[0])
    predicted_noise = model(xt)
    loss = nn.MSELoss()(predicted_noise, noise)
    # loss.backward() 및 optimizer.step()...