|
import os, sys |
|
from pathlib import Path |
|
import tempfile |
|
import gradio as gr |
|
from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call |
|
from modules.shared import opts, OptionInfo |
|
from modules import shared, paths, script_callbacks |
|
import launch |
|
import glob |
|
from huggingface_hub import snapshot_download |
|
|
|
|
|
|
|
def check_all_files_safetensor(current_dir): |
|
kv = { |
|
"SadTalker_V0.0.2_256.safetensors": "sadtalker-256", |
|
"SadTalker_V0.0.2_512.safetensors": "sadtalker-512", |
|
"mapping_00109-model.pth.tar" : "mapping-109" , |
|
"mapping_00229-model.pth.tar" : "mapping-229" , |
|
} |
|
|
|
if not os.path.isdir(current_dir): |
|
return False |
|
|
|
dirs = os.listdir(current_dir) |
|
|
|
for f in dirs: |
|
if f in kv.keys(): |
|
del kv[f] |
|
|
|
return len(kv.keys()) == 0 |
|
|
|
def check_all_files(current_dir): |
|
kv = { |
|
"auido2exp_00300-model.pth": "audio2exp", |
|
"auido2pose_00140-model.pth": "audio2pose", |
|
"epoch_20.pth": "face_recon", |
|
"facevid2vid_00189-model.pth.tar": "face-render", |
|
"mapping_00109-model.pth.tar" : "mapping-109" , |
|
"mapping_00229-model.pth.tar" : "mapping-229" , |
|
"wav2lip.pth": "wav2lip", |
|
"shape_predictor_68_face_landmarks.dat": "dlib", |
|
} |
|
|
|
if not os.path.isdir(current_dir): |
|
return False |
|
|
|
dirs = os.listdir(current_dir) |
|
|
|
for f in dirs: |
|
if f in kv.keys(): |
|
del kv[f] |
|
|
|
return len(kv.keys()) == 0 |
|
|
|
|
|
|
|
def download_model(local_dir='./checkpoints'): |
|
REPO_ID = 'vinthony/SadTalker' |
|
snapshot_download(repo_id=REPO_ID, local_dir=local_dir, local_dir_use_symlinks=False) |
|
|
|
def get_source_image(image): |
|
return image |
|
|
|
def get_img_from_txt2img(x): |
|
talker_path = Path(paths.script_path) / "outputs" |
|
imgs_from_txt_dir = str(talker_path / "txt2img-images/") |
|
imgs = glob.glob(imgs_from_txt_dir+'/*/*.png') |
|
imgs.sort(key=lambda x:os.path.getmtime(os.path.join(imgs_from_txt_dir, x))) |
|
img_from_txt_path = os.path.join(imgs_from_txt_dir, imgs[-1]) |
|
return img_from_txt_path, img_from_txt_path |
|
|
|
def get_img_from_img2img(x): |
|
talker_path = Path(paths.script_path) / "outputs" |
|
imgs_from_img_dir = str(talker_path / "img2img-images/") |
|
imgs = glob.glob(imgs_from_img_dir+'/*/*.png') |
|
imgs.sort(key=lambda x:os.path.getmtime(os.path.join(imgs_from_img_dir, x))) |
|
img_from_img_path = os.path.join(imgs_from_img_dir, imgs[-1]) |
|
return img_from_img_path, img_from_img_path |
|
|
|
def get_default_checkpoint_path(): |
|
|
|
checkpoint_path = Path(paths.script_path) / "models"/ "SadTalker" |
|
extension_checkpoint_path = Path(paths.script_path) / "extensions"/ "SadTalker" / "checkpoints" |
|
|
|
if check_all_files_safetensor(checkpoint_path): |
|
|
|
return checkpoint_path |
|
|
|
if check_all_files_safetensor(extension_checkpoint_path): |
|
|
|
return extension_checkpoint_path |
|
|
|
if check_all_files(checkpoint_path): |
|
|
|
return checkpoint_path |
|
|
|
if check_all_files(extension_checkpoint_path): |
|
|
|
return extension_checkpoint_path |
|
|
|
return None |
|
|
|
|
|
|
|
def install(): |
|
|
|
kv = { |
|
"face_alignment": "face-alignment==1.3.5", |
|
"imageio": "imageio==2.19.3", |
|
"imageio_ffmpeg": "imageio-ffmpeg==0.4.7", |
|
"librosa":"librosa==0.8.0", |
|
"pydub":"pydub==0.25.1", |
|
"scipy":"scipy==1.8.1", |
|
"tqdm": "tqdm", |
|
"yacs":"yacs==0.1.8", |
|
"yaml": "pyyaml", |
|
"av":"av", |
|
"gfpgan": "gfpgan", |
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
for k,v in kv.items(): |
|
if not launch.is_installed(k): |
|
print(k, launch.is_installed(k)) |
|
launch.run_pip("install "+ v, "requirements for SadTalker") |
|
|
|
if os.getenv('SADTALKER_CHECKPOINTS'): |
|
print('load Sadtalker Checkpoints from '+ os.getenv('SADTALKER_CHECKPOINTS')) |
|
|
|
elif get_default_checkpoint_path() is not None: |
|
os.environ['SADTALKER_CHECKPOINTS'] = str(get_default_checkpoint_path()) |
|
else: |
|
|
|
print( |
|
"""" |
|
SadTalker will not support download all the files from hugging face, which will take a long time. |
|
|
|
please manually set the SADTALKER_CHECKPOINTS in `webui_user.bat`(windows) or `webui_user.sh`(linux) |
|
""" |
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def on_ui_tabs(): |
|
install() |
|
|
|
sys.path.extend([paths.script_path+'/extensions/SadTalker']) |
|
|
|
repo_dir = paths.script_path+'/extensions/SadTalker/' |
|
|
|
result_dir = opts.sadtalker_result_dir |
|
os.makedirs(result_dir, exist_ok=True) |
|
|
|
from app_sadtalker import sadtalker_demo |
|
|
|
if os.getenv('SADTALKER_CHECKPOINTS'): |
|
checkpoint_path = os.getenv('SADTALKER_CHECKPOINTS') |
|
else: |
|
checkpoint_path = repo_dir+'checkpoints/' |
|
|
|
audio_to_video = sadtalker_demo(checkpoint_path=checkpoint_path, config_path=repo_dir+'src/config', warpfn = wrap_queued_call) |
|
|
|
return [(audio_to_video, "SadTalker", "extension")] |
|
|
|
def on_ui_settings(): |
|
talker_path = Path(paths.script_path) / "outputs" |
|
section = ('extension', "SadTalker") |
|
opts.add_option("sadtalker_result_dir", OptionInfo(str(talker_path / "SadTalker/"), "Path to save results of sadtalker", section=section)) |
|
|
|
script_callbacks.on_ui_settings(on_ui_settings) |
|
script_callbacks.on_ui_tabs(on_ui_tabs) |
|
|