├── train └── .gitkeep ├── _in ├── dlats │ ├── ffhq-1024-f-1.txt │ ├── ffhq-1024-f-2.txt │ ├── ffhq-1024-f-1.npy │ └── ffhq-1024-f-2.npy ├── mask.jpg ├── blonde458.npy ├── photo │ └── blonde458.jpg └── vectors_ffhq │ ├── age.npy │ ├── bald.npy │ ├── blur.npy │ ├── roll.npy │ ├── yaw.npy │ ├── beard.npy │ ├── gender.npy │ ├── glasses.npy │ ├── noise.npy │ ├── pitch.npy │ ├── pupils.npy │ ├── smile.npy │ ├── squint.npy │ ├── exposure.npy │ ├── eye_open.npy │ ├── eye_ratio.npy │ ├── eye_size.npy │ ├── eyes_open.npy │ ├── feminine.npy │ ├── hair_gray.npy │ ├── hair_red.npy │ ├── headwear.npy │ ├── lip_ratio.npy │ ├── moustache.npy │ ├── nose_size.npy │ ├── nose_tip.npy │ ├── sideburns.npy │ ├── eye_distance.npy │ ├── eye_makeup.npy │ ├── eye_occluded.npy │ ├── hair_black.npy │ ├── hair_blond.npy │ ├── hair_brown.npy │ ├── hair_other.npy │ ├── jaw_height.npy │ ├── lip_height.npy │ ├── lip_makeup.npy │ ├── mouth_open.npy │ ├── mouth_ratio.npy │ ├── mouth_size.npy │ ├── nose_ratio.npy │ ├── sunglasses.npy │ ├── mouth_occluded.npy │ ├── forehead_height.npy │ ├── forehead_occluded.npy │ ├── nose_mouth_distance.npy │ └── eye_eyebrow_distance.npy ├── requirements.txt ├── .gitattributes ├── README.md ├── models_mix.bat ├── models └── .gitignore ├── train.bat ├── model_convert.bat ├── _out ├── palekh-512-1536x512-3x1.jpg ├── ffhq-1024-1024x1024-2x2-digress.jpg └── ffhq-1024-3072x1024-3x1-digress-1.jpg ├── train_resume.bat ├── models_blend.bat ├── project.bat ├── .gitignore ├── prepare_dataset.bat ├── src ├── training │ ├── __init__.py │ ├── loss_diffaug.py │ ├── DiffAugment_tf.py │ ├── misc.py │ ├── dataset_tool.py │ ├── dataset.py │ └── loss.py ├── dnnlib │ ├── submission │ │ ├── internal │ │ │ ├── __init__.py │ │ │ └── local.py │ │ ├── __init__.py │ │ ├── run_context.py │ │ └── submit.py │ ├── tflib │ │ ├── ops │ │ │ ├── __init__.py │ │ │ ├── fused_bias_act.cu │ │ │ └── fused_bias_act.py │ │ ├── __init__.py │ │ ├── custom_ops.py │ │ ├── autosummary.py │ │ └── tfutil.py │ └── __init__.py ├── models_swa.py ├── project_latent.py ├── util │ ├── multicrop.py │ └── progress_bar.py ├── model_pt2pkl.py ├── models_blend.py ├── _play_dlatents.py ├── _play_vectors.py ├── _genSGAN2.py ├── train.py ├── projector.py └── model_convert.py ├── play_vectors.bat ├── play_dlatents.bat ├── gen.bat ├── data └── multicrop.bat └── LICENSE.txt /train/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /_in/dlats/ffhq-1024-f-1.txt: -------------------------------------------------------------------------------- 1 | 5,8 -------------------------------------------------------------------------------- /_in/dlats/ffhq-1024-f-2.txt: -------------------------------------------------------------------------------- 1 | 11,14 -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy==1.19.5 2 | pyturbojpeg 3 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | # *.pkl filter=lfs diff=lfs merge=lfs -text 2 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/README.md -------------------------------------------------------------------------------- /_in/mask.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/mask.jpg -------------------------------------------------------------------------------- /models_mix.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | python src/models_swa.py --in_dir %1 4 | 5 | -------------------------------------------------------------------------------- /_in/blonde458.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/blonde458.npy -------------------------------------------------------------------------------- /_in/photo/blonde458.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/photo/blonde458.jpg -------------------------------------------------------------------------------- /_in/vectors_ffhq/age.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/age.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/bald.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/bald.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/blur.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/blur.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/roll.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/roll.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/yaw.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/yaw.npy -------------------------------------------------------------------------------- /models/.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore everything in this directory 2 | * 3 | # Except this file 4 | !.gitignore -------------------------------------------------------------------------------- /train.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | python src/train.py --data data/%1 ^ 4 | %2 %3 %4 %5 %6 %7 %8 %9 5 | 6 | -------------------------------------------------------------------------------- /_in/dlats/ffhq-1024-f-1.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/dlats/ffhq-1024-f-1.npy -------------------------------------------------------------------------------- /_in/dlats/ffhq-1024-f-2.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/dlats/ffhq-1024-f-2.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/beard.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/beard.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/gender.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/gender.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/glasses.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/glasses.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/noise.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/noise.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/pitch.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/pitch.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/pupils.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/pupils.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/smile.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/smile.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/squint.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/squint.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/exposure.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/exposure.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/eye_open.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/eye_open.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/eye_ratio.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/eye_ratio.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/eye_size.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/eye_size.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/eyes_open.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/eyes_open.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/feminine.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/feminine.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/hair_gray.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/hair_gray.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/hair_red.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/hair_red.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/headwear.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/headwear.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/lip_ratio.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/lip_ratio.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/moustache.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/moustache.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/nose_size.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/nose_size.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/nose_tip.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/nose_tip.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/sideburns.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/sideburns.npy -------------------------------------------------------------------------------- /model_convert.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | python src/model_convert.py --source %1 %2 %3 %4 %5 %6 %7 %8 %9 4 | 5 | -------------------------------------------------------------------------------- /_in/vectors_ffhq/eye_distance.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/eye_distance.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/eye_makeup.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/eye_makeup.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/eye_occluded.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/eye_occluded.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/hair_black.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/hair_black.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/hair_blond.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/hair_blond.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/hair_brown.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/hair_brown.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/hair_other.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/hair_other.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/jaw_height.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/jaw_height.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/lip_height.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/lip_height.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/lip_makeup.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/lip_makeup.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/mouth_open.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/mouth_open.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/mouth_ratio.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/mouth_ratio.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/mouth_size.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/mouth_size.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/nose_ratio.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/nose_ratio.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/sunglasses.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/sunglasses.npy -------------------------------------------------------------------------------- /_out/palekh-512-1536x512-3x1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_out/palekh-512-1536x512-3x1.jpg -------------------------------------------------------------------------------- /_in/vectors_ffhq/mouth_occluded.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/mouth_occluded.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/forehead_height.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/forehead_height.npy -------------------------------------------------------------------------------- /_in/vectors_ffhq/forehead_occluded.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/forehead_occluded.npy -------------------------------------------------------------------------------- /train_resume.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | python src/train.py --data data/%1 --resume train/%2 ^ 4 | %3 %4 %5 %6 %7 %8 %9 5 | 6 | -------------------------------------------------------------------------------- /_in/vectors_ffhq/nose_mouth_distance.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/nose_mouth_distance.npy -------------------------------------------------------------------------------- /_out/ffhq-1024-1024x1024-2x2-digress.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_out/ffhq-1024-1024x1024-2x2-digress.jpg -------------------------------------------------------------------------------- /_in/vectors_ffhq/eye_eyebrow_distance.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_in/vectors_ffhq/eye_eyebrow_distance.npy -------------------------------------------------------------------------------- /_out/ffhq-1024-3072x1024-3x1-digress-1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/eps696/stylegan2/HEAD/_out/ffhq-1024-3072x1024-3x1-digress-1.jpg -------------------------------------------------------------------------------- /models_blend.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | python src/models_blend.py --out_dir ./ --pkl1 %1 --pkl2 %2 --res %3 --level %4 %5 %6 %7 %8 %9 4 | -------------------------------------------------------------------------------- /project.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | python src/project_latent.py --model=models/%1 --in_dir=_in/%2 --out_dir=_out/proj/%2 ^ 4 | %3 %4 %5 %6 %7 %8 %9 5 | 6 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | _cudacache/ 4 | 5 | # Jupyter Notebook 6 | .ipynb_checkpoints/ 7 | 8 | *.pkl 9 | *.avi 10 | *.mp4 11 | -------------------------------------------------------------------------------- /prepare_dataset.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | if "%1"=="" goto help 3 | 4 | python src/training/dataset_tool.py --data data/%1 ^ 5 | %2 %3 %4 %5 %6 %7 %8 %9 6 | goto end 7 | 8 | :help 9 | echo Usage: prepare_dataset imagedir 10 | 11 | :end 12 | -------------------------------------------------------------------------------- /src/training/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | # empty 8 | -------------------------------------------------------------------------------- /src/dnnlib/submission/internal/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | from . import local 8 | -------------------------------------------------------------------------------- /src/dnnlib/submission/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | from . import run_context 8 | from . import submit 9 | -------------------------------------------------------------------------------- /play_vectors.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | if "%1"=="" goto help 3 | 4 | python src/_play_vectors.py --model models/%1 --base_lat _in/%2 --vector_dir _in/%3 --out_dir _out/%~n1-%~n3 ^ 5 | %4 %5 %6 %7 %8 %9 6 | 7 | ffmpeg -y -v warning -i _out\%~n1-%~n3\%%06d.jpg -c:v mjpeg -q:v 2 %~n1-%~n3.avi 8 | rmdir /s /q _out\%~n1-%~n3 9 | 10 | goto end 11 | 12 | :help 13 | echo Usage: play_vectors model latents vector 14 | echo e.g.: play_vectors ffhq-1024-f npy age.npy 15 | 16 | :end -------------------------------------------------------------------------------- /play_dlatents.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | if "%1"=="" goto help 3 | 4 | python src/_play_dlatents.py --model models/%1 --dlatents _in/%2 --out_dir _out/%~n1-%~n2 --fstep %3 --size %4 ^ 5 | %5 %6 %7 %8 %9 6 | 7 | ffmpeg -y -v warning -r 25 -i _out\%~n1-%~n2\%%06d.jpg -c:v mjpeg -q:v 2 _out/%~n1-%~n2.avi 8 | rmdir /s /q _out\%~n1-%~n2 9 | 10 | goto end 11 | 12 | :help 13 | echo Usage: play_dlatents model latents fstep size 14 | echo e.g.: play_dlatents ffhq-1024-f npy 25 1920-1080 15 | 16 | :end -------------------------------------------------------------------------------- /src/dnnlib/tflib/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # NVIDIA CORPORATION and its licensors retain all intellectual property 4 | # and proprietary rights in and to this software, related documentation 5 | # and any modifications thereto. Any use, reproduction, disclosure or 6 | # distribution of this software and related documentation without an express 7 | # license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | # empty 10 | -------------------------------------------------------------------------------- /src/dnnlib/tflib/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | from . import autosummary 8 | from . import network 9 | from . import optimizer 10 | from . import tfutil 11 | from . import custom_ops 12 | 13 | from .tfutil import * 14 | from .network import Network 15 | 16 | from .optimizer import Optimizer 17 | 18 | from .custom_ops import get_plugin 19 | -------------------------------------------------------------------------------- /gen.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | if "%1"=="" goto help 3 | 4 | if "%2"=="" goto test 5 | if "%2"=="1" goto test 6 | 7 | python src/_genSGAN2.py --model models/%1 --out_dir _out/%~n1 --size %2 --frames %3 ^ 8 | %4 %5 %6 %7 %8 %9 9 | goto ff 10 | 11 | :test 12 | python src/_genSGAN2.py --model models/%~n1.pkl --out_dir _out/%~n1 --frames 200-20 ^ 13 | %3 %4 %5 %6 %7 %8 %9 14 | 15 | :ff 16 | ffmpeg -y -v warning -i _out\%~n1\%%06d.jpg -c:v mjpeg -q:v 2 _out/%~n1-%2.avi 17 | rem rmdir /s /q _out\%~n1 18 | 19 | goto end 20 | 21 | 22 | :help 23 | echo Usage: gen model x-y framecount-transit 24 | echo e.g.: gen ffhq-1024 1280-720 100-25 25 | 26 | :end 27 | -------------------------------------------------------------------------------- /data/multicrop.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | 3 | : multicrop images to squares of specific size 4 | : %1 = video file or directory with images 5 | : %2 = size 6 | : %3 = shift step 7 | 8 | if exist %1\* goto dir 9 | 10 | :video 11 | echo .. cropping video 12 | if not exist "%~dp1\%~n1-tmp" md "%~dp1\%~n1-tmp" 13 | ffmpeg -y -v error -i %1 -q:v 2 "%~dp1\%~n1-tmp\%~n1-c-%%06d.png" 14 | python ../src/util/multicrop.py --in_dir %~dp1/%~n1-tmp --out_dir %~dp1/%~n1-sub --size %2 --step %3 15 | rmdir /s /q %~dp1\%~n1-tmp 16 | goto end 17 | 18 | :dir 19 | echo .. cropping images 20 | python ../src/util/multicrop.py --in_dir %1 --out_dir %~dp1/%~n1-sub --size %2 --step %3 21 | goto end 22 | 23 | :end -------------------------------------------------------------------------------- /src/dnnlib/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | from . import submission 8 | 9 | from .submission.run_context import RunContext 10 | 11 | from .submission.submit import SubmitTarget 12 | from .submission.submit import PathType 13 | from .submission.submit import SubmitConfig 14 | from .submission.submit import submit_run 15 | from .submission.submit import get_path_from_template 16 | from .submission.submit import convert_path 17 | from .submission.submit import make_run_dir_path 18 | 19 | from .util import EasyDict, make_cache_dir_path 20 | 21 | # submit_config: SubmitConfig = None # Package level variable for SubmitConfig which is only valid when inside the run function. 22 | -------------------------------------------------------------------------------- /src/dnnlib/submission/internal/local.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | class TargetOptions(): 8 | def __init__(self): 9 | self.do_not_copy_source_files = False 10 | 11 | class Target(): 12 | def __init__(self): 13 | pass 14 | 15 | def finalize_submit_config(self, submit_config, host_run_dir): 16 | print (' Local submit ', end='', flush=True) 17 | submit_config.run_dir = host_run_dir 18 | 19 | def submit(self, submit_config, host_run_dir): 20 | from ..submit import run_wrapper, convert_path 21 | # print('- run_dir: %s' % convert_path(submit_config.run_dir), flush=True) 22 | print(':: %s' % convert_path(submit_config.run_dir), flush=True) 23 | return run_wrapper(submit_config) 24 | -------------------------------------------------------------------------------- /src/models_swa.py: -------------------------------------------------------------------------------- 1 | """ 2 | https://github.com/arfafax/StyleGAN2_experiments 3 | Stochastic Weight Averaging: https://arxiv.org/abs/1803.05407 4 | See: https://github.com/kristpapadopoulos/keras-stochastic-weight-averaging 5 | """ 6 | import os 7 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 8 | import glob 9 | import pickle 10 | import argparse 11 | 12 | import dnnlib 13 | import dnnlib.tflib as tflib 14 | 15 | parser = argparse.ArgumentParser(description='Perform stochastic weight averaging', formatter_class=argparse.ArgumentDefaultsHelpFormatter) 16 | parser.add_argument('--in_dir', default='models', help='Directory with network checkpoints for weight averaging') 17 | parser.add_argument('--output', default='output-mixed.pkl', help='The averaged model to output') 18 | parser.add_argument('--count', default=None, help='Average the last n checkpoints', type=int) 19 | # args, other_args = parser.parse_known_args() 20 | args = parser.parse_args() 21 | 22 | def fetch_models_from_files(model_list): 23 | for fn in model_list: 24 | with open(fn, 'rb') as f: 25 | yield pickle.load(f) 26 | 27 | def apply_swa_to_checkpoints(models): 28 | next_models = next(models) 29 | try: 30 | gen, dis, gs = next_models 31 | mod_gen, mod_dis, mod_gs = gen, dis, gs 32 | except: 33 | gs = next_models 34 | mod_gs = gs 35 | print('Loading 1 ', end='', flush=True) 36 | epoch = 0 37 | try: 38 | while True: 39 | epoch += 1 40 | print('. ', end='', flush=True) 41 | next_models = next(models) 42 | if next_models is None: 43 | print('') 44 | break 45 | try: 46 | gen, dis, gs = next_models 47 | # mod_gen.apply_swa(gen, epoch) 48 | # mod_dis.apply_swa(dis, epoch) 49 | except: 50 | gs = next_models 51 | mod_gs.apply_swa(gs, epoch) 52 | print('%d '%(epoch+1), end='', flush=True) 53 | except: 54 | print("end") 55 | # return (mod_gen, mod_dis, mod_gs) 56 | return mod_gs 57 | 58 | def main(): 59 | files = glob.glob(os.path.join(args.in_dir, '*.pkl')) 60 | files.sort() 61 | if args.count is not None: 62 | if (len(files) > args.count): 63 | files = files[-args.count:] 64 | print(files) 65 | 66 | tflib.init_tf() 67 | models = fetch_models_from_files(files) 68 | swa_models = apply_swa_to_checkpoints(models) 69 | 70 | with open(args.output, 'wb') as f: 71 | pickle.dump(swa_models, f) 72 | print('Output model with stochastic averaged weights saved as', args.output) 73 | 74 | main() -------------------------------------------------------------------------------- /src/training/loss_diffaug.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # This work is made available under the Nvidia Source Code License-NC. 3 | # https://nvlabs.github.io/stylegan2/license.html 4 | """Loss functions.""" 5 | 6 | import tensorflow as tf 7 | 8 | from dnnlib.tflib.autosummary import autosummary 9 | from training.DiffAugment_tf import DiffAugment 10 | 11 | def ns_DiffAugment_r1(G, D, training_set, minibatch_size, reals, gamma=10, policy='', **kwargs): 12 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 13 | labels = training_set.get_random_labels_tf(minibatch_size) 14 | fakes = G.get_output_for(latents, labels, is_training=True) 15 | 16 | real_scores = D.get_output_for(DiffAugment(reals, policy=policy, channels_first=True), labels, is_training=True) 17 | fake_scores = D.get_output_for(DiffAugment(fakes, policy=policy, channels_first=True), labels, is_training=True) 18 | real_scores = autosummary('Loss/scores/real', real_scores) 19 | fake_scores = autosummary('Loss/scores/fake', fake_scores) 20 | 21 | G_loss = tf.nn.softplus(-fake_scores) 22 | G_loss = autosummary('Loss/G_loss', G_loss) 23 | D_loss = tf.nn.softplus(fake_scores) + tf.nn.softplus(-real_scores) 24 | D_loss = autosummary('Loss/D_loss', D_loss) 25 | 26 | with tf.name_scope('GradientPenalty'): 27 | real_grads = tf.gradients(tf.reduce_sum(real_scores), [reals])[0] 28 | gradient_penalty = tf.reduce_sum(tf.square(real_grads), axis=[1, 2, 3]) 29 | gradient_penalty = autosummary('Loss/gradient_penalty', gradient_penalty) 30 | D_reg = gradient_penalty * (gamma * 0.5) 31 | return G_loss, D_loss, D_reg 32 | 33 | def ns_r1_DiffAugment(G, D, training_set, minibatch_size, reals, gamma=10, policy='', **kwargs): 34 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 35 | labels = training_set.get_random_labels_tf(minibatch_size) 36 | fakes = G.get_output_for(latents, labels, is_training=True) 37 | 38 | reals = DiffAugment(reals, policy=policy, channels_first=True) 39 | fakes = DiffAugment(fakes, policy=policy, channels_first=True) 40 | real_scores = D.get_output_for(reals, labels, is_training=True) 41 | fake_scores = D.get_output_for(fakes, labels, is_training=True) 42 | real_scores = autosummary('Loss/scores/real', real_scores) 43 | fake_scores = autosummary('Loss/scores/fake', fake_scores) 44 | 45 | G_loss = tf.nn.softplus(-fake_scores) 46 | G_loss = autosummary('Loss/G_loss', G_loss) 47 | D_loss = tf.nn.softplus(fake_scores) + tf.nn.softplus(-real_scores) 48 | D_loss = autosummary('Loss/D_loss', D_loss) 49 | 50 | with tf.name_scope('GradientPenalty'): 51 | real_grads = tf.gradients(tf.reduce_sum(real_scores), [reals])[0] 52 | gradient_penalty = tf.reduce_sum(tf.square(real_grads), axis=[1, 2, 3]) 53 | gradient_penalty = autosummary('Loss/gradient_penalty', gradient_penalty) 54 | D_reg = gradient_penalty * (gamma * 0.5) 55 | return G_loss, D_loss, D_reg 56 | 57 | -------------------------------------------------------------------------------- /src/training/DiffAugment_tf.py: -------------------------------------------------------------------------------- 1 | # Differentiable Augmentation for Data-Efficient GAN Training 2 | # Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song Han 3 | # https://arxiv.org/pdf/2006.10738 4 | 5 | import tensorflow; tf = tensorflow.compat.v1 if hasattr(tensorflow.compat, 'v1') else tensorflow 6 | 7 | def DiffAugment(x, policy='', channels_first=False): 8 | if policy: 9 | if channels_first: 10 | x = tf.transpose(x, [0, 2, 3, 1]) 11 | for p in policy.split(','): 12 | for f in AUGMENT_FNS[p]: 13 | x = f(x) 14 | if channels_first: 15 | x = tf.transpose(x, [0, 3, 1, 2]) 16 | return x 17 | 18 | def rand_brightness(x): 19 | magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) - 0.5 20 | x = x + magnitude 21 | return x 22 | 23 | def rand_saturation(x): 24 | magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) * 2 25 | x_mean = tf.reduce_mean(x, axis=3, keepdims=True) 26 | x = (x - x_mean) * magnitude + x_mean 27 | return x 28 | 29 | def rand_contrast(x): 30 | magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) + 0.5 31 | x_mean = tf.reduce_mean(x, axis=[1, 2, 3], keepdims=True) 32 | x = (x - x_mean) * magnitude + x_mean 33 | return x 34 | 35 | def rand_translation(x, ratio=0.125): 36 | batch_size = tf.shape(x)[0] 37 | image_size = tf.shape(x)[1:3] 38 | shift = tf.cast(tf.cast(image_size, tf.float32) * ratio + 0.5, tf.int32) 39 | translation_x = tf.random.uniform([batch_size, 1], -shift[0], shift[0] + 1, dtype=tf.int32) 40 | translation_y = tf.random.uniform([batch_size, 1], -shift[1], shift[1] + 1, dtype=tf.int32) 41 | grid_x = tf.clip_by_value(tf.expand_dims(tf.range(image_size[0], dtype=tf.int32), 0) + translation_x + 1, 0, image_size[0] + 1) 42 | grid_y = tf.clip_by_value(tf.expand_dims(tf.range(image_size[1], dtype=tf.int32), 0) + translation_y + 1, 0, image_size[1] + 1) 43 | x = tf.gather_nd(tf.pad(x, [[0, 0], [1, 1], [0, 0], [0, 0]]), tf.expand_dims(grid_x, -1), batch_dims=1) 44 | x = tf.transpose(tf.gather_nd(tf.pad(tf.transpose(x, [0, 2, 1, 3]), [[0, 0], [1, 1], [0, 0], [0, 0]]), tf.expand_dims(grid_y, -1), batch_dims=1), [0, 2, 1, 3]) 45 | return x 46 | 47 | def rand_cutout(x, ratio=0.5): 48 | batch_size = tf.shape(x)[0] 49 | image_size = tf.shape(x)[1:3] 50 | cutout_size = tf.cast(tf.cast(image_size, tf.float32) * ratio + 0.5, tf.int32) 51 | offset_x = tf.random.uniform([tf.shape(x)[0], 1, 1], maxval=image_size[0] + (1 - cutout_size[0] % 2), dtype=tf.int32) 52 | offset_y = tf.random.uniform([tf.shape(x)[0], 1, 1], maxval=image_size[1] + (1 - cutout_size[1] % 2), dtype=tf.int32) 53 | grid_batch, grid_x, grid_y = tf.meshgrid(tf.range(batch_size, dtype=tf.int32), tf.range(cutout_size[0], dtype=tf.int32), tf.range(cutout_size[1], dtype=tf.int32), indexing='ij') 54 | cutout_grid = tf.stack([grid_batch, grid_x + offset_x - cutout_size[0] // 2, grid_y + offset_y - cutout_size[1] // 2], axis=-1) 55 | mask_shape = tf.stack([batch_size, image_size[0], image_size[1]]) 56 | cutout_grid = tf.maximum(cutout_grid, 0) 57 | cutout_grid = tf.minimum(cutout_grid, tf.reshape(mask_shape - 1, [1, 1, 1, 3])) 58 | mask = tf.maximum(1 - tf.scatter_nd(cutout_grid, tf.ones([batch_size, cutout_size[0], cutout_size[1]], dtype=tf.float32), mask_shape), 0) 59 | x = x * tf.expand_dims(mask, axis=3) 60 | return x 61 | 62 | AUGMENT_FNS = { 63 | 'color': [rand_brightness, rand_saturation, rand_contrast], 64 | 'translation': [rand_translation], 65 | 'cutout': [rand_cutout], 66 | } 67 | 68 | -------------------------------------------------------------------------------- /src/project_latent.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # This work is made available under the Nvidia Source Code License-NC. 3 | # https://nvlabs.github.io/stylegan2/license.html 4 | 5 | import os 6 | import os.path as osp 7 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 8 | import time 9 | import argparse 10 | import numpy as np 11 | import dnnlib 12 | import dnnlib.tflib as tflib 13 | import re 14 | import sys 15 | import cv2, PIL 16 | import pickle 17 | 18 | import projector 19 | from training import misc 20 | 21 | from util.utilgan import img_list, img_read, basename 22 | try: # progress bar for notebooks 23 | get_ipython().__class__.__name__ 24 | from util.progress_bar import ProgressIPy as ProgressBar 25 | except: # normal console 26 | from util.progress_bar import ProgressBar 27 | 28 | parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter) 29 | parser.add_argument('--model', help='Network pickle filename', required=True) 30 | parser.add_argument('--in_dir', help='Dataset root directory', required=True) 31 | parser.add_argument('--out_dir', help='Root directory for run results (default: %(default)s)', default='_out', metavar='DIR') 32 | parser.add_argument('--steps', type=int, default=1000, help='Number of iterations (default: %(default)s)') # 1000 33 | parser.add_argument('--num_snapshots', type=int, default=10, help='Number of snapshots (default: %(default)s)') 34 | a = parser.parse_args() 35 | 36 | def write_video_frame(proj, video_out): 37 | img = proj.get_images()[0] 38 | img = misc.convert_to_pil_image(img, drange=[-1, 1]) 39 | video_frame = img # .resize((512, 512)) 40 | video_out.write(cv2.cvtColor(np.array(video_frame).astype('uint8'), cv2.COLOR_RGB2BGR)) 41 | 42 | def project_image(proj, targets, work_dir, resolution, num_snapshots): 43 | filename = osp.join(work_dir, basename(work_dir)) 44 | video_out = cv2.VideoWriter(filename + '.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 25, resolution) 45 | 46 | snapshot_steps = set(proj.num_steps - np.linspace(0, proj.num_steps, num_snapshots, endpoint=False, dtype=int)) 47 | misc.save_image_grid(targets, filename + '.jpg', drange=[-1,1]) 48 | proj.start(targets) 49 | pbar = ProgressBar(proj.num_steps) 50 | while proj.get_cur_step() < proj.num_steps: 51 | proj.step() 52 | write_video_frame(proj, video_out) 53 | if proj.get_cur_step() in snapshot_steps: 54 | misc.save_image_grid(proj.get_images(), filename + '-%04d.jpg' % proj.get_cur_step(), drange=[-1,1]) 55 | pbar.upd() 56 | 57 | dlats = proj.get_dlatents() 58 | np.save(filename + '-%04d.npy' % proj.get_cur_step(), dlats) 59 | video_out.release() 60 | 61 | def main(): 62 | print('Loading networks from "%s"...' % a.model) 63 | sess = tflib.init_tf() 64 | with open(a.model, 'rb') as file: 65 | network = pickle.load(file, encoding='latin1') 66 | try: _, _, Gs = network 67 | except: Gs = network 68 | resolution = tuple(Gs.output_shape[2:]) 69 | proj = projector.Projector(a.steps) 70 | proj.set_network(Gs) 71 | 72 | img_files = img_list(a.in_dir) 73 | num_images = len(img_files) 74 | for image_idx in range(num_images): 75 | print('Projecting image %d/%d ...' % (image_idx+1, num_images)) 76 | images = img_read(img_files[image_idx]) 77 | images = np.expand_dims(np.transpose(images, [2,0,1]), 0) 78 | images = misc.adjust_dynamic_range(images, [0, 255], [-1, 1]) 79 | work_dir = osp.join(a.out_dir, basename(img_files[image_idx])) 80 | os.makedirs(work_dir, exist_ok=True) 81 | project_image(proj, images, work_dir, resolution, a.num_snapshots) 82 | 83 | 84 | if __name__ == "__main__": 85 | main() 86 | -------------------------------------------------------------------------------- /src/util/multicrop.py: -------------------------------------------------------------------------------- 1 | import os 2 | import warnings 3 | warnings.filterwarnings("ignore") 4 | import sys 5 | import time 6 | import argparse 7 | from multiprocessing import Pool 8 | 9 | import numpy as np 10 | import cv2 11 | 12 | from utilgan import img_list, basename 13 | try: # progress bar for notebooks 14 | get_ipython().__class__.__name__ 15 | from progress_bar import ProgressIPy as ProgressBar 16 | except: # normal console 17 | from progress_bar import ProgressBar 18 | 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument('-i', '--in_dir', help='Input directory') 21 | parser.add_argument('-o', '--out_dir', help='Output directory') 22 | parser.add_argument('-s', '--size', type=int, default=512, help='Output size in pixels') 23 | parser.add_argument('--step', type=int, default=None, help='Step to shift between crops') 24 | parser.add_argument('--workers', type=int, default=4, help='number of workers') 25 | parser.add_argument('--png_compression', type=int, default=1, help='png compression (0 to 9; 0 = uncompressed, fast)') 26 | parser.add_argument('--jpg_quality', type=int, default=95, help='jpeg quality (0 to 100; 95 = max reasonable)') 27 | parser.add_argument('-d', '--down', action='store_true', help='Downscale before crop? (smaller side to size)') 28 | parser.add_argument('--ext', default=None, help='Override output format') 29 | a = parser.parse_args() 30 | 31 | # https://pillow.readthedocs.io/en/3.0.x/handbook/image-file-formats.html#jpeg 32 | # image quality = from 1 (worst) to 95 (best); default 75. Values above 95 should be avoided; 33 | # 100 disables portions of the JPEG compression algorithm => results in large files with hardly any gain in image quality. 34 | 35 | # CV_IMWRITE_PNG_COMPRESSION from 0 to 9. A higher value means a smaller size and longer 36 | # compression time. If read raw images during training, use 0 for faster IO speed. 37 | 38 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 39 | 40 | def crop_step(path, out_dir, out_size, step, min_step): 41 | img_name = basename(path) 42 | img = cv2.imread(path, cv2.IMREAD_UNCHANGED) 43 | 44 | # convert monochrome to RGB if needed 45 | if len(img.shape) == 2: 46 | img = img[:,:,np.newaxis] 47 | if img.shape[2] == 1: 48 | img = img[:, :, (0,0,0)] 49 | h, w, c = img.shape 50 | 51 | ext = a.ext if a.ext is not None else 'png' if img.shape[2]==4 else 'jpg' 52 | compr = [cv2.IMWRITE_PNG_COMPRESSION, a.png_compression] if ext=='png' else [cv2.IMWRITE_JPEG_QUALITY, a.jpg_quality] 53 | 54 | min_size = min(h,w) 55 | if min_size < out_size: 56 | h = int(h * out_size/min_size) 57 | w = int(w * out_size/min_size) 58 | img = cv2.resize(img, (w,h), interpolation = cv2.INTER_AREA) 59 | elif min_size > out_size and a.down is True: 60 | h = int(h * out_size/min_size) 61 | w = int(w * out_size/min_size) 62 | img = cv2.resize(img, (w,h), interpolation = cv2.INTER_AREA) 63 | 64 | h_space = np.arange(0, h - out_size + 1, step) 65 | if h - (h_space[-1] + out_size) < min_step: 66 | h_space = h_space[:-1] 67 | h_space = np.append(h_space, h - out_size) 68 | w_space = np.arange(0, w - out_size + 1, step) 69 | if w - (w_space[-1] + out_size) < min_step: 70 | w_space = w_space[:-1] 71 | w_space = np.append(w_space, w - out_size) 72 | 73 | index = 0 74 | for x in h_space: 75 | for y in w_space: 76 | index += 1 77 | crop_img = img[x:x + out_size, y:y + out_size, :] 78 | crop_img = np.ascontiguousarray(crop_img) 79 | cv2.imwrite(os.path.join(out_dir, '%s-s%03d.%s' % (img_name, index, ext)), crop_img, compr) 80 | return 'Processing {:s} ...'.format(img_name) 81 | 82 | def main(): 83 | """A multi-thread tool to crop sub images.""" 84 | os.makedirs(a.out_dir, exist_ok=True) 85 | images = img_list(a.in_dir, subdir=True) 86 | 87 | step = a.size // 2 if a.step is None else a.step 88 | min_step = step // 4 89 | 90 | def update(arg): 91 | pbar.upd(arg) 92 | 93 | pbar = ProgressBar(len(images)) 94 | pool = Pool(a.workers) 95 | for path in images: 96 | pool.apply_async(crop_step, args=(path, a.out_dir, a.size, step, min_step), callback=update) 97 | pool.close() 98 | pool.join() 99 | print('All done') 100 | 101 | 102 | if __name__ == '__main__': 103 | # workaround for multithreading in jupyter console 104 | __spec__ = "ModuleSpec(name='builtins', loader=)" 105 | main() 106 | -------------------------------------------------------------------------------- /src/model_pt2pkl.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 3 | import warnings 4 | warnings.filterwarnings("ignore") 5 | import sys 6 | import math 7 | import argparse 8 | import numpy as np 9 | import pickle 10 | 11 | import torch 12 | from torchvision import utils 13 | 14 | # sys.path.append(args.dnnlib) 15 | import dnnlib 16 | from dnnlib import tflib 17 | 18 | try: # progress bar for notebooks 19 | get_ipython().__class__.__name__ 20 | from util.progress_bar import ProgressIPy as ProgressBar 21 | except: # normal console 22 | from util.progress_bar import ProgressBar 23 | 24 | parser = argparse.ArgumentParser(description="Rosinality (pytorch) to Nvidia (pkl) checkpoint converter") 25 | parser.add_argument("--model_pkl", metavar="PATH", help="path to the source tensorflow weights") 26 | parser.add_argument("--model_pt", metavar="PATH", help="path to the updated pytorch weights") 27 | args = parser.parse_args() 28 | 29 | def save_pkl(networks, filepath): 30 | with open(filepath, 'wb') as file: 31 | pickle.dump(networks, file, protocol=pickle.HIGHEST_PROTOCOL) 32 | 33 | def update(network, name, value): 34 | network.set_var(name, value.cpu().numpy()) 35 | 36 | def convert_modconv(tgt_net, src_dict, target_name, source_name, flip=False): 37 | conv_weight = src_dict[source_name + ".conv.weight"].squeeze(0).permute(2,3,1,0) 38 | if flip: 39 | conv_weight = torch.flip(conv_weight, [0,1]) 40 | update(tgt_net, target_name + "/weight", conv_weight) 41 | update(tgt_net, target_name + "/mod_weight", src_dict[source_name + ".conv.modulation.weight"].permute(1,0)) 42 | update(tgt_net, target_name + "/mod_bias", src_dict[source_name + ".conv.modulation.bias"] - 1.) 43 | update(tgt_net, target_name + "/noise_strength", src_dict[source_name + ".noise.weight"].squeeze()) 44 | update(tgt_net, target_name + "/bias", src_dict[source_name + ".activate.bias"].squeeze()) 45 | 46 | def convert_torgb(tgt_net, src_dict, target_name, source_name): 47 | update(tgt_net, target_name + "/weight", src_dict[source_name + ".conv.weight"].squeeze(0).permute(2,3,1,0)) 48 | update(tgt_net, target_name + "/mod_weight", src_dict[source_name + ".conv.modulation.weight"].permute(1,0)) 49 | update(tgt_net, target_name + "/mod_bias", src_dict[source_name + ".conv.modulation.bias"] - 1.) 50 | update(tgt_net, target_name + "/bias", src_dict[source_name + ".bias"].squeeze()) 51 | 52 | def convert_dense(tgt_net, src_dict, target_name, source_name): 53 | update(tgt_net, target_name + "/weight", src_dict[source_name + ".weight"].permute(1,0)) 54 | update(tgt_net, target_name + "/bias", src_dict[source_name + ".bias"]) 55 | 56 | def update_G(src_dict, tgt_net, size, n_mlp): 57 | log_size = int(math.log(size, 2)) 58 | 59 | pbar = ProgressBar(n_mlp + log_size-2 + log_size-2 + (log_size-2)*2+1 + 2) 60 | for i in range(n_mlp): 61 | convert_dense(tgt_net, src_dict, f"G_mapping/Dense{i}", f"style.{i+1}") 62 | pbar.upd() 63 | update(tgt_net, "G_synthesis/4x4/Const/const", src_dict["input.input"]) 64 | convert_torgb(tgt_net, src_dict, "G_synthesis/4x4/ToRGB", "to_rgb1") 65 | pbar.upd() 66 | 67 | for i in range(log_size - 2): 68 | reso = 4 * 2 ** (i + 1) 69 | convert_torgb(tgt_net, src_dict, f"G_synthesis/{reso}x{reso}/ToRGB", f"to_rgbs.{i}") 70 | pbar.upd() 71 | convert_modconv(tgt_net, src_dict, "G_synthesis/4x4/Conv", "conv1") 72 | pbar.upd() 73 | 74 | conv_i = 0 75 | for i in range(log_size - 2): 76 | reso = 4 * 2 ** (i + 1) 77 | convert_modconv(tgt_net, src_dict, f"G_synthesis/{reso}x{reso}/Conv0_up", f"convs.{conv_i}", flip=True) 78 | convert_modconv(tgt_net, src_dict, f"G_synthesis/{reso}x{reso}/Conv1", f"convs.{conv_i + 1}") 79 | conv_i += 2 80 | pbar.upd() 81 | 82 | for i in range(0, (log_size - 2) * 2 + 1): 83 | update(tgt_net, f"G_synthesis/noise{i}", src_dict[f"noises.noise_{i}"]) 84 | pbar.upd() 85 | 86 | 87 | if __name__ == "__main__": 88 | tflib.init_tf() 89 | 90 | with open(args.model_pkl, "rb") as f: 91 | nets = pickle.load(f) 92 | try: 93 | G_in, D_in, Gs = nets 94 | except: 95 | Gs = nets 96 | 97 | src_dict = torch.load(args.model_pt) 98 | 99 | mapping_layers = Gs.components.mapping.list_layers() 100 | n_mlp = len([l for l in mapping_layers if l[0].startswith('Dense')]) 101 | size = Gs.output_shape[2] # 1024 102 | 103 | update_G(src_dict['g_ema'], Gs, size, n_mlp) 104 | 105 | out_name = args.model_pt.replace('.pt', '.pkl') 106 | save_pkl(Gs, out_name) 107 | 108 | -------------------------------------------------------------------------------- /src/dnnlib/submission/run_context.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | """Helpers for managing the run/training loop.""" 8 | 9 | import datetime 10 | import json 11 | import os 12 | import pprint 13 | import time 14 | import types 15 | 16 | from typing import Any 17 | 18 | from . import submit 19 | 20 | # Singleton RunContext 21 | _run_context = None 22 | 23 | class RunContext(object): 24 | """Helper class for managing the run/training loop. 25 | 26 | The context will hide the implementation details of a basic run/training loop. 27 | It will set things up properly, tell if run should be stopped, and then cleans up. 28 | User should call update periodically and use should_stop to determine if run should be stopped. 29 | 30 | Args: 31 | submit_config: The SubmitConfig that is used for the current run. 32 | config_module: (deprecated) The whole config module that is used for the current run. 33 | """ 34 | 35 | def __init__(self, submit_config: submit.SubmitConfig, config_module: types.ModuleType = None): 36 | global _run_context 37 | # Only a single RunContext can be alive 38 | assert _run_context is None 39 | _run_context = self 40 | self.submit_config = submit_config 41 | self.should_stop_flag = False 42 | self.has_closed = False 43 | self.start_time = time.time() 44 | self.last_update_time = time.time() 45 | self.last_update_interval = 0.0 46 | self.progress_monitor_file_path = None 47 | 48 | # vestigial config_module support just prints a warning 49 | if config_module is not None: 50 | print("RunContext.config_module parameter support has been removed.") 51 | 52 | # write out details about the run to a text file 53 | self.run_txt_data = {"task_name": submit_config.task_name, "host_name": submit_config.host_name, "start_time": datetime.datetime.now().isoformat(sep=" ")} 54 | with open(os.path.join(submit_config.run_dir, "run.txt"), "w") as f: 55 | pprint.pprint(self.run_txt_data, stream=f, indent=4, width=200, compact=False) 56 | 57 | def __enter__(self) -> "RunContext": 58 | return self 59 | 60 | def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: 61 | self.close() 62 | 63 | def update(self, loss: Any = 0, cur_epoch: Any = 0, max_epoch: Any = None) -> None: 64 | """Do general housekeeping and keep the state of the context up-to-date. 65 | Should be called often enough but not in a tight loop.""" 66 | assert not self.has_closed 67 | 68 | self.last_update_interval = time.time() - self.last_update_time 69 | self.last_update_time = time.time() 70 | 71 | if os.path.exists(os.path.join(self.submit_config.run_dir, "abort.txt")): 72 | self.should_stop_flag = True 73 | 74 | def should_stop(self) -> bool: 75 | """Tell whether a stopping condition has been triggered one way or another.""" 76 | return self.should_stop_flag 77 | 78 | def get_time_since_start(self) -> float: 79 | """How much time has passed since the creation of the context.""" 80 | return time.time() - self.start_time 81 | 82 | def get_time_since_last_update(self) -> float: 83 | """How much time has passed since the last call to update.""" 84 | return time.time() - self.last_update_time 85 | 86 | def get_last_update_interval(self) -> float: 87 | """How much time passed between the previous two calls to update.""" 88 | return self.last_update_interval 89 | 90 | def close(self) -> None: 91 | """Close the context and clean up. 92 | Should only be called once.""" 93 | if not self.has_closed: 94 | # update the run.txt with stopping time 95 | self.run_txt_data["stop_time"] = datetime.datetime.now().isoformat(sep=" ") 96 | with open(os.path.join(self.submit_config.run_dir, "run.txt"), "w") as f: 97 | pprint.pprint(self.run_txt_data, stream=f, indent=4, width=200, compact=False) 98 | self.has_closed = True 99 | 100 | # detach the global singleton 101 | global _run_context 102 | if _run_context is self: 103 | _run_context = None 104 | 105 | @staticmethod 106 | def get(): 107 | import dnnlib 108 | if _run_context is not None: 109 | return _run_context 110 | return RunContext(dnnlib.submit_config) 111 | -------------------------------------------------------------------------------- /src/models_blend.py: -------------------------------------------------------------------------------- 1 | """ 2 | originally from https://github.com/justinpinkney/stylegan2/blob/master/blend_models.py 3 | """ 4 | 5 | import os 6 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 7 | import warnings 8 | warnings.filterwarnings("ignore") 9 | import glob 10 | import argparse 11 | import math 12 | import numpy as np 13 | import pickle 14 | from imageio import imsave 15 | 16 | import tensorflow; tf = tensorflow.compat.v1 if hasattr(tensorflow.compat, 'v1') else tensorflow 17 | tf.logging.set_verbosity(tf.logging.ERROR) 18 | 19 | import dnnlib 20 | import dnnlib.tflib as tflib 21 | from dnnlib.tflib import tfutil 22 | 23 | from util.utilgan import basename 24 | 25 | parser = argparse.ArgumentParser() 26 | parser.add_argument('--out_dir', default='./', help='Output directory') 27 | parser.add_argument('--pkl1', required=True, help='PKL for low res layers') 28 | parser.add_argument('--pkl2', required=True, help='PKL for hi res layers') 29 | parser.add_argument('--res', type=int, default=64, help='Resolution level at which to switch between models') 30 | parser.add_argument('--level', type=int, default=0, help='Switch at Conv block 0 or 1?') 31 | parser.add_argument('--blend_width', type=float, default=None, help='None = hard switch, float = smooth switch (logistic) with given width') 32 | parser.add_argument('-v', '--verbose', action='store_true', help='Print out blended layers') 33 | a = parser.parse_args() 34 | 35 | def load_pkl(filepath): 36 | with open(filepath, 'rb') as f: 37 | nets = pickle.load(f, encoding='latin1') 38 | return nets 39 | 40 | def save_pkl(networks, filepath): 41 | with open(filepath, 'wb') as file: 42 | pickle.dump(networks, file, protocol=pickle.HIGHEST_PROTOCOL) 43 | 44 | # list of (name, resolution, level, position) 45 | def extract_conv_names(model, type='G'): 46 | model_names = list(model.trainables.keys()) 47 | conv_names = [] 48 | resolutions = [4*2**x for x in range(9)] 49 | 50 | if type=='G': 51 | level_names = [["Conv0_up", "Const"], ["Conv1", "ToRGB"]] 52 | var_names = "G_synthesis/{}x{}/" 53 | else: # D 54 | level_names = [["Conv1_down", "Skip"], ["Conv0", "FromRGB"]] 55 | var_names = "{}x{}/" 56 | model_names = model_names[::-1] 57 | 58 | position = 0 59 | for res in resolutions: 60 | root_name = var_names.format(res, res) 61 | for level, level_suffixes in enumerate(level_names): 62 | for suffix in level_suffixes: 63 | search_name = root_name + suffix 64 | matched_names = [x for x in model_names if x.startswith(search_name)] 65 | to_add = [(name, "{}x{}".format(res, res), level, position) for name in matched_names] 66 | conv_names.extend(to_add) 67 | position += 1 68 | 69 | return conv_names 70 | 71 | def blend_layers(Net_lo, Net_hi, type='G'): 72 | print(' blending', type) 73 | resolution = "{}x{}".format(a.res, a.res) 74 | 75 | model_1_names = extract_conv_names(Net_lo, type) 76 | model_2_names = extract_conv_names(Net_hi, type) 77 | assert all((x == y for x, y in zip(model_1_names, model_2_names))) 78 | 79 | Net_out = Net_lo.clone() 80 | 81 | short_names = [(x[1:3]) for x in model_1_names] 82 | full_names = [(x[0]) for x in model_1_names] 83 | mid_point_idx = short_names.index((resolution, a.level)) 84 | mid_point_pos = model_1_names[mid_point_idx][3] 85 | print(' boundary ::', mid_point_idx, mid_point_pos, model_1_names[mid_point_idx]) 86 | 87 | ys = [] 88 | for name, resolution, level, position in model_1_names: 89 | # print(name, resolution, level, position) 90 | # add small x offset for smoother blend animations ? 91 | x = position - mid_point_pos 92 | if a.blend_width is not None: 93 | exponent = -x / a.blend_width 94 | y = 1 / (1 + math.exp(exponent)) 95 | else: 96 | y = 1 if x > 1 else 0 97 | ys.append(y) 98 | if a.verbose and y > 0: 99 | print(" .. {} *{}".format(name, y)) 100 | 101 | tfutil.set_vars(tfutil.run({ 102 | Net_out.vars[name]: (Net_hi.vars[name] * y + Net_lo.vars[name] * (1-y)) 103 | for name, y in zip(full_names, ys)} )) 104 | return Net_out 105 | 106 | def main(): 107 | os.makedirs(a.out_dir, exist_ok=True) 108 | 109 | tflib.init_tf() 110 | with tf.Session() as sess, tf.device('/gpu:0'): 111 | Net_lo = load_pkl(a.pkl1) 112 | Net_hi = load_pkl(a.pkl2) 113 | 114 | try: # full model 115 | G_lo, D_lo, Gs_lo = Net_lo 116 | G_hi, D_hi, Gs_hi = Net_hi 117 | G_out = blend_layers(G_lo, G_hi) 118 | Gs_out = blend_layers(Gs_lo, Gs_hi) 119 | D_out = blend_layers(D_lo, D_hi, type='D') 120 | Net_out = G_out, D_out, Gs_out 121 | except: # Gs only 122 | Gs_out = blend_layers(Net_lo, Net_hi) # only Gs 123 | Net_out = Gs_out 124 | 125 | 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)) 126 | save_pkl(Net_out, '%s.pkl' % out_name) 127 | 128 | rnd = np.random.RandomState(696) 129 | grid_latents = rnd.randn(4, *Gs_out.input_shape[1:]) 130 | grid_fakes = Gs_out.run(grid_latents, [None], is_validation=True, minibatch_size=1) 131 | grid_fakes = np.hstack(np.transpose(grid_fakes, [0,2,3,1])) 132 | imsave('%s.jpg' % out_name, ((grid_fakes+1)*127.5).astype(np.uint8)) 133 | 134 | print('\n All done') 135 | 136 | 137 | if __name__ == '__main__': 138 | main() 139 | 140 | -------------------------------------------------------------------------------- /src/util/progress_bar.py: -------------------------------------------------------------------------------- 1 | """ 2 | from progress_bar import ProgressBar 3 | 4 | pbar = ProgressBar(steps) 5 | pbar.upd() 6 | """ 7 | 8 | import os 9 | import sys 10 | import math 11 | os.system('') #enable VT100 Escape Sequence for WINDOWS 10 Ver. 1607 12 | 13 | from shutil import get_terminal_size 14 | import time 15 | 16 | import ipywidgets as ipy 17 | import IPython 18 | class ProgressIPy(object): 19 | def __init__(self, task_num=10): 20 | self.pbar = ipy.IntProgress(min=0, max=task_num, bar_style='') # (value=0, min=0, max=max, step=1, description=description, bar_style='') 21 | self.labl = ipy.Label() 22 | IPython.display.display(ipy.HBox([self.pbar, self.labl])) 23 | self.task_num = task_num 24 | self.completed = 0 25 | self.start() 26 | 27 | def start(self, task_num=None): 28 | if task_num is not None: 29 | self.task_num = task_num 30 | if self.task_num > 0: 31 | self.labl.value = '0/{}'.format(self.task_num) 32 | else: 33 | self.labl.value = 'completed: 0, elapsed: 0s' 34 | self.start_time = time.time() 35 | 36 | def upd(self, *p, **kw): 37 | self.completed += 1 38 | elapsed = time.time() - self.start_time + 0.0000000000001 39 | fps = self.completed / elapsed if elapsed>0 else 0 40 | if self.task_num > 0: 41 | finaltime = time.asctime(time.localtime(self.start_time + self.task_num * elapsed / float(self.completed))) 42 | fin = ' end %s' % finaltime[11:16] 43 | percentage = self.completed / float(self.task_num) 44 | eta = int(elapsed * (1 - percentage) / percentage + 0.5) 45 | self.labl.value = '{}/{}, rate {:.3g}s, time {}s, left {}s, {}'.format(self.completed, self.task_num, 1./fps, shortime(elapsed), shortime(eta), fin) 46 | else: 47 | self.labl.value = 'completed {}, time {}s, {:.1f} steps/s'.format(self.completed, int(elapsed + 0.5), fps) 48 | self.pbar.value += 1 49 | if self.completed == self.task_num: self.pbar.bar_style = 'success' 50 | return self.completed 51 | 52 | 53 | class ProgressBar(object): 54 | '''A progress bar which can print the progress 55 | modified from https://github.com/hellock/cvbase/blob/master/cvbase/progress.py 56 | ''' 57 | def __init__(self, task_num=0, bar_width=50, start=True): 58 | self.task_num = task_num 59 | max_bar_width = self._get_max_bar_width() 60 | self.bar_width = (bar_width if bar_width <= max_bar_width else max_bar_width) 61 | self.completed = 0 62 | if start: 63 | self.start() 64 | 65 | def _get_max_bar_width(self): 66 | terminal_width, _ = get_terminal_size() 67 | max_bar_width = min(int(terminal_width * 0.6), terminal_width - 50) 68 | if max_bar_width < 10: 69 | print('terminal is small ({}), make it bigger for proper visualization'.format(terminal_width)) 70 | max_bar_width = 10 71 | return max_bar_width 72 | 73 | def start(self, task_num=None): 74 | if task_num is not None: 75 | self.task_num = task_num 76 | if self.task_num > 0: 77 | sys.stdout.write('[{}] 0/{}, elapsed: 0s, ETA:\n{}\n'.format(' ' * self.bar_width, self.task_num, 'Start...')) 78 | else: 79 | sys.stdout.write('completed: 0, elapsed: 0s') 80 | sys.stdout.flush() 81 | self.start_time = time.time() 82 | 83 | def upd(self, msg=None): 84 | self.completed += 1 85 | elapsed = time.time() - self.start_time + 0.0000000000001 86 | fps = self.completed / elapsed if elapsed>0 else 0 87 | if self.task_num > 0: 88 | percentage = self.completed / float(self.task_num) 89 | eta = int(elapsed * (1 - percentage) / percentage + 0.5) 90 | finaltime = time.asctime(time.localtime(self.start_time + self.task_num * elapsed / float(self.completed))) 91 | fin_msg = ' %ss left, end %s' % (shortime(eta), finaltime[11:16]) 92 | if msg is not None: fin_msg += ' ' + str(msg) 93 | mark_width = int(self.bar_width * percentage) 94 | bar_chars = 'X' * mark_width + '-' * (self.bar_width - mark_width) # ▒ ▓ █ 95 | sys.stdout.write('\033[2A') # cursor up 2 lines 96 | sys.stdout.write('\033[J') # clean the output (remove extra chars since last display) 97 | try: 98 | sys.stdout.write('[{}] {}/{}, rate {:.3g}s, time {}s, left {}s \n{}\n'.format( 99 | bar_chars, self.completed, self.task_num, 1./fps, shortime(elapsed), shortime(eta), fin_msg)) 100 | except: 101 | sys.stdout.write('[{}] {}/{}, rate {:.3g}s, time {}s, left {}s \n{}\n'.format( 102 | bar_chars, self.completed, self.task_num, 1./fps, shortime(elapsed), shortime(eta), '<< unprintable >>')) 103 | else: 104 | sys.stdout.write('completed {}, time {}s, {:.1f} steps/s'.format(self.completed, int(elapsed + 0.5), fps)) 105 | sys.stdout.flush() 106 | 107 | def reset(self, count=None, newline=False): 108 | self.start_time = time.time() 109 | if count is not None: 110 | self.task_num = count 111 | if newline is True: 112 | sys.stdout.write('\n\n') 113 | 114 | def time_days(sec): 115 | return '%dd %d:%02d:%02d' % (sec/86400, (sec/3600)%24, (sec/60)%60, sec%60) 116 | def time_hrs(sec): 117 | return '%d:%02d:%02d' % (sec/3600, (sec/60)%60, sec%60) 118 | def shortime(sec): 119 | if sec < 60: 120 | time_short = '%d' % (sec) 121 | elif sec < 3600: 122 | time_short = '%d:%02d' % ((sec/60)%60, sec%60) 123 | elif sec < 86400: 124 | time_short = time_hrs(sec) 125 | else: 126 | time_short = time_days(sec) 127 | return time_short 128 | 129 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | --------------------------- LICENSE FOR Diff Augment -------------------------------- 2 | 3 | Copyright (c) 2020, Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song Han 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | * Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | * Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 17 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 18 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 19 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 20 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 21 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 22 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 23 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 24 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | 27 | 28 | --------------------------- LICENSE FOR stylegan2 -------------------------------- 29 | 30 | Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 31 | 32 | 33 | Nvidia Source Code License-NC 34 | 35 | ======================================================================= 36 | 37 | 1. Definitions 38 | 39 | "Licensor" means any person or entity that distributes its Work. 40 | 41 | "Software" means the original work of authorship made available under 42 | this License. 43 | 44 | "Work" means the Software and any additions to or derivative works of 45 | the Software that are made available under this License. 46 | 47 | "Nvidia Processors" means any central processing unit (CPU), graphics 48 | processing unit (GPU), field-programmable gate array (FPGA), 49 | application-specific integrated circuit (ASIC) or any combination 50 | thereof designed, made, sold, or provided by Nvidia or its affiliates. 51 | 52 | The terms "reproduce," "reproduction," "derivative works," and 53 | "distribution" have the meaning as provided under U.S. copyright law; 54 | provided, however, that for the purposes of this License, derivative 55 | works shall not include works that remain separable from, or merely 56 | link (or bind by name) to the interfaces of, the Work. 57 | 58 | Works, including the Software, are "made available" under this License 59 | by including in or with the Work either (a) a copyright notice 60 | referencing the applicability of this License to the Work, or (b) a 61 | copy of this License. 62 | 63 | 2. License Grants 64 | 65 | 2.1 Copyright Grant. Subject to the terms and conditions of this 66 | License, each Licensor grants to you a perpetual, worldwide, 67 | non-exclusive, royalty-free, copyright license to reproduce, 68 | prepare derivative works of, publicly display, publicly perform, 69 | sublicense and distribute its Work and any resulting derivative 70 | works in any form. 71 | 72 | 3. Limitations 73 | 74 | 3.1 Redistribution. You may reproduce or distribute the Work only 75 | if (a) you do so under this License, (b) you include a complete 76 | copy of this License with your distribution, and (c) you retain 77 | without modification any copyright, patent, trademark, or 78 | attribution notices that are present in the Work. 79 | 80 | 3.2 Derivative Works. You may specify that additional or different 81 | terms apply to the use, reproduction, and distribution of your 82 | derivative works of the Work ("Your Terms") only if (a) Your Terms 83 | provide that the use limitation in Section 3.3 applies to your 84 | derivative works, and (b) you identify the specific derivative 85 | works that are subject to Your Terms. Notwithstanding Your Terms, 86 | this License (including the redistribution requirements in Section 87 | 3.1) will continue to apply to the Work itself. 88 | 89 | 3.3 Use Limitation. The Work and any derivative works thereof only 90 | may be used or intended for use non-commercially. The Work or 91 | derivative works thereof may be used or intended for use by Nvidia 92 | or its affiliates commercially or non-commercially. As used herein, 93 | "non-commercially" means for research or evaluation purposes only. 94 | 95 | 3.4 Patent Claims. If you bring or threaten to bring a patent claim 96 | against any Licensor (including any claim, cross-claim or 97 | counterclaim in a lawsuit) to enforce any patents that you allege 98 | are infringed by any Work, then your rights under this License from 99 | such Licensor (including the grants in Sections 2.1 and 2.2) will 100 | terminate immediately. 101 | 102 | 3.5 Trademarks. This License does not grant any rights to use any 103 | Licensor's or its affiliates' names, logos, or trademarks, except 104 | as necessary to reproduce the notices described in this License. 105 | 106 | 3.6 Termination. If you violate any term of this License, then your 107 | rights under this License (including the grants in Sections 2.1 and 108 | 2.2) will terminate immediately. 109 | 110 | 4. Disclaimer of Warranty. 111 | 112 | THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY 113 | KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF 114 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR 115 | NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER 116 | THIS LICENSE. 117 | 118 | 5. Limitation of Liability. 119 | 120 | EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL 121 | THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE 122 | SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, 123 | INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF 124 | OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK 125 | (INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, 126 | LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER 127 | COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF 128 | THE POSSIBILITY OF SUCH DAMAGES. 129 | 130 | ======================================================================= 131 | -------------------------------------------------------------------------------- /src/training/misc.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # This work is made available under the Nvidia Source Code License-NC. 3 | # https://nvlabs.github.io/stylegan2/license.html 4 | """Miscellaneous utility functions.""" 5 | 6 | import os 7 | import glob 8 | import pickle 9 | import numpy as np 10 | import PIL.Image 11 | import dnnlib 12 | 13 | # ---------------------------------------------------------------------------- 14 | # Convenience wrappers for pickle that are able to load data produced by 15 | # older versions of the code, and from external URLs. 16 | 17 | def open_file_or_url(file_or_url): 18 | if dnnlib.util.is_url(file_or_url): 19 | return dnnlib.util.open_url(file_or_url, cache_dir='.stylegan2-cache') 20 | return open(file_or_url, 'rb') 21 | 22 | def load_pkl(file_or_url): 23 | with open_file_or_url(file_or_url) as file: 24 | return pickle.load(file, encoding='latin1') 25 | 26 | def locate_latest_pkl(train_dir): 27 | allpickles = sorted(glob.glob(os.path.join(train_dir, 'snapshot-*.pkl'))) 28 | if len(allpickles) == 0: 29 | return None, 0. 30 | latest_pickle = allpickles[-1] 31 | kimg = int(os.path.splitext(latest_pickle)[0].split('-')[-1]) 32 | return latest_pickle, float(kimg) 33 | 34 | def save_pkl(obj, filename): 35 | with open(filename, 'wb') as file: 36 | pickle.dump(obj, file, protocol=pickle.HIGHEST_PROTOCOL) 37 | 38 | # ---------------------------------------------------------------------------- 39 | # Image utils. 40 | 41 | def adjust_dynamic_range(data, drange_in, drange_out): 42 | if drange_in != drange_out: 43 | scale = (np.float32(drange_out[1]) - np.float32(drange_out[0])) / (np.float32(drange_in[1]) - np.float32(drange_in[0])) 44 | bias = (np.float32(drange_out[0]) - np.float32(drange_in[0]) * scale) 45 | data = data * scale + bias 46 | return data 47 | 48 | def create_image_grid(images, grid_size=None): 49 | assert images.ndim == 3 or images.ndim == 4 50 | num, img_w, img_h = images.shape[0], images.shape[-1], images.shape[-2] 51 | 52 | if grid_size is not None: 53 | grid_w, grid_h = tuple(grid_size) 54 | else: 55 | grid_w = max(int(np.ceil(np.sqrt(num))), 1) 56 | grid_h = max((num - 1) // grid_w + 1, 1) 57 | 58 | grid = np.zeros(list(images.shape[1:-2]) + [grid_h * img_h, grid_w * img_w], dtype=images.dtype) 59 | for idx in range(num): 60 | x = (idx % grid_w) * img_w 61 | y = (idx // grid_w) * img_h 62 | grid[..., y: y + img_h, x: x + img_w] = images[idx] 63 | return grid 64 | 65 | def convert_to_pil_image(image, drange=[0, 1]): 66 | assert image.ndim == 2 or image.ndim == 3 67 | if image.ndim == 3: 68 | if image.shape[0] == 1: 69 | image = image[0] # grayscale CHW => HW 70 | else: 71 | image = image.transpose(1, 2, 0) # CHW -> HWC 72 | 73 | image = adjust_dynamic_range(image, drange, [0, 255]) 74 | image = np.rint(image).clip(0, 255).astype(np.uint8) 75 | if image.ndim == 2: 76 | fmt = 'L' 77 | elif image.shape[2] == 3: 78 | fmt = 'RGB' 79 | elif image.shape[2] == 4: 80 | fmt = 'RGBA' 81 | return PIL.Image.fromarray(image, fmt) 82 | 83 | def save_image_grid(images, filename, drange=[0, 1], grid_size=None): 84 | convert_to_pil_image(create_image_grid(images, grid_size), drange).save(filename) 85 | 86 | def apply_mirror_augment(minibatch): 87 | mask = np.random.rand(minibatch.shape[0]) < 0.5 88 | minibatch = np.array(minibatch) 89 | minibatch[mask] = minibatch[mask, :, :, ::-1] 90 | return minibatch 91 | 92 | def apply_mirror_augment_v(minibatch): 93 | mask = np.random.rand(minibatch.shape[0]) < 0.5 94 | minibatch = np.array(minibatch) 95 | minibatch[mask] = minibatch[mask, :, ::-1, :] 96 | return minibatch 97 | 98 | # Loading data from previous training runs. 99 | def parse_config_for_previous_run(run_dir): 100 | with open(os.path.join(run_dir, 'submit_config.pkl'), 'rb') as f: 101 | data = pickle.load(f) 102 | data = data.get('run_func_kwargs', {}) 103 | return dict(train=data, dataset=data.get('dataset_args', {})) 104 | 105 | # Size and contents of the image snapshot grids that are exported periodically during training. 106 | def setup_snapshot_image_grid(training_set, 107 | size='1080p', # '1080p' = to be viewed on 1080p display, '4k' = to be viewed on 4k display. 108 | layout='random'): # 'random' = grid contents are selected randomly, 'row_per_class' = each row corresponds to one class label. 109 | 110 | # Select size. 111 | gw = 1; gh = 1 112 | if size == '1080p': 113 | gw = np.clip(1920 // training_set.shape[2], 3, 32) 114 | gh = np.clip(1080 // training_set.shape[1], 2, 32) 115 | if size == '4k': 116 | gw = np.clip(3840 // training_set.shape[2], 7, 32) 117 | gh = np.clip(2160 // training_set.shape[1], 4, 32) 118 | if size == '8k': 119 | gw = np.clip(7680 // training_set.shape[2], 7, 32) 120 | gh = np.clip(4320 // training_set.shape[1], 4, 32) 121 | 122 | # Initialize data arrays. 123 | reals = np.zeros([gw * gh] + training_set.shape, dtype=training_set.dtype) 124 | labels = np.zeros([gw * gh, training_set.label_size], dtype=training_set.label_dtype) 125 | 126 | # Random layout. 127 | if layout == 'random': 128 | reals[:], labels[:] = training_set.get_minibatch_np(gw * gh) 129 | 130 | # Class-conditional layouts. 131 | class_layouts = dict(row_per_class=[gw,1], col_per_class=[1,gh], class4x4=[4,4]) 132 | if layout in class_layouts: 133 | bw, bh = class_layouts[layout] 134 | nw = (gw - 1) // bw + 1 135 | nh = (gh - 1) // bh + 1 136 | blocks = [[] for _i in range(nw * nh)] 137 | for _iter in range(1000000): 138 | real, label = training_set.get_minibatch_np(1) 139 | idx = np.argmax(label[0]) 140 | while idx < len(blocks) and len(blocks[idx]) >= bw * bh: 141 | idx += training_set.label_size 142 | if idx < len(blocks): 143 | blocks[idx].append((real, label)) 144 | if all(len(block) >= bw * bh for block in blocks): 145 | break 146 | for i, block in enumerate(blocks): 147 | for j, (real, label) in enumerate(block): 148 | x = (i % nw) * bw + j % bw 149 | y = (i // nw) * bh + j // bw 150 | if x < gw and y < gh: 151 | reals [x + y * gw] = real[0] 152 | labels[x + y * gw] = label[0] 153 | 154 | return (gw, gh), reals, labels 155 | 156 | -------------------------------------------------------------------------------- /src/_play_dlatents.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 3 | import os.path as osp 4 | import argparse 5 | import numpy as np 6 | from imageio import imsave 7 | import pickle 8 | 9 | import dnnlib 10 | import dnnlib.tflib as tflib 11 | 12 | from util.utilgan import latent_anima, load_latents, file_list, basename 13 | try: # progress bar for notebooks 14 | get_ipython().__class__.__name__ 15 | from util.progress_bar import ProgressIPy as ProgressBar 16 | except: # normal console 17 | from util.progress_bar import ProgressBar 18 | 19 | desc = "Customized StyleGAN2 on Tensorflow" 20 | parser = argparse.ArgumentParser(description=desc) 21 | parser.add_argument('--dlatents', default=None, help='Saved dlatent vectors in single *.npy file or directory with such files') 22 | parser.add_argument('--style_dlat', default=None, help='Saved latent vector for hi res (style) features') 23 | parser.add_argument('--out_dir', default='_out', help='Output directory') 24 | parser.add_argument('--model', default='models/ffhq-1024-f.pkl', help='path to checkpoint file') 25 | parser.add_argument('--size', default=None, help='Output resolution') 26 | parser.add_argument('--scale_type', default='pad', help="main types: pad, padside, symm, symmside") 27 | parser.add_argument('--trunc', type=float, default=1, help='Truncation psi 0..1 (lower = stable, higher = various)') 28 | parser.add_argument('--digress', type=float, default=0, help='distortion technique by Aydao (strength of the effect)') 29 | parser.add_argument('--verbose', action='store_true') 30 | parser.add_argument('--ops', default='cuda', help='custom op implementation (cuda or ref)') 31 | # animation 32 | parser.add_argument("--fstep", type=int, default=25, help="Number of frames for smooth interpolation") 33 | parser.add_argument("--cubic", action='store_true', help="Use cubic splines for smoothing") 34 | a = parser.parse_args() 35 | 36 | if a.size is not None: a.size = [int(s) for s in a.size.split('-')][::-1] 37 | 38 | def main(): 39 | os.makedirs(a.out_dir, exist_ok=True) 40 | 41 | # setup generator 42 | fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) 43 | Gs_kwargs = dnnlib.EasyDict() 44 | Gs_kwargs.func_name = 'training.stylegan2_multi.G_main' 45 | Gs_kwargs.verbose = a.verbose 46 | Gs_kwargs.size = a.size 47 | Gs_kwargs.scale_type = a.scale_type 48 | Gs_kwargs.impl = a.ops 49 | 50 | # load model with arguments 51 | sess = tflib.init_tf({'allow_soft_placement':True}) 52 | pkl_name = osp.splitext(a.model)[0] 53 | with open(pkl_name + '.pkl', 'rb') as file: 54 | network = pickle.load(file, encoding='latin1') 55 | try: _, _, network = network 56 | except: pass 57 | for k in list(network.static_kwargs.keys()): 58 | Gs_kwargs[k] = network.static_kwargs[k] 59 | 60 | # reload custom network, if needed 61 | if '.pkl' in a.model.lower(): 62 | print(' .. Gs from pkl ..', basename(a.model)) 63 | Gs = network 64 | else: # reconstruct network 65 | print(' .. Gs custom ..', basename(a.model)) 66 | Gs = tflib.Network('Gs', **Gs_kwargs) 67 | Gs.copy_vars_from(network) 68 | 69 | z_dim = Gs.input_shape[1] 70 | dz_dim = 512 # dlatent_size 71 | try: dl_dim = 2 * (int(np.floor(np.log2(Gs_kwargs.resolution))) - 1) 72 | except: print(' Resave model, no resolution kwarg found!'); exit(1) 73 | dlat_shape = (1, dl_dim, dz_dim) # [1,18,512] 74 | 75 | # read saved latents 76 | if a.dlatents is not None and osp.isfile(a.dlatents): 77 | key_dlatents = load_latents(a.dlatents) 78 | if len(key_dlatents.shape) == 2: key_dlatents = np.expand_dims(key_dlatents, 0) 79 | elif a.dlatents is not None and osp.isdir(a.dlatents): 80 | # if a.dlatents.endswith('/') or a.dlatents.endswith('\\'): a.dlatents = a.dlatents[:-1] 81 | key_dlatents = [] 82 | npy_list = file_list(a.dlatents, 'npy') 83 | for npy in npy_list: 84 | key_dlatent = load_latents(npy) 85 | if len(key_dlatent.shape) == 2: key_dlatent = np.expand_dims(key_dlatent, 0) 86 | key_dlatents.append(key_dlatent) 87 | key_dlatents = np.concatenate(key_dlatents) # [frm,18,512] 88 | else: 89 | print(' No input dlatents found'); exit() 90 | key_dlatents = key_dlatents[:, np.newaxis] # [frm,1,18,512] 91 | print(' key dlatents', key_dlatents.shape) 92 | 93 | # replace higher layers with single (style) latent 94 | if a.style_dlat is not None: 95 | print(' styling with dlatent', a.style_dlat) 96 | style_dlatent = load_latents(a.style_dlat) 97 | while len(style_dlatent.shape) < 4: style_dlatent = np.expand_dims(style_dlatent, 0) 98 | # try replacing 5 by other value, less than dl_dim 99 | key_dlatents[:, :, range(5,dl_dim), :] = style_dlatent[:, :, range(5,dl_dim), :] 100 | 101 | frames = key_dlatents.shape[0] * a.fstep 102 | 103 | dlatents = latent_anima(dlat_shape, frames, a.fstep, key_latents=key_dlatents, cubic=a.cubic, verbose=True) # [frm,1,512] 104 | print(' dlatents', dlatents.shape) 105 | frame_count = dlatents.shape[0] 106 | 107 | # truncation trick 108 | dlatent_avg = Gs.get_var('dlatent_avg') # (512,) 109 | tr_range = range(0,8) 110 | dlatents[:,:,tr_range,:] = dlatent_avg + (dlatents[:,:,tr_range,:] - dlatent_avg) * a.trunc 111 | 112 | # distort image by tweaking initial const layer 113 | if a.digress > 0: 114 | try: latent_size = Gs.static_kwargs['latent_size'] 115 | except: latent_size = 512 # default latent size 116 | try: init_res = Gs.static_kwargs['init_res'] 117 | except: init_res = (4,4) # default initial layer size 118 | dconst = a.digress * latent_anima([1, latent_size, *init_res], frames, a.fstep, cubic=True, verbose=False) 119 | else: 120 | dconst = np.zeros([frame_count, 1, 1, 1, 1]) 121 | 122 | # generate images from latent timeline 123 | pbar = ProgressBar(frame_count) 124 | for i in range(frame_count): 125 | 126 | # generate multi-latent result 127 | if Gs.num_inputs == 2: 128 | output = Gs.components.synthesis.run(dlatents[i], randomize_noise=False, output_transform=fmt, minibatch_size=1) 129 | else: 130 | output = Gs.components.synthesis.run(dlatents[i], [None], dconst[i], randomize_noise=False, output_transform=fmt, minibatch_size=1) 131 | 132 | ext = 'png' if output.shape[3]==4 else 'jpg' 133 | filename = osp.join(a.out_dir, "%06d.%s" % (i,ext)) 134 | imsave(filename, output[0]) 135 | pbar.upd() 136 | 137 | 138 | if __name__ == '__main__': 139 | main() 140 | 141 | -------------------------------------------------------------------------------- /src/_play_vectors.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 3 | import os.path as osp 4 | import argparse 5 | import numpy as np 6 | from imageio import imsave 7 | import pickle 8 | import cv2 9 | 10 | import dnnlib 11 | import dnnlib.tflib as tflib 12 | 13 | from util.utilgan import load_latents, file_list, basename 14 | try: # progress bar for notebooks 15 | get_ipython().__class__.__name__ 16 | from util.progress_bar import ProgressIPy as ProgressBar 17 | except: # normal console 18 | from util.progress_bar import ProgressBar 19 | 20 | desc = "Customized StyleGAN2 on Tensorflow" 21 | parser = argparse.ArgumentParser(description=desc) 22 | parser.add_argument('--vector_dir', default=None, help='Saved latent directions in *.npy format') 23 | parser.add_argument('--base_lat', default=None, help='Saved latent vector as *.npy file') 24 | parser.add_argument('--out_dir', default='_out/ttt', help='Output directory') 25 | parser.add_argument('--model', default='models/ffhq-1024.pkl', help='path to checkpoint file') 26 | parser.add_argument('--size', default=None, help='output resolution, set in X-Y format') 27 | parser.add_argument('--scale_type', default='pad', help="main types: pad, padside, symm, symmside") 28 | parser.add_argument('--trunc', type=float, default=0.8, help='truncation psi 0..1 (lower = stable, higher = various)') 29 | parser.add_argument('--verbose', action='store_true') 30 | parser.add_argument('--ops', default='cuda', help='custom op implementation (cuda or ref)') 31 | # animation 32 | parser.add_argument("--fstep", type=int, default=25, help="Number of frames for interpolation step") 33 | a = parser.parse_args() 34 | 35 | if a.size is not None: a.size = [int(s) for s in a.size.split('-')][::-1] 36 | 37 | def generate_image(latent): 38 | fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) 39 | args = [] if Gs.num_inputs==2 else [[None], [0]] # custom model? 40 | if use_d: 41 | img = Gs.components.synthesis.run(latent, *args, randomize_noise=False, output_transform=fmt)[0] 42 | else: 43 | img = Gs.run(latent, None, *args, truncation_psi=a.trunc, randomize_noise=False, output_transform=fmt)[0] 44 | return img 45 | 46 | def render_latent_dir(latent, direction, coeff): 47 | new_latent = latent + coeff*direction 48 | img = generate_image(new_latent) 49 | return img 50 | 51 | def render_latent_mix(latent1, latent2, coeff): 52 | new_latent = latent1 * (1-coeff) + latent2 * coeff 53 | img = generate_image(new_latent) 54 | return img 55 | 56 | def pingpong(x, delta): 57 | x = (x + delta) % 2 58 | if x > 1: 59 | x = 1 - (x%1) 60 | delta = -delta 61 | return x, delta 62 | 63 | def get_coeffs_dir(lrange, count): 64 | dx = 1 / count 65 | x = -lrange[0] / (lrange[1] - lrange[0]) 66 | xs = [0] 67 | for _ in range(count*2): 68 | x, dx = pingpong(x, dx) 69 | xs.append( x * (lrange[1] - lrange[0]) + lrange[0] ) 70 | return xs 71 | 72 | def make_loop(base_latent, direction, lrange, fcount, start_frame=0): 73 | coeffs = get_coeffs_dir(lrange, fcount//2) 74 | # pbar = ProgressBar(fcount) 75 | for i in range(fcount): 76 | img = render_latent_dir(base_latent, direction, coeffs[i]) 77 | fname1 = os.path.join(a.out_dir, "%06d.jpg" % (i+start_frame)) 78 | if i%2==0 and a.verbose is True: 79 | cv2.imshow('latent', img[:,:,::-1]) 80 | cv2.waitKey(10) 81 | imsave(fname1, img) 82 | # pbar.upd() 83 | 84 | def make_transit(lat1, lat2, fcount, start_frame=0): 85 | # pbar = ProgressBar(fcount) 86 | for i in range(fcount): 87 | img = render_latent_mix(lat1, lat2, i/fcount) 88 | fname = os.path.join(a.out_dir, "%06d.jpg" % (i+start_frame)) 89 | if i%2==0 and a.verbose is True: 90 | cv2.imshow('latent', img[:,:,::-1]) 91 | cv2.waitKey(10) 92 | imsave(fname, img) 93 | # pbar.upd() 94 | 95 | def main(): 96 | if a.vector_dir is not None: 97 | if a.vector_dir.endswith('/') or a.vector_dir.endswith('\\'): a.vector_dir = a.vector_dir[:-1] 98 | os.makedirs(a.out_dir, exist_ok=True) 99 | 100 | global Gs, use_d 101 | 102 | # setup generator 103 | Gs_kwargs = dnnlib.EasyDict() 104 | Gs_kwargs.func_name = 'training.stylegan2_multi.G_main' 105 | Gs_kwargs.verbose = a.verbose 106 | Gs_kwargs.size = a.size 107 | Gs_kwargs.scale_type = a.scale_type 108 | Gs_kwargs.impl = a.ops 109 | 110 | # load model with arguments 111 | sess = tflib.init_tf({'allow_soft_placement':True}) 112 | pkl_name = osp.splitext(a.model)[0] 113 | with open(pkl_name + '.pkl', 'rb') as file: 114 | network = pickle.load(file, encoding='latin1') 115 | try: _, _, network = network 116 | except: pass 117 | for k in list(network.static_kwargs.keys()): 118 | Gs_kwargs[k] = network.static_kwargs[k] 119 | 120 | # reload custom network, if needed 121 | if '.pkl' in a.model.lower(): 122 | print(' .. Gs from pkl ..', basename(a.model)) 123 | Gs = network 124 | else: # reconstruct network 125 | print(' .. Gs custom ..', basename(a.model)) 126 | Gs = tflib.Network('Gs', **Gs_kwargs) 127 | Gs.copy_vars_from(network) 128 | 129 | # load directions 130 | if a.vector_dir is not None: 131 | directions = [] 132 | vector_list = file_list(a.vector_dir, 'npy') 133 | for v in vector_list: 134 | direction = load_latents(v) 135 | if len(direction.shape) == 2: direction = np.expand_dims(direction, 0) 136 | directions.append(direction) 137 | directions = np.concatenate(directions)[:, np.newaxis] # [frm,1,18,512] 138 | else: 139 | print(' No vectors found'); exit() 140 | 141 | if len(direction[0].shape) > 1 and direction[0].shape[0] > 1: 142 | use_d = True 143 | print(' directions', directions.shape, 'using d' if use_d else 'using w') 144 | 145 | # latent direction range 146 | lrange = [-0.5, 0.5] 147 | 148 | # load saved latents 149 | if a.base_lat is not None: 150 | base_latent = load_latents(a.base_lat) 151 | else: 152 | print(' No NPY input given, making random') 153 | z_dim = Gs.input_shape[1] 154 | shape = (1, z_dim) 155 | base_latent = np.random.randn(*shape) 156 | if use_d: 157 | base_latent = Gs.components.mapping.run(base_latent, None) # [frm,18,512] 158 | 159 | pbar = ProgressBar(len(directions)) 160 | for i, direction in enumerate(directions): 161 | make_loop(base_latent, direction, lrange, a.fstep*2, a.fstep*2 * i) 162 | pbar.upd() 163 | 164 | # make_transit(base_lats[i], base_lats[(i+1)%len(base_lats)], n, 2*n*i + n) 165 | 166 | 167 | if __name__ == '__main__': 168 | main() 169 | 170 | -------------------------------------------------------------------------------- /src/dnnlib/tflib/custom_ops.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # NVIDIA CORPORATION and its licensors retain all intellectual property 4 | # and proprietary rights in and to this software, related documentation 5 | # and any modifications thereto. Any use, reproduction, disclosure or 6 | # distribution of this software and related documentation without an express 7 | # license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | """TensorFlow custom ops builder. 10 | """ 11 | 12 | import glob 13 | import os 14 | import re 15 | import uuid 16 | import hashlib 17 | import tempfile 18 | import shutil 19 | import tensorflow as tf 20 | from tensorflow.python.client import device_lib # pylint: disable=no-name-in-module 21 | 22 | from .. import util 23 | 24 | #---------------------------------------------------------------------------- 25 | # Global options. 26 | 27 | cuda_cache_path = os.path.join(os.path.dirname(__file__), '_cudacache') 28 | cuda_cache_version_tag = 'v1' 29 | do_not_hash_included_headers = True # Speed up compilation by assuming that headers included by the CUDA code never change. 30 | verbose = False # Print status messages to stdout. 31 | 32 | #---------------------------------------------------------------------------- 33 | # Internal helper funcs. 34 | 35 | def _find_compiler_bindir(): 36 | hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) 37 | if hostx64_paths != []: 38 | return hostx64_paths[0] 39 | hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) 40 | if hostx64_paths != []: 41 | return hostx64_paths[0] 42 | hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) 43 | if hostx64_paths != []: 44 | return hostx64_paths[0] 45 | vc_bin_dir = 'C:/Program Files (x86)/Microsoft Visual Studio 14.0/vc/bin' 46 | if os.path.isdir(vc_bin_dir): 47 | return vc_bin_dir 48 | return None 49 | 50 | def _get_compute_cap(device): 51 | caps_str = device.physical_device_desc 52 | m = re.search('compute capability: (\\d+).(\\d+)', caps_str) 53 | major = m.group(1) 54 | minor = m.group(2) 55 | return (major, minor) 56 | 57 | def _get_cuda_gpu_arch_string(): 58 | gpus = [x for x in device_lib.list_local_devices() if x.device_type == 'GPU'] 59 | if len(gpus) == 0: 60 | raise RuntimeError('No GPU devices found') 61 | (major, minor) = _get_compute_cap(gpus[0]) 62 | return 'sm_%s%s' % (major, minor) 63 | 64 | def _run_cmd(cmd): 65 | with os.popen(cmd) as pipe: 66 | output = pipe.read() 67 | status = pipe.close() 68 | if status is not None: 69 | raise RuntimeError('NVCC returned an error. See below for full command line and output log:\n\n%s\n\n%s' % (cmd, output)) 70 | 71 | def _prepare_nvcc_cli(opts): 72 | cmd = 'nvcc ' + opts.strip() 73 | cmd += ' --disable-warnings' 74 | cmd += ' --include-path "%s"' % tf.sysconfig.get_include() 75 | cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'protobuf_archive', 'src') 76 | cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'com_google_absl') 77 | cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'eigen_archive') 78 | 79 | compiler_bindir = _find_compiler_bindir() 80 | if compiler_bindir is None: 81 | # Require that _find_compiler_bindir succeeds on Windows. Allow 82 | # nvcc to use whatever is the default on Linux. 83 | if os.name == 'nt': 84 | raise RuntimeError('Could not find MSVC/GCC/CLANG installation on this computer. Check compiler_bindir_search_path list in "%s".' % __file__) 85 | else: 86 | cmd += ' --compiler-bindir "%s"' % compiler_bindir 87 | cmd += ' 2>&1' 88 | return cmd 89 | 90 | #---------------------------------------------------------------------------- 91 | # Main entry point. 92 | 93 | _plugin_cache = dict() 94 | 95 | def get_plugin(cuda_file, extra_nvcc_options=[]): 96 | cuda_file_base = os.path.basename(cuda_file) 97 | cuda_file_name, cuda_file_ext = os.path.splitext(cuda_file_base) 98 | 99 | # Already in cache? 100 | if cuda_file in _plugin_cache: 101 | return _plugin_cache[cuda_file] 102 | 103 | # Setup plugin. 104 | if verbose: 105 | print('Setting up TensorFlow plugin "%s": ' % cuda_file_base, end='', flush=True) 106 | try: 107 | # Hash CUDA source. 108 | md5 = hashlib.md5() 109 | with open(cuda_file, 'rb') as f: 110 | md5.update(f.read()) 111 | md5.update(b'\n') 112 | 113 | # Hash headers included by the CUDA code by running it through the preprocessor. 114 | if not do_not_hash_included_headers: 115 | if verbose: 116 | print('Preprocessing... ', end='', flush=True) 117 | with tempfile.TemporaryDirectory() as tmp_dir: 118 | tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + cuda_file_ext) 119 | _run_cmd(_prepare_nvcc_cli('"%s" --preprocess -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir))) 120 | with open(tmp_file, 'rb') as f: 121 | bad_file_str = ('"' + cuda_file.replace('\\', '/') + '"').encode('utf-8') # __FILE__ in error check macros 122 | good_file_str = ('"' + cuda_file_base + '"').encode('utf-8') 123 | for ln in f: 124 | if not ln.startswith(b'# ') and not ln.startswith(b'#line '): # ignore line number pragmas 125 | ln = ln.replace(bad_file_str, good_file_str) 126 | md5.update(ln) 127 | md5.update(b'\n') 128 | 129 | # Select compiler options. 130 | compile_opts = '' 131 | if os.name == 'nt': 132 | compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.lib') 133 | elif os.name == 'posix': 134 | compile_opts += ' --compiler-options \'-fPIC\'' 135 | compile_opts += ' --compiler-options \'%s\'' % " ".join(tf.sysconfig.get_compile_flags()) 136 | compile_opts += ' --linker-options \'%s\'' % " ".join(tf.sysconfig.get_link_flags()) 137 | else: 138 | assert False # not Windows or Linux, w00t? 139 | compile_opts += ' --gpu-architecture=%s' % _get_cuda_gpu_arch_string() 140 | compile_opts += ' --use_fast_math' 141 | for opt in extra_nvcc_options: 142 | compile_opts += ' ' + opt 143 | nvcc_cmd = _prepare_nvcc_cli(compile_opts) 144 | 145 | # Hash build configuration. 146 | md5.update(('nvcc_cmd: ' + nvcc_cmd).encode('utf-8') + b'\n') 147 | md5.update(('tf.VERSION: ' + tf.VERSION).encode('utf-8') + b'\n') 148 | md5.update(('cuda_cache_version_tag: ' + cuda_cache_version_tag).encode('utf-8') + b'\n') 149 | 150 | # Compile if not already compiled. 151 | cache_dir = util.make_cache_dir_path('tflib-cudacache') if cuda_cache_path is None else cuda_cache_path 152 | bin_file_ext = '.dll' if os.name == 'nt' else '.so' 153 | bin_file = os.path.join(cache_dir, cuda_file_name + '_' + md5.hexdigest() + bin_file_ext) 154 | if not os.path.isfile(bin_file): 155 | if verbose: 156 | print('Compiling... ', end='', flush=True) 157 | with tempfile.TemporaryDirectory() as tmp_dir: 158 | tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + bin_file_ext) 159 | _run_cmd(nvcc_cmd + ' "%s" --shared -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir)) 160 | os.makedirs(cache_dir, exist_ok=True) 161 | intermediate_file = os.path.join(cache_dir, cuda_file_name + '_' + uuid.uuid4().hex + '_tmp' + bin_file_ext) 162 | shutil.copyfile(tmp_file, intermediate_file) 163 | os.rename(intermediate_file, bin_file) # atomic 164 | 165 | # Load. 166 | if verbose: 167 | print('Loading... ', end='', flush=True) 168 | plugin = tf.load_op_library(bin_file) 169 | 170 | # Add to cache. 171 | _plugin_cache[cuda_file] = plugin 172 | if verbose: 173 | print('Done.', flush=True) 174 | return plugin 175 | 176 | except: 177 | if verbose: 178 | print('Failed!', flush=True) 179 | raise 180 | 181 | #---------------------------------------------------------------------------- 182 | -------------------------------------------------------------------------------- /src/_genSGAN2.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 3 | import os.path as osp 4 | import random 5 | import argparse 6 | import numpy as np 7 | from imageio import imread, imsave 8 | import pickle 9 | 10 | import dnnlib 11 | import dnnlib.tflib as tflib 12 | 13 | from util.utilgan import latent_anima, basename, img_read, img_list 14 | try: # progress bar for notebooks 15 | get_ipython().__class__.__name__ 16 | from util.progress_bar import ProgressIPy as ProgressBar 17 | except: # normal console 18 | from util.progress_bar import ProgressBar 19 | 20 | desc = "Customized StyleGAN2 on Tensorflow" 21 | parser = argparse.ArgumentParser(description=desc) 22 | parser.add_argument('--out_dir', default='_out', help='output directory') 23 | parser.add_argument('--model', default='models/ffhq-1024.pkl', help='path to pkl checkpoint file') 24 | parser.add_argument('--size', default=None, help='output resolution, set in X-Y format') 25 | parser.add_argument('--scale_type', default='pad', help="main types: pad, padside, symm, symmside") 26 | parser.add_argument('--latmask', default=None, help='external mask file (or directory) for multi latent blending (overriding frame split method)') 27 | parser.add_argument('--nXY', '-n', default='1-1', help='multi latent frame split count by X (width) and Y (height)') 28 | parser.add_argument('--splitfine', type=float, default=0, help='multi latent frame split edge sharpness (0 = smooth, higher => finer)') 29 | parser.add_argument('--trunc', type=float, default=0.8, help='truncation psi 0..1 (lower = stable, higher = various)') 30 | parser.add_argument('--labels', default=None, help='labels/categories for conditioning') 31 | parser.add_argument('--digress', type=float, default=0, help='distortion technique by Aydao (strength of the effect)') 32 | parser.add_argument('--save_lat', action='store_true', help='save latent vectors to file') 33 | parser.add_argument('--seed', type=int, default=None) 34 | parser.add_argument('--verbose', '-v', action='store_true') 35 | parser.add_argument('--ops', default='cuda', help='custom op implementation (cuda or ref)') 36 | # animation 37 | parser.add_argument('--frames', default='200-25', help='total frames to generate, length of interpolation step') 38 | parser.add_argument("--cubic", action='store_true', help="use cubic splines for smoothing") 39 | parser.add_argument("--gauss", action='store_true', help="use Gaussian smoothing") 40 | a = parser.parse_args() 41 | 42 | if a.size is not None: a.size = [int(s) for s in a.size.split('-')][::-1] 43 | [a.frames, a.fstep] = [int(s) for s in a.frames.split('-')] 44 | 45 | def main(): 46 | os.makedirs(a.out_dir, exist_ok=True) 47 | if a.seed==0: a.seed = None 48 | np.random.seed(seed=a.seed) 49 | 50 | # setup generator 51 | fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True) 52 | Gs_kwargs = dnnlib.EasyDict() 53 | Gs_kwargs.func_name = 'training.stylegan2_multi.G_main' 54 | Gs_kwargs.verbose = a.verbose 55 | Gs_kwargs.size = a.size 56 | Gs_kwargs.scale_type = a.scale_type 57 | Gs_kwargs.impl = a.ops 58 | 59 | # mask/blend latents with external latmask or by splitting the frame 60 | if a.latmask is None: 61 | nHW = [int(s) for s in a.nXY.split('-')][::-1] 62 | assert len(nHW)==2, ' Wrong count nXY: %d (must be 2)' % len(nHW) 63 | n_mult = nHW[0] * nHW[1] 64 | if a.verbose is True and n_mult > 1: print(' Latent blending w/split frame %d x %d' % (nHW[1], nHW[0])) 65 | lmask = np.tile(np.asarray([[[[None]]]]), (1,n_mult,1,1)) 66 | Gs_kwargs.countHW = nHW 67 | Gs_kwargs.splitfine = a.splitfine 68 | else: 69 | if a.verbose is True: print(' Latent blending with mask', a.latmask) 70 | n_mult = 2 71 | if osp.isfile(a.latmask): # single file 72 | lmask = np.asarray([[img_read(a.latmask)[:,:,0] / 255.]]) # [1,1,h,w] 73 | elif osp.isdir(a.latmask): # directory with frame sequence 74 | lmask = np.expand_dims(np.asarray([img_read(f)[:,:,0] / 255. for f in img_list(a.latmask)]), 1) # [n,1,h,w] 75 | else: 76 | print(' !! Blending mask not found:', a.latmask); exit(1) 77 | if a.verbose is True: print(' latmask shape', lmask.shape) 78 | lmask = np.concatenate((lmask, 1 - lmask), 1) # [n,2,h,w] 79 | Gs_kwargs.latmask_res = lmask.shape[2:] 80 | 81 | # load model with arguments 82 | sess = tflib.init_tf({'allow_soft_placement':True}) 83 | pkl_name = osp.splitext(a.model)[0] 84 | with open(pkl_name + '.pkl', 'rb') as file: 85 | network = pickle.load(file, encoding='latin1') 86 | try: _, _, network = network 87 | except: pass 88 | for k in list(network.static_kwargs.keys()): 89 | Gs_kwargs[k] = network.static_kwargs[k] 90 | 91 | # reload custom network, if needed 92 | if '.pkl' in a.model.lower(): 93 | print(' .. Gs from pkl ..', basename(a.model)) 94 | Gs = network 95 | else: # reconstruct network 96 | print(' .. Gs custom ..', basename(a.model)) 97 | # print(Gs_kwargs) 98 | Gs = tflib.Network('Gs', **Gs_kwargs) 99 | Gs.copy_vars_from(network) 100 | if a.verbose is True: print('kwargs:', ['%s: %s'%(kv[0],kv[1]) for kv in sorted(Gs.static_kwargs.items())]) 101 | 102 | if a.verbose is True: print(' out shape', Gs.output_shape[1:]) 103 | if a.size is None: a.size = Gs.output_shape[2:] 104 | 105 | if a.verbose is True: print(' making timeline..') 106 | lats = [] # list of [frm,1,512] 107 | for i in range(n_mult): 108 | lat_tmp = latent_anima((1, Gs.input_shape[1]), a.frames, a.fstep, cubic=a.cubic, gauss=a.gauss, seed=a.seed, verbose=False) # [frm,1,512] 109 | lats.append(lat_tmp) # list of [frm,1,512] 110 | latents = np.concatenate(lats, 1) # [frm,X,512] 111 | print(' latents', latents.shape) 112 | frame_count = latents.shape[0] 113 | 114 | # distort image by tweaking initial const layer 115 | if a.digress > 0: 116 | try: latent_size = Gs.static_kwargs['latent_size'] 117 | except: latent_size = 512 # default latent size 118 | try: init_res = Gs.static_kwargs['init_res'] 119 | except: init_res = (4,4) # default initial layer size 120 | dconst = [] 121 | for i in range(n_mult): 122 | dc_tmp = a.digress * latent_anima([1, latent_size, *init_res], a.frames, a.fstep, cubic=True, seed=a.seed, verbose=False) 123 | dconst.append(dc_tmp) 124 | dconst = np.concatenate(dconst, 1) 125 | else: 126 | dconst = np.zeros([frame_count, 1, 1, 1, 1]) 127 | 128 | # labels / conditions 129 | try: 130 | label_size = Gs_kwargs.label_size 131 | except: 132 | label_size = 0 133 | if label_size > 0: 134 | labels = np.zeros((frame_count, n_mult, label_size)) # [frm,X,lbl] 135 | if a.labels is None: 136 | label_ids = [] 137 | for i in range(n_mult): 138 | label_ids.append(random.randint(0, label_size-1)) 139 | else: 140 | label_ids = [int(x) for x in a.labels.split('-')] 141 | label_ids = label_ids[:n_mult] # ensure we have enough labels 142 | for i, l in enumerate(label_ids): 143 | labels[:,i,l] = 1 144 | else: 145 | labels = [None] 146 | 147 | # generate images from latent timeline 148 | pbar = ProgressBar(frame_count) 149 | for i in range(frame_count): 150 | 151 | latent = latents[i] # [X,512] 152 | label = labels[i % len(labels)] 153 | latmask = lmask[i % len(lmask)] if lmask is not None else [None] # [X,h,w] 154 | dc = dconst[i % len(dconst)] # [X,512,4,4] 155 | 156 | # generate multi-latent result 157 | if Gs.num_inputs == 2: 158 | output = Gs.run(latent, label, truncation_psi=a.trunc, randomize_noise=False, output_transform=fmt) 159 | else: 160 | output = Gs.run(latent, label, latmask, dc, truncation_psi=a.trunc, randomize_noise=False, output_transform=fmt) 161 | 162 | # save image 163 | ext = 'png' if output.shape[3]==4 else 'jpg' 164 | filename = osp.join(a.out_dir, "%06d.%s" % (i,ext)) 165 | imsave(filename, output[0]) 166 | pbar.upd() 167 | 168 | # convert latents to dlatents, save them 169 | if a.save_lat is True: 170 | latents = latents.squeeze(1) # [frm,512] 171 | dlatents = Gs.components.mapping.run(latents, label, dtype='float16') # [frm,18,512] 172 | filename = '{}-{}-{}.npy'.format(basename(a.model), a.size[1], a.size[0]) 173 | filename = osp.join(osp.dirname(a.out_dir), filename) 174 | np.save(filename, dlatents) 175 | print('saved dlatents', dlatents.shape, 'to', filename) 176 | 177 | if __name__ == '__main__': 178 | main() 179 | -------------------------------------------------------------------------------- /src/dnnlib/tflib/autosummary.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | """Helper for adding automatically tracked values to Tensorboard. 8 | 9 | Autosummary creates an identity op that internally keeps track of the input 10 | values and automatically shows up in TensorBoard. The reported value 11 | represents an average over input components. The average is accumulated 12 | constantly over time and flushed when save_summaries() is called. 13 | 14 | Notes: 15 | - The output tensor must be used as an input for something else in the 16 | graph. Otherwise, the autosummary op will not get executed, and the average 17 | value will not get accumulated. 18 | - It is perfectly fine to include autosummaries with the same name in 19 | several places throughout the graph, even if they are executed concurrently. 20 | - It is ok to also pass in a python scalar or numpy array. In this case, it 21 | is added to the average immediately. 22 | """ 23 | 24 | from collections import OrderedDict 25 | import numpy as np 26 | import tensorflow as tf 27 | from tensorboard import summary as summary_lib 28 | from tensorboard.plugins.custom_scalar import layout_pb2 29 | 30 | from . import tfutil 31 | from .tfutil import TfExpression 32 | from .tfutil import TfExpressionEx 33 | 34 | # Enable "Custom scalars" tab in TensorBoard for advanced formatting. 35 | # Disabled by default to reduce tfevents file size. 36 | enable_custom_scalars = False 37 | 38 | _dtype = tf.float64 39 | _vars = OrderedDict() # name => [var, ...] 40 | _immediate = OrderedDict() # name => update_op, update_value 41 | _finalized = False 42 | _merge_op = None 43 | 44 | 45 | def _create_var(name: str, value_expr: TfExpression) -> TfExpression: 46 | """Internal helper for creating autosummary accumulators.""" 47 | assert not _finalized 48 | name_id = name.replace("/", "_") 49 | v = tf.cast(value_expr, _dtype) 50 | 51 | if v.shape.is_fully_defined(): 52 | size = np.prod(v.shape.as_list()) 53 | size_expr = tf.constant(size, dtype=_dtype) 54 | else: 55 | size = None 56 | size_expr = tf.reduce_prod(tf.cast(tf.shape(v), _dtype)) 57 | 58 | if size == 1: 59 | if v.shape.ndims != 0: 60 | v = tf.reshape(v, []) 61 | v = [size_expr, v, tf.square(v)] 62 | else: 63 | v = [size_expr, tf.reduce_sum(v), tf.reduce_sum(tf.square(v))] 64 | v = tf.cond(tf.is_finite(v[1]), lambda: tf.stack(v), lambda: tf.zeros(3, dtype=_dtype)) 65 | 66 | with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.control_dependencies(None): 67 | var = tf.Variable(tf.zeros(3, dtype=_dtype), trainable=False) # [sum(1), sum(x), sum(x**2)] 68 | update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v)) 69 | 70 | if name in _vars: 71 | _vars[name].append(var) 72 | else: 73 | _vars[name] = [var] 74 | return update_op 75 | 76 | 77 | def autosummary(name: str, value: TfExpressionEx, passthru: TfExpressionEx = None, condition: TfExpressionEx = True) -> TfExpressionEx: 78 | """Create a new autosummary. 79 | 80 | Args: 81 | name: Name to use in TensorBoard 82 | value: TensorFlow expression or python value to track 83 | passthru: Optionally return this TF node without modifications but tack an autosummary update side-effect to this node. 84 | 85 | Example use of the passthru mechanism: 86 | 87 | n = autosummary('l2loss', loss, passthru=n) 88 | 89 | This is a shorthand for the following code: 90 | 91 | with tf.control_dependencies([autosummary('l2loss', loss)]): 92 | n = tf.identity(n) 93 | """ 94 | tfutil.assert_tf_initialized() 95 | name_id = name.replace("/", "_") 96 | 97 | if tfutil.is_tf_expression(value): 98 | with tf.name_scope("summary_" + name_id), tf.device(value.device): 99 | condition = tf.convert_to_tensor(condition, name='condition') 100 | update_op = tf.cond(condition, lambda: tf.group(_create_var(name, value)), tf.no_op) 101 | with tf.control_dependencies([update_op]): 102 | return tf.identity(value if passthru is None else passthru) 103 | 104 | else: # python scalar or numpy array 105 | assert not tfutil.is_tf_expression(passthru) 106 | assert not tfutil.is_tf_expression(condition) 107 | if condition: 108 | if name not in _immediate: 109 | with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.device(None), tf.control_dependencies(None): 110 | update_value = tf.placeholder(_dtype) 111 | update_op = _create_var(name, update_value) 112 | _immediate[name] = update_op, update_value 113 | update_op, update_value = _immediate[name] 114 | tfutil.run(update_op, {update_value: value}) 115 | return value if passthru is None else passthru 116 | 117 | 118 | def finalize_autosummaries() -> None: 119 | """Create the necessary ops to include autosummaries in TensorBoard report. 120 | Note: This should be done only once per graph. 121 | """ 122 | global _finalized 123 | tfutil.assert_tf_initialized() 124 | 125 | if _finalized: 126 | return None 127 | 128 | _finalized = True 129 | tfutil.init_uninitialized_vars([var for vars_list in _vars.values() for var in vars_list]) 130 | 131 | # Create summary ops. 132 | with tf.device(None), tf.control_dependencies(None): 133 | for name, vars_list in _vars.items(): 134 | name_id = name.replace("/", "_") 135 | with tfutil.absolute_name_scope("Autosummary/" + name_id): 136 | moments = tf.add_n(vars_list) 137 | moments /= moments[0] 138 | with tf.control_dependencies([moments]): # read before resetting 139 | reset_ops = [tf.assign(var, tf.zeros(3, dtype=_dtype)) for var in vars_list] 140 | with tf.name_scope(None), tf.control_dependencies(reset_ops): # reset before reporting 141 | mean = moments[1] 142 | std = tf.sqrt(moments[2] - tf.square(moments[1])) 143 | tf.summary.scalar(name, mean) 144 | if enable_custom_scalars: 145 | tf.summary.scalar("xCustomScalars/" + name + "/margin_lo", mean - std) 146 | tf.summary.scalar("xCustomScalars/" + name + "/margin_hi", mean + std) 147 | 148 | # Setup layout for custom scalars. 149 | layout = None 150 | if enable_custom_scalars: 151 | cat_dict = OrderedDict() 152 | for series_name in sorted(_vars.keys()): 153 | p = series_name.split("/") 154 | cat = p[0] if len(p) >= 2 else "" 155 | chart = "/".join(p[1:-1]) if len(p) >= 3 else p[-1] 156 | if cat not in cat_dict: 157 | cat_dict[cat] = OrderedDict() 158 | if chart not in cat_dict[cat]: 159 | cat_dict[cat][chart] = [] 160 | cat_dict[cat][chart].append(series_name) 161 | categories = [] 162 | for cat_name, chart_dict in cat_dict.items(): 163 | charts = [] 164 | for chart_name, series_names in chart_dict.items(): 165 | series = [] 166 | for series_name in series_names: 167 | series.append(layout_pb2.MarginChartContent.Series( 168 | value=series_name, 169 | lower="xCustomScalars/" + series_name + "/margin_lo", 170 | upper="xCustomScalars/" + series_name + "/margin_hi")) 171 | margin = layout_pb2.MarginChartContent(series=series) 172 | charts.append(layout_pb2.Chart(title=chart_name, margin=margin)) 173 | categories.append(layout_pb2.Category(title=cat_name, chart=charts)) 174 | layout = summary_lib.custom_scalar_pb(layout_pb2.Layout(category=categories)) 175 | return layout 176 | 177 | def save_summaries(file_writer, global_step=None): 178 | """Call FileWriter.add_summary() with all summaries in the default graph, 179 | automatically finalizing and merging them on the first call. 180 | """ 181 | global _merge_op 182 | tfutil.assert_tf_initialized() 183 | 184 | if _merge_op is None: 185 | layout = finalize_autosummaries() 186 | if layout is not None: 187 | file_writer.add_summary(layout) 188 | with tf.device(None), tf.control_dependencies(None): 189 | _merge_op = tf.summary.merge_all() 190 | 191 | file_writer.add_summary(_merge_op.eval(), global_step) 192 | -------------------------------------------------------------------------------- /src/training/dataset_tool.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # This work is made available under the Nvidia Source Code License-NC 3 | # https://nvlabs.github.io/stylegan2/license.html 4 | """Tool for creating multi-resolution TFRecords datasets.""" 5 | 6 | import os 7 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 8 | import sys 9 | import argparse 10 | import numpy as np 11 | import PIL.Image 12 | from turbojpeg import TurboJPEG 13 | jpeg = TurboJPEG() 14 | 15 | import tensorflow; tf = tensorflow.compat.v1 if hasattr(tensorflow.compat, 'v1') else tensorflow 16 | 17 | sys.path.append(os.path.dirname(os.path.dirname(__file__))) # upper dir 18 | 19 | import dnnlib 20 | 21 | try: # progress bar for notebooks 22 | get_ipython().__class__.__name__ 23 | from util.progress_bar import ProgressIPy as ProgressBar 24 | except: # normal console 25 | from util.progress_bar import ProgressBar 26 | 27 | class TFRecordExporter: 28 | def __init__(self, data_dir, expected_images, print_progress=False, progress_interval=10): 29 | # self.tfr_prefix = os.path.join(self.data_dir, os.path.basename(self.data_dir)) 30 | self.tfr_prefix = os.path.splitext(data_dir)[0] 31 | self.expected_images = expected_images 32 | self.cur_images = 0 33 | self.shape = None 34 | self.res_log2 = None 35 | self.tfr_writer = None 36 | self.print_progress = print_progress 37 | self.progress_interval = progress_interval 38 | 39 | if self.print_progress: 40 | print('Creating dataset "%s"' % data_dir) 41 | 42 | def close(self): 43 | if self.print_progress: 44 | print('%-40s\r' % 'Flushing data...', end='', flush=True) 45 | if self.tfr_writer is not None: 46 | self.tfr_writer.close() 47 | if self.print_progress: 48 | print('%-40s\r' % '', end='', flush=True) 49 | print('Added %d images.' % self.cur_images) 50 | 51 | def choose_shuffled_order(self): # Note: Images and labels must be added in shuffled order. 52 | order = np.arange(self.expected_images) 53 | np.random.RandomState(123).shuffle(order) 54 | return order 55 | 56 | def set_shape(self, shape): 57 | self.shape = shape # [c,h,w] 58 | assert self.shape[0] in [1,3,4] 59 | tfr_opt = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.NONE) 60 | self.tfr_file = self.tfr_prefix + '-%dx%d.tfr' % (self.shape[2], self.shape[1]) 61 | self.tfr_writer = tf.python_io.TFRecordWriter(self.tfr_file, tfr_opt) 62 | 63 | def add_image(self, img_path, jpg=False, size=None): 64 | if self.print_progress and self.cur_images % self.progress_interval == 0: 65 | print('%d / %d\r' % (self.cur_images, self.expected_images), end='', flush=True) 66 | 67 | def get_img(img_path): 68 | img = np.asarray(PIL.Image.open(img_path)) 69 | if img.shape[2] == 1: # monochrome 70 | img = img[:, :, np.newaxis] # HW => HWC 71 | return img.transpose([2,0,1]) # HWC => CHW 72 | 73 | if self.shape is None: 74 | img = get_img(img_path) 75 | self.set_shape(img.shape) 76 | assert img.shape == self.shape 77 | 78 | if jpg is True: 79 | with tf.gfile.GFile(img_path, 'rb') as jpg_file: 80 | raw_jpg = jpg_file.read() 81 | (width, height, jpeg_subsample, jpeg_colorspace) = jpeg.decode_header(raw_jpg) 82 | jpg_shape = [self.shape[0], height, width] 83 | ex = tf.train.Example(features = tf.train.Features(feature={ 84 | "shape": tf.train.Feature(int64_list=tf.train.Int64List(value = jpg_shape)), 85 | "data": tf.train.Feature(bytes_list=tf.train.BytesList(value = [raw_jpg]))})) 86 | self.tfr_writer.write(ex.SerializeToString()) 87 | else: 88 | img = get_img(img_path) 89 | if size is not None: 90 | img = img.resize(size, PIL.Image.ANTIALIAS) 91 | assert img.shape == self.shape, ' Image %s has shape %s (must be %s)' % (os.path.basename(img_path), str(img.shape), str(self.shape)) 92 | quant = np.rint(img).clip(0, 255).astype(np.uint8) 93 | ex = tf.train.Example(features=tf.train.Features(feature={ 94 | 'shape': tf.train.Feature(int64_list=tf.train.Int64List(value = quant.shape)), 95 | # 'data': tf.train.Feature(bytes_list=tf.train.BytesList(value = [quant.tostring()]))})) 96 | 'data': tf.train.Feature(bytes_list=tf.train.BytesList(value = [quant.tobytes()]))})) 97 | self.tfr_writer.write(ex.SerializeToString()) 98 | 99 | self.cur_images += 1 100 | 101 | def add_labels(self, labels): 102 | if self.print_progress: 103 | print('%-40s\r' % 'Saving labels...', end='', flush=True) 104 | assert labels.shape[0] == self.cur_images 105 | with open(self.tfr_prefix + '-rxx.labels', 'wb') as f: 106 | # np.save(f, labels.astype(np.int32)) 107 | np.save(f, labels.astype(np.float32)) 108 | 109 | def __enter__(self): 110 | return self 111 | 112 | def __exit__(self, *args): 113 | self.close() 114 | 115 | def img_list(path, subdir=None): 116 | if subdir is True: 117 | files = [os.path.join(dp, f) for dp, dn, fn in os.walk(path) for f in fn] 118 | else: 119 | files = [os.path.join(path, f) for f in os.listdir(path)] 120 | files = [f for f in files if os.path.splitext(f.lower())[1][1:] in ['jpg', 'jpeg', 'png', 'ppm', 'tif']] 121 | files = [f for f in files if not '/__MACOSX/' in f.replace('\\', '/')] # workaround fix for macos phantom files 122 | return sorted([f for f in files if os.path.isfile(f)]) 123 | 124 | def create_from_images(datadir, shuffle=True, size=None): 125 | assert os.path.isdir(datadir) 126 | imgs = sorted(img_list(datadir, subdir=True)) 127 | assert len(imgs) > 0, ' No input images found!' 128 | 129 | sample_img = np.asarray(PIL.Image.open(imgs[0])) 130 | sample_shape = sample_img.shape 131 | channels = sample_shape[2] if sample_img.ndim == 3 else 1 132 | assert channels in [1,3,4], ' Weird color dim: %d' % channels 133 | print(' Making dataset ..', datadir, sample_shape) 134 | jpg = channels < 4 135 | if jpg is True: print(' Loading JPG as is!') 136 | 137 | with TFRecordExporter(datadir, len(imgs)) as tfr: 138 | order = tfr.choose_shuffled_order() if shuffle else np.arange(len(imgs)) 139 | pbar = ProgressBar(order.size) 140 | for idx in range(order.size): 141 | img_path = imgs[order[idx]] 142 | tfr.add_image(img_path, jpg=jpg, size=size) 143 | pbar.upd() 144 | return tfr.tfr_file, len(imgs) 145 | 146 | def create_from_image_folders(datadir, shuffle=True, size=None): 147 | assert os.path.isdir(datadir) 148 | imgs = [] 149 | labels = [] 150 | for root, subdirs, files in os.walk(datadir): 151 | for i, subdir in enumerate(subdirs): 152 | tmp_list = img_list(os.path.join(root, subdir)) 153 | imgs = imgs + tmp_list 154 | labels = labels + [i] * len(tmp_list) 155 | labels = np.array(labels) 156 | onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32) 157 | onehot[np.arange(labels.size), labels] = 1. 158 | 159 | assert len(imgs) > 0, ' No input images found!' 160 | sample_img = np.asarray(PIL.Image.open(imgs[0])) 161 | sample_shape = sample_img.shape 162 | channels = sample_shape[2] if sample_img.ndim == 3 else 1 163 | assert channels in [1,3,4], ' Weird color dim: %d' % channels 164 | print(' Making dataset ..', datadir, sample_shape, '%d labels' % (np.max(labels)+1)) 165 | jpg = channels < 4 166 | if jpg is True: print(' Loading JPG as is!') 167 | 168 | with TFRecordExporter(datadir, len(imgs)) as tfr: 169 | order = tfr.choose_shuffled_order() if shuffle else np.arange(len(imgs)) 170 | pbar = ProgressBar(order.size) 171 | for idx in range(order.size): 172 | img_path = imgs[order[idx]] 173 | tfr.add_image(img_path, jpg=jpg, size=size) 174 | pbar.upd() 175 | tfr.add_labels(onehot[order]) 176 | return tfr.tfr_file, len(imgs) 177 | 178 | def main(): 179 | parser = argparse.ArgumentParser() 180 | parser.add_argument('--data', required=True, help='Directory containing the images') 181 | parser.add_argument('--shuffle', type=bool, default=True, help='Randomize image order (default: 1)') 182 | parser.add_argument('--labels', action='store_true', help='use folders to generate labels') 183 | args = parser.parse_args() 184 | 185 | if args.labels is True: 186 | create_from_image_folders(args.data, args.shuffle) 187 | else: 188 | create_from_images(args.data, args.shuffle) 189 | 190 | if __name__ == "__main__": 191 | main() 192 | 193 | -------------------------------------------------------------------------------- /src/train.py: -------------------------------------------------------------------------------- 1 | # Original copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # This work is made available under the Nvidia Source Code License-NC 3 | # https://nvlabs.github.io/stylegan2/license.html 4 | 5 | import os 6 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 7 | import warnings 8 | warnings.filterwarnings("ignore") 9 | import sys 10 | import argparse 11 | import copy 12 | import numpy as np 13 | import tensorflow as tf 14 | 15 | import dnnlib 16 | import dnnlib.tflib as tflib 17 | from dnnlib import EasyDict 18 | from training import dataset 19 | from training.dataset_tool import create_from_images, create_from_image_folders 20 | 21 | from util.utilgan import basename, file_list 22 | 23 | def run(data, train_dir, config, d_aug, diffaug_policy, cond, ops, mirror, mirror_v, \ 24 | kimg, batch_size, lrate, resume, resume_kimg, num_gpus, ema_kimg, gamma, freezeD): 25 | 26 | # training functions 27 | if d_aug: # https://github.com/mit-han-lab/data-efficient-gans 28 | train = EasyDict(run_func_name='training.training_loop_diffaug.training_loop') # Options for training loop (Diff Augment method) 29 | loss_args = EasyDict(func_name='training.loss_diffaug.ns_DiffAugment_r1', policy=diffaug_policy) # Options for loss (Diff Augment method) 30 | else: # original nvidia 31 | train = EasyDict(run_func_name='training.training_loop.training_loop') # Options for training loop (original from NVidia) 32 | G_loss = EasyDict(func_name='training.loss.G_logistic_ns_pathreg') # Options for generator loss. 33 | D_loss = EasyDict(func_name='training.loss.D_logistic_r1') # Options for discriminator loss. 34 | 35 | # network functions 36 | G = EasyDict(func_name='training.networks_stylegan2.G_main') # Options for generator network. 37 | D = EasyDict(func_name='training.networks_stylegan2.D_stylegan2') # Options for discriminator network. 38 | G_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for generator optimizer. 39 | D_opt = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8) # Options for discriminator optimizer. 40 | sched = EasyDict() # Options for TrainingSchedule. 41 | grid = EasyDict(size='1080p', layout='random') # Options for setup_snapshot_image_grid(). 42 | sc = dnnlib.SubmitConfig() # Options for dnnlib.submit_run(). 43 | tf_config = {'rnd.np_random_seed': 1000} # Options for tflib.init_tf(). 44 | G.impl = D.impl = ops 45 | 46 | # dataset (tfrecords) - get or create 47 | tfr_files = file_list(os.path.dirname(data), 'tfr') 48 | tfr_files = [f for f in tfr_files if basename(data) == basename(f).split('-')[0]] 49 | if len(tfr_files) == 0 or os.stat(tfr_files[0]).st_size == 0: 50 | tfr_file, total_samples = create_from_image_folders(data) if cond is True else create_from_images(data) 51 | else: 52 | tfr_file = tfr_files[0] 53 | dataset_args = EasyDict(tfrecord=tfr_file) 54 | 55 | # resolutions 56 | with tf.Graph().as_default(), tflib.create_session().as_default(): # pylint: disable=not-context-manager 57 | dataset_obj = dataset.load_dataset(**dataset_args) # loading the data to see what comes out 58 | resolution = dataset_obj.resolution 59 | init_res = dataset_obj.init_res 60 | res_log2 = dataset_obj.res_log2 61 | dataset_obj.close() 62 | dataset_obj = None 63 | 64 | if list(init_res) == [4,4]: 65 | desc = '%s-%d' % (basename(data), resolution) 66 | else: 67 | print(' custom init resolution', init_res) 68 | desc = basename(tfr_file) 69 | G.init_res = D.init_res = list(init_res) 70 | 71 | train.savenames = [desc.replace(basename(data), 'snapshot'), desc] 72 | desc += '-%s' % config 73 | 74 | # training schedule 75 | train.total_kimg = kimg 76 | train.image_snapshot_ticks = 1 * num_gpus if kimg <= 1000 else 4 * num_gpus 77 | train.network_snapshot_ticks = 5 78 | train.mirror_augment = mirror 79 | train.mirror_augment_v = mirror_v 80 | sched.tick_kimg_base = 2 if train.total_kimg < 2000 else 4 81 | 82 | # learning rate 83 | if config == 'e': 84 | sched.G_lrate_base = 0.001 85 | sched.G_lrate_dict = {0:0.001, 1:0.0007, 2:0.0005, 3:0.0003} 86 | sched.lrate_step = 1500 # period for stepping to next lrate, in kimg 87 | if config == 'f': 88 | sched.G_lrate_base = lrate # 0.001 for big datasets, 0.0003 for few-shot 89 | sched.D_lrate_base = sched.G_lrate_base # *2 - not used anyway 90 | 91 | # batch size (for 16gb memory GPU) 92 | sched.minibatch_gpu_base = 4096 // resolution if batch_size is None else batch_size 93 | print(' Batch size', sched.minibatch_gpu_base) 94 | sched.minibatch_size_base = num_gpus * sched.minibatch_gpu_base 95 | sc.num_gpus = num_gpus 96 | 97 | if config == 'e': 98 | G.fmap_base = D.fmap_base = 8 << 10 99 | if d_aug: loss_args.gamma = 100 if gamma is None else gamma 100 | else: D_loss.gamma = 100 if gamma is None else gamma 101 | elif config == 'f': 102 | G.fmap_base = D.fmap_base = 16 << 10 103 | else: 104 | print(' Only configs E and F are implemented'); exit() 105 | 106 | if cond: 107 | desc += '-cond'; dataset_args.max_label_size = 'full' # conditioned on full label 108 | 109 | if freezeD: 110 | D.freezeD = True 111 | train.resume_with_new_nets = True 112 | 113 | if d_aug: 114 | desc += '-daug' 115 | 116 | sc.submit_target = dnnlib.SubmitTarget.LOCAL 117 | sc.local.do_not_copy_source_files = True 118 | kwargs = EasyDict(train) 119 | kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt) 120 | kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, tf_config=tf_config) 121 | kwargs.update(resume_pkl=resume, resume_kimg=resume_kimg, resume_with_new_nets=True) 122 | if ema_kimg is not None: 123 | kwargs.update(G_ema_kimg=ema_kimg) 124 | if d_aug: 125 | kwargs.update(loss_args=loss_args) 126 | else: 127 | kwargs.update(G_loss_args=G_loss, D_loss_args=D_loss) 128 | kwargs.submit_config = copy.deepcopy(sc) 129 | kwargs.submit_config.run_dir_root = train_dir 130 | kwargs.submit_config.run_desc = desc 131 | dnnlib.submit_run(**kwargs) 132 | 133 | def _str_to_bool(v): 134 | if isinstance(v, bool): 135 | return v 136 | if v.lower() in ('yes', 'true', 't', 'y', '1'): 137 | return True 138 | elif v.lower() in ('no', 'false', 'f', 'n', '0'): 139 | return False 140 | else: 141 | raise argparse.ArgumentTypeError('Boolean value expected.') 142 | 143 | 144 | def main(): 145 | parser = argparse.ArgumentParser(description='StyleGAN2 for practice', formatter_class=argparse.RawDescriptionHelpFormatter) 146 | # main 147 | parser.add_argument('--data', required=True, help='Training dataset path', metavar='DIR') 148 | parser.add_argument('--train_dir', default='train', help='Root directory for training results (default: %(default)s)', metavar='DIR') 149 | parser.add_argument('--resume', default=None, help='Resume checkpoint path. None = from scratch') 150 | parser.add_argument('--resume_kimg', type=int, default=0, help='Resume training from (in thousands of images)', metavar='N') 151 | parser.add_argument('--kimg', type=int, default=None, help='Override total training duration', metavar='N') 152 | # network 153 | parser.add_argument('--config', default='F', help='Training config E (shrink) or F (large) (default: %(default)s)', metavar='CONFIG') 154 | parser.add_argument('--ops', default='cuda', help='Custom op implementation (cuda or ref, default: %(default)s)') 155 | parser.add_argument('--gamma', default=None, type=float, help='R1 regularization weight') 156 | # special 157 | parser.add_argument('--d_aug', action='store_true', help='Use Diff Augment training for small datasets') 158 | parser.add_argument('--diffaug_policy', default='translation,cutout', help='Comma-separated list of DiffAugment policies (default: %(default)s)', metavar='..') # color 159 | parser.add_argument('--ema_kimg', default=None, type=int, help='Half-life of exponential moving average (for Diff Augment)', metavar='N') 160 | parser.add_argument('--freezeD', action='store_true', help='freeze lower D layers for better finetuning') 161 | parser.add_argument('--cond', action='store_true', help='conditional model') 162 | # training 163 | parser.add_argument('--batch_size', default=None, type=int, help='Batch size per GPU (default: %(default)s)', metavar='N') 164 | parser.add_argument('-lr', '--lrate', default=0.001, type=float, help='Learning rate for F config (default: %(default)s)') 165 | parser.add_argument('--mirror', help='Mirror augment (default: %(default)s)', default=True, metavar='BOOL', type=_str_to_bool) 166 | parser.add_argument('--mirror_v', help='Mirror augment vertically (default: %(default)s)', default=False, metavar='BOOL', type=_str_to_bool) 167 | parser.add_argument('--num_gpus', help='Number of GPUs (default: %(default)s)', default=1, type=int, metavar='N') 168 | args = parser.parse_args() 169 | 170 | args.config = args.config.lower() 171 | 172 | run(**vars(args)) 173 | 174 | 175 | if __name__ == "__main__": 176 | main() 177 | -------------------------------------------------------------------------------- /src/dnnlib/tflib/ops/fused_bias_act.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. 2 | // 3 | // NVIDIA CORPORATION and its licensors retain all intellectual property 4 | // and proprietary rights in and to this software, related documentation 5 | // and any modifications thereto. Any use, reproduction, disclosure or 6 | // distribution of this software and related documentation without an express 7 | // license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | #define EIGEN_USE_GPU 10 | #define __CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__ 11 | #include "tensorflow/core/framework/op.h" 12 | #include "tensorflow/core/framework/op_kernel.h" 13 | #include "tensorflow/core/framework/shape_inference.h" 14 | #include 15 | 16 | using namespace tensorflow; 17 | using namespace tensorflow::shape_inference; 18 | 19 | #define OP_CHECK_CUDA_ERROR(CTX, CUDA_CALL) do { cudaError_t err = CUDA_CALL; OP_REQUIRES(CTX, err == cudaSuccess, errors::Internal(cudaGetErrorName(err))); } while (false) 20 | 21 | //------------------------------------------------------------------------ 22 | // CUDA kernel. 23 | 24 | template 25 | struct FusedBiasActKernelParams 26 | { 27 | const T* x; // [sizeX] 28 | const T* b; // [sizeB] or NULL 29 | const T* xref; // [sizeX] or NULL 30 | const T* yref; // [sizeX] or NULL 31 | T* y; // [sizeX] 32 | 33 | int grad; 34 | int axis; 35 | int act; 36 | float alpha; 37 | float gain; 38 | float clamp; 39 | 40 | int sizeX; 41 | int sizeB; 42 | int stepB; 43 | int loopX; 44 | }; 45 | 46 | template 47 | static __global__ void FusedBiasActKernel(const FusedBiasActKernelParams p) 48 | { 49 | const float expRange = 80.0f; 50 | const float halfExpRange = 40.0f; 51 | const float seluScale = 1.0507009873554804934193349852946f; 52 | const float seluAlpha = 1.6732632423543772848170429916717f; 53 | 54 | // Loop over elements. 55 | int xi = blockIdx.x * p.loopX * blockDim.x + threadIdx.x; 56 | for (int loopIdx = 0; loopIdx < p.loopX && xi < p.sizeX; loopIdx++, xi += blockDim.x) 57 | { 58 | // Load and apply bias. 59 | float x = (float)p.x[xi]; 60 | if (p.b) 61 | x += (float)p.b[(xi / p.stepB) % p.sizeB]; 62 | float xref = (p.xref) ? (float)p.xref[xi] : 0.0f; 63 | float yref = (p.yref) ? (float)p.yref[xi] : 0.0f; 64 | float yy = (p.gain != 0.0f) ? yref / p.gain : 0.0f; 65 | 66 | // Evaluate activation func. 67 | float y; 68 | switch (p.act * 10 + p.grad) 69 | { 70 | // linear 71 | default: 72 | case 10: y = x; break; 73 | case 11: y = x; break; 74 | case 12: y = 0.0f; break; 75 | 76 | // relu 77 | case 20: y = (x > 0.0f) ? x : 0.0f; break; 78 | case 21: y = (yy > 0.0f) ? x : 0.0f; break; 79 | case 22: y = 0.0f; break; 80 | 81 | // lrelu 82 | case 30: y = (x > 0.0f) ? x : x * p.alpha; break; 83 | case 31: y = (yy > 0.0f) ? x : x * p.alpha; break; 84 | case 32: y = 0.0f; break; 85 | 86 | // tanh 87 | case 40: { float c = expf(x); float d = 1.0f / c; y = (x < -expRange) ? -1.0f : (x > expRange) ? 1.0f : (c - d) / (c + d); } break; 88 | case 41: y = x * (1.0f - yy * yy); break; 89 | case 42: y = x * (1.0f - yy * yy) * (-2.0f * yy); break; 90 | 91 | // sigmoid 92 | case 50: y = (x < -expRange) ? 0.0f : 1.0f / (expf(-x) + 1.0f); break; 93 | case 51: y = x * yy * (1.0f - yy); break; 94 | case 52: y = x * yy * (1.0f - yy) * (1.0f - 2.0f * yy); break; 95 | 96 | // elu 97 | case 60: y = (x >= 0.0f) ? x : expf(x) - 1.0f; break; 98 | case 61: y = (yy >= 0.0f) ? x : x * (yy + 1.0f); break; 99 | case 62: y = (yy >= 0.0f) ? 0.0f : x * (yy + 1.0f); break; 100 | 101 | // selu 102 | case 70: y = (x >= 0.0f) ? seluScale * x : (seluScale * seluAlpha) * (expf(x) - 1.0f); break; 103 | case 71: y = (yy >= 0.0f) ? x * seluScale : x * (yy + seluScale * seluAlpha); break; 104 | case 72: y = (yy >= 0.0f) ? 0.0f : x * (yy + seluScale * seluAlpha); break; 105 | 106 | // softplus 107 | case 80: y = (x > expRange) ? x : logf(expf(x) + 1.0f); break; 108 | case 81: y = x * (1.0f - expf(-yy)); break; 109 | case 82: { float c = expf(-yy); y = x * c * (1.0f - c); } break; 110 | 111 | // swish 112 | case 90: y = (x < -expRange) ? 0.0f : x / (expf(-x) + 1.0f); break; 113 | case 91: 114 | case 92: 115 | { 116 | float c = expf(xref); 117 | float d = c + 1.0f; 118 | if (p.grad == 1) 119 | y = (xref > halfExpRange) ? x : x * c * (xref + d) / (d * d); 120 | else 121 | y = (xref > halfExpRange) ? 0.0f : x * c * (xref * (2.0f - d) + 2.0f * d) / (d * d * d); 122 | yref = (xref < -expRange) ? 0.0f : xref / (expf(-xref) + 1.0f) * p.gain; 123 | } 124 | break; 125 | } 126 | 127 | // Apply gain. 128 | y *= p.gain; 129 | 130 | // Clamp. 131 | if (p.clamp >= 0.0f) 132 | { 133 | if (p.grad == 0) 134 | y = (fabsf(y) < p.clamp) ? y : (y >= 0.0f) ? p.clamp : -p.clamp; 135 | else 136 | y = (fabsf(yref) < p.clamp) ? y : 0.0f; 137 | } 138 | 139 | // Store. 140 | p.y[xi] = (T)y; 141 | } 142 | } 143 | 144 | //------------------------------------------------------------------------ 145 | // TensorFlow op. 146 | 147 | template 148 | struct FusedBiasActOp : public OpKernel 149 | { 150 | FusedBiasActKernelParams m_attribs; 151 | 152 | FusedBiasActOp(OpKernelConstruction* ctx) : OpKernel(ctx) 153 | { 154 | memset(&m_attribs, 0, sizeof(m_attribs)); 155 | OP_REQUIRES_OK(ctx, ctx->GetAttr("grad", &m_attribs.grad)); 156 | OP_REQUIRES_OK(ctx, ctx->GetAttr("axis", &m_attribs.axis)); 157 | OP_REQUIRES_OK(ctx, ctx->GetAttr("act", &m_attribs.act)); 158 | OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha", &m_attribs.alpha)); 159 | OP_REQUIRES_OK(ctx, ctx->GetAttr("gain", &m_attribs.gain)); 160 | OP_REQUIRES_OK(ctx, ctx->GetAttr("clamp", &m_attribs.clamp)); 161 | OP_REQUIRES(ctx, m_attribs.grad >= 0, errors::InvalidArgument("grad must be non-negative")); 162 | OP_REQUIRES(ctx, m_attribs.axis >= 0, errors::InvalidArgument("axis must be non-negative")); 163 | OP_REQUIRES(ctx, m_attribs.act >= 0, errors::InvalidArgument("act must be non-negative")); 164 | } 165 | 166 | void Compute(OpKernelContext* ctx) 167 | { 168 | FusedBiasActKernelParams p = m_attribs; 169 | cudaStream_t stream = ctx->eigen_device().stream(); 170 | 171 | const Tensor& x = ctx->input(0); // [...] 172 | const Tensor& b = ctx->input(1); // [sizeB] or [0] 173 | const Tensor& xref = ctx->input(2); // x.shape or [0] 174 | const Tensor& yref = ctx->input(3); // x.shape or [0] 175 | p.x = x.flat().data(); 176 | p.b = (b.NumElements()) ? b.flat().data() : NULL; 177 | p.xref = (xref.NumElements()) ? xref.flat().data() : NULL; 178 | p.yref = (yref.NumElements()) ? yref.flat().data() : NULL; 179 | OP_REQUIRES(ctx, b.NumElements() == 0 || m_attribs.axis < x.dims(), errors::InvalidArgument("axis out of bounds")); 180 | OP_REQUIRES(ctx, b.dims() == 1, errors::InvalidArgument("b must have rank 1")); 181 | OP_REQUIRES(ctx, b.NumElements() == 0 || b.NumElements() == x.dim_size(m_attribs.axis), errors::InvalidArgument("b has wrong number of elements")); 182 | OP_REQUIRES(ctx, xref.NumElements() == 0 || xref.NumElements() == x.NumElements(), errors::InvalidArgument("xref has wrong number of elements")); 183 | OP_REQUIRES(ctx, yref.NumElements() == 0 || yref.NumElements() == x.NumElements(), errors::InvalidArgument("yref has wrong number of elements")); 184 | OP_REQUIRES(ctx, x.NumElements() <= kint32max, errors::InvalidArgument("x is too large")); 185 | 186 | p.sizeX = (int)x.NumElements(); 187 | p.sizeB = (int)b.NumElements(); 188 | p.stepB = 1; 189 | for (int i = m_attribs.axis + 1; i < x.dims(); i++) 190 | p.stepB *= (int)x.dim_size(i); 191 | 192 | Tensor* y = NULL; // x.shape 193 | OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &y)); 194 | p.y = y->flat().data(); 195 | 196 | p.loopX = 4; 197 | int blockSize = 4 * 32; 198 | int gridSize = (p.sizeX - 1) / (p.loopX * blockSize) + 1; 199 | void* args[] = {&p}; 200 | OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel((void*)FusedBiasActKernel, gridSize, blockSize, args, 0, stream)); 201 | } 202 | }; 203 | 204 | REGISTER_OP("FusedBiasAct") 205 | .Input ("x: T") 206 | .Input ("b: T") 207 | .Input ("xref: T") 208 | .Input ("yref: T") 209 | .Output ("y: T") 210 | .Attr ("T: {float, half}") 211 | .Attr ("grad: int = 0") 212 | .Attr ("axis: int = 1") 213 | .Attr ("act: int = 0") 214 | .Attr ("alpha: float = 0.0") 215 | .Attr ("gain: float = 1.0") 216 | .Attr ("clamp: float = -1.0"); 217 | REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint("T"), FusedBiasActOp); 218 | REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint("T"), FusedBiasActOp); 219 | 220 | //------------------------------------------------------------------------ 221 | -------------------------------------------------------------------------------- /src/projector.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | import os 8 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 9 | import numpy as np 10 | import tensorflow as tf 11 | import dnnlib 12 | import dnnlib.tflib as tflib 13 | 14 | from training import misc 15 | 16 | #---------------------------------------------------------------------------- 17 | 18 | class Projector: 19 | def __init__(self, steps=1000): 20 | self.num_steps = steps 21 | self.dlatent_avg_samples = 10000 22 | self.initial_learning_rate = 0.1 23 | self.initial_noise_factor = 0.05 24 | self.lr_rampdown_length = 0.25 25 | self.lr_rampup_length = 0.05 26 | self.noise_ramp_length = 0.75 27 | self.regularize_noise_weight = 1e5 28 | self.verbose = False 29 | self.clone_net = True 30 | 31 | self._Gs = None 32 | self._minibatch_size = None 33 | self._dlatent_avg = None 34 | self._dlatent_std = None 35 | self._noise_vars = None 36 | self._noise_init_op = None 37 | self._noise_normalize_op = None 38 | self._dlatents_var = None 39 | self._noise_in = None 40 | self._dlatents_expr = None 41 | self._images_expr = None 42 | self._target_images_var = None 43 | self._lpips = None 44 | self._dist = None 45 | self._loss = None 46 | self._reg_sizes = None 47 | self._lrate_in = None 48 | self._opt = None 49 | self._opt_step = None 50 | self._cur_step = None 51 | 52 | def _info(self, *args): 53 | if self.verbose: 54 | print('Projector:', *args) 55 | 56 | def set_network(self, Gs, minibatch_size=1): 57 | assert minibatch_size == 1 58 | self._Gs = Gs 59 | self._minibatch_size = minibatch_size 60 | if self._Gs is None: 61 | return 62 | if self.clone_net: 63 | self._Gs = self._Gs.clone() 64 | 65 | # Find dlatent stats. 66 | self._info('Finding W midpoint and stddev using %d samples...' % self.dlatent_avg_samples) 67 | latent_samples = np.random.RandomState(123).randn(self.dlatent_avg_samples, *self._Gs.input_shapes[0][1:]) 68 | dlatent_samples = self._Gs.components.mapping.run(latent_samples, None) # [N, 1, 512] 69 | self._dlatent_avg = np.mean(dlatent_samples, axis=0, keepdims=True) # [1, 1, 512] 70 | self._dlatent_std = (np.sum((dlatent_samples - self._dlatent_avg) ** 2) / self.dlatent_avg_samples) ** 0.5 71 | self._info('std = %g' % self._dlatent_std) 72 | 73 | # Find noise inputs. 74 | self._info('Setting up noise inputs...') 75 | self._noise_vars = [] 76 | noise_init_ops = [] 77 | noise_normalize_ops = [] 78 | while True: 79 | n = 'G_synthesis/noise%d' % len(self._noise_vars) 80 | if not n in self._Gs.vars: 81 | break 82 | v = self._Gs.vars[n] 83 | self._noise_vars.append(v) 84 | noise_init_ops.append(tf.assign(v, tf.random_normal(tf.shape(v), dtype=tf.float32))) 85 | noise_mean = tf.reduce_mean(v) 86 | noise_std = tf.reduce_mean((v - noise_mean)**2)**0.5 87 | noise_normalize_ops.append(tf.assign(v, (v - noise_mean) / noise_std)) 88 | self._info(n, v) 89 | self._noise_init_op = tf.group(*noise_init_ops) 90 | self._noise_normalize_op = tf.group(*noise_normalize_ops) 91 | 92 | # Image output graph. 93 | self._info('Building image output graph...') 94 | self._dlatents_var = tf.Variable(tf.zeros([self._minibatch_size] + list(self._dlatent_avg.shape[1:])), name='dlatents_var') 95 | self._noise_in = tf.placeholder(tf.float32, [], name='noise_in') 96 | dlatents_noise = tf.random.normal(shape=self._dlatents_var.shape) * self._noise_in 97 | # self._dlatents_expr = tf.tile(self._dlatents_var + dlatents_noise, [1, self._Gs.components.synthesis.input_shape[1], 1]) 98 | self._dlatents_expr = self._dlatents_var + dlatents_noise 99 | self._images_expr = self._Gs.components.synthesis.get_output_for(self._dlatents_expr, randomize_noise=False) 100 | 101 | # Downsample image to 256x256 if it's larger than that. VGG was built for 224x224 images. 102 | proc_images_expr = (self._images_expr[:,:3,:,:] + 1) * (255 / 2) # go uint range, fix to rgb colospace 103 | sh = proc_images_expr.shape.as_list() 104 | if sh[2] > 256: 105 | factor = sh[2] // 256 106 | proc_images_expr = tf.reduce_mean(tf.reshape(proc_images_expr, [-1, sh[1], sh[2] // factor, factor, sh[2] // factor, factor]), axis=[3,5]) 107 | 108 | # Loss graph. 109 | self._info('Building loss graph...') 110 | self._target_images_var = tf.Variable(tf.zeros(proc_images_expr.shape), name='target_images_var') 111 | if self._lpips is None: 112 | vgg_file = 'models/vgg/vgg16_zhang_perceptual.pkl' 113 | if os.path.isfile(vgg_file) and os.stat(vgg_file).st_size == 58871973: 114 | self._lpips = misc.load_pkl(vgg_file) 115 | else: 116 | self._lpips = misc.load_pkl('https://drive.google.com/uc?id=1N2-m9qszOeVC9Tq77WxsLnuWwOedQiD2') 117 | 118 | self._dist = self._lpips.get_output_for(proc_images_expr, self._target_images_var) 119 | self._loss = tf.reduce_sum(self._dist) 120 | 121 | # Noise regularization graph. 122 | self._info('Building noise regularization graph...') 123 | reg_loss = 0.0 124 | for v in self._noise_vars: 125 | sz = v.shape[2] 126 | while True: 127 | reg_loss += tf.reduce_mean(v * tf.roll(v, shift=1, axis=3))**2 + tf.reduce_mean(v * tf.roll(v, shift=1, axis=2))**2 128 | if sz <= 8: 129 | break # Small enough already 130 | v = tf.reshape(v, [1, 1, sz//2, 2, sz//2, 2]) # Downscale 131 | v = tf.reduce_mean(v, axis=[3, 5]) 132 | sz = sz // 2 133 | self._loss += reg_loss * self.regularize_noise_weight 134 | 135 | # Optimizer. 136 | self._info('Setting up optimizer...') 137 | self._lrate_in = tf.placeholder(tf.float32, [], name='lrate_in') 138 | self._opt = dnnlib.tflib.Optimizer(learning_rate=self._lrate_in) 139 | self._opt.register_gradients(self._loss, [self._dlatents_var] + self._noise_vars) 140 | self._opt_step = self._opt.apply_updates() 141 | 142 | def run(self, target_images): 143 | # Run to completion. 144 | self.start(target_images) 145 | while self._cur_step < self.num_steps: 146 | self.step() 147 | 148 | # Collect results. 149 | pres = dnnlib.EasyDict() 150 | pres.dlatents = self.get_dlatents() 151 | pres.noises = self.get_noises() 152 | pres.images = self.get_images() 153 | return pres 154 | 155 | def start(self, target_images): 156 | assert self._Gs is not None 157 | 158 | # Prepare target images. 159 | self._info('Preparing target images...') 160 | target_images = np.asarray(target_images, dtype='float32') 161 | target_images = (target_images + 1) * (255 / 2) 162 | sh = target_images.shape 163 | assert sh[0] == self._minibatch_size 164 | if sh[2] > self._target_images_var.shape[2]: 165 | factor = sh[2] // self._target_images_var.shape[2] 166 | target_images = np.reshape(target_images, [-1, sh[1], sh[2] // factor, factor, sh[3] // factor, factor]).mean((3, 5)) 167 | 168 | # Initialize optimization state. 169 | self._info('Initializing optimization state...') 170 | tflib.set_vars({self._target_images_var: target_images, self._dlatents_var: np.tile(self._dlatent_avg, [self._minibatch_size, 1, 1])}) 171 | tflib.run(self._noise_init_op) 172 | self._opt.reset_optimizer_state() 173 | self._cur_step = 0 174 | 175 | def step(self): 176 | assert self._cur_step is not None 177 | if self._cur_step >= self.num_steps: 178 | return 179 | if self._cur_step == 0: 180 | self._info('Running...') 181 | 182 | # Hyperparameters. 183 | t = self._cur_step / self.num_steps 184 | noise_strength = self._dlatent_std * self.initial_noise_factor * max(0.0, 1.0 - t / self.noise_ramp_length) ** 2 185 | lr_ramp = min(1.0, (1.0 - t) / self.lr_rampdown_length) 186 | lr_ramp = 0.5 - 0.5 * np.cos(lr_ramp * np.pi) 187 | lr_ramp = lr_ramp * min(1.0, t / self.lr_rampup_length) 188 | learning_rate = self.initial_learning_rate * lr_ramp 189 | 190 | # Train. 191 | feed_dict = {self._noise_in: noise_strength, self._lrate_in: learning_rate} 192 | _, dist_value, loss_value = tflib.run([self._opt_step, self._dist, self._loss], feed_dict) 193 | tflib.run(self._noise_normalize_op) 194 | 195 | # Print status. 196 | self._cur_step += 1 197 | if self._cur_step == self.num_steps or self._cur_step % 10 == 0: 198 | self._info('%-8d%-12g%-12g' % (self._cur_step, dist_value, loss_value)) 199 | if self._cur_step == self.num_steps: 200 | self._info('Done.') 201 | 202 | def get_cur_step(self): 203 | return self._cur_step 204 | 205 | def get_dlatents(self): 206 | return tflib.run(self._dlatents_expr, {self._noise_in: 0}) 207 | 208 | def get_noises(self): 209 | return tflib.run(self._noise_vars) 210 | 211 | def get_images(self): 212 | return tflib.run(self._images_expr, {self._noise_in: 0}) 213 | 214 | #---------------------------------------------------------------------------- 215 | -------------------------------------------------------------------------------- /src/dnnlib/tflib/ops/fused_bias_act.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. 2 | # 3 | # NVIDIA CORPORATION and its licensors retain all intellectual property 4 | # and proprietary rights in and to this software, related documentation 5 | # and any modifications thereto. Any use, reproduction, disclosure or 6 | # distribution of this software and related documentation without an express 7 | # license agreement from NVIDIA CORPORATION is strictly prohibited. 8 | 9 | """Custom TensorFlow ops for efficient bias and activation.""" 10 | 11 | import os 12 | import numpy as np 13 | import tensorflow as tf 14 | from .. import custom_ops 15 | from ...util import EasyDict 16 | 17 | def _get_plugin(): 18 | return custom_ops.get_plugin(os.path.splitext(__file__)[0] + '.cu') 19 | 20 | #---------------------------------------------------------------------------- 21 | 22 | activation_funcs = { 23 | 'linear': EasyDict(func=lambda x, **_: x, def_alpha=None, def_gain=1.0, cuda_idx=1, ref='y', zero_2nd_grad=True), 24 | 'relu': EasyDict(func=lambda x, **_: tf.nn.relu(x), def_alpha=None, def_gain=np.sqrt(2), cuda_idx=2, ref='y', zero_2nd_grad=True), 25 | 'lrelu': EasyDict(func=lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha), def_alpha=0.2, def_gain=np.sqrt(2), cuda_idx=3, ref='y', zero_2nd_grad=True), 26 | 'tanh': EasyDict(func=lambda x, **_: tf.nn.tanh(x), def_alpha=None, def_gain=1.0, cuda_idx=4, ref='y', zero_2nd_grad=False), 27 | 'sigmoid': EasyDict(func=lambda x, **_: tf.nn.sigmoid(x), def_alpha=None, def_gain=1.0, cuda_idx=5, ref='y', zero_2nd_grad=False), 28 | 'elu': EasyDict(func=lambda x, **_: tf.nn.elu(x), def_alpha=None, def_gain=1.0, cuda_idx=6, ref='y', zero_2nd_grad=False), 29 | 'selu': EasyDict(func=lambda x, **_: tf.nn.selu(x), def_alpha=None, def_gain=1.0, cuda_idx=7, ref='y', zero_2nd_grad=False), 30 | 'softplus': EasyDict(func=lambda x, **_: tf.nn.softplus(x), def_alpha=None, def_gain=1.0, cuda_idx=8, ref='y', zero_2nd_grad=False), 31 | 'swish': EasyDict(func=lambda x, **_: tf.nn.sigmoid(x) * x, def_alpha=None, def_gain=np.sqrt(2), cuda_idx=9, ref='x', zero_2nd_grad=False), 32 | } 33 | 34 | #---------------------------------------------------------------------------- 35 | 36 | def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, clamp=None, impl='cuda'): 37 | r"""Fused bias and activation function. 38 | 39 | Adds bias `b` to activation tensor `x`, evaluates activation function `act`, 40 | and scales the result by `gain`. Each of the steps is optional. In most cases, 41 | the fused op is considerably more efficient than performing the same calculation 42 | using standard TensorFlow ops. It supports first and second order gradients, 43 | but not third order gradients. 44 | 45 | Args: 46 | x: Input activation tensor. Can have any shape, but if `b` is defined, the 47 | dimension corresponding to `axis`, as well as the rank, must be known. 48 | b: Bias vector, or `None` to disable. Must be a 1D tensor of the same type 49 | as `x`. The shape must be known, and it must match the dimension of `x` 50 | corresponding to `axis`. 51 | axis: The dimension in `x` corresponding to the elements of `b`. 52 | The value of `axis` is ignored if `b` is not specified. 53 | act: Name of the activation function to evaluate, or `"linear"` to disable. 54 | Can be e.g. `"relu"`, `"lrelu"`, `"tanh"`, `"sigmoid"`, `"swish"`, etc. 55 | See `activation_funcs` for a full list. `None` is not allowed. 56 | alpha: Shape parameter for the activation function, or `None` to use the default. 57 | gain: Scaling factor for the output tensor, or `None` to use default. 58 | See `activation_funcs` for the default scaling of each activation function. 59 | If unsure, consider specifying `1.0`. 60 | clamp: Clamp the output values to `[-clamp, +clamp]`, or `None` to disable 61 | the clamping (default). 62 | impl: Name of the implementation to use. Can be `"ref"` or `"cuda"` (default). 63 | 64 | Returns: 65 | Tensor of the same shape and datatype as `x`. 66 | """ 67 | 68 | impl_dict = { 69 | 'ref': _fused_bias_act_ref, 70 | 'cuda': _fused_bias_act_cuda, 71 | } 72 | return impl_dict[impl](x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain, clamp=clamp) 73 | 74 | #---------------------------------------------------------------------------- 75 | 76 | def _fused_bias_act_ref(x, b, axis, act, alpha, gain, clamp): 77 | """Slow reference implementation of `fused_bias_act()` using standard TensorFlow ops.""" 78 | 79 | # Validate arguments. 80 | x = tf.convert_to_tensor(x) 81 | b = tf.convert_to_tensor(b) if b is not None else tf.constant([], dtype=x.dtype) 82 | act_spec = activation_funcs[act] 83 | assert b.shape.rank == 1 and (b.shape[0] == 0 or b.shape[0] == x.shape[axis]) 84 | assert b.shape[0] == 0 or 0 <= axis < x.shape.rank 85 | if alpha is None: 86 | alpha = act_spec.def_alpha 87 | if gain is None: 88 | gain = act_spec.def_gain 89 | 90 | # Add bias. 91 | if b.shape[0] != 0: 92 | x += tf.reshape(b, [-1 if i == axis else 1 for i in range(x.shape.rank)]) 93 | 94 | # Evaluate activation function. 95 | x = act_spec.func(x, alpha=alpha) 96 | 97 | # Scale by gain. 98 | if gain != 1: 99 | x *= gain 100 | 101 | # Clamp. 102 | if clamp is not None: 103 | clamp = np.asarray(clamp, dtype=x.dtype.name) 104 | assert clamp.shape == () and clamp >= 0 105 | x = tf.clip_by_value(x, -clamp, clamp) 106 | return x 107 | 108 | #---------------------------------------------------------------------------- 109 | 110 | def _fused_bias_act_cuda(x, b, axis, act, alpha, gain, clamp): 111 | """Fast CUDA implementation of `fused_bias_act()` using custom ops.""" 112 | 113 | # Validate arguments. 114 | x = tf.convert_to_tensor(x) 115 | empty_tensor = tf.constant([], dtype=x.dtype) 116 | b = tf.convert_to_tensor(b) if b is not None else empty_tensor 117 | act_spec = activation_funcs[act] 118 | assert b.shape.rank == 1 and (b.shape[0] == 0 or b.shape[0] == x.shape[axis]) 119 | assert b.shape[0] == 0 or 0 <= axis < x.shape.rank 120 | if alpha is None: 121 | alpha = act_spec.def_alpha 122 | if gain is None: 123 | gain = act_spec.def_gain 124 | 125 | # Special cases. 126 | if act == 'linear' and b is None and gain == 1.0: 127 | return x 128 | if act_spec.cuda_idx is None: 129 | return _fused_bias_act_ref(x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain, clamp=clamp) 130 | 131 | # CUDA op. 132 | cuda_op = _get_plugin().fused_bias_act 133 | cuda_kwargs = dict(axis=int(axis), act=int(act_spec.cuda_idx), gain=float(gain)) 134 | if alpha is not None: 135 | cuda_kwargs['alpha'] = float(alpha) 136 | if clamp is not None: 137 | clamp = np.asarray(clamp, dtype=x.dtype.name) 138 | assert clamp.shape == () and clamp >= 0 139 | cuda_kwargs['clamp'] = float(clamp.astype(np.float32)) 140 | def ref(tensor, name): 141 | return tensor if act_spec.ref == name else empty_tensor 142 | 143 | # Forward pass: y = func(x, b). 144 | def func_y(x, b): 145 | y = cuda_op(x=x, b=b, xref=empty_tensor, yref=empty_tensor, grad=0, **cuda_kwargs) 146 | y.set_shape(x.shape) 147 | return y 148 | 149 | # Backward pass: dx, db = grad(dy, x, y) 150 | def grad_dx(dy, x, y): 151 | dx = cuda_op(x=dy, b=empty_tensor, xref=ref(x,'x'), yref=ref(y,'y'), grad=1, **cuda_kwargs) 152 | dx.set_shape(x.shape) 153 | return dx 154 | def grad_db(dx): 155 | if b.shape[0] == 0: 156 | return empty_tensor 157 | db = dx 158 | if axis < x.shape.rank - 1: 159 | db = tf.reduce_sum(db, list(range(axis + 1, x.shape.rank))) 160 | if axis > 0: 161 | db = tf.reduce_sum(db, list(range(axis))) 162 | db.set_shape(b.shape) 163 | return db 164 | 165 | # Second order gradients: d_dy, d_x = grad2(d_dx, d_db, x, y) 166 | def grad2_d_dy(d_dx, d_db, x, y): 167 | d_dy = cuda_op(x=d_dx, b=d_db, xref=ref(x,'x'), yref=ref(y,'y'), grad=1, **cuda_kwargs) 168 | d_dy.set_shape(x.shape) 169 | return d_dy 170 | def grad2_d_x(d_dx, d_db, x, y): 171 | d_x = cuda_op(x=d_dx, b=d_db, xref=ref(x,'x'), yref=ref(y,'y'), grad=2, **cuda_kwargs) 172 | d_x.set_shape(x.shape) 173 | return d_x 174 | 175 | # Fast version for piecewise-linear activation funcs. 176 | @tf.custom_gradient 177 | def func_zero_2nd_grad(x, b): 178 | y = func_y(x, b) 179 | @tf.custom_gradient 180 | def grad(dy): 181 | dx = grad_dx(dy, x, y) 182 | db = grad_db(dx) 183 | def grad2(d_dx, d_db): 184 | d_dy = grad2_d_dy(d_dx, d_db, x, y) 185 | return d_dy 186 | return (dx, db), grad2 187 | return y, grad 188 | 189 | # Slow version for general activation funcs. 190 | @tf.custom_gradient 191 | def func_nonzero_2nd_grad(x, b): 192 | y = func_y(x, b) 193 | def grad_wrap(dy): 194 | @tf.custom_gradient 195 | def grad_impl(dy, x): 196 | dx = grad_dx(dy, x, y) 197 | db = grad_db(dx) 198 | def grad2(d_dx, d_db): 199 | d_dy = grad2_d_dy(d_dx, d_db, x, y) 200 | d_x = grad2_d_x(d_dx, d_db, x, y) 201 | return d_dy, d_x 202 | return (dx, db), grad2 203 | return grad_impl(dy, x) 204 | return y, grad_wrap 205 | 206 | # Which version to use? 207 | if act_spec.zero_2nd_grad: 208 | return func_zero_2nd_grad(x, b) 209 | return func_nonzero_2nd_grad(x, b) 210 | 211 | #---------------------------------------------------------------------------- 212 | -------------------------------------------------------------------------------- /src/training/dataset.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # This work is made available under the Nvidia Source Code License-NC 3 | # https://nvlabs.github.io/stylegan2/license.html 4 | """Multi-resolution input data pipeline.""" 5 | 6 | import os 7 | import glob 8 | import numpy as np 9 | import functools 10 | import tensorflow as tf 11 | import dnnlib 12 | import dnnlib.tflib as tflib 13 | 14 | from util.utilgan import calc_res, basename, file_list 15 | 16 | # Dataset class that loads data from tfrecords files. 17 | class TFRecordDataset: 18 | def __init__(self, 19 | tfrecord, # tfrecords file 20 | resolution=None, # Dataset resolution, None = autodetect. 21 | label_file=None, # Relative path of the labels file, None = autodetect. 22 | max_label_size=0, # 0 = no labels, 'full' = full labels, = N first label components. 23 | repeat=True, # Repeat dataset indefinitely? 24 | shuffle_mb=4096, # Shuffle data within specified window (megabytes), 0 = disable shuffling. 25 | prefetch_mb=2048, # Amount of data to prefetch (megabytes), 0 = disable prefetching. 26 | buffer_mb=256, # Read buffer size (megabytes). 27 | num_threads=2): # Number of concurrent threads. 28 | 29 | self.tfrecord = tfrecord 30 | self.resolution = None 31 | self.res_log2 = None 32 | self.shape = [] # [channels, height, width] 33 | self.dtype = 'uint8' 34 | self.dynamic_range = [0, 255] 35 | self.label_file = label_file 36 | self.label_size = None # components 37 | self.label_dtype = None 38 | self._np_labels = None 39 | self._tf_minibatch_in = None 40 | self._tf_labels_var = None 41 | self._tf_labels_dataset = None 42 | self._tf_dataset = None 43 | self._tf_iterator = None 44 | self._tf_init_op = None 45 | self._tf_minibatch_np = None 46 | self._cur_minibatch = -1 47 | 48 | # List tfrecords files and inspect their shapes. 49 | assert os.path.isfile(self.tfrecord) 50 | 51 | tfr_file = self.tfrecord 52 | data_dir = os.path.dirname(tfr_file) 53 | 54 | tfr_opt = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.NONE) 55 | for record in tf.python_io.tf_record_iterator(tfr_file, tfr_opt): # in fact only one 56 | tfr_shape = self.parse_tfrecord_shape(record) # [c,h,w] 57 | jpg_data = tfr_shape[0] < 4 58 | break 59 | 60 | # Autodetect label filename. 61 | if self.label_file is None: 62 | # guess = sorted(glob.glob(os.path.join(data_dir, '*.labels'))) 63 | guess = [ff for ff in file_list(data_dir, 'labels') if basename(ff).split('-')[0] == basename(tfrecord).split('-')[0]] 64 | if len(guess): 65 | self.label_file = guess[0] 66 | elif not os.path.isfile(self.label_file): 67 | guess = os.path.join(data_dir, self.label_file) 68 | if os.path.isfile(guess): 69 | self.label_file = guess 70 | 71 | # Determine shape and resolution 72 | self.shape = list(tfr_shape) 73 | max_res = calc_res(tfr_shape[1:]) 74 | self.resolution = resolution if resolution is not None else max_res 75 | self.res_log2 = int(np.ceil(np.log2(self.resolution))) 76 | self.init_res = [int(s * 2**(2-self.res_log2)) for s in self.shape[1:]] 77 | 78 | # Load labels 79 | assert max_label_size == 'full' or max_label_size >= 0 80 | self._np_labels = np.zeros([1<<30, 0], dtype=np.float32) 81 | if self.label_file is not None and max_label_size != 0: 82 | self._np_labels = np.load(self.label_file) 83 | assert self._np_labels.ndim == 2 84 | if max_label_size != 'full' and self._np_labels.shape[1] > max_label_size: 85 | self._np_labels = self._np_labels[:, :max_label_size] 86 | self.label_size = self._np_labels.shape[1] 87 | self.label_dtype = self._np_labels.dtype.name 88 | 89 | # Build TF expressions. 90 | with tf.name_scope('Dataset'), tf.device('/cpu:0'): 91 | self._tf_minibatch_in = tf.placeholder(tf.int64, name='minibatch_in', shape=[]) 92 | self._tf_labels_var = tflib.create_var_with_large_initial_value(self._np_labels, name='labels_var') 93 | self._tf_labels_dataset = tf.data.Dataset.from_tensor_slices(self._tf_labels_var) 94 | 95 | dset = tf.data.TFRecordDataset(tfr_file, compression_type='', buffer_size=buffer_mb << 20) 96 | if jpg_data is True: 97 | dset = dset.map(self.parse_tfrecord_tf_jpg, num_parallel_calls=num_threads) 98 | else: 99 | dset = dset.map(self.parse_tfrecord_tf, num_parallel_calls=num_threads) 100 | dset = tf.data.Dataset.zip((dset, self._tf_labels_dataset)) 101 | 102 | bytes_per_item = np.prod(tfr_shape) * np.dtype(self.dtype).itemsize 103 | if shuffle_mb > 0: 104 | dset = dset.shuffle(((shuffle_mb << 20) - 1) // bytes_per_item + 1) 105 | if repeat: 106 | dset = dset.repeat() 107 | if prefetch_mb > 0: 108 | dset = dset.prefetch(((prefetch_mb << 20) - 1) // bytes_per_item + 1) 109 | dset = dset.batch(self._tf_minibatch_in) 110 | self._tf_dataset = dset 111 | 112 | # self._tf_iterator = tf.data.Iterator.from_structure(tf.data.get_output_types(self._tf_dataset), tf.data.get_output_shapes(self._tf_dataset),) 113 | self._tf_iterator = tf.data.Iterator.from_structure(self._tf_dataset.output_types, self._tf_dataset.output_shapes) 114 | self._tf_init_op = self._tf_iterator.make_initializer(self._tf_dataset) 115 | 116 | def close(self): 117 | pass 118 | 119 | # Parse individual image from a JPG tfrecords file into NumPy array. 120 | @staticmethod 121 | def parse_tfrecord_tf_jpg(record): 122 | features = tf.parse_single_example(record, features={ 123 | "shape": tf.FixedLenFeature([3], tf.int64), 124 | "data": tf.FixedLenFeature([], tf.string),},) 125 | image = tf.image.decode_image(features['data']) 126 | return tf.transpose(image, [2,0,1]) 127 | # return tf.reshape(data, features["shape"]) 128 | 129 | # Parse image shape from a tfrecords file into NumPy array. 130 | @staticmethod 131 | def parse_tfrecord_shape(record): 132 | ex = tf.train.Example() 133 | ex.ParseFromString(record) 134 | shape = ex.features.feature['shape'].int64_list.value # pylint: disable=no-member 135 | return shape 136 | 137 | # Parse individual image from a tfrecords file into TensorFlow expression. 138 | @staticmethod 139 | def parse_tfrecord_tf(record): 140 | features = tf.parse_single_example(record, features={ 141 | 'shape': tf.FixedLenFeature([3], tf.int64), 142 | 'data': tf.FixedLenFeature([], tf.string)}) 143 | data = tf.decode_raw(features['data'], tf.uint8) 144 | return tf.reshape(data, features['shape']) 145 | 146 | # Parse individual image from a tfrecords file into NumPy array. 147 | @staticmethod 148 | def parse_tfrecord_np(record): 149 | ex = tf.train.Example() 150 | ex.ParseFromString(record) 151 | shape = ex.features.feature['shape'].int64_list.value # pylint: disable=no-member 152 | data = ex.features.feature['data'].bytes_list.value[0] # pylint: disable=no-member 153 | return np.fromstring(data, np.uint8).reshape(shape) 154 | 155 | # Use the given minibatch size and level-of-detail for the data returned by get_minibatch_tf(). 156 | def configure(self, minibatch_size): 157 | assert minibatch_size >= 1 158 | if self._cur_minibatch != minibatch_size: 159 | self._tf_init_op.run({self._tf_minibatch_in: minibatch_size}) 160 | self._cur_minibatch = minibatch_size 161 | 162 | # Get next minibatch as TensorFlow expressions 163 | def get_minibatch_tf(self): # => images, labels 164 | return self._tf_iterator.get_next() 165 | 166 | # Get next minibatch as NumPy arrays 167 | def get_minibatch_np(self, minibatch_size): # => images, labels 168 | self.configure(minibatch_size) 169 | with tf.name_scope('Dataset'): 170 | if self._tf_minibatch_np is None: 171 | self._tf_minibatch_np = self.get_minibatch_tf() 172 | return tflib.run(self._tf_minibatch_np) 173 | 174 | # Get random labels as TensorFlow expression. 175 | def get_random_labels_tf(self, minibatch_size): # => labels 176 | with tf.name_scope('Dataset'): 177 | if self.label_size > 0: 178 | with tf.device('/cpu:0'): 179 | return tf.gather(self._tf_labels_var, tf.random_uniform([minibatch_size], 0, self._np_labels.shape[0], dtype=tf.int32)) 180 | return tf.zeros([minibatch_size, 0], self.label_dtype) 181 | 182 | # Get random labels as NumPy array. 183 | def get_random_labels_np(self, minibatch_size): # => labels 184 | if self.label_size > 0: 185 | return self._np_labels[np.random.randint(self._np_labels.shape[0], size=[minibatch_size])] 186 | return np.zeros([minibatch_size, 0], self.label_dtype) 187 | 188 | # Helper func for constructing a dataset object using the given options. 189 | def load_dataset(class_name=None, data_dir=None, verbose=False, **kwargs): 190 | kwargs = dict(kwargs) 191 | if 'tfrecord' in kwargs: 192 | if class_name is None: 193 | class_name = __name__ + '.TFRecordDataset' 194 | if data_dir is not None: 195 | kwargs['tfrecord_dir'] = os.path.join(data_dir, kwargs['tfrecord_dir']) 196 | 197 | assert class_name is not None 198 | # if verbose: 199 | # print('Streaming data using %s...' % class_name) 200 | dataset = dnnlib.util.get_obj_by_name(class_name)(**kwargs) 201 | if verbose: 202 | print('Dataset shape =', np.int32(dataset.shape).tolist()) 203 | # print('Dynamic range =', dataset.dynamic_range) 204 | print('Label size =', dataset.label_size) 205 | return dataset 206 | 207 | -------------------------------------------------------------------------------- /src/dnnlib/tflib/tfutil.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | """Miscellaneous helper utils for Tensorflow.""" 8 | 9 | import os 10 | import numpy as np 11 | import tensorflow as tf 12 | 13 | # Silence deprecation warnings from TensorFlow 1.13 onwards 14 | import logging 15 | logging.getLogger('tensorflow').setLevel(logging.ERROR) 16 | import tensorflow.contrib # requires TensorFlow 1.x! 17 | tf.contrib = tensorflow.contrib 18 | 19 | from typing import Any, Iterable, List, Union 20 | 21 | TfExpression = Union[tf.Tensor, tf.Variable, tf.Operation] 22 | """A type that represents a valid Tensorflow expression.""" 23 | 24 | TfExpressionEx = Union[TfExpression, int, float, np.ndarray] 25 | """A type that can be converted to a valid Tensorflow expression.""" 26 | 27 | 28 | def run(*args, **kwargs) -> Any: 29 | """Run the specified ops in the default session.""" 30 | assert_tf_initialized() 31 | return tf.get_default_session().run(*args, **kwargs) 32 | 33 | 34 | def is_tf_expression(x: Any) -> bool: 35 | """Check whether the input is a valid Tensorflow expression, i.e., Tensorflow Tensor, Variable, or Operation.""" 36 | return isinstance(x, (tf.Tensor, tf.Variable, tf.Operation)) 37 | 38 | 39 | def shape_to_list(shape: Iterable[tf.Dimension]) -> List[Union[int, None]]: 40 | """Convert a Tensorflow shape to a list of ints. Retained for backwards compatibility -- use TensorShape.as_list() in new code.""" 41 | return [dim.value for dim in shape] 42 | 43 | 44 | def flatten(x: TfExpressionEx) -> TfExpression: 45 | """Shortcut function for flattening a tensor.""" 46 | with tf.name_scope("Flatten"): 47 | return tf.reshape(x, [-1]) 48 | 49 | 50 | def log2(x: TfExpressionEx) -> TfExpression: 51 | """Logarithm in base 2.""" 52 | with tf.name_scope("Log2"): 53 | return tf.log(x) * np.float32(1.0 / np.log(2.0)) 54 | 55 | 56 | def exp2(x: TfExpressionEx) -> TfExpression: 57 | """Exponent in base 2.""" 58 | with tf.name_scope("Exp2"): 59 | return tf.exp(x * np.float32(np.log(2.0))) 60 | 61 | 62 | def lerp(a: TfExpressionEx, b: TfExpressionEx, t: TfExpressionEx) -> TfExpressionEx: 63 | """Linear interpolation.""" 64 | with tf.name_scope("Lerp"): 65 | return a + (b - a) * t 66 | 67 | 68 | def lerp_clip(a: TfExpressionEx, b: TfExpressionEx, t: TfExpressionEx) -> TfExpression: 69 | """Linear interpolation with clip.""" 70 | with tf.name_scope("LerpClip"): 71 | return a + (b - a) * tf.clip_by_value(t, 0.0, 1.0) 72 | 73 | 74 | def absolute_name_scope(scope: str) -> tf.name_scope: 75 | """Forcefully enter the specified name scope, ignoring any surrounding scopes.""" 76 | return tf.name_scope(scope + "/") 77 | 78 | 79 | def absolute_variable_scope(scope: str, **kwargs) -> tf.variable_scope: 80 | """Forcefully enter the specified variable scope, ignoring any surrounding scopes.""" 81 | return tf.variable_scope(tf.VariableScope(name=scope, **kwargs), auxiliary_name_scope=False) 82 | 83 | 84 | def _sanitize_tf_config(config_dict: dict = None) -> dict: 85 | # Defaults. 86 | cfg = dict() 87 | cfg["rnd.np_random_seed"] = None # Random seed for NumPy. None = keep as is. 88 | cfg["rnd.tf_random_seed"] = "auto" # Random seed for TensorFlow. 'auto' = derive from NumPy random state. None = keep as is. 89 | cfg["env.TF_CPP_MIN_LOG_LEVEL"] = "1" # 0 = Print all available debug info from TensorFlow. 1 = Print warnings and errors, but disable debug info. 90 | cfg["graph_options.place_pruned_graph"] = True # False = Check that all ops are available on the designated device. True = Skip the check for ops that are not used. 91 | cfg["gpu_options.allow_growth"] = True # False = Allocate all GPU memory at the beginning. True = Allocate only as much GPU memory as needed. 92 | 93 | # Remove defaults for environment variables that are already set. 94 | for key in list(cfg): 95 | fields = key.split(".") 96 | if fields[0] == "env": 97 | assert len(fields) == 2 98 | if fields[1] in os.environ: 99 | del cfg[key] 100 | 101 | # User overrides. 102 | if config_dict is not None: 103 | cfg.update(config_dict) 104 | return cfg 105 | 106 | 107 | def init_tf(config_dict: dict = None) -> None: 108 | """Initialize TensorFlow session using good default settings.""" 109 | # Skip if already initialized. 110 | if tf.get_default_session() is not None: 111 | return 112 | 113 | # Setup config dict and random seeds. 114 | cfg = _sanitize_tf_config(config_dict) 115 | np_random_seed = cfg["rnd.np_random_seed"] 116 | if np_random_seed is not None: 117 | np.random.seed(np_random_seed) 118 | tf_random_seed = cfg["rnd.tf_random_seed"] 119 | if tf_random_seed == "auto": 120 | tf_random_seed = np.random.randint(1 << 31) 121 | if tf_random_seed is not None: 122 | tf.set_random_seed(tf_random_seed) 123 | 124 | # Setup environment variables. 125 | for key, value in cfg.items(): 126 | fields = key.split(".") 127 | if fields[0] == "env": 128 | assert len(fields) == 2 129 | os.environ[fields[1]] = str(value) 130 | 131 | # Create default TensorFlow session. 132 | sess = create_session(cfg, force_as_default=True) 133 | return sess 134 | 135 | 136 | def assert_tf_initialized(): 137 | """Check that TensorFlow session has been initialized.""" 138 | if tf.get_default_session() is None: 139 | raise RuntimeError("No default TensorFlow session found. Please call dnnlib.tflib.init_tf().") 140 | 141 | 142 | def create_session(config_dict: dict = None, force_as_default: bool = False) -> tf.Session: 143 | """Create tf.Session based on config dict.""" 144 | # Setup TensorFlow config proto. 145 | cfg = _sanitize_tf_config(config_dict) 146 | config_proto = tf.ConfigProto() 147 | for key, value in cfg.items(): 148 | fields = key.split(".") 149 | if fields[0] not in ["rnd", "env"]: 150 | obj = config_proto 151 | for field in fields[:-1]: 152 | obj = getattr(obj, field) 153 | setattr(obj, fields[-1], value) 154 | 155 | # Create session. 156 | session = tf.Session(config=config_proto) 157 | if force_as_default: 158 | # pylint: disable=protected-access 159 | session._default_session = session.as_default() 160 | session._default_session.enforce_nesting = False 161 | session._default_session.__enter__() 162 | return session 163 | 164 | 165 | def init_uninitialized_vars(target_vars: List[tf.Variable] = None) -> None: 166 | """Initialize all tf.Variables that have not already been initialized. 167 | 168 | Equivalent to the following, but more efficient and does not bloat the tf graph: 169 | tf.variables_initializer(tf.report_uninitialized_variables()).run() 170 | """ 171 | assert_tf_initialized() 172 | if target_vars is None: 173 | target_vars = tf.global_variables() 174 | 175 | test_vars = [] 176 | test_ops = [] 177 | 178 | with tf.control_dependencies(None): # ignore surrounding control_dependencies 179 | for var in target_vars: 180 | assert is_tf_expression(var) 181 | 182 | try: 183 | tf.get_default_graph().get_tensor_by_name(var.name.replace(":0", "/IsVariableInitialized:0")) 184 | except KeyError: 185 | # Op does not exist => variable may be uninitialized. 186 | test_vars.append(var) 187 | 188 | with absolute_name_scope(var.name.split(":")[0]): 189 | test_ops.append(tf.is_variable_initialized(var)) 190 | 191 | init_vars = [var for var, inited in zip(test_vars, run(test_ops)) if not inited] 192 | run([var.initializer for var in init_vars]) 193 | 194 | 195 | def set_vars(var_to_value_dict: dict) -> None: 196 | """Set the values of given tf.Variables. 197 | 198 | Equivalent to the following, but more efficient and does not bloat the tf graph: 199 | tflib.run([tf.assign(var, value) for var, value in var_to_value_dict.items()] 200 | """ 201 | assert_tf_initialized() 202 | ops = [] 203 | feed_dict = {} 204 | 205 | for var, value in var_to_value_dict.items(): 206 | assert is_tf_expression(var) 207 | 208 | try: 209 | setter = tf.get_default_graph().get_tensor_by_name(var.name.replace(":0", "/setter:0")) # look for existing op 210 | except KeyError: 211 | with absolute_name_scope(var.name.split(":")[0]): 212 | with tf.control_dependencies(None): # ignore surrounding control_dependencies 213 | setter = tf.assign(var, tf.placeholder(var.dtype, var.shape, "new_value"), name="setter") # create new setter 214 | 215 | ops.append(setter) 216 | feed_dict[setter.op.inputs[1]] = value 217 | 218 | run(ops, feed_dict) 219 | 220 | 221 | def create_var_with_large_initial_value(initial_value: np.ndarray, *args, **kwargs): 222 | """Create tf.Variable with large initial value without bloating the tf graph.""" 223 | assert_tf_initialized() 224 | assert isinstance(initial_value, np.ndarray) 225 | zeros = tf.zeros(initial_value.shape, initial_value.dtype) 226 | var = tf.Variable(zeros, *args, **kwargs) 227 | set_vars({var: initial_value}) 228 | return var 229 | 230 | 231 | def convert_images_from_uint8(images, drange=[-1,1], nhwc_to_nchw=False): 232 | """Convert a minibatch of images from uint8 to float32 with configurable dynamic range. 233 | Can be used as an input transformation for Network.run(). 234 | """ 235 | images = tf.cast(images, tf.float32) 236 | if nhwc_to_nchw: 237 | images = tf.transpose(images, [0, 3, 1, 2]) 238 | return images * ((drange[1] - drange[0]) / 255) + drange[0] 239 | 240 | 241 | def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False, shrink=1, uint8_cast=True): 242 | """Convert a minibatch of images from float32 to uint8 with configurable dynamic range. 243 | Can be used as an output transformation for Network.run(). 244 | """ 245 | images = tf.cast(images, tf.float32) 246 | if shrink > 1: 247 | ksize = [1, 1, shrink, shrink] 248 | images = tf.nn.avg_pool(images, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW") 249 | if nchw_to_nhwc: 250 | images = tf.transpose(images, [0, 2, 3, 1]) 251 | scale = 255 / (drange[1] - drange[0]) 252 | images = images * scale + (0.5 - drange[0] * scale) 253 | if uint8_cast: 254 | images = tf.saturate_cast(images, tf.uint8) 255 | return images 256 | -------------------------------------------------------------------------------- /src/training/loss.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # This work is made available under the Nvidia Source Code License-NC. 3 | # https://nvlabs.github.io/stylegan2/license.html 4 | """Loss functions.""" 5 | 6 | import numpy as np 7 | import tensorflow as tf 8 | 9 | import dnnlib.tflib as tflib 10 | from dnnlib.tflib.autosummary import autosummary 11 | 12 | #---------------------------------------------------------------------------- 13 | # Logistic loss from the paper 14 | # "Generative Adversarial Nets", Goodfellow et al. 2014 15 | 16 | def G_logistic(G, D, opt, training_set, minibatch_size): 17 | _ = opt 18 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 19 | labels = training_set.get_random_labels_tf(minibatch_size) 20 | fake_images_out = G.get_output_for(latents, labels, is_training=True) 21 | fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True) 22 | loss = -tf.nn.softplus(fake_scores_out) # log(1-sigmoid(fake_scores_out)) # pylint: disable=invalid-unary-operand-type 23 | return loss, None 24 | 25 | def G_logistic_ns(G, D, opt, training_set, minibatch_size): 26 | _ = opt 27 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 28 | labels = training_set.get_random_labels_tf(minibatch_size) 29 | fake_images_out = G.get_output_for(latents, labels, is_training=True) 30 | fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True) 31 | loss = tf.nn.softplus(-fake_scores_out) # -log(sigmoid(fake_scores_out)) 32 | return loss, None 33 | 34 | def D_logistic(G, D, opt, training_set, minibatch_size, reals, labels): 35 | _ = opt, training_set 36 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 37 | fake_images_out = G.get_output_for(latents, labels, is_training=True) 38 | real_scores_out = D.get_output_for(reals, labels, is_training=True) 39 | fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True) 40 | real_scores_out = autosummary('Loss/scores/real', real_scores_out) 41 | fake_scores_out = autosummary('Loss/scores/fake', fake_scores_out) 42 | loss = tf.nn.softplus(fake_scores_out) # -log(1-sigmoid(fake_scores_out)) 43 | loss += tf.nn.softplus(-real_scores_out) # -log(sigmoid(real_scores_out)) # pylint: disable=invalid-unary-operand-type 44 | return loss, None 45 | 46 | #---------------------------------------------------------------------------- 47 | # R1 and R2 regularizers from the paper 48 | # "Which Training Methods for GANs do actually Converge?", Mescheder et al. 2018 49 | 50 | def D_logistic_r1(G, D, opt, training_set, minibatch_size, reals, labels, gamma=10.0): 51 | _ = opt, training_set 52 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 53 | fake_images_out = G.get_output_for(latents, labels, is_training=True) 54 | real_scores_out = D.get_output_for(reals, labels, is_training=True) 55 | fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True) 56 | real_scores_out = autosummary('Loss/scores/real', real_scores_out) 57 | fake_scores_out = autosummary('Loss/scores/fake', fake_scores_out) 58 | loss = tf.nn.softplus(fake_scores_out) # -log(1-sigmoid(fake_scores_out)) 59 | loss += tf.nn.softplus(-real_scores_out) # -log(sigmoid(real_scores_out)) # pylint: disable=invalid-unary-operand-type 60 | 61 | with tf.name_scope('GradientPenalty'): 62 | real_grads = tf.gradients(tf.reduce_sum(real_scores_out), [reals])[0] 63 | gradient_penalty = tf.reduce_sum(tf.square(real_grads), axis=[1,2,3]) 64 | gradient_penalty = autosummary('Loss/gradient_penalty', gradient_penalty) 65 | reg = gradient_penalty * (gamma * 0.5) 66 | return loss, reg 67 | 68 | def D_logistic_r2(G, D, opt, training_set, minibatch_size, reals, labels, gamma=10.0): 69 | _ = opt, training_set 70 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 71 | fake_images_out = G.get_output_for(latents, labels, is_training=True) 72 | real_scores_out = D.get_output_for(reals, labels, is_training=True) 73 | fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True) 74 | real_scores_out = autosummary('Loss/scores/real', real_scores_out) 75 | fake_scores_out = autosummary('Loss/scores/fake', fake_scores_out) 76 | loss = tf.nn.softplus(fake_scores_out) # -log(1-sigmoid(fake_scores_out)) 77 | loss += tf.nn.softplus(-real_scores_out) # -log(sigmoid(real_scores_out)) # pylint: disable=invalid-unary-operand-type 78 | 79 | with tf.name_scope('GradientPenalty'): 80 | fake_grads = tf.gradients(tf.reduce_sum(fake_scores_out), [fake_images_out])[0] 81 | gradient_penalty = tf.reduce_sum(tf.square(fake_grads), axis=[1,2,3]) 82 | gradient_penalty = autosummary('Loss/gradient_penalty', gradient_penalty) 83 | reg = gradient_penalty * (gamma * 0.5) 84 | return loss, reg 85 | 86 | #---------------------------------------------------------------------------- 87 | # WGAN loss from the paper 88 | # "Wasserstein Generative Adversarial Networks", Arjovsky et al. 2017 89 | 90 | def G_wgan(G, D, opt, training_set, minibatch_size): 91 | _ = opt 92 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 93 | labels = training_set.get_random_labels_tf(minibatch_size) 94 | fake_images_out = G.get_output_for(latents, labels, is_training=True) 95 | fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True) 96 | loss = -fake_scores_out 97 | return loss, None 98 | 99 | def D_wgan(G, D, opt, training_set, minibatch_size, reals, labels, wgan_epsilon=0.001): 100 | _ = opt, training_set 101 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 102 | fake_images_out = G.get_output_for(latents, labels, is_training=True) 103 | real_scores_out = D.get_output_for(reals, labels, is_training=True) 104 | fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True) 105 | real_scores_out = autosummary('Loss/scores/real', real_scores_out) 106 | fake_scores_out = autosummary('Loss/scores/fake', fake_scores_out) 107 | loss = fake_scores_out - real_scores_out 108 | with tf.name_scope('EpsilonPenalty'): 109 | epsilon_penalty = autosummary('Loss/epsilon_penalty', tf.square(real_scores_out)) 110 | loss += epsilon_penalty * wgan_epsilon 111 | return loss, None 112 | 113 | #---------------------------------------------------------------------------- 114 | # WGAN-GP loss from the paper 115 | # "Improved Training of Wasserstein GANs", Gulrajani et al. 2017 116 | 117 | def D_wgan_gp(G, D, opt, training_set, minibatch_size, reals, labels, wgan_lambda=10.0, wgan_epsilon=0.001, wgan_target=1.0): 118 | _ = opt, training_set 119 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 120 | fake_images_out = G.get_output_for(latents, labels, is_training=True) 121 | real_scores_out = D.get_output_for(reals, labels, is_training=True) 122 | fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True) 123 | real_scores_out = autosummary('Loss/scores/real', real_scores_out) 124 | fake_scores_out = autosummary('Loss/scores/fake', fake_scores_out) 125 | loss = fake_scores_out - real_scores_out 126 | with tf.name_scope('EpsilonPenalty'): 127 | epsilon_penalty = autosummary('Loss/epsilon_penalty', tf.square(real_scores_out)) 128 | loss += epsilon_penalty * wgan_epsilon 129 | 130 | with tf.name_scope('GradientPenalty'): 131 | mixing_factors = tf.random_uniform([minibatch_size, 1, 1, 1], 0.0, 1.0, dtype=fake_images_out.dtype) 132 | mixed_images_out = tflib.lerp(tf.cast(reals, fake_images_out.dtype), fake_images_out, mixing_factors) 133 | mixed_scores_out = D.get_output_for(mixed_images_out, labels, is_training=True) 134 | mixed_scores_out = autosummary('Loss/scores/mixed', mixed_scores_out) 135 | mixed_grads = tf.gradients(tf.reduce_sum(mixed_scores_out), [mixed_images_out])[0] 136 | mixed_norms = tf.sqrt(tf.reduce_sum(tf.square(mixed_grads), axis=[1,2,3])) 137 | mixed_norms = autosummary('Loss/mixed_norms', mixed_norms) 138 | gradient_penalty = tf.square(mixed_norms - wgan_target) 139 | reg = gradient_penalty * (wgan_lambda / (wgan_target**2)) 140 | return loss, reg 141 | 142 | #---------------------------------------------------------------------------- 143 | # Non-saturating logistic loss with path length regularizer from the paper 144 | # "Analyzing and Improving the Image Quality of StyleGAN", Karras et al. 2019 145 | 146 | def G_logistic_ns_pathreg(G, D, opt, training_set, minibatch_size, pl_minibatch_shrink=2, pl_decay=0.01, pl_weight=2.0): 147 | _ = opt 148 | latents = tf.random_normal([minibatch_size] + G.input_shapes[0][1:]) 149 | labels = training_set.get_random_labels_tf(minibatch_size) 150 | fake_images_out, fake_dlatents_out = G.get_output_for(latents, labels, is_training=True, return_dlatents=True) 151 | fake_scores_out = D.get_output_for(fake_images_out, labels, is_training=True) 152 | loss = tf.nn.softplus(-fake_scores_out) # -log(sigmoid(fake_scores_out)) 153 | 154 | # Path length regularization. 155 | with tf.name_scope('PathReg'): 156 | 157 | # Evaluate the regularization term using a smaller minibatch to conserve memory. 158 | if pl_minibatch_shrink > 1: 159 | pl_minibatch = minibatch_size // pl_minibatch_shrink 160 | pl_latents = tf.random_normal([pl_minibatch] + G.input_shapes[0][1:]) 161 | pl_labels = training_set.get_random_labels_tf(pl_minibatch) 162 | fake_images_out, fake_dlatents_out = G.get_output_for(pl_latents, pl_labels, is_training=True, return_dlatents=True) 163 | 164 | # Compute |J*y|. 165 | pl_noise = tf.random_normal(tf.shape(fake_images_out)) / np.sqrt(np.prod(G.output_shape[2:])) 166 | pl_grads = tf.gradients(tf.reduce_sum(fake_images_out * pl_noise), [fake_dlatents_out])[0] 167 | pl_lengths = tf.sqrt(tf.reduce_mean(tf.reduce_sum(tf.square(pl_grads), axis=2), axis=1)) 168 | pl_lengths = autosummary('Loss/pl_lengths', pl_lengths) 169 | 170 | # Track exponential moving average of |J*y|. 171 | with tf.control_dependencies(None): 172 | pl_mean_var = tf.Variable(name='pl_mean', trainable=False, initial_value=0.0, dtype=tf.float32) 173 | pl_mean = pl_mean_var + pl_decay * (tf.reduce_mean(pl_lengths) - pl_mean_var) 174 | pl_update = tf.assign(pl_mean_var, pl_mean) 175 | 176 | # Calculate (|J*y|-a)^2. 177 | with tf.control_dependencies([pl_update]): 178 | pl_penalty = tf.square(pl_lengths - pl_mean) 179 | pl_penalty = autosummary('Loss/pl_penalty', pl_penalty) 180 | 181 | # Apply weight. 182 | # 183 | # Note: The division in pl_noise decreases the weight by num_pixels, and the reduce_mean 184 | # in pl_lengths decreases it by num_affine_layers. The effective weight then becomes: 185 | # 186 | # gamma_pl = pl_weight / num_pixels / num_affine_layers 187 | # = 2 / (r^2) / (log2(r) * 2 - 2) 188 | # = 1 / (r^2 * (log2(r) - 1)) 189 | # = ln(2) / (r^2 * (ln(r) - ln(2)) 190 | # 191 | reg = pl_penalty * pl_weight 192 | 193 | return loss, reg 194 | 195 | #---------------------------------------------------------------------------- 196 | -------------------------------------------------------------------------------- /src/model_convert.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 3 | import warnings 4 | warnings.filterwarnings("ignore") 5 | import argparse 6 | import pickle 7 | from collections import OrderedDict 8 | import numpy as np 9 | 10 | import dnnlib 11 | import dnnlib.tflib as tflib 12 | from dnnlib.tflib import tfutil 13 | 14 | import tensorflow; tf = tensorflow.compat.v1 if hasattr(tensorflow.compat, 'v1') else tensorflow 15 | tf.logging.set_verbosity(tf.logging.ERROR) 16 | 17 | from util.utilgan import basename, calc_init_res 18 | try: # progress bar for notebooks 19 | get_ipython().__class__.__name__ 20 | from util.progress_bar import ProgressIPy as ProgressBar 21 | except: # normal console 22 | from util.progress_bar import ProgressBar 23 | 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument('--source', required=True, help='Source model path') 26 | parser.add_argument('--out_dir', default='./', help='Output directory for reduced/reconstructed model') 27 | parser.add_argument('-r', '--reconstruct', action='store_true', help='Reconstruct model (add internal arguments)') 28 | parser.add_argument('-s', '--res', default=None, help='Target resolution in format X-Y') 29 | parser.add_argument('-a', '--alpha', action='store_true', help='Add alpha channel for RGBA processing') 30 | parser.add_argument('-l', '--labels', default=None, type=int, help='Labels for conditional model') 31 | parser.add_argument('-f', '--full', action='store_true', help='Save full model') 32 | parser.add_argument('-v', '--verbose', action='store_true') 33 | a = parser.parse_args() 34 | 35 | if a.res is not None: 36 | a.res = [int(s) for s in a.res.split('-')][::-1] 37 | if len(a.res) == 1: a.res = a.res + a.res 38 | 39 | def load_pkl(filepath): 40 | with open(filepath, 'rb') as f: 41 | networks = pickle.load(f, encoding='latin1') 42 | try: 43 | G, D, Gs = networks 44 | except: 45 | Gs = networks 46 | G = D = None 47 | return G, D, Gs 48 | 49 | def save_pkl(networks, filepath): 50 | with open(filepath, 'wb') as file: 51 | pickle.dump(networks, file, protocol=pickle.HIGHEST_PROTOCOL) 52 | 53 | def create_model(data_shape, full=False, labels=None, kwargs_in=None): 54 | init_res, resolution, res_log2 = calc_init_res(data_shape[1:]) 55 | kwargs_out = dnnlib.EasyDict() 56 | kwargs_out.num_channels = data_shape[0] 57 | if kwargs_in is not None: 58 | for k in list(kwargs_in.keys()): 59 | kwargs_out[k] = kwargs_in[k] 60 | if labels is not None: kwargs_out.label_size = labels 61 | kwargs_out.resolution = resolution 62 | kwargs_out.init_res = init_res 63 | if a.verbose is True: print(['%s: %s'%(kv[0],kv[1]) for kv in sorted(kwargs_out.items())]) 64 | if full is True: 65 | G = tflib.Network('G', func_name='training.networks_stylegan2.G_main', **kwargs_out) 66 | D = tflib.Network('D', func_name='training.networks_stylegan2.D_stylegan2', **kwargs_out) 67 | Gs = G.clone('Gs') 68 | else: 69 | Gs = tflib.Network('Gs', func_name='training.networks_stylegan2.G_main', **kwargs_out) 70 | G = D = None 71 | return G, D, Gs 72 | 73 | def copy_vars(src_net, tgt_net, D=False): 74 | names = [name for name in tgt_net.trainables.keys() if name in src_net.trainables.keys()] 75 | var_dict = OrderedDict() 76 | 77 | for name in names: 78 | if tgt_net.vars[name].shape == src_net.vars[name].shape: # fixing rgb-to-rgba only !! 79 | var_dict[name] = src_net.vars[name] 80 | else: 81 | var_dict[name] = add_channel(src_net.vars[name], D=D) 82 | 83 | weights_to_copy = {tgt_net.vars[name]: var_dict[name] for name in names} 84 | tfutil.set_vars(tfutil.run(weights_to_copy)) 85 | 86 | def add_channel(x, D=False): # [BCHW] 87 | if D is True: # pad dim before last [-2] 88 | padding = [[0,1],[0,0]] 89 | for i in range(len(x.shape)-2): 90 | padding.insert(0, [0,0]) 91 | else: # pad last dim [-1] 92 | padding = [[0,1]] 93 | for i in range(len(x.shape)-1): 94 | padding.insert(0, [0,0]) 95 | y = tf.pad(x, padding, 'symmetric') # symmetric reflect 96 | return y 97 | 98 | ### [edited] from https://github.com/aydao/stylegan2-surgery/blob/master/copy_crop_weights.py 99 | 100 | def copy_and_crop_or_pad_trainables(src_net, tgt_net) -> None: 101 | source_trainables = src_net.trainables.keys() 102 | target_trainables = tgt_net.trainables.keys() 103 | names = [pair for pair in zip(source_trainables, target_trainables)] 104 | 105 | skip = [] 106 | pbar = ProgressBar(len(names)) 107 | for pair in names: 108 | source_name, target_name = pair 109 | log = source_name 110 | x = src_net.get_var(source_name) 111 | y = tgt_net.get_var(target_name) 112 | source_shape = x.shape 113 | target_shape = y.shape 114 | if source_shape != target_shape: 115 | update = x 116 | index = None 117 | if 'Dense' in source_name: 118 | if source_shape[0] > target_shape[0]: 119 | gap = source_shape[0] - target_shape[0] 120 | start = abs(gap) // 2 121 | end = start + target_shape[0] 122 | update = update[start:end,:] 123 | else: 124 | update = pad_symm_np(update, target_shape) 125 | log = (log, source_shape, '=>', target_shape) 126 | else: 127 | try: 128 | if source_shape[2] > target_shape[2]: 129 | index = 2 130 | gap = source_shape[index] - target_shape[index] 131 | start = abs(gap) // 2 132 | end = start + target_shape[index] 133 | update = update[:,:,start:end,:] 134 | if source_shape[3] > target_shape[3]: 135 | index = 3 136 | gap = source_shape[index] - target_shape[index] 137 | start = abs(gap) // 2 138 | end = start + target_shape[index] 139 | update = update[:,:,:,start:end] 140 | except: 141 | print(' Wrong var pair?', source_name, source_shape, target_name, target_shape) 142 | exit(1) 143 | 144 | if source_shape[2] < target_shape[2] or source_shape[3] < target_shape[3]: 145 | update = pad_symm_np(update, target_shape[2:]) 146 | log = (log, source_shape, '=>', target_shape) 147 | # print(pair, source_shape, target_shape) 148 | 149 | tgt_net.set_var(target_name, update) 150 | skip.append(source_name) 151 | pbar.upd(pair) 152 | 153 | weights_to_copy = {tgt_net.vars[pair[1]]: src_net.vars[pair[0]] for pair in names if pair[0] not in skip} 154 | tfutil.set_vars(tfutil.run(weights_to_copy)) 155 | 156 | def pad_symm_np(x, size): 157 | sh = x.shape[-len(size):] 158 | padding = [[0,0]] * (len(x.shape)-len(size)) 159 | for i, s in enumerate(size): 160 | p0 = (s-sh[i]) // 2 161 | p1 = s-sh[i] - p0 162 | padding.append([p0,p1]) 163 | return np.pad(x, padding, mode='symmetric') 164 | 165 | def copy_and_fill_trainables(src_net, tgt_net) -> None: # model => conditional 166 | train_vars = [name for name in src_net.trainables.keys() if name in tgt_net.trainables.keys()] 167 | skip = [] 168 | pbar = ProgressBar(len(train_vars)) 169 | for name in train_vars: 170 | x = src_net.get_var(name) 171 | y = tgt_net.get_var(name) 172 | src_shape = x.shape 173 | tgt_shape = y.shape 174 | if src_shape != tgt_shape: 175 | assert len(src_shape) == len(tgt_shape), "Different shapes: %s %s" % (str(src_shape), str(tgt_shape)) 176 | if np.less(tgt_shape, src_shape).any(): # kill labels: [1024,512] => [512,512] 177 | try: 178 | update = x[:tgt_shape[0], :tgt_shape[1], ...] # !!! corrects only first two dims 179 | except: 180 | update = x[:tgt_shape[0]] 181 | elif np.greater(tgt_shape, src_shape).any(): # add labels: [512,512] => [1024,512] 182 | tile_count = [tgt_shape[i] // src_shape[i] for i in range(len(src_shape))] 183 | if a.verbose is True: print(name, tile_count, src_shape, '=>', tgt_shape, '\n\n') # G_mapping/Dense0, D/Output 184 | update = np.tile(x, tile_count) 185 | tgt_net.set_var(name, update) 186 | skip.append(name) 187 | pbar.upd(name) 188 | weights_to_copy = {tgt_net.vars[name]: src_net.vars[name] for name in train_vars if name not in skip} 189 | tfutil.set_vars(tfutil.run(weights_to_copy)) 190 | 191 | 192 | def main(): 193 | tflib.init_tf({'allow_soft_placement':True}) 194 | 195 | G_in, D_in, Gs_in = load_pkl(a.source) 196 | print(' Loading model', a.source, Gs_in.output_shape) 197 | _, res_in, _ = calc_init_res(Gs_in.output_shape[1:]) 198 | 199 | if a.res is not None or a.alpha is True: 200 | if a.res is None: a.res = Gs_in.output_shape[2:] 201 | colors = 4 if a.alpha is True else Gs_in.output_shape[1] # EXPERIMENTAL 202 | _, res_out, _ = calc_init_res([colors, *a.res]) 203 | 204 | if res_in != res_out or a.alpha is True: # add or remove layers 205 | assert G_in is not None and D_in is not None, " !! G/D subnets not found in source model !!" 206 | data_shape = [colors, res_out, res_out] 207 | print(' Reconstructing full model with shape', data_shape) 208 | G_out, D_out, Gs_out = create_model(data_shape, True, 0, Gs_in.static_kwargs) 209 | copy_vars(Gs_in, Gs_out) 210 | copy_vars(G_in, G_out) 211 | copy_vars(D_in, D_out, D=True) 212 | G_in, D_in, Gs_in = G_out, D_out, Gs_out 213 | a.full = True 214 | 215 | if a.res[0] != res_out or a.res[1] != res_out: # crop or pad layers 216 | data_shape = [colors, *a.res] 217 | G_out, D_out, Gs_out = create_model(data_shape, True, 0, Gs_in.static_kwargs) 218 | if G_in is not None and D_in is not None: 219 | print(' Reconstructing full model with shape', data_shape) 220 | copy_and_crop_or_pad_trainables(G_in, G_out) 221 | copy_and_crop_or_pad_trainables(D_in, D_out) 222 | G_in, D_in = G_out, D_out 223 | a.full = True 224 | else: 225 | print(' Reconstructing Gs model with shape', data_shape) 226 | copy_and_crop_or_pad_trainables(Gs_in, Gs_out) 227 | Gs_in = Gs_out 228 | 229 | if a.labels is not None: 230 | assert G_in is not None and D_in is not None, " !! G/D subnets not found in source model !!" 231 | print(' Reconstructing full model with labels', a.labels) 232 | data_shape = Gs_in.output_shape[1:] 233 | G_out, D_out, Gs_out = create_model(data_shape, True, a.labels, Gs_in.static_kwargs) 234 | if a.verbose is True: D_out.print_layers() 235 | if a.verbose is True: G_out.print_layers() 236 | copy_and_fill_trainables(G_in, G_out) 237 | copy_and_fill_trainables(D_in, D_out) 238 | copy_and_fill_trainables(Gs_in, Gs_out) 239 | a.full = True 240 | 241 | if a.labels is None and a.res is None and a.alpha is not True: 242 | if a.reconstruct is True: 243 | print(' Reconstructing model with same size /', 'full' if a.full else 'Gs') 244 | data_shape = Gs_in.output_shape[1:] 245 | G_out, D_out, Gs_out = create_model(data_shape, a.full, 0, Gs_in.static_kwargs) 246 | Gs_out.copy_vars_from(Gs_in) 247 | if a.full is True and G_in is not None and D_in is not None: 248 | G_out.copy_vars_from(G_in) 249 | D_out.copy_vars_from(D_in) 250 | else: 251 | Gs_out = Gs_in 252 | 253 | out_name = basename(a.source) 254 | if a.res is not None: out_name += '-%dx%d' % (a.res[1], a.res[0]) 255 | if a.alpha is True: out_name += 'a' 256 | if a.labels is not None: out_name += '-c%d' % a.labels 257 | 258 | if a.full is True: # G_in is not None and D_in is not None 259 | save_pkl((G_out, D_out, Gs_out), os.path.join(a.out_dir, '%s.pkl' % out_name)) 260 | else: 261 | save_pkl(Gs_out, os.path.join(a.out_dir, '%s-Gs.pkl' % out_name)) 262 | 263 | print(' Done') 264 | 265 | 266 | if __name__ == '__main__': 267 | main() 268 | -------------------------------------------------------------------------------- /src/dnnlib/submission/submit.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2019, NVIDIA Corporation. All rights reserved. 2 | # 3 | # This work is made available under the Nvidia Source Code License-NC. 4 | # To view a copy of this license, visit 5 | # https://nvlabs.github.io/stylegan2/license.html 6 | 7 | """Submit a function to be run either locally or in a computing cluster.""" 8 | 9 | import copy 10 | import inspect 11 | import os 12 | import pathlib 13 | import pickle 14 | import platform 15 | import pprint 16 | import re 17 | import shutil 18 | import sys 19 | import time 20 | import traceback 21 | 22 | from enum import Enum 23 | 24 | from .. import util 25 | from ..util import EasyDict 26 | 27 | from . import internal 28 | 29 | class SubmitTarget(Enum): 30 | """The target where the function should be run. 31 | 32 | LOCAL: Run it locally. 33 | """ 34 | LOCAL = 1 35 | 36 | 37 | class PathType(Enum): 38 | """Determines in which format should a path be formatted. 39 | 40 | WINDOWS: Format with Windows style. 41 | LINUX: Format with Linux/Posix style. 42 | AUTO: Use current OS type to select either WINDOWS or LINUX. 43 | """ 44 | WINDOWS = 1 45 | LINUX = 2 46 | AUTO = 3 47 | 48 | 49 | class PlatformExtras: 50 | """A mixed bag of values used by dnnlib heuristics. 51 | 52 | Attributes: 53 | 54 | data_reader_buffer_size: Used by DataReader to size internal shared memory buffers. 55 | data_reader_process_count: Number of worker processes to spawn (zero for single thread operation) 56 | """ 57 | def __init__(self): 58 | self.data_reader_buffer_size = 1<<30 # 1 GB 59 | self.data_reader_process_count = 0 # single threaded default 60 | 61 | 62 | _user_name_override = None 63 | 64 | class SubmitConfig(util.EasyDict): 65 | """Strongly typed config dict needed to submit runs. 66 | 67 | Attributes: 68 | run_dir_root: Path to the run dir root. Can be optionally templated with tags. Needs to always be run through get_path_from_template. 69 | run_desc: Description of the run. Will be used in the run dir and task name. 70 | run_dir_ignore: List of file patterns used to ignore files when copying files to the run dir. 71 | run_dir_extra_files: List of (abs_path, rel_path) tuples of file paths. rel_path root will be the src directory inside the run dir. 72 | submit_target: Submit target enum value. Used to select where the run is actually launched. 73 | num_gpus: Number of GPUs used/requested for the run. 74 | print_info: Whether to print debug information when submitting. 75 | local.do_not_copy_source_files: Do not copy source files from the working directory to the run dir. 76 | run_id: Automatically populated value during submit. 77 | run_name: Automatically populated value during submit. 78 | run_dir: Automatically populated value during submit. 79 | run_func_name: Automatically populated value during submit. 80 | run_func_kwargs: Automatically populated value during submit. 81 | user_name: Automatically populated value during submit. Can be set by the user which will then override the automatic value. 82 | task_name: Automatically populated value during submit. 83 | host_name: Automatically populated value during submit. 84 | platform_extras: Automatically populated values during submit. Used by various dnnlib libraries such as the DataReader class. 85 | """ 86 | 87 | def __init__(self): 88 | super().__init__() 89 | 90 | # run (set these) 91 | self.run_dir_root = "" # should always be passed through get_path_from_template 92 | self.run_desc = "" 93 | self.run_dir_ignore = ["__pycache__", "*.pyproj", "*.sln", "*.suo", ".cache", ".idea", ".vs", ".vscode", "_cudacache"] 94 | self.run_dir_extra_files = [] 95 | 96 | # submit (set these) 97 | self.submit_target = SubmitTarget.LOCAL 98 | self.num_gpus = 1 99 | self.print_info = False 100 | self.nvprof = False 101 | self.local = internal.local.TargetOptions() 102 | self.datasets = [] 103 | 104 | # (automatically populated) 105 | self.run_id = None 106 | self.run_name = None 107 | self.run_dir = None 108 | self.run_func_name = None 109 | self.run_func_kwargs = None 110 | self.user_name = None 111 | self.task_name = None 112 | self.host_name = "localhost" 113 | self.platform_extras = PlatformExtras() 114 | 115 | 116 | def get_path_from_template(path_template: str, path_type: PathType = PathType.AUTO) -> str: 117 | """Replace tags in the given path template and return either Windows or Linux formatted path.""" 118 | # automatically select path type depending on running OS 119 | if path_type == PathType.AUTO: 120 | if platform.system() == "Windows": 121 | path_type = PathType.WINDOWS 122 | elif platform.system() == "Linux": 123 | path_type = PathType.LINUX 124 | else: 125 | raise RuntimeError("Unknown platform") 126 | 127 | path_template = path_template.replace("", get_user_name()) 128 | 129 | # return correctly formatted path 130 | if path_type == PathType.WINDOWS: 131 | return str(pathlib.PureWindowsPath(path_template)) 132 | elif path_type == PathType.LINUX: 133 | return str(pathlib.PurePosixPath(path_template)) 134 | else: 135 | raise RuntimeError("Unknown platform") 136 | 137 | 138 | def get_template_from_path(path: str) -> str: 139 | """Convert a normal path back to its template representation.""" 140 | path = path.replace("\\", "/") 141 | return path 142 | 143 | 144 | def convert_path(path: str, path_type: PathType = PathType.AUTO) -> str: 145 | """Convert a normal path to template and the convert it back to a normal path with given path type.""" 146 | path_template = get_template_from_path(path) 147 | path = get_path_from_template(path_template, path_type) 148 | return path 149 | 150 | 151 | def set_user_name_override(name: str) -> None: 152 | """Set the global username override value.""" 153 | global _user_name_override 154 | _user_name_override = name 155 | 156 | 157 | def get_user_name(): 158 | """Get the current user name.""" 159 | if _user_name_override is not None: 160 | return _user_name_override 161 | elif platform.system() == "Windows": 162 | return os.getlogin() 163 | elif platform.system() == "Linux": 164 | try: 165 | import pwd 166 | return pwd.getpwuid(os.geteuid()).pw_name 167 | except: 168 | return "unknown" 169 | else: 170 | raise RuntimeError("Unknown platform") 171 | 172 | 173 | def make_run_dir_path(*paths): 174 | """Make a path/filename that resides under the current submit run_dir. 175 | 176 | Args: 177 | *paths: Path components to be passed to os.path.join 178 | 179 | Returns: 180 | A file/dirname rooted at submit_config.run_dir. If there's no 181 | submit_config or run_dir, the base directory is the current 182 | working directory. 183 | 184 | E.g., `os.path.join(dnnlib.submit_config.run_dir, "output.txt"))` 185 | """ 186 | import dnnlib 187 | if (dnnlib.submit_config is None) or (dnnlib.submit_config.run_dir is None): 188 | return os.path.join(os.getcwd(), *paths) 189 | return os.path.join(dnnlib.submit_config.run_dir, *paths) 190 | 191 | 192 | def _create_run_dir_local(submit_config: SubmitConfig) -> str: 193 | """Create a new run dir with increasing ID number at the start.""" 194 | run_dir_root = get_path_from_template(submit_config.run_dir_root, PathType.AUTO) 195 | 196 | if not os.path.exists(run_dir_root): 197 | os.makedirs(run_dir_root) 198 | 199 | submit_config.run_id = _get_next_run_id_local(run_dir_root) 200 | submit_config.run_name = "{0:03d}-{1}".format(submit_config.run_id, submit_config.run_desc) 201 | run_dir = os.path.join(run_dir_root, submit_config.run_name) 202 | 203 | if os.path.exists(run_dir): 204 | raise RuntimeError("The run dir already exists! ({0})".format(run_dir)) 205 | 206 | os.makedirs(run_dir) 207 | 208 | return run_dir 209 | 210 | 211 | def _get_next_run_id_local(run_dir_root: str) -> int: 212 | """Reads all directory names in a given directory (non-recursive) and returns the next (increasing) run id. Assumes IDs are numbers at the start of the directory names.""" 213 | dir_names = [d for d in os.listdir(run_dir_root) if os.path.isdir(os.path.join(run_dir_root, d))] 214 | r = re.compile("^\\d+") # match one or more digits at the start of the string 215 | run_id = 0 216 | 217 | for dir_name in dir_names: 218 | m = r.match(dir_name) 219 | 220 | if m is not None: 221 | i = int(m.group()) 222 | run_id = max(run_id, i + 1) 223 | 224 | return run_id 225 | 226 | 227 | def _populate_run_dir(submit_config: SubmitConfig, run_dir: str) -> None: 228 | """Copy all necessary files into the run dir. Assumes that the dir exists, is local, and is writable.""" 229 | pickle.dump(submit_config, open(os.path.join(run_dir, "submit_config.pkl"), "wb")) 230 | with open(os.path.join(run_dir, "submit_config.txt"), "w") as f: 231 | pprint.pprint(submit_config, stream=f, indent=4, width=200, compact=False) 232 | 233 | if (submit_config.submit_target == SubmitTarget.LOCAL) and submit_config.local.do_not_copy_source_files: 234 | return 235 | 236 | files = [] 237 | 238 | run_func_module_dir_path = util.get_module_dir_by_obj_name(submit_config.run_func_name) 239 | assert '.' in submit_config.run_func_name 240 | for _idx in range(submit_config.run_func_name.count('.') - 1): 241 | run_func_module_dir_path = os.path.dirname(run_func_module_dir_path) 242 | files += util.list_dir_recursively_with_ignore(run_func_module_dir_path, ignores=submit_config.run_dir_ignore, add_base_to_relative=False) 243 | 244 | dnnlib_module_dir_path = util.get_module_dir_by_obj_name("dnnlib") 245 | files += util.list_dir_recursively_with_ignore(dnnlib_module_dir_path, ignores=submit_config.run_dir_ignore, add_base_to_relative=True) 246 | 247 | files += submit_config.run_dir_extra_files 248 | 249 | files = [(f[0], os.path.join(run_dir, "src", f[1])) for f in files] 250 | files += [(os.path.join(dnnlib_module_dir_path, "submission", "internal", "run.py"), os.path.join(run_dir, "run.py"))] 251 | 252 | util.copy_files_and_create_dirs(files) 253 | 254 | 255 | 256 | def run_wrapper(submit_config: SubmitConfig) -> None: 257 | """Wrap the actual run function call for handling logging, exceptions, typing, etc.""" 258 | is_local = submit_config.submit_target == SubmitTarget.LOCAL 259 | 260 | # when running locally, redirect stderr to stdout, log stdout to a file, and force flushing 261 | if is_local: 262 | logger = util.Logger(file_name=os.path.join(submit_config.run_dir, "log.txt"), file_mode="w", should_flush=True) 263 | else: # when running in a cluster, redirect stderr to stdout, and just force flushing (log writing is handled by run.sh) 264 | logger = util.Logger(file_name=None, should_flush=True) 265 | 266 | import dnnlib 267 | dnnlib.submit_config = submit_config 268 | 269 | exit_with_errcode = False 270 | try: 271 | print("dnnlib: Running {0}() on {1}...".format(submit_config.run_func_name, submit_config.host_name)) 272 | start_time = time.time() 273 | 274 | run_func_obj = util.get_obj_by_name(submit_config.run_func_name) 275 | assert callable(run_func_obj) 276 | sig = inspect.signature(run_func_obj) 277 | if 'submit_config' in sig.parameters: 278 | run_func_obj(submit_config=submit_config, **submit_config.run_func_kwargs) 279 | else: 280 | run_func_obj(**submit_config.run_func_kwargs) 281 | 282 | print("dnnlib: Finished {0}() in {1}.".format(submit_config.run_func_name, util.format_time(time.time() - start_time))) 283 | except: 284 | if is_local: 285 | raise 286 | else: 287 | traceback.print_exc() 288 | 289 | log_src = os.path.join(submit_config.run_dir, "log.txt") 290 | log_dst = os.path.join(get_path_from_template(submit_config.run_dir_root), "{0}-error.txt".format(submit_config.run_name)) 291 | shutil.copyfile(log_src, log_dst) 292 | 293 | # Defer sys.exit(1) to happen after we close the logs and create a _finished.txt 294 | exit_with_errcode = True 295 | finally: 296 | open(os.path.join(submit_config.run_dir, "_finished.txt"), "w").close() 297 | 298 | dnnlib.RunContext.get().close() 299 | dnnlib.submit_config = None 300 | logger.close() 301 | 302 | # If we hit an error, get out of the script now and signal the error 303 | # to whatever process that started this script. 304 | if exit_with_errcode: 305 | sys.exit(1) 306 | 307 | return submit_config 308 | 309 | 310 | def submit_run(submit_config: SubmitConfig, run_func_name: str, **run_func_kwargs) -> None: 311 | """Create a run dir, gather files related to the run, copy files to the run dir, and launch the run in appropriate place.""" 312 | submit_config = copy.deepcopy(submit_config) 313 | 314 | submit_target = submit_config.submit_target 315 | farm = None 316 | if submit_target == SubmitTarget.LOCAL: 317 | farm = internal.local.Target() 318 | assert farm is not None # unknown target 319 | 320 | # Disallow submitting jobs with zero num_gpus. 321 | if (submit_config.num_gpus is None) or (submit_config.num_gpus == 0): 322 | raise RuntimeError("submit_config.num_gpus must be set to a non-zero value") 323 | 324 | if submit_config.user_name is None: 325 | submit_config.user_name = get_user_name() 326 | 327 | submit_config.run_func_name = run_func_name 328 | submit_config.run_func_kwargs = run_func_kwargs 329 | 330 | #-------------------------------------------------------------------- 331 | # Prepare submission by populating the run dir 332 | #-------------------------------------------------------------------- 333 | host_run_dir = _create_run_dir_local(submit_config) 334 | 335 | submit_config.task_name = "{0}-{1:03d}-{2}".format(submit_config.user_name, submit_config.run_id, submit_config.run_desc) 336 | docker_valid_name_regex = "^[a-zA-Z0-9][a-zA-Z0-9_.-]+$" 337 | if not re.match(docker_valid_name_regex, submit_config.task_name): 338 | raise RuntimeError("Invalid task name. Probable reason: unacceptable characters in your submit_config.run_desc. Task name must be accepted by the following regex: " + docker_valid_name_regex + ", got " + submit_config.task_name) 339 | 340 | # Farm specific preparations for a submit 341 | farm.finalize_submit_config(submit_config, host_run_dir) 342 | _populate_run_dir(submit_config, host_run_dir) 343 | return farm.submit(submit_config, host_run_dir) 344 | --------------------------------------------------------------------------------