diff --git a/.gitignore b/.gitignore index bd6a181ac8d66af3c486ec82538d5f297645b990..400ce70510811cd910fdd17d2f2ce1fb97123562 100644 --- a/.gitignore +++ b/.gitignore @@ -152,4 +152,8 @@ dmypy.json # Cython debug symbols cython_debug/ -results/ \ No newline at end of file +results/ +checkpoints/ +gradio_cached_examples/ +gfpgan/ +start.sh \ No newline at end of file diff --git a/app.py b/app.py index 1fe16489b6bc2dd39d2de5ac1f89c04f5da49eff..8f3fded14dc1aac1932f302c031c725634212ee0 100644 --- a/app.py +++ b/app.py @@ -3,15 +3,20 @@ import tempfile import gradio as gr from src.gradio_demo import SadTalker from src.utils.text2speech import TTSTalker +from huggingface_hub import snapshot_download def get_source_image(image): return image - +def download_model(): + REPO_ID = 'vinthony/SadTalker' + snapshot_download(repo_id=REPO_ID, local_dir='./checkpoints', local_dir_use_symlinks=True) def sadtalker_demo(): - sad_talker = SadTalker() + download_model() + + sad_talker = SadTalker(lazy_load=True) tts_talker = TTSTalker() with gr.Blocks(analytics_enabled=False) as sadtalker_interface: @@ -42,7 +47,8 @@ def sadtalker_demo(): with gr.Tabs(elem_id="sadtalker_checkbox"): with gr.TabItem('Settings'): with gr.Column(variant='panel'): - is_still_mode = gr.Checkbox(label="w/ Still Mode (fewer hand motion, works on full body)") + preprocess_type = gr.Radio(['crop','resize','full'], value='crop', label='preprocess', info="How to handle input image?") + is_still_mode = gr.Checkbox(label="w/ Still Mode (fewer hand motion, works with preprocess `full`)") enhancer = gr.Checkbox(label="w/ GFPGAN as Face enhancer") submit = gr.Button('Generate', elem_id="sadtalker_generate", variant='primary') @@ -54,42 +60,83 @@ def sadtalker_demo(): [ 'examples/source_image/full_body_1.png', 'examples/driven_audio/bus_chinese.wav', + 'crop', True, False ], [ 'examples/source_image/full_body_2.png', - 'examples/driven_audio/itosinger1.wav', + 'examples/driven_audio/japanese.wav', + 'crop', + False, + False + ], + [ + 'examples/source_image/full3.png', + 'examples/driven_audio/deyu.wav', + 'crop', + False, + True + ], + [ + 'examples/source_image/full4.jpeg', + 'examples/driven_audio/eluosi.wav', + 'full', + False, + True + ], + [ + 'examples/source_image/full4.jpeg', + 'examples/driven_audio/imagine.wav', + 'full', + True, + True + ], + [ + 'examples/source_image/full_body_1.png', + 'examples/driven_audio/bus_chinese.wav', + 'full', True, False ], [ 'examples/source_image/art_13.png', 'examples/driven_audio/fayu.wav', + 'resize', True, False ], [ 'examples/source_image/art_5.png', 'examples/driven_audio/chinese_news.wav', - True, + 'resize', + False, False ], + [ + 'examples/source_image/art_5.png', + 'examples/driven_audio/RD_Radio31_000.wav', + 'resize', + True, + True + ], ] gr.Examples(examples=examples, inputs=[ source_image, driven_audio, + preprocess_type, is_still_mode, enhancer], outputs=[gen_video], fn=sad_talker.test, - cache_examples=os.getenv('SYSTEM') == 'spaces') + cache_examples=True) # os.getenv('SYSTEM') == 'spaces') submit.click( fn=sad_talker.test, inputs=[source_image, driven_audio, + preprocess_type, is_still_mode, enhancer], outputs=[gen_video] diff --git a/checkpoints/BFM_Fitting.zip b/checkpoints/BFM_Fitting.zip deleted file mode 100644 index 895479e053ea9f18c12cf68217cd58543b1d2d84..0000000000000000000000000000000000000000 --- a/checkpoints/BFM_Fitting.zip +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:785f77f3de288568e76666cd419dcf40646d3f74eae6d4fa3b766c933087a9d8 -size 404051745 diff --git a/checkpoints/BFM_Fitting/01_MorphableModel.mat b/checkpoints/BFM_Fitting/01_MorphableModel.mat deleted file mode 100644 index f251485b55d35adac0ad4f1622a47d7a39a1502c..0000000000000000000000000000000000000000 --- a/checkpoints/BFM_Fitting/01_MorphableModel.mat +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:37b1f0742db356a3b1568a8365a06f5b0fe0ab687ac1c3068c803666cbd4d8e2 -size 240875364 diff --git a/checkpoints/BFM_Fitting/01_MorphableModel.mat b/checkpoints/BFM_Fitting/01_MorphableModel.mat new file mode 120000 index 0000000000000000000000000000000000000000..9754188b2735b609804e725c86baa0749934874b --- /dev/null +++ b/checkpoints/BFM_Fitting/01_MorphableModel.mat @@ -0,0 +1 @@ +../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/37b1f0742db356a3b1568a8365a06f5b0fe0ab687ac1c3068c803666cbd4d8e2 \ No newline at end of file diff --git a/checkpoints/BFM_Fitting/BFM09_model_info.mat b/checkpoints/BFM_Fitting/BFM09_model_info.mat deleted file mode 100644 index 605b1aa60286236b4041d15fccdd978b9d89761d..0000000000000000000000000000000000000000 --- a/checkpoints/BFM_Fitting/BFM09_model_info.mat +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:db8d00544f0b0182f1b8430a3bb87662b3ff674eb33c84e6f52dbe2971adb81b -size 127170280 diff --git a/checkpoints/BFM_Fitting/BFM09_model_info.mat b/checkpoints/BFM_Fitting/BFM09_model_info.mat new file mode 120000 index 0000000000000000000000000000000000000000..271cc275c5f24fdca95cb58084bae7923c5bda16 --- /dev/null +++ b/checkpoints/BFM_Fitting/BFM09_model_info.mat @@ -0,0 +1 @@ +../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/db8d00544f0b0182f1b8430a3bb87662b3ff674eb33c84e6f52dbe2971adb81b \ No newline at end of file diff --git a/checkpoints/BFM_Fitting/BFM_exp_idx.mat b/checkpoints/BFM_Fitting/BFM_exp_idx.mat deleted file mode 100644 index 1146e4e9c3bef303a497383aa7974c014fe945c7..0000000000000000000000000000000000000000 Binary files a/checkpoints/BFM_Fitting/BFM_exp_idx.mat and /dev/null differ diff --git a/checkpoints/BFM_Fitting/BFM_exp_idx.mat b/checkpoints/BFM_Fitting/BFM_exp_idx.mat new file mode 120000 index 0000000000000000000000000000000000000000..00c157e4f8fd641ae92ee08940a1e9a694409ad0 --- /dev/null +++ b/checkpoints/BFM_Fitting/BFM_exp_idx.mat @@ -0,0 +1 @@ +../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/1146e4e9c3bef303a497383aa7974c014fe945c7 \ No newline at end of file diff --git a/checkpoints/BFM_Fitting/BFM_front_idx.mat b/checkpoints/BFM_Fitting/BFM_front_idx.mat deleted file mode 100644 index b9d7b0953dd1dc5b1e28144610485409ac321f9b..0000000000000000000000000000000000000000 Binary files a/checkpoints/BFM_Fitting/BFM_front_idx.mat and /dev/null differ diff --git a/checkpoints/BFM_Fitting/BFM_front_idx.mat b/checkpoints/BFM_Fitting/BFM_front_idx.mat new file mode 120000 index 0000000000000000000000000000000000000000..f9bf14dd9c2f772e45e0db99b2dc050cd6897aab --- /dev/null +++ b/checkpoints/BFM_Fitting/BFM_front_idx.mat @@ -0,0 +1 @@ +../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/b9d7b0953dd1dc5b1e28144610485409ac321f9b \ No newline at end of file diff --git a/checkpoints/BFM_Fitting/Exp_Pca.bin b/checkpoints/BFM_Fitting/Exp_Pca.bin deleted file mode 100644 index 3c1785e6abc52b13e54a573f9f3ebc099915b1e0..0000000000000000000000000000000000000000 --- a/checkpoints/BFM_Fitting/Exp_Pca.bin +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e7f31380e6cbdaf2aeec698db220bac4f221946e4d551d88c092d47ec49b1726 -size 51086404 diff --git a/checkpoints/BFM_Fitting/Exp_Pca.bin b/checkpoints/BFM_Fitting/Exp_Pca.bin new file mode 120000 index 0000000000000000000000000000000000000000..e90269a678fa1a5e85b393cb1698a337679c7664 --- /dev/null +++ b/checkpoints/BFM_Fitting/Exp_Pca.bin @@ -0,0 +1 @@ +../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/e7f31380e6cbdaf2aeec698db220bac4f221946e4d551d88c092d47ec49b1726 \ No newline at end of file diff --git a/checkpoints/BFM_Fitting/facemodel_info.mat b/checkpoints/BFM_Fitting/facemodel_info.mat deleted file mode 100644 index 3e516ec7297fa3248098f49ecea10579f4831c0a..0000000000000000000000000000000000000000 Binary files a/checkpoints/BFM_Fitting/facemodel_info.mat and /dev/null differ diff --git a/checkpoints/BFM_Fitting/facemodel_info.mat b/checkpoints/BFM_Fitting/facemodel_info.mat new file mode 120000 index 0000000000000000000000000000000000000000..a805d0177a3b2edbeb7649756882ec2948ecfe56 --- /dev/null +++ b/checkpoints/BFM_Fitting/facemodel_info.mat @@ -0,0 +1 @@ +../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/3e516ec7297fa3248098f49ecea10579f4831c0a \ No newline at end of file diff --git a/checkpoints/BFM_Fitting/select_vertex_id.mat b/checkpoints/BFM_Fitting/select_vertex_id.mat deleted file mode 100644 index 5b8b220093d93b133acc94ffed159f31a74854cd..0000000000000000000000000000000000000000 Binary files a/checkpoints/BFM_Fitting/select_vertex_id.mat and /dev/null differ diff --git a/checkpoints/BFM_Fitting/select_vertex_id.mat b/checkpoints/BFM_Fitting/select_vertex_id.mat new file mode 120000 index 0000000000000000000000000000000000000000..f4b8cc6db350c42879a84de5e8aa89ce01688017 --- /dev/null +++ b/checkpoints/BFM_Fitting/select_vertex_id.mat @@ -0,0 +1 @@ +../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/5b8b220093d93b133acc94ffed159f31a74854cd \ No newline at end of file diff --git a/checkpoints/BFM_Fitting/similarity_Lm3D_all.mat b/checkpoints/BFM_Fitting/similarity_Lm3D_all.mat deleted file mode 100644 index a0e23588302bc71fc899eef53ff06df5f4df4c1d..0000000000000000000000000000000000000000 Binary files a/checkpoints/BFM_Fitting/similarity_Lm3D_all.mat and /dev/null differ diff --git a/checkpoints/BFM_Fitting/similarity_Lm3D_all.mat b/checkpoints/BFM_Fitting/similarity_Lm3D_all.mat new file mode 120000 index 0000000000000000000000000000000000000000..525f130df9377a61500618bd99d7c0e28ea80382 --- /dev/null +++ b/checkpoints/BFM_Fitting/similarity_Lm3D_all.mat @@ -0,0 +1 @@ +../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/a0e23588302bc71fc899eef53ff06df5f4df4c1d \ No newline at end of file diff --git a/checkpoints/BFM_Fitting/std_exp.txt b/checkpoints/BFM_Fitting/std_exp.txt deleted file mode 100644 index 767b8de4ea1ca78b6f22b98ff2dee4fa345500bb..0000000000000000000000000000000000000000 --- a/checkpoints/BFM_Fitting/std_exp.txt +++ /dev/null @@ -1 +0,0 @@ -453980 257264 263068 211890 135873 184721 47055.6 72732 62787.4 106226 56708.5 51439.8 34887.1 44378.7 51813.4 31030.7 23354.9 23128.1 19400 21827.6 22767.7 22057.4 19894.3 16172.8 17142.7 10035.3 14727.5 12972.5 10763.8 8953.93 8682.62 8941.81 6342.3 5205.3 7065.65 6083.35 6678.88 4666.63 5082.89 5134.76 4908.16 3964.93 3739.95 3180.09 2470.45 1866.62 1624.71 2423.74 1668.53 1471.65 1194.52 782.102 815.044 835.782 834.937 744.496 575.146 633.76 705.685 753.409 620.306 673.326 766.189 619.866 559.93 357.264 396.472 556.849 455.048 460.592 400.735 326.702 279.428 291.535 326.584 305.664 287.816 283.642 276.19 \ No newline at end of file diff --git a/checkpoints/BFM_Fitting/std_exp.txt b/checkpoints/BFM_Fitting/std_exp.txt new file mode 120000 index 0000000000000000000000000000000000000000..140ca790f5f0abc6aae8b5fee68547be819d811c --- /dev/null +++ b/checkpoints/BFM_Fitting/std_exp.txt @@ -0,0 +1 @@ +../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/767b8de4ea1ca78b6f22b98ff2dee4fa345500bb \ No newline at end of file diff --git a/checkpoints/auido2exp_00300-model.pth b/checkpoints/auido2exp_00300-model.pth deleted file mode 100644 index 927072b6ea4aafde874e6cc0f51594f20e8dac17..0000000000000000000000000000000000000000 --- a/checkpoints/auido2exp_00300-model.pth +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b7608f0e6b477e50e03ca569ac5b04a841b9217f89d502862fc78fda4e46dec4 -size 34278319 diff --git a/checkpoints/auido2exp_00300-model.pth b/checkpoints/auido2exp_00300-model.pth new file mode 120000 index 0000000000000000000000000000000000000000..dad279ea2719742f2f194560da013c902551a67c --- /dev/null +++ b/checkpoints/auido2exp_00300-model.pth @@ -0,0 +1 @@ +../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/b7608f0e6b477e50e03ca569ac5b04a841b9217f89d502862fc78fda4e46dec4 \ No newline at end of file diff --git a/checkpoints/auido2pose_00140-model.pth b/checkpoints/auido2pose_00140-model.pth deleted file mode 100644 index db44aee66c9246710511e59b552fc041aebe5d8a..0000000000000000000000000000000000000000 --- a/checkpoints/auido2pose_00140-model.pth +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4fba6701852dc57efbed25b1e4276e4ff752941860d69fc4429f08a02326ebce -size 95916155 diff --git a/checkpoints/auido2pose_00140-model.pth b/checkpoints/auido2pose_00140-model.pth new file mode 120000 index 0000000000000000000000000000000000000000..bf80808d4a6efb07e73062176320c831c2bd4b98 --- /dev/null +++ b/checkpoints/auido2pose_00140-model.pth @@ -0,0 +1 @@ +../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/4fba6701852dc57efbed25b1e4276e4ff752941860d69fc4429f08a02326ebce \ No newline at end of file diff --git a/checkpoints/epoch_20.pth b/checkpoints/epoch_20.pth deleted file mode 100644 index 97ebd6753f7ca4bcd39d3b82e7109b66a2dbc1fb..0000000000000000000000000000000000000000 --- a/checkpoints/epoch_20.pth +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:6d17a6b23457b521801baae583cb6a58f7238fe6721fc3d65d76407460e9149b -size 288860037 diff --git a/checkpoints/epoch_20.pth b/checkpoints/epoch_20.pth new file mode 120000 index 0000000000000000000000000000000000000000..e1b14dc95869be3eb0ce72d28807076d6b153321 --- /dev/null +++ b/checkpoints/epoch_20.pth @@ -0,0 +1 @@ +../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/6d17a6b23457b521801baae583cb6a58f7238fe6721fc3d65d76407460e9149b \ No newline at end of file diff --git a/checkpoints/facevid2vid_00189-model.pth.tar b/checkpoints/facevid2vid_00189-model.pth.tar deleted file mode 100644 index 6c676eb119186e3ae866188f5b0ab2cff10473bc..0000000000000000000000000000000000000000 --- a/checkpoints/facevid2vid_00189-model.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fbad01d46f0510276dc4521322dde6824a873a4222cd0740c85762e7067ea71d -size 2112619148 diff --git a/checkpoints/facevid2vid_00189-model.pth.tar b/checkpoints/facevid2vid_00189-model.pth.tar new file mode 120000 index 0000000000000000000000000000000000000000..42186b1541038664306ea59e22e78b5a4affe0a1 --- /dev/null +++ b/checkpoints/facevid2vid_00189-model.pth.tar @@ -0,0 +1 @@ +../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/fbad01d46f0510276dc4521322dde6824a873a4222cd0740c85762e7067ea71d \ No newline at end of file diff --git a/checkpoints/hub/checkpoints/2DFAN4-cd938726ad.zip b/checkpoints/hub/checkpoints/2DFAN4-cd938726ad.zip deleted file mode 100644 index 6bb44e2ec4c154607a919de8e3d5c5448e86b586..0000000000000000000000000000000000000000 --- a/checkpoints/hub/checkpoints/2DFAN4-cd938726ad.zip +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cd938726adb1f15f361263cce2db9cb820c42585fa8796ec72ce19107f369a46 -size 96316515 diff --git a/checkpoints/hub/checkpoints/2DFAN4-cd938726ad.zip b/checkpoints/hub/checkpoints/2DFAN4-cd938726ad.zip new file mode 120000 index 0000000000000000000000000000000000000000..6b37c7e06d458bfd1d3a909fe37f1bac023db1e0 --- /dev/null +++ b/checkpoints/hub/checkpoints/2DFAN4-cd938726ad.zip @@ -0,0 +1 @@ +../../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/cd938726adb1f15f361263cce2db9cb820c42585fa8796ec72ce19107f369a46 \ No newline at end of file diff --git a/checkpoints/hub/checkpoints/s3fd-619a316812.pth b/checkpoints/hub/checkpoints/s3fd-619a316812.pth deleted file mode 100644 index 895538e7fb6df3ad6e0e80d6d48b3c4e60cd9e6c..0000000000000000000000000000000000000000 --- a/checkpoints/hub/checkpoints/s3fd-619a316812.pth +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:619a31681264d3f7f7fc7a16a42cbbe8b23f31a256f75a366e5a1bcd59b33543 -size 89843225 diff --git a/checkpoints/hub/checkpoints/s3fd-619a316812.pth b/checkpoints/hub/checkpoints/s3fd-619a316812.pth new file mode 120000 index 0000000000000000000000000000000000000000..f91ea38229a5e06a1b7f2d046317ab938a033f51 --- /dev/null +++ b/checkpoints/hub/checkpoints/s3fd-619a316812.pth @@ -0,0 +1 @@ +../../../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/619a31681264d3f7f7fc7a16a42cbbe8b23f31a256f75a366e5a1bcd59b33543 \ No newline at end of file diff --git a/checkpoints/mapping_00229-model.pth.tar b/checkpoints/mapping_00229-model.pth.tar deleted file mode 100644 index 6400233ae3fa5ff9426800ef761fd6c830bc0cd7..0000000000000000000000000000000000000000 --- a/checkpoints/mapping_00229-model.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:62a1e06006cc963220f6477438518ed86e9788226c62ae382ddc42fbcefb83f1 -size 155521183 diff --git a/checkpoints/mapping_00229-model.pth.tar b/checkpoints/mapping_00229-model.pth.tar new file mode 120000 index 0000000000000000000000000000000000000000..0c9d8ea35d9e471d338d989d235234e4481ab3d9 --- /dev/null +++ b/checkpoints/mapping_00229-model.pth.tar @@ -0,0 +1 @@ +../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/62a1e06006cc963220f6477438518ed86e9788226c62ae382ddc42fbcefb83f1 \ No newline at end of file diff --git a/checkpoints/shape_predictor_68_face_landmarks.dat b/checkpoints/shape_predictor_68_face_landmarks.dat deleted file mode 100644 index 1e5da4f9a556bec8582e6c55b89b3e6bfdd60021..0000000000000000000000000000000000000000 --- a/checkpoints/shape_predictor_68_face_landmarks.dat +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fbdc2cb80eb9aa7a758672cbfdda32ba6300efe9b6e6c7a299ff7e736b11b92f -size 99693937 diff --git a/checkpoints/shape_predictor_68_face_landmarks.dat b/checkpoints/shape_predictor_68_face_landmarks.dat new file mode 120000 index 0000000000000000000000000000000000000000..ab6f865c4370dcad68dd6dfb896cc3bcdb7aba6e --- /dev/null +++ b/checkpoints/shape_predictor_68_face_landmarks.dat @@ -0,0 +1 @@ +../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/fbdc2cb80eb9aa7a758672cbfdda32ba6300efe9b6e6c7a299ff7e736b11b92f \ No newline at end of file diff --git a/checkpoints/wav2lip.pth b/checkpoints/wav2lip.pth deleted file mode 100644 index c575a07ac4e62abfd60cb8681ebb6df241cf31e6..0000000000000000000000000000000000000000 --- a/checkpoints/wav2lip.pth +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b78b681b68ad9fe6c6fb1debc6ff43ad05834a8af8a62ffc4167b7b34ef63c37 -size 435807851 diff --git a/checkpoints/wav2lip.pth b/checkpoints/wav2lip.pth new file mode 120000 index 0000000000000000000000000000000000000000..16beaee282514c3ae6425b25e044256b44fb4d38 --- /dev/null +++ b/checkpoints/wav2lip.pth @@ -0,0 +1 @@ +../../../../root/.cache/huggingface/hub/models--vinthony--SadTalker/blobs/b78b681b68ad9fe6c6fb1debc6ff43ad05834a8af8a62ffc4167b7b34ef63c37 \ No newline at end of file diff --git a/docs/sadtalker_logo.png b/docs/sadtalker_logo.png new file mode 100644 index 0000000000000000000000000000000000000000..107aebfab5c2cdf842d7fc29ceed657e8c83ad56 Binary files /dev/null and b/docs/sadtalker_logo.png differ diff --git a/examples/source_image/full3.png b/examples/source_image/full3.png new file mode 100644 index 0000000000000000000000000000000000000000..40cd6d6d3c5b95c29d6648c2ba7d7e27c9781970 Binary files /dev/null and b/examples/source_image/full3.png differ diff --git a/examples/source_image/full4.jpeg b/examples/source_image/full4.jpeg new file mode 100644 index 0000000000000000000000000000000000000000..79f17f52123e8d173600e0df138a30e98ba2c6f3 Binary files /dev/null and b/examples/source_image/full4.jpeg differ diff --git a/inference.py b/inference.py deleted file mode 100644 index 3582c98cf92b34c93fdf3df585aeb84b2c7d77f4..0000000000000000000000000000000000000000 --- a/inference.py +++ /dev/null @@ -1,134 +0,0 @@ -import torch -from time import strftime -import os, sys, time -from argparse import ArgumentParser - -from src.utils.preprocess import CropAndExtract -from src.test_audio2coeff import Audio2Coeff -from src.facerender.animate import AnimateFromCoeff -from src.generate_batch import get_data -from src.generate_facerender_batch import get_facerender_data - -def main(args): - #torch.backends.cudnn.enabled = False - - pic_path = args.source_image - audio_path = args.driven_audio - save_dir = os.path.join(args.result_dir, strftime("%Y_%m_%d_%H.%M.%S")) - os.makedirs(save_dir, exist_ok=True) - pose_style = args.pose_style - device = args.device - batch_size = args.batch_size - camera_yaw_list = args.camera_yaw - camera_pitch_list = args.camera_pitch - camera_roll_list = args.camera_roll - - current_code_path = sys.argv[0] - current_root_path = os.path.split(current_code_path)[0] - - os.environ['TORCH_HOME']=os.path.join(current_root_path, args.checkpoint_dir) - - path_of_lm_croper = os.path.join(current_root_path, args.checkpoint_dir, 'shape_predictor_68_face_landmarks.dat') - path_of_net_recon_model = os.path.join(current_root_path, args.checkpoint_dir, 'epoch_20.pth') - dir_of_BFM_fitting = os.path.join(current_root_path, args.checkpoint_dir, 'BFM_Fitting') - wav2lip_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'wav2lip.pth') - - audio2pose_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'auido2pose_00140-model.pth') - audio2pose_yaml_path = os.path.join(current_root_path, 'src', 'config', 'auido2pose.yaml') - - audio2exp_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'auido2exp_00300-model.pth') - audio2exp_yaml_path = os.path.join(current_root_path, 'src', 'config', 'auido2exp.yaml') - - free_view_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'facevid2vid_00189-model.pth.tar') - mapping_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'mapping_00229-model.pth.tar') - facerender_yaml_path = os.path.join(current_root_path, 'src', 'config', 'facerender.yaml') - - #init model - print(path_of_net_recon_model) - preprocess_model = CropAndExtract(path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting, device) - - print(audio2pose_checkpoint) - print(audio2exp_checkpoint) - audio_to_coeff = Audio2Coeff(audio2pose_checkpoint, audio2pose_yaml_path, - audio2exp_checkpoint, audio2exp_yaml_path, - wav2lip_checkpoint, device) - - print(free_view_checkpoint) - print(mapping_checkpoint) - animate_from_coeff = AnimateFromCoeff(free_view_checkpoint, mapping_checkpoint, - facerender_yaml_path, device) - - #crop image and extract 3dmm from image - first_frame_dir = os.path.join(save_dir, 'first_frame_dir') - os.makedirs(first_frame_dir, exist_ok=True) - first_coeff_path, crop_pic_path = preprocess_model.generate(pic_path, first_frame_dir) - if first_coeff_path is None: - print("Can't get the coeffs of the input") - return - - #audio2ceoff - batch = get_data(first_coeff_path, audio_path, device) - coeff_path = audio_to_coeff.generate(batch, save_dir, pose_style) - - # 3dface render - if args.face3dvis: - from src.face3d.visualize import gen_composed_video - gen_composed_video(args, device, first_coeff_path, coeff_path, audio_path, os.path.join(save_dir, '3dface.mp4')) - - #coeff2video - data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, - batch_size, camera_yaw_list, camera_pitch_list, camera_roll_list, - expression_scale=args.expression_scale, still_mode=args.still) - - animate_from_coeff.generate(data, save_dir, enhancer=args.enhancer) - video_name = data['video_name'] - - if args.enhancer is not None: - print(f'The generated video is named {video_name}_enhanced in {save_dir}') - else: - print(f'The generated video is named {video_name} in {save_dir}') - - return os.path.join(save_dir, video_name+'.mp4'), os.path.join(save_dir, video_name+'.mp4') - - -if __name__ == '__main__': - - parser = ArgumentParser() - parser.add_argument("--driven_audio", default='./examples/driven_audio/japanese.wav', help="path to driven audio") - parser.add_argument("--source_image", default='./examples/source_image/art_0.png', help="path to source image") - parser.add_argument("--checkpoint_dir", default='./checkpoints', help="path to output") - parser.add_argument("--result_dir", default='./results', help="path to output") - parser.add_argument("--pose_style", type=int, default=0, help="input pose style from [0, 46)") - parser.add_argument("--batch_size", type=int, default=2, help="the batch size of facerender") - parser.add_argument("--expression_scale", type=float, default=1., help="the batch size of facerender") - parser.add_argument('--camera_yaw', nargs='+', type=int, default=[0], help="the camera yaw degree") - parser.add_argument('--camera_pitch', nargs='+', type=int, default=[0], help="the camera pitch degree") - parser.add_argument('--camera_roll', nargs='+', type=int, default=[0], help="the camera roll degree") - parser.add_argument('--enhancer', type=str, default=None, help="Face enhancer, [GFPGAN]") - parser.add_argument("--cpu", dest="cpu", action="store_true") - parser.add_argument("--face3dvis", action="store_true", help="generate 3d face and 3d landmarks") - parser.add_argument("--still", action="store_true") - - # net structure and parameters - parser.add_argument('--net_recon', type=str, default='resnet50', choices=['resnet18', 'resnet34', 'resnet50'], help='not use') - parser.add_argument('--init_path', type=str, default=None, help='not Use') - parser.add_argument('--use_last_fc',default=False, help='zero initialize the last fc') - parser.add_argument('--bfm_folder', type=str, default='./checkpoints/BFM_Fitting/') - parser.add_argument('--bfm_model', type=str, default='BFM_model_front.mat', help='bfm model') - - # default renderer parameters - parser.add_argument('--focal', type=float, default=1015.) - parser.add_argument('--center', type=float, default=112.) - parser.add_argument('--camera_d', type=float, default=10.) - parser.add_argument('--z_near', type=float, default=5.) - parser.add_argument('--z_far', type=float, default=15.) - - args = parser.parse_args() - - if torch.cuda.is_available() and not args.cpu: - args.device = "cuda" - else: - args.device = "cpu" - - main(args) - diff --git a/modules/__pycache__/gfpgan_inference.cpython-38.pyc b/modules/__pycache__/gfpgan_inference.cpython-38.pyc deleted file mode 100644 index 83f5c7d7953f80bd717ae9a965fbe7b7a6d0e3a3..0000000000000000000000000000000000000000 Binary files a/modules/__pycache__/gfpgan_inference.cpython-38.pyc and /dev/null differ diff --git a/modules/__pycache__/gfpgan_inference.cpython-39.pyc b/modules/__pycache__/gfpgan_inference.cpython-39.pyc deleted file mode 100644 index e0ecebb67b685d1cd74fd4266ddb33d75f3c60b2..0000000000000000000000000000000000000000 Binary files a/modules/__pycache__/gfpgan_inference.cpython-39.pyc and /dev/null differ diff --git a/modules/__pycache__/sadtalker_test.cpython-38.pyc b/modules/__pycache__/sadtalker_test.cpython-38.pyc deleted file mode 100644 index a96311c6eee958b442fec8776d088b74e7b8b3a2..0000000000000000000000000000000000000000 Binary files a/modules/__pycache__/sadtalker_test.cpython-38.pyc and /dev/null differ diff --git a/modules/__pycache__/sadtalker_test.cpython-39.pyc b/modules/__pycache__/sadtalker_test.cpython-39.pyc deleted file mode 100644 index 9b73e50c9766b517d9d765d1e3e58b8a7153b8dd..0000000000000000000000000000000000000000 Binary files a/modules/__pycache__/sadtalker_test.cpython-39.pyc and /dev/null differ diff --git a/modules/__pycache__/text2speech.cpython-38.pyc b/modules/__pycache__/text2speech.cpython-38.pyc deleted file mode 100644 index 90ad4127ce0050c2215bdb797974ad849d12a96c..0000000000000000000000000000000000000000 Binary files a/modules/__pycache__/text2speech.cpython-38.pyc and /dev/null differ diff --git a/modules/__pycache__/text2speech.cpython-39.pyc b/modules/__pycache__/text2speech.cpython-39.pyc deleted file mode 100644 index 74134d79b4e1171fed217853c1457430705b6616..0000000000000000000000000000000000000000 Binary files a/modules/__pycache__/text2speech.cpython-39.pyc and /dev/null differ diff --git a/modules/gfpgan_inference.py b/modules/gfpgan_inference.py deleted file mode 100644 index f4e7dc80eac012906b797843aa6019c2c4a39b3b..0000000000000000000000000000000000000000 --- a/modules/gfpgan_inference.py +++ /dev/null @@ -1,36 +0,0 @@ -import os,sys - -def gfpgan(scale, origin_mp4_path): - current_code_path = sys.argv[0] - current_root_path = os.path.split(current_code_path)[0] - print(current_root_path) - gfpgan_code_path = current_root_path+'/repositories/GFPGAN/inference_gfpgan.py' - print(gfpgan_code_path) - - #video2pic - result_dir = os.path.split(origin_mp4_path)[0] - video_name = os.path.split(origin_mp4_path)[1] - video_name = video_name.split('.')[0] - print(video_name) - str_scale = str(scale).replace('.', '_') - output_mp4_path = os.path.join(result_dir, video_name+'##'+str_scale+'.mp4') - temp_output_mp4_path = os.path.join(result_dir, 'temp_'+video_name+'##'+str_scale+'.mp4') - - audio_name = video_name.split('##')[-1] - audio_path = os.path.join(result_dir, audio_name+'.wav') - temp_pic_dir1 = os.path.join(result_dir, video_name) - temp_pic_dir2 = os.path.join(result_dir, video_name+'##'+str_scale) - os.makedirs(temp_pic_dir1, exist_ok=True) - os.makedirs(temp_pic_dir2, exist_ok=True) - cmd1 = 'ffmpeg -i \"{}\" -start_number 0 \"{}\"/%06d.png -loglevel error -y'.format(origin_mp4_path, temp_pic_dir1) - os.system(cmd1) - cmd2 = f'python {gfpgan_code_path} -i {temp_pic_dir1} -o {temp_pic_dir2} -s {scale}' - os.system(cmd2) - cmd3 = f'ffmpeg -r 25 -f image2 -i {temp_pic_dir2}/%06d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p {temp_output_mp4_path}' - os.system(cmd3) - cmd4 = f'ffmpeg -y -i {temp_output_mp4_path} -i {audio_path} -vcodec copy {output_mp4_path}' - os.system(cmd4) - #shutil.rmtree(temp_pic_dir1) - #shutil.rmtree(temp_pic_dir2) - - return output_mp4_path diff --git a/modules/sadtalker_test.py b/modules/sadtalker_test.py deleted file mode 100644 index 34d9699f71fcd6d8f413f9cc96926dd6ceff36b1..0000000000000000000000000000000000000000 --- a/modules/sadtalker_test.py +++ /dev/null @@ -1,118 +0,0 @@ -import torch -import os, sys, shutil -from src.utils.preprocess import CropAndExtract -from src.test_audio2coeff import Audio2Coeff -from src.facerender.animate import AnimateFromCoeff -from src.generate_batch import get_data -from src.generate_facerender_batch import get_facerender_data -import uuid - -from pydub import AudioSegment - -def mp3_to_wav(mp3_filename,wav_filename,frame_rate): - mp3_file = AudioSegment.from_file(file=mp3_filename) - mp3_file.set_frame_rate(frame_rate).export(wav_filename,format="wav") - -from modules.text2speech import text2speech - -class SadTalker(): - - def __init__(self, checkpoint_path='checkpoints'): - - if torch.cuda.is_available() : - device = "cuda" - else: - device = "cpu" - - # current_code_path = sys.argv[0] - # modules_path = os.path.split(current_code_path)[0] - - current_root_path = './' - - os.environ['TORCH_HOME']=os.path.join(current_root_path, 'checkpoints') - - path_of_lm_croper = os.path.join(current_root_path, 'checkpoints', 'shape_predictor_68_face_landmarks.dat') - path_of_net_recon_model = os.path.join(current_root_path, 'checkpoints', 'epoch_20.pth') - dir_of_BFM_fitting = os.path.join(current_root_path, 'checkpoints', 'BFM_Fitting') - wav2lip_checkpoint = os.path.join(current_root_path, 'checkpoints', 'wav2lip.pth') - - audio2pose_checkpoint = os.path.join(current_root_path, 'checkpoints', 'auido2pose_00140-model.pth') - audio2pose_yaml_path = os.path.join(current_root_path, 'config', 'auido2pose.yaml') - - audio2exp_checkpoint = os.path.join(current_root_path, 'checkpoints', 'auido2exp_00300-model.pth') - audio2exp_yaml_path = os.path.join(current_root_path, 'config', 'auido2exp.yaml') - - free_view_checkpoint = os.path.join(current_root_path, 'checkpoints', 'facevid2vid_00189-model.pth.tar') - mapping_checkpoint = os.path.join(current_root_path, 'checkpoints', 'mapping_00229-model.pth.tar') - facerender_yaml_path = os.path.join(current_root_path, 'config', 'facerender.yaml') - - #init model - print(path_of_lm_croper) - self.preprocess_model = CropAndExtract(path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting, device) - - print(audio2pose_checkpoint) - self.audio_to_coeff = Audio2Coeff(audio2pose_checkpoint, audio2pose_yaml_path, - audio2exp_checkpoint, audio2exp_yaml_path, wav2lip_checkpoint, device) - print(free_view_checkpoint) - self.animate_from_coeff = AnimateFromCoeff(free_view_checkpoint, mapping_checkpoint, - facerender_yaml_path, device) - self.device = device - - def test(self, source_image, driven_audio, still_mode, resize_mode, use_enhancer, result_dir='./'): - - time_tag = str(uuid.uuid4()) # strftime("%Y_%m_%d_%H.%M.%S") - save_dir = os.path.join(result_dir, time_tag) - os.makedirs(save_dir, exist_ok=True) - - input_dir = os.path.join(save_dir, 'input') - os.makedirs(input_dir, exist_ok=True) - - print(source_image) - pic_path = os.path.join(input_dir, os.path.basename(source_image)) - shutil.move(source_image, input_dir) - - if os.path.isfile(driven_audio): - audio_path = os.path.join(input_dir, os.path.basename(driven_audio)) - - #### mp3 to wav - if '.mp3' in audio_path: - mp3_to_wav(driven_audio, audio_path.replace('.mp3', '.wav'), 16000) - audio_path = audio_path.replace('.mp3', '.wav') - else: - shutil.move(driven_audio, input_dir) - else: - text2speech - - - os.makedirs(save_dir, exist_ok=True) - pose_style = 0 - #crop image and extract 3dmm from image - first_frame_dir = os.path.join(save_dir, 'first_frame_dir') - os.makedirs(first_frame_dir, exist_ok=True) - first_coeff_path, crop_pic_path, original_size = self.preprocess_model.generate(pic_path, first_frame_dir, crop_or_resize= 'resize' if resize_mode else 'crop') - if first_coeff_path is None: - raise AttributeError("No face is detected") - - #audio2ceoff - batch = get_data(first_coeff_path, audio_path, self.device) - coeff_path = self.audio_to_coeff.generate(batch, save_dir, pose_style) - #coeff2video - batch_size = 4 - data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode) - self.animate_from_coeff.generate(data, save_dir, enhancer='gfpgan' if use_enhancer else None, original_size=original_size) - video_name = data['video_name'] - print(f'The generated video is named {video_name} in {save_dir}') - - torch.cuda.empty_cache() - torch.cuda.synchronize() - - import gc; gc.collect() - - if use_enhancer: - return os.path.join(save_dir, video_name+'_enhanced.mp4'), os.path.join(save_dir, video_name+'_enhanced.mp4') - - else: - return os.path.join(save_dir, video_name+'.mp4'), os.path.join(save_dir, video_name+'.mp4') - - - \ No newline at end of file diff --git a/modules/text2speech.py b/modules/text2speech.py deleted file mode 100644 index 3ecaef36961494c8b2b1f5771a70b997efa04ffd..0000000000000000000000000000000000000000 --- a/modules/text2speech.py +++ /dev/null @@ -1,12 +0,0 @@ -import os - -def text2speech(txt, audio_path): - print(txt) - cmd = f'tts --text "{txt}" --out_path {audio_path}' - print(cmd) - try: - os.system(cmd) - return audio_path - except: - print("Error: Failed convert txt to audio") - return None \ No newline at end of file diff --git a/src/__pycache__/generate_batch.cpython-38.pyc b/src/__pycache__/generate_batch.cpython-38.pyc index dc3eb4726e9835d34c08362da995941fef530b8f..dbd150e74cc1628497672277155e495938f87231 100644 Binary files a/src/__pycache__/generate_batch.cpython-38.pyc and b/src/__pycache__/generate_batch.cpython-38.pyc differ diff --git a/src/__pycache__/generate_facerender_batch.cpython-38.pyc b/src/__pycache__/generate_facerender_batch.cpython-38.pyc index cc944270498549b70e901f5b1c764d1d832eb49e..2b6fb645533cd9c22624032834274954dc4fdacb 100644 Binary files a/src/__pycache__/generate_facerender_batch.cpython-38.pyc and b/src/__pycache__/generate_facerender_batch.cpython-38.pyc differ diff --git a/src/__pycache__/test_audio2coeff.cpython-38.pyc b/src/__pycache__/test_audio2coeff.cpython-38.pyc index a6d261868c02b57145618adcd583481cf623e391..fec3aa3f8b75436e96d84d5d88014ba7bd2a79fe 100644 Binary files a/src/__pycache__/test_audio2coeff.cpython-38.pyc and b/src/__pycache__/test_audio2coeff.cpython-38.pyc differ diff --git a/src/audio2exp_models/__pycache__/audio2exp.cpython-38.pyc b/src/audio2exp_models/__pycache__/audio2exp.cpython-38.pyc index de88551314f6c19ad1f5b5b33704f1303f51e029..330ddec8004f4d5b50200d658aa8377ee44628f2 100644 Binary files a/src/audio2exp_models/__pycache__/audio2exp.cpython-38.pyc and b/src/audio2exp_models/__pycache__/audio2exp.cpython-38.pyc differ diff --git a/src/audio2exp_models/__pycache__/networks.cpython-38.pyc b/src/audio2exp_models/__pycache__/networks.cpython-38.pyc index d703bd9e8f3d0c27c16fa713bba3d0969e984ad3..e99c0c180aba2c3bdcff2c5711d4e245533283b2 100644 Binary files a/src/audio2exp_models/__pycache__/networks.cpython-38.pyc and b/src/audio2exp_models/__pycache__/networks.cpython-38.pyc differ diff --git a/src/audio2pose_models/__pycache__/audio2pose.cpython-38.pyc b/src/audio2pose_models/__pycache__/audio2pose.cpython-38.pyc index 5b2dcc996a73224e972148e252fb4e2deedd69a5..9d513edddc0887e16d5329d014db72cae9e7dff7 100644 Binary files a/src/audio2pose_models/__pycache__/audio2pose.cpython-38.pyc and b/src/audio2pose_models/__pycache__/audio2pose.cpython-38.pyc differ diff --git a/src/audio2pose_models/__pycache__/audio_encoder.cpython-38.pyc b/src/audio2pose_models/__pycache__/audio_encoder.cpython-38.pyc index b0f11a59fea18ee93c30da5cd4c94d04897ea010..3e22cb39b5f8e425ce332c1c6378a3d637e88cd4 100644 Binary files a/src/audio2pose_models/__pycache__/audio_encoder.cpython-38.pyc and b/src/audio2pose_models/__pycache__/audio_encoder.cpython-38.pyc differ diff --git a/src/audio2pose_models/__pycache__/cvae.cpython-38.pyc b/src/audio2pose_models/__pycache__/cvae.cpython-38.pyc index 1aa0e494be950e6ca972390b27f2dddc8be6d193..748551cd87803ab44d896402a00ff59099bdd40b 100644 Binary files a/src/audio2pose_models/__pycache__/cvae.cpython-38.pyc and b/src/audio2pose_models/__pycache__/cvae.cpython-38.pyc differ diff --git a/src/audio2pose_models/__pycache__/discriminator.cpython-38.pyc b/src/audio2pose_models/__pycache__/discriminator.cpython-38.pyc index 817b8836123ed1a3b5795d912d84c3ff54d7accc..67438ac9ff2f17c052a28f37f109bcd9ab1a6638 100644 Binary files a/src/audio2pose_models/__pycache__/discriminator.cpython-38.pyc and b/src/audio2pose_models/__pycache__/discriminator.cpython-38.pyc differ diff --git a/src/audio2pose_models/__pycache__/networks.cpython-38.pyc b/src/audio2pose_models/__pycache__/networks.cpython-38.pyc index d18f56064377373a8f4f400c59379b0b79d9f649..858cd56d12672e39b957152070c77b92f0929260 100644 Binary files a/src/audio2pose_models/__pycache__/networks.cpython-38.pyc and b/src/audio2pose_models/__pycache__/networks.cpython-38.pyc differ diff --git a/src/audio2pose_models/__pycache__/res_unet.cpython-38.pyc b/src/audio2pose_models/__pycache__/res_unet.cpython-38.pyc index 5aa2863a646a6eb8b44e0ebdebc5c21b562c2f39..8c8ea5b2ed612c1e68c73e63763219998e8cd2ab 100644 Binary files a/src/audio2pose_models/__pycache__/res_unet.cpython-38.pyc and b/src/audio2pose_models/__pycache__/res_unet.cpython-38.pyc differ diff --git a/src/config/facerender_still.yaml b/src/config/facerender_still.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6b4d66dade3e655ac4cfc25a994ca28e53821d80 --- /dev/null +++ b/src/config/facerender_still.yaml @@ -0,0 +1,45 @@ +model_params: + common_params: + num_kp: 15 + image_channel: 3 + feature_channel: 32 + estimate_jacobian: False # True + kp_detector_params: + temperature: 0.1 + block_expansion: 32 + max_features: 1024 + scale_factor: 0.25 # 0.25 + num_blocks: 5 + reshape_channel: 16384 # 16384 = 1024 * 16 + reshape_depth: 16 + he_estimator_params: + block_expansion: 64 + max_features: 2048 + num_bins: 66 + generator_params: + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + reshape_channel: 32 + reshape_depth: 16 # 512 = 32 * 16 + num_resblocks: 6 + estimate_occlusion_map: True + dense_motion_params: + block_expansion: 32 + max_features: 1024 + num_blocks: 5 + reshape_depth: 16 + compress: 4 + discriminator_params: + scales: [1] + block_expansion: 32 + max_features: 512 + num_blocks: 4 + sn: True + mapping_params: + coeff_nc: 73 + descriptor_nc: 1024 + layer: 3 + num_kp: 15 + num_bins: 66 + diff --git a/src/face3d/__pycache__/extract_kp_videos.cpython-38.pyc b/src/face3d/__pycache__/extract_kp_videos.cpython-38.pyc index 25b9b1377b35ea7231f4d3b44d81aab8d44f4b5b..57bd4c2e63f9c3685c9b29100d32d407283b8377 100644 Binary files a/src/face3d/__pycache__/extract_kp_videos.cpython-38.pyc and b/src/face3d/__pycache__/extract_kp_videos.cpython-38.pyc differ diff --git a/src/face3d/models/__pycache__/__init__.cpython-38.pyc b/src/face3d/models/__pycache__/__init__.cpython-38.pyc index 023f4afb376ad418cc6e3cdd9e821cfa0bcd33f3..d30126e21e09597ba773cf390c2756a2c1957351 100644 Binary files a/src/face3d/models/__pycache__/__init__.cpython-38.pyc and b/src/face3d/models/__pycache__/__init__.cpython-38.pyc differ diff --git a/src/face3d/models/__pycache__/base_model.cpython-38.pyc b/src/face3d/models/__pycache__/base_model.cpython-38.pyc index 1076d15ca87eb8922a4fb3706a3aff777187b612..ae28165f30ebacc1f7c7386e9394d566242c3e03 100644 Binary files a/src/face3d/models/__pycache__/base_model.cpython-38.pyc and b/src/face3d/models/__pycache__/base_model.cpython-38.pyc differ diff --git a/src/face3d/models/__pycache__/networks.cpython-38.pyc b/src/face3d/models/__pycache__/networks.cpython-38.pyc index e52b5dac3ce0e017ed844aed711ddfb94223be98..7645026a4700a47fc0f6110f57de74d30d5a95d0 100644 Binary files a/src/face3d/models/__pycache__/networks.cpython-38.pyc and b/src/face3d/models/__pycache__/networks.cpython-38.pyc differ diff --git a/src/face3d/models/arcface_torch/backbones/__pycache__/__init__.cpython-38.pyc b/src/face3d/models/arcface_torch/backbones/__pycache__/__init__.cpython-38.pyc index a891077dd80e455e762875f37b16ff11e58441e7..dfbb4a34bcdfed39732dc25091c2c3e0d0aa2843 100644 Binary files a/src/face3d/models/arcface_torch/backbones/__pycache__/__init__.cpython-38.pyc and b/src/face3d/models/arcface_torch/backbones/__pycache__/__init__.cpython-38.pyc differ diff --git a/src/face3d/models/arcface_torch/backbones/__pycache__/iresnet.cpython-38.pyc b/src/face3d/models/arcface_torch/backbones/__pycache__/iresnet.cpython-38.pyc index e7d3278234555217f1055e02d930d1cd8731afa1..97a4dff040618d8ae69b2f7beab87fa9060a7c3b 100644 Binary files a/src/face3d/models/arcface_torch/backbones/__pycache__/iresnet.cpython-38.pyc and b/src/face3d/models/arcface_torch/backbones/__pycache__/iresnet.cpython-38.pyc differ diff --git a/src/face3d/models/arcface_torch/backbones/__pycache__/mobilefacenet.cpython-38.pyc b/src/face3d/models/arcface_torch/backbones/__pycache__/mobilefacenet.cpython-38.pyc index db57e8b41e4fe5bdbee04db62986c15c0e4bffb1..4707dd4605a497e1419d51c2af6d5dd4c9bbff0e 100644 Binary files a/src/face3d/models/arcface_torch/backbones/__pycache__/mobilefacenet.cpython-38.pyc and b/src/face3d/models/arcface_torch/backbones/__pycache__/mobilefacenet.cpython-38.pyc differ diff --git a/src/face3d/util/__pycache__/__init__.cpython-38.pyc b/src/face3d/util/__pycache__/__init__.cpython-38.pyc index 2671705d02bed0a099b4a375070d0949c1450b7b..ad535758cbcaa062f6b139651995e2e3393fef1b 100644 Binary files a/src/face3d/util/__pycache__/__init__.cpython-38.pyc and b/src/face3d/util/__pycache__/__init__.cpython-38.pyc differ diff --git a/src/face3d/util/__pycache__/load_mats.cpython-38.pyc b/src/face3d/util/__pycache__/load_mats.cpython-38.pyc index f44224c0f7c12afc3590f10b9f5ac570b6b668bb..a870c8eb0fc0ea69f1c1fe858cabb3a0c66a0bae 100644 Binary files a/src/face3d/util/__pycache__/load_mats.cpython-38.pyc and b/src/face3d/util/__pycache__/load_mats.cpython-38.pyc differ diff --git a/src/face3d/util/__pycache__/preprocess.cpython-38.pyc b/src/face3d/util/__pycache__/preprocess.cpython-38.pyc index 90eb37261ae38ab925f149db62d91a1d0078bfcf..2912472ca888e6abfd4ee489a987a3aa24314086 100644 Binary files a/src/face3d/util/__pycache__/preprocess.cpython-38.pyc and b/src/face3d/util/__pycache__/preprocess.cpython-38.pyc differ diff --git a/src/facerender/__pycache__/animate.cpython-38.pyc b/src/facerender/__pycache__/animate.cpython-38.pyc index 1f8003ddb550fc6e235abccfb5f8481ee8c16afa..dd7578cb508809973b2609304e8abfb3d1877983 100644 Binary files a/src/facerender/__pycache__/animate.cpython-38.pyc and b/src/facerender/__pycache__/animate.cpython-38.pyc differ diff --git a/src/facerender/animate.py b/src/facerender/animate.py index 1bd221ad4c99d911222fdf1eb087ebb626afc867..9a2953d10367f8094842fafb2dccb11746ba29b7 100644 --- a/src/facerender/animate.py +++ b/src/facerender/animate.py @@ -4,6 +4,7 @@ import yaml import numpy as np import warnings from skimage import img_as_ubyte + warnings.filterwarnings('ignore') import imageio @@ -17,7 +18,7 @@ from src.facerender.modules.make_animation import make_animation from pydub import AudioSegment from src.utils.face_enhancer import enhancer as face_enhancer from src.utils.paste_pic import paste_pic - +from src.utils.videoio import save_video_with_watermark class AnimateFromCoeff(): @@ -116,7 +117,7 @@ class AnimateFromCoeff(): return checkpoint['epoch'] - def generate(self, x, video_save_dir, pic_path, crop_info, enhancer=None, full_img_enhancer=None): + def generate(self, x, video_save_dir, pic_path, crop_info, enhancer=None, background_enhancer=None, preprocess='crop'): source_image=x['source_image'].type(torch.FloatTensor) source_semantics=x['source_semantics'].type(torch.FloatTensor) @@ -165,17 +166,6 @@ class AnimateFromCoeff(): path = os.path.join(video_save_dir, 'temp_'+video_name) imageio.mimsave(path, result, fps=float(25)) - if enhancer: - video_name_enhancer = x['video_name'] + '_enhanced.mp4' - av_path_enhancer = os.path.join(video_save_dir, video_name_enhancer) - enhanced_path = os.path.join(video_save_dir, 'temp_'+video_name_enhancer) - enhanced_images = face_enhancer(result, method=enhancer) - - if original_size: - enhanced_images = [ cv2.resize(result_i,(256, int(256.0 * original_size[1]/original_size[0]) )) for result_i in enhanced_images ] - - imageio.mimsave(enhanced_path, enhanced_images, fps=float(25)) - av_path = os.path.join(video_save_dir, video_name) return_path = av_path @@ -190,27 +180,31 @@ class AnimateFromCoeff(): word = word1[start_time:end_time] word.export(new_audio_path, format="wav") - cmd = r'ffmpeg -y -i "%s" -i "%s" -vcodec copy "%s"' % (path, new_audio_path, av_path) - os.system(cmd) + save_video_with_watermark(path, new_audio_path, av_path, watermark= None) print(f'The generated video is named {video_name} in {video_save_dir}') - if enhancer: - return_path = av_path_enhancer - cmd = r'ffmpeg -y -i "%s" -i "%s" -vcodec copy "%s"' % (enhanced_path, new_audio_path, av_path_enhancer) - os.system(cmd) - os.remove(enhanced_path) - print(f'The generated video is named {video_name_enhancer} in {video_save_dir}') - - if len(crop_info) == 3: + if preprocess.lower() == 'full': + # only add watermark to the full image. video_name_full = x['video_name'] + '_full.mp4' full_video_path = os.path.join(video_save_dir, video_name_full) return_path = full_video_path - if enhancer: - paste_pic(av_path_enhancer, pic_path, crop_info, new_audio_path, full_video_path) - else: - paste_pic(path, pic_path, crop_info, new_audio_path, full_video_path) - print(f'The generated video is named {video_name_full} in {video_save_dir}') + paste_pic(path, pic_path, crop_info, new_audio_path, full_video_path) + print(f'The generated video is named {video_save_dir}/{video_name_full}') + else: + full_video_path = av_path + #### paste back then enhancers + if enhancer: + video_name_enhancer = x['video_name'] + '_enhanced.mp4' + enhanced_path = os.path.join(video_save_dir, 'temp_'+video_name_enhancer) + av_path_enhancer = os.path.join(video_save_dir, video_name_enhancer) + return_path = av_path_enhancer + enhanced_images = face_enhancer(full_video_path, method=enhancer, bg_upsampler=background_enhancer) + imageio.mimsave(enhanced_path, enhanced_images, fps=float(25)) + + save_video_with_watermark(enhanced_path, new_audio_path, av_path_enhancer, watermark= None) + print(f'The generated video is named {video_save_dir}/{video_name_enhancer}') + os.remove(enhanced_path) os.remove(path) os.remove(new_audio_path) diff --git a/src/facerender/modules/__pycache__/dense_motion.cpython-38.pyc b/src/facerender/modules/__pycache__/dense_motion.cpython-38.pyc index 7558dbc6512fceb2147fd1fae031212d07e4449d..0493b55d3eac4360fb537c992f33f732abdf6980 100644 Binary files a/src/facerender/modules/__pycache__/dense_motion.cpython-38.pyc and b/src/facerender/modules/__pycache__/dense_motion.cpython-38.pyc differ diff --git a/src/facerender/modules/__pycache__/generator.cpython-38.pyc b/src/facerender/modules/__pycache__/generator.cpython-38.pyc index 11aa36c10f79820e84d8a275234b85b0371cc050..e571c3490dd5726d1d872bb3b6bc45416f647609 100644 Binary files a/src/facerender/modules/__pycache__/generator.cpython-38.pyc and b/src/facerender/modules/__pycache__/generator.cpython-38.pyc differ diff --git a/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc b/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc index e0bd1dcd3e98a316628449370f08dc8bd2dde4b9..6b24ae7eb35dcbdd5db23dba48ae71385e03ab85 100644 Binary files a/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc and b/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc differ diff --git a/src/facerender/modules/__pycache__/make_animation.cpython-38.pyc b/src/facerender/modules/__pycache__/make_animation.cpython-38.pyc index 76e338a936f0354c81abaa5fc677c5622db16eb3..d09b8930133c93a61d8a44c114d625d1d786a57e 100644 Binary files a/src/facerender/modules/__pycache__/make_animation.cpython-38.pyc and b/src/facerender/modules/__pycache__/make_animation.cpython-38.pyc differ diff --git a/src/facerender/modules/__pycache__/mapping.cpython-38.pyc b/src/facerender/modules/__pycache__/mapping.cpython-38.pyc index b464c917a4d3feb94fa629b3390c000af89ceb9a..49f9cef9b41c97dde5b19e2c8e0865644315437e 100644 Binary files a/src/facerender/modules/__pycache__/mapping.cpython-38.pyc and b/src/facerender/modules/__pycache__/mapping.cpython-38.pyc differ diff --git a/src/facerender/modules/__pycache__/util.cpython-38.pyc b/src/facerender/modules/__pycache__/util.cpython-38.pyc index 4f4d1a6d0e3797390e942821e1e2c238e1c8a8d2..c38231ac3ea3cbe7267ef898cefd138283e111a1 100644 Binary files a/src/facerender/modules/__pycache__/util.cpython-38.pyc and b/src/facerender/modules/__pycache__/util.cpython-38.pyc differ diff --git a/src/facerender/modules/generator.py b/src/facerender/modules/generator.py index 1b5f8d26b18a1fa5cb1d8cbe9d1fa2413bf39f01..5a9edcb3b328d3afc99072b2461d7ca69919f813 100644 --- a/src/facerender/modules/generator.py +++ b/src/facerender/modules/generator.py @@ -225,6 +225,8 @@ class OcclusionAwareSPADEGenerator(nn.Module): kp_source=kp_source) output_dict['mask'] = dense_motion['mask'] + # import pdb; pdb.set_trace() + if 'occlusion_map' in dense_motion: occlusion_map = dense_motion['occlusion_map'] output_dict['occlusion_map'] = occlusion_map @@ -238,6 +240,8 @@ class OcclusionAwareSPADEGenerator(nn.Module): out = self.third(out) out = self.fourth(out) + # occlusion_map = torch.where(occlusion_map < 0.95, 0, occlusion_map) + if occlusion_map is not None: if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]: occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear') diff --git a/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc b/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc index a08f1284e68bb6251119739bc46a2dab9f5a171b..6d153a115f154ffec2d271c718046c3883e734ac 100644 Binary files a/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc and b/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc differ diff --git a/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc b/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc index f1a96eace36b537e5cfc85be1be94616151aca85..235345f1d6928bbae647a8fa1fef4df7b7bd3e29 100644 Binary files a/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc and b/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc differ diff --git a/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc b/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc index e6578b03a7060d9b9b31681e6f7ef27e4251f52e..5489fed1d07118271b9bcb6ae4c889b072dee8f4 100644 Binary files a/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc and b/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc differ diff --git a/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc b/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc index 90f775d27997dc8659edde9eb763d0f8b4007ace..e5e2164c9413b0e07882b849fa77a3609be7a103 100644 Binary files a/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc and b/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc differ diff --git a/src/generate_batch.py b/src/generate_batch.py index 8bf580e49427527bfd1c2ff533de45ee91e3872e..294047cb10495a66f124f3b8ce0869fc5e22447c 100644 --- a/src/generate_batch.py +++ b/src/generate_batch.py @@ -48,7 +48,7 @@ def generate_blink_seq_randomly(num_frames): break return ratio -def get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path): +def get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path, still=False): syncnet_mel_step_size = 16 fps = 25 @@ -95,7 +95,12 @@ def get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path): ref_coeff[:, :64] = refeyeblink_coeff[:num_frames, :64] indiv_mels = torch.FloatTensor(indiv_mels).unsqueeze(1).unsqueeze(0) # bs T 1 80 16 - ratio = torch.FloatTensor(ratio).unsqueeze(0) # bs T + + if still: + ratio = torch.FloatTensor(ratio).unsqueeze(0).fill_(0.) # bs T + else: + ratio = torch.FloatTensor(ratio).unsqueeze(0) + # bs T ref_coeff = torch.FloatTensor(ref_coeff).unsqueeze(0) # bs 1 70 indiv_mels = indiv_mels.to(device) diff --git a/src/generate_facerender_batch.py b/src/generate_facerender_batch.py index 53b8cf7ada396907a77702c264616c3a8cdc05ab..9ec7a169706e9e4697f0f847d4e3d46101bd55d9 100644 --- a/src/generate_facerender_batch.py +++ b/src/generate_facerender_batch.py @@ -7,7 +7,7 @@ import scipy.io as scio def get_facerender_data(coeff_path, pic_path, first_coeff_path, audio_path, batch_size, input_yaw_list=None, input_pitch_list=None, input_roll_list=None, - expression_scale=1.0, still_mode = False): + expression_scale=1.0, still_mode = False, preprocess='crop'): semantic_radius = 13 video_name = os.path.splitext(os.path.split(coeff_path)[-1])[0] @@ -25,7 +25,12 @@ def get_facerender_data(coeff_path, pic_path, first_coeff_path, audio_path, data['source_image'] = source_image_ts source_semantics_dict = scio.loadmat(first_coeff_path) - source_semantics = source_semantics_dict['coeff_3dmm'][:1,:70] #1 70 + + if preprocess.lower() != 'full': + source_semantics = source_semantics_dict['coeff_3dmm'][:1,:70] #1 70 + else: + source_semantics = source_semantics_dict['coeff_3dmm'][:1,:73] #1 70 + source_semantics_new = transform_semantic_1(source_semantics, semantic_radius) source_semantics_ts = torch.FloatTensor(source_semantics_new).unsqueeze(0) source_semantics_ts = source_semantics_ts.repeat(batch_size, 1, 1) @@ -36,6 +41,9 @@ def get_facerender_data(coeff_path, pic_path, first_coeff_path, audio_path, generated_3dmm = generated_dict['coeff_3dmm'] generated_3dmm[:, :64] = generated_3dmm[:, :64] * expression_scale + if preprocess.lower() == 'full': + generated_3dmm = np.concatenate([generated_3dmm, np.repeat(source_semantics[:,70:], generated_3dmm.shape[0], axis=0)], axis=1) + if still_mode: generated_3dmm[:, 64:] = np.repeat(source_semantics[:, 64:], generated_3dmm.shape[0], axis=0) @@ -82,7 +90,7 @@ def transform_semantic_1(semantic, semantic_radius): def transform_semantic_target(coeff_3dmm, frame_index, semantic_radius): num_frames = coeff_3dmm.shape[0] - seq = list(range(frame_index- semantic_radius, frame_index+ semantic_radius+1)) + seq = list(range(frame_index- semantic_radius, frame_index + semantic_radius+1)) index = [ min(max(item, 0), num_frames-1) for item in seq ] coeff_3dmm_g = coeff_3dmm[index, :] return coeff_3dmm_g.transpose(1,0) @@ -124,8 +132,3 @@ def gen_camera_pose(camera_degree_list, frame_num, batch_size): new_degree_np = np.array(new_degree_list).reshape(batch_size, -1) return new_degree_np - - - - - diff --git a/src/gradio_demo.py b/src/gradio_demo.py index d2310d7323b05f8ef08eccbeeb64c329f2072d01..b67fc2fd98e2faa0064e0db74beb0dc682483b69 100644 --- a/src/gradio_demo.py +++ b/src/gradio_demo.py @@ -1,106 +1,137 @@ -import torch, uuid -import os, sys, shutil -from src.utils.preprocess import CropAndExtract -from src.test_audio2coeff import Audio2Coeff -from src.facerender.animate import AnimateFromCoeff -from src.generate_batch import get_data -from src.generate_facerender_batch import get_facerender_data - -from pydub import AudioSegment - -def mp3_to_wav(mp3_filename,wav_filename,frame_rate): - mp3_file = AudioSegment.from_file(file=mp3_filename) - mp3_file.set_frame_rate(frame_rate).export(wav_filename,format="wav") - - -class SadTalker(): - - def __init__(self, checkpoint_path='checkpoints', config_path='src/config'): - - if torch.cuda.is_available() : - device = "cuda" - else: - device = "cpu" - - os.environ['TORCH_HOME']= checkpoint_path - - path_of_lm_croper = os.path.join( checkpoint_path, 'shape_predictor_68_face_landmarks.dat') - path_of_net_recon_model = os.path.join( checkpoint_path, 'epoch_20.pth') - dir_of_BFM_fitting = os.path.join( checkpoint_path, 'BFM_Fitting') - wav2lip_checkpoint = os.path.join( checkpoint_path, 'wav2lip.pth') - - audio2pose_checkpoint = os.path.join( checkpoint_path, 'auido2pose_00140-model.pth') - audio2pose_yaml_path = os.path.join( config_path, 'auido2pose.yaml') - - audio2exp_checkpoint = os.path.join( checkpoint_path, 'auido2exp_00300-model.pth') - audio2exp_yaml_path = os.path.join( config_path, 'auido2exp.yaml') - - free_view_checkpoint = os.path.join( checkpoint_path, 'facevid2vid_00189-model.pth.tar') - mapping_checkpoint = os.path.join( checkpoint_path, 'mapping_00229-model.pth.tar') - facerender_yaml_path = os.path.join( config_path, 'facerender.yaml') - - #init model - print(path_of_lm_croper) - self.preprocess_model = CropAndExtract(path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting, device) - - print(audio2pose_checkpoint) - self.audio_to_coeff = Audio2Coeff(audio2pose_checkpoint, audio2pose_yaml_path, - audio2exp_checkpoint, audio2exp_yaml_path, wav2lip_checkpoint, device) - print(free_view_checkpoint) - self.animate_from_coeff = AnimateFromCoeff(free_view_checkpoint, mapping_checkpoint, - facerender_yaml_path, device) - self.device = device - - def test(self, source_image, driven_audio, still_mode, use_enhancer, result_dir='./results/'): - - time_tag = str(uuid.uuid4()) - save_dir = os.path.join(result_dir, time_tag) - os.makedirs(save_dir, exist_ok=True) - - input_dir = os.path.join(save_dir, 'input') - os.makedirs(input_dir, exist_ok=True) - - print(source_image) - pic_path = os.path.join(input_dir, os.path.basename(source_image)) - shutil.move(source_image, input_dir) - - if os.path.isfile(driven_audio): - audio_path = os.path.join(input_dir, os.path.basename(driven_audio)) - - #### mp3 to wav - if '.mp3' in audio_path: - mp3_to_wav(driven_audio, audio_path.replace('.mp3', '.wav'), 16000) - audio_path = audio_path.replace('.mp3', '.wav') - else: - shutil.move(driven_audio, input_dir) - else: - raise AttributeError("error audio") - - - os.makedirs(save_dir, exist_ok=True) - pose_style = 0 - #crop image and extract 3dmm from image - first_frame_dir = os.path.join(save_dir, 'first_frame_dir') - os.makedirs(first_frame_dir, exist_ok=True) - first_coeff_path, crop_pic_path, crop_info = self.preprocess_model.generate(pic_path, first_frame_dir) - - if first_coeff_path is None: - raise AttributeError("No face is detected") - - #audio2ceoff - batch = get_data(first_coeff_path, audio_path, self.device, None) # longer audio? - coeff_path = self.audio_to_coeff.generate(batch, save_dir, pose_style) - #coeff2video - batch_size = 2 - data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode) - return_path = self.animate_from_coeff.generate(data, save_dir, pic_path, crop_info, enhancer='gfpgan' if use_enhancer else None) - video_name = data['video_name'] - print(f'The generated video is named {video_name} in {save_dir}') - - torch.cuda.empty_cache() - torch.cuda.synchronize() - import gc; gc.collect() - - return return_path - +import torch, uuid +import os, sys, shutil +from src.utils.preprocess import CropAndExtract +from src.test_audio2coeff import Audio2Coeff +from src.facerender.animate import AnimateFromCoeff +from src.generate_batch import get_data +from src.generate_facerender_batch import get_facerender_data + +from pydub import AudioSegment + +def mp3_to_wav(mp3_filename,wav_filename,frame_rate): + mp3_file = AudioSegment.from_file(file=mp3_filename) + mp3_file.set_frame_rate(frame_rate).export(wav_filename,format="wav") + + +class SadTalker(): + + def __init__(self, checkpoint_path='checkpoints', config_path='src/config', lazy_load=False): + + if torch.cuda.is_available() : + device = "cuda" + else: + device = "cpu" + + self.device = device + + os.environ['TORCH_HOME']= checkpoint_path + + self.checkpoint_path = checkpoint_path + self.config_path = config_path + + self.path_of_lm_croper = os.path.join( checkpoint_path, 'shape_predictor_68_face_landmarks.dat') + self.path_of_net_recon_model = os.path.join( checkpoint_path, 'epoch_20.pth') + self.dir_of_BFM_fitting = os.path.join( checkpoint_path, 'BFM_Fitting') + self.wav2lip_checkpoint = os.path.join( checkpoint_path, 'wav2lip.pth') + + self.audio2pose_checkpoint = os.path.join( checkpoint_path, 'auido2pose_00140-model.pth') + self.audio2pose_yaml_path = os.path.join( config_path, 'auido2pose.yaml') + + self.audio2exp_checkpoint = os.path.join( checkpoint_path, 'auido2exp_00300-model.pth') + self.audio2exp_yaml_path = os.path.join( config_path, 'auido2exp.yaml') + + self.free_view_checkpoint = os.path.join( checkpoint_path, 'facevid2vid_00189-model.pth.tar') + + + + self.lazy_load = lazy_load + + if not self.lazy_load: + #init model + print(self.path_of_lm_croper) + self.preprocess_model = CropAndExtract(self.path_of_lm_croper, self.path_of_net_recon_model, self.dir_of_BFM_fitting, self.device) + + print(self.audio2pose_checkpoint) + self.audio_to_coeff = Audio2Coeff(self.audio2pose_checkpoint, self.audio2pose_yaml_path, + self.audio2exp_checkpoint, self.audio2exp_yaml_path, self.wav2lip_checkpoint, self.device) + + def test(self, source_image, driven_audio, preprocess='crop', still_mode=False, use_enhancer=False, result_dir='./results/'): + + ### crop: only model, + + if self.lazy_load: + #init model + print(self.path_of_lm_croper) + self.preprocess_model = CropAndExtract(self.path_of_lm_croper, self.path_of_net_recon_model, self.dir_of_BFM_fitting, self.device) + + print(self.audio2pose_checkpoint) + self.audio_to_coeff = Audio2Coeff(self.audio2pose_checkpoint, self.audio2pose_yaml_path, + self.audio2exp_checkpoint, self.audio2exp_yaml_path, self.wav2lip_checkpoint, self.device) + + if preprocess == 'full': + self.mapping_checkpoint = os.path.join(self.checkpoint_path, 'mapping_00109-model.pth.tar') + self.facerender_yaml_path = os.path.join(self.config_path, 'facerender_still.yaml') + else: + self.mapping_checkpoint = os.path.join(self.checkpoint_path, 'mapping_00229-model.pth.tar') + self.facerender_yaml_path = os.path.join(self.config_path, 'facerender.yaml') + + print(self.free_view_checkpoint) + self.animate_from_coeff = AnimateFromCoeff(self.free_view_checkpoint, self.mapping_checkpoint, + self.facerender_yaml_path, self.device) + + time_tag = str(uuid.uuid4()) + save_dir = os.path.join(result_dir, time_tag) + os.makedirs(save_dir, exist_ok=True) + + input_dir = os.path.join(save_dir, 'input') + os.makedirs(input_dir, exist_ok=True) + + print(source_image) + pic_path = os.path.join(input_dir, os.path.basename(source_image)) + shutil.move(source_image, input_dir) + + if os.path.isfile(driven_audio): + audio_path = os.path.join(input_dir, os.path.basename(driven_audio)) + + #### mp3 to wav + if '.mp3' in audio_path: + mp3_to_wav(driven_audio, audio_path.replace('.mp3', '.wav'), 16000) + audio_path = audio_path.replace('.mp3', '.wav') + else: + shutil.move(driven_audio, input_dir) + else: + raise AttributeError("error audio") + + + os.makedirs(save_dir, exist_ok=True) + pose_style = 0 + #crop image and extract 3dmm from image + first_frame_dir = os.path.join(save_dir, 'first_frame_dir') + os.makedirs(first_frame_dir, exist_ok=True) + first_coeff_path, crop_pic_path, crop_info = self.preprocess_model.generate(pic_path, first_frame_dir,preprocess) + + if first_coeff_path is None: + raise AttributeError("No face is detected") + + #audio2ceoff + batch = get_data(first_coeff_path, audio_path, self.device, ref_eyeblink_coeff_path=None, still=still_mode) # longer audio? + coeff_path = self.audio_to_coeff.generate(batch, save_dir, pose_style) + #coeff2video + batch_size = 2 + data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode, preprocess=preprocess) + return_path = self.animate_from_coeff.generate(data, save_dir, pic_path, crop_info, enhancer='gfpgan' if use_enhancer else None, preprocess=preprocess) + video_name = data['video_name'] + print(f'The generated video is named {video_name} in {save_dir}') + + if self.lazy_load: + del self.preprocess_model + del self.audio_to_coeff + del self.animate_from_coeff + + torch.cuda.empty_cache() + torch.cuda.synchronize() + import gc; gc.collect() + + return return_path + \ No newline at end of file diff --git a/src/utils/__pycache__/audio.cpython-38.pyc b/src/utils/__pycache__/audio.cpython-38.pyc index 9f4fe6227f50165dfd5ef7458765d3c806e571c9..78b5f23ab07bacf1e8d9b9bc74de0bb3fadecb5a 100644 Binary files a/src/utils/__pycache__/audio.cpython-38.pyc and b/src/utils/__pycache__/audio.cpython-38.pyc differ diff --git a/src/utils/__pycache__/croper.cpython-38.pyc b/src/utils/__pycache__/croper.cpython-38.pyc index e9eacf099f5dd124f8d00eb7fa9076dddde74df7..e563c8ef7068b98f5f81756269cece9dbfee8ca1 100644 Binary files a/src/utils/__pycache__/croper.cpython-38.pyc and b/src/utils/__pycache__/croper.cpython-38.pyc differ diff --git a/src/utils/__pycache__/face_enhancer.cpython-38.pyc b/src/utils/__pycache__/face_enhancer.cpython-38.pyc index a46dc8b648c6ed407e9486a5452677166f6ed6ea..ed3df35701c4ead91cc1e83f83a0c9307272bcf6 100644 Binary files a/src/utils/__pycache__/face_enhancer.cpython-38.pyc and b/src/utils/__pycache__/face_enhancer.cpython-38.pyc differ diff --git a/src/utils/__pycache__/hparams.cpython-38.pyc b/src/utils/__pycache__/hparams.cpython-38.pyc index bab0f32bce885e5d2fb8061dd81b4767fb987df8..e393d49ddeefa6c304fbe547421128c8671340be 100644 Binary files a/src/utils/__pycache__/hparams.cpython-38.pyc and b/src/utils/__pycache__/hparams.cpython-38.pyc differ diff --git a/src/utils/__pycache__/preprocess.cpython-38.pyc b/src/utils/__pycache__/preprocess.cpython-38.pyc index 6a7c7064a0cf9edb8e4a8af6a41ba3020daecd88..52ca52954d43b2f2c6c6945e084008768e9aa20b 100644 Binary files a/src/utils/__pycache__/preprocess.cpython-38.pyc and b/src/utils/__pycache__/preprocess.cpython-38.pyc differ diff --git a/src/utils/croper.py b/src/utils/croper.py index e68d280ee4bd83db2089c226af5d4be714fcca9d..1c9b68459ac6ab7d39ea1bc692e127b064eeeca8 100644 --- a/src/utils/croper.py +++ b/src/utils/croper.py @@ -161,7 +161,7 @@ class Croper: # img_np_list[_i] = _inp # return img_np_list - def crop(self, img_np_list, xsize=512): # first frame for all video + def crop(self, img_np_list, still=False, xsize=512): # first frame for all video img_np = img_np_list[0] lm = self.get_landmark(img_np) if lm is None: @@ -174,7 +174,8 @@ class Croper: _inp = img_np_list[_i] _inp = _inp[cly:cry, clx:crx] # cv2.imwrite('test1.jpg', _inp) - _inp = _inp[ly:ry, lx:rx] + if not still: + _inp = _inp[ly:ry, lx:rx] # cv2.imwrite('test2.jpg', _inp) img_np_list[_i] = _inp return img_np_list, crop, quad @@ -292,4 +293,4 @@ if __name__ == '__main__': device_ids = opt.device_ids.split(",") device_ids = cycle(device_ids) for data in tqdm(pool.imap_unordered(run, zip(filenames, args_list, device_ids))): - None + None \ No newline at end of file diff --git a/src/utils/face_enhancer.py b/src/utils/face_enhancer.py index ecf3587ced450273f05388e6cb9d8d6e5b11a51f..e089d140907498ad6c75c253f9719b992b5aff4b 100644 --- a/src/utils/face_enhancer.py +++ b/src/utils/face_enhancer.py @@ -1,11 +1,18 @@ import os -from basicsr.utils import imwrite +import torch from gfpgan import GFPGANer from tqdm import tqdm -def enhancer(images, method='gfpgan'): +from src.utils.videoio import load_video_to_cv2 + + + +def enhancer(images, method='gfpgan', bg_upsampler='realesrgan'): + print('face enhancer....') + if os.path.isfile(images): # handle video to images + images = load_video_to_cv2(images) # ------------------------ set up GFPGAN restorer ------------------------ if method == 'gfpgan': @@ -26,8 +33,31 @@ def enhancer(images, method='gfpgan'): else: raise ValueError(f'Wrong model version {method}.') + + # ------------------------ set up background upsampler ------------------------ + if bg_upsampler == 'realesrgan': + if not torch.cuda.is_available(): # CPU + import warnings + warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. ' + 'If you really want to use it, please modify the corresponding codes.') + bg_upsampler = None + else: + from basicsr.archs.rrdbnet_arch import RRDBNet + from realesrgan import RealESRGANer + model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2) + bg_upsampler = RealESRGANer( + scale=2, + model_path='/apdcephfs/private_shadowcun/SadTalker/gfpgan/weights/RealESRGAN_x2plus.pth', + model=model, + tile=400, + tile_pad=10, + pre_pad=0, + half=True) # need to set False in CPU mode + else: + bg_upsampler = None + # determine model paths - model_path = os.path.join('experiments/pretrained_models', model_name + '.pth') + model_path = os.path.join('gfpgan/weights', model_name + '.pth') if not os.path.isfile(model_path): model_path = os.path.join('checkpoints', model_name + '.pth') @@ -41,19 +71,19 @@ def enhancer(images, method='gfpgan'): upscale=2, arch=arch, channel_multiplier=channel_multiplier, - bg_upsampler=None) + bg_upsampler=bg_upsampler) # ------------------------ restore ------------------------ restored_img = [] for idx in tqdm(range(len(images)), 'Face Enhancer:'): # restore faces and background if necessary - cropped_faces, restored_faces, _ = restorer.enhance( + cropped_faces, restored_faces, r_img = restorer.enhance( images[idx], - has_aligned=True, + has_aligned=False, only_center_face=False, paste_back=True) - restored_img += restored_faces + restored_img += [r_img] return restored_img \ No newline at end of file diff --git a/src/utils/paste_pic.py b/src/utils/paste_pic.py index 508aa88d5f12aa81ab9a40b71913f6fabfc6332f..a05a55caeea190d2af32f2341e3a96d1fc417b09 100644 --- a/src/utils/paste_pic.py +++ b/src/utils/paste_pic.py @@ -1,12 +1,13 @@ import cv2, os import numpy as np from tqdm import tqdm -import uuid +import uuid + +from src.utils.videoio import save_video_with_watermark def paste_pic(video_path, pic_path, crop_info, new_audio_path, full_video_path): full_img = cv2.imread(pic_path) - print(full_img.dtype) frame_h = full_img.shape[0] frame_w = full_img.shape[1] @@ -28,20 +29,22 @@ def paste_pic(video_path, pic_path, crop_info, new_audio_path, full_video_path): clx, cly, crx, cry = crop_info[1] lx, ly, rx, ry = crop_info[2] lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry) - oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx + # oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx + # oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx + oy1, oy2, ox1, ox2 = cly, cry, clx, crx + tmp_path = str(uuid.uuid4())+'.mp4' out_tmp = cv2.VideoWriter(tmp_path, cv2.VideoWriter_fourcc(*'MP4V'), fps, (frame_w, frame_h)) for crop_frame in tqdm(crop_frames, 'seamlessClone:'): - p = cv2.resize(crop_frame.astype(np.uint8), (r_w, r_h)) + p = cv2.resize(crop_frame.astype(np.uint8), (crx-clx, cry - cly)) mask = 255*np.ones(p.shape, p.dtype) location = ((ox1+ox2) // 2, (oy1+oy2) // 2) gen_img = cv2.seamlessClone(p, full_img, mask, location, cv2.NORMAL_CLONE) - - #full_img[oy1:oy2, ox1:ox2] = p out_tmp.write(gen_img) + out_tmp.release() - cmd = r'ffmpeg -y -i "%s" -i "%s" "%s"' % (tmp_path, new_audio_path, full_video_path) - os.system(cmd) + + save_video_with_watermark(tmp_path, new_audio_path, full_video_path) os.remove(tmp_path) diff --git a/src/utils/preprocess.py b/src/utils/preprocess.py index d12cada70ccefcc905bc6727a304a6e974c4f25c..454e26b2fd1a3b662399700c7805c02a63384301 100644 --- a/src/utils/preprocess.py +++ b/src/utils/preprocess.py @@ -63,7 +63,7 @@ class CropAndExtract(): #load input if not os.path.isfile(input_path): raise ValueError('input_path must be a valid path to video/image file') - elif input_path.split('.')[1] in ['jpg', 'png', 'jpeg']: + elif input_path.split('.')[-1] in ['jpg', 'png', 'jpeg']: # loader for first frame full_frames = [cv2.imread(input_path)] fps = 25 @@ -78,8 +78,10 @@ class CropAndExtract(): video_stream.release() break full_frames.append(frame) - x_full_frames = [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in full_frames] + x_full_frames= [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in full_frames] + + #### crop images as the if crop_or_resize.lower() == 'crop': # default crop x_full_frames, crop, quad = self.croper.crop(x_full_frames, xsize=pic_size) clx, cly, crx, cry = crop @@ -87,9 +89,16 @@ class CropAndExtract(): lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry) oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx crop_info = ((ox2 - ox1, oy2 - oy1), crop, quad) - else: + elif crop_or_resize.lower() == 'full': + x_full_frames, crop, quad = self.croper.crop(x_full_frames, still=True, xsize=pic_size) + clx, cly, crx, cry = crop + lx, ly, rx, ry = quad + lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry) + oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx + crop_info = ((ox2 - ox1, oy2 - oy1), crop, quad) + else: # resize mode oy1, oy2, ox1, ox2 = 0, x_full_frames[0].shape[0], 0, x_full_frames[0].shape[1] - crop_info = ((ox2 - ox1, oy2 - oy1)) + crop_info = ((ox2 - ox1, oy2 - oy1), None, None) frames_pil = [Image.fromarray(cv2.resize(frame,(pic_size, pic_size))) for frame in x_full_frames] if len(frames_pil) == 0: @@ -128,7 +137,7 @@ class CropAndExtract(): trans_params = np.array([float(item) for item in np.hsplit(trans_params, 5)]).astype(np.float32) im_t = torch.tensor(np.array(im1)/255., dtype=torch.float32).permute(2, 0, 1).to(self.device).unsqueeze(0) - + with torch.no_grad(): full_coeff = self.net_recon(im_t) coeffs = split_coeff(full_coeff) @@ -148,4 +157,4 @@ class CropAndExtract(): savemat(coeff_path, {'coeff_3dmm': semantic_npy, 'full_3dmm': np.array(full_coeffs)[0]}) - return coeff_path, png_path, crop_info \ No newline at end of file + return coeff_path, png_path, crop_info diff --git a/src/utils/videoio.py b/src/utils/videoio.py new file mode 100644 index 0000000000000000000000000000000000000000..53998919fc7fd24c7bff85daf248de435f329622 --- /dev/null +++ b/src/utils/videoio.py @@ -0,0 +1,31 @@ +import shutil +import uuid + +import os + +import cv2 + +def load_video_to_cv2(input_path): + video_stream = cv2.VideoCapture(input_path) + fps = video_stream.get(cv2.CAP_PROP_FPS) + full_frames = [] + while 1: + still_reading, frame = video_stream.read() + if not still_reading: + video_stream.release() + break + full_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + return full_frames + +def save_video_with_watermark(video, audio, save_path, watermark='docs/sadtalker_logo.png'): + temp_file = str(uuid.uuid4())+'.mp4' + cmd = r'ffmpeg -y -i "%s" -i "%s" -vcodec copy "%s"' % (video, audio, temp_file) + os.system(cmd) + + if not watermark: + shutil.move(temp_file, save_path) + else: + # watermark + cmd = r'ffmpeg -y -hide_banner -i "%s" -i "%s" -filter_complex "[1]scale=100:-1[wm];[0][wm]overlay=(main_w-overlay_w)-10:10" "%s"' % (temp_file, watermark, save_path) + os.system(cmd) + os.remove(temp_file) \ No newline at end of file