(self, diffusion_schedule, enforce_zero_snr=True)
| 38 | return _pred |
| 39 | |
| 40 | def register_sdv2_schedule(self, diffusion_schedule, enforce_zero_snr=True): |
| 41 | # SDV2 schedule |
| 42 | linear_start = 0.00085 |
| 43 | linear_end = 0.0120 |
| 44 | |
| 45 | betas = make_beta_schedule( |
| 46 | diffusion_schedule, |
| 47 | n_timestep=1000, |
| 48 | linear_start=linear_start, |
| 49 | linear_end=linear_end, |
| 50 | ) |
| 51 | if enforce_zero_snr: |
| 52 | betas = enforce_zero_terminal_snr(betas) |
| 53 | alphas = 1. - betas |
| 54 | alphas_cumprod = np.cumprod(alphas, axis=0) |
| 55 | alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) |
| 56 | alphas_cumprod_full = np.append(1., alphas_cumprod) |
| 57 | |
| 58 | timesteps, = betas.shape |
| 59 | self.num_timesteps = int(timesteps) |
| 60 | self.linear_start = linear_start |
| 61 | self.linear_end = linear_end |
| 62 | assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep' |
| 63 | |
| 64 | to_torch = partial(torch.tensor, dtype=torch.float32) |
| 65 | |
| 66 | self.register_buffer('betas', to_torch(betas)) |
| 67 | self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) |
| 68 | # self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) |
| 69 | self.register_buffer('alphas_cumprod_full', to_torch(alphas_cumprod_full)) |
| 70 | |
| 71 | self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) |
| 72 | self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) |
| 73 | self.register_buffer('sqrt_alphas_cumprod_full', to_torch(np.sqrt(alphas_cumprod_full))) |
| 74 | self.register_buffer('sqrt_one_minus_alphas_cumprod_full', to_torch(np.sqrt(1. - alphas_cumprod_full))) |
| 75 | |
| 76 | self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) |
| 77 | self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) |
| 78 | |
| 79 | self.register_buffer('rectified_alphas_cumprod_full', self.sqrt_alphas_cumprod_full / (self.sqrt_alphas_cumprod_full + self.sqrt_one_minus_alphas_cumprod_full)) |
| 80 | self.register_buffer('rectified_sqrt_alphas_cumprod_full', self.sqrt_one_minus_alphas_cumprod_full / (self.sqrt_alphas_cumprod_full + self.sqrt_one_minus_alphas_cumprod_full)) |
| 81 | |
| 82 | def sample_vt(self, fm_x, fm_t, **kwargs): |
| 83 | """ |
no test coverage detected