Initial Commit

5f061c9 over 1 year ago

13.8 kB

import torch
import os
import shutil
import subprocess
import json
import tempfile
import base64
from typing import Dict, Any, Optional, List
import soundfile as sf
from datetime import datetime
from huggingface_hub import snapshot_download


class EndpointHandler:
    def __init__(self, path=""):
        """
        Initialize the handler for Meigen MultiTalk inference.
        Args:
            path: Path to the model directory (not used in this implementation)
        """
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.setup_models()
        self.setup_gpu_config()

    def setup_models(self):
        """Download and setup all required models."""
        print("Setting up models...")

        # Download Wan2.1-I2V-14B-480P model
        self.wan_model_path = snapshot_download(
            repo_id="Wan-AI/Wan2.1-I2V-14B-480P",
            local_dir="./weights/Wan2.1-I2V-14B-480P",
        )

        # Download Chinese wav2vec2 model
        self.wav2vec_path = snapshot_download(
            repo_id="TencentGameMate/chinese-wav2vec2-base",
            local_dir="./weights/chinese-wav2vec2-base",
        )

        # Download MeiGen MultiTalk weights
        self.multitalk_path = snapshot_download(
            repo_id="MeiGen-AI/MeiGen-MultiTalk",
            local_dir="./weights/MeiGen-MultiTalk",
        )

        self.setup_model_files()

    def setup_model_files(self):
        """Setup model files by copying and renaming as needed."""
        base_model_dir = "./weights/Wan2.1-I2V-14B-480P"
        multitalk_dir = "./weights/MeiGen-MultiTalk"

        # File to rename
        original_index = os.path.join(
            base_model_dir, "diffusion_pytorch_model.safetensors.index.json"
        )
        backup_index = os.path.join(
            base_model_dir, "diffusion_pytorch_model.safetensors.index.json_old"
        )

        # Rename the original index file
        if os.path.exists(original_index):
            os.rename(original_index, backup_index)
            print("Renamed original index file to .json_old")

        # Copy updated index file from MultiTalk
        shutil.copy2(
            os.path.join(
                multitalk_dir, "diffusion_pytorch_model.safetensors.index.json"
            ),
            base_model_dir,
        )

        # Copy MultiTalk model weights
        shutil.copy2(
            os.path.join(multitalk_dir, "multitalk.safetensors"), base_model_dir
        )

        print("Copied MultiTalk files into base model directory.")

    def setup_gpu_config(self):
        """Setup GPU configuration and VRAM parameters."""
        self.num_gpus = torch.cuda.device_count()
        print(f"GPU AVAILABLE: {self.num_gpus}")

        if torch.cuda.is_available():
            self.gpu_name = torch.cuda.get_device_name(torch.cuda.current_device())
            print(f"Current GPU: {self.gpu_name}")

            # GPU to VRAM mapping
            GPU_TO_VRAM_PARAMS = {
                "NVIDIA A100": 11000000000,
                "NVIDIA A100-SXM4-40GB": 11000000000,
                "NVIDIA A100-SXM4-80GB": 22000000000,
                "NVIDIA L4": 5000000000,
                "NVIDIA L40S": 11000000000,
            }

            self.used_vram_params = GPU_TO_VRAM_PARAMS.get(self.gpu_name, 5000000000)
            print(f"Using {self.used_vram_params} for num_persistent_param_in_dit")
        else:
            raise RuntimeError("No CUDA-compatible GPU found. A GPU is required.")

    def trim_audio_to_5s_temp(self, audio_path: str, sample_rate: int = 16000) -> str:
        """
        Trim audio to 5 seconds and save to temporary file.
        Args:
            audio_path: Path to input audio file
            sample_rate: Target sample rate
        Returns:
            Path to trimmed audio file
        """
        max_duration_sec = 5
        audio, sr = sf.read(audio_path)

        if sr != sample_rate:
            sample_rate = sr

        max_samples = max_duration_sec * sample_rate
        if len(audio) > max_samples:
            audio = audio[:max_samples]

        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
        base_name = os.path.splitext(os.path.basename(audio_path))[0]
        temp_filename = f"{base_name}_trimmed_{timestamp}.wav"
        temp_path = os.path.join(tempfile.gettempdir(), temp_filename)

        sf.write(temp_path, audio, samplerate=sample_rate)
        return temp_path

    def save_base64_file(self, base64_data: str, file_extension: str) -> str:
        """
        Save base64 encoded file to temporary location.
        Args:
            base64_data: Base64 encoded file data
            file_extension: File extension (e.g., '.png', '.wav')
        Returns:
            Path to saved file
        """
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
        temp_filename = f"temp_file_{timestamp}{file_extension}"
        temp_path = os.path.join(tempfile.gettempdir(), temp_filename)

        with open(temp_path, "wb") as f:
            f.write(base64.b64decode(base64_data))

        return temp_path

    def create_temp_input_json(
        self,
        prompt: str,
        cond_image_path: str,
        cond_audio_path_spk1: str,
        cond_audio_path_spk2: Optional[str] = None,
    ) -> str:
        """
        Create a temporary JSON file with the user-provided prompt, image, and audio paths.
        Args:
            prompt: Text prompt for generation
            cond_image_path: Path to conditioning image
            cond_audio_path_spk1: Path to first speaker audio
            cond_audio_path_spk2: Path to second speaker audio (optional)
        Returns:
            Path to temporary JSON file
        """
        if cond_audio_path_spk2 is None:
            data = {
                "prompt": prompt,
                "cond_image": cond_image_path,
                "cond_audio": {"person1": cond_audio_path_spk1},
            }
        else:
            data = {
                "prompt": prompt,
                "cond_image": cond_image_path,
                "audio_type": "para",
                "cond_audio": {
                    "person1": cond_audio_path_spk1,
                    "person2": cond_audio_path_spk2,
                },
            }

        temp_json = tempfile.NamedTemporaryFile(
            delete=False, suffix=".json", mode="w", encoding="utf-8"
        )
        json.dump(data, temp_json, indent=4)
        temp_json_path = temp_json.name
        temp_json.close()

        print(f"Temporary input JSON saved to: {temp_json_path}")
        return temp_json_path

    def run_inference(
        self,
        prompt: str,
        cond_image_path: str,
        cond_audio_path_spk1: str,
        cond_audio_path_spk2: Optional[str] = None,
        sample_steps: int = 12,
    ) -> str:
        """
        Run the MultiTalk inference.
        Args:
            prompt: Text prompt
            cond_image_path: Path to conditioning image
            cond_audio_path_spk1: Path to first speaker audio
            cond_audio_path_spk2: Path to second speaker audio (optional)
            sample_steps: Number of sampling steps
        Returns:
            Path to generated video
        """
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
        result_filename = f"meigen_multitalk_result_{sample_steps}_steps_{timestamp}"
        temp_files_to_cleanup = []

        try:
            # Trim audio files to 5 seconds
            trimmed_audio_path_spk1 = self.trim_audio_to_5s_temp(cond_audio_path_spk1)
            if trimmed_audio_path_spk1 != cond_audio_path_spk1:
                cond_audio_path_spk1 = trimmed_audio_path_spk1
                temp_files_to_cleanup.append(trimmed_audio_path_spk1)

            if cond_audio_path_spk2 is not None:
                trimmed_audio_path_spk2 = self.trim_audio_to_5s_temp(
                    cond_audio_path_spk2
                )
                if trimmed_audio_path_spk2 != cond_audio_path_spk2:
                    cond_audio_path_spk2 = trimmed_audio_path_spk2
                    temp_files_to_cleanup.append(trimmed_audio_path_spk2)

            # Prepare input JSON
            input_json_path = self.create_temp_input_json(
                prompt, cond_image_path, cond_audio_path_spk1, cond_audio_path_spk2
            )
            temp_files_to_cleanup.append(input_json_path)

            # Base command arguments
            common_args = [
                "--ckpt_dir",
                "weights/Wan2.1-I2V-14B-480P",
                "--wav2vec_dir",
                "weights/chinese-wav2vec2-base",
                "--input_json",
                input_json_path,
                "--sample_steps",
                str(sample_steps),
                "--mode",
                "streaming",
                "--use_teacache",
                "--save_file",
                result_filename,
            ]

            # Choose command based on number of GPUs
            if self.num_gpus > 1:
                cmd = [
                    "torchrun",
                    f"--nproc_per_node={self.num_gpus}",
                    "--standalone",
                    "generate_multitalk.py",
                    "--dit_fsdp",
                    "--t5_fsdp",
                    "--ulysses_size",
                    str(self.num_gpus),
                ] + common_args
            else:
                cmd = [
                    "python3",
                    "generate_multitalk.py",
                    "--num_persistent_param_in_dit",
                    str(self.used_vram_params),
                ] + common_args

            # Run inference
            with open("inference.log", "w") as log_file:
                process = subprocess.Popen(
                    cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                )
                for line in process.stdout:
                    print(line, end="")
                    log_file.write(line)
                process.wait()

            if process.returncode != 0:
                raise RuntimeError("Inference failed. Check inference.log for details.")

            return f"{result_filename}.mp4"

        finally:
            # Clean up temporary files
            for f in temp_files_to_cleanup:
                try:
                    if os.path.exists(f):
                        os.remove(f)
                        print(f"[INFO] Removed temporary file: {f}")
                except Exception as e:
                    print(f"[WARNING] Could not remove {f}: {e}")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Main inference method called by Hugging Face Inference API.
        Args:
            data: Input data dictionary containing:
                - inputs: Main input data
                - parameters: Optional parameters
        Returns:
            Dictionary containing the generated video
        """
        try:
            # Extract inputs
            inputs = data.get("inputs", {})
            parameters = data.get("parameters", {})

            # Required inputs
            prompt = inputs.get("prompt", "")
            if not prompt:
                raise ValueError("Prompt is required")

            # Image input (base64 encoded)
            image_data = inputs.get("image")
            if not image_data:
                raise ValueError("Image is required")

            # Audio inputs (base64 encoded)
            audio_spk1_data = inputs.get("audio_spk1")
            if not audio_spk1_data:
                raise ValueError("Audio for speaker 1 is required")

            audio_spk2_data = inputs.get("audio_spk2")  # Optional

            # Parameters
            sample_steps = parameters.get("sample_steps", 12)
            sample_steps = max(2, min(25, sample_steps))  # Clamp between 2 and 25

            # Save base64 files to temporary locations
            temp_files_to_cleanup = []

            # Save image
            image_path = self.save_base64_file(image_data, ".png")
            temp_files_to_cleanup.append(image_path)

            # Save audio files
            audio_spk1_path = self.save_base64_file(audio_spk1_data, ".wav")
            temp_files_to_cleanup.append(audio_spk1_path)

            audio_spk2_path = None
            if audio_spk2_data:
                audio_spk2_path = self.save_base64_file(audio_spk2_data, ".wav")
                temp_files_to_cleanup.append(audio_spk2_path)

            # Run inference
            output_video_path = self.run_inference(
                prompt=prompt,
                cond_image_path=image_path,
                cond_audio_path_spk1=audio_spk1_path,
                cond_audio_path_spk2=audio_spk2_path,
                sample_steps=sample_steps,
            )

            # Read and encode output video
            with open(output_video_path, "rb") as f:
                video_data = base64.b64encode(f.read()).decode()

            # Clean up input files
            for f in temp_files_to_cleanup:
                try:
                    if os.path.exists(f):
                        os.remove(f)
                except Exception as e:
                    print(f"[WARNING] Could not remove {f}: {e}")

            # Clean up output video
            try:
                if os.path.exists(output_video_path):
                    os.remove(output_video_path)
            except Exception as e:
                print(f"[WARNING] Could not remove output video: {e}")

            return {"generated_video": video_data, "content_type": "video/mp4"}

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

	import torch
	import os
	import shutil
	import subprocess
	import json
	import tempfile
	import base64
	from typing import Dict, Any, Optional, List
	import soundfile as sf
	from datetime import datetime
	from huggingface_hub import snapshot_download


	class EndpointHandler:
	def __init__(self, path=""):
	"""
	Initialize the handler for Meigen MultiTalk inference.

	Args:
	path: Path to the model directory (not used in this implementation)
	"""
	self.device = "cuda" if torch.cuda.is_available() else "cpu"
	self.setup_models()
	self.setup_gpu_config()

	def setup_models(self):
	"""Download and setup all required models."""
	print("Setting up models...")

	# Download Wan2.1-I2V-14B-480P model
	self.wan_model_path = snapshot_download(
	repo_id="Wan-AI/Wan2.1-I2V-14B-480P",
	local_dir="./weights/Wan2.1-I2V-14B-480P",
	)

	# Download Chinese wav2vec2 model
	self.wav2vec_path = snapshot_download(
	repo_id="TencentGameMate/chinese-wav2vec2-base",
	local_dir="./weights/chinese-wav2vec2-base",
	)

	# Download MeiGen MultiTalk weights
	self.multitalk_path = snapshot_download(
	repo_id="MeiGen-AI/MeiGen-MultiTalk",
	local_dir="./weights/MeiGen-MultiTalk",
	)

	self.setup_model_files()

	def setup_model_files(self):
	"""Setup model files by copying and renaming as needed."""
	base_model_dir = "./weights/Wan2.1-I2V-14B-480P"
	multitalk_dir = "./weights/MeiGen-MultiTalk"

	# File to rename
	original_index = os.path.join(
	base_model_dir, "diffusion_pytorch_model.safetensors.index.json"
	)
	backup_index = os.path.join(
	base_model_dir, "diffusion_pytorch_model.safetensors.index.json_old"
	)

	# Rename the original index file
	if os.path.exists(original_index):
	os.rename(original_index, backup_index)
	print("Renamed original index file to .json_old")

	# Copy updated index file from MultiTalk
	shutil.copy2(
	os.path.join(
	multitalk_dir, "diffusion_pytorch_model.safetensors.index.json"
	),
	base_model_dir,
	)

	# Copy MultiTalk model weights
	shutil.copy2(
	os.path.join(multitalk_dir, "multitalk.safetensors"), base_model_dir
	)

	print("Copied MultiTalk files into base model directory.")

	def setup_gpu_config(self):
	"""Setup GPU configuration and VRAM parameters."""
	self.num_gpus = torch.cuda.device_count()
	print(f"GPU AVAILABLE: {self.num_gpus}")

	if torch.cuda.is_available():
	self.gpu_name = torch.cuda.get_device_name(torch.cuda.current_device())
	print(f"Current GPU: {self.gpu_name}")

	# GPU to VRAM mapping
	GPU_TO_VRAM_PARAMS = {
	"NVIDIA A100": 11000000000,
	"NVIDIA A100-SXM4-40GB": 11000000000,
	"NVIDIA A100-SXM4-80GB": 22000000000,
	"NVIDIA L4": 5000000000,
	"NVIDIA L40S": 11000000000,
	}

	self.used_vram_params = GPU_TO_VRAM_PARAMS.get(self.gpu_name, 5000000000)
	print(f"Using {self.used_vram_params} for num_persistent_param_in_dit")
	else:
	raise RuntimeError("No CUDA-compatible GPU found. A GPU is required.")

	def trim_audio_to_5s_temp(self, audio_path: str, sample_rate: int = 16000) -> str:
	"""
	Trim audio to 5 seconds and save to temporary file.

	Args:
	audio_path: Path to input audio file
	sample_rate: Target sample rate

	Returns:
	Path to trimmed audio file
	"""
	max_duration_sec = 5
	audio, sr = sf.read(audio_path)

	if sr != sample_rate:
	sample_rate = sr

	max_samples = max_duration_sec * sample_rate
	if len(audio) > max_samples:
	audio = audio[:max_samples]

	timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
	base_name = os.path.splitext(os.path.basename(audio_path))[0]
	temp_filename = f"{base_name}_trimmed_{timestamp}.wav"
	temp_path = os.path.join(tempfile.gettempdir(), temp_filename)

	sf.write(temp_path, audio, samplerate=sample_rate)
	return temp_path

	def save_base64_file(self, base64_data: str, file_extension: str) -> str:
	"""
	Save base64 encoded file to temporary location.

	Args:
	base64_data: Base64 encoded file data
	file_extension: File extension (e.g., '.png', '.wav')

	Returns:
	Path to saved file
	"""
	timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
	temp_filename = f"temp_file_{timestamp}{file_extension}"
	temp_path = os.path.join(tempfile.gettempdir(), temp_filename)

	with open(temp_path, "wb") as f:
	f.write(base64.b64decode(base64_data))

	return temp_path

	def create_temp_input_json(
	self,
	prompt: str,
	cond_image_path: str,
	cond_audio_path_spk1: str,
	cond_audio_path_spk2: Optional[str] = None,
	) -> str:
	"""
	Create a temporary JSON file with the user-provided prompt, image, and audio paths.

	Args:
	prompt: Text prompt for generation
	cond_image_path: Path to conditioning image
	cond_audio_path_spk1: Path to first speaker audio
	cond_audio_path_spk2: Path to second speaker audio (optional)

	Returns:
	Path to temporary JSON file
	"""
	if cond_audio_path_spk2 is None:
	data = {
	"prompt": prompt,
	"cond_image": cond_image_path,
	"cond_audio": {"person1": cond_audio_path_spk1},
	}
	else:
	data = {
	"prompt": prompt,
	"cond_image": cond_image_path,
	"audio_type": "para",
	"cond_audio": {
	"person1": cond_audio_path_spk1,
	"person2": cond_audio_path_spk2,
	},
	}

	temp_json = tempfile.NamedTemporaryFile(
	delete=False, suffix=".json", mode="w", encoding="utf-8"
	)
	json.dump(data, temp_json, indent=4)
	temp_json_path = temp_json.name
	temp_json.close()

	print(f"Temporary input JSON saved to: {temp_json_path}")
	return temp_json_path

	def run_inference(
	self,
	prompt: str,
	cond_image_path: str,
	cond_audio_path_spk1: str,
	cond_audio_path_spk2: Optional[str] = None,
	sample_steps: int = 12,
	) -> str:
	"""
	Run the MultiTalk inference.

	Args:
	prompt: Text prompt
	cond_image_path: Path to conditioning image
	cond_audio_path_spk1: Path to first speaker audio
	cond_audio_path_spk2: Path to second speaker audio (optional)
	sample_steps: Number of sampling steps

	Returns:
	Path to generated video
	"""
	timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
	result_filename = f"meigen_multitalk_result_{sample_steps}_steps_{timestamp}"
	temp_files_to_cleanup = []

	try:
	# Trim audio files to 5 seconds
	trimmed_audio_path_spk1 = self.trim_audio_to_5s_temp(cond_audio_path_spk1)
	if trimmed_audio_path_spk1 != cond_audio_path_spk1:
	cond_audio_path_spk1 = trimmed_audio_path_spk1
	temp_files_to_cleanup.append(trimmed_audio_path_spk1)

	if cond_audio_path_spk2 is not None:
	trimmed_audio_path_spk2 = self.trim_audio_to_5s_temp(
	cond_audio_path_spk2
	)
	if trimmed_audio_path_spk2 != cond_audio_path_spk2:
	cond_audio_path_spk2 = trimmed_audio_path_spk2
	temp_files_to_cleanup.append(trimmed_audio_path_spk2)

	# Prepare input JSON
	input_json_path = self.create_temp_input_json(
	prompt, cond_image_path, cond_audio_path_spk1, cond_audio_path_spk2
	)
	temp_files_to_cleanup.append(input_json_path)

	# Base command arguments
	common_args = [
	"--ckpt_dir",
	"weights/Wan2.1-I2V-14B-480P",
	"--wav2vec_dir",
	"weights/chinese-wav2vec2-base",
	"--input_json",
	input_json_path,
	"--sample_steps",
	str(sample_steps),
	"--mode",
	"streaming",
	"--use_teacache",
	"--save_file",
	result_filename,
	]

	# Choose command based on number of GPUs
	if self.num_gpus > 1:
	cmd = [
	"torchrun",
	f"--nproc_per_node={self.num_gpus}",
	"--standalone",
	"generate_multitalk.py",
	"--dit_fsdp",
	"--t5_fsdp",
	"--ulysses_size",
	str(self.num_gpus),
	] + common_args
	else:
	cmd = [
	"python3",
	"generate_multitalk.py",
	"--num_persistent_param_in_dit",
	str(self.used_vram_params),
	] + common_args

	# Run inference
	with open("inference.log", "w") as log_file:
	process = subprocess.Popen(
	cmd,
	stdout=subprocess.PIPE,
	stderr=subprocess.STDOUT,
	text=True,
	bufsize=1,
	)
	for line in process.stdout:
	print(line, end="")
	log_file.write(line)
	process.wait()

	if process.returncode != 0:
	raise RuntimeError("Inference failed. Check inference.log for details.")

	return f"{result_filename}.mp4"

	finally:
	# Clean up temporary files
	for f in temp_files_to_cleanup:
	try:
	if os.path.exists(f):
	os.remove(f)
	print(f"[INFO] Removed temporary file: {f}")
	except Exception as e:
	print(f"[WARNING] Could not remove {f}: {e}")

	def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
	"""
	Main inference method called by Hugging Face Inference API.

	Args:
	data: Input data dictionary containing:
	- inputs: Main input data
	- parameters: Optional parameters

	Returns:
	Dictionary containing the generated video
	"""
	try:
	# Extract inputs
	inputs = data.get("inputs", {})
	parameters = data.get("parameters", {})

	# Required inputs
	prompt = inputs.get("prompt", "")
	if not prompt:
	raise ValueError("Prompt is required")

	# Image input (base64 encoded)
	image_data = inputs.get("image")
	if not image_data:
	raise ValueError("Image is required")

	# Audio inputs (base64 encoded)
	audio_spk1_data = inputs.get("audio_spk1")
	if not audio_spk1_data:
	raise ValueError("Audio for speaker 1 is required")

	audio_spk2_data = inputs.get("audio_spk2") # Optional

	# Parameters
	sample_steps = parameters.get("sample_steps", 12)
	sample_steps = max(2, min(25, sample_steps)) # Clamp between 2 and 25

	# Save base64 files to temporary locations
	temp_files_to_cleanup = []

	# Save image
	image_path = self.save_base64_file(image_data, ".png")
	temp_files_to_cleanup.append(image_path)

	# Save audio files
	audio_spk1_path = self.save_base64_file(audio_spk1_data, ".wav")
	temp_files_to_cleanup.append(audio_spk1_path)

	audio_spk2_path = None
	if audio_spk2_data:
	audio_spk2_path = self.save_base64_file(audio_spk2_data, ".wav")
	temp_files_to_cleanup.append(audio_spk2_path)

	# Run inference
	output_video_path = self.run_inference(
	prompt=prompt,
	cond_image_path=image_path,
	cond_audio_path_spk1=audio_spk1_path,
	cond_audio_path_spk2=audio_spk2_path,
	sample_steps=sample_steps,
	)

	# Read and encode output video
	with open(output_video_path, "rb") as f:
	video_data = base64.b64encode(f.read()).decode()

	# Clean up input files
	for f in temp_files_to_cleanup:
	try:
	if os.path.exists(f):
	os.remove(f)
	except Exception as e:
	print(f"[WARNING] Could not remove {f}: {e}")

	# Clean up output video
	try:
	if os.path.exists(output_video_path):
	os.remove(output_video_path)
	except Exception as e:
	print(f"[WARNING] Could not remove output video: {e}")

	return {"generated_video": video_data, "content_type": "video/mp4"}

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