| import torch |
| from torch.utils.data import DataLoader |
| import numpy as np |
| import imageio |
| import os |
| import shutil |
| from helpers.utils import is_main_process, get_rank, get_world_size, safe_barrier |
|
|
|
|
| def delete_content_of_dir(folder): |
| os.makedirs(folder, exist_ok=True) |
| for filename in os.listdir(folder): |
| file_path = os.path.join(folder, filename) |
| try: |
| if os.path.isfile(file_path) or os.path.islink(file_path): |
| os.unlink(file_path) |
| elif os.path.isdir(file_path): |
| shutil.rmtree(file_path) |
| except Exception as e: |
| print('Failed to delete %s. Reason: %s' % (file_path, e)) |
|
|
| def get_sample_for_visualization(data, preprocess_fn, num, dataset): |
| for x in DataLoader(data, batch_size=num): |
| break |
| orig_image = (x[0] * 255.0).to(torch.uint8).permute(0, 2, 3, 1) if dataset == 'ffhq_1024' else x[0] |
| preprocessed = preprocess_fn(x)[0] |
| return orig_image, preprocessed |
|
|
|
|
| def _resolve_snoise(sampler, batch_size, use_snoise): |
| if use_snoise and hasattr(sampler, 'snoise_tmp'): |
| return [s[:batch_size].normal_() for s in sampler.snoise_tmp] |
| if hasattr(sampler, 'neutral_snoise'): |
| return [s[:batch_size] for s in sampler.neutral_snoise] |
| return None |
|
|
|
|
|
|
| def generate_for_NN(sampler, orig, initial, *args): |
| if len(args) == 4: |
| snoise = None |
| shape, ema_imle, fname, logprint = args |
| elif len(args) == 5: |
| snoise, shape, ema_imle, fname, logprint = args |
| else: |
| raise TypeError("generate_for_NN expected 4 or 5 trailing args") |
|
|
| mb = shape[0] |
| initial = initial[:mb].cuda() |
| nns = sampler.sample(initial, ema_imle, snoise) |
| batches = [orig[:mb], nns] |
| n_rows = len(batches) |
| im = np.concatenate(batches, axis=0).reshape((n_rows, mb, *shape[1:])).transpose([0, 2, 1, 3, 4]).reshape( |
| [n_rows * shape[1], mb * shape[2], 3]) |
|
|
| logprint(f'printing samples to {fname}') |
| imageio.imwrite(fname, im) |
|
|
|
|
| def generate_images_initial(H, sampler, orig, initial, snoise, shape, imle, ema_imle, fname, logprint, experiment=None): |
| mb = shape[0] |
| initial = initial[:mb] |
| batches = [orig[:mb], sampler.sample(initial, imle, snoise)] |
|
|
| temp_latent_rnds = torch.randn([mb, H.latent_dim], dtype=torch.float32).cuda() |
| for t in range(H.num_rows_visualize + 4): |
| temp_latent_rnds.normal_() |
| tmp_snoise = _resolve_snoise(sampler, mb, H.use_snoise) |
| batches.append(sampler.sample(temp_latent_rnds, imle, tmp_snoise)) |
|
|
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
|
|
| |
| |
|
|
| |
| |
| |
|
|
| n_rows = len(batches) |
| im = np.concatenate(batches, axis=0).reshape((n_rows, mb, *shape[1:])).transpose([0, 2, 1, 3, 4]).reshape( |
| [n_rows * shape[1], mb * shape[2], 3]) |
|
|
| logprint(f'printing samples to {fname}') |
| imageio.imwrite(fname, im) |
| if(experiment): |
| experiment.log_image(fname, overwrite=True) |
|
|
|
|
| def generate_visualization(H, sampler, orig, initial, last_initial, latent_for_visualization, |
| shape, imle, fname, logprint, experiment=None): |
| _ = last_initial |
| _ = latent_for_visualization |
| generate_images_initial(H, sampler, orig, initial, None, shape, imle, imle, fname, logprint, experiment) |
|
|
| def generate_and_save(H, imle, sampler, n_samp, subdir='fid'): |
| rank = get_rank() |
| world_size = get_world_size() |
|
|
| if is_main_process(): |
| delete_content_of_dir(f'{H.save_dir}/{subdir}') |
| safe_barrier() |
|
|
| per_rank = (n_samp + world_size - 1) // world_size |
| start_idx = rank * per_rank |
| end_idx = min(start_idx + per_rank, n_samp) |
| local_n = end_idx - start_idx |
|
|
| |
| latent_std = float(getattr(H, 'eval_latent_std', 1.0) or 1.0) |
|
|
| with torch.no_grad(): |
| temp_latent_rnds = torch.randn([H.imle_batch, H.latent_dim], dtype=torch.float32).cuda() |
| generated = 0 |
| while generated < local_n: |
| batch_size = min(H.imle_batch, local_n - generated) |
|
|
| temp_latent_rnds.normal_(mean=0.0, std=latent_std) |
| tmp_snoise = _resolve_snoise(sampler, H.imle_batch, H.use_snoise) |
| samp = sampler.sample(temp_latent_rnds, imle, tmp_snoise) |
|
|
| for j in range(batch_size): |
| imageio.imwrite(f'{H.save_dir}/{subdir}/{start_idx + generated + j}.png', samp[j]) |
| generated += batch_size |
| safe_barrier() |
|
|
