| import torch |
| from samplers.general_solver import ODESolver |
|
|
|
|
|
|
| class Euler(ODESolver): |
| def __init__(self, noise_schedule, algorithm_type="data_prediction"): |
| ''' |
| algorithm_type needs to be data_prediction |
| ''' |
| super().__init__(noise_schedule, algorithm_type) |
| self.noise_schedule = noise_schedule |
| self.predict_x0 = algorithm_type == "data_prediction" |
| assert self.predict_x0, "Only data prediction is supported for now." |
| |
| |
| def sample( |
| self, |
| model_fn, |
| x, |
| steps=20, |
| t_start=0.002, |
| t_end=80., |
| skip_type="edm", flags=None, |
| ): |
| self.model = lambda x, t: model_fn(x, t.expand((x.shape[0]))) |
| t_0 = t_end |
| t_T = t_start |
| device = x.device |
| |
| timesteps, timesteps2 = self.prepare_timesteps(steps=steps, t_start=t_T, t_end=t_0, skip_type=skip_type, device=device, load_from=flags.load_from) |
|
|
| with torch.no_grad(): |
| return self.sample_simple(model_fn, x, timesteps, timesteps2, NFEs=steps) |
| |
| def sample_simple(self, model_fn, x, timesteps, timesteps2, NFEs=20, condition=None, unconditional_condition=None, **kwargs): |
| self.model = lambda x, t: model_fn(x, t.expand((x.shape[0])), condition, unconditional_condition) |
| steps = NFEs |
| |
| x_next = x |
| for step in range(steps): |
| t_cur1, t_next1 = timesteps[step], timesteps[step + 1] |
| t_cur2 = timesteps2[step] |
| x_cur = x_next |
| |
| d_cur = self.dx_dt_for_blackbox_solvers(x_cur, t_cur1, t_cur2) |
| x_next = x_cur + (t_next1 - t_cur1) * d_cur |
| return x_next |
|
|
