def project_images(G, images, name_prefix, args):
device = torch.device(args.gpu[0] if args.gpu else 'cpu')
if device.index is not None:
torch.cuda.set_device(device.index)
if len(args.gpu) > 1:
warnings.warn(
'Multi GPU is not available for projection. ' + \
'Using device {}'.format(device)
)
G = utils.unwrap_module(G).to(device)
lpips_model = stylegan2.external_models.lpips.LPIPS_VGG16(
pixel_min=args.pixel_min, pixel_max=args.pixel_max)
proj = stylegan2.project.Projector(G=G,
dlatent_avg_samples=10000,
dlatent_avg_label=args.label,
dlatent_device=device,
dlatent_batch_size=1024,
lpips_model=lpips_model,
lpips_size=256)
for i in range(0, len(images), args.batch_size):
target = images[i:i + args.batch_size]
proj.start(target=target,
num_steps=args.num_steps + 1,
initial_learning_rate=args.initial_learning_rate,
initial_noise_factor=args.initial_noise_factor,
lr_rampdown_length=args.lr_rampdown_length,
lr_rampup_length=args.lr_rampup_length,
noise_ramp_length=args.noise_ramp_length,
regularize_noise_weight=args.regularize_noise_weight,
verbose=True,
verbose_prefix='Projecting image(s) {}/{}'.format(
i * args.batch_size + len(target), len(images)))
snapshot_steps = set(args.num_steps - np.linspace(
0, args.num_steps, args.num_snapshots, endpoint=False, dtype=int))
for k, image in enumerate(
utils.tensor_to_PIL(target,
pixel_min=args.pixel_min,
pixel_max=args.pixel_max)):
image.save(
os.path.join(args.output, name_prefix[i + k] + 'target.png'))
for j in range(args.num_steps):
proj.step()
if j in snapshot_steps:
generated = utils.tensor_to_PIL(proj.generate(),
pixel_min=args.pixel_min,
pixel_max=args.pixel_max)
for k, image in enumerate(generated):
torch.save(
proj.get_dlatent(),
os.path.join(args.output,
name_prefix[i + k] + 'step%04d.pt' % j))
image.save(
os.path.join(args.output,
name_prefix[i + k] + 'step%04d.png' % j))
def convert(model, inputs):
truncation = inputs['truncation']
Gs.set_truncation(truncation_psi=truncation)
qlatents = torch.Tensor(inputs['z']).reshape(1, 512).to(device=device, dtype=torch.float32)
dlatents = Gs.G_mapping(qlatents)
swifted_dlatents = shift_latents(dlatents, inputs)
generated = Gs(dlatents=swifted_dlatents)
images = utils.tensor_to_PIL(generated)
return {'image': images[0]}
def synthesis(G_file, latent_file):
device = torch.device('cpu')
G = stylegan2.models.load(G_file).G_synthesis
latent = np.load(latent_file, allow_pickle=True)
G.to(device)
latent = torch.tensor(latent[np.newaxis, ...]).to(device)
out = G(latent)
out = utils.tensor_to_PIL(out, pixel_min=-1, pixel_max=1)[0]
return out
def generate_images(G, args):
latent_size, label_size = G.latent_size, G.label_size
device = torch.device('cpu')
G.to(device)
if args['truncation_psi'] != 1.0:
G.set_truncation(truncation_psi=args['truncation_psi'])
noise_reference = G.static_noise()
def get_batch(seeds):
latents = []
labels = []
noise_tensors = [[] for _ in noise_reference]
for seed in seeds:
rnd = np.random.RandomState(seed)
latents.append(torch.from_numpy(rnd.randn(latent_size)))
for i, ref in enumerate(noise_reference):
noise_tensors[i].append(torch.from_numpy(rnd.randn(*ref.size()[1:])))
if label_size:
labels.append(torch.tensor([rnd.randint(0, label_size)]))
latents = torch.stack(latents, dim=0).to(device=device, dtype=torch.float32)
if labels:
labels = torch.cat(labels, dim=0).to(device=device, dtype=torch.int64)
else:
labels = None
noise_tensors = [
torch.stack(noise, dim=0).to(device=device, dtype=torch.float32)
for noise in noise_tensors
]
return latents, labels, noise_tensors
for i in range(0, len(args['seed'])):
latents, labels, noise_tensors = get_batch(args['seed'][i: i + 1])
if noise_tensors is not None:
G.static_noise(noise_tensors=noise_tensors)
with torch.no_grad():
generated = G(latents, labels=labels)
images = utils.tensor_to_PIL(
generated, pixel_min=-1.0, pixel_max=1.0)
for seed, img in zip(args['seed'][i: i + 1], images):
pass
return img
def generate_image(G, weights):
latent_size, label_size = G.latent_size, G.label_size
device = torch.device('cpu')
if device.index is not None:
torch.cuda.set_device(device.index)
G.to(device)
G.set_truncation(0.5)
noise_reference = G.static_noise()
noise_tensors = [[] for _ in noise_reference]
latents = []
labels = []
rnd = np.random.RandomState(6600)
latents.append(torch.from_numpy(weights))
for i, ref in enumerate(noise_reference):
noise_tensors[i].append(torch.from_numpy(rnd.randn(*ref.size()[1:])))
if label_size:
labels.append(torch.tensor([rnd.randint(0, label_size)]))
latents = torch.stack(latents, dim=0).to(device=device,
dtype=torch.float32)
if labels:
labels = torch.cat(labels, dim=0).to(device=device, dtype=torch.int64)
else:
labels = None
noise_tensors = [
torch.stack(noise, dim=0).to(device=device, dtype=torch.float32)
for noise in noise_tensors
]
if noise_tensors is not None:
G.static_noise(noise_tensors=noise_tensors)
with torch.no_grad():
generated = G(latents, labels=labels)
images = utils.tensor_to_PIL(generated, pixel_min=-1, pixel_max=1)
for img in images:
return img
def generate_images(type_, G, args):
device = set_device(G, args)
# Set truncation_psi
if args.truncation_psi != 1:
G.set_truncation(truncation_psi=args.truncation_psi)
# Set random noise (instead of noise generated from seed)
G.random_noise()
# Get labels (deactivated for the moment)
labels = None
#if G.label_size:
# labels = [torch.tensor([np.random.randint(0, G.label_size)]) for _ in range(nb_images)]
# labels = torch.cat(labels, dim=0).to(device=device, dtype=torch.int64)
# Get latent factors and names of images
latents, names = build_latents_names(type_, G, args)
latents.to(device=device, dtype=torch.float32)
# Generate images per batch
for i in tqdm(range(0, latents.shape[0], args.batch_size)):
batch_latents = latents[i:i + args.batch_size]
batch_names = names[i:i + args.batch_size]
with torch.no_grad():
if args.seeds:
# In this case the latents factors must be fed first into the mapping network
generated = G(batch_latents, labels=labels)
if args.latents:
# In this case we assume that the latent factors are output of mapping network (because that's the output of the projection)
# Labels are not handled here
generated = G.G_synthesis(latents=batch_latents)
images = utils.tensor_to_PIL(generated,
pixel_min=args.pixel_min,
pixel_max=args.pixel_max)
for name, img in zip(batch_names, images):
img.save(os.path.join(args.output, name + '.png'))
def generate_images(G, args):
latent_size, label_size = G.latent_size, G.label_size
device = torch.device(args.gpu[0] if args.gpu else 'cpu')
if device.index is not None:
torch.cuda.set_device(device.index)
G.to(device)
if args.truncation_psi != 1:
G.set_truncation(truncation_psi=args.truncation_psi)
if len(args.gpu) > 1:
warnings.warn(
'Noise can not be randomized based on the seed ' + \
'when using more than 1 GPU device. Noise will ' + \
'now be randomized from default random state.'
)
G.random_noise()
G = torch.nn.DataParallel(G, device_ids=args.gpu)
else:
noise_reference = G.static_noise()
def get_batch(seeds):
latents = []
labels = []
if len(args.gpu) <= 1:
noise_tensors = [[] for _ in noise_reference]
for seed in seeds:
rnd = np.random.RandomState(seed)
latents.append(torch.from_numpy(rnd.randn(latent_size)))
if len(args.gpu) <= 1:
for i, ref in enumerate(noise_reference):
noise_tensors[i].append(
torch.from_numpy(rnd.randn(*ref.size()[1:])))
if label_size:
labels.append(torch.tensor([rnd.randint(0, label_size)]))
latents = torch.stack(latents, dim=0).to(device=device,
dtype=torch.float32)
if labels:
labels = torch.cat(labels, dim=0).to(device=device,
dtype=torch.int64)
else:
labels = None
if len(args.gpu) <= 1:
noise_tensors = [
torch.stack(noise, dim=0).to(device=device,
dtype=torch.float32)
for noise in noise_tensors
]
else:
noise_tensors = None
return latents, labels, noise_tensors
progress = utils.ProgressWriter(len(args.seeds))
progress.write('Generating images...', step=False)
for i in range(0, len(args.seeds), args.batch_size):
latents, labels, noise_tensors = get_batch(args.seeds[i:i +
args.batch_size])
if noise_tensors is not None:
G.static_noise(noise_tensors=noise_tensors)
with torch.no_grad():
generated = G(latents, labels=labels)
images = utils.tensor_to_PIL(generated,
pixel_min=args.pixel_min,
pixel_max=args.pixel_max)
for seed, img in zip(args.seeds[i:i + args.batch_size], images):
img.save(os.path.join(args.output, 'seed%04d.png' % seed))
progress.step()
progress.write('Done!', step=False)
progress.close()
def style_mixing_example(G, args):
assert max(args.style_layers) < len(G), \
'Style layer indices can not be larger than ' + \
'number of style layers ({}) of the generator.'.format(len(G))
device = torch.device(args.gpu[0] if args.gpu else 'cpu')
if device.index is not None:
torch.cuda.set_device(device.index)
if len(args.gpu) > 1:
warnings.warn(
'Multi GPU is not available for style mixing example. Using device {}'
.format(device))
G.to(device)
G.static_noise()
latent_size, label_size = G.latent_size, G.label_size
G_mapping, G_synthesis = G.G_mapping, G.G_synthesis
all_seeds = list(set(args.row_seeds + args.col_seeds))
all_z = torch.stack([
torch.from_numpy(np.random.RandomState(seed).randn(latent_size))
for seed in all_seeds
])
all_z = all_z.to(device=device, dtype=torch.float32)
if label_size:
labels = torch.zeros(len(all_z), dtype=torch.int64, device=device)
else:
labels = None
print('Generating disentangled latents...')
#print(all_z.shape) #[-1,512]
with torch.no_grad():
all_w = G_mapping(latents=all_z, labels=labels)
all_w = all_w.unsqueeze(1).repeat(1, len(G_synthesis), 1) #[-1,18,512]
w_avg = G.dlatent_avg # [512]
if args.truncation_psi != 1:
all_w = w_avg + args.truncation_psi * (all_w - w_avg)
w_dict = {seed: w for seed, w in zip(all_seeds, all_w)}
all_images = []
progress = utils.ProgressWriter(len(all_w))
progress.write('Generating images...', step=False)
with torch.no_grad():
for w in all_w:
all_images.append(G_synthesis(w.unsqueeze(0)))
progress.step()
progress.write('Done!', step=False)
progress.close()
all_images = torch.cat(all_images, dim=0)
image_dict = {(seed, seed): image
for seed, image in zip(all_seeds, all_images)}
progress = utils.ProgressWriter(len(args.row_seeds) * len(args.col_seeds))
progress.write('Generating style-mixed images...', step=False)
for row_seed in args.row_seeds:
for col_seed in args.col_seeds:
w = w_dict[row_seed].clone()
w[args.style_layers] = w_dict[col_seed][args.style_layers]
with torch.no_grad():
image_dict[(row_seed,
col_seed)] = G_synthesis(w.unsqueeze(0)).squeeze(0)
progress.step()
progress.write('Done!', step=False)
progress.close()
progress = utils.ProgressWriter(len(image_dict))
progress.write('Saving images...', step=False)
for (row_seed, col_seed), image in list(image_dict.items()):
image = utils.tensor_to_PIL(image,
pixel_min=args.pixel_min,
pixel_max=args.pixel_max)
image_dict[(row_seed, col_seed)] = image
image.save(
os.path.join(args.output, '%d-%d.png' % (row_seed, col_seed)))
progress.step()
progress.write('Done!', step=False)
progress.close()
if args.grid:
print('\n\nSaving style-mixed grid...')
H, W = all_images.size()[2:]
canvas = Image.new('RGB', (W * (len(args.col_seeds) + 1), H *
(len(args.row_seeds) + 1)), 'black')
for row_idx, row_seed in enumerate([None] + args.row_seeds):
for col_idx, col_seed in enumerate([None] + args.col_seeds):
if row_seed is None and col_seed is None:
continue
key = (row_seed, col_seed)
if row_seed is None:
key = (col_seed, col_seed)
if col_seed is None:
key = (row_seed, row_seed)
canvas.paste(image_dict[key], (W * col_idx, H * row_idx))
canvas.save(os.path.join(args.output, 'grid.png'))
print('Done!')
def generate_images(G, args):
args.seeds = [int('0x%s' % args.tokenid, 16) % ((1 << 31) - 1)]
print('generating ', args.tokenid, args.seeds[0])
latent_size, label_size = G.latent_size, G.label_size
device = torch.device(args.gpu[0] if args.gpu else 'cpu')
if device.index is not None:
torch.cuda.set_device(device.index)
G.to(device)
if args.truncation_psi != 1:
G.set_truncation(truncation_psi=args.truncation_psi)
if len(args.gpu) > 1:
warnings.warn(
'Noise can not be randomized based on the seed ' + \
'when using more than 1 GPU device. Noise will ' + \
'now be randomized from default random state.'
)
G.random_noise()
G = torch.nn.DataParallel(G, device_ids=args.gpu)
else:
noise_reference = G.static_noise()
def get_batch(seeds):
latents = []
labels = []
if len(args.gpu) <= 1:
noise_tensors = [[] for _ in noise_reference]
for seed in seeds:
rnd = np.random.RandomState(seed)
latents.append(torch.from_numpy(rnd.randn(latent_size)))
if len(args.gpu) <= 1:
for i, ref in enumerate(noise_reference):
noise_tensors[i].append(
torch.from_numpy(rnd.randn(*ref.size()[1:])))
if label_size:
labels.append(torch.tensor([rnd.randint(0, label_size)]))
latents = torch.stack(latents, dim=0).to(device=device,
dtype=torch.float32)
if labels:
labels = torch.cat(labels, dim=0).to(device=device,
dtype=torch.int64)
else:
labels = None
if len(args.gpu) <= 1:
noise_tensors = [
torch.stack(noise, dim=0).to(device=device,
dtype=torch.float32)
for noise in noise_tensors
]
else:
noise_tensors = None
return latents, labels, noise_tensors
for i in range(0, len(args.seeds), args.batch_size):
latents, labels, noise_tensors = get_batch(args.seeds[i:i +
args.batch_size])
if noise_tensors is not None:
G.static_noise(noise_tensors=noise_tensors)
with torch.no_grad():
generated = G(latents, labels=labels)
images = utils.tensor_to_PIL(generated,
pixel_min=args.pixel_min,
pixel_max=args.pixel_max)
for seed, img in zip(args.seeds[i:i + args.batch_size], images):
path = os.path.join(args.output, 'original',
'%s.png' % args.tokenid)
img.save(path)
print('saved', path)
for d in [32, 64, 128]:
rimg = img.resize((d, d))
path = os.path.join(args.output, str(d),
'%s.png' % args.tokenid)
rimg.save(path)
print('saved', path)
def interpolate(G, args):
latent_size, label_size = G.latent_size, G.label_size
device = torch.device(args.gpu[0] if args.gpu else 'cpu')
if device.index is not None:
torch.cuda.set_device(device.index)
G.to(device)
if args.truncation_psi != 1:
G.set_truncation(truncation_psi=args.truncation_psi)
if len(args.gpu) > 1:
warnings.warn(
'Noise can not be randomized based on the seed ' + \
'when using more than 1 GPU device. Noise will ' + \
'now be randomized from default random state.'
)
G.random_noise()
G = torch.nn.DataParallel(G, device_ids=args.gpu)
else:
noise_reference = G.static_noise()
noise_tensors = None
if noise_tensors is not None:
G.static_noise(noise_tensors=noise_tensors)
def gen_latent(seed):
return torch.from_numpy(np.random.RandomState(seed).randn(latent_size))
def interpolate_generator(seed, step):
if len(args.gpu) <= 1:
noise_tensors = [[] for _ in noise_reference]
for i, ref in enumerate(noise_reference):
noise_tensors[i].append(
torch.from_numpy(
np.random.RandomState(seed).randn(*ref.size()[1:])))
noise_tensors = [
torch.stack(noise, dim=0).to(device=device,
dtype=torch.float32)
for noise in noise_tensors
]
else:
noise_tensors = None
latent1 = gen_latent(seed)
latent2 = gen_latent(seed + 1)
d_latents = (latent2 - latent1) / float((step - 1))
for i in range(step):
yield latent1 + i * d_latents, noise_tensors
progress = utils.ProgressWriter(len(args.seeds) * args.interpolation_step)
progress.write('Generating images...', step=False)
if args.interpolation_step:
fourcc_ = cv2.VideoWriter_fourcc(*'avc1')
video_writer = cv2.VideoWriter(filename=os.path.join(
args.output, args.animation_filename),
fourcc=fourcc_,
fps=args.animation_fps,
apiPreference=cv2.CAP_ANY,
frameSize=args.animation_frame_size)
for seed in args.seeds:
for i, (latent, noise_tensors) in enumerate(
interpolate_generator(seed, args.interpolation_step)):
latents = torch.stack([latent], dim=0).to(device=device,
dtype=torch.float32)
if noise_tensors is not None:
G.static_noise(noise_tensors=noise_tensors)
with torch.no_grad():
generated = G(latents, labels=None)
images = utils.tensor_to_PIL(generated,
pixel_min=args.pixel_min,
pixel_max=args.pixel_max)
for img in images: # args.seeds[i: i + args.batch_size]
img.save(os.path.join(args.output, f'{seed}_{i}.png'))
if args.interpolation_step:
img = np.array(img)
# Convert RGB to BGR
img = img[:, :, ::-1].copy()
video_writer.write(img)
progress.step()
progress.write('Done!', step=False)
progress.close()
def project_images(G, images, name_prefix, args):
device = torch.device(args.gpu[0] if args.gpu else 'cpu')
if device.index is not None:
torch.cuda.set_device(device.index)
if len(args.gpu) > 1:
warnings.warn(
'Multi GPU is not available for projection. ' + \
'Using device {}'.format(device)
)
G = utils.unwrap_module(G).to(device)
lpips_model = stylegan2.external_models.lpips.LPIPS_VGG16(
pixel_min=args.pixel_min, pixel_max=args.pixel_max)
proj = stylegan2.project.Projector(
G=G,
dlatent_avg_samples=10000,
dlatent_avg_label=args.label,
dlatent_device=device,
dlatent_batch_size=1024,
lpips_model=lpips_model,
lpips_size=256
)
for i in range(0, len(images), args.batch_size):
target = images[i: i + args.batch_size]
proj.start(
target=target,
num_steps=args.num_steps,
initial_learning_rate=args.initial_learning_rate,
initial_noise_factor=args.initial_noise_factor,
lr_rampdown_length=args.lr_rampdown_length,
lr_rampup_length=args.lr_rampup_length,
noise_ramp_length=args.noise_ramp_length,
regularize_noise_weight=args.regularize_noise_weight,
verbose=True,
verbose_prefix='Projecting image(s) {}/{}'.format(
i * args.batch_size + len(target), len(images)),
noise_layers=5
)
# snapshot_steps = set(
# args.num_steps - np.linspace(
# 0, args.num_steps, args.num_snapshots, endpoint=False, dtype=int))
for k, image in enumerate(utils.tensor_to_PIL(target, pixel_min=args.pixel_min, pixel_max=args.pixel_max)):
os.makedirs(os.path.join(args.output, 'target+BEST'), exist_ok=True)
image.save(os.path.join(args.output, 'target+BEST', name_prefix[i + k] + 'target.png'))
for j in range(args.num_steps):
current_image, loss_dict_step, best_output = proj.step()
# if j in snapshot_steps:
# generated = utils.tensor_to_PIL(
# proj.generate(), pixel_min=args.pixel_min, pixel_max=args.pixel_max)
# for k, image in enumerate(generated):
# image.save(os.path.join(
# args.output, name_prefix[i + k] + 'step%04d.png' % (j + 1)))
current_image = utils.tensor_to_PIL(current_image, pixel_min=args.pixel_min, pixel_max=args.pixel_max)
best_output = utils.tensor_to_PIL(best_output, pixel_min=args.pixel_min, pixel_max=args.pixel_min)
for j % SAVE_PER == 0 and j != 0 or j == args.num_steps -1:
for k, (image, best_image) in enumerate(zip(current_image, best_output)):
L2, GEOCROSS = loss_dict_stepp['L2'], loss_dict_stepp['GEOCROSS']
save_name = f'{name_prefix[i + k]}-loss-{L2:.2f}+GEOCROSS-{GEOCROSS:.2f}-{j}.png'
if j != args.num_steps - 1:
image.save(os.path.join(args.output, save_name))
else:
save_name = f'{name_prefix[i + k]}-BEST-{j}.png'
best_image.save(os.path.join(args.output, 'target+BEST', save_name))
