65 lines
2.7 KiB
Python
65 lines
2.7 KiB
Python
import tensorflow as tf
|
|
|
|
|
|
def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=True):
|
|
if custom_cuda:
|
|
from stylegan2.generator import Generator
|
|
else:
|
|
from stylegan2_ref.generator import Generator
|
|
|
|
if g_params is None:
|
|
g_params = {
|
|
'z_dim': 512,
|
|
'w_dim': 512,
|
|
'labels_dim': 0,
|
|
'n_mapping': 8,
|
|
'resolutions': [4, 8, 16, 32, 64, 128, 256, 512, 1024],
|
|
'featuremaps': [512, 512, 512, 512, 512, 256, 128, 64, 32],
|
|
}
|
|
|
|
test_latent = tf.ones((1, g_params['z_dim']), dtype=tf.float32)
|
|
# test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32), tf.ones((1, 512, 256), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32))
|
|
# test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32))
|
|
test_conditions = { 'pose_shape': tf.ones((1, 83), dtype=tf.float32), 'uv_static': tf.ones((1, 3, 512, 256), dtype=tf.float32), 'uv_mask': tf.ones((1, 512, 256), dtype=tf.float32), 'uv_shaped_body': tf.ones((1, 3, 512, 256), dtype=tf.float32) }
|
|
|
|
# build generator model
|
|
generator = Generator(g_params)
|
|
_ = generator([test_latent, test_conditions])
|
|
|
|
if ckpt_dir is not None:
|
|
if is_g_clone:
|
|
ckpt = tf.train.Checkpoint(g_clone=generator)
|
|
else:
|
|
ckpt = tf.train.Checkpoint(generator=generator)
|
|
manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=1)
|
|
ckpt.restore(manager.latest_checkpoint).expect_partial()
|
|
if manager.latest_checkpoint:
|
|
print(f'Generator restored from {manager.latest_checkpoint}')
|
|
return generator
|
|
|
|
|
|
def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
|
|
# if custom_cuda:
|
|
# from stylegan2.discriminator import Discriminator
|
|
# else:
|
|
# from stylegan2_ref.discriminator import Discriminator
|
|
|
|
from custom_utils.graph_utils import GNNDiscriminator as Discriminator
|
|
|
|
|
|
res = d_params['resolutions'][-1]
|
|
batch_size = d_params['batch_size']
|
|
test_images = tf.ones((batch_size, 3, res, int(res/2)), dtype=tf.float32)
|
|
|
|
# build discriminator model
|
|
discriminator = Discriminator(batch_size, 32)
|
|
_ = discriminator(test_images)
|
|
|
|
if ckpt_dir is not None:
|
|
ckpt = tf.train.Checkpoint(discriminator=discriminator)
|
|
manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=1)
|
|
ckpt.restore(manager.latest_checkpoint).expect_partial()
|
|
if manager.latest_checkpoint:
|
|
print('Discriminator restored from {}'.format(manager.latest_checkpoint))
|
|
return discriminator
|