loop_windows moves beside loop_window and takes the model; a test pins it on a real take with the two skeletons

The worker's window path loaded the 8B model at import, so no test reached it and swapped atan2 arguments or
the 30-joint skeleton went unnoticed. loop_windows(output, model) needs only model.output_skeleton and
model.fps. The test uses a committed 240-frame somaskel77 take (fixtures/take.npz, joints from a live Vili BVH
by forward kinematics) and stand-ins that carry the 77- and 30-joint skeletons' names and parents.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

scripts/kimodo-resident-worker.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 import argparse
 import json
-import math
 import os
 import shutil
 import sys
@@ -15,7 +14,7 @@ import torch
 from huggingface_hub.constants import HF_HOME
 
 from kimodo import DEFAULT_MODEL, load_model
-from kimodo.constraints import compute_global_heading, load_constraints_lst
+from kimodo.constraints import load_constraints_lst
 from kimodo.exports.motion_io import save_kimodo_npz
 from kimodo.meta import parse_prompts_from_meta
 from kimodo.model.load_model import TEXT_ENCODER_PRESETS
@@ -186,20 +185,6 @@ def save_outputs(output, output_base, export_bvh, bvh_standard_tpose):
     return files
 
 
-def loop_windows(output):
-    """One loop window per sample, chosen from Kimodo's own posed joints (always the somaskel77 output skeleton) and the
-    heading it computes from them. Frame indices into the take, last inclusive."""
-    skeleton = model.output_skeleton
-    hips = skeleton.bone_index["Hips"]
-    windows = []
-    for joints in output["posed_joints"]:
-        heading = compute_global_heading(torch.from_numpy(joints[None]), skeleton)[0]
-        angle = torch.atan2(heading[:, 1], heading[:, 0]).numpy()
-        first, last, pose_cost, velocity_cost = loop_window.loop_window(joints, angle, model.fps, math.ceil(len(joints) / 2), hips)
-        windows.append({"first": first, "last": last, "poseCost": pose_cost, "velocityCost": velocity_cost})
-    return windows
-
-
 def save_example_dirs(output_base, texts, num_frames, num_samples, constraints_path):
     output_stem = os.path.splitext(output_base)[0].rstrip(os.sep)
     base_name = os.path.basename(output_stem)
@@ -335,7 +320,7 @@ def generate(payload):
             request_bool(payload, "bvh"),
             request_bool(payload, "bvhStandardTpose") or request_bool(payload, "bvh_standard_tpose"),
         )
-        loops = loop_windows(output) if request_bool(payload, "loop") else None
+        loops = loop_window.loop_windows(output, model) if request_bool(payload, "loop") else None
         if request_bool(payload, "saveExampleDir") or request_bool(payload, "save_example_dir"):
             files.extend(save_example_dirs(output_base, texts, num_frames, num_samples, constraints_path))
         generation_count += 1

scripts/loop_window.py

@@ -3,9 +3,13 @@
 A seam is the jump from frame b back to frame a. Its cost is the pose difference plus 0.1 s times the velocity
 difference (Kovar, Gleicher and Pighin 2002, motion-graph transitions), both as RMS over joints, with each frame
 taken relative to its own Hips and turned to face a fixed way, so where the take wanders and which way it faces
-do not count. Numpy only.
+do not count. loop_window is numpy only; loop_windows needs torch and kimodo's heading, and no model weights.
 """
+import math
+
 import numpy as np
+import torch
+from kimodo.constraints import compute_global_heading
 
 VELOCITY_WEIGHT_SECONDS = 0.1
 
@@ -39,3 +43,18 @@ def loop_window(joints, heading, fps, min_frames, hips=0):
     best = np.flatnonzero(cost == cost.min())
     pick = best[np.lexsort((firsts[best], -(lasts[best] - firsts[best])))[0]]
     return int(firsts[pick]), int(lasts[pick]), float(pose_cost[pick]), float(velocity_cost[pick])
+
+
+def loop_windows(output, model):
+    """One loop window per sample of a Kimodo output, as the worker answers them: {first, last, poseCost, velocityCost},
+    frame indices into the take, last inclusive. The joints are posed_joints, always in model.output_skeleton (the
+    somaskel77, not model.skeleton), and the heading is Kimodo's own from those joints."""
+    skeleton = model.output_skeleton
+    hips = skeleton.bone_index["Hips"]
+    windows = []
+    for joints in output["posed_joints"]:
+        heading = compute_global_heading(torch.from_numpy(joints[None]), skeleton)[0]
+        angle = torch.atan2(heading[:, 1], heading[:, 0]).numpy()
+        first, last, pose_cost, velocity_cost = loop_window(joints, angle, model.fps, math.ceil(len(joints) / 2), hips)
+        windows.append({"first": first, "last": last, "poseCost": pose_cost, "velocityCost": velocity_cost})
+    return windows

