init

src/utils/paste_pic.py

+import cv2, os
+import numpy as np
+from tqdm import tqdm
+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, extended_crop=False):
+
+    if not os.path.isfile(pic_path):
+        raise ValueError('pic_path must be a valid path to video/image file')
+    elif pic_path.split('.')[-1] in ['jpg', 'png', 'jpeg']:
+        # loader for first frame
+        full_img = cv2.imread(pic_path)
+    else:
+        # loader for videos
+        video_stream = cv2.VideoCapture(pic_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 
+            break 
+        full_img = frame
+    frame_h = full_img.shape[0]
+    frame_w = full_img.shape[1]
+
+    video_stream = cv2.VideoCapture(video_path)
+    fps = video_stream.get(cv2.CAP_PROP_FPS)
+    crop_frames = []
+    while 1:
+        still_reading, frame = video_stream.read()
+        if not still_reading:
+            video_stream.release()
+            break
+        crop_frames.append(frame)
+    
+    if len(crop_info) != 3:
+        print("you didn't crop the image")
+        return
+    else:
+        r_w, r_h = crop_info[0]
+        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
+
+        if extended_crop:
+            oy1, oy2, ox1, ox2 = cly, cry, clx, crx
+        else:
+            oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx
+
+    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), (ox2-ox1, oy2 - oy1)) 
+
+        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)
+        out_tmp.write(gen_img)
+
+    out_tmp.release()
+
+    save_video_with_watermark(tmp_path, new_audio_path, full_video_path, watermark=False)
+    os.remove(tmp_path)

src/utils/preprocess.py

+import numpy as np
+import cv2, os, sys, torch
+from tqdm import tqdm
+from PIL import Image 
+
+# 3dmm extraction
+import safetensors
+import safetensors.torch 
+from src.face3d.util.preprocess import align_img
+from src.face3d.util.load_mats import load_lm3d
+from src.face3d.models import networks
+
+from scipy.io import loadmat, savemat
+from src.utils.croper import Preprocesser
+
+
+import warnings
+
+from src.utils.safetensor_helper import load_x_from_safetensor 
+warnings.filterwarnings("ignore")
+
+def split_coeff(coeffs):
+        """
+        Return:
+            coeffs_dict     -- a dict of torch.tensors
+
+        Parameters:
+            coeffs          -- torch.tensor, size (B, 256)
+        """
+        id_coeffs = coeffs[:, :80]
+        exp_coeffs = coeffs[:, 80: 144]
+        tex_coeffs = coeffs[:, 144: 224]
+        angles = coeffs[:, 224: 227]
+        gammas = coeffs[:, 227: 254]
+        translations = coeffs[:, 254:]
+        return {
+            'id': id_coeffs,
+            'exp': exp_coeffs,
+            'tex': tex_coeffs,
+            'angle': angles,
+            'gamma': gammas,
+            'trans': translations
+        }
+
+
+class CropAndExtract():
+    def __init__(self, sadtalker_path, device):
+
+        self.propress = Preprocesser(device)
+        self.net_recon = networks.define_net_recon(net_recon='resnet50', use_last_fc=False, init_path='').to(device)
+        
+        if sadtalker_path['use_safetensor']:
+            checkpoint = safetensors.torch.load_file(sadtalker_path['checkpoint'])    
+            self.net_recon.load_state_dict(load_x_from_safetensor(checkpoint, 'face_3drecon'))
+        else:
+            checkpoint = torch.load(sadtalker_path['path_of_net_recon_model'], map_location=torch.device(device))    
+            self.net_recon.load_state_dict(checkpoint['net_recon'])
+
+        self.net_recon.eval()
+        self.lm3d_std = load_lm3d(sadtalker_path['dir_of_BFM_fitting'])
+        self.device = device
+    
+    def generate(self, input_path, save_dir, crop_or_resize='crop', source_image_flag=False, pic_size=256):
+
+        pic_name = os.path.splitext(os.path.split(input_path)[-1])[0]  
+
+        landmarks_path =  os.path.join(save_dir, pic_name+'_landmarks.txt') 
+        coeff_path =  os.path.join(save_dir, pic_name+'.mat')  
+        png_path =  os.path.join(save_dir, pic_name+'.png')  
+
+        #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']:
+            # loader for first frame
+            full_frames = [cv2.imread(input_path)]
+            fps = 25
+        else:
+            # loader for videos
+            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(frame) 
+                if source_image_flag:
+                    break
+
+        x_full_frames= [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  for frame in full_frames] 
+
+        #### crop images as the 
+        if 'crop' in crop_or_resize.lower(): # default crop
+            x_full_frames, crop, quad = self.propress.crop(x_full_frames, still=True if 'ext' in crop_or_resize.lower() else False, xsize=512)
+            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)
+        elif 'full' in crop_or_resize.lower():
+            x_full_frames, crop, quad = self.propress.crop(x_full_frames, still=True if 'ext' in crop_or_resize.lower() else False, xsize=512)
+            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), None, None)
+
+        frames_pil = [Image.fromarray(cv2.resize(frame,(pic_size, pic_size))) for frame in x_full_frames]
+        if len(frames_pil) == 0:
+            print('No face is detected in the input file')
+            return None, None
+
+        # save crop info
+        for frame in frames_pil:
+            cv2.imwrite(png_path, cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR))
+
+        # 2. get the landmark according to the detected face. 
+        if not os.path.isfile(landmarks_path): 
+            lm = self.propress.predictor.extract_keypoint(frames_pil, landmarks_path)
+        else:
+            print(' Using saved landmarks.')
+            lm = np.loadtxt(landmarks_path).astype(np.float32)
+            lm = lm.reshape([len(x_full_frames), -1, 2])
+
+        if not os.path.isfile(coeff_path):
+            # load 3dmm paramter generator from Deep3DFaceRecon_pytorch 
+            video_coeffs, full_coeffs = [],  []
+            for idx in tqdm(range(len(frames_pil)), desc='3DMM Extraction In Video:'):
+                frame = frames_pil[idx]
+                W,H = frame.size
+                lm1 = lm[idx].reshape([-1, 2])
+            
+                if np.mean(lm1) == -1:
+                    lm1 = (self.lm3d_std[:, :2]+1)/2.
+                    lm1 = np.concatenate(
+                        [lm1[:, :1]*W, lm1[:, 1:2]*H], 1
+                    )
+                else:
+                    lm1[:, -1] = H - 1 - lm1[:, -1]
+
+                trans_params, im1, lm1, _ = align_img(frame, lm1, self.lm3d_std)
+ 
+                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)
+
+                pred_coeff = {key:coeffs[key].cpu().numpy() for key in coeffs}
+ 
+                pred_coeff = np.concatenate([
+                    pred_coeff['exp'], 
+                    pred_coeff['angle'],
+                    pred_coeff['trans'],
+                    trans_params[2:][None],
+                    ], 1)
+                video_coeffs.append(pred_coeff)
+                full_coeffs.append(full_coeff.cpu().numpy())
+
+            semantic_npy = np.array(video_coeffs)[:,0] 
+
+            savemat(coeff_path, {'coeff_3dmm': semantic_npy, 'full_3dmm': np.array(full_coeffs)[0]})
+
+        return coeff_path, png_path, crop_info

