0.4_test_inference.py

"""Real-time streaming test: video -> predicted landmarks -> VTube Studio.

Streams chunk-by-chunk. For each STREAM_CHUNK frames:
  1. decode + feature-extract that chunk's audio (mel + FCPE F0 + RMS)
  2. one NAT forward over the chunk
  3. yield frames to vts_driver_stream, which sends them paced at 12 fps
The driver plays chunk N-1 over ~STREAM_CHUNK/12 s while this script computes
chunk N. As long as per-chunk compute < playback time (it is, by a wide
margin: a 5.3 s chunk needs <1 s to extract+predict), playback is continuous.

Causal model + causal smoothing = true live latency: frame t needs only
audio[<=t] within the current 64-frame window.

Plays the video's audio via ffplay in parallel, synced approximately with
the VTS animation. Waits for you to accept the VTS connection, then starts
both audio and landmark injection at the same offset.
"""
import os
import argparse
import asyncio
import subprocess
import shutil
import importlib.util

import numpy as np
import torch
import librosa
import cv2
from torchfcpe import spawn_bundled_infer_model

_spec_vts = importlib.util.spec_from_file_location("vts_playback_driver", "vts_playback_driver.py")
_vts = importlib.util.module_from_spec(_spec_vts); _spec_vts.loader.exec_module(_vts)
vts_driver_stream = _vts.vts_driver_stream
_STREAM_END = _vts._STREAM_END

_spec_train = importlib.util.spec_from_file_location("train_nat", "0.3_train_nat.py")
_train = importlib.util.module_from_spec(_spec_train); _spec_train.loader.exec_module(_train)
AudioToLandmarkNAT = _train.AudioToLandmarkNAT

SR = 16000
HOP = 1333            # 16000/12 -> one mel frame per 12fps frame
USE_BF16 = torch.cuda.is_available()
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_KW = dict(input_dim=82, d_model=192, nhead=3, num_layers=8, output_dim=12, causal=True)
SEQ_LEN = 64          # MUST match training (positional embeddings learned at this len)
STREAM_CHUNK = SEQ_LEN   # one model forward per chunk; chunk == window -> no padding
# ponytail: bump STREAM_CHUNK to a multiple of SEQ_LEN (e.g. 128 = 2 windows/chunk)
# if you want fewer ffmpeg seeks at the cost of slightly coarser chunking.

MAX_BATCH = 4
AUDIO_DIR = "dataset/audio_features"   # for loading cached audio norm stats
# Standardize the 82-dim features with the SAME mean/std used in training
# (see 0.3 compute_audio_norm). Without this F0 (std ~162) swamps the first
# layer and inference output drifts from what the model was trained on.
AUDIO_MEAN, AUDIO_STD = _train.compute_audio_norm(AUDIO_DIR)

# PCA output reconstruction: model predicts 12 blendshape coefficients;
# reconstruct to 56-dim landmarks via a single matrix multiply.
_pca = np.load("dataset/pca.npz")
PCA_MEAN = torch.from_numpy(_pca["mean"].astype(np.float32))
PCA_COMP = torch.from_numpy(_pca["components"].astype(np.float32))  # [12, 56]


def parse_start(s):
    h, m, sec = map(int, s.split(":"))
    return h * 3600 + m * 60 + sec


def features_for_segment(y_seg, V, fcpe_model):
    """Build [V, 82] features from a pre-decoded audio segment (mel+FCPE F0+RMS)."""
    mel = librosa.feature.melspectrogram(y=y_seg, sr=SR, n_fft=2048, hop_length=HOP, n_mels=80)
    mel_db = librosa.power_to_db(mel, ref=np.max).T
    if mel_db.shape[0] > V:
        mel_db = mel_db[:V]
    elif mel_db.shape[0] < V:
        mel_db = np.pad(mel_db, ((0, V - mel_db.shape[0]), (0, 0)), mode="edge")

    # FCPE: for segments <= 30s run directly (avoids 8x zero-padding for
    # 5s streaming chunks). For longer segments, chunk at 60s to avoid OOM.
    seg_sec = len(y_seg) / SR
    if seg_sec <= 30:
        b = torch.from_numpy(y_seg.astype(np.float32)).unsqueeze(0).unsqueeze(-1).to(DEVICE)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=USE_BF16):
            f0 = fcpe_model.infer(b, sr=SR, decoder_mode="local_argmax", threshold=0.006,
                                  f0_min=80, f0_max=880, interp_uv=False,
                                  output_interp_target_length=V)
        f0_np = f0.reshape(-1, 1).cpu().numpy()[:V]
    else:
        chunk_samples = 60 * SR
        frames_per_chunk = 60 * 12
        y_pad = np.pad(y_seg, (0, (-len(y_seg)) % chunk_samples), mode="constant")
        chunks = y_pad.reshape(len(y_pad) // chunk_samples, chunk_samples).astype(np.float32)
        f0_parts = []
        for i in range(0, chunks.shape[0], MAX_BATCH):
            b = torch.from_numpy(chunks[i:i + MAX_BATCH]).to(DEVICE).unsqueeze(-1)
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=USE_BF16):
                f0 = fcpe_model.infer(b, sr=SR, decoder_mode="local_argmax", threshold=0.006,
                                      f0_min=80, f0_max=880, interp_uv=False,
                                      output_interp_target_length=frames_per_chunk)
            f0_parts.append(f0.reshape(-1, 1).cpu().numpy())
        f0_np = np.concatenate(f0_parts, axis=0)[:V]
    if f0_np.shape[0] < V:
        f0_np = np.pad(f0_np, ((0, V - f0_np.shape[0]), (0, 0)), mode="edge")

    rms = librosa.feature.rms(y=y_seg, frame_length=2048, hop_length=HOP).T
    if rms.shape[0] > V:
        rms = rms[:V]
    elif rms.shape[0] < V:
        rms = np.pad(rms, ((0, V - rms.shape[0]), (0, 0)), mode="edge")

    return np.concatenate([mel_db, f0_np, rms], axis=-1)


