| 19 | #samples using classifier free guidance and only enables the cfg_val when the sigma is within the interval. |
| 20 | #idead from this paper: https://arxiv.org/pdf/2404.07724 |
| 21 | def sample_images_cfg(nr_images, cfg_val, cfg_interval, model_ema, model_config, nr_iters=100, extra_args={}, callback=None): |
| 22 | model_ema.eval() |
| 23 | sigma_min = model_config['sigma_min'] |
| 24 | sigma_max = model_config['sigma_max'] |
| 25 | size = model_config['input_size'] |
| 26 | n_per_proc = nr_images |
| 27 | x = torch.randn([1, n_per_proc, model_config['input_channels'], size[0], size[1]]).cuda() |
| 28 | x = x[0] * sigma_max |
| 29 | model_fn = model_ema |
| 30 | sigmas = K.sampling.get_sigmas_karras(nr_iters, sigma_min, sigma_max, rho=7., device="cuda") |
| 31 | x_0 = K.sampling.sample_dpmpp_2m_sde_cfg(model_fn, x, sigmas, cfg_val, cfg_interval, extra_args=extra_args, eta=0.0, solver_type='heun', disable=False, callback=callback) |
| 32 | return x_0 |