Add initial RunwayML support from @genekogan's branch

Author: pbaylies

README.md

@@ -6,6 +6,9 @@
 * port of my encoder from @robertluxemburg
 * fp16 branch from @veqtor
 * tpu and swarm branches from @shawwn
+* runwayml support from @genekogan
+
+<a href="https://open-app.runwayml.com/?model=pbaylies/stylegan2" target="_blank"><img src="https://open-app.runwayml.com/gh-badge.svg" /></a>
 
 **Various Improvements from skyflynil to make StyleGAN2 more suitible to be trained on Google Colab**
 * Supports Non-Square images, for example, 768x512, which basically as 6x4 (x2^7), or 640x384 as 5x3 (x2^7), etc.
@@ -95,6 +98,7 @@ This is an experimental port of [pbaylies/stylegan-encoder](https://github.com/p
 To test this, try out the notebook.
 
 ### Original Readme
+
 ![Teaser image](./docs/stylegan2-teaser-1024x256.png)
 
 **Analyzing and Improving the Image Quality of StyleGAN**<br>

dnnlib/tflib/custom_ops.py

@@ -61,7 +61,7 @@ def _run_cmd(cmd):
         raise RuntimeError('NVCC returned an error. See below for full command line and output log:\n\n%s\n\n%s' % (cmd, output))
 
 def _prepare_nvcc_cli(opts):
-    cmd = 'nvcc ' + opts.strip()
+    cmd = 'nvcc --std=c++11 -DNDEBUG ' + opts.strip()
     cmd += ' --disable-warnings'
     cmd += ' --include-path "%s"' % tf.sysconfig.get_include()
     cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'protobuf_archive', 'src')

requirements.txt

@@ -0,0 +1,32 @@
+absl-py==0.7.0
+astor==0.7.1
+certifi==2018.11.29
+chardet==3.0.4
+Click==7.0
+Flask==1.0.2
+Flask-Cors==3.0.7
+gast==0.2.2
+gevent==1.4.0
+greenlet==0.4.15
+grpcio==1.19.0
+h5py==2.9.0
+idna==2.8
+itsdangerous==1.1.0
+Jinja2==2.10
+Keras-Applications==1.0.7
+Keras-Preprocessing==1.0.9
+Markdown==3.0.1
+MarkupSafe==1.1.1
+mock==2.0.0
+numpy==1.16.2
+pbr==5.1.2
+Pillow==5.4.1
+protobuf==3.6.1
+requests==2.21.0
+six==1.12.0
+tensorflow-gpu==1.15.0
+termcolor==1.1.0
+urllib3==1.24.1
+Werkzeug==0.14.1
+wget==3.2
+runway-python

run_training.py

@@ -33,7 +33,7 @@
 
 #----------------------------------------------------------------------------
 
-def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, mirror_augment_v, metrics, min_h, min_w, res_log2, lr, cond):
+def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, mirror_augment, mirror_augment_v, metrics, min_h, min_w, res_log2, lr, cond, resume_pkl, resume_kimg):
     train     = EasyDict(run_func_name='training.training_loop.training_loop') # Options for training loop.
     G         = EasyDict(func_name='training.networks_stylegan2.G_main')       # Options for generator network.
     D         = EasyDict(func_name='training.networks_stylegan2.D_stylegan2')  # Options for discriminator network.
@@ -122,6 +122,7 @@ def run(dataset, data_dir, result_dir, config_id, num_gpus, total_kimg, gamma, m
     kwargs = EasyDict(train)
     kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss)
     kwargs.update(dataset_args=dataset_args, sched_args=sched, grid_args=grid, metric_arg_list=metrics, tf_config=tf_config)
+    kwargs.update(resume_pkl=resume_pkl, resume_kimg=resume_kimg)
     kwargs.submit_config = copy.deepcopy(sc)
     kwargs.submit_config.run_dir_root = result_dir
     kwargs.submit_config.run_desc = desc
@@ -182,6 +183,8 @@ def main():
     parser.add_argument('--res-log2', help='multiplier for image size, the training image size (height, width) should be (min_h * 2**res_log2, min_w * 2**res_log2)', default=4, type=int)
     parser.add_argument('--lr', help='base learning rate', default=0.003, type=float)
     parser.add_argument('--cond', help='conditional model', default=False, metavar='BOOL', type=_str_to_bool)
+    parser.add_argument('--resume-pkl', help='pkl to resume training from: None)', default=None, type=str)
+    parser.add_argument('--resume-kimg', help='kimg to resume training from" (default: 0)', default=0, type=int)
 
     args = parser.parse_args()
 
@@ -207,3 +210,4 @@ def main():
 
 #----------------------------------------------------------------------------
 
+

runway.yml

@@ -0,0 +1,11 @@
+entrypoint: python runway_model.py
+python: 3.6
+cuda: 10.0
+spec:
+  gpu: True
+  cpu: False
+files:
+  ignore:
+    - checkpoints/*
+build_steps:
+  - pip install -r requirements.txt

runway_model.py

@@ -0,0 +1,37 @@
+import pickle
+import numpy as np
+import tensorflow as tf
+import dnnlib.tflib as tflib
+import runway
+
+fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
+
+@runway.setup(options={'checkpoint': runway.file(extension='.pkl')})
+def setup(opts):
+    global Gs
+    tflib.init_tf()
+    with open(opts['checkpoint'], 'rb') as file:
+        _G, _D, Gs = pickle.load(file, encoding='latin1')
+    noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
+    rnd = np.random.RandomState()
+    tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars})
+    return Gs
+
+
+generate_inputs = {
+    'z': runway.vector(512, sampling_std=0.5),
+    'truncation': runway.number(min=0, max=1, default=0.8, step=0.01)
+}
+
+@runway.command('generate', inputs=generate_inputs, outputs={'image': runway.image})
+def convert(model, inputs):
+    z = inputs['z']
+    truncation = inputs['truncation']
+    latents = z.reshape((1, 512))
+    images = model.run(latents, None, truncation_psi=truncation, randomize_noise=False, output_transform=fmt)
+    output = np.clip(images[0], 0, 255).astype(np.uint8)
+    return {'image': output}
+
+
+if __name__ == '__main__':
+    runway.run()