forked from NVlabs/stylegan2
-
Notifications
You must be signed in to change notification settings - Fork 55
/
runway_model.py
37 lines (30 loc) · 1.19 KB
/
runway_model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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()