Spaces:
Runtime error
Runtime error
| import os, sys | |
| import importlib | |
| import argparse | |
| import numpy as np | |
| import torch | |
| import matplotlib.pyplot as plt | |
| from PIL import Image | |
| sys.path.append("./rome/") | |
| from rome.src.utils import args as args_utils | |
| from rome.src.utils.processing import process_black_shape, tensor2image | |
| # loading models ---- create model repo | |
| from huggingface_hub import hf_hub_url | |
| default_modnet_path = hf_hub_url('Pie31415/rome','modnet_photographic_portrait_matting.ckpt') | |
| default_model_path = hf_hub_url('Pie31415/rome','models/rome.pth') | |
| # parser configurations | |
| parser = argparse.ArgumentParser(conflict_handler='resolve') | |
| parser.add_argument('--save_dir', default='.', type=str) | |
| parser.add_argument('--save_render', default='True', type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--model_checkpoint', default=default_model_path, type=str) | |
| parser.add_argument('--modnet_path', default=default_modnet_path, type=str) | |
| parser.add_argument('--random_seed', default=0, type=int) | |
| parser.add_argument('--debug', action='store_true') | |
| parser.add_argument('--verbose', default='False', type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--model_image_size', default=256, type=int) | |
| parser.add_argument('--align_source', default='True', type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--align_target', default='False', type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--align_scale', default=1.25, type=float) | |
| parser.add_argument('--use_mesh_deformations', default='False', type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--subdivide_mesh', default='False', type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--renderer_sigma', default=1e-8, type=float) | |
| parser.add_argument('--renderer_zfar', default=100.0, type=float) | |
| parser.add_argument('--renderer_type', default='soft_mesh') | |
| parser.add_argument('--renderer_texture_type', default='texture_uv') | |
| parser.add_argument('--renderer_normalized_alphas', default='False', type=args_utils.str2bool, | |
| choices=[True, False]) | |
| parser.add_argument('--deca_path', default='') | |
| parser.add_argument('--rome_data_dir', default='') | |
| parser.add_argument('--autoenc_cat_alphas', default='False', type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--autoenc_align_inputs', default='False', type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--autoenc_use_warp', default='False', type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--autoenc_num_channels', default=64, type=int) | |
| parser.add_argument('--autoenc_max_channels', default=512, type=int) | |
| parser.add_argument('--autoenc_num_groups', default=4, type=int) | |
| parser.add_argument('--autoenc_num_bottleneck_groups', default=0, type=int) | |
| parser.add_argument('--autoenc_num_blocks', default=2, type=int) | |
| parser.add_argument('--autoenc_num_layers', default=4, type=int) | |
| parser.add_argument('--autoenc_block_type', default='bottleneck') | |
| parser.add_argument('--neural_texture_channels', default=8, type=int) | |
| parser.add_argument('--num_harmonic_encoding_funcs', default=6, type=int) | |
| parser.add_argument('--unet_num_channels', default=64, type=int) | |
| parser.add_argument('--unet_max_channels', default=512, type=int) | |
| parser.add_argument('--unet_num_groups', default=4, type=int) | |
| parser.add_argument('--unet_num_blocks', default=1, type=int) | |
| parser.add_argument('--unet_num_layers', default=2, type=int) | |
| parser.add_argument('--unet_block_type', default='conv') | |
| parser.add_argument('--unet_skip_connection_type', default='cat') | |
| parser.add_argument('--unet_use_normals_cond', default=True, action='store_true') | |
| parser.add_argument('--unet_use_vertex_cond', action='store_true') | |
| parser.add_argument('--unet_use_uvs_cond', action='store_true') | |
| parser.add_argument('--unet_pred_mask', action='store_true') | |
| parser.add_argument('--use_separate_seg_unet', default='True', type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--norm_layer_type', default='gn', type=str, choices=['bn', 'sync_bn', 'in', 'gn']) | |
| parser.add_argument('--activation_type', default='relu', type=str, choices=['relu', 'lrelu']) | |
| parser.add_argument('--conv_layer_type', default='ws_conv', type=str, choices=['conv', 'ws_conv']) | |
| parser.add_argument('--deform_norm_layer_type', default='gn', type=str, choices=['bn', 'sync_bn', 'in', 'gn']) | |
| parser.add_argument('--deform_activation_type', default='relu', type=str, choices=['relu', 'lrelu']) | |
| parser.add_argument('--deform_conv_layer_type', default='ws_conv', type=str, choices=['conv', 'ws_conv']) | |
| parser.add_argument('--unet_seg_weight', default=0.0, type=float) | |
| parser.add_argument('--unet_seg_type', default='bce_with_logits', type=str, choices=['bce_with_logits', 'dice']) | |
| parser.add_argument('--deform_face_tightness', default=0.0, type=float) | |
| parser.add_argument('--use_whole_segmentation', action='store_true') | |
| parser.add_argument('--mask_hair_for_neck', action='store_true') | |
| parser.add_argument('--use_hair_from_avatar', action='store_true') | |
| # Basis deformations | |
| parser.add_argument('--use_scalp_deforms', default='True', type=args_utils.str2bool, | |
| choices=[True, False], help='') | |
| parser.add_argument('--use_neck_deforms', default='True', type=args_utils.str2bool, | |
| choices=[True, False], help='') | |
| parser.add_argument('--use_basis_deformer', default='False', type=args_utils.str2bool, | |
| choices=[True, False], help='') | |
| parser.add_argument('--use_unet_deformer', default='True', type=args_utils.str2bool, | |
| choices=[True, False], help='') | |
| parser.add_argument('--pretrained_encoder_basis_path', default='') | |
| parser.add_argument('--pretrained_vertex_basis_path', default='') | |
| parser.add_argument('--num_basis', default=50, type=int) | |
| parser.add_argument('--basis_init', default='pca', type=str, choices=['random', 'pca']) | |
| parser.add_argument('--num_vertex', default=5023, type=int) | |
| parser.add_argument('--train_basis', default=True, type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--path_to_deca', default='DECA') | |
| parser.add_argument('--path_to_linear_hair_model', | |
| default='data/linear_hair.pth') | |
| parser.add_argument('--path_to_mobile_model', | |
| default='data/disp_model.pth') | |
| parser.add_argument('--n_scalp', default=60, type=int) | |
| parser.add_argument('--use_distill', default=False, type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--use_mobile_version', default=False, type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--deformer_path', default='data/rome.pth') | |
| parser.add_argument('--output_unet_deformer_feats', default=32, type=int, | |
| help='output features in the UNet') | |
| parser.add_argument('--use_deca_details', default=False, type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--use_flametex', default=False, type=args_utils.str2bool, choices=[True, False]) | |
| parser.add_argument('--upsample_type', default='nearest', type=str, | |
| choices=['nearest', 'bilinear', 'bicubic']) | |
| parser.add_argument('--num_frequencies', default=6, type=int, help='frequency for harmonic encoding') | |
| parser.add_argument('--deform_face_scale_coef', default=0.0, type=float) | |
| parser.add_argument('--device', default='cpu', type=str) | |
| # args, _ = parser.parse_known_args() | |
| # parser = importlib.import_module(f'src.rome').ROME.add_argparse_args(parser) | |
| args = parser.parse_args() | |
| args.deca_path = 'DECA' | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| from infer import Infer | |
| infer = Infer(args) | |
| infer = infer.to(device) | |
| def predict(source_img, driver_img): | |
| out = infer.evaluate(source_img, driver_img, crop_center=False) | |
| res = tensor2image(torch.cat([out['source_information']['data_dict']['source_img'][0].cpu(), | |
| out['source_information']['data_dict']['target_img'][0].cpu(), | |
| out['render_masked'].cpu(), out['pred_target_shape_img'][0].cpu()], dim=2)) | |
| return res[..., ::-1] | |
| import gradio as gr | |
| gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.Image(type="pil"), | |
| gr.Image(type="pil") | |
| ], | |
| outputs=gr.Image(), | |
| examples=[]).launch() |