diff --git a/download_models.cmd b/download_models.cmd new file mode 100644 index 0000000..7acadf3 --- /dev/null +++ b/download_models.cmd @@ -0,0 +1,24 @@ +# load arcface +wget -P ./arcface_model https://github.com/sberbank-ai/sber-swap/releases/download/arcface/backbone.pth +wget -P ./arcface_model https://github.com/sberbank-ai/sber-swap/releases/download/arcface/iresnet.py + +# load landmarks detector +wget -P ./insightface_func/models/antelope https://github.com/sberbank-ai/sber-swap/releases/download/antelope/glintr100.onnx +wget -P ./insightface_func/models/antelope https://github.com/sberbank-ai/sber-swap/releases/download/antelope/scrfd_10g_bnkps.onnx + +# load G and D models with 1, 2, 3 blocks +# model with 2 blocks is main +wget -P ./weights https://github.com/sberbank-ai/sber-swap/releases/download/sber-swap-v2.0/G_unet_2blocks.pth +wget -P ./weights https://github.com/sberbank-ai/sber-swap/releases/download/sber-swap-v2.0/D_unet_2blocks.pth + +wget -P ./weights https://github.com/sberbank-ai/sber-swap/releases/download/sber-swap-v2.0/G_unet_1block.pth +wget -P ./weights https://github.com/sberbank-ai/sber-swap/releases/download/sber-swap-v2.0/D_unet_1block.pth + +wget -P ./weights https://github.com/sberbank-ai/sber-swap/releases/download/sber-swap-v2.0/G_unet_3blocks.pth +wget -P ./weights https://github.com/sberbank-ai/sber-swap/releases/download/sber-swap-v2.0/D_unet_3blocks.pth + +# load model for eyes loss +wget -P ./AdaptiveWingLoss/AWL_detector https://github.com/sberbank-ai/sber-swap/releases/download/awl_detector/WFLW_4HG.pth + +# load super res model +wget -P ./weights https://github.com/sberbank-ai/sber-swap/releases/download/super-res/10_net_G.pth diff --git a/inference.py b/inference.py index c3a824c..d37e2c4 100644 --- a/inference.py +++ b/inference.py @@ -114,7 +114,8 @@ def main(args): fps, handler) - add_audio_from_another_video(args.target_video, args.out_video_name, "audio") + if not args.ignore_audio: + add_audio_from_another_video(args.target_video, args.out_video_name, "audio") print(f"Video saved with path {args.out_video_name}") else: result = get_final_image(final_frames_list, crop_frames_list, full_frames[0], tfm_array_list, handler) @@ -149,5 +150,7 @@ def main(args): parser.add_argument('--target_image', default='examples/images/beckham.jpg', type=str, help="It's necessary for image to image swap") parser.add_argument('--out_image_name', default='examples/results/result.png', type=str,help="It's necessary for image to image swap") + parser.add_argument('--ignore_audio', default=False, type=bool, help="Do not include the audio on the video result") + args = parser.parse_args() main(args) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index d256a97..1c8b5da 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,7 +9,7 @@ onnxruntime-gpu==1.4.0 mxnet-cu101mkl scikit-image insightface==0.2.1 -requests==2.25.1 +requests kornia==0.5.4 dill wandb \ No newline at end of file