def decode_clip_audio(video_path, start_sec, duration_sec):
    """Decode the clip's audio ONCE via ffmpeg (native mp4 seek) to a float32
    array. librosa.load on mp4 falls back to audioread which does NOT seek --
    it decodes from the start of the file up to the offset every call, so
    per-chunk decode cost GROWS WITH OFFSET (~12s per chunk at 1h into a 2h
    video). That starved the producer, stalled the animation, and burst frames
    when a chunk finally landed (the 'API sent at the same time' + 'not synced'
    symptoms). ffmpeg -ss before -i does a real seek: ~350ms regardless of
    offset. Slicing the predecoded array per chunk is then ~free.

    ponytail: one-shot whole-clip decode adds ~clip_duration/30s of startup
    time + ~115 MB RAM for 2h @ 16k mono. Acceptable for a test script; if you
    run multi-hour clips cold, switch to per-chunk ffmpeg subprocess calls."""
    lookahead = 2.0  # seconds, matches features_for_segment's extra context
    cmd = ["ffmpeg", "-y", "-ss", f"{start_sec:.3f}",
           "-t", f"{duration_sec + lookahead + 0.5:.3f}",
           "-i", video_path, "-vn", "-ac", "1", "-ar", str(SR),
           "-f", "f32le", "-"]
    proc = subprocess.run(cmd, capture_output=True, check=True)
    return np.frombuffer(proc.stdout, dtype=np.float32).copy()