src/utils/safetensor_helper.py

+
+
+def load_x_from_safetensor(checkpoint, key):
+    x_generator = {}
+    for k,v in checkpoint.items():
+        if key in k:
+            x_generator[k.replace(key+'.', '')] = v
+    return x_generator
\ No newline at end of file

src/utils/text2speech.py

+import os
+import tempfile
+from TTS.api import TTS
+
+
+class TTSTalker():
+    def __init__(self) -> None:
+        model_name = TTS().list_models()[0]
+        self.tts = TTS(model_name)
+
+    def test(self, text, language='en'):
+
+        tempf  = tempfile.NamedTemporaryFile(
+                delete = False,
+                suffix = ('.'+'wav'),
+            )
+
+        self.tts.tts_to_file(text, speaker=self.tts.speakers[0], language=language, file_path=tempf.name)
+
+        return tempf.name

src/utils/videoio.py

+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=False):
+    temp_file = str(uuid.uuid4())+'.mp4'
+    cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -vcodec copy "%s"' % (video, audio, temp_file)
+    os.system(cmd)
+
+    if watermark is False:
+        shutil.move(temp_file, save_path)
+    else:
+        # watermark
+        try:
+            ##### check if stable-diffusion-webui
+            import webui
+            from modules import paths
+            watarmark_path = paths.script_path+"/extensions/SadTalker/docs/sadtalker_logo.png"
+        except:
+            # get the root path of sadtalker.
+            dir_path = os.path.dirname(os.path.realpath(__file__))
+            watarmark_path = dir_path+"/../../docs/sadtalker_logo.png"
+
+        cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -filter_complex "[1]scale=100:-1[wm];[0][wm]overlay=(main_w-overlay_w)-10:10" "%s"' % (temp_file, watarmark_path, save_path)
+        os.system(cmd)
+        os.remove(temp_file)
\ No newline at end of file

webui.bat

+@echo off
+
+IF NOT EXIST venv (
+python -m venv venv
+) ELSE (
+echo venv folder already exists, skipping creation...
+)
+call .\venv\Scripts\activate.bat
+
+set PYTHON="venv\Scripts\Python.exe"
+echo venv %PYTHON%
+
+%PYTHON% Launcher.py
+
+echo.
+echo Launch unsuccessful. Exiting.
+pause
\ No newline at end of file

webui.sh

