app.py

"""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)