File size: 5,179 Bytes
3ce19a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | 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))
# if(H.use_snoise == True):
# tmp_snoise = [s[:mb].normal_() for s in sampler.snoise_tmp]
# else:
# tmp_snoise = [s[:mb] for s in sampler.neutral_snoise]
# batches.append(sampler.sample(temp_latent_rnds, imle, tmp_snoise))
# if(H.use_snoise == True):
# tmp_snoise = [s[:mb].normal_() for s in sampler.snoise_tmp]
# else:
# tmp_snoise = [s[:mb] for s in sampler.neutral_snoise]
# batches.append(sampler.sample(temp_latent_rnds, imle, tmp_snoise))
# tmp_snoise = [s[:mb] for s in sampler.neutral_snoise]
# batches.append(sampler.sample(temp_latent_rnds, imle, tmp_snoise))
# tmp_snoise = [s[:mb] for s in sampler.neutral_snoise]
# temp_latent_rnds.normal_()
# 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
# Inference-time latent-noise amplitude (1.0 = standard N(0, I)).
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()
|