+#!/usr/bin/env bash
+
+
+# If run from macOS, load defaults from webui-macos-env.sh
+if [[ "$OSTYPE" == "darwin"* ]]; then
+    export TORCH_COMMAND="pip install torch==1.12.1 torchvision==0.13.1"
+fi
+
+# python3 executable
+if [[ -z "${python_cmd}" ]]
+then
+    python_cmd="python3"
+fi
+
+# git executable
+if [[ -z "${GIT}" ]]
+then
+    export GIT="git"
+fi
+
+# python3 venv without trailing slash (defaults to ${install_dir}/${clone_dir}/venv)
+if [[ -z "${venv_dir}" ]]
+then
+    venv_dir="venv"
+fi
+
+if [[ -z "${LAUNCH_SCRIPT}" ]]
+then
+    LAUNCH_SCRIPT="launcher.py"
+fi
+
+# this script cannot be run as root by default
+can_run_as_root=1
+
+# read any command line flags to the webui.sh script
+while getopts "f" flag > /dev/null 2>&1
+do
+    case ${flag} in
+        f) can_run_as_root=1;;
+        *) break;;
+    esac
+done
+
+# Disable sentry logging
+export ERROR_REPORTING=FALSE
+
+# Do not reinstall existing pip packages on Debian/Ubuntu
+export PIP_IGNORE_INSTALLED=0
+
+# Pretty print
+delimiter="################################################################"
+
+printf "\n%s\n" "${delimiter}"
+printf "\e[1m\e[32mInstall script for SadTalker + Web UI\n"
+printf "\e[1m\e[34mTested on Debian 11 (Bullseye)\e[0m"
+printf "\n%s\n" "${delimiter}"
+
+# Do not run as root
+if [[ $(id -u) -eq 0 && can_run_as_root -eq 0 ]]
+then
+    printf "\n%s\n" "${delimiter}"
+    printf "\e[1m\e[31mERROR: This script must not be launched as root, aborting...\e[0m"
+    printf "\n%s\n" "${delimiter}"
+    exit 1
+else
+    printf "\n%s\n" "${delimiter}"
+    printf "Running on \e[1m\e[32m%s\e[0m user" "$(whoami)"
+    printf "\n%s\n" "${delimiter}"
+fi
+
+if [[ -d .git ]]
+then
+    printf "\n%s\n" "${delimiter}"
+    printf "Repo already cloned, using it as install directory"
+    printf "\n%s\n" "${delimiter}"
+    install_dir="${PWD}/../"
+    clone_dir="${PWD##*/}"
+fi
+
+# Check prerequisites
+gpu_info=$(lspci 2>/dev/null | grep VGA)
+case "$gpu_info" in
+    *"Navi 1"*|*"Navi 2"*) export HSA_OVERRIDE_GFX_VERSION=10.3.0
+    ;;
+    *"Renoir"*) export HSA_OVERRIDE_GFX_VERSION=9.0.0
+        printf "\n%s\n" "${delimiter}"
+        printf "Experimental support for Renoir: make sure to have at least 4GB of VRAM and 10GB of RAM or enable cpu mode: --use-cpu all --no-half"
+        printf "\n%s\n" "${delimiter}"
+    ;;
+    *) 
+    ;;
+esac
+if echo "$gpu_info" | grep -q "AMD" && [[ -z "${TORCH_COMMAND}" ]]
+then
+    export TORCH_COMMAND="pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/rocm5.2"
+fi  
+
+for preq in "${GIT}" "${python_cmd}"
+do
+    if ! hash "${preq}" &>/dev/null
+    then
+        printf "\n%s\n" "${delimiter}"
+        printf "\e[1m\e[31mERROR: %s is not installed, aborting...\e[0m" "${preq}"
+        printf "\n%s\n" "${delimiter}"
+        exit 1
+    fi
+done
+
+if ! "${python_cmd}" -c "import venv" &>/dev/null
+then
+    printf "\n%s\n" "${delimiter}"
+    printf "\e[1m\e[31mERROR: python3-venv is not installed, aborting...\e[0m"
+    printf "\n%s\n" "${delimiter}"
+    exit 1
+fi
+
+printf "\n%s\n" "${delimiter}"
+printf "Create and activate python venv"
+printf "\n%s\n" "${delimiter}"
+cd "${install_dir}"/"${clone_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/%s/, aborting...\e[0m" "${install_dir}" "${clone_dir}"; exit 1; }
+if [[ ! -d "${venv_dir}" ]]
+then
+    "${python_cmd}" -m venv "${venv_dir}"
+    first_launch=1
+fi
+# shellcheck source=/dev/null
+if [[ -f "${venv_dir}"/bin/activate ]]
+then
+    source "${venv_dir}"/bin/activate
+else
+    printf "\n%s\n" "${delimiter}"
+    printf "\e[1m\e[31mERROR: Cannot activate python venv, aborting...\e[0m"
+    printf "\n%s\n" "${delimiter}"
+    exit 1
+fi
+
+printf "\n%s\n" "${delimiter}"
+printf "Launching launcher.py..."
+printf "\n%s\n" "${delimiter}"      
+exec "${python_cmd}" "${LAUNCH_SCRIPT}" "$@"
\ No newline at end of file