stylegan2-ada_ffhq_gen_pairs

"""Generate images using pretrained network pickle."""

import argparse
import sys
import os
import subprocess
import pickle
import re

import scipy
import numpy as np
import PIL.Image

import dnnlib
import dnnlib.tflib as tflib

os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"
import moviepy.editor
from opensimplex import OpenSimplex

import warnings # mostly numpy warnings for me
warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=DeprecationWarning)
tflib.init_tf()

# Load pre-trained network.
# See here: https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/
# and here: https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/
url_1 = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-ffhq-config-f.pkl'
url_2 = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/metfaces.pkl'    #Replace with your own FFHQ fine tuned model
# Here you can find my MEGA pkl collection with a whole lot of pre-trained ones: https://mega.nz/folder/OtllzJwa#C947mCCdEfMCRTWnDcs4qw

with dnnlib.util.open_url(url_1, cache_dir='/Models/') as f_1:
    _G_1, _D_1, Gs_1 = pickle.load(f_1)
    # _G = Instantaneous snapshot of the generator. Mainly useful for resuming a previous training run.
    # _D = Instantaneous snapshot of the discriminator. Mainly useful for resuming a previous training run.
    # Gs = Long-term average of the generator. Yields higher-quality results than the instantaneous snapshot.
with dnnlib.util.open_url(url_2, cache_dir='/Models/') as f_2:
    _G_2, _D_2, Gs_2 = pickle.load(f_2)

import os

out_path = './out'

if os.path.isdir(out_path) == False:
    os.mkdir(out_path)

from PIL import Image

with dnnlib.util.open_url(url_1, cache_dir='/Models/') as f_1:
    _G_1, _D_1, Gs_1 = pickle.load(f_1)
    # _G = Instantaneous snapshot of the generator. Mainly useful for resuming a previous training run.
    # _D = Instantaneous snapshot of the discriminator. Mainly useful for resuming a previous training run.
    # Gs = Long-term average of the generator. Yields higher-quality results t
model_1 = 'ffhq'
model_2 = 'metfaces'

random_seed = 1042

rnd = np.random.RandomState(random_seed)
seed_path = out_path + '/' + str(random_seed)
if os.path.isdir(seed_path) == False:
    os.mkdir(seed_path)

model_1_path = seed_path + '/' + model_1
if os.path.isdir(model_1_path) == False:
    os.mkdir(model_1_path)

model_2_path = seed_path + '/' + model_2
if os.path.isdir(model_2_path) == False:
    os.mkdir(model_2_path)

# Generate images
import skimage
import skimage.io
import skimage.io._plugins.pil_plugin as pp

number_of_gen_cycles = 1750     #1750 * 4 = 7K image pairs, good for p2s2p, e4e. For Pix2PixHD use a higher number like 6500 (*4 = 26000)
number_of_images_per_gen = 4
number_of_gen_cycles_for_noise = 3  #3: Every 12 images, generate 1 with resize blurry artifect and one with dirty poisson noise
image_output_size = (512, 512)      #Uncomment to resize to the size you need
#image_output_size = (1024, 1024)
psi = 0.7
fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)

for gen_cycle in range(number_of_gen_cycles):
    latents = rnd.randn(number_of_images_per_gen, Gs_1.input_shape[1])
    images = Gs_1.run(latents, None, minibatch = 1, truncation_psi=psi, randomize_noise=True, output_transform=fmt)
    for i in range(number_of_images_per_gen):
        im = Image.fromarray(images[i]).resize(image_output_size)
        if (number_of_gen_cycles_for_noise > 0) and ((gen_cycle % number_of_gen_cycles_for_noise) == 0):
            if i == 0:
                im64x64 = Image.fromarray(images[i]).resize((96, 96))
                im = im64x64.resize(image_output_size)
            if i == 1:
                im = pp.ndarray_to_pil(skimage.util.random_noise(pp.pil_to_ndarray(im), mode="poisson"))
        im.save(model_1_path + '/StyleGanImg_' + str(random_seed) + '_' + str(gen_cycle*number_of_images_per_gen + i) + '_' + str(i) + '.jpg')
    images = Gs_2.run(latents, None, minibatch = 1, truncation_psi=psi, randomize_noise=True, output_transform=fmt)
    for i in range(number_of_images_per_gen):
        im = Image.fromarray(images[i]).resize(image_output_size)
        im.save(model_2_path + '/StyleGanImg_' + str(random_seed) + '_' + str(gen_cycle*number_of_images_per_gen + i) + '_' + str(i) + '.jpg')