Blame for handler.py

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