Compute the intermediate time steps for sampling. Args: skip_type: A `str`. The type for the spacing of the time steps. We support three types: - 'logSNR': uniform logSNR for the time steps. - 'time_uniform': uniform time for the time steps. (**Rec
(self, skip_type, t_T, t_0, N, device)
| 374 | return self.noise_prediction_fn(x, t) |
| 375 | |
| 376 | def get_time_steps(self, skip_type, t_T, t_0, N, device): |
| 377 | """Compute the intermediate time steps for sampling. |
| 378 | Args: |
| 379 | skip_type: A `str`. The type for the spacing of the time steps. We support three types: |
| 380 | - 'logSNR': uniform logSNR for the time steps. |
| 381 | - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.) |
| 382 | - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.) |
| 383 | t_T: A `float`. The starting time of the sampling (default is T). |
| 384 | t_0: A `float`. The ending time of the sampling (default is epsilon). |
| 385 | N: A `int`. The total number of the spacing of the time steps. |
| 386 | device: A torch device. |
| 387 | Returns: |
| 388 | A pytorch tensor of the time steps, with the shape (N + 1,). |
| 389 | """ |
| 390 | if skip_type == 'logSNR': |
| 391 | lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device)) |
| 392 | lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device)) |
| 393 | logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device) |
| 394 | return self.noise_schedule.inverse_lambda(logSNR_steps) |
| 395 | elif skip_type == 'time_uniform': |
| 396 | return torch.linspace(t_T, t_0, N + 1).to(device) |
| 397 | elif skip_type == 'time_quadratic': |
| 398 | t_order = 2 |
| 399 | t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device) |
| 400 | return t |
| 401 | else: |
| 402 | raise ValueError( |
| 403 | "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type)) |
| 404 | |
| 405 | def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device): |
| 406 | """ |
no test coverage detected