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("""
""")
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'''
To make it work, duplicate the Space and run it on your own profile using a private GPU (L40s recommended).US$1.80/h .
to start experimenting with this demo
''', elem_id="warning-duplicate")
else:
if(is_gpu_associated):
top_description = gr.HTML(f'''
You will be billed by the minute from when you activated the GPU until when it is turned off.
''', elem_id="warning-ready")
else:
top_description = gr.HTML(f'''
''', 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) 
to start experimenting with this demo