Spaces:
Runtime error
Runtime error
File size: 4,672 Bytes
a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://nvlabs.github.io/stylegan2/license.html
## this script is for generating images from pre-trained network based on StyleGAN1 (TensorFlow) and StyleGAN2-ada (PyTorch) ##
import os
import click
import dnnlib
import numpy as np
import PIL.Image
import legacy
from typing import List, Optional
"""
Generate images using pretrained network pickle.
Examples:
\b
# Generate human full-body images without truncation
python generate.py --outdir=outputs/generate/stylegan_human_v2_1024 --trunc=1 --seeds=1,3,5,7 \\
--network=pretrained_models/stylegan_human_v2_1024.pkl --version 2
\b
# Generate human full-body images with truncation
python generate.py --outdir=outputs/generate/stylegan_human_v2_1024 --trunc=0.8 --seeds=0-100\\
--network=pretrained_models/stylegan_human_v2_1024.pkl --version 2
# \b
# Generate human full-body images using stylegan V1
# python generate.py --outdir=outputs/generate/stylegan_human_v1_1024 \\
# --network=pretrained_models/stylegan_human_v1_1024.pkl --version 1
"""
@click.command()
@click.pass_context
@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
@click.option('--seeds', type=legacy.num_range, help='List of random seeds')
@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)
@click.option('--noise-mode', help='Noise mode', type=click.Choice(['const', 'random', 'none']), default='const', show_default=True)
@click.option('--outdir', help='Where to save the output images', default='outputs/generate/', type=str, required=True, metavar='DIR')
@click.option('--version', help="stylegan version, 1, 2 or 3", type=int, default=2)
def generate_images(
ctx: click.Context,
network_pkl: str,
seeds: Optional[List[int]],
truncation_psi: float,
noise_mode: str,
outdir: str,
version: int
):
print('Loading networks from "%s"...' % network_pkl)
if version == 1:
import dnnlib.tflib as tflib
tflib.init_tf()
G, D, Gs = legacy.load_pkl(network_pkl)
else:
import torch
device = torch.device('cuda')
with dnnlib.util.open_url(network_pkl) as f:
G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
os.makedirs(outdir, exist_ok=True)
if seeds is None:
ctx.fail('--seeds option is required.')
# Generate images.
target_z = np.array([])
target_w = np.array([])
latent_out = outdir.replace('/images/', '')
for seed_idx, seed in enumerate(seeds):
if seed % 5000 == 0:
print('Generating image for seed %d (%d/%d) ...' %
(seed, seed_idx, len(seeds)))
if version == 1: # stylegan v1
z = np.random.RandomState(seed).randn(1, Gs.input_shape[1])
# Generate image.
fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
if noise_mode == 'const':
randomize_noise = False
else:
randomize_noise = True
images = Gs.run(z, None, truncation_psi=truncation_psi,
randomize_noise=randomize_noise, output_transform=fmt)
PIL.Image.fromarray(images[0], 'RGB').save(
f'{outdir}/seed{seed:04d}.png')
else: # stylegan v2/v3
label = torch.zeros([1, G.c_dim], device=device)
z = torch.from_numpy(np.random.RandomState(
seed).randn(1, G.z_dim)).to(device)
if target_z.size == 0:
target_z = z.cpu()
else:
target_z = np.append(target_z, z.cpu(), axis=0)
w = G.mapping(z, label, truncation_psi=truncation_psi)
img = G.synthesis(w, noise_mode=noise_mode, force_fp32=True)
if target_w.size == 0:
target_w = w.cpu()
else:
target_w = np.append(target_w, w.cpu(), axis=0)
img = (img.permute(0, 2, 3, 1) * 127.5 +
128).clamp(0, 255).to(torch.uint8)
PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(
f'{outdir}/seed{seed:04d}.png')
# print(target_z)
# print(target_z.shape,target_w.shape)
# ----------------------------------------------------------------------------
if __name__ == "__main__":
generate_images()
# ----------------------------------------------------------------------------
|