(self, noise_steps=1000, beta_start=1e-4, beta_end=0.02, img_size=256, device="cuda")
| 14 | |
| 15 | class Diffusion: |
| 16 | def __init__(self, noise_steps=1000, beta_start=1e-4, beta_end=0.02, img_size=256, device="cuda"): |
| 17 | self.noise_steps = noise_steps |
| 18 | self.beta_start = beta_start |
| 19 | self.beta_end = beta_end |
| 20 | self.img_size = img_size |
| 21 | self.device = device |
| 22 | |
| 23 | self.beta = self.prepare_noise_schedule().to(device) |
| 24 | self.alpha = 1. - self.beta |
| 25 | self.alpha_hat = torch.cumprod(self.alpha, dim=0) |
| 26 | |
| 27 | def prepare_noise_schedule(self): |
| 28 | return torch.linspace(self.beta_start, self.beta_end, self.noise_steps) |
nothing calls this directly
no test coverage detected