Initial

516b9a9 over 1 year ago

13.3 kB

import os
import time
import json
import base64
import tempfile
import shutil
import subprocess
from io import BytesIO
from typing import Dict, Any, Optional

import torch
import numpy as np
import cv2
import glob
import pickle
from tqdm import tqdm
import copy
from argparse import Namespace
import requests
import gdown
from huggingface_hub import snapshot_download

from musetalk.utils.utils import get_file_type, get_video_fps, datagen
from musetalk.utils.preprocessing import (
    get_landmark_and_bbox,
    read_imgs,
    coord_placeholder,
)
from musetalk.utils.blending import get_image
from musetalk.utils.utils import load_all_model


class EndpointHandler:
    def __init__(self, path=""):
        """
        Initialize the handler with model loading and checkpoint downloading.
        """
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.project_dir = path or os.path.abspath(os.path.dirname(__file__))
        self.checkpoints_dir = os.path.join(self.project_dir, "checkpoints")

        # Download and load models
        self._download_model()
        self._load_models()

        # Create temp directory for processing
        self.temp_dir = tempfile.mkdtemp()

    def _download_model(self):
        """Download model checkpoints if they don't exist."""
        if not os.path.exists(self.checkpoints_dir):
            os.makedirs(self.checkpoints_dir)
            print("Checkpoint Not Downloaded, start downloading...")
            tic = time.time()

            # Main model
            snapshot_download(
                repo_id="TMElyralab/MuseTalk",
                local_dir=self.checkpoints_dir,
                max_workers=8,
                local_dir_use_symlinks=True,
            )

            # VAE weights
            snapshot_download(
                repo_id="stabilityai/sd-vae-ft-mse",
                local_dir=self.checkpoints_dir,
                max_workers=8,
                local_dir_use_symlinks=True,
            )

            # DWPose
            snapshot_download(
                repo_id="yzd-v/DWPose",
                local_dir=self.checkpoints_dir,
                max_workers=8,
                local_dir_use_symlinks=True,
            )

            # Whisper model
            self._download_whisper()

            # Face parser
            self._download_face_parser()

            # ResNet
            self._download_resnet()

            toc = time.time()
            print(f"Download cost {toc-tic} seconds")
        else:
            print("Already downloaded the model.")

    def _download_whisper(self):
        """Download Whisper model."""
        url = "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt"
        response = requests.get(url)
        if response.status_code == 200:
            file_path = f"{self.checkpoints_dir}/whisper/tiny.pt"
            os.makedirs(f"{self.checkpoints_dir}/whisper/", exist_ok=True)
            with open(file_path, "wb") as f:
                f.write(response.content)
        else:
            print(f"Whisper download failed, status code: {response.status_code}")

    def _download_face_parser(self):
        """Download face parser model."""
        url = "https://drive.google.com/uc?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812"
        os.makedirs(f"{self.checkpoints_dir}/face-parse-bisent/", exist_ok=True)
        output_path = f"{self.checkpoints_dir}/face-parse-bisent/79999_iter.pth"
        gdown.download(url, output_path, quiet=False)

    def _download_resnet(self):
        """Download ResNet model."""
        url = "https://download.pytorch.org/models/resnet18-5c106cde.pth"
        response = requests.get(url)
        if response.status_code == 200:
            file_path = (
                f"{self.checkpoints_dir}/face-parse-bisent/resnet18-5c106cde.pth"
            )
            with open(file_path, "wb") as f:
                f.write(response.content)
        else:
            print(f"ResNet download failed, status code: {response.status_code}")

    def _load_models(self):
        """Load all required models."""
        print("Loading models...")
        self.audio_processor, self.vae, self.unet, self.pe = load_all_model()
        self.timesteps = torch.tensor([0], device=self.device)
        print("Models loaded successfully!")

    def _save_file_from_base64(self, base64_data: str, filename: str) -> str:
        """Save base64 encoded file to temporary location."""
        file_data = base64.b64decode(base64_data)
        file_path = os.path.join(self.temp_dir, filename)
        with open(file_path, "wb") as f:
            f.write(file_data)
        return file_path

    def _preprocess_video(self, video_path: str) -> str:
        """Preprocess video to ensure 25fps."""
        dir_path, file_name = os.path.split(video_path)
        if file_name.startswith("processed_"):
            return video_path

        output_file_name = "processed_" + file_name
        output_video = os.path.join(dir_path, output_file_name)

        # Convert to 25fps
        command = f"ffmpeg -i {video_path} -r 25 {output_video} -y"
        subprocess.run(command, shell=True, check=True)
        return output_video

    @torch.no_grad()
    def _generate_video(
        self, audio_path: str, video_path: str, bbox_shift: int = -1
    ) -> str:
        """Generate lip-synced video."""
        # Setup arguments
        args_dict = {
            "result_dir": self.temp_dir,
            "fps": 25,
            "batch_size": 8,
            "output_vid_name": "",
            "use_saved_coord": False,
        }
        args = Namespace(**args_dict)

        # Generate output paths
        input_basename = os.path.basename(video_path).split(".")[0]
        audio_basename = os.path.basename(audio_path).split(".")[0]
        output_basename = f"{input_basename}_{audio_basename}"
        result_img_save_path = os.path.join(args.result_dir, output_basename)
        crop_coord_save_path = os.path.join(
            result_img_save_path, input_basename + ".pkl"
        )
        os.makedirs(result_img_save_path, exist_ok=True)

        output_vid_name = os.path.join(args.result_dir, output_basename + ".mp4")

        # Extract frames from video
        if get_file_type(video_path) == "video":
            save_dir_full = os.path.join(args.result_dir, input_basename)
            os.makedirs(save_dir_full, exist_ok=True)
            cmd = f"ffmpeg -v fatal -i {video_path} -start_number 0 {save_dir_full}/%08d.png"
            os.system(cmd)
            input_img_list = sorted(
                glob.glob(os.path.join(save_dir_full, "*.[jpJP][pnPN]*[gG]"))
            )
            fps = get_video_fps(video_path)
        else:
            input_img_list = glob.glob(os.path.join(video_path, "*.[jpJP][pnPN]*[gG]"))
            input_img_list = sorted(
                input_img_list,
                key=lambda x: int(os.path.splitext(os.path.basename(x))[0]),
            )
            fps = args.fps

        # Extract audio features
        whisper_feature = self.audio_processor.audio2feat(audio_path)
        whisper_chunks = self.audio_processor.feature2chunks(
            feature_array=whisper_feature, fps=fps
        )

        # Preprocess input images
        if os.path.exists(crop_coord_save_path) and args.use_saved_coord:
            print("Using extracted coordinates")
            with open(crop_coord_save_path, "rb") as f:
                coord_list = pickle.load(f)
            frame_list = read_imgs(input_img_list)
        else:
            print("Extracting landmarks...")
            coord_list, frame_list = get_landmark_and_bbox(input_img_list, bbox_shift)
            with open(crop_coord_save_path, "wb") as f:
                pickle.dump(coord_list, f)

        # Prepare input latents
        input_latent_list = []
        for bbox, frame in zip(coord_list, frame_list):
            if bbox == coord_placeholder:
                continue
            x1, y1, x2, y2 = bbox
            crop_frame = frame[y1:y2, x1:x2]
            crop_frame = cv2.resize(
                crop_frame, (256, 256), interpolation=cv2.INTER_LANCZOS4
            )
            latents = self.vae.get_latents_for_unet(crop_frame)
            input_latent_list.append(latents)

        # Create cyclic lists for smoothing
        frame_list_cycle = frame_list + frame_list[::-1]
        coord_list_cycle = coord_list + coord_list[::-1]
        input_latent_list_cycle = input_latent_list + input_latent_list[::-1]

        # Inference batch by batch
        print("Starting inference...")
        video_num = len(whisper_chunks)
        batch_size = args.batch_size
        gen = datagen(whisper_chunks, input_latent_list_cycle, batch_size)
        res_frame_list = []

        for i, (whisper_batch, latent_batch) in enumerate(
            tqdm(gen, total=int(np.ceil(float(video_num) / batch_size)))
        ):
            tensor_list = [torch.FloatTensor(arr) for arr in whisper_batch]
            audio_feature_batch = torch.stack(tensor_list).to(self.unet.device)
            audio_feature_batch = self.pe(audio_feature_batch)

            pred_latents = self.unet.model(
                latent_batch, self.timesteps, encoder_hidden_states=audio_feature_batch
            ).sample
            recon = self.vae.decode_latents(pred_latents)
            for res_frame in recon:
                res_frame_list.append(res_frame)

        # Pad to full image
        print("Padding talking image to original video...")
        for i, res_frame in enumerate(tqdm(res_frame_list)):
            bbox = coord_list_cycle[i % (len(coord_list_cycle))]
            ori_frame = copy.deepcopy(frame_list_cycle[i % (len(frame_list_cycle))])
            x1, y1, x2, y2 = bbox
            try:
                res_frame = cv2.resize(res_frame.astype(np.uint8), (x2 - x1, y2 - y1))
            except:
                continue

            combine_frame = get_image(ori_frame, res_frame, bbox)
            cv2.imwrite(f"{result_img_save_path}/{str(i).zfill(8)}.png", combine_frame)

        # Create video from images
        cmd_img2video = f"ffmpeg -y -v fatal -r {fps} -f image2 -i {result_img_save_path}/%08d.png -vcodec libx264 -vf format=rgb24,scale=out_color_matrix=bt709,format=yuv420p -crf 18 {self.temp_dir}/temp.mp4"
        os.system(cmd_img2video)

        # Combine with audio
        cmd_combine_audio = f"ffmpeg -y -v fatal -i {audio_path} -i {self.temp_dir}/temp.mp4 {output_vid_name}"
        os.system(cmd_combine_audio)

        # Cleanup
        os.remove(f"{self.temp_dir}/temp.mp4")
        shutil.rmtree(result_img_save_path)

        print(f"Result saved to {output_vid_name}")
        return output_vid_name

    def _encode_video_to_base64(self, video_path: str) -> str:
        """Encode video file to base64."""
        with open(video_path, "rb") as video_file:
            video_bytes = video_file.read()
            return base64.b64encode(video_bytes).decode("utf-8")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Handle inference request.
        Expected input format:
        {
            "inputs": {
                "audio": "base64_encoded_audio_file",
                "video": "base64_encoded_video_file",
                "bbox_shift": -1  # optional, defaults to -1
            }
        }
        Returns:
        {
            "video": "base64_encoded_output_video",
            "message": "success"
        }
        """
        try:
            # Extract inputs
            inputs = data.get("inputs", {})
            audio_b64 = inputs.get("audio")
            video_b64 = inputs.get("video")
            bbox_shift = inputs.get("bbox_shift", -1)

            if not audio_b64 or not video_b64:
                return {"error": "Both audio and video inputs are required"}

            # Save input files
            audio_path = self._save_file_from_base64(audio_b64, "input_audio.wav")
            video_path = self._save_file_from_base64(video_b64, "input_video.mp4")

            # Preprocess video
            processed_video_path = self._preprocess_video(video_path)

            # Generate output video
            output_video_path = self._generate_video(
                audio_path, processed_video_path, bbox_shift
            )

            # Encode result to base64
            result_b64 = self._encode_video_to_base64(output_video_path)

            return {"video": result_b64, "message": "success"}

        except Exception as e:
            return {"error": f"Processing failed: {str(e)}"}

        finally:
            # Cleanup temporary files
            if hasattr(self, "temp_dir") and os.path.exists(self.temp_dir):
                try:
                    shutil.rmtree(self.temp_dir)
                    self.temp_dir = (
                        tempfile.mkdtemp()
                    )  # Create new temp dir for next request
                except:
                    pass

    def __del__(self):
        """Cleanup on destruction."""
        if hasattr(self, "temp_dir") and os.path.exists(self.temp_dir):
            try:
                shutil.rmtree(self.temp_dir)
            except:
                pass

	import os
	import time
	import json
	import base64
	import tempfile
	import shutil
	import subprocess
	from io import BytesIO
	from typing import Dict, Any, Optional

	import torch
	import numpy as np
	import cv2
	import glob
	import pickle
	from tqdm import tqdm
	import copy
	from argparse import Namespace
	import requests
	import gdown
	from huggingface_hub import snapshot_download

	from musetalk.utils.utils import get_file_type, get_video_fps, datagen
	from musetalk.utils.preprocessing import (
	get_landmark_and_bbox,
	read_imgs,
	coord_placeholder,
	)
	from musetalk.utils.blending import get_image
	from musetalk.utils.utils import load_all_model


	class EndpointHandler:
	def __init__(self, path=""):
	"""
	Initialize the handler with model loading and checkpoint downloading.
	"""
	self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	self.project_dir = path or os.path.abspath(os.path.dirname(__file__))
	self.checkpoints_dir = os.path.join(self.project_dir, "checkpoints")

	# Download and load models
	self._download_model()
	self._load_models()

	# Create temp directory for processing
	self.temp_dir = tempfile.mkdtemp()

	def _download_model(self):
	"""Download model checkpoints if they don't exist."""
	if not os.path.exists(self.checkpoints_dir):
	os.makedirs(self.checkpoints_dir)
	print("Checkpoint Not Downloaded, start downloading...")
	tic = time.time()

	# Main model
	snapshot_download(
	repo_id="TMElyralab/MuseTalk",
	local_dir=self.checkpoints_dir,
	max_workers=8,
	local_dir_use_symlinks=True,
	)

	# VAE weights
	snapshot_download(
	repo_id="stabilityai/sd-vae-ft-mse",
	local_dir=self.checkpoints_dir,
	max_workers=8,
	local_dir_use_symlinks=True,
	)

	# DWPose
	snapshot_download(
	repo_id="yzd-v/DWPose",
	local_dir=self.checkpoints_dir,
	max_workers=8,
	local_dir_use_symlinks=True,
	)

	# Whisper model
	self._download_whisper()

	# Face parser
	self._download_face_parser()

	# ResNet
	self._download_resnet()

	toc = time.time()
	print(f"Download cost {toc-tic} seconds")
	else:
	print("Already downloaded the model.")

	def _download_whisper(self):
	"""Download Whisper model."""
	url = "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt"
	response = requests.get(url)
	if response.status_code == 200:
	file_path = f"{self.checkpoints_dir}/whisper/tiny.pt"
	os.makedirs(f"{self.checkpoints_dir}/whisper/", exist_ok=True)
	with open(file_path, "wb") as f:
	f.write(response.content)
	else:
	print(f"Whisper download failed, status code: {response.status_code}")

	def _download_face_parser(self):
	"""Download face parser model."""
	url = "https://drive.google.com/uc?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812"
	os.makedirs(f"{self.checkpoints_dir}/face-parse-bisent/", exist_ok=True)
	output_path = f"{self.checkpoints_dir}/face-parse-bisent/79999_iter.pth"
	gdown.download(url, output_path, quiet=False)

	def _download_resnet(self):
	"""Download ResNet model."""
	url = "https://download.pytorch.org/models/resnet18-5c106cde.pth"
	response = requests.get(url)
	if response.status_code == 200:
	file_path = (
	f"{self.checkpoints_dir}/face-parse-bisent/resnet18-5c106cde.pth"
	)
	with open(file_path, "wb") as f:
	f.write(response.content)
	else:
	print(f"ResNet download failed, status code: {response.status_code}")

	def _load_models(self):
	"""Load all required models."""
	print("Loading models...")
	self.audio_processor, self.vae, self.unet, self.pe = load_all_model()
	self.timesteps = torch.tensor([0], device=self.device)
	print("Models loaded successfully!")

	def _save_file_from_base64(self, base64_data: str, filename: str) -> str:
	"""Save base64 encoded file to temporary location."""
	file_data = base64.b64decode(base64_data)
	file_path = os.path.join(self.temp_dir, filename)
	with open(file_path, "wb") as f:
	f.write(file_data)
	return file_path

	def _preprocess_video(self, video_path: str) -> str:
	"""Preprocess video to ensure 25fps."""
	dir_path, file_name = os.path.split(video_path)
	if file_name.startswith("processed_"):
	return video_path

	output_file_name = "processed_" + file_name
	output_video = os.path.join(dir_path, output_file_name)

	# Convert to 25fps
	command = f"ffmpeg -i {video_path} -r 25 {output_video} -y"
	subprocess.run(command, shell=True, check=True)
	return output_video

	@torch.no_grad()
	def _generate_video(
	self, audio_path: str, video_path: str, bbox_shift: int = -1
	) -> str:
	"""Generate lip-synced video."""
	# Setup arguments
	args_dict = {
	"result_dir": self.temp_dir,
	"fps": 25,
	"batch_size": 8,
	"output_vid_name": "",
	"use_saved_coord": False,
	}
	args = Namespace(**args_dict)

	# Generate output paths
	input_basename = os.path.basename(video_path).split(".")[0]
	audio_basename = os.path.basename(audio_path).split(".")[0]
	output_basename = f"{input_basename}_{audio_basename}"
	result_img_save_path = os.path.join(args.result_dir, output_basename)
	crop_coord_save_path = os.path.join(
	result_img_save_path, input_basename + ".pkl"
	)
	os.makedirs(result_img_save_path, exist_ok=True)

	output_vid_name = os.path.join(args.result_dir, output_basename + ".mp4")

	# Extract frames from video
	if get_file_type(video_path) == "video":
	save_dir_full = os.path.join(args.result_dir, input_basename)
	os.makedirs(save_dir_full, exist_ok=True)
	cmd = f"ffmpeg -v fatal -i {video_path} -start_number 0 {save_dir_full}/%08d.png"
	os.system(cmd)
	input_img_list = sorted(
	glob.glob(os.path.join(save_dir_full, ".[jpJP][pnPN][gG]"))
	)
	fps = get_video_fps(video_path)
	else:
	input_img_list = glob.glob(os.path.join(video_path, ".[jpJP][pnPN][gG]"))
	input_img_list = sorted(
	input_img_list,
	key=lambda x: int(os.path.splitext(os.path.basename(x))[0]),
	)
	fps = args.fps

	# Extract audio features
	whisper_feature = self.audio_processor.audio2feat(audio_path)
	whisper_chunks = self.audio_processor.feature2chunks(
	feature_array=whisper_feature, fps=fps
	)

	# Preprocess input images
	if os.path.exists(crop_coord_save_path) and args.use_saved_coord:
	print("Using extracted coordinates")
	with open(crop_coord_save_path, "rb") as f:
	coord_list = pickle.load(f)
	frame_list = read_imgs(input_img_list)
	else:
	print("Extracting landmarks...")
	coord_list, frame_list = get_landmark_and_bbox(input_img_list, bbox_shift)
	with open(crop_coord_save_path, "wb") as f:
	pickle.dump(coord_list, f)

	# Prepare input latents
	input_latent_list = []
	for bbox, frame in zip(coord_list, frame_list):
	if bbox == coord_placeholder:
	continue
	x1, y1, x2, y2 = bbox
	crop_frame = frame[y1:y2, x1:x2]
	crop_frame = cv2.resize(
	crop_frame, (256, 256), interpolation=cv2.INTER_LANCZOS4
	)
	latents = self.vae.get_latents_for_unet(crop_frame)
	input_latent_list.append(latents)

	# Create cyclic lists for smoothing
	frame_list_cycle = frame_list + frame_list[::-1]
	coord_list_cycle = coord_list + coord_list[::-1]
	input_latent_list_cycle = input_latent_list + input_latent_list[::-1]

	# Inference batch by batch
	print("Starting inference...")
	video_num = len(whisper_chunks)
	batch_size = args.batch_size
	gen = datagen(whisper_chunks, input_latent_list_cycle, batch_size)
	res_frame_list = []

	for i, (whisper_batch, latent_batch) in enumerate(
	tqdm(gen, total=int(np.ceil(float(video_num) / batch_size)))
	):
	tensor_list = [torch.FloatTensor(arr) for arr in whisper_batch]
	audio_feature_batch = torch.stack(tensor_list).to(self.unet.device)
	audio_feature_batch = self.pe(audio_feature_batch)

	pred_latents = self.unet.model(
	latent_batch, self.timesteps, encoder_hidden_states=audio_feature_batch
	).sample
	recon = self.vae.decode_latents(pred_latents)
	for res_frame in recon:
	res_frame_list.append(res_frame)

	# Pad to full image
	print("Padding talking image to original video...")
	for i, res_frame in enumerate(tqdm(res_frame_list)):
	bbox = coord_list_cycle[i % (len(coord_list_cycle))]
	ori_frame = copy.deepcopy(frame_list_cycle[i % (len(frame_list_cycle))])
	x1, y1, x2, y2 = bbox
	try:
	res_frame = cv2.resize(res_frame.astype(np.uint8), (x2 - x1, y2 - y1))
	except:
	continue

	combine_frame = get_image(ori_frame, res_frame, bbox)
	cv2.imwrite(f"{result_img_save_path}/{str(i).zfill(8)}.png", combine_frame)

	# Create video from images
	cmd_img2video = f"ffmpeg -y -v fatal -r {fps} -f image2 -i {result_img_save_path}/%08d.png -vcodec libx264 -vf format=rgb24,scale=out_color_matrix=bt709,format=yuv420p -crf 18 {self.temp_dir}/temp.mp4"
	os.system(cmd_img2video)

	# Combine with audio
	cmd_combine_audio = f"ffmpeg -y -v fatal -i {audio_path} -i {self.temp_dir}/temp.mp4 {output_vid_name}"
	os.system(cmd_combine_audio)

	# Cleanup
	os.remove(f"{self.temp_dir}/temp.mp4")
	shutil.rmtree(result_img_save_path)

	print(f"Result saved to {output_vid_name}")
	return output_vid_name

	def _encode_video_to_base64(self, video_path: str) -> str:
	"""Encode video file to base64."""
	with open(video_path, "rb") as video_file:
	video_bytes = video_file.read()
	return base64.b64encode(video_bytes).decode("utf-8")

	def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
	"""
	Handle inference request.

	Expected input format:
	{
	"inputs": {
	"audio": "base64_encoded_audio_file",
	"video": "base64_encoded_video_file",
	"bbox_shift": -1 # optional, defaults to -1
	}
	}

	Returns:
	{
	"video": "base64_encoded_output_video",
	"message": "success"
	}
	"""
	try:
	# Extract inputs
	inputs = data.get("inputs", {})
	audio_b64 = inputs.get("audio")
	video_b64 = inputs.get("video")
	bbox_shift = inputs.get("bbox_shift", -1)

	if not audio_b64 or not video_b64:
	return {"error": "Both audio and video inputs are required"}

	# Save input files
	audio_path = self._save_file_from_base64(audio_b64, "input_audio.wav")
	video_path = self._save_file_from_base64(video_b64, "input_video.mp4")

	# Preprocess video
	processed_video_path = self._preprocess_video(video_path)

	# Generate output video
	output_video_path = self._generate_video(
	audio_path, processed_video_path, bbox_shift
	)

	# Encode result to base64
	result_b64 = self._encode_video_to_base64(output_video_path)

	return {"video": result_b64, "message": "success"}

	except Exception as e:
	return {"error": f"Processing failed: {str(e)}"}

	finally:
	# Cleanup temporary files
	if hasattr(self, "temp_dir") and os.path.exists(self.temp_dir):
	try:
	shutil.rmtree(self.temp_dir)
	self.temp_dir = (
	tempfile.mkdtemp()
	) # Create new temp dir for next request
	except:
	pass

	def __del__(self):
	"""Cleanup on destruction."""
	if hasattr(self, "temp_dir") and os.path.exists(self.temp_dir):
	try:
	shutil.rmtree(self.temp_dir)
	except:
	pass