import unittest import numpy as np import torch from body_engine import BodyEngine, BodyStream, basis_to_humanoid_space def _quaternion_delta_degrees(qa, qb): qa = qa / np.linalg.norm(qa, axis=-1, keepdims=True) qb = qb / np.linalg.norm(qb, axis=-1, keepdims=True) dot = np.clip(np.abs(np.sum(qa * qb, axis=-1)), 0.0, 1.0) return np.rad2deg(2.0 * np.arccos(dot)) class _FakeEngine: PRE = 2 SEQ = 16 WIN_SAMP = 20 bone_names = ("hips", "spine") def __init__(self, rollout_bias_degrees=6.0, negate_calls=()): self.calls = [] self.rollout_bias_degrees = rollout_bias_degrees self.negate_calls = set(negate_calls) def seed_for(self, emotion): return torch.full((1, self.PRE, 1), -1.0) def gen_window(self, onset, seed, emotion, style, guidance): self.calls.append((onset.copy(), seed.clone(), emotion)) offset = 100 * (len(self.calls) - 1) return torch.arange(offset, offset + self.SEQ, dtype=torch.float32).reshape(1, -1, 1) def rollout_quaternions(self, call_index, signed=True): # Regular test rollouts advance by 24 poses. A call-specific bias makes # duplicate predictions differ, so endpoint ownership is observable. frame = np.arange(self.SEQ * 4, dtype=np.float32) degrees = ( call_index * 24.0 + frame + call_index * self.rollout_bias_degrees ) half = np.deg2rad(degrees) / 2.0 q = np.zeros((self.SEQ * 4, len(self.bone_names), 4), dtype=np.float32) q[..., 2] = np.sin(half)[:, None] q[..., 3] = np.cos(half)[:, None] if signed and call_index in self.negate_calls: q = -q return q def _decode_latent(self, g): call_index = int(g[0, 0, 0].item() // 100) return self.rollout_quaternions(call_index) class BodyStreamTests(unittest.TestCase): def test_feature_window_is_left_padded_so_new_audio_is_at_now(self): eng = _FakeEngine() stream = BodyStream(eng, "neutral", emit_frames=8) chunk = np.ones((5, 2), dtype=np.float32) stream.tick(chunk) onset = eng.calls[0][0] self.assertTrue(np.all(onset[:15] == 0)) self.assertTrue(np.all(onset[-5:] == 1)) def test_next_seed_uses_time_aligned_slice_not_previous_tail(self): eng = _FakeEngine() stream = BodyStream(eng, "neutral", emit_frames=8) stream.tick(np.ones((4, 2), dtype=np.float32)) stream.tick(np.ones((4, 2), dtype=np.float32)) second_seed = eng.calls[1][1].numpy().reshape(-1) np.testing.assert_array_equal(second_seed, np.array([2.0, 3.0], dtype=np.float32)) def test_pcm_steps_emit_codec_aligned_frame_count_and_flush_tail(self): eng = _FakeEngine() feature_fn = lambda audio, _sr: np.stack([audio, np.zeros_like(audio)], axis=-1) stream = BodyStream(eng, "joy", emit_frames=8, feature_fn=feature_fn) first_n = stream._samples_for_step(0) out = stream.push_pcm(np.ones(first_n, dtype=np.float32)) self.assertEqual(out.shape, (0, 2, 4)) tail = stream.push_pcm(np.ones(800, dtype=np.float32), flush=True) expected = round((first_n + 800) * 30 / 16_000) self.assertEqual(tail.shape, (expected, 2, 4)) self.assertEqual(stream.frames_emitted, expected) def test_regular_packets_are_held_for_one_rollout_of_lookahead(self): eng = _FakeEngine() feature_fn = lambda audio, _sr: np.stack([audio, np.zeros_like(audio)], axis=-1) stream = BodyStream(eng, "neutral", emit_frames=24, feature_fn=feature_fn) first = stream.push_pcm( np.ones(stream._samples_for_step(0), dtype=np.float32) ) second = stream.push_pcm( np.ones(stream._samples_for_step(1), dtype=np.float32) ) self.assertEqual(first.shape, (0, 2, 4)) self.assertEqual(second.shape, (24, 2, 4)) def test_total_frames_match_30fps_for_tiny_and_long_final_remainders(self): feature_fn = lambda audio, _sr: np.stack([audio, np.zeros_like(audio)], axis=-1) for n_samples in (1, 100, 267, 12_799, 12_800, 12_900, 160_100): with self.subTest(n_samples=n_samples): eng = _FakeEngine() stream = BodyStream(eng, "neutral", emit_frames=24, feature_fn=feature_fn) out = stream.push_pcm(np.ones(n_samples, dtype=np.float32), flush=True) self.assertEqual(len(out), round(n_samples * 30 / 16_000)) self.assertEqual(stream.frames_emitted, len(out)) def test_partial_flush_uses_regular_seed_shift_and_adjacent_tail_prefix(self): feature_fn = lambda audio, _sr: np.stack([audio, np.zeros_like(audio)], axis=-1) for remainder_frames in (1, 2, 3, 5, 23): with self.subTest(remainder_frames=remainder_frames): eng = _FakeEngine() stream = BodyStream( eng, "neutral", emit_frames=24, feature_fn=feature_fn ) regular_samples = stream._samples_for_step(0) stream.push_pcm(np.ones(regular_samples, dtype=np.float32)) remainder_samples = round(remainder_frames * 16_000 / 30) packet = stream.push_pcm( np.ones(remainder_samples, dtype=np.float32), flush=True ) self.assertEqual(packet.shape, (24 + remainder_frames, 2, 4)) # The final rollout always advances by the regular six latents, # independently of the 1..23 output-frame remainder. second_seed = eng.calls[1][1].numpy().reshape(-1) np.testing.assert_array_equal( second_seed, np.array([6.0, 7.0], dtype=np.float32) ) # Its audio step was padded on the right, not treated as a # shorter, differently aligned rolling window. self.assertTrue(np.all(eng.calls[1][0] == 0)) rollout_one = eng.rollout_quaternions(1, signed=False) # The stitched regular packet ends one frame before the real # prefix retained from the newest 24-frame rollout interval. np.testing.assert_allclose(packet[23], rollout_one[39], atol=1e-6) np.testing.assert_allclose(packet[24], rollout_one[40], atol=1e-6) np.testing.assert_allclose( packet[-1], rollout_one[40 + remainder_frames - 1], atol=1e-6 ) np.testing.assert_allclose( _quaternion_delta_degrees(packet[23], packet[24]), _quaternion_delta_degrees(rollout_one[39], rollout_one[40]), atol=1e-5, ) def test_packet_seam_uses_adjacent_frames_from_the_same_rollout(self): eng = _FakeEngine() feature_fn = lambda audio, _sr: np.stack([audio, np.zeros_like(audio)], axis=-1) stream = BodyStream(eng, "neutral", emit_frames=24, feature_fn=feature_fn) packets = [] for step in range(3): packets.append(stream.push_pcm( np.ones(stream._samples_for_step(step), dtype=np.float32) )) rollout_one = eng.rollout_quaternions(1, signed=False) rollout_zero = eng.rollout_quaternions(0, signed=False) # Smoothstep starts at zero, so the pending rollout owns the first frame. np.testing.assert_allclose(packets[1][0], rollout_zero[40], atol=1e-6) # Packet 0 ends at the last frame of rollout one's overlap prediction. np.testing.assert_allclose(packets[1][-1], rollout_one[39], atol=1e-6) # Packet 1 starts at the immediately following frame of rollout one. np.testing.assert_allclose(packets[2][0], rollout_one[40], atol=1e-6) # The cross-packet angular step is therefore exactly rollout one's own # adjacent-frame angular step, rather than a new packet-start correction. seam_delta = _quaternion_delta_degrees( packets[1][-1], packets[2][0] ) rollout_delta = _quaternion_delta_degrees( rollout_one[39], rollout_one[40] ) np.testing.assert_allclose(seam_delta, rollout_delta, atol=1e-5) def test_lookahead_slerp_normalizes_and_aligns_quaternion_signs(self): eng = _FakeEngine(rollout_bias_degrees=0.0, negate_calls=(1,)) feature_fn = lambda audio, _sr: np.stack([audio, np.zeros_like(audio)], axis=-1) stream = BodyStream(eng, "neutral", emit_frames=24, feature_fn=feature_fn) stream.push_pcm(np.ones(stream._samples_for_step(0), dtype=np.float32)) packet = stream.push_pcm( np.ones(stream._samples_for_step(1), dtype=np.float32) ) expected = eng.rollout_quaternions(0, signed=False)[-24:] np.testing.assert_allclose(packet, expected, atol=1e-6) np.testing.assert_allclose( np.linalg.norm(packet, axis=-1), np.ones(packet.shape[:-1], dtype=np.float32), atol=1e-6, ) class BodyRigSpaceTests(unittest.TestCase): def test_source_world_rest_conjugation_rotates_basis_axis_for_normalized_rig(self): """A source-rest Y rotation maps local X arm-down into normalized -Z.""" n_bones = len(BodyEngine.bone_names) identity = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32) world_rest = np.tile(identity, (n_bones, 1)) basis = np.tile(identity, (1, n_bones, 1)) arm = BodyEngine.bone_names.index("leftUpperArm") sin45 = np.sqrt(0.5) world_rest[arm] = [0.0, sin45, 0.0, sin45] # +90 degrees around Y sin30, cos30 = 0.5, np.sqrt(3.0) / 2.0 basis[0, arm] = [sin30, 0.0, 0.0, cos30] # +60 degrees around local X converted = basis_to_humanoid_space(basis, world_rest) np.testing.assert_allclose( converted[0, arm], np.array([0.0, 0.0, -sin30, cos30], dtype=np.float32), atol=1e-6, ) def test_stream_contract_omits_unmapped_upper_chest(self): self.assertIn("upperChest", BodyEngine.bone_names) self.assertNotIn("upperChest", BodyEngine.stream_bone_names) self.assertEqual(len(BodyEngine.stream_bone_names), 43) eng = BodyEngine.__new__(BodyEngine) eng._WORLD_REST_ARR = np.tile( np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float64), (len(BodyEngine.bone_names), 1), ) q = np.tile( np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32), (2, len(BodyEngine.bone_names), 1), ) mapped = eng.basis_to_humanoid(q, mapped_only=True) self.assertEqual(mapped.shape, (2, 43, 4)) if __name__ == "__main__": unittest.main()