| 17 | |
| 18 | |
| 19 | class FlowModelObj(FlowModel): |
| 20 | def __init__( |
| 21 | self, |
| 22 | enforce_zero_snr: bool = True, |
| 23 | diffusion_parameterization: str = 'v', |
| 24 | diffusion_schedule: str = 'linear', |
| 25 | *args, |
| 26 | **kwargs |
| 27 | ): |
| 28 | super().__init__(*args, **kwargs) |
| 29 | self.register_sdv2_schedule(diffusion_schedule, enforce_zero_snr) |
| 30 | assert diffusion_parameterization in ['v', 'eps'], 'Diffusion parameterization has to be either v or eps' |
| 31 | self.diffusion_parameterization = diffusion_parameterization |
| 32 | self.diffusion_schedule = diffusion_schedule |
| 33 | |
| 34 | def ode_fn(self, t, x, **kwargs): |
| 35 | if t.numel() == 1: |
| 36 | t = t.expand(x.shape[0]) |
| 37 | _pred = self.sample_vt(x, t, **kwargs) |
| 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))) |
nothing calls this directly
no outgoing calls
no test coverage detected