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))