Canopus51/Cosmic-Large-Scale-Structure-Generator

1import torch
2
3def get_diffusion_params(T=400, beta_start=1e-4, beta_end=0.02, device="cpu"):
4    """Culculate the diffusion parameters"""
5    betas = torch.linspace(beta_start, beta_end, T).to(device)
6    alphas = 1.0 - betas
7    alphas_cumprod = torch.cumprod(alphas, dim=0)
8    
9    return {
10        "betas": betas,
11        "alphas_cumprod": alphas_cumprod,
12        "sqrt_one_minus_alphas_cumprod": torch.sqrt(1.0 - alphas_cumprod),
13        "sqrt_recip_alphas": torch.sqrt(1.0 / alphas),
14        "T": T
15    }
16
17@torch.no_grad()
18def p_sample(model, x_t, t, params):
19    """DDPM single-step reverse sampling: x_t -> x_{t-1}"""
20    eps_theta = model(x_t, t)
21    
22    beta_t = params["betas"][t].view(-1, 1, 1, 1)
23    sqrt_recip_alpha_t = params["sqrt_recip_alphas"][t].view(-1, 1, 1, 1)
24    sqrt_om_ac_t = params["sqrt_one_minus_alphas_cumprod"][t].view(-1, 1, 1, 1)
25
26    # DDPM reverse mean formula
27    mean = sqrt_recip_alpha_t * (x_t - (beta_t / sqrt_om_ac_t) * eps_theta)
28
29    # Do not add noise when t=0
30    if (t == 0).all():
31        return mean
32    
33    noise = torch.randn_like(x_t)
34    return mean + torch.sqrt(beta_t) * noise
35