app.py

import torch
import os
import shutil
import subprocess
import gradio as gr 
import json
import tempfile
from huggingface_hub import snapshot_download

import soundfile as sf
import tempfile
from datetime import datetime

is_shared_ui = True if "fffiloni/Meigen-MultiTalk" in os.environ['SPACE_ID'] else False
is_gpu_associated = torch.cuda.is_available()

def trim_audio_to_5s_temp(audio_path, sample_rate=16000):
    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

if is_gpu_associated:
    num_gpus = torch.cuda.device_count()
    print(f"GPU AVAILABLE: {num_gpus}")


if not is_shared_ui and is_gpu_associated:
    # Download All Required Models using `snapshot_download`

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

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

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

    # Define paths
    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.")

if not is_shared_ui:

    # Check if CUDA-compatible GPU is available
    if torch.cuda.is_available():
        # Get current GPU name
        gpu_name = torch.cuda.get_device_name(torch.cuda.current_device())
        print(f"Current GPU: {gpu_name}")

        # Enforce GPU requirement
        if "A100" not in gpu_name and "L4" not in gpu_name:
           #raise RuntimeError(f"This notebook requires an A100 or L4 GPU. Found: {gpu_name}")
            print(f"This notebook requires an A100 or L4 GPU. Found: {gpu_name}")
        elif "L4" in gpu_name:
            print("Warning: L4 or L40S is supported, but A100 is recommended for faster inference.")
    else:
        #raise RuntimeError("No CUDA-compatible GPU found. An A100, L4 or L40S GPU is required.")
        print("No CUDA-compatible GPU found. An A100, L4 or L40S GPU is required.")


    GPU_TO_VRAM_PARAMS = {
        "NVIDIA A100": 11000000000,
        "NVIDIA A100-SXM4-40GB": 11000000000,
        "NVIDIA A100-SXM4-80GB": 22000000000,
        "NVIDIA L4": 5000000000,
        "NVIDIA L40S": 11000000000
    }
    USED_VRAM_PARAMS = GPU_TO_VRAM_PARAMS[gpu_name]
    print("Using", USED_VRAM_PARAMS, "for num_persistent_param_in_dit")



def create_temp_input_json(prompt: str, cond_image_path: str, cond_audio_path_spk1: str, cond_audio_path_spk2: str) -> str:
    """
    Create a temporary JSON file with the user-provided prompt, image, and audio paths.
    Returns the path to the temporary JSON file.
    """
    # Structure based on your original JSON format
    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
            }
        }

    # Create a temp file
    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 infer(prompt, cond_image_path, cond_audio_path_spk1, cond_audio_path_spk2, sample_steps):

    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 = []
    
    if is_shared_ui:
        trimmed_audio_path_spk1 = 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 = 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 = 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 args
    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
    ]

    if num_gpus > 1:
        cmd = [
            "torchrun",
            f"--nproc_per_node={num_gpus}",
            "--standalone",
            "generate_multitalk.py",
            #"--num_persistent_param_in_dit", "22000000000", # On 4xL40S
            "--dit_fsdp", "--t5_fsdp",
            "--ulysses_size", str(num_gpus),
        ] + common_args
    else:
        cmd = [
            "python3", 
            "generate_multitalk.py",
            "--num_persistent_param_in_dit", str(USED_VRAM_PARAMS),
        ] + common_args

    try:
        # Log to file and stream
        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:
        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 load_prerendered_examples(prompt, cond_image_path, cond_audio_path_spk1, cond_audio_path_spk2, sample_steps):
    output_video = None
    
    if cond_image_path == "examples/single/single1.png":
        output_video = "examples/results/multitalk_single_example_1.mp4"
    elif cond_image_path == "examples/multi/3/multi3.png":
        output_video = "examples/results/multitalk_multi_example_2.mp4"

    return output_video

