def main(): os.makedirs(a.out_dir, exist_ok=True) tflib.init_tf() with tf.Session() as sess, tf.device('/gpu:0'): Net_lo = load_pkl(a.pkl1) Net_hi = load_pkl(a.pkl2) try: # full model G_lo, D_lo, Gs_lo = Net_lo G_hi, D_hi, Gs_hi = Net_hi G_out = blend_layers(G_lo, G_hi) Gs_out = blend_layers(Gs_lo, Gs_hi) D_out = blend_layers(D_lo, D_hi, type='D') Net_out = G_out, D_out, Gs_out except: # Gs only Gs_out = blend_layers(Net_lo, Net_hi) # only Gs Net_out = Gs_out out_name = os.path.join( a.out_dir, '%s-%s-%d-%d' % (basename(a.pkl1).split('-')[0], basename( a.pkl2).split('-')[0], a.res, a.level)) save_pkl(Net_out, '%s.pkl' % out_name) rnd = np.random.RandomState(696) grid_latents = rnd.randn(4, *Gs_out.input_shape[1:]) grid_fakes = Gs_out.run(grid_latents, [None], is_validation=True, minibatch_size=1) grid_fakes = np.hstack(np.transpose(grid_fakes, [0, 2, 3, 1])) imsave('%s.jpg' % out_name, ((grid_fakes + 1) * 127.5).astype(np.uint8)) print('\n All done')
def locate_latest_pkl(train_dir): allpickles = sorted([f for f in file_list(train_dir, 'pkl') if 'snapshot' in f]) if len(allpickles) == 0: return None, 0.0 latest_pickle = allpickles[-1] kimg = float(basename(latest_pickle).split('-')[-1]) return latest_pickle, kimg
def main(): tflib.init_tf({'allow_soft_placement':True}) print(' Loading model', a.source) G_in, D_in, Gs_in = load_pkl(a.source) if a.res is not None: print(' Reconstructing model with size', a.res) data_shape = [Gs_in.output_shape[1], *a.res] G_out, D_out, Gs_out, res_out_log2 = create_model(data_shape, True, Gs_in.static_kwargs) if a.res[0] == a.res[1]: assert G_in is not None and D_in is not None, " !! G/D subnets not found in source model !!" res_in_log2 = np.log2(get_model_res(Gs_in)) lod_diff = res_out_log2 - res_in_log2 Gs_in_names_to_copy = update_dict_keys(Gs_in, 'ToRGB_lod', 'generator', lod_diff) G_in_names_to_copy = update_dict_keys(G_in, 'ToRGB_lod', 'generator', lod_diff) D_in_names_to_copy = update_dict_keys(D_in, 'FromRGB_lod', 'discriminator', lod_diff) copy_weights(Gs_in, Gs_out, Gs_in_names_to_copy) copy_weights(G_in, G_out, G_in_names_to_copy) copy_weights(D_in, D_out, D_in_names_to_copy) else: # EXPERIMENTAL .. check source repo if G_in is not None and D_in is not None: copy_and_crop_trainables(G_in, G_out) copy_and_crop_trainables(D_in, D_out) copy_and_crop_trainables(Gs_in, Gs_out) elif a.reconstruct is True: print(' Reconstructing model with same size') data_shape = Gs_in.output_shape[1:] _, _, Gs_out, _ = create_model(data_shape, False, Gs_in.static_kwargs) Gs_out.copy_vars_from(Gs_in) else: Gs_out = Gs_in if a.res is not None: if G_in is not None and D_in is not None: save_pkl((G_out, D_out, Gs_out), os.path.join(a.out_dir, '%s-%dx%d.pkl' % (basename(a.source), *a.res[::-1]))) else: save_pkl(Gs_out, os.path.join(a.out_dir, '%s-%dx%d.pkl' % (basename(a.source), *a.res[::-1]))) else: save_pkl(Gs_out, os.path.join(a.out_dir, '%s-Gs.pkl' % basename(a.source))) print(' Done')
def main(): os.makedirs(a.out_dir, exist_ok=True) tflib.init_tf() with tf.Session() as sess, tf.device('/gpu:0'): Gs_lo = load_pkl(a.pkl1) Gs_hi = load_pkl(a.pkl2) # TODO add small x offset for smoother blend animations resolution = "{}x{}".format(a.res, a.res) model_1_names = extract_conv_names(Gs_lo) model_2_names = extract_conv_names(Gs_hi) assert all((x == y for x, y in zip(model_1_names, model_2_names))) Gs_out = Gs_lo.clone() short_names = [(x[1:3]) for x in model_1_names] full_names = [(x[0]) for x in model_1_names] mid_point_idx = short_names.index((resolution, a.level)) mid_point_pos = model_1_names[mid_point_idx][3] ys = [] for name, resolution, level, position in model_1_names: x = position - mid_point_pos if a.blend_width is not None: exponent = -x / a.blend_width y = 1 / (1 + math.exp(exponent)) else: y = 1 if x > 1 else 0 ys.append(y) if a.verbose is True: print("Blending {} by {}".format(name, y)) tfutil.set_vars(tfutil.run({ Gs_out.vars[name]: (Gs_hi.vars[name] * y + Gs_lo.vars[name] * (1-y)) for name, y in zip(full_names, ys)} )) out_name = os.path.join(a.out_dir, '%s-%s-%d-%d' % (basename(a.pkl1).split('-')[0], basename(a.pkl2).split('-')[0], a.res, a.level)) save_pkl(Gs_out, '%s.pkl' % out_name) rnd = np.random.RandomState(696) grid_latents = rnd.randn(4, *Gs_lo.input_shape[1:]) grid_fakes = Gs_out.run(grid_latents, [None], is_validation=True, minibatch_size=1) grid_fakes = np.hstack(np.transpose(grid_fakes, [0,2,3,1])) imsave('%s.jpg' % out_name, ((grid_fakes+1)*127.5).astype(np.uint8))
def project_image(proj, targets, work_dir, resolution, num_snapshots): filename = osp.join(work_dir, basename(work_dir)) video_out = cv2.VideoWriter(filename + '.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 25, resolution) snapshot_steps = set(proj.num_steps - np.linspace(0, proj.num_steps, num_snapshots, endpoint=False, dtype=int)) misc.save_image_grid(targets, filename + '.jpg', drange=[-1,1]) proj.start(targets) pbar = ProgressBar(proj.num_steps) while proj.get_cur_step() < proj.num_steps: proj.step() write_video_frame(proj, video_out) if proj.get_cur_step() in snapshot_steps: misc.save_image_grid(proj.get_images(), filename + '-%04d.jpg' % proj.get_cur_step(), drange=[-1,1]) pbar.upd() dlats = proj.get_dlatents() np.save(filename + '-%04d.npy' % proj.get_cur_step(), dlats) video_out.release()
def main(): print('Loading networks from "%s"...' % a.model) sess = tflib.init_tf() with open(a.model, 'rb') as file: network = pickle.load(file, encoding='latin1') try: _, _, Gs = network except: Gs = network resolution = tuple(Gs.output_shape[2:]) proj = projector.Projector(a.steps) proj.set_network(Gs) img_files = img_list(a.in_dir) num_images = len(img_files) for image_idx in range(num_images): print('Projecting image %d/%d ...' % (image_idx+1, num_images)) images = img_read(img_files[image_idx]) images = np.expand_dims(np.transpose(images, [2,0,1]), 0) images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1]) work_dir = osp.join(a.out_dir, basename(img_files[image_idx])) os.makedirs(work_dir, exist_ok=True) project_image(proj, images, work_dir, resolution, a.num_snapshots)
def main(): if a.vector_dir is not None: if a.vector_dir.endswith('/') or a.vector_dir.endswith('\\'): a.vector_dir = a.vector_dir[:-1] os.makedirs(osp.join(a.out_dir, 'ttt'), exist_ok=True) global Gs, use_d # parse filename to model parameters mparams = basename(a.model).split('-') res = int(mparams[1]) cfg = mparams[2] # setup generator Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.func_name = 'training.stylegan2_custom.G_main' Gs_kwargs.verbose = False Gs_kwargs.resolution = res Gs_kwargs.size = a.size Gs_kwargs.scale_type = a.scale_type Gs_kwargs.latent_size = a.latent_size if cfg.lower() == 'f': Gs_kwargs.synthesis_func = 'G_synthesis_stylegan2' elif cfg.lower() == 'e': Gs_kwargs.synthesis_func = 'G_synthesis_stylegan2' Gs_kwargs.fmap_base = 8 << 10 else: print(' old modes [A-D] not implemented'); exit() # check initial model resolution if len(mparams) > 3: if 'x' in mparams[3].lower(): init_res = [int(x) for x in mparams[3].lower().split('x')] Gs_kwargs.init_res = list(reversed(init_res)) # [H,W] !!! custom res # load model, check channels sess = tflib.init_tf({'allow_soft_placement':True}) pkl_name = osp.splitext(a.model)[0] with open(pkl_name + '.pkl', 'rb') as file: network = pickle.load(file, encoding='latin1') try: _, _, Gs = network except: Gs = network Gs_kwargs.num_channels = Gs.output_shape[1] # reload custom network, if needed if '.pkl' in a.model.lower(): print(' .. Gs from pkl ..') else: print(' .. Gs custom ..') Gs = tflib.Network('Gs', **Gs_kwargs) Gs.copy_vars_from(network) # load directions if a.vector_dir is not None: directions = [] vector_list = file_list(a.vector_dir, 'npy') for v in vector_list: direction = load_latents(v) if len(direction.shape) == 2: direction = np.expand_dims(direction, 0) directions.append(direction) directions = np.concatenate(directions)[:, np.newaxis] # [frm,1,18,512] else: print(' No vectors found'); exit() if len(direction[0].shape) > 1 and direction[0].shape[0] > 1: use_d = True print(' directions', directions.shape, 'using d' if use_d else 'using w') # latent direction range lrange = [-0.5, 0.5] # load saved latents if a.npy_file is not None: base_latent = load_latents(a.npy_file) else: print(' No NPY input given, making random') z_dim = Gs.input_shape[1] shape = (1, z_dim) base_latent = np.random.randn(*shape) if use_d: base_latent = Gs.components.mapping.run(base_latent, None) # [frm,18,512] for i, direction in enumerate(directions): make_loop(base_latent, direction, lrange, a.fstep*2, a.fstep*2 * i)
def generate(): os.makedirs(a.out_dir, exist_ok=True) np.random.seed(seed=696) device = torch.device('cuda') # setup generator Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.verbose = a.verbose Gs_kwargs.size = a.size Gs_kwargs.scale_type = a.scale_type # mask/blend latents with external latmask or by splitting the frame if a.latmask is None: nHW = [int(s) for s in a.nXY.split('-')][::-1] assert len(nHW) == 2, ' Wrong count nXY: %d (must be 2)' % len(nHW) n_mult = nHW[0] * nHW[1] if a.verbose is True and n_mult > 1: print(' Latent blending w/split frame %d x %d' % (nHW[1], nHW[0])) lmask = np.tile(np.asarray([[[[1]]]]), (1, n_mult, 1, 1)) Gs_kwargs.countHW = nHW Gs_kwargs.splitfine = a.splitfine else: if a.verbose is True: print(' Latent blending with mask', a.latmask) n_mult = 2 if os.path.isfile(a.latmask): # single file lmask = np.asarray([[img_read(a.latmask)[:, :, 0] / 255.] ]) # [h,w] elif os.path.isdir(a.latmask): # directory with frame sequence lmask = np.asarray([[ img_read(f)[:, :, 0] / 255. for f in img_list(a.latmask) ]]) # [h,w] else: print(' !! Blending mask not found:', a.latmask) exit(1) lmask = np.concatenate((lmask, 1 - lmask), 1) # [frm,2,h,w] lmask = torch.from_numpy(lmask).to(device) # load base or custom network pkl_name = osp.splitext(a.model)[0] if '.pkl' in a.model.lower(): custom = False print(' .. Gs from pkl ..', basename(a.model)) else: custom = True print(' .. Gs custom ..', basename(a.model)) with dnnlib.util.open_url(pkl_name + '.pkl') as f: Gs = legacy.load_network_pkl(f, custom=custom, **Gs_kwargs)['G_ema'].to( device) # type: ignore if a.verbose is True: print(' out shape', Gs.output_shape[1:]) if a.verbose is True: print(' making timeline..') lats = [] # list of [frm,1,512] for i in range(n_mult): lat_tmp = latent_anima((1, Gs.z_dim), a.frames, a.fstep, cubic=a.cubic, gauss=a.gauss, verbose=False) # [frm,1,512] lats.append(lat_tmp) # list of [frm,1,512] latents = np.concatenate(lats, 1) # [frm,X,512] print(' latents', latents.shape) latents = torch.from_numpy(latents).to(device) frame_count = latents.shape[0] # distort image by tweaking initial const layer if a.digress > 0: try: init_res = Gs.init_res except: init_res = (4, 4) # default initial layer size dconst = [] for i in range(n_mult): dc_tmp = a.digress * latent_anima([1, Gs.z_dim, *init_res], a.frames, a.fstep, cubic=True, verbose=False) dconst.append(dc_tmp) dconst = np.concatenate(dconst, 1) else: dconst = np.zeros([frame_count, 1, 1, 1, 1]) dconst = torch.from_numpy(dconst).to(device) # labels / conditions label_size = Gs.c_dim if label_size > 0: labels = torch.zeros((frame_count, n_mult, label_size), device=device) # [frm,X,lbl] if a.labels is None: label_ids = [] for i in range(n_mult): label_ids.append(random.randint(0, label_size - 1)) else: label_ids = [int(x) for x in a.labels.split('-')] label_ids = label_ids[:n_mult] # ensure we have enough labels for i, l in enumerate(label_ids): labels[:, i, l] = 1 else: labels = [None] # generate images from latent timeline pbar = ProgressBar(frame_count) for i in range(frame_count): latent = latents[i] # [X,512] label = labels[i % len(labels)] latmask = lmask[i % len(lmask)] if lmask is not None else [None] # [X,h,w] dc = dconst[i % len(dconst)] # [X,512,4,4] # generate multi-latent result if custom: output = Gs(latent, label, latmask, dc, truncation_psi=a.trunc, noise_mode='const') else: output = Gs(latent, label, truncation_psi=a.trunc, noise_mode='const') output = (output.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to( torch.uint8).cpu().numpy() # save image ext = 'png' if output.shape[3] == 4 else 'jpg' filename = osp.join(a.out_dir, "%06d.%s" % (i, ext)) imsave(filename, output[0]) pbar.upd() # convert latents to dlatents, save them if a.save_lat is True: latents = latents.squeeze(1) # [frm,512] dlatents = Gs.mapping(latents, label) # [frm,18,512] if a.size is None: a.size = [''] * 2 filename = '{}-{}-{}.npy'.format(basename(a.model), a.size[1], a.size[0]) filename = osp.join(osp.dirname(a.out_dir), filename) dlatents = dlatents.cpu().numpy() np.save(filename, dlatents) print('saved dlatents', dlatents.shape, 'to', filename)
def main(): if a.vector_dir is not None: if a.vector_dir.endswith('/') or a.vector_dir.endswith('\\'): a.vector_dir = a.vector_dir[:-1] os.makedirs(osp.join(a.out_dir, 'ttt'), exist_ok=True) global Gs, use_d # setup generator Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.func_name = 'training.stylegan2_multi.G_main' Gs_kwargs.verbose = a.verbose Gs_kwargs.size = a.size Gs_kwargs.scale_type = a.scale_type Gs_kwargs.impl = a.ops # load model with arguments sess = tflib.init_tf({'allow_soft_placement': True}) pkl_name = osp.splitext(a.model)[0] with open(pkl_name + '.pkl', 'rb') as file: network = pickle.load(file, encoding='latin1') try: _, _, network = network except: pass for k in list(network.static_kwargs.keys()): Gs_kwargs[k] = network.static_kwargs[k] # reload custom network, if needed if '.pkl' in a.model.lower(): print(' .. Gs from pkl ..', basename(a.model)) Gs = network else: # reconstruct network print(' .. Gs custom ..', basename(a.model)) Gs = tflib.Network('Gs', **Gs_kwargs) Gs.copy_vars_from(network) # load directions if a.vector_dir is not None: directions = [] vector_list = file_list(a.vector_dir, 'npy') for v in vector_list: direction = load_latents(v) if len(direction.shape) == 2: direction = np.expand_dims(direction, 0) directions.append(direction) directions = np.concatenate(directions)[:, np.newaxis] # [frm,1,18,512] else: print(' No vectors found') exit() if len(direction[0].shape) > 1 and direction[0].shape[0] > 1: use_d = True print(' directions', directions.shape, 'using d' if use_d else 'using w') # latent direction range lrange = [-0.5, 0.5] # load saved latents if a.npy_file is not None: base_latent = load_latents(a.npy_file) else: print(' No NPY input given, making random') z_dim = Gs.input_shape[1] shape = (1, z_dim) base_latent = np.random.randn(*shape) if use_d: base_latent = Gs.components.mapping.run(base_latent, None) # [frm,18,512] for i, direction in enumerate(directions): make_loop(base_latent, direction, lrange, a.fstep * 2, a.fstep * 2 * i)
def __init__( self, tfrecord, # tfrecords file resolution=None, # Dataset resolution, None = autodetect. label_file=None, # Relative path of the labels file, None = autodetect. max_label_size=0, # 0 = no labels, 'full' = full labels, <int> = N first label components. repeat=True, # Repeat dataset indefinitely? shuffle_mb=4096, # Shuffle data within specified window (megabytes), 0 = disable shuffling. prefetch_mb=2048, # Amount of data to prefetch (megabytes), 0 = disable prefetching. buffer_mb=256, # Read buffer size (megabytes). num_threads=2): # Number of concurrent threads. self.tfrecord = tfrecord self.resolution = None self.res_log2 = None self.shape = [] # [channels, height, width] self.dtype = 'uint8' self.dynamic_range = [0, 255] self.label_file = label_file self.label_size = None # components self.label_dtype = None self._np_labels = None self._tf_minibatch_in = None self._tf_labels_var = None self._tf_labels_dataset = None self._tf_dataset = None self._tf_iterator = None self._tf_init_op = None self._tf_minibatch_np = None self._cur_minibatch = -1 # List tfrecords files and inspect their shapes. assert os.path.isfile(self.tfrecord) tfr_file = self.tfrecord data_dir = os.path.dirname(tfr_file) tfr_opt = tf.python_io.TFRecordOptions( tf.python_io.TFRecordCompressionType.NONE) for record in tf.python_io.tf_record_iterator( tfr_file, tfr_opt): # in fact only one tfr_shape = self.parse_tfrecord_shape(record) # [c,h,w] jpg_data = tfr_shape[0] < 4 break # Autodetect label filename. if self.label_file is None: # guess = sorted(glob.glob(os.path.join(data_dir, '*.labels'))) guess = [ ff for ff in file_list(data_dir, 'labels') if basename( ff).split('-')[0] == basename(tfrecord).split('-')[0] ] if len(guess): self.label_file = guess[0] elif not os.path.isfile(self.label_file): guess = os.path.join(data_dir, self.label_file) if os.path.isfile(guess): self.label_file = guess # Determine shape and resolution self.shape = list(tfr_shape) max_res = calc_res(tfr_shape[1:]) self.resolution = resolution if resolution is not None else max_res self.res_log2 = int(np.ceil(np.log2(self.resolution))) self.init_res = [ int(s * 2**(2 - self.res_log2)) for s in self.shape[1:] ] # Load labels assert max_label_size == 'full' or max_label_size >= 0 self._np_labels = np.zeros([1 << 30, 0], dtype=np.float32) if self.label_file is not None and max_label_size != 0: self._np_labels = np.load(self.label_file) assert self._np_labels.ndim == 2 if max_label_size != 'full' and self._np_labels.shape[ 1] > max_label_size: self._np_labels = self._np_labels[:, :max_label_size] self.label_size = self._np_labels.shape[1] self.label_dtype = self._np_labels.dtype.name # Build TF expressions. with tf.name_scope('Dataset'), tf.device('/cpu:0'): self._tf_minibatch_in = tf.placeholder(tf.int64, name='minibatch_in', shape=[]) self._tf_labels_var = tflib.create_var_with_large_initial_value( self._np_labels, name='labels_var') self._tf_labels_dataset = tf.data.Dataset.from_tensor_slices( self._tf_labels_var) dset = tf.data.TFRecordDataset(tfr_file, compression_type='', buffer_size=buffer_mb << 20) if jpg_data is True: dset = dset.map(self.parse_tfrecord_tf_jpg, num_parallel_calls=num_threads) else: dset = dset.map(self.parse_tfrecord_tf, num_parallel_calls=num_threads) dset = tf.data.Dataset.zip((dset, self._tf_labels_dataset)) bytes_per_item = np.prod(tfr_shape) * np.dtype(self.dtype).itemsize if shuffle_mb > 0: dset = dset.shuffle(( (shuffle_mb << 20) - 1) // bytes_per_item + 1) if repeat: dset = dset.repeat() if prefetch_mb > 0: dset = dset.prefetch(( (prefetch_mb << 20) - 1) // bytes_per_item + 1) dset = dset.batch(self._tf_minibatch_in) self._tf_dataset = dset # self._tf_iterator = tf.data.Iterator.from_structure(tf.data.get_output_types(self._tf_dataset), tf.data.get_output_shapes(self._tf_dataset),) self._tf_iterator = tf.data.Iterator.from_structure( self._tf_dataset.output_types, self._tf_dataset.output_shapes) self._tf_init_op = self._tf_iterator.make_initializer( self._tf_dataset)
def main(): os.makedirs(a.out_dir, exist_ok=True) device = torch.device('cuda') # setup generator Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.verbose = a.verbose Gs_kwargs.size = a.size Gs_kwargs.scale_type = a.scale_type # load base or custom network pkl_name = osp.splitext(a.model)[0] if '.pkl' in a.model.lower(): custom = False print(' .. Gs from pkl ..', basename(a.model)) else: custom = True print(' .. Gs custom ..', basename(a.model)) with dnnlib.util.open_url(pkl_name + '.pkl') as f: Gs = legacy.load_network_pkl(f, custom=custom, **Gs_kwargs)['G_ema'].to( device) # type: ignore dlat_shape = (1, Gs.num_ws, Gs.w_dim) # [1,18,512] # read saved latents if a.dlatents is not None and osp.isfile(a.dlatents): key_dlatents = load_latents(a.dlatents) if len(key_dlatents.shape) == 2: key_dlatents = np.expand_dims(key_dlatents, 0) elif a.dlatents is not None and osp.isdir(a.dlatents): # if a.dlatents.endswith('/') or a.dlatents.endswith('\\'): a.dlatents = a.dlatents[:-1] key_dlatents = [] npy_list = file_list(a.dlatents, 'npy') for npy in npy_list: key_dlatent = load_latents(npy) if len(key_dlatent.shape) == 2: key_dlatent = np.expand_dims(key_dlatent, 0) key_dlatents.append(key_dlatent) key_dlatents = np.concatenate(key_dlatents) # [frm,18,512] else: print(' No input dlatents found') exit() key_dlatents = key_dlatents[:, np.newaxis] # [frm,1,18,512] print(' key dlatents', key_dlatents.shape) # replace higher layers with single (style) latent if a.style_dlat is not None: print(' styling with dlatent', a.style_dlat) style_dlatent = load_latents(a.style_dlat) while len(style_dlatent.shape) < 4: style_dlatent = np.expand_dims(style_dlatent, 0) # try replacing 5 by other value, less than Gs.num_ws key_dlatents[:, :, range(5, Gs.num_ws ), :] = style_dlatent[:, :, range(5, Gs.num_ws), :] frames = key_dlatents.shape[0] * a.fstep dlatents = latent_anima(dlat_shape, frames, a.fstep, key_latents=key_dlatents, cubic=a.cubic, verbose=True) # [frm,1,512] print(' dlatents', dlatents.shape) frame_count = dlatents.shape[0] dlatents = torch.from_numpy(dlatents).to(device) # distort image by tweaking initial const layer if a.digress > 0: try: init_res = Gs.init_res except Exception: init_res = (4, 4) # default initial layer size dconst = a.digress * latent_anima([1, Gs.z_dim, *init_res], frame_count, a.fstep, cubic=True, verbose=False) else: dconst = np.zeros([frame_count, 1, 1, 1, 1]) dconst = torch.from_numpy(dconst).to(device) # generate images from latent timeline pbar = ProgressBar(frame_count) for i in range(frame_count): # generate multi-latent result if custom: output = Gs.synthesis(dlatents[i], None, dconst[i], noise_mode='const') else: output = Gs.synthesis(dlatents[i], noise_mode='const') output = (output.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to( torch.uint8).cpu().numpy() ext = 'png' if output.shape[3] == 4 else 'jpg' filename = osp.join(a.out_dir, "%06d.%s" % (i, ext)) imsave(filename, output[0]) pbar.upd()
def main(): os.makedirs(a.out_dir, exist_ok=True) np.random.seed(seed=696) # setup generator fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.func_name = 'training.stylegan2_multi.G_main' Gs_kwargs.verbose = a.verbose Gs_kwargs.size = a.size Gs_kwargs.scale_type = a.scale_type Gs_kwargs.impl = a.ops # mask/blend latents with external latmask or by splitting the frame if a.latmask is None: nHW = [int(s) for s in a.nXY.split('-')][::-1] assert len(nHW)==2, ' Wrong count nXY: %d (must be 2)' % len(nHW) n_mult = nHW[0] * nHW[1] if a.verbose is True and n_mult > 1: print(' Latent blending w/split frame %d x %d' % (nHW[1], nHW[0])) lmask = np.tile(np.asarray([[[[None]]]]), (1,n_mult,1,1)) Gs_kwargs.countHW = nHW Gs_kwargs.splitfine = a.splitfine else: if a.verbose is True: print(' Latent blending with mask', a.latmask) n_mult = 2 if os.path.isfile(a.latmask): # single file lmask = np.asarray([[img_read(a.latmask)[:,:,0] / 255.]]) # [h,w] elif os.path.isdir(a.latmask): # directory with frame sequence lmask = np.asarray([[img_read(f)[:,:,0] / 255. for f in img_list(a.latmask)]]) # [h,w] else: print(' !! Blending mask not found:', a.latmask); exit(1) lmask = np.concatenate((lmask, 1 - lmask), 1) # [frm,2,h,w] Gs_kwargs.latmask_res = lmask.shape[2:] # load model with arguments sess = tflib.init_tf({'allow_soft_placement':True}) pkl_name = osp.splitext(a.model)[0] with open(pkl_name + '.pkl', 'rb') as file: network = pickle.load(file, encoding='latin1') try: _, _, network = network except: pass for k in list(network.static_kwargs.keys()): Gs_kwargs[k] = network.static_kwargs[k] # reload custom network, if needed if '.pkl' in a.model.lower(): print(' .. Gs from pkl ..', basename(a.model)) Gs = network else: # reconstruct network print(' .. Gs custom ..', basename(a.model)) # print(Gs_kwargs) Gs = tflib.Network('Gs', **Gs_kwargs) Gs.copy_vars_from(network) if a.verbose is True: print('kwargs:', ['%s: %s'%(kv[0],kv[1]) for kv in sorted(Gs.static_kwargs.items())]) if a.verbose is True: print(' out shape', Gs.output_shape[1:]) if a.size is None: a.size = Gs.output_shape[2:] if a.verbose is True: print(' making timeline..') lats = [] # list of [frm,1,512] for i in range(n_mult): lat_tmp = latent_anima((1, Gs.input_shape[1]), a.frames, a.fstep, cubic=a.cubic, gauss=a.gauss, verbose=False) # [frm,1,512] lats.append(lat_tmp) # list of [frm,1,512] latents = np.concatenate(lats, 1) # [frm,X,512] print(' latents', latents.shape) frame_count = latents.shape[0] # distort image by tweaking initial const layer if a.digress > 0: try: latent_size = Gs.static_kwargs['latent_size'] except: latent_size = 512 # default latent size try: init_res = Gs.static_kwargs['init_res'] except: init_res = (4,4) # default initial layer size dconst = [] for i in range(n_mult): dc_tmp = a.digress * latent_anima([1, latent_size, *init_res], a.frames, a.fstep, cubic=True, verbose=False) dconst.append(dc_tmp) dconst = np.concatenate(dconst, 1) else: dconst = np.zeros([frame_count, 1, 1, 1, 1]) # labels / conditions try: label_size = Gs_kwargs.label_size except: label_size = 0 if label_size > 0: labels = np.zeros((frame_count, n_mult, label_size)) # [frm,X,lbl] if a.labels is None: label_ids = [] for i in range(n_mult): label_ids.append(random.randint(0, label_size-1)) else: label_ids = [int(x) for x in a.labels.split('-')] label_ids = label_ids[:n_mult] # ensure we have enough labels for i, l in enumerate(label_ids): labels[:,i,l] = 1 else: labels = [None] # generate images from latent timeline pbar = ProgressBar(frame_count) for i in range(frame_count): latent = latents[i] # [X,512] label = labels[i % len(labels)] latmask = lmask[i % len(lmask)] if lmask is not None else [None] # [X,h,w] dc = dconst[i % len(dconst)] # [X,512,4,4] # generate multi-latent result if Gs.num_inputs == 2: output = Gs.run(latent, label, truncation_psi=a.trunc, randomize_noise=False, output_transform=fmt) else: output = Gs.run(latent, label, latmask, dc, truncation_psi=a.trunc, randomize_noise=False, output_transform=fmt) # save image ext = 'png' if output.shape[3]==4 else 'jpg' filename = osp.join(a.out_dir, "%06d.%s" % (i,ext)) imsave(filename, output[0]) pbar.upd() # convert latents to dlatents, save them if a.save_lat is True: latents = latents.squeeze(1) # [frm,512] dlatents = Gs.components.mapping.run(latents, label, dtype='float16') # [frm,18,512] filename = '{}-{}-{}.npy'.format(basename(a.model), a.size[1], a.size[0]) filename = osp.join(osp.dirname(a.out_dir), filename) np.save(filename, dlatents) print('saved dlatents', dlatents.shape, 'to', filename)
def main(): os.makedirs(a.out_dir, exist_ok=True) # parse filename to model parameters mparams = basename(a.model).split('-') res = int(mparams[1]) cfg = mparams[2] # setup generator fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.func_name = 'training.stylegan2_custom.G_main' Gs_kwargs.verbose = False Gs_kwargs.resolution = res Gs_kwargs.size = a.size Gs_kwargs.scale_type = a.scale_type Gs_kwargs.latent_size = a.latent_size if cfg.lower() == 'f': Gs_kwargs.synthesis_func = 'G_synthesis_stylegan2' elif cfg.lower() == 'e': Gs_kwargs.synthesis_func = 'G_synthesis_stylegan2' Gs_kwargs.fmap_base = 8 << 10 else: print(' old modes [A-D] not implemented') exit() # check initial model resolution if len(mparams) > 3: if 'x' in mparams[3].lower(): init_res = [int(x) for x in mparams[3].lower().split('x')] Gs_kwargs.init_res = list(reversed(init_res)) # [H,W] # load model, check channels sess = tflib.init_tf({'allow_soft_placement': True}) pkl_name = osp.splitext(a.model)[0] with open(pkl_name + '.pkl', 'rb') as file: network = pickle.load(file, encoding='latin1') try: _, _, Gs = network except: Gs = network Gs_kwargs.num_channels = Gs.output_shape[1] # reload custom network, if needed if '.pkl' in a.model.lower(): print(' .. Gs from pkl ..') else: print(' .. Gs custom ..') Gs = tflib.Network('Gs', **Gs_kwargs) Gs.copy_vars_from(network) z_dim = Gs.input_shape[1] dz_dim = a.dlatent_size # 512 dl_dim = 2 * (int(np.floor(np.log2(res))) - 1) dlat_shape = (1, dl_dim, dz_dim) # [1,18,512] # read saved latents if a.dlatents is not None and osp.isfile(a.dlatents): key_dlatents = load_latents(a.dlatents) if len(key_dlatents.shape) == 2: key_dlatents = np.expand_dims(key_dlatents, 0) elif a.dlatents is not None and osp.isdir(a.dlatents): # if a.dlatents.endswith('/') or a.dlatents.endswith('\\'): a.dlatents = a.dlatents[:-1] key_dlatents = [] npy_list = file_list(a.dlatents, 'npy') for npy in npy_list: key_dlatent = load_latents(npy) if len(key_dlatent.shape) == 2: key_dlatent = np.expand_dims(key_dlatent, 0) key_dlatents.append(key_dlatent) key_dlatents = np.concatenate(key_dlatents) # [frm,18,512] else: print(' No input dlatents found') exit() key_dlatents = key_dlatents[:, np.newaxis] # [frm,1,18,512] print(' key dlatents', key_dlatents.shape) # replace higher layers with single (style) latent if a.style_npy_file is not None: print(' styling with latent', a.style_npy_file) style_dlatent = load_latents(a.style_npy_file) while len(style_dlatent.shape) < 4: style_dlatent = np.expand_dims(style_dlatent, 0) # try other values < dl_dim besides 5 key_dlatents[:, :, range(5, dl_dim), :] = style_dlatent[:, :, range(5, dl_dim), :] frames = key_dlatents.shape[0] * a.fstep dlatents = latent_anima(dlat_shape, frames, a.fstep, key_latents=key_dlatents, cubic=a.cubic, verbose=True) # [frm,1,512] print(' dlatents', dlatents.shape) # truncation trick dlatent_avg = Gs.get_var('dlatent_avg') # (512,) tr_range = range(0, 8) dlatents[:, :, tr_range, :] = dlatent_avg + (dlatents[:, :, tr_range, :] - dlatent_avg) * a.trunc # loop for graph frame by frame frame_count = dlatents.shape[0] pbar = ProgressBar(frame_count) for i in range(frame_count): dlatent = dlatents[i] output = Gs.components.synthesis.run(dlatent, randomize_noise=False, output_transform=fmt, minibatch_size=1) ext = 'png' if output.shape[3] == 4 else 'jpg' filename = osp.join(a.out_dir, "%05d.%s" % (i, ext)) imsave(filename, output[0]) pbar.upd()
def main(): os.makedirs(a.out_dir, exist_ok=True) np.random.seed(seed=696) # parse filename to model parameters mparams = basename(a.model).split('-') res = int(mparams[1]) cfg = mparams[2] # setup generator fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.func_name = 'training.stylegan2_custom.G_main' Gs_kwargs.verbose = False Gs_kwargs.resolution = res Gs_kwargs.size = a.size Gs_kwargs.scale_type = a.scale_type Gs_kwargs.latent_size = a.latent_size Gs_kwargs.impl = a.ops if cfg.lower() == 'f': Gs_kwargs.synthesis_func = 'G_synthesis_stylegan2' elif cfg.lower() == 'e': Gs_kwargs.synthesis_func = 'G_synthesis_stylegan2' Gs_kwargs.fmap_base = 8 << 10 else: print(' old modes [A-D] not implemented'); exit() # check initial model resolution if len(mparams) > 3: if 'x' in mparams[3].lower(): init_res = [int(x) for x in mparams[3].lower().split('x')] Gs_kwargs.init_res = list(reversed(init_res)) # [H,W] # load model, check channels sess = tflib.init_tf({'allow_soft_placement':True}) pkl_name = osp.splitext(a.model)[0] with open(pkl_name + '.pkl', 'rb') as file: network = pickle.load(file, encoding='latin1') try: _, _, Gs = network except: Gs = network Gs_kwargs.num_channels = Gs.output_shape[1] # reload custom network, if needed if '.pkl' in a.model.lower(): print(' .. Gs from pkl ..') else: print(' .. Gs custom ..') Gs = tflib.Network('Gs', **Gs_kwargs) Gs.copy_vars_from(network) # Gs.print_layers() print(' out shape', Gs.output_shape[1:]) if a.size is None: a.size = Gs.output_shape[2:] z_dim = Gs.input_shape[1] shape = (1, z_dim) print(' making timeline..') latents = latent_anima(shape, a.frames, a.fstep, cubic=a.cubic, gauss=a.gauss, verbose=True) # [frm,1,512] print(' latents', latents.shape) # generate images from latent timeline frame_count = latents.shape[0] pbar = ProgressBar(frame_count) for i in range(frame_count): output = Gs.run(latents[i], [None], truncation_psi=a.trunc, randomize_noise=False, output_transform=fmt) ext = 'png' if output.shape[3]==4 else 'jpg' filename = osp.join(a.out_dir, "%05d.%s" % (i,ext)) imsave(filename, output[0]) pbar.upd() # convert latents to dlatents, save them latents = latents.squeeze(1) # [frm,512] dlatents = Gs.components.mapping.run(latents, None, latent_size=z_dim, dtype='float16') # [frm,18,512] filename = '{}-{}-{}.npy'.format(basename(a.model), a.size[1], a.size[0]) filename = osp.join(osp.dirname(a.out_dir), filename) np.save(filename, dlatents) print('saved dlatents', dlatents.shape, 'to', filename)
def main(): os.makedirs(a.out_dir, exist_ok=True) # setup generator fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.func_name = 'training.stylegan2_multi.G_main' Gs_kwargs.verbose = a.verbose Gs_kwargs.size = a.size Gs_kwargs.scale_type = a.scale_type Gs_kwargs.impl = a.ops # load model with arguments sess = tflib.init_tf({'allow_soft_placement': True}) pkl_name = osp.splitext(a.model)[0] with open(pkl_name + '.pkl', 'rb') as file: network = pickle.load(file, encoding='latin1') try: _, _, network = network except: pass for k in list(network.static_kwargs.keys()): Gs_kwargs[k] = network.static_kwargs[k] # reload custom network, if needed if '.pkl' in a.model.lower(): print(' .. Gs from pkl ..', basename(a.model)) Gs = network else: # reconstruct network print(' .. Gs custom ..', basename(a.model)) Gs = tflib.Network('Gs', **Gs_kwargs) Gs.copy_vars_from(network) z_dim = Gs.input_shape[1] dz_dim = 512 # dlatent_size try: dl_dim = 2 * (int(np.floor(np.log2(Gs_kwargs.resolution))) - 1) except: print(' Resave model, no resolution kwarg found!') exit(1) dlat_shape = (1, dl_dim, dz_dim) # [1,18,512] # read saved latents if a.dlatents is not None and osp.isfile(a.dlatents): key_dlatents = load_latents(a.dlatents) if len(key_dlatents.shape) == 2: key_dlatents = np.expand_dims(key_dlatents, 0) elif a.dlatents is not None and osp.isdir(a.dlatents): # if a.dlatents.endswith('/') or a.dlatents.endswith('\\'): a.dlatents = a.dlatents[:-1] key_dlatents = [] npy_list = file_list(a.dlatents, 'npy') for npy in npy_list: key_dlatent = load_latents(npy) if len(key_dlatent.shape) == 2: key_dlatent = np.expand_dims(key_dlatent, 0) key_dlatents.append(key_dlatent) key_dlatents = np.concatenate(key_dlatents) # [frm,18,512] else: print(' No input dlatents found') exit() key_dlatents = key_dlatents[:, np.newaxis] # [frm,1,18,512] print(' key dlatents', key_dlatents.shape) # replace higher layers with single (style) latent if a.style_npy_file is not None: print(' styling with latent', a.style_npy_file) style_dlatent = load_latents(a.style_npy_file) while len(style_dlatent.shape) < 4: style_dlatent = np.expand_dims(style_dlatent, 0) # try replacing 5 by other value, less than dl_dim key_dlatents[:, :, range(5, dl_dim), :] = style_dlatent[:, :, range(5, dl_dim), :] frames = key_dlatents.shape[0] * a.fstep dlatents = latent_anima(dlat_shape, frames, a.fstep, key_latents=key_dlatents, cubic=a.cubic, verbose=True) # [frm,1,512] print(' dlatents', dlatents.shape) frame_count = dlatents.shape[0] # truncation trick dlatent_avg = Gs.get_var('dlatent_avg') # (512,) tr_range = range(0, 8) dlatents[:, :, tr_range, :] = dlatent_avg + (dlatents[:, :, tr_range, :] - dlatent_avg) * a.trunc # distort image by tweaking initial const layer if a.digress > 0: try: latent_size = Gs.static_kwargs['latent_size'] except: latent_size = 512 # default latent size try: init_res = Gs.static_kwargs['init_res'] except: init_res = (4, 4) # default initial layer size dconst = a.digress * latent_anima([1, latent_size, *init_res], frames, a.fstep, cubic=True, verbose=False) else: dconst = np.zeros([frame_count, 1, 1, 1, 1]) # generate images from latent timeline pbar = ProgressBar(frame_count) for i in range(frame_count): if a.digress is True: tf.get_default_session().run(tf.assign(wvars[0], wts[i])) # generate multi-latent result if Gs.num_inputs == 2: output = Gs.components.synthesis.run(dlatents[i], randomize_noise=False, output_transform=fmt, minibatch_size=1) else: output = Gs.components.synthesis.run(dlatents[i], [None], dconst[i], randomize_noise=False, output_transform=fmt, minibatch_size=1) ext = 'png' if output.shape[3] == 4 else 'jpg' filename = osp.join(a.out_dir, "%06d.%s" % (i, ext)) imsave(filename, output[0]) pbar.upd()
def run_projection(network_pkl: str, in_dir: str, out_dir: str, save_video: bool, seed: int, steps: int): """Project given image to the latent space of pretrained network pickle. Examples: python projector.py --outdir=out --target=~/mytargetimg.png \\ --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl """ np.random.seed(seed) torch.manual_seed(seed) # Load networks. print('Loading networks from "%s"...' % network_pkl) device = torch.device('cuda') with dnnlib.util.open_url(network_pkl) as fp: G = legacy.load_network_pkl(fp)['G_ema'].requires_grad_(False).to( device) # type: ignore img_files = img_list(in_dir) num_images = len(img_files) for image_idx in range(num_images): fname = basename(img_files[image_idx]) print('Projecting image %d/%d .. %s' % (image_idx + 1, num_images, basename(img_files[image_idx]))) work_dir = os.path.join(out_dir, fname) os.makedirs(work_dir, exist_ok=True) # Load target image. target_pil = PIL.Image.open(img_files[image_idx]).convert('RGB') w, h = target_pil.size s = min(w, h) target_pil = target_pil.crop( ((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2)) target_pil = target_pil.resize((G.img_resolution, G.img_resolution), PIL.Image.LANCZOS) target_uint8 = np.array(target_pil, dtype=np.uint8) # Optimize projection. # start_time = perf_counter() projected_w_steps = project( G, target=torch.tensor(target_uint8.transpose([2, 0, 1]), device=device), # pylint: disable=not-callable num_steps=steps, device=device, verbose=True) # print (f'Elapsed: {(perf_counter()-start_time):.1f} s') # Render debug output: optional video and projected image and W vector. os.makedirs(out_dir, exist_ok=True) if save_video: vfile = '%s/proj.mp4' % work_dir video = imageio.get_writer(vfile, mode='I', fps=25, codec='libx264', bitrate='16M') print('Saving optimization progress video %s' % vfile) for projected_w in projected_w_steps: synth_image = G.synthesis(projected_w.unsqueeze(0), noise_mode='const') synth_image = (synth_image + 1) * (255 / 2) synth_image = synth_image.permute(0, 2, 3, 1).clamp(0, 255).to( torch.uint8)[0].cpu().numpy() video.append_data( np.concatenate([target_uint8, synth_image], axis=1)) video.close() # Save final projected frame and W vector. target_pil.save('%s/target.jpg' % work_dir) projected_w = projected_w_steps[-1] synth_image = G.synthesis(projected_w.unsqueeze(0), noise_mode='const') synth_image = (synth_image + 1) * (255 / 2) synth_image = synth_image.permute(0, 2, 3, 1).clamp(0, 255).to( torch.uint8)[0].cpu().numpy() PIL.Image.fromarray(synth_image, 'RGB').save('%s/%s.jpg' % (work_dir, fname)) np.savez('%s/%s.npz' % (work_dir, fname), w=projected_w.unsqueeze(0).cpu().numpy())
def run(dataset, train_dir, config, d_aug, diffaug_policy, cond, ops, jpg_data, mirror, mirror_v, \ lod_step_kimg, batch_size, resume, resume_kimg, finetune, num_gpus, ema_kimg, gamma, freezeD): # dataset (tfrecords) - preprocess or get tfr_files = file_list(os.path.dirname(dataset), 'tfr') tfr_files = [f for f in tfr_files if basename(dataset) in f] if len(tfr_files) == 0: tfr_file, total_samples = create_from_images(dataset, jpg=jpg_data) else: tfr_file = tfr_files[0] dataset_args = EasyDict(tfrecord=tfr_file, jpg_data=jpg_data) desc = basename(tfr_file).split('-')[0] # training functions if d_aug: # https://github.com/mit-han-lab/data-efficient-gans train = EasyDict( run_func_name='training.training_loop_diffaug.training_loop' ) # Options for training loop (Diff Augment method) loss_args = EasyDict( func_name='training.loss_diffaug.ns_DiffAugment_r1', policy=diffaug_policy) # Options for loss (Diff Augment method) else: # original nvidia train = EasyDict(run_func_name='training.training_loop.training_loop' ) # Options for training loop (original from NVidia) G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg' ) # Options for generator loss. D_loss = EasyDict(func_name='training.loss.D_logistic_r1' ) # Options for discriminator loss. # network functions G = EasyDict(func_name='training.networks_stylegan2.G_main' ) # Options for generator network. D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2' ) # Options for discriminator network. G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer. D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='1080p', layout='random') # Options for setup_snapshot_image_grid(). sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). G.impl = D.impl = ops # resolutions data_res = basename(tfr_file).split('-')[-1].split( 'x') # get resolution from dataset filename data_res = list(reversed([int(x) for x in data_res])) # convert to int list init_res, resolution, res_log2 = calc_init_res(data_res) if init_res != [4, 4]: print(' custom init resolution', init_res) G.init_res = D.init_res = list(init_res) train.setname = desc + config desc = '%s-%d-%s' % (desc, resolution, config) # training schedule sched.lod_training_kimg = lod_step_kimg sched.lod_transition_kimg = lod_step_kimg train.total_kimg = lod_step_kimg * res_log2 * 2 # a la ProGAN if finetune: train.total_kimg = 15000 # should start from ~10k kimg train.image_snapshot_ticks = 1 train.network_snapshot_ticks = 5 train.mirror_augment = mirror train.mirror_augment_v = mirror_v # learning rate if config == 'e': if finetune: # uptrain 1024 sched.G_lrate_base = 0.001 else: # train 1024 sched.G_lrate_base = 0.001 sched.G_lrate_dict = {0: 0.001, 1: 0.0007, 2: 0.0005, 3: 0.0003} sched.lrate_step = 1500 # period for stepping to next lrate, in kimg if config == 'f': # sched.G_lrate_base = 0.0003 sched.G_lrate_base = 0.001 sched.D_lrate_base = sched.G_lrate_base # *2 - not used anyway sched.minibatch_gpu_base = batch_size sched.minibatch_size_base = num_gpus * sched.minibatch_gpu_base sc.num_gpus = num_gpus if config == 'e': G.fmap_base = D.fmap_base = 8 << 10 if d_aug: loss_args.gamma = 100 if gamma is None else gamma else: D_loss.gamma = 100 if gamma is None else gamma elif config == 'f': G.fmap_base = D.fmap_base = 16 << 10 else: print(' Only configs E and F are implemented') exit() if cond: desc += '-cond' dataset_args.max_label_size = 'full' # conditioned on full label if freezeD: D.freezeD = True train.resume_with_new_nets = True if d_aug: desc += '-daug' sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, tf_config=tf_config) kwargs.update(resume_pkl=resume, resume_kimg=resume_kimg, resume_with_new_nets=True) if ema_kimg is not None: kwargs.update(G_ema_kimg=ema_kimg) if d_aug: kwargs.update(loss_args=loss_args) else: kwargs.update(G_loss_args=G_loss, D_loss_args=D_loss) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = train_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def run(data, train_dir, config, d_aug, diffaug_policy, cond, ops, mirror, mirror_v, \ kimg, batch_size, lrate, resume, resume_kimg, num_gpus, ema_kimg, gamma, freezeD): # training functions if d_aug: # https://github.com/mit-han-lab/data-efficient-gans train = EasyDict( run_func_name='training.training_loop_diffaug.training_loop' ) # Options for training loop (Diff Augment method) loss_args = EasyDict( func_name='training.loss_diffaug.ns_DiffAugment_r1', policy=diffaug_policy) # Options for loss (Diff Augment method) else: # original nvidia train = EasyDict(run_func_name='training.training_loop.training_loop' ) # Options for training loop (original from NVidia) G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg' ) # Options for generator loss. D_loss = EasyDict(func_name='training.loss.D_logistic_r1' ) # Options for discriminator loss. # network functions G = EasyDict(func_name='training.networks_stylegan2.G_main' ) # Options for generator network. D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2' ) # Options for discriminator network. G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer. D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. sched = EasyDict() # Options for TrainingSchedule. grid = EasyDict( size='1080p', layout='random') # Options for setup_snapshot_image_grid(). sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). G.impl = D.impl = ops # dataset (tfrecords) - get or create tfr_files = file_list(os.path.dirname(data), 'tfr') tfr_files = [ f for f in tfr_files if basename(data) == basename(f).split('-')[0] ] if len(tfr_files) == 0 or os.stat(tfr_files[0]).st_size == 0: tfr_file, total_samples = create_from_image_folders( data) if cond is True else create_from_images(data) else: tfr_file = tfr_files[0] dataset_args = EasyDict(tfrecord=tfr_file) # resolutions with tf.Graph().as_default(), tflib.create_session().as_default(): # pylint: disable=not-context-manager dataset_obj = dataset.load_dataset( **dataset_args) # loading the data to see what comes out resolution = dataset_obj.resolution init_res = dataset_obj.init_res res_log2 = dataset_obj.res_log2 dataset_obj.close() dataset_obj = None if list(init_res) == [4, 4]: desc = '%s-%d' % (basename(data), resolution) else: print(' custom init resolution', init_res) desc = basename(tfr_file) G.init_res = D.init_res = list(init_res) train.savenames = [desc.replace(basename(data), 'snapshot'), desc] desc += '-%s' % config # training schedule train.total_kimg = kimg train.image_snapshot_ticks = 1 * num_gpus if kimg <= 1000 else 4 * num_gpus train.network_snapshot_ticks = 5 train.mirror_augment = mirror train.mirror_augment_v = mirror_v sched.tick_kimg_base = 2 if train.total_kimg < 2000 else 4 # learning rate if config == 'e': sched.G_lrate_base = 0.001 sched.G_lrate_dict = {0: 0.001, 1: 0.0007, 2: 0.0005, 3: 0.0003} sched.lrate_step = 1500 # period for stepping to next lrate, in kimg if config == 'f': sched.G_lrate_base = lrate # 0.001 for big datasets, 0.0003 for few-shot sched.D_lrate_base = sched.G_lrate_base # *2 - not used anyway # batch size (for 16gb memory GPU) sched.minibatch_gpu_base = 4096 // resolution if batch_size is None else batch_size print(' Batch size', sched.minibatch_gpu_base) sched.minibatch_size_base = num_gpus * sched.minibatch_gpu_base sc.num_gpus = num_gpus if config == 'e': G.fmap_base = D.fmap_base = 8 << 10 if d_aug: loss_args.gamma = 100 if gamma is None else gamma else: D_loss.gamma = 100 if gamma is None else gamma elif config == 'f': G.fmap_base = D.fmap_base = 16 << 10 else: print(' Only configs E and F are implemented') exit() if cond: desc += '-cond' dataset_args.max_label_size = 'full' # conditioned on full label if freezeD: D.freezeD = True train.resume_with_new_nets = True if d_aug: desc += '-daug' sc.submit_target = dnnlib.SubmitTarget.LOCAL sc.local.do_not_copy_source_files = True kwargs = EasyDict(train) kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt) kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, tf_config=tf_config) kwargs.update(resume_pkl=resume, resume_kimg=resume_kimg, resume_with_new_nets=True) if ema_kimg is not None: kwargs.update(G_ema_kimg=ema_kimg) if d_aug: kwargs.update(loss_args=loss_args) else: kwargs.update(G_loss_args=G_loss, D_loss_args=D_loss) kwargs.submit_config = copy.deepcopy(sc) kwargs.submit_config.run_dir_root = train_dir kwargs.submit_config.run_desc = desc dnnlib.submit_run(**kwargs)
def main(): if a.vector_dir is not None: if a.vector_dir.endswith('/') or a.vector_dir.endswith('\\'): a.vector_dir = a.vector_dir[:-1] os.makedirs(a.out_dir, exist_ok=True) device = torch.device('cuda') global Gs, use_d, custom # setup generator Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.verbose = a.verbose Gs_kwargs.size = a.size Gs_kwargs.scale_type = a.scale_type # load base or custom network pkl_name = osp.splitext(a.model)[0] if '.pkl' in a.model.lower(): custom = False print(' .. Gs from pkl ..', basename(a.model)) else: custom = True print(' .. Gs custom ..', basename(a.model)) with dnnlib.util.open_url(pkl_name + '.pkl') as f: Gs = legacy.load_network_pkl(f, custom=custom, **Gs_kwargs)['G_ema'].to( device) # type: ignore # load directions if a.vector_dir is not None: directions = [] vector_list = file_list(a.vector_dir, 'npy') for v in vector_list: direction = load_latents(v) if len(direction.shape) == 2: direction = np.expand_dims(direction, 0) directions.append(direction) directions = np.concatenate(directions)[:, np.newaxis] # [frm,1,18,512] else: print(' No vectors found') exit() if len(direction[0].shape) > 1 and direction[0].shape[0] > 1: use_d = True print(' directions', directions.shape, 'using d' if use_d else 'using w') directions = torch.from_numpy(directions).to(device) # latent direction range lrange = [-0.5, 0.5] # load saved latents if a.base_lat is not None: base_latent = load_latents(a.base_lat) base_latent = torch.from_numpy(base_latent).to(device) else: print(' No NPY input given, making random') base_latent = np.random.randn(1, Gs.z_dim) if use_d: base_latent = Gs.mapping(base_latent, None) # [frm,18,512] pbar = ProgressBar(len(directions)) for i, direction in enumerate(directions): make_loop(base_latent, direction, lrange, a.fstep * 2, a.fstep * 2 * i) pbar.upd()
def main(): tflib.init_tf({'allow_soft_placement': True}) G_in, D_in, Gs_in = load_pkl(a.source) print(' Loading model', a.source, Gs_in.output_shape) _, res_in, _ = calc_init_res(Gs_in.output_shape[1:]) if a.res is not None or a.alpha is True: if a.res is None: a.res = Gs_in.output_shape[2:] colors = 4 if a.alpha is True else Gs_in.output_shape[ 1] # EXPERIMENTAL _, res_out, _ = calc_init_res([colors, *a.res]) if res_in != res_out or a.alpha is True: # add or remove layers assert G_in is not None and D_in is not None, " !! G/D subnets not found in source model !!" data_shape = [colors, res_out, res_out] print(' Reconstructing full model with shape', data_shape) G_out, D_out, Gs_out = create_model(data_shape, True, 0, Gs_in.static_kwargs) copy_vars(Gs_in, Gs_out) copy_vars(G_in, G_out) copy_vars(D_in, D_out, D=True) G_in, D_in, Gs_in = G_out, D_out, Gs_out a.full = True if a.res[0] != res_out or a.res[1] != res_out: # crop or pad layers data_shape = [colors, *a.res] G_out, D_out, Gs_out = create_model(data_shape, True, 0, Gs_in.static_kwargs) if G_in is not None and D_in is not None: print(' Reconstructing full model with shape', data_shape) copy_and_crop_or_pad_trainables(G_in, G_out) copy_and_crop_or_pad_trainables(D_in, D_out) G_in, D_in = G_out, D_out a.full = True else: print(' Reconstructing Gs model with shape', data_shape) copy_and_crop_or_pad_trainables(Gs_in, Gs_out) Gs_in = Gs_out if a.labels is not None: assert G_in is not None and D_in is not None, " !! G/D subnets not found in source model !!" print(' Reconstructing full model with labels', a.labels) data_shape = Gs_in.output_shape[1:] G_out, D_out, Gs_out = create_model(data_shape, True, a.labels, Gs_in.static_kwargs) if a.verbose is True: D_out.print_layers() if a.verbose is True: G_out.print_layers() copy_and_fill_trainables(G_in, G_out) copy_and_fill_trainables(D_in, D_out) copy_and_fill_trainables(Gs_in, Gs_out) a.full = True if a.labels is None and a.res is None and a.alpha is not True: if a.reconstruct is True: print(' Reconstructing model with same size /', 'full' if a.full else 'Gs') data_shape = Gs_in.output_shape[1:] G_out, D_out, Gs_out = create_model(data_shape, a.full, 0, Gs_in.static_kwargs) Gs_out.copy_vars_from(Gs_in) if a.full is True and G_in is not None and D_in is not None: G_out.copy_vars_from(G_in) D_out.copy_vars_from(D_in) else: Gs_out = Gs_in out_name = basename(a.source) if a.res is not None: out_name += '-%dx%d' % (a.res[1], a.res[0]) if a.alpha is True: out_name += 'a' if a.labels is not None: out_name += '-c%d' % a.labels if a.full is True: # G_in is not None and D_in is not None save_pkl((G_out, D_out, Gs_out), os.path.join(a.out_dir, '%s.pkl' % out_name)) else: save_pkl(Gs_out, os.path.join(a.out_dir, '%s-Gs.pkl' % out_name)) print(' Done')