"""InfiniteTalk on ZeroGPU: portrait image + speech audio -> lip-synced talking-head video. Audio-driven, so the spoken language does not matter (Russian, Kazakh, ... all work). Uses the fp8-quantized T5 and a quantized DiT so the whole pipeline fits a 48 GB ZeroGPU slot without ever materializing the 65 GB fp32 base checkpoint. """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # noqa: E402 must precede torch import logging import math import tempfile import time from pathlib import Path from types import SimpleNamespace import gradio as gr import librosa import numpy as np import pyloudnorm as pyln import soundfile as sf import torch from einops import rearrange from huggingface_hub import hf_hub_download, snapshot_download from transformers import Wav2Vec2FeatureExtractor import wan from src.audio_analysis.wav2vec2 import Wav2Vec2Model from wan.configs import WAN_CONFIGS from wan.utils.multitalk_utils import save_video_ffmpeg logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(levelname)s: %(message)s") log = logging.getLogger("infinitetalk-space") # --- configuration ------------------------------------------------------------------------- # DiT variant from MeiGen-AI/InfiniteTalk/quant_models. "single_int8_lora" ships with an # acceleration LoRA fused in (8 steps, cfg 1.0/2.0, shift 2); "single_fp8" is the plain # model (40 steps, cfg 5.0/4.0, shift 7). Switch via the Space's env settings. DIT_VARIANT = os.environ.get("DIT_VARIANT", "single_int8_lora") FUSED_LORA = DIT_VARIANT.endswith("_lora") MAX_AUDIO_SEC = float(os.environ.get("MAX_AUDIO_SEC", "20")) FRAME_NUM, MOTION_FRAME, FPS = 81, 9, 25 DEFAULTS = ( dict(steps=8, text_cfg=1.0, audio_cfg=2.0, shift=2.0) if FUSED_LORA else dict(steps=40, text_cfg=5.0, audio_cfg=4.0, shift=7.0) ) WEIGHTS = Path("weights") # --- weights --------------------------------------------------------------------------------- t0 = time.perf_counter() base_dir = snapshot_download( "Wan-AI/Wan2.1-I2V-14B-480P", local_dir=WEIGHTS / "Wan2.1-I2V-14B-480P", allow_patterns=[ "config.json", "Wan2.1_VAE.pth", "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", "google/**", "xlm-roberta-large/**", ], ) it_dir = snapshot_download( "MeiGen-AI/InfiniteTalk", local_dir=WEIGHTS / "InfiniteTalk", allow_patterns=[ f"quant_models/infinitetalk_{DIT_VARIANT}.safetensors", f"quant_models/infinitetalk_{DIT_VARIANT}.json", "quant_models/t5_fp8.safetensors", "quant_models/t5_map_fp8.json", ], ) wav2vec_dir = snapshot_download( "TencentGameMate/chinese-wav2vec2-base", local_dir=WEIGHTS / "chinese-wav2vec2-base" ) hf_hub_download( "TencentGameMate/chinese-wav2vec2-base", "model.safetensors", revision="refs/pr/1", local_dir=WEIGHTS / "chinese-wav2vec2-base", ) log.info("weights ready in %.0fs", time.perf_counter() - t0) # --- models (module scope, moved to cuda eagerly so ZeroGPU can pack them) ----------------- t0 = time.perf_counter() feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec_dir, local_files_only=True) audio_encoder = Wav2Vec2Model.from_pretrained(wav2vec_dir, local_files_only=True).eval() audio_encoder.feature_extractor._freeze_parameters() cfg = WAN_CONFIGS["infinitetalk-14B"] pipe = wan.InfiniteTalkPipeline( config=cfg, checkpoint_dir=base_dir, quant_dir=os.path.join(it_dir, "quant_models", f"infinitetalk_{DIT_VARIANT}.safetensors"), device_id=0, rank=0, t5_cpu=False, quant="fp8", # selects t5_fp8; the DiT file above carries its own quantization map infinitetalk_dir=None, ) pipe.model.to("cuda") pipe.text_encoder.model.to("cuda") log.info("models built in %.0fs", time.perf_counter() - t0) # Only the flags generate_infinitetalk() reads; TeaCache/APG stay off. EXTRA = SimpleNamespace( use_teacache=False, teacache_thresh=0.2, size="infinitetalk-480", use_apg=False, apg_momentum=-0.75, apg_norm_threshold=55, ) # --- helpers --------------------------------------------------------------------------------- def prepare_audio(path: str, sr: int = 16000) -> np.ndarray: """Load any audio/video file as 16 kHz mono, loudness-normalized, trimmed to MAX_AUDIO_SEC.""" audio, _ = librosa.load(path, sr=sr) audio = audio[: int(MAX_AUDIO_SEC * sr)] meter = pyln.Meter(sr) loudness = meter.integrated_loudness(audio) if abs(loudness) < 100: audio = pyln.normalize.loudness(audio, loudness, -23) return audio def audio_embedding(audio: np.ndarray, sr: int = 16000) -> torch.Tensor: """wav2vec2 hidden states resampled to 25 fps, as the DiT expects (S, layers, 768).""" frames = int(len(audio) / sr * FPS) feats = np.squeeze(feature_extractor(audio, sampling_rate=sr).input_values) feats = torch.from_numpy(feats).float().unsqueeze(0) with torch.no_grad(): out = audio_encoder(feats, seq_len=frames, output_hidden_states=True) emb = torch.stack(out.hidden_states[1:], dim=1).squeeze(0) return rearrange(emb, "b s d -> s b d").cpu() def n_chunks(audio_sec: float) -> int: frames = int(audio_sec * FPS) return 1 + max(0, math.ceil((frames - FRAME_NUM) / (FRAME_NUM - MOTION_FRAME))) def estimate_duration(image, audio, steps, *args, **kwargs) -> int: try: sec = min(librosa.get_duration(path=audio), MAX_AUDIO_SEC) except Exception: sec = MAX_AUDIO_SEC per_chunk = 20 + int(steps) * 4 # seconds; refined from measured runs return min(600, 45 + n_chunks(sec) * per_chunk) # --- inference ------------------------------------------------------------------------------- @spaces.GPU(duration=estimate_duration) def generate( image: str, audio: str, steps: int = DEFAULTS["steps"], audio_cfg: float = DEFAULTS["audio_cfg"], text_cfg: float = DEFAULTS["text_cfg"], prompt: str = "A person is talking to the camera.", seed: int = -1, progress=gr.Progress(track_tqdm=True), ) -> str: """Animate a portrait so it speaks the given audio (any language) with synced lips. Args: image: Path to a front-facing portrait (PNG/JPG). audio: Path to speech audio (wav/mp3/…); trimmed to MAX_AUDIO_SEC. steps: Diffusion steps per chunk. audio_cfg: Audio guidance scale (higher = tighter lip sync). text_cfg: Text guidance scale. prompt: Short scene description. seed: RNG seed, -1 for random. Returns: Path to the generated mp4 with the audio muxed in. """ if image is None or audio is None: raise gr.Error("Please provide both a portrait image and an audio file.") t_start = time.perf_counter() work = Path(tempfile.mkdtemp(prefix="infinitetalk_")) speech = prepare_audio(audio) wav_path = work / "speech.wav" sf.write(wav_path, speech, 16000) emb_path = work / "speech.pt" torch.save(audio_embedding(speech), emb_path) if seed is None or int(seed) < 0: seed = int(torch.randint(0, 2**31 - 1, (1,)).item()) input_data = { "prompt": prompt or "", "cond_video": image, "cond_audio": {"person1": str(emb_path)}, "video_audio": str(wav_path), } video = pipe.generate_infinitetalk( input_data, size_buckget="infinitetalk-480", motion_frame=MOTION_FRAME, frame_num=FRAME_NUM, shift=DEFAULTS["shift"], sampling_steps=int(steps), text_guide_scale=float(text_cfg), audio_guide_scale=float(audio_cfg), seed=int(seed), offload_model=False, max_frames_num=1000, color_correction_strength=1.0, extra_args=EXTRA, ) out_stem = str(work / "result") save_video_ffmpeg(video, out_stem, [str(wav_path)], high_quality_save=False) log.info( "generated %.1fs of video (%d chunks, %d steps) in %.0fs", len(speech) / 16000, n_chunks(len(speech) / 16000), int(steps), time.perf_counter() - t_start, ) return out_stem + ".mp4" # --- UI -------------------------------------------------------------------------------------- with gr.Blocks(title="InfiniteTalk — talking-head from image + audio") as demo: gr.Markdown( f""" # 🗣️ InfiniteTalk — image + audio → lip-synced video Upload a portrait and speech audio in **any language** (the model is audio-driven, so Russian, Kazakh, English… all work). Audio is trimmed to **{MAX_AUDIO_SEC:.0f} s**; output is 480p @ 25 fps. Model: [MeiGen-AI/InfiniteTalk](https://huggingface.co/MeiGen-AI/InfiniteTalk) · variant `{DIT_VARIANT}`. """ ) with gr.Row(): with gr.Column(): image_in = gr.Image(type="filepath", label="Portrait image") audio_in = gr.Audio(type="filepath", label="Speech audio") with gr.Accordion("Advanced", open=False): steps_in = gr.Slider(4, 40, value=DEFAULTS["steps"], step=1, label="Diffusion steps") audio_cfg_in = gr.Slider(1.0, 6.0, value=DEFAULTS["audio_cfg"], step=0.5, label="Audio guidance") text_cfg_in = gr.Slider(1.0, 6.0, value=DEFAULTS["text_cfg"], step=0.5, label="Text guidance") prompt_in = gr.Textbox(value="A person is talking to the camera.", label="Prompt") seed_in = gr.Number(value=-1, precision=0, label="Seed (-1 = random)") run_btn = gr.Button("Generate", variant="primary") with gr.Column(): video_out = gr.Video(label="Result") run_btn.click( fn=generate, inputs=[image_in, audio_in, steps_in, audio_cfg_in, text_cfg_in, prompt_in, seed_in], outputs=video_out, api_name="generate", ) gr.Examples( examples=[["examples/ref_image.png", "examples/1.wav"]], inputs=[image_in, audio_in], outputs=video_out, fn=generate, cache_examples=True, cache_mode="lazy", ) demo.launch(mcp_server=True)