Args: - sampling_method: type of sampler used in solving the SDE; default to be Euler-Maruyama - diffusion_form: function form of diffusion coefficient; default to be matching SBDM - diffusion_norm: function magnitude of diffusion coefficient; default to 1 -
(
self,
init,
model,
sampling_method="euler",
diffusion_form="sigma",
diffusion_norm=1.0,
last_step="Mean",
last_step_size=0.04,
num_steps=250,
progress=True,
return_intermediates=False,
cfg_scale=1.0,
uc_cond=None,
cond_key="y",
**model_kwargs
)
| 301 | raise NotImplementedError(f"Last step '{last_step}' not implemented.") |
| 302 | |
| 303 | def sample( |
| 304 | self, |
| 305 | init, |
| 306 | model, |
| 307 | sampling_method="euler", |
| 308 | diffusion_form="sigma", |
| 309 | diffusion_norm=1.0, |
| 310 | last_step="Mean", |
| 311 | last_step_size=0.04, |
| 312 | num_steps=250, |
| 313 | progress=True, |
| 314 | return_intermediates=False, |
| 315 | cfg_scale=1.0, |
| 316 | uc_cond=None, |
| 317 | cond_key="y", |
| 318 | **model_kwargs |
| 319 | ): |
| 320 | """ |
| 321 | Args: |
| 322 | - sampling_method: type of sampler used in solving the SDE; default to be Euler-Maruyama |
| 323 | - diffusion_form: function form of diffusion coefficient; default to be matching SBDM |
| 324 | - diffusion_norm: function magnitude of diffusion coefficient; default to 1 |
| 325 | - last_step: type of the last step; default to identity |
| 326 | - last_step_size: size of the last step; default to match the stride of 250 steps over [0,1] |
| 327 | - num_steps: total integration step of SDE |
| 328 | """ |
| 329 | if last_step is None: |
| 330 | last_step_size = 0.0 |
| 331 | |
| 332 | sde_drift, sde_diffusion = self.__get_sde_diffusion_and_drift( |
| 333 | diffusion_form=diffusion_form, diffusion_norm=diffusion_norm, |
| 334 | ) |
| 335 | |
| 336 | t0, t1 = self.check_interval(diffusion_form=diffusion_form, reverse=False, last_step_size=last_step_size) |
| 337 | ts = torch.linspace(t0, t1, num_steps).to(init.device) |
| 338 | dt = ts[1] - ts[0] |
| 339 | |
| 340 | # enable classifier-free guidance |
| 341 | model_forward_fn = partial(forward_with_cfg, model=model, cfg_scale=cfg_scale, uc_cond=uc_cond, cond_key=cond_key) |
| 342 | |
| 343 | """ forward loop of sde """ |
| 344 | sampler = StepSDE(dt=dt, drift=sde_drift, diffusion=sde_diffusion, sampler_type=sampling_method) |
| 345 | |
| 346 | # sample |
| 347 | x = init |
| 348 | mean_x = init |
| 349 | xs = [] |
| 350 | for ti in tqdm(ts[:-1], disable=not progress, desc="SDE sampling", total=num_steps, initial=1): |
| 351 | with torch.no_grad(): |
| 352 | x, mean_x = sampler(x, mean_x, ti, model_forward_fn, **model_kwargs) |
| 353 | xs.append(x) |
| 354 | |
| 355 | # make last step |
| 356 | t_last = torch.ones(x.size(0), device=x.device) * t1 |
| 357 | x = self.last_step( |
| 358 | x=xs[-1], t=t_last, |
| 359 | model=model_forward_fn, |
| 360 | sde_drift=sde_drift, |
no test coverage detected