def __init__(self,
image_size: int = 256,
image_channels: int = 4,
num_expression_params: int = 67,
start_channels: int = 16,
bottleneck_image_size=4,
num_bottleneck_blocks=3,
max_channels: int = 512,
block_args: Optional[BlockArgs] = None):
self.max_channels = max_channels
self.num_bottleneck_blocks = num_bottleneck_blocks
assert bottleneck_image_size > 1
self.bottleneck_image_size = bottleneck_image_size
self.start_channels = start_channels
self.image_channels = image_channels
self.num_expression_params = num_expression_params
self.image_size = image_size
if block_args is None:
self.block_args = BlockArgs(
normalization_layer_factory=InstanceNorm2dFactory(),
nonlinearity_factory=LeakyReLUFactory(negative_slope=0.2,
inplace=True))
else:
self.block_args = block_args
def create_conv3_from_block_args(in_channels: int,
out_channels: int,
bias: bool = False,
block_args: Optional[BlockArgs] = None):
if block_args is None:
block_args = BlockArgs()
return create_conv3(in_channels, out_channels, bias,
block_args.initialization_method,
block_args.use_spectral_norm)
def create_conv3_block_from_block_args(in_channels: int,
out_channels: int,
block_args: Optional[BlockArgs] = None):
if block_args is None:
block_args = BlockArgs()
return create_conv3_block(in_channels, out_channels,
block_args.initialization_method,
block_args.nonlinearity_factory,
block_args.normalization_layer_factory,
block_args.use_spectral_norm)
def create(num_channels: int,
is1x1: bool = False,
use_scale_parameters: bool = False,
block_args: Optional[BlockArgs] = None):
if block_args is None:
block_args = BlockArgs()
return ResnetBlock(num_channels, is1x1,
block_args.initialization_method,
block_args.nonlinearity_factory,
block_args.normalization_layer_factory,
block_args.use_spectral_norm, use_scale_parameters)
def create_downsample_block_from_block_args(
in_channels: int,
out_channels: int,
is_output_1x1: bool = False,
block_args: Optional[BlockArgs] = None):
if block_args is None:
block_args = BlockArgs()
return create_downsample_block(in_channels, out_channels, is_output_1x1,
block_args.initialization_method,
block_args.nonlinearity_factory,
block_args.normalization_layer_factory,
block_args.use_spectral_norm)
def __init__(self,
image_size: int,
input_image_channels: int,
output_image_channels: int,
start_channels: int,
num_pose_params: int,
block_args: Optional[BlockArgs] = None):
self.num_pose_params = num_pose_params
self.start_channels = start_channels
self.output_image_channels = output_image_channels
self.input_image_channels = input_image_channels
self.image_size = image_size
if block_args is None:
self.block_args = BlockArgs(
normalization_layer_factory=InstanceNorm2dFactory(),
nonlinearity_factory=ReLUFactory(inplace=True))
else:
self.block_args = block_args
def load_eyebrow_decomposer(file_name: str):
factory = EyebrowDecomposer00Factory(
EyebrowDecomposer00Args(
image_size=128,
image_channels=4,
start_channels=64,
bottleneck_image_size=16,
num_bottleneck_blocks=6,
max_channels=512,
block_args=BlockArgs(
initialization_method='he',
use_spectral_norm=False,
normalization_layer_factory=InstanceNorm2dFactory(),
nonlinearity_factory=ReLUFactory(inplace=True))))
#print("Loading the eyebrow decomposer ... ", end="")
module = factory.create()
module.load_state_dict(torch_load(file_name))
#print("DONE!!!")
return module
def load_face_morpher(file_name: str):
factory = FaceMorpher08Factory(
FaceMorpher08Args(
image_size=192,
image_channels=4,
num_expression_params=27,
start_channels=64,
bottleneck_image_size=24,
num_bottleneck_blocks=6,
max_channels=512,
block_args=BlockArgs(
initialization_method='he',
use_spectral_norm=False,
normalization_layer_factory=InstanceNorm2dFactory(),
nonlinearity_factory=ReLUFactory(inplace=False))))
#print("Loading the face morpher ... ", end="")
module = factory.create()
module.load_state_dict(torch_load(file_name))
#print("DONE")
return module
def __init__(self,
image_size: int,
input_image_channels: int,
output_image_channels: int,
num_pose_params: int,
start_channels: int,
bottleneck_image_size,
num_bottleneck_blocks,
max_channels: int,
block_args: Optional[BlockArgs] = None):
super().__init__(image_size, input_image_channels,
output_image_channels, start_channels,
num_pose_params, block_args)
self.max_channels = max_channels
self.num_bottleneck_blocks = num_bottleneck_blocks
self.bottleneck_image_size = bottleneck_image_size
assert bottleneck_image_size > 1
if block_args is None:
self.block_args = BlockArgs(
normalization_layer_factory=InstanceNorm2dFactory(),
nonlinearity_factory=ReLUFactory(inplace=True))
else:
self.block_args = block_args
class EyebrowDecomposer00Factory(BatchInputModuleFactory):
def __init__(self, args: EyebrowDecomposer00Args):
super().__init__()
self.args = args
def create(self) -> BatchInputModule:
return EyebrowDecomposer00(self.args)
if __name__ == "__main__":
cuda = torch.device('cuda')
args = EyebrowDecomposer00Args(
image_size=128,
image_channels=4,
start_channels=64,
bottleneck_image_size=16,
num_bottleneck_blocks=3,
block_args=BlockArgs(
initialization_method='xavier',
use_spectral_norm=False,
normalization_layer_factory=InstanceNorm2dFactory(),
nonlinearity_factory=ReLUFactory(inplace=True)))
face_morpher = EyebrowDecomposer00(args).to(cuda)
image = torch.randn(8, 4, 128, 128, device=cuda)
outputs = face_morpher.forward(image)
for i in range(len(outputs)):
print(i, outputs[i].shape)