"""Resident (warm) body-motion engine — the kemix velocity codec+flow, loaded ONCE. Why: `infer_turn` currently shells out to gen_fulltake_tag3.py per turn, which reloads the whole model every time (~4.7s, mostly load). This keeps the model resident so a turn pays only real inference, and — crucially — exposes `gen_window()`, the single-window primitive that v2 real-time streaming will call every ~500ms. It reuses gen_fulltake_tag3.py's *exact* classes/constants/decode by importing that module (the top-level imports run; main() does not). The generation here is the audio_manifest path of gen_fulltake (arbitrary wav + emotion + proto seed): the same window loop (fixed 4.27s window, autoregressive seed carry `seed = g[:, -PRE:, :]`, latent overlap-add stitch) — just callable and warm instead of a one-shot subprocess. eng = BodyEngine(ckpt=..., vqvae_dir=..., meanstd_dir=..., npz=..., proto_seed=...) quats, T = eng.generate_take(wav_path, "excitement") # (T,44,4) basis quats eng.write_vrma(quats, "out.vrma") """ import os import json import pickle import sys import librosa import numpy as np import torch from numpy.lib import stride_tricks HERE = os.path.dirname(os.path.abspath(__file__)) WH = os.path.abspath(os.path.join(HERE, "..", "motion-blender", "experiments", "wave-hands")) sys.path.insert(0, WH) sys.path.insert(0, os.path.join(WH, "..", "..", "lib")) # vrma_writer, vrma_to_npz import gen_fulltake_tag3 as G # runs G's top-level imports (torch, Tag3VrmLSM, …), NOT main() import scripts.infer_vrma as iv # GVRM is on sys.path after importing G from models.vrm_lsm import VrmLSM as BaseVrmLSM # noqa: E402 from omegaconf import OmegaConf # noqa: E402 from vrma_writer import ( # noqa: E402 CHILDREN as _VRMA_CHILDREN, REST_OFFSETS as _VRMA_REST_OFFSETS, write_vrma as _write_vrma44, write_vrma_from_template as _write_vrma_tpl, _read_glb_json as _read_glb, ) from vrma_to_npz import ( # noqa: E402 FINGER_BONES as _FB, _qinv as _v2n_qinv, _qmul as _v2n_qmul, ) FPS, AUDIO_SR = 30, 16000 _BONES44 = ["hips", "spine", "chest", "upperChest", "neck", "head", "leftShoulder", "leftUpperArm", "leftLowerArm", "leftHand", "rightShoulder", "rightUpperArm", "rightLowerArm", "rightHand", "leftUpperLeg", "leftLowerLeg", "leftFoot", "rightUpperLeg", "rightLowerLeg", "rightFoot"] + list(_FB) # vrma_writer's VRMC_vrm_animation humanoid map contains exactly the bones # present in REST_OFFSETS. In particular, its source rig has no upperChest: # that model channel is intentionally ignored by the batch VRMA playback path. _STREAM_BONE_INDICES = tuple( i for i, bone in enumerate(_BONES44) if bone in _VRMA_REST_OFFSETS ) _STREAM_BONES = tuple(_BONES44[i] for i in _STREAM_BONE_INDICES) def _source_world_rest_quaternions(local_rest, bones=None): """Build the source VRMA rig's cumulative rest rotations in ``bones`` order.""" bones = list(bones) if bones is not None else _BONES44 parent = { child: bone for bone, children in _VRMA_CHILDREN.items() for child in children } identity = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float64) cache = {} def world(bone): if bone in cache: return cache[bone] local = np.asarray(local_rest.get(bone, identity), dtype=np.float64) local /= max(float(np.linalg.norm(local)), 1e-12) parent_bone = parent.get(bone) value = _v2n_qmul(world(parent_bone), local) if parent_bone else local value /= max(float(np.linalg.norm(value)), 1e-12) cache[bone] = value return value return np.stack([world(bone) for bone in bones]) def basis_to_humanoid_space(quats, world_rest_quats): """Convert local source-rig basis rotations to VRM normalized-humanoid rotations. GestureVRM training stores ``basis = inverse(localRest) * channel``. The VRMAnimation loader converts the reconstructed source channel into normalized humanoid space as ``worldRest * basis * inverse(worldRest)``. Streaming skips that loader, so it must perform this conjugation explicitly before the browser writes quaternions to normalized bone nodes. """ q = np.asarray(quats, dtype=np.float64) world = np.asarray(world_rest_quats, dtype=np.float64) if q.shape[-2:] != (world.shape[-2], 4) or world.shape[-1] != 4: raise ValueError( f"quaternion/world-rest shapes do not match: {q.shape} vs {world.shape}" ) out = _v2n_qmul(_v2n_qmul(world, q), _v2n_qinv(world)) out /= np.maximum(np.linalg.norm(out, axis=-1, keepdims=True), 1e-12) return out.astype(np.float32) def audio_to_onset_amplitude(audio_data, audio_sr=AUDIO_SR): """In-memory equivalent of GestureVRM's ``load_onset_amplitude``. Streaming audio never exists as a complete wav file, so the worker keeps a rolling PCM window and computes the same two conditioning channels directly. The returned channel order intentionally matches the training loader: ``[amplitude, onset]`` at audio sample rate. """ audio_data = np.asarray(audio_data, dtype=np.float32).reshape(-1) if audio_data.size == 0: return np.zeros((0, 2), dtype=np.float32) frame_length = 1024 if audio_data.size < frame_length: padded = np.pad(audio_data, (0, frame_length - audio_data.size)) amplitude = np.full(audio_data.size, np.max(np.abs(padded)), dtype=np.float32) else: shape = (audio_data.size - frame_length + 1, frame_length) strides = (audio_data.strides[-1], audio_data.strides[-1]) rolling = stride_tricks.as_strided(audio_data, shape=shape, strides=strides) amplitude = np.max(np.abs(rolling), axis=1) amplitude = np.pad(amplitude, (0, frame_length - 1), mode="constant", constant_values=amplitude[-1]) onset_frames = librosa.onset.onset_detect(y=audio_data, sr=audio_sr, units="frames") onset_samples = librosa.frames_to_samples(onset_frames) onset_samples = onset_samples[onset_samples < audio_data.size] onset = np.zeros(audio_data.size, dtype=np.float32) onset[onset_samples] = 1.0 return np.stack([amplitude, onset], axis=-1).astype(np.float32) class BodyEngine: bone_names = tuple(_BONES44) stream_bone_indices = _STREAM_BONE_INDICES stream_bone_names = _STREAM_BONES def __init__(self, ckpt, vqvae_dir, meanstd_dir, npz, proto_seed=None, fingers=True, steps=8, split="test", device="cuda:0", model_kind="tag3", codec_selection=None, continuous_decode=False, template_vrma=None, denoiser_params=None): self.device = torch.device(device) self.fingers = fingers self.model_kind = model_kind self.continuous_decode = continuous_decode self.template_vrma = template_vrma if model_kind not in ("tag3", "base"): raise ValueError(f"unknown body model kind: {model_kind}") # ── 본 스키마는 mean/std 차원에서 유도한다 (2026-08-05) ── # 종전엔 44본·손가락 24본이 상수였다. 솔로 코퍼스가 엄지 6본을 더해 # 50본이 되면서 이 경로만 44에 묶여 조용히 어긋났다 — 렌더러 두 개에서 # 오늘 같은 잔재를 고쳤다. FINGER_BONES는 엄지를 끝에 APPEND했으므로 # 앞에서 n개를 자르면 구 24본 코퍼스와 순서가 그대로 일치한다. self.n_fingers = 0 if fingers: fdim = int(np.load(os.path.join(meanstd_dir, "vrm_fingers_mean.npy")).shape[0]) if fdim % 6: raise RuntimeError(f"fingers mean dim {fdim} is not a multiple of 6") self.n_fingers = fdim // 6 if self.n_fingers > len(_FB): raise RuntimeError(f"fingers mean dim {fdim} exceeds known bones {len(_FB)}") self.bone_names = tuple(_BONES44[:20] + list(_FB)[:self.n_fingers]) self.n_bones = len(self.bone_names) # 스트리밍 인덱스도 같은 스키마에서 다시 뽑는다 — 모듈 상수(44본)를 쓰면 # 50본 모델에서 엉뚱한 본을 잘라낸다. self.stream_bone_indices = tuple( i for i, bone in enumerate(self.bone_names) if bone in _VRMA_REST_OFFSETS) self.stream_bone_names = tuple(self.bone_names[i] for i in self.stream_bone_indices) # ── rest table (basis→channel) ── # template_vrma가 주어지면 **그 데이터셋이 실제로 벗겨낸 rest**를 쓴다. # 구 kemix_rest_table.npy는 솔로 베이크 템플릿과 어깨 120°/hips 90°가 # 다르다(43본 중 13본이 1° 초과) — 남의 테이블로 되입히면 제3의 잘못된 # 자세가 나온다. 기존 v4 프로파일은 template_vrma 없이 종전 경로 유지. self._hips_rest_trans = None if template_vrma: gltf, _ = _read_glb(template_vrma) hb = gltf["extensions"]["VRMC_vrm_animation"]["humanoid"]["humanBones"] nodes = gltf["nodes"] self._REST = {b: [float(v) for v in nodes[e["node"]].get("rotation", [0, 0, 0, 1])] for b, e in hb.items() if e.get("node") is not None} # 이동 적분의 기준점 = 리그 rest 골반 위치. 코덱은 이동 **속도**만 # 만들어서 절대 높이를 모른다 — 원점(0,0,0)에서 적분하면 골반이 # 발목 높이에 붙어 아바타가 바닥 아래로 잠긴다(실측 64.1 -> 1.8cm, # 사용자 보고 "모션이 바닥 아래로 잠겨있고", 2026-08-05). hips = (hb.get("hips") or {}).get("node") if hips is not None: self._hips_rest_trans = np.asarray( nodes[hips].get("translation", [0.0, 0.0, 0.0]), dtype=np.float32) else: self._REST = np.load(os.path.join(WH, "outputs", "kemix_rest_table.npy"), allow_pickle=True).item() self._REST_ARR = np.stack([np.asarray(self._REST[b], np.float64) if b in self._REST else np.array([0., 0., 0., 1.]) for b in self.bone_names]) self._WORLD_REST_ARR = _source_world_rest_quaternions(self._REST, self.bone_names) # ── config (mirrors gen_fulltake main() lines 124-148) ── cfg = OmegaConf.load(os.path.join(G.GVRM, "configs/shortcut_vrm.yaml")) cfg.gpus = [0] if fingers: cfg.model.use_exp = True # 구조 스위치가 켜진 채 학습된 체크포인트는 yaml 기본값(off)으로 만든 모델과 # 키가 다르다 — r04b가 rope_emb.inv_freq 3개로 로드에 실패했다. 호출자가 # 런의 구조를 알려줄 통로를 둔다 (2026-08-06). if denoiser_params: for k, v in dict(denoiser_params).items(): cfg.model.denoiser.params[k] = v codec_paths = {} if codec_selection: with open(codec_selection) as f: selection = json.load(f) codec_paths = selection["checkpoints"] for p in ("spine", "arms", "legs"): cfg[f"vqvae_{p}_path"] = codec_paths.get( p, os.path.join(vqvae_dir, f"best_{p}.pth")) cfg[f"mean_{p}_path"] = os.path.join(meanstd_dir, f"vrm_{p}_mean.npy") cfg[f"std_{p}_path"] = os.path.join(meanstd_dir, f"vrm_{p}_std.npy") if fingers: cfg["vqvae_fingers_path"] = codec_paths.get( "fingers", os.path.join(vqvae_dir, "best_fingers.pth")) cfg["mean_fingers_path"] = os.path.join(meanstd_dir, "vrm_fingers_mean.npy") cfg["std_fingers_path"] = os.path.join(meanstd_dir, "vrm_fingers_std.npy") if steps > 0: cfg.model.n_steps = steps _lang = os.path.join(WH, "outputs", "dummy_lang") + os.sep cfg.model.modality_encoder.params.data_path = _lang cfg.data.beat_data_path = _lang cfg.data.ours_mode = True cfg.data.our_npz_path = npz cfg.data.our_wave_path = os.path.join(WH, "outputs", "kemix_wave16k") cfg.data.our_speaker_id = 2 self.cfg = cfg # ── flow model (loaded ONCE) ── model_cls = G.Tag3VrmLSM if model_kind == "tag3" else BaseVrmLSM self.model = model_cls(cfg).to(self.device) ck = torch.load(ckpt, map_location=self.device) sd = {k.replace("module.", ""): v for k, v in ck.get("model_state_dict", ck).items()} self.model.load_state_dict(sd, strict=(model_kind == "base")) self.model.eval() # ── window geometry (gen_fulltake lines 160-166) ── self.SEQ = self.model.seq_len self.PRE = int(cfg.pre_frames) self.HOP_LAT = self.SEQ - self.PRE self.HOP_POSE = self.HOP_LAT * 4 self.WIN_POSE = self.SEQ * 4 self.WIN_SAMP = int(self.WIN_POSE * AUDIO_SR / FPS) # ── VQ codecs + mean/std (loaded ONCE) ── self._parts = [("spine", 36), ("arms", 48), ("legs", 39)] + \ ([("fingers", self.n_fingers * 6)] if fingers else []) self.vq = {} for part, dim in self._parts: m = iv.create_vqvae(dim) m.load_state_dict(torch.load(cfg[f"vqvae_{part}_path"], map_location="cpu")["net"]) self.vq[part] = m.eval().to(self.device) self.msd = {p: (torch.from_numpy(np.load(cfg[f"mean_{p}_path"])).to(self.device), torch.from_numpy(np.load(cfg[f"std_{p}_path"])).to(self.device)) for p, _ in self._parts} self.scale = cfg.vqvae_latent_scale # ── PAD token (words are disabled, but the modality encoder still needs an id). ── with open(os.path.join(_lang, "weights", "vocab.pkl"), "rb") as f: self.PAD = pickle.load(f).PAD_token # ── optional per-emotion prototype seed bank ── self.proto = None if proto_seed: npzf = np.load(proto_seed) self.proto = {k: torch.from_numpy(npzf[k]).float().unsqueeze(0).to(self.device) for k in npzf.files} self.C = 128 * len(self._parts) # ── decode helpers (verbatim from gen_fulltake main()) ── def _dec(self, p, lat): if self.continuous_decode: return self.vq[p].decoder((lat * self.scale).permute(0, 2, 1)) return self.vq[p].latent2origin(lat * self.scale)[0] def _decode_latent(self, g, with_trans=False): """-> (T,B,4) basis quats. with_trans=True also returns (T,3) 이동 속도. legs 코덱의 마지막 3차원은 hips 이동 속도인데, 종전엔 `legs[:, :-3]`로 **버리고 있었다** — 배우가 실제로 걷는데(구간 median 84cm, 최대 4m) 루트가 고정돼 발이 러닝머신처럼 쓸린다. 호출자가 적분해서 쓴다. """ parts = {"spine": g[..., :128], "arms": g[..., 128:256], "legs": g[..., 256:384]} if self.fingers and g.shape[-1] >= 512: parts["fingers"] = g[..., 384:512] rec = {p: self._dec(p, parts[p]) * self.msd[p][1] + self.msd[p][0] for p in parts} legs = rec["legs"][0].detach().cpu().numpy() T = rec["spine"].shape[1] nb = self.n_bones if "fingers" in rec else 20 r6 = np.zeros((T, nb, 6), np.float32) r6[:, G.SPINE] = rec["spine"][0].detach().cpu().numpy().reshape(T, 6, 6) r6[:, G.ARMS] = rec["arms"][0].detach().cpu().numpy().reshape(T, 8, 6) r6[:, G.LEGS] = legs[:, :-3].reshape(T, 6, 6) if "fingers" in rec: r6[:, 20:nb] = rec["fingers"][0].detach().cpu().numpy().reshape( T, self.n_fingers, 6) quats = iv.rotation_6d_to_quaternion(r6) # (T,B,4) basis quats return (quats, legs[:, -3:].astype(np.float32)) if with_trans else quats def _basis_to_channel(self, q): return _v2n_qmul(self._REST_ARR[None], q.astype(np.float64)).astype(np.float32) def basis_to_humanoid(self, q, mapped_only=False): """Return loader-equivalent rotations for normalized VRM humanoid nodes. ``mapped_only=True`` also removes channels absent from the source VRMA humanoid map. Today that is only ``upperChest``; batch VRMA playback ignores the same channel. """ out = basis_to_humanoid_space(q, self._WORLD_REST_ARR) if mapped_only: out = out[..., self.stream_bone_indices, :] return out def seed_for(self, emotion): """First-window seed: prototype opener for the emotion, else zeros.""" if self.proto is not None and emotion in self.proto: return self.proto[emotion] return torch.zeros((1, self.PRE, self.C), device=self.device) # ── the streaming primitive: one 4.27s window (gen_fulltake loop body 363-391) ── @torch.no_grad() def gen_window(self, onset_win, seed, emotion, style="general", guidance=None): """onset_win: (WIN_SAMP, 2) np; seed: (1,PRE,C). Returns latent g (1,SEQ,C). ``guidance=None`` → 모델 기본값(base는 yaml ``guidance_scale``, tag3은 모듈 ``GUIDANCE``). 종전 기본값은 3.0이었지만 **base 경로에서는 이 인자가 읽히지도 않아** 실제로는 늘 2.0(yaml)이 쓰였다 (2026-08-06 진단). 기본값을 3.0으로 둔 채 배관만 고치면 기존 발행물이 전부 달라지므로 None으로 바꾼다. """ onset_win = np.asarray(onset_win, dtype=np.float32) if onset_win.size == 0: onset_win = np.zeros((self.WIN_SAMP, 2), dtype=np.float32) if len(onset_win) < self.WIN_SAMP: onset_win = np.pad(onset_win, ((0, self.WIN_SAMP - len(onset_win)), (0, 0)), mode="edge") a_t = torch.from_numpy(onset_win).float().unsqueeze(0).to(self.device) word = torch.full((1, self.WIN_POSE), self.PAD, dtype=torch.long, device=self.device) sid = torch.full((1, self.WIN_POSE), 2, dtype=torch.long, device=self.device) body_cond = {"audio_onset": a_t, "word": word, "id": sid, "seed": seed, "style_feature": None} if guidance is not None: # base: VrmLSM.forward가 y['guidance']를 우선한다 body_cond["guidance"] = float(guidance) if self.model_kind == "tag3": body_cond.update({ "style_tag": torch.tensor( [G.STYLE_ID.get(style, 1)], device=self.device), "emo_tag": torch.tensor( [G.EMO_ID.get(emotion, 0)], device=self.device), }) if guidance is not None: body_cond["tag_guidance"] = float(guidance) self.model._cur_energy = None self.model._cur_rest = torch.from_numpy( np.full(self.SEQ, 2, np.int64))[None].to(self.device) cond = {"y": body_cond} x_t = self.model(cond)["latents"] # (1, C, 1, seq) return x_t.squeeze(2).permute(0, 2, 1) # (1, seq, C) @torch.no_grad() def effective_guidance(self, guidance=None): """실제로 적용되는 guidance — 발행 메타에 거짓값이 남지 않게 한다. tag3 기본값은 ``Tag2VrmLSM.forward``가 읽는 **그 모듈의** 전역 ``GUIDANCE``다(train_ours_tag2, 4.0). ``G``(gen_fulltake_tag3)에는 없어서 초판이 조용히 nan을 돌려줬다 — 그 nan이 샘플러까지 흘러가 latent 전체를 nan으로 만들었다. 정의된 자리에서 직접 찾는다. """ if guidance is not None: return float(guidance) if self.model_kind == "tag3": mod = sys.modules.get(type(self.model).forward.__module__) g = getattr(mod, "GUIDANCE", None) if g is not None: return float(g) return float(getattr(self.model, "guidance_scale", float("nan"))) @torch.no_grad() def generate_take(self, wav_path, emotion, T_full=None, style="general", guidance=None): """Batch (turn-based) full-take generation — warm. Returns (quats (T,B,4) basis, T, root (T,3)). `root`은 코덱이 만든 이동 속도를 원점에서 적분한 hips 위치 — 종전엔 이 값을 버려 발이 쓸렸다. """ onset = G.load_onset_amplitude(wav_path, AUDIO_SR) # (N,2) if T_full is None: T_full = int(len(onset) / AUDIO_SR * FPS) n_win = max(1, int(np.ceil((T_full - self.WIN_POSE) / self.HOP_POSE)) + 1) seed = self.seed_for(emotion) wins = [] for w in range(n_win): astart = int(w * self.HOP_POSE * AUDIO_SR / FPS) g = self.gen_window(onset[astart:astart + self.WIN_SAMP], seed, emotion, style, guidance) wins.append(g) seed = g[:, -self.PRE:, :].detach() # autoregressive continuity # latent overlap-add crossfade (gen_fulltake 394-406) C = wins[0].shape[-1] Ltot = self.HOP_LAT * (n_win - 1) + self.SEQ acc = torch.zeros((1, Ltot, C), device=self.device) wsum = torch.zeros((1, Ltot, 1), device=self.device) ramp = torch.arange(1, self.PRE + 1, device=self.device).float() / (self.PRE + 1) for w, g in enumerate(wins): wt = torch.ones(self.SEQ, device=self.device) if w > 0: wt[:self.PRE] = ramp if w < n_win - 1: wt[-self.PRE:] = ramp.flip(0) s = w * self.HOP_LAT acc[:, s:s + self.SEQ] += g * wt[None, :, None] wsum[:, s:s + self.SEQ] += wt[None, :, None] g_full = acc / wsum.clamp_min(1e-6) q, trans_v = self._decode_latent(g_full, with_trans=True) q, trans_v = q[:T_full], trans_v[:T_full] root = np.zeros_like(trans_v) root[1:] = np.cumsum(trans_v[1:], axis=0) if self._hips_rest_trans is not None: # 리그 rest 골반에서 적분 root = root + self._hips_rest_trans return q, q.shape[0], root def write_vrma(self, quats, out_path, root_translation=None): """basis 회전을 VRMA로 발행. template_vrma가 있으면 **그 데이터셋의 베이크 템플릿 구조 위에** 쓴다 (노드 계층·humanoid 맵·노드 rest가 베이크와 동일 — 베이크 산출물 채널을 1.2e-07로 재현 검증). 없으면 종전 제네릭 리그 경로 그대로. """ if self.template_vrma: return _write_vrma_tpl(self.template_vrma, quats, list(self.bone_names), FPS, out_path, root_translation=root_translation, quats_are_basis=True) _write_vrma44(self._basis_to_channel(quats), list(self.bone_names), FPS, out_path, root_translation=root_translation, rest_quats=self._REST) def new_stream(self, emotion, style="general", guidance=3.0, emit_frames=None): return BodyStream(self, emotion, style, guidance, emit_frames) class BodyStream: """Stateful real-time streaming with one generated-window of lookahead. v2 primitive. Audio is kept in a 4.27s rolling window whose RIGHT edge is "now"; missing history is padded on the LEFT. Each regular step advances by ``emit_frames / 4`` latent frames. A rollout's newest poses are not emitted immediately. The following rollout predicts the same global interval in its interior, so the two predictions are quaternion-SLERPed over that interval. Holding one rollout makes the end of one output packet and the beginning of the next packet adjacent frames from the same rollout, instead of unrelated tail predictions from two rollouts. The seed must be aligned with the *new window's beginning*. For a ``shift_lat`` rolling shift that is ``previous_g[:, shift_lat:shift_lat+PRE]`` — not the previous tail (the tail is only correct for the original 28-latent batch hop). This distinction is what keeps the generated motion and the rolling audio on the same timeline. """ def __init__(self, engine, emotion, style="general", guidance=3.0, emit_frames=None, blend=0, feature_fn=audio_to_onset_amplitude): self.e = engine self.emotion, self.style, self.guidance = emotion, style, guidance self.emit = int(emit_frames or engine.PRE * 4) if self.emit <= 0 or self.emit % 4: raise ValueError("emit_frames must be a positive multiple of the codec stride (4)") self.shift_lat = self.emit // 4 if self.shift_lat + engine.PRE > engine.SEQ: raise ValueError("emit_frames is too large to carry an aligned PRE-frame seed") # Kept as an accepted legacy argument so existing callers do not break. # Fixed-last-pose blending is intentionally superseded by rollout lookahead. _ = blend self.feature_fn = feature_fn self._pcm_pending = np.zeros(0, np.float32) self._pcm_window = np.zeros(0, np.float32) self._feature_window = np.zeros((0, 2), np.float32) self._step_index = 0 self._previous_g = None self._pending_segment = None # decoded poses awaiting lookahead self._last = None # last emitted frame (44,4) self._samples_received = 0 self._frames_scheduled = 0 self.frames_emitted = 0 def set_emotion(self, emotion): self.emotion = emotion # mid-stream emotion change (new set_emotion) def tick(self, onset_chunk): """Low-level feature input used by tests; production uses :meth:`push_pcm`.""" chunk = np.asarray(onset_chunk, dtype=np.float32).reshape(-1, 2) self._feature_window = np.concatenate([self._feature_window, chunk])[-self.e.WIN_SAMP:] window = self._left_pad_features(self._feature_window) out = self._generate(window, self.emit, self.shift_lat) self.frames_emitted += len(out) return out def push_pcm(self, pcm_chunk, flush=False): """Consume mono float32 PCM at 16kHz and return newly generated basis quaternions. Arbitrary input chunk sizes are accepted. Regular generation is quantized to a codec-aligned pose count. Output is delayed by one rollout so the next rollout can smooth the pending interval. ``flush=True`` right-pads a short final audio step to preserve codec alignment, then emits only its real prefix. """ pcm = np.asarray(pcm_chunk, dtype=np.float32).reshape(-1) if pcm.size: self._pcm_pending = np.concatenate([self._pcm_pending, pcm]) self._samples_received += pcm.size outputs = [] while True: need = self._samples_for_step(self._step_index) if self._pcm_pending.size < need: break segment = self._pcm_pending[:need] self._pcm_pending = self._pcm_pending[need:] outputs.append(self._push_segment(segment, self.emit, self.shift_lat)) self._step_index += 1 if flush and self._pcm_pending.size: segment = self._pcm_pending self._pcm_pending = np.zeros(0, np.float32) target_frames = int(round(self._samples_received * FPS / AUDIO_SR)) n_frames = min(self.emit, max(0, target_frames - self._frames_scheduled)) # A sub-half-frame transport remainder has zero duration on the 30fps # output grid. Consuming it without forcing a frame keeps the total # motion length equal to round(total_samples * FPS / AUDIO_SR). if n_frames > 0: # The codec/seed timeline cannot advance by a non-multiple of four # poses. Pad the final audio on the right and perform one regular # rollout advance; only the first real ``n_frames`` of that newest # 24-pose interval are retained below. regular_samples = self._samples_for_step(self._step_index) padded = np.pad( segment, (0, regular_samples - segment.size), mode="constant" ) outputs.append(self._push_segment( padded, n_frames, self.shift_lat, advance_frames=self.emit )) if flush: outputs.append(self._flush_pending()) if not outputs: return np.zeros((0, len(self.e.bone_names), 4), dtype=np.float32) out = np.concatenate(outputs, axis=0) self.frames_emitted += len(out) return out def _samples_for_step(self, step_index): samples_per_step = self.emit * AUDIO_SR / FPS lo = int(round(step_index * samples_per_step)) hi = int(round((step_index + 1) * samples_per_step)) return hi - lo def _left_pad_features(self, features): if len(features) >= self.e.WIN_SAMP: return features[-self.e.WIN_SAMP:] return np.pad(features, ((self.e.WIN_SAMP - len(features), 0), (0, 0)), mode="constant") def _push_segment(self, segment, out_frames, shift_lat, advance_frames=None): self._pcm_window = np.concatenate([self._pcm_window, segment])[-self.e.WIN_SAMP:] if self._pcm_window.size < self.e.WIN_SAMP: audio = np.pad(self._pcm_window, (self.e.WIN_SAMP - self._pcm_window.size, 0), mode="constant") else: audio = self._pcm_window out = self._generate( self.feature_fn(audio, AUDIO_SR), out_frames, shift_lat, advance_frames=advance_frames, ) self._frames_scheduled += out_frames return out def _generate(self, onset_window, out_frames, shift_lat, advance_frames=None): advance_frames = int(advance_frames or out_frames) if out_frames <= 0 or out_frames > advance_frames: raise ValueError( f"retained poses must be in 1..{advance_frames}, got {out_frames}" ) if self._previous_g is None: seed = self.e.seed_for(self.emotion) else: # The next rolling window begins `shift_lat` positions later than the # previous one. Carry the PRE latents at exactly that aligned location. seed = self._previous_g[:, shift_lat:shift_lat + self.e.PRE, :].detach() g = self.e.gen_window(onset_window, seed, self.emotion, self.style, self.guidance) self._previous_g = g.detach() q = np.asarray(self.e._decode_latent(g), dtype=np.float32) # (128,44,4) if advance_frames > len(q): raise ValueError( f"rollout has {len(q)} poses, cannot advance by {advance_frames} poses" ) if self._pending_segment is None: out = np.zeros((0, len(self.e.bone_names), 4), dtype=np.float32) else: out = self._stitch_pending( self._pending_segment, q, advance_frames ) newest_start = len(q) - advance_frames self._pending_segment = q[ newest_start:newest_start + out_frames ].copy() return out def _stitch_pending(self, previous, current_full, current_advance_frames): """Emit the pending interval using its duplicate prediction as lookahead. If the pending segment has duration ``P`` and the current rollout advances by ``C`` poses, these slices describe the same global time: ``previous == current_full[-(P+C):-C]``. """ p, c = len(previous), int(current_advance_frames) overlap_end = len(current_full) - c overlap_start = overlap_end - p if overlap_start < 0: raise ValueError( "rollout is too short for lookahead overlap: " f"need {p + c} poses, got {len(current_full)}" ) previous = _normalize_quaternions(previous.copy()) current = _normalize_quaternions( current_full[overlap_start:overlap_end].copy() ) # Smoothstep is exactly 0 at the pending interval's first frame and # exactly 1 at its last. Consequently the following packet begins at # the immediately-next frame of this same current rollout. if p == 1: alpha = np.ones((1, 1, 1), dtype=np.float32) else: u = np.linspace(0.0, 1.0, p, dtype=np.float32) alpha = (u * u * (3.0 - 2.0 * u))[:, None, None] return self._finish_seam(_slerp_frames(previous, current, alpha)) def _flush_pending(self): if self._pending_segment is None: return np.zeros((0, len(self.e.bone_names), 4), dtype=np.float32) segment = self._pending_segment self._pending_segment = None return self._finish_seam(segment.copy()) def _finish_seam(self, out): if not len(out): return out.astype(np.float32) out = _normalize_quaternions(out) # Quaternion signs q/-q encode the same rotation but interpolate very # differently. Keep every bone in the same hemisphere across the stream. prev = self._last for i in range(len(out)): ref = prev if i == 0 else out[i - 1] if ref is not None: out[i] = np.where((ref * out[i]).sum(-1, keepdims=True) < 0, -out[i], out[i]) self._last = out[-1].copy() return out.astype(np.float32) def _normalize_quaternions(q): q = np.asarray(q, dtype=np.float32) return q / np.maximum(np.linalg.norm(q, axis=-1, keepdims=True), 1e-8) def _slerp_frames(qa, qb, alpha): """Framewise, per-bone shortest-path quaternion SLERP.""" qa = _normalize_quaternions(qa) qb = _normalize_quaternions(qb) qb = np.where((qa * qb).sum(-1, keepdims=True) < 0, -qb, qb) dot = np.clip((qa * qb).sum(-1, keepdims=True), -1.0, 1.0) theta = np.arccos(dot) sin_theta = np.sin(theta) safe_sin = np.where(np.abs(sin_theta) > 1e-6, sin_theta, 1.0) spherical = ( np.sin((1.0 - alpha) * theta) / safe_sin * qa + np.sin(alpha * theta) / safe_sin * qb ) linear = (1.0 - alpha) * qa + alpha * qb out = np.where(np.abs(sin_theta) > 1e-6, spherical, linear) return _normalize_quaternions(out)