scripts/test_loop_window.py

@@ -1,4 +1,4 @@
-"""python3 -m unittest scripts/test_loop_window.py   (numpy only)"""
+"""python3 -m unittest scripts/test_loop_window.py   (numpy, torch and kimodo; no model weights)"""
 import math
 import pathlib
 import sys
@@ -7,7 +7,8 @@ import unittest
 import numpy as np
 
 sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
-from loop_window import loop_window  # noqa: E402
+from loop_window import loop_window, loop_windows  # noqa: E402
+from types import SimpleNamespace  # noqa: E402
 
 FPS = 30
 
@@ -132,5 +133,57 @@ class LoopWindow(unittest.TestCase):
             self.assertNotIn("4321", str(caught.exception))
 
 
+class StandInSkeleton:
+    """What kimodo's heading and the worker read of a skeleton: the names and parents of the real one, its Hips index
+    and the [right, left] hip joints (RightLeg, LeftLeg) that heading is taken from."""
+
+    def __init__(self, names, parents):
+        self.bone_order_names = [str(name) for name in names]
+        self.bone_index = {name: index for index, name in enumerate(self.bone_order_names)}
+        self.bone_parents = {name: (None if parent < 0 else self.bone_order_names[parent]) for name, parent in zip(self.bone_order_names, parents)}
+        self.hip_joint_idx = [self.bone_index["RightLeg"], self.bone_index["LeftLeg"]]
+
+
+class WorkerLoopWindowsTest(unittest.TestCase):
+    """loop_windows on a real generated take (the somaskel77 joints of a live Vili take, 240 frames at 30 fps), with
+    stand-ins for the model's two skeletons: the 77-joint output skeleton and the 30-joint one the model also carries."""
+
+    @classmethod
+    def setUpClass(cls):
+        data = np.load(pathlib.Path(__file__).resolve().parent / "fixtures" / "take.npz")
+        cls.joints = data["joints"]
+        cls.model = SimpleNamespace(
+            output_skeleton=StandInSkeleton(data["names77"], data["parents77"]),
+            skeleton=StandInSkeleton(data["names30"], data["parents30"]),
+            fps=FPS,
+        )
+
+    def window(self, joints):
+        (answer,) = loop_windows({"posed_joints": [joints]}, self.model)
+        return answer["first"], answer["last"]
+
+    def test_the_real_take_gets_the_window_of_its_own_heading_in_the_output_skeleton(self):
+        self.assertEqual(self.joints.shape, (240, 77, 3))
+        self.assertEqual(self.model.output_skeleton.hip_joint_idx, [72, 67])
+        self.assertEqual(self.model.skeleton.hip_joint_idx, [26, 22])
+        self.assertEqual(self.window(self.joints), (35, 222))
+        # The 30-joint skeleton's hip indices read the wrong joints of a 77-joint take and move the window.
+        wrong = SimpleNamespace(output_skeleton=self.model.skeleton, fps=FPS)
+        self.assertNotEqual(loop_windows({"posed_joints": [self.joints]}, wrong)[0]["first"], 35)
+
+    def test_turning_the_body_about_its_hips_by_a_smooth_angle_leaves_the_window(self):
+        turned = np.stack([
+            turn(self.joints[f] - self.joints[f, 0], 0.8 * math.sin(0.05 * f)) + self.joints[f, 0] for f in range(len(self.joints))
+        ]).astype(np.float32)
+        self.assertEqual(self.window(turned), (35, 222))
+
+    def test_each_sample_gets_its_own_window_with_its_costs(self):
+        still = np.repeat(self.joints[:1], 240, axis=0)
+        answers = loop_windows({"posed_joints": [self.joints, still]}, self.model)
+        self.assertEqual([(a["first"], a["last"]) for a in answers], [(35, 222), (0, 239)])
+        self.assertEqual(set(answers[0]), {"first", "last", "poseCost", "velocityCost"})
+        self.assertAlmostEqual(answers[1]["poseCost"], 0.0, places=6)
+
+
 if __name__ == "__main__":
     unittest.main()