(nr_images)
| 303 | @torch.no_grad() |
| 304 | @K.utils.eval_mode(model_ema) |
| 305 | def sample_images(nr_images): |
| 306 | if accelerator.is_main_process: |
| 307 | tqdm.write('Sampling...') |
| 308 | n_per_proc = math.ceil(nr_images / accelerator.num_processes) |
| 309 | x = torch.randn([accelerator.num_processes, n_per_proc, model_config['input_channels'], size[0], size[1]], generator=demo_gen).to(device) |
| 310 | dist.broadcast(x, 0) |
| 311 | x = x[accelerator.process_index] * sigma_max |
| 312 | model_fn, extra_args = model_ema, {} |
| 313 | #Not really relevent for our case currently since we don't have classes |
| 314 | if num_classes: |
| 315 | class_cond = torch.randint(0, num_classes, [accelerator.num_processes, n_per_proc], generator=demo_gen).to(device) |
| 316 | dist.broadcast(class_cond, 0) |
| 317 | extra_args['class_cond'] = class_cond[accelerator.process_index] |
| 318 | model_fn = make_cfg_model_fn(model_ema) |
| 319 | sigmas = K.sampling.get_sigmas_karras(100, sigma_min, sigma_max, rho=7., device=device) |
| 320 | x_0 = K.sampling.sample_dpmpp_2m_sde(model_fn, x, sigmas, extra_args=extra_args, eta=0.0, solver_type='heun', disable=not accelerator.is_main_process) |
| 321 | x_0 = accelerator.gather(x_0)[:nr_images] |
| 322 | return x_0 |
| 323 | |
| 324 | |
| 325 |
no test coverage detected