css = """
div#warning-duplicate {
    background-color: #ebf5ff;
    padding: 0 16px 16px;
    margin: 0px 0;
    color: #030303!important;
}
div#warning-duplicate > .gr-prose > h2, div#warning-duplicate > .gr-prose > p {
    color: #0f4592!important;
}
div#warning-duplicate strong {
    color: #0f4592;
}
p.actions {
    display: flex;
    align-items: center;
    margin: 20px 0;
}
div#warning-duplicate .actions a {
    display: inline-block;
    margin-right: 10px;
}
div#warning-setgpu {
    background-color: #fff4eb;
    padding: 0 16px 16px;
    margin: 0px 0;
    color: #030303!important;
}
div#warning-setgpu > .gr-prose > h2, div#warning-setgpu > .gr-prose > p {
    color: #92220f!important;
}
div#warning-setgpu a, div#warning-setgpu b {
    color: #91230f;
}
div#warning-setgpu p.actions > a {
    display: inline-block;
    background: #1f1f23;
    border-radius: 40px;
    padding: 6px 24px;
    color: antiquewhite;
    text-decoration: none;
    font-weight: 600;
    font-size: 1.2em;
}
div#warning-ready {
    background-color: #ecfdf5;
    padding: 0 16px 16px;
    margin: 0px 0;
    color: #030303!important;
}
div#warning-ready > .gr-prose > h2, div#warning-ready > .gr-prose > p {
    color: #057857!important;
}
.custom-color {
    color: #030303 !important;
}
"""

with gr.Blocks(title="MultiTalk Inference", css=css) as demo:
    gr.Markdown("## 🎤 Meigen MultiTalk Inference Demo")
    gr.Markdown("Let Them Talk: Audio-Driven Multi-Person Conversational Video Generation")
    if is_shared_ui:
        gr.Markdown("Audio will be trimmed to max 5 seconds on fffiloni's shared UI. Sample steps are limited to 12. Gradio queue size is set to 4. Generating a 5 seconds video will take approximatively 20 minutes. Duplicate to skip the queue and work with longer audio inference. ")
    gr.HTML("""
    <div style="display:flex;column-gap:4px;">
        <a href="https://github.com/MeiGen-AI/MultiTalk">
            <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
        </a>
        <a href='https://meigen-ai.github.io/multi-talk/'><img src='https://img.shields.io/badge/Project-Page-blue'></a>
        <a href='https://huggingface.co/MeiGen-AI/MeiGen-MultiTalk'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Model-yellow'></a>
        <a href='https://arxiv.org/abs/2505.22647'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>
        <a href="https://huggingface.co/spaces/fffiloni/Meigen-MultiTalk?duplicate=true">
            <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
        </a>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Text Prompt",
                placeholder="Describe the scene...",
            )

            image_input = gr.Image(
                type="filepath",
                label="Conditioning Image"
            )

            audio_input_spk1 = gr.Audio(
                type="filepath",
                label="Conditioning Audio for speaker 1(.wav)"
            )

            audio_input_spk2 = gr.Audio(
                type="filepath",
                label="Conditioning Audio for speaker 2(.wav) (Optional)"
            )

            with gr.Accordion("Advanced settings", open=False):
                sample_steps = gr.Slider(
                    label="sample steps",
                    value=12,
                    minimum=2,
                    maximum=25,
                    step=1,
                    interactive=False if is_shared_ui else True
                )

            submit_btn = gr.Button("Generate", interactive=False if is_shared_ui else True)

        with gr.Column(scale=3):
            if is_shared_ui:
                 top_description = gr.HTML(f'''
                    <div class="gr-prose">
                        <h2 class="custom-color"><svg xmlns="http://www.w3.org/2000/svg" width="18px" height="18px" style="margin-right: 0px;display: inline-block;"fill="none"><path fill="#fff" d="M7 13.2a6.3 6.3 0 0 0 4.4-10.7A6.3 6.3 0 0 0 .6 6.9 6.3 6.3 0 0 0 7 13.2Z"/><path fill="#fff" fill-rule="evenodd" d="M7 0a6.9 6.9 0 0 1 4.8 11.8A6.9 6.9 0 0 1 0 7 6.9 6.9 0 0 1 7 0Zm0 0v.7V0ZM0 7h.6H0Zm7 6.8v-.6.6ZM13.7 7h-.6.6ZM9.1 1.7c-.7-.3-1.4-.4-2.2-.4a5.6 5.6 0 0 0-4 1.6 5.6 5.6 0 0 0-1.6 4 5.6 5.6 0 0 0 1.6 4 5.6 5.6 0 0 0 4 1.7 5.6 5.6 0 0 0 4-1.7 5.6 5.6 0 0 0 1.7-4 5.6 5.6 0 0 0-1.7-4c-.5-.5-1.1-.9-1.8-1.2Z" clip-rule="evenodd"/><path fill="#000" fill-rule="evenodd" d="M7 2.9a.8.8 0 1 1 0 1.5A.8.8 0 0 1 7 3ZM5.8 5.7c0-.4.3-.6.6-.6h.7c.3 0 .6.2.6.6v3.7h.5a.6.6 0 0 1 0 1.3H6a.6.6 0 0 1 0-1.3h.4v-3a.6.6 0 0 1-.6-.7Z" clip-rule="evenodd"/></svg>
                        Attention: this Space need to be duplicated to work</h2>
                        <p class="main-message custom-color">
                            To make it work, <strong>duplicate the Space</strong> and run it on your own profile using a <strong>private</strong> GPU (L40s recommended).<br />
                            A L40s costs <strong>US$1.80/h</strong>. 
                        </p>
                        <p class="actions custom-color">
                            <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}?duplicate=true">
                                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg-dark.svg" alt="Duplicate this Space" />
                            </a>
                            to start experimenting with this demo
                        </p>
                    </div>
                ''', elem_id="warning-duplicate")
            else:
                if(is_gpu_associated):
                    top_description = gr.HTML(f'''
                        <div class="gr-prose">
                            <h2 class="custom-color"><svg xmlns="http://www.w3.org/2000/svg" width="18px" height="18px" style="margin-right: 0px;display: inline-block;"fill="none"><path fill="#fff" d="M7 13.2a6.3 6.3 0 0 0 4.4-10.7A6.3 6.3 0 0 0 .6 6.9 6.3 6.3 0 0 0 7 13.2Z"/><path fill="#fff" fill-rule="evenodd" d="M7 0a6.9 6.9 0 0 1 4.8 11.8A6.9 6.9 0 0 1 0 7 6.9 6.9 0 0 1 7 0Zm0 0v.7V0ZM0 7h.6H0Zm7 6.8v-.6.6ZM13.7 7h-.6.6ZM9.1 1.7c-.7-.3-1.4-.4-2.2-.4a5.6 5.6 0 0 0-4 1.6 5.6 5.6 0 0 0-1.6 4 5.6 5.6 0 0 0 1.6 4 5.6 5.6 0 0 0 4 1.7 5.6 5.6 0 0 0 4-1.7 5.6 5.6 0 0 0 1.7-4 5.6 5.6 0 0 0-1.7-4c-.5-.5-1.1-.9-1.8-1.2Z" clip-rule="evenodd"/><path fill="#000" fill-rule="evenodd" d="M7 2.9a.8.8 0 1 1 0 1.5A.8.8 0 0 1 7 3ZM5.8 5.7c0-.4.3-.6.6-.6h.7c.3 0 .6.2.6.6v3.7h.5a.6.6 0 0 1 0 1.3H6a.6.6 0 0 1 0-1.3h.4v-3a.6.6 0 0 1-.6-.7Z" clip-rule="evenodd"/></svg>
                            You have successfully associated a GPU to this Space 🎉</h2>
                            <p class="custom-color">
                                You will be billed by the minute from when you activated the GPU until when it is turned off.
                            </p> 
                        </div>
                    ''', elem_id="warning-ready")
                else:
                    top_description = gr.HTML(f'''
                        <div class="gr-prose">
                            <h2 class="custom-color"><svg xmlns="http://www.w3.org/2000/svg" width="18px" height="18px" style="margin-right: 0px;display: inline-block;"fill="none"><path fill="#fff" d="M7 13.2a6.3 6.3 0 0 0 4.4-10.7A6.3 6.3 0 0 0 .6 6.9 6.3 6.3 0 0 0 7 13.2Z"/><path fill="#fff" fill-rule="evenodd" d="M7 0a6.9 6.9 0 0 1 4.8 11.8A6.9 6.9 0 0 1 0 7 6.9 6.9 0 0 1 7 0Zm0 0v.7V0ZM0 7h.6H0Zm7 6.8v-.6.6ZM13.7 7h-.6.6ZM9.1 1.7c-.7-.3-1.4-.4-2.2-.4a5.6 5.6 0 0 0-4 1.6 5.6 5.6 0 0 0-1.6 4 5.6 5.6 0 0 0 1.6 4 5.6 5.6 0 0 0 4 1.7 5.6 5.6 0 0 0 4-1.7 5.6 5.6 0 0 0 1.7-4 5.6 5.6 0 0 0-1.7-4c-.5-.5-1.1-.9-1.8-1.2Z" clip-rule="evenodd"/><path fill="#000" fill-rule="evenodd" d="M7 2.9a.8.8 0 1 1 0 1.5A.8.8 0 0 1 7 3ZM5.8 5.7c0-.4.3-.6.6-.6h.7c.3 0 .6.2.6.6v3.7h.5a.6.6 0 0 1 0 1.3H6a.6.6 0 0 1 0-1.3h.4v-3a.6.6 0 0 1-.6-.7Z" clip-rule="evenodd"/></svg>
                            You have successfully duplicated the MimicMotion Space 🎉</h2>
                            <p class="custom-color">There's only one step left before you can properly play with this demo: <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings" style="text-decoration: underline" target="_blank">attribute a GPU</b> to it (via the Settings tab)</a> and run the app below.
                            You will be billed by the minute from when you activate the GPU until when it is turned off.</p> 
                            <p class="actions custom-color">
                                <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings">🔥 &nbsp; Set recommended GPU</a>
                            </p>
                        </div>
                    ''', elem_id="warning-setgpu")
            output_video = gr.Video(label="Generated Video", interactive=False)

            gr.Examples(
                examples = [
                    ["A woman sings passionately in a dimly lit studio.", "examples/single/single1.png", "examples/single/1.wav", None, 12, "examples/results/multitalk_single_example_1.mp4"],
                    ["In a cozy recording studio, a man and a woman are singing together. The man, with tousled brown hair, stands to the left, wearing a light green button-down shirt. His gaze is directed towards the woman, who is smiling warmly. She, with wavy dark hair, is dressed in a black floral dress and stands to the right, her eyes closed in enjoyment. Between them is a professional microphone, capturing their harmonious voices. The background features wooden panels and various audio equipment, creating an intimate and focused atmosphere. The lighting is soft and warm, highlighting their expressions and the intimate setting. A medium shot captures their interaction closely.", "examples/multi/3/multi3.png", "examples/multi/3/1-man.WAV", "examples/multi/3/1-woman.WAV", 12, "examples/results/multitalk_multi_example_2.mp4"],
                ],
                inputs = [prompt_input, image_input, audio_input_spk1, audio_input_spk2, sample_steps, output_video],
            )

    submit_btn.click(
        fn=infer,
        inputs=[prompt_input, image_input, audio_input_spk1, audio_input_spk2, sample_steps],
        outputs=output_video
    )

demo.queue(max_size=4).launch(ssr_mode=False, show_error=True, show_api=False)