"""StreamTalk — streaming co-speech gesture generation, ZeroGPU Gradio demo. Faithful port of the authors' reference inference script (`Scripts/FM/TestFixedExpressions_keyposes.py` in Xiangyue-Zhang/StreamTalk): WavLM features -> part-aware DiT flow sampling -> body forward kinematics -> speaker-specific key-pose retrieval -> anchored refinement, window by window. The only deviations from the reference are operational: * the reference drives an SMPL-X neutral body, which may not be redistributed by third parties. Forward kinematics and rendering therefore run on NVIDIA's Apache-2.0 SOMA-X body instead (see `soma_body.py`); the motion StreamTalk produces is still SMPL-X-parameterised and is exported as such; * the speaker-2 retrieval database (joint positions + global rot6d poses for the 95 BEAT2 `2_scott_0` training clips) is precomputed offline and shipped as `assets/retrieval_db_2_scott_0.npz` instead of being rebuilt at every start; * the resulting sequence is rendered to video with PyTorch3D (the repo only ships a Blender add-on workflow). """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # noqa: E402 — must precede any CUDA-touching import import random # noqa: E402 import subprocess # noqa: E402 import sys # noqa: E402 import tempfile # noqa: E402 import time # noqa: E402 from pathlib import Path # noqa: E402 import gradio as gr # noqa: E402 import numpy as np # noqa: E402 import torch # noqa: E402 import torch.nn as nn # noqa: E402 import torch.nn.functional as F # noqa: E402 from huggingface_hub import hf_hub_download # noqa: E402 ROOT = Path(__file__).resolve().parent sys.path.insert(0, str(ROOT / "Scripts")) import librosa # noqa: E402 import soundfile as sf # noqa: E402 from transformers import Wav2Vec2FeatureExtractor, WavLMModel # noqa: E402 from Common.SMPLXConfig import GetBodyPartIndices # noqa: E402 from Common.SMPLX_Util import ( # noqa: E402 LocalToGlobalAA, global_to_local_batch, ) from Common.Util import ( # noqa: E402 NormalizeRot6d, SplitIntoOverlapedWindows, axis_angle_to_rotation_6d, matrix_to_rotation_6d, rotation_6d_to_matrix, rotation_matrix_to_axis_angle, ) from FM.Net import ( # noqa: E402 DiffusionDITNetPartsFixedExpressions2PostNormInteraction2 as MotionModel, ) from FM.Net import SimpleSpeechModel # noqa: E402 from soma_body import SomaBody # noqa: E402 # --------------------------------------------------------------------------- # # Constants # --------------------------------------------------------------------------- # WAVLM_ID = "patrickvonplaten/wavlm-libri-clean-100h-large" STREAMTALK_REPO = "X-Zhang/StreamTalk" STREAMTALK_CKPT = "streamtalk_speaker2_combined_e0946_cfg3.pt" # Body model: NVIDIA SOMA-X neutral (Apache-2.0, redistributable) — see # `soma_body.py` / `LICENSE-SOMA-X`. SOMA_RIG = "assets/soma_rig.npz" PID = "2_scott_0" FPS = 30 N_SEED = 8 MAX_SECONDS = 20.0 RENDER_SIZE = 512 CACHE_VERSION = 1 # --------------------------------------------------------------------------- # # Module-scope model loading (ZeroGPU packs these into VRAM on first GPU call) # --------------------------------------------------------------------------- # print("[startup] downloading assets ...", flush=True) CKPT_PATH = hf_hub_download(STREAMTALK_REPO, STREAMTALK_CKPT) print("[startup] building models ...", flush=True) face_model = SimpleSpeechModel().requires_grad_(False).eval() face_model.load_state_dict( torch.load( ROOT / "Scripts/FM/ckpt/split/SimpleSpeechModel/best.pt", map_location="cpu" ) ) face_model = face_model.to("cuda") net = MotionModel().requires_grad_(False).eval() net.load_state_dict(torch.load(CKPT_PATH, map_location="cpu"), strict=True) net = net.to("cuda") body_model = SomaBody(ROOT / SOMA_RIG).requires_grad_(False).eval().to("cuda") BODY_FACES = body_model.faces # Rest-pose pelvis and mesh centre, used to place/frame the body (metres). REST_PELVIS = body_model.local_t[1].clone() REST_CENTRE = ( body_model.bind_shape.amin(dim=0) + body_model.bind_shape.amax(dim=0) ) / 2.0 wavlm_processor = Wav2Vec2FeatureExtractor.from_pretrained(WAVLM_ID) wavlm = WavLMModel.from_pretrained(WAVLM_ID).requires_grad_(False).eval().to("cuda") class _RetrievalDB(nn.Module): """Speaker-specific key-pose database, held as buffers so ZeroGPU packs it.""" def __init__(self, path: Path): super().__init__() blob = np.load(path) self.register_buffer( "joints", torch.from_numpy(blob["joints"]).half(), persistent=False ) self.register_buffer( "poses", torch.from_numpy(blob["poses_rot6d"]).float(), persistent=False ) print("[startup] loading retrieval database ...", flush=True) retrieval_db = _RetrievalDB(ROOT / "assets/retrieval_db_2_scott_0.npz").eval().to("cuda") print(f"[startup] retrieval frames: {retrieval_db.joints.shape}", flush=True) _prompt_blob = np.load(ROOT / "assets/prompt_2_scott_0.npz") PROMPT_BETAS = _prompt_blob["betas"].astype(np.float32) PROMPT_POSES = _prompt_blob["poses"].astype(np.float32) PROMPT_TRANS = _prompt_blob["trans"].astype(np.float32) HAND_IDX = GetBodyPartIndices("hand", "rot6d") UPPER_IDX = GetBodyPartIndices("upper", "rot6d") LOWER_IDX = GetBodyPartIndices("lower", "rot6d") FACE_IDX = GetBodyPartIndices("face", "rot6d") TRANS_IDX = np.array([330, 331, 332]) print("[startup] ready", flush=True) # --------------------------------------------------------------------------- # # Helpers # --------------------------------------------------------------------------- # def _setup_seed(seed: int) -> None: torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) np.random.seed(seed) random.seed(seed) def _extract_wavlm(raw_audio: np.ndarray, overlap_frame_size: int = N_SEED): """Reference `AudioFeatureExtractor.extract`: 2 s windows -> 60 motion frames.""" sr = 16000 clips = SplitIntoOverlapedWindows( raw_audio, win_size=sr * 2, stride_step=sr * 2 - int(sr / FPS * overlap_frame_size), keep_last=False, ) feats = [] for clip in clips: values = wavlm_processor( clip, return_tensors="pt", sampling_rate=sr ).input_values.to("cuda") with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16): emb = wavlm(values).last_hidden_state emb = F.interpolate( emb.permute((0, 2, 1)), 60, mode="linear", align_corners=True ).permute((0, 2, 1)) feats.append(emb) return feats def _fk_joints(local_rotmats: torch.Tensor) -> torch.Tensor: """Reference `EvaluateSMPLX.__call__` — canonical-orientation joints. `local_rotmats` is (B, 55, 3, 3) of SMPL-X *local* rotations; the global orientation is zeroed exactly as in the reference so retrieval matches on pose alone. The retrieval database in `assets/` was rebuilt with the very same forward kinematics, so query and database stay consistent. """ R = local_rotmats.clone() R[:, 0] = torch.eye(3, device=R.device, dtype=R.dtype) with torch.no_grad(): return body_model.joints(R) def _nearest_pose_ids(query_joints: torch.Tensor): """L1 nearest neighbour against the speaker database, chunked over the DB.""" db = retrieval_db.joints # (N, 165) fp16 q = query_joints.half() # (L, 165) best_d = None best_i = None chunk = 20000 for start in range(0, db.shape[0], chunk): sub = db[start : start + chunk] d = torch.abs(sub.unsqueeze(0) - q.unsqueeze(1)).sum(dim=-1) # (L, chunk) d_min, i_min = torch.min(d, dim=-1) i_min = i_min + start if best_d is None: best_d, best_i = d_min, i_min else: take = d_min < best_d best_d = torch.where(take, d_min, best_d) best_i = torch.where(take, i_min, best_i) return best_d, best_i def _render_video( local_rotmats: torch.Tensor, trans: torch.Tensor, audio_path: str, out_path: str, ) -> None: """Rasterise the body sequence to an mp4 and mux the driving audio in.""" import imageio import imageio_ffmpeg from pytorch3d.renderer import ( BlendParams, FoVOrthographicCameras, HardPhongShader, MeshRasterizer, MeshRenderer, PointLights, RasterizationSettings, TexturesVertex, look_at_view_transform, ) from pytorch3d.structures import Meshes device = local_rotmats.device n_frames = local_rotmats.shape[0] # BEAT2 translations position the SMPL-X pelvis in world space; offset them # so the SOMA-X pelvis lands in the same place. transl = trans - REST_PELVIS centre = transl.mean(dim=0) + REST_CENTRE centre_x = float(centre[0]) centre_y = float(centre[1]) # Camera sits on +Z looking back at the body, matching the BEAT2/EMAGE # orthographic setup (xmag = ymag = 1). R, T = look_at_view_transform( dist=5.0, elev=0.0, azim=0.0, at=((centre_x, centre_y, 0.0),) ) cameras = FoVOrthographicCameras( device=device, R=R, T=T, znear=0.01, zfar=50.0, min_x=-1.05, max_x=1.05, min_y=-1.05, max_y=1.05, ) lights = PointLights(device=device, location=[[centre_x, centre_y + 1.5, 3.0]]) raster = RasterizationSettings( image_size=RENDER_SIZE, blur_radius=0.0, faces_per_pixel=1 ) renderer = MeshRenderer( rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster), shader=HardPhongShader( device=device, cameras=cameras, lights=lights, blend_params=BlendParams(background_color=(1.0, 1.0, 1.0)), ), ) faces = torch.from_numpy(BODY_FACES).to(device) silent_video = out_path + ".silent.mp4" writer = imageio.get_writer( silent_video, fps=FPS, codec="libx264", quality=8, macro_block_size=1, ffmpeg_log_level="error", ) chunk = 16 try: for start in range(0, n_frames, chunk): end = min(start + chunk, n_frames) with torch.no_grad(): verts = body_model.vertices( local_rotmats[start:end], transl[start:end] ) b = verts.shape[0] textures = TexturesVertex( verts_features=torch.full_like(verts, 0.82) ) meshes = Meshes( verts=verts, faces=faces.unsqueeze(0).expand(b, -1, -1), textures=textures ) images = renderer(meshes)[..., :3].clamp(0, 1) frames = (images * 255).to(torch.uint8).cpu().numpy() for frame in frames: writer.append_data(frame) finally: writer.close() ffmpeg = imageio_ffmpeg.get_ffmpeg_exe() subprocess.run( [ffmpeg, "-y", "-loglevel", "error", "-i", silent_video, "-i", audio_path, "-c:v", "copy", "-c:a", "aac", "-b:a", "128k", "-shortest", out_path], check=True, ) os.remove(silent_video) def _estimate_duration(audio_path, cfg_scale=3.0, seed=0, *args, **kwargs): try: seconds = min(MAX_SECONDS, librosa.get_duration(path=audio_path)) except Exception: seconds = MAX_SECONDS # Measured on ZeroGPU: ~10.7 s for the 20 s cap, ~7.8 s for 14 s. return int(min(60, 8 + 0.7 * seconds)) # --------------------------------------------------------------------------- # # Inference # --------------------------------------------------------------------------- # @spaces.GPU(duration=_estimate_duration) def generate( audio_path: str, cfg_scale: float = 3.0, seed: int = 0, progress=gr.Progress(track_tqdm=False), ) -> tuple: """Generate a co-speech 3D gesture animation from a speech recording. Args: audio_path: path to a speech audio file (any format ffmpeg/librosa reads). cfg_scale: classifier-free guidance scale; the released checkpoint is tuned for 3.0. 1.0 reproduces the original unguided forward path. seed: RNG seed for the flow-matching noise. Returns: A tuple of (rendered mp4 video path, SMPL-X .npz path, status text). """ if not audio_path: raise gr.Error("Please provide a speech recording first.") t_start = time.perf_counter() _setup_seed(int(seed)) audio = librosa.load(audio_path, sr=16000)[0] if audio.size < 16000 * 2: raise gr.Error("Audio is too short — StreamTalk needs at least 2 seconds.") audio = audio[: int(16000 * MAX_SECONDS)] trimmed_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name sf.write(trimmed_wav, audio, 16000) progress(0.05, desc="Extracting WavLM features") huberts = _extract_wavlm(audio) num_segments = len(huberts) if num_segments == 0: raise gr.Error("Audio is too short — StreamTalk needs at least 2 seconds.") pose_prompt = ( torch.from_numpy(axis_angle_to_rotation_6d(LocalToGlobalAA(PROMPT_POSES))) .float() .unsqueeze(0) .to("cuda") ) trans_prompt = torch.from_numpy(PROMPT_TRANS).float().unsqueeze(0).to("cuda") pid = torch.LongTensor([int(PID.split("_")[0])]).to("cuda") B, has_trans = 1, True motions = [] for n in range(num_segments): progress( 0.1 + 0.6 * n / max(1, num_segments), desc=f"Streaming window {n + 1}/{num_segments}", ) prompt_mask = torch.full((B, 60), True, device="cuda") prompt_mask[:, :N_SEED] = False pred_expressions = face_model(huberts[n], pid) pred_expressions_uncond = None if cfg_scale != 1.0: pred_expressions_uncond = face_model(huberts[n], torch.zeros_like(pid)) trans_mask = torch.full((B, 60), True, device="cuda") trans_mask[:, :N_SEED] = False def _kwargs(pmask): kw = { "pid": pid, "poses_prompt": pose_prompt, "prompt_mask": pmask, "trans_prompt": trans_prompt, "trans_mask": trans_mask, "hand_indices": HAND_IDX, "upper_indices": UPPER_IDX, "lower_indices": LOWER_IDX, "face_joint_indices": FACE_IDX, "trans_indices": TRANS_IDX, "hubert": huberts[n], "expressions": pred_expressions, } if pred_expressions_uncond is not None: kw["expressions_uncond"] = pred_expressions_uncond return kw noise = torch.randn((B, 60, 55 * 6 + 3 * has_trans), device="cuda") with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16): xt_list = net.sample( noise, normalize_func=NormalizeRot6d, cfg_scale=cfg_scale, **_kwargs(prompt_mask), ) xt = xt_list[-1] xt[..., :330] = NormalizeRot6d(xt[..., :330]) local_poses = rotation_6d_to_matrix(xt[..., :330].reshape(-1, 55, 6)) local_poses = global_to_local_batch(local_poses) joints = _fk_joints(local_poses).reshape((60, -1)) dists, nearest_ids = _nearest_pose_ids(joints) dists = dists[-N_SEED:] nearest_ids = nearest_ids[-N_SEED:] pick = torch.argmin(dists) nearest_poses = retrieval_db.poses[nearest_ids[pick]].reshape(1, 55, 6) nearest_poses = rotation_6d_to_matrix(nearest_poses) global_orientation = rotation_6d_to_matrix(xt[0, 52 + pick, :6].reshape(1, 6)) nearest_poses = torch.einsum( "bij,bnjk->bnik", global_orientation, nearest_poses ) nearest_poses = matrix_to_rotation_6d(nearest_poses).reshape(1, 55 * 6) prompt_mask = torch.full((B, 60), True, device="cuda") prompt_mask[:, :N_SEED] = False prompt_mask[:, 52 + pick] = False pose_prompt[:, 52 + pick] = nearest_poses xt[:, 52 + pick, :330] = nearest_poses with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16): xt = net.sample_with_kp_bk( noise, xt_list[5], 5, normalize_func=NormalizeRot6d, cfg_scale=cfg_scale, **_kwargs(prompt_mask), ) xt[..., :330] = NormalizeRot6d(xt[..., :330]) pose_prompt[:, :N_SEED] = xt[:, -N_SEED:][..., :330] if has_trans: trans_prompt[:, :N_SEED] = xt[:, -N_SEED:][..., 330:333] motions.append(torch.cat([xt, pred_expressions], dim=-1)) for i in range(1, len(motions)): motions[i - 1][:, -N_SEED:] = ( motions[i - 1][:, -N_SEED:] + motions[i][:, :N_SEED] ) / 2.0 motions[i] = motions[i][:, N_SEED:] sample_rot6d = torch.cat(motions, dim=1) pred_expressions = sample_rot6d[..., -100:] pred_rotmats = rotation_6d_to_matrix(sample_rot6d[..., :330].reshape(-1, 55, 6)) pred_rotmats = global_to_local_batch(pred_rotmats) pred_poses = rotation_matrix_to_axis_angle(pred_rotmats).reshape(B, -1, 55 * 3) pred_trans = sample_rot6d[..., 330 : 330 + 3] gen_seconds = time.perf_counter() - t_start n_frames = pred_poses.shape[1] npz_path = tempfile.NamedTemporaryFile(suffix=".npz", delete=False).name np.savez( npz_path, betas=PROMPT_BETAS, poses=pred_poses[0].float().cpu().numpy(), expressions=pred_expressions[0].float().cpu().numpy(), trans=pred_trans[0].float().cpu().numpy(), gender=np.array("neutral"), mocap_frame_rate=np.array(FPS), model="smplx2020", ) progress(0.75, desc="Rendering body mesh") video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name _render_video( pred_rotmats.float(), pred_trans[0].float(), trimmed_wav, video_path, ) total = time.perf_counter() - t_start status = ( f"{n_frames} frames @ {FPS} fps ({n_frames / FPS:.1f} s) from " f"{num_segments} streaming windows · motion {gen_seconds:.1f} s " f"({n_frames / max(gen_seconds, 1e-6):.0f} fps) · total {total:.1f} s" ) print("[timing] " + status, flush=True) return video_path, npz_path, status # --------------------------------------------------------------------------- # # UI # --------------------------------------------------------------------------- # CSS = """ #col-container { max-width: 1080px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ DESCRIPTION = """ # 🗣️ StreamTalk — Streaming Co-Speech Gesture Generation Speech in, full-body **SMPL-X gesture animation** out. StreamTalk generates motion one 2-second window at a time, retrieves a plausible destination pose from a speaker-specific motion database, and refines the window toward that anchor before it becomes context for the next one — which is what keeps long sequences from drifting. Weights: [`X-Zhang/StreamTalk`](https://huggingface.co/X-Zhang/StreamTalk) (speaker-2 / Scott checkpoint, CFG 3) · [paper](https://huggingface.co/papers/2608.01643) · [code](https://github.com/Xiangyue-Zhang/StreamTalk) · [project page](https://xiangyuezhang.com/StreamTalk/) """ with gr.Blocks(title="StreamTalk") as demo: with gr.Column(elem_id="col-container"): gr.Markdown(DESCRIPTION) with gr.Row(): with gr.Column(scale=1): audio_in = gr.Audio( label="Speech", sources=["upload", "microphone"], type="filepath", ) run = gr.Button("Generate gestures", variant="primary") with gr.Accordion("Advanced settings", open=False): cfg = gr.Slider( label="CFG scale", minimum=1.0, maximum=7.0, step=0.5, value=3.0, info="The released checkpoint is tuned for 3.0; 1.0 is the " "original unguided path.", ) seed = gr.Number(label="Seed", value=0, precision=0) gr.Markdown( f"Audio longer than {int(MAX_SECONDS)} s is trimmed. " "Motion is generated for BEAT2 speaker 2 (Scott)." ) with gr.Column(scale=1): video_out = gr.Video(label="Gesture animation", autoplay=True) status_out = gr.Textbox(label="Run info", lines=2) npz_out = gr.File(label="SMPL-X sequence (.npz, 30 fps)") gr.Examples( examples=[ ["examples/scott_quickstart.wav"], ["examples/scott_beat2_test.wav"], ], inputs=[audio_in], outputs=[video_out, npz_out, status_out], fn=generate, cache_examples=True, cache_mode="lazy", ) gr.Markdown( "Example audio: BEAT2 speaker 2 (Scott) — the clip bundled with the " "StreamTalk repository and one BEAT2 test clip, from " "[H-Liu1997/BEAT2](https://huggingface.co/datasets/H-Liu1997/BEAT2)." ) run.click( fn=generate, inputs=[audio_in, cfg, seed], outputs=[video_out, npz_out, status_out], api_name="generate", ) if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)