async def stream_landmarks(audio_full, start_sec, V, model, fcpe_model, queue,
                           audio_mean=AUDIO_MEAN, audio_std=AUDIO_STD,
                           start_chunk=0, first_chunk_only=False):
    """Producer: slices the predecoded *audio_full* array per chunk, extracts
    features, runs one NAT forward, in a thread (run_in_executor) so it never
    blocks the event loop. Puts frames into *queue*.

    *start_chunk*: first chunk index to process (use 1 if chunk 0 was prefilled).
    *first_chunk_only*: if True, only process one chunk and return (no END sentinel)."""
    loop = asyncio.get_event_loop()
    chunk_sec = STREAM_CHUNK / 12.0
    chunk_samples = int(round(chunk_sec * SR))
    look_samples = int(round(2.0 * SR))
    total_chunks = (V + STREAM_CHUNK - 1) // STREAM_CHUNK

    def compute_chunk(c):
        n = min(STREAM_CHUNK, V - c * STREAM_CHUNK)
        n_samples = int(round(n / 12.0 * SR))
        s = c * chunk_samples
        e = s + n_samples + look_samples
        y_seg = audio_full[s:e]
        need = n_samples + look_samples
        if len(y_seg) < need:
            y_seg = np.pad(y_seg, (0, need - len(y_seg)))
        feats = features_for_segment(y_seg, n, fcpe_model)
        if audio_mean is not None:
            feats = (feats - audio_mean) / audio_std
        pad = (-feats.shape[0]) % SEQ_LEN
        if pad:
            feats = np.pad(feats, ((0, pad), (0, 0)), mode="edge")
        n_win = feats.shape[0] // SEQ_LEN
        batch = torch.from_numpy(feats.reshape(n_win, SEQ_LEN, 82)).to(DEVICE)
        with torch.no_grad():
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=USE_BF16):
                coeffs = model(batch)                           # [N, SEQ_LEN, 12]
                pred = coeffs @ PCA_COMP.to(DEVICE) + PCA_MEAN.to(DEVICE)  # -> 56
        return pred.cpu().float().numpy().reshape(n_win, SEQ_LEN, 28, 2).reshape(-1, 28, 2)[:n]

    n_chunks = 1 if first_chunk_only else total_chunks
    for c in range(start_chunk, start_chunk + n_chunks):
        frames = await loop.run_in_executor(None, compute_chunk, c)
        for lm in frames:
            await queue.put(lm)
    if not first_chunk_only:
        await queue.put(_STREAM_END)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--video", required=True, help="Path to a 12 fps input video")
    ap.add_argument("--start", type=str, default="00:00:00", help="Start time HH:MM:SS")
    ap.add_argument("--checkpoint", type=str, default="checkpoints/best_nat_model.pth")
    ap.add_argument("--max-frames", type=int, default=600, help="Cap predicted frames (default 600=50s)")
    ap.add_argument("--no-audio", action="store_true", help="Skip audio playback")
    args = ap.parse_args()
    start_sec = parse_start(args.start)

    cap = cv2.VideoCapture(args.video)
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    print(f"Video: {args.video}  fps={fps}  total_frames={total_frames}")
    if not (11.9 <= fps <= 12.1):
        raise SystemExit(f"Expected ~12 fps video (got {fps}). Re-encode to 12 fps first.")
    start_frame = int(round(start_sec * fps))
    V = total_frames - start_frame
    if args.max_frames is not None:
        V = min(V, args.max_frames)
    clip_duration = V / 12.0
    print(f"Streaming {V} frames ({clip_duration:.1f}s) from {start_sec}s (frame {start_frame}).")

    print("Loading FCPE + NAT model...")
    fcpe_model = spawn_bundled_infer_model(device=DEVICE)
    model = AudioToLandmarkNAT(**MODEL_KW).to(DEVICE)
    model.load_state_dict(torch.load(args.checkpoint, map_location=DEVICE))
    model.eval()

    # Warmup both models so the first user-facing chunk doesn't pay cold-start
    # (makes the first-frame lag ~0.24s instead of ~2.3s).
    print("Warming up models...")
    warmup_y = np.zeros(16000, dtype=np.float32)
    warmup_t = torch.from_numpy(warmup_y).unsqueeze(0).unsqueeze(-1).to(DEVICE)
    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=USE_BF16):
        fcpe_model.infer(warmup_t, sr=SR, decoder_mode="local_argmax", threshold=0.006)
    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=USE_BF16):
        model(torch.zeros(1, SEQ_LEN, 82, device=DEVICE))
    print("Done.")

    if not args.no_audio and not shutil.which("ffplay"):
        print("ffplay not found; audio playback disabled. Use `--no-audio` to silence this.")
        args.no_audio = True

    # One-shot async orchestrator: connect VTS, prompt, then start audio+stream.
    # The handshake happens BEFORE the prompt so both audio and animation start
    # at the same Enter press (no desync from handshake time).
    async def run():
        # 1. Connect to VTube Studio first (handshake + param listing).
        ws = await _vts.vts_connect()
        if ws is None:
            raise SystemExit("VTS connection failed.")

        # 2. Predecode the whole clip's audio ONCE via ffmpeg native seek.
        #    Also precompute the VERY FIRST chunk now (while the user reads
        #    the prompt) so the consumer has frames immediately at Enter.
        #    Without this there's a 4.2s dead-air gap (chunk-0 compute time).
        print("Predecoding clip audio...")
        audio_full = await asyncio.to_thread(decode_clip_audio,
                                              args.video, start_sec, clip_duration)
        print(f"  audio: {len(audio_full)} samples ({len(audio_full)/SR:.1f}s)")
        print("Precomputing first chunk (warm consumer)...")
        first_chunk_frames = V if V <= STREAM_CHUNK else STREAM_CHUNK
        queue = asyncio.Queue(maxsize=STREAM_CHUNK * 2)  # 2 chunks' worth
        # prefill queue with first chunk's frames so they're ready at Enter
        await stream_landmarks(
            audio_full, start_sec, first_chunk_frames, model, fcpe_model, queue,
            audio_mean=AUDIO_MEAN, audio_std=AUDIO_STD, first_chunk_only=True)

        # 3. Wait for user confirmation.
        print("\nVTS connected. Press Enter to start audio + landmark streaming...")
        await asyncio.to_thread(input)

        # 4. Launch audio (wall-clock playback via ffplay).
        ffplay = None
        if not args.no_audio:
            ffplay = subprocess.Popen(
                ["ffplay", "-nodisp", "-autoexit",
                 "-ss", f"{start_sec:.3f}", "-t", f"{clip_duration:.3f}",
                 args.video],
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

        # 5. Producer/consumer. Chunk 0 was prefilled above; producer starts
        #    from chunk 1. If the clip is only one chunk (V <= STREAM_CHUNK),
        #    skip production and just signal end-of-stream.
        async def producer():
            if V > STREAM_CHUNK:
                await stream_landmarks(audio_full, start_sec, V, model, fcpe_model, queue,
                                       start_chunk=1)
            await queue.put(_STREAM_END)

        async def get_frame():
            return await queue.get()

        try:
            await asyncio.gather(producer(), vts_driver_stream(get_frame, ws=ws))
        finally:
            if ffplay is not None:
                ffplay.terminate()

    asyncio.run(run())


if __name__ == "__main__":
    main()