cozyclay 1.0.0

package/tools/ardy/cclay_sequence_generate.py

@@ -0,0 +1,595 @@
+# SPDX-License-Identifier: GPL-3.0-or-later
+# SPDX-FileCopyrightText: CozyClay contributors
+#
+# cclay-owned script. Source of truth lives in the CozyClay repo at
+# scripts/ardy/cclay_sequence_generate.py and is deployed to the ARDY box at
+# ~/ardy/scripts/cclay_sequence_generate.py. Do not edit the deployed copy.
+"""Segment-chained text-to-motion generation for the cclay ARDY pipeline.
+
+Generates ONE continuous rollout across multiple text segments ("run forward,
+then bow"): each segment is generated by the model's autoregressive sampling
+loop with its own prompt, conditioned on the tail of the previously generated
+motion via ``init_history_sequence`` — the same history-conditioning mechanism
+the interactive demo's replan path uses. No npz splicing, no crossfading, no
+array blending: segment continuity comes purely from history conditioning.
+
+Two empirically load-bearing choices (tuned on ARDY-Core-RP-20FPS-Horizon40):
+
+* Each segment is generated with a full ``model(...)`` call whose pad mask
+  marks the segment's exact clip length. The model was trained on clips that
+  start/end near rest, so it settles the action before the segment boundary —
+  which is what makes the next segment's transition smooth.
+* Only a SHORT history tail (default 0.6 s) is passed into the next segment.
+  Long histories dominate the window and suppress the new segment's prompt
+  (e.g. a 3 s running history makes the model ignore "bows deeply").
+
+Segment tails are additionally gated on continuity: if the last frames of a
+non-final segment still move faster than --max_boundary_jump, the segment is
+regenerated with a derived seed (bounded attempts, deterministic given
+--seed); the calmest attempt wins. This is regeneration, never blending.
+
+Root 2D waypoints (--root-2d FRAME X Z HEADING, frames in the ROLLOUT-GLOBAL
+0..total-1 range) ride the same chained calls: the constraint set is built
+once over the whole rollout, and each segment call receives the slice covering
+its history tail + new frames with the history frames zeroed — exactly the
+interactive demo's streaming-constraint pattern. The model re-anchors observed
+roots by its running global translation, so rollout-global coordinates stay
+valid across every seam. Waypoint conditioning and history conditioning are
+first-class together in ``Ardy.__call__``; nothing here is off-label.
+
+Usage (from ~/ardy):
+    .venv/bin/python scripts/cclay_sequence_generate.py \
+        --segment "A person runs forward." 3 \
+        --segment "A person bows deeply." 4 \
+        --output outputs/cclay/seq --seed 7
+
+The last stdout line is a single JSON object:
+    {"frames": int, "fps": int,
+     "segments": [{"prompt", "requested_s", "start_frame", "end_frame"}],
+     "continuity": {"mean_jump_m", "max_jump_m", "max_jump_frame",
+                    "boundary_max_jump_m"}}
+start_frame/end_frame are 0-based inclusive npz frame indices. A "jump" is the
+max per-frame L2 displacement across all posed joints between consecutive
+frames (meters); jump index j measures the transition from frame j to j+1 and
+is reported as the landing frame j+1 in max_jump_frame.
+
+The output npz has exactly the same keys/shapes/conventions as
+scripts/generate.py single-prompt output (Y-up meters, cskel27), with
+text = " then ".join(segment prompts).
+"""
+import argparse
+import json
+import math
+import os
+
+import numpy as np
+import torch
+
+from ardy.model import DEFAULT_MODEL, load_model
+from ardy.model.loading import get_env_var
+from ardy.model.registry import resolve_model_name
+from ardy.motion_rep.tools import length_to_mask
+from ardy.postprocess import post_process_motion
+from ardy.skeleton import SOMASkeleton30
+from ardy.tools import seed_everything, to_numpy
+
+# Prime stride between per-segment base seeds so retry seeds never collide.
+_SEED_STRIDE = 9973
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Segment-chained autoregressive text-to-motion generation (cclay)"
+    )
+    parser.add_argument(
+        "--segment",
+        action="append",
+        nargs=2,
+        metavar=("PROMPT", "SECONDS"),
+        required=True,
+        help="Repeatable: a text prompt and its duration in seconds.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="output",
+        help="Output stem name; bare names are placed under outputs/.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=DEFAULT_MODEL,
+        help="Model nickname or full folder name (default: %(default)s).",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Seed for reproducible results.",
+    )
+    parser.add_argument(
+        "--diffusion_steps",
+        type=int,
+        default=None,
+        help="Denoising steps, at most the model's num_base_steps (the default).",
+    )
+    parser.add_argument(
+        "--history_frames",
+        type=int,
+        default=None,
+        help=(
+            "History crop budget within the FIRST segment (multiple of the model's "
+            "token size), as in scripts/generate.py. Default: the longest history "
+            "fitting the trained 10 s window together with the generation horizon."
+        ),
+    )
+    parser.add_argument(
+        "--transition_frames",
+        type=int,
+        default=None,
+        help=(
+            "Frames of previous-motion tail passed as init history into each "
+            "subsequent segment (multiple of the model's token size). Short tails "
+            "keep boundaries continuous while letting the new prompt take over. "
+            "Default: ~0.6 s at the model fps."
+        ),
+    )
+    parser.add_argument(
+        "--root-2d",
+        dest="root_2d",
+        action="append",
+        nargs=4,
+        metavar=("FRAME", "X", "Z", "HEADING"),
+        help=(
+            "Repeatable: pin the root's ground position (meters) at a rollout-global "
+            "frame, optionally with a heading in radians ('none' to omit). Headings "
+            "are all-or-none across the set. Frame 0 is the rollout origin facing +Z."
+        ),
+    )
+    parser.add_argument(
+        "--max_boundary_jump",
+        type=float,
+        default=0.08,
+        help=(
+            "Continuity gate (meters/frame): a non-final segment whose tail still "
+            "moves faster than this is regenerated with a derived seed (default: "
+            "%(default)s)."
+        ),
+    )
+    parser.add_argument(
+        "--max_segment_attempts",
+        type=int,
+        default=6,
+        help="Max generation attempts per non-final segment for the continuity gate (default: %(default)s).",
+    )
+    parser.add_argument(
+        "--cfg_weight",
+        type=float,
+        nargs="+",
+        default=[2.0, 2.0],
+        help="CFG scale(s): one float (text) or two floats (text, constraint).",
+    )
+    parser.add_argument(
+        "--no-postprocess",
+        action="store_true",
+        help="Don't apply motion post-processing (foot-skate reduction).",
+    )
+    parser.add_argument(
+        "--checkpoints_dir",
+        type=str,
+        default=None,
+        help="Local dir holding released model folders (falls back to CHECKPOINTS_DIR env).",
+    )
+    return parser.parse_args()
+
+
+def _default_history_frames(fps: float, gen_horizon_len: int, num_frames_per_token: int) -> int:
+    """Longest history that, together with the generation horizon, fits the trained 10 s window.
+
+    Same budget as scripts/generate.py and the interactive demo: unbounded
+    history degrades long generations into jitter.
+    """
+    max_window_len = (int(10 * fps) // num_frames_per_token) * num_frames_per_token
+    return ((max_window_len - gen_horizon_len) // num_frames_per_token) * num_frames_per_token
+
+
+def _resolve_output_base(path: str, default_dir: str = "outputs") -> str:
+    """Place bare output names under ``default_dir``; honor explicit paths (as generate.py)."""
+    if os.path.dirname(path):
+        return path
+    return os.path.join(default_dir, path)
+
+
+def _single_file_path(path: str, ext: str) -> str:
+    """Return path for a single output file; add ext if missing, create parent dirs."""
+    if not path.endswith(ext):
+        path = path.rstrip(os.sep) + ext
+    parent = os.path.dirname(path)
+    if parent:
+        os.makedirs(parent, exist_ok=True)
+    return path
+
+
+def save_motion_npz(path: str, motion_dict: dict, fps: float, text: str) -> None:
+    """Save a motion output dict to ``.npz`` along with fps and the prompt (as generate.py)."""
+    arrays = {k: np.asarray(v) for k, v in motion_dict.items()}
+    arrays["fps"] = np.asarray(fps)
+    arrays["text"] = np.asarray(text)
+    np.savez(path, **arrays)
+
+
+def _parse_frame(raw, num_frames: int, flag: str) -> int:
+    try:
+        frame = int(raw)
+    except ValueError:
+        raise ValueError(f"{flag} frame must be an integer, got {raw!r}.")
+    if not 0 <= frame < num_frames:
+        raise ValueError(
+            f"{flag} frame {frame} is outside the rollout (0..{num_frames - 1} for "
+            f"these segments); lengthen a segment or move the constraint."
+        )
+    return frame
+
+
+def _parse_floats(raw_values, flag: str, names: str) -> list:
+    try:
+        return [float(value) for value in raw_values]
+    except ValueError:
+        raise ValueError(f"{flag} {names} must be numbers, got {tuple(raw_values)!r}.")
+
+
+def parse_root_waypoints(raw_waypoints, num_frames: int) -> list:
+    """Validate --root-2d. HEADING is radians or the literal 'none'.
+
+    Same contract as cclay_constrained_generate.py, with frames measured in
+    the rollout-global 0..total-1 range.
+    """
+    if not raw_waypoints:
+        return []
+    waypoints = []
+    seen = set()
+    for raw_frame, raw_x, raw_z, raw_heading in raw_waypoints:
+        frame = _parse_frame(raw_frame, num_frames, "--root-2d")
+        if frame in seen:
+            raise ValueError(
+                f"duplicate --root-2d for frame {frame}; one waypoint per frame."
+            )
+        seen.add(frame)
+        x, z = _parse_floats((raw_x, raw_z), "--root-2d", "X Z")
+        if str(raw_heading).lower() == "none":
+            heading = None
+        else:
+            heading = _parse_floats((raw_heading,), "--root-2d", "HEADING")[0]
+        waypoints.append({"frame": frame, "xz": [x, z], "heading": heading})
+    waypoints.sort(key=lambda entry: entry["frame"])
+    if len({entry["heading"] is None for entry in waypoints}) > 1:
+        # ARDY conditions the whole waypoint set on one heading tensor, so a
+        # partly-headed request would silently invent headings for the rest.
+        raise ValueError(
+            "--root-2d heading must be given for every waypoint or for none of them; "
+            "ARDY conditions the whole waypoint set on one heading tensor."
+        )
+    return waypoints
+
+
+def measure_waypoints(waypoints: list, generated_joints, skeleton) -> list:
+    """Horizontal distance from each requested root waypoint to the rollout's root."""
+    report = []
+    root_index = skeleton.root_idx
+    for entry in waypoints:
+        root = np.asarray(generated_joints[entry["frame"], root_index], dtype=np.float64)
+        achieved = [float(root[0]), float(root[2])]
+        error = float(np.linalg.norm(np.asarray(achieved) - np.asarray(entry["xz"])))
+        report.append(
+            {
+                "frame": entry["frame"],
+                "requested_xz": [round(value, 4) for value in entry["xz"]],
+                "achieved_xz": [round(value, 4) for value in achieved],
+                "achieved_error_m": round(error, 4),
+                "heading_rad": entry["heading"],
+            }
+        )
+    return report
+
+
+def _posed_joint_jumps(posed_joints: np.ndarray) -> np.ndarray:
+    """jump[j] = max over joints of the L2 displacement between frames j and j+1 (meters)."""
+    disp = np.linalg.norm(posed_joints[1:] - posed_joints[:-1], axis=-1)  # (F-1, J)
+    return disp.max(axis=-1)  # (F-1,)
+
+
+def _continuity_metrics(posed_joints: np.ndarray, boundaries: list) -> dict:
+    """Continuity metrics from posed joints (F, J, 3), meters.
+
+    boundary_max_jump_m = max jump within +/-2 frames of every segment boundary
+    (jumps whose either endpoint frame is within 2 frames of a boundary start).
+    """
+    jumps = _posed_joint_jumps(posed_joints)
+    boundary_max = 0.0
+    for b in boundaries:
+        lo = max(0, b - 3)
+        hi = min(len(jumps), b + 2)
+        if lo < hi:
+            boundary_max = max(boundary_max, float(jumps[lo:hi].max()))
+    return {
+        "mean_jump_m": float(jumps.mean()),
+        "max_jump_m": float(jumps.max()),
+        "max_jump_frame": int(jumps.argmax()) + 1,
+        "boundary_max_jump_m": boundary_max,
+    }
+
+
+def main():
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+    args = parse_args()
+
+    segments = []
+    for prompt, seconds_str in args.segment:
+        prompt = prompt.strip()
+        try:
+            seconds = float(seconds_str)
+        except ValueError:
+            raise ValueError(f"--segment duration must be a number, got {seconds_str!r}.")
+        if not prompt:
+            raise ValueError("--segment prompt must be non-empty.")
+        if seconds <= 0:
+            raise ValueError(f"--segment duration must be > 0 seconds, got {seconds}.")
+        segments.append((prompt, seconds))
+
+    if len(args.cfg_weight) == 1:
+        cfg_weight = float(args.cfg_weight[0])
+    elif len(args.cfg_weight) == 2:
+        cfg_weight = (float(args.cfg_weight[0]), float(args.cfg_weight[1]))
+    else:
+        raise ValueError("--cfg_weight expects one float (text) or two floats (text, constraint).")
+    if args.max_segment_attempts < 1:
+        raise ValueError(f"--max_segment_attempts must be >= 1, got {args.max_segment_attempts}.")
+
+    # Load model (same path as scripts/generate.py)
+    checkpoints_dir = args.checkpoints_dir or get_env_var("CHECKPOINTS_DIR")
+    resolved_model = resolve_model_name(args.model, checkpoints_dir=checkpoints_dir)
+    model = load_model(resolved_model, device=device, checkpoints_dir=checkpoints_dir)
+    print(f"Loaded model: {resolved_model}")
+
+    fps = model.motion_rep.fps
+    patch = model.num_frames_per_token
+    gen_horizon = model.gen_horizon_len
+    max_window_len = (int(10 * fps) // patch) * patch  # trained-window budget
+
+    num_base_steps = int(model.diffusion.num_base_steps)
+    diffusion_steps = args.diffusion_steps if args.diffusion_steps is not None else num_base_steps
+    if not 1 <= diffusion_steps <= num_base_steps:
+        raise ValueError(
+            f"--diffusion_steps must be between 1 and {num_base_steps} "
+            f"(this model's num_base_steps); got {diffusion_steps}."
+        )
+
+    history_budget = args.history_frames
+    if history_budget is None:
+        history_budget = _default_history_frames(fps, gen_horizon, patch)
+    elif history_budget < patch or history_budget % patch != 0:
+        raise ValueError(f"--history_frames must be a positive multiple of {patch} (this model's token size).")
+
+    transition_frames = args.transition_frames
+    if transition_frames is None:
+        transition_frames = max(patch, (int(0.6 * fps) // patch) * patch)
+    elif transition_frames < patch or transition_frames % patch != 0:
+        raise ValueError(f"--transition_frames must be a positive multiple of {patch} (this model's token size).")
+    print(
+        f"Using {history_budget} history frames (first-segment crop), "
+        f"{transition_frames} transition frames, {diffusion_steps} denoising steps"
+    )
+
+    # Per-segment frame counts. Later segments must fit the trained window
+    # together with the transition history (init history cannot be cropped
+    # mid-generation by the model).
+    seg_bounds = []
+    cum = 0
+    for i, (prompt, seconds) in enumerate(segments):
+        n = max(1, int(round(seconds * fps)))
+        if n < patch:
+            raise ValueError(
+                f"segment {prompt!r}: {seconds}s = {n} frames is shorter than one model "
+                f"token ({patch} frames at {fps} fps)."
+            )
+        if i > 0:
+            window = transition_frames + math.ceil(n / gen_horizon) * gen_horizon
+            if window > max_window_len:
+                max_s = (max_window_len - transition_frames) / fps
+                raise ValueError(
+                    f"segment {prompt!r}: {seconds}s exceeds the trained window for a "
+                    f"chained segment (max ~{max_s:.1f}s); split it into shorter segments."
+                )
+        seg_bounds.append({"prompt": prompt, "requested_s": seconds, "start_frame": cum, "end_frame": cum + n - 1})
+        cum += n
+    total_frames = cum
+    print(f"Will generate {total_frames} frames ({total_frames / fps:.2f}s at {fps} fps) over {len(segments)} segments")
+
+    # Root waypoints: ONE constraint set over the whole rollout, built in the
+    # rollout-global frame (frame 0 at the origin facing +Z — the first
+    # segment's canonical frame). Each chained call slices its own span below.
+    waypoints = parse_root_waypoints(args.root_2d or [], total_frames)
+    observed_full = None
+    mask_full = None
+    constraint_lst = []
+    if waypoints:
+        from ardy.constraints import Root2DConstraintSet
+
+        headings = [entry["heading"] for entry in waypoints]
+        constraint_lst = [
+            Root2DConstraintSet(
+                model.skeleton,
+                frame_indices=torch.tensor([entry["frame"] for entry in waypoints]),
+                root_2d=torch.tensor(
+                    [entry["xz"] for entry in waypoints], device=device, dtype=torch.float32
+                ),
+                global_root_heading=(
+                    None
+                    if headings[0] is None
+                    else torch.tensor(headings, device=device, dtype=torch.float32)
+                ),
+            )
+        ]
+        observed_full, mask_full = model.motion_rep.create_conditions_from_constraints_batched(
+            constraint_lst,
+            torch.tensor([total_frames], device=device),
+            to_normalize=True,
+            device=device,
+        )
+        print(
+            f"Root2DConstraintSet on rollout frames "
+            f"{[entry['frame'] for entry in waypoints]}"
+        )
+
+    def generate_segment(prompt: str, n: int, history: torch.Tensor | None, seg_start: int) -> torch.Tensor:
+        """One segment: full sampling loop conditioned on the (optional) history tail.
+
+        ``seg_start`` is the segment's first frame in rollout-global numbering;
+        with waypoints present the call receives the observed-motion slice for
+        its history tail + new frames, history frames zeroed (the demo's
+        streaming pattern — constraints never fight the conditioning history).
+        Returns only the n newly generated NORMALIZED frames.
+        """
+        history_len = 0 if history is None else history.shape[1]
+        num_frames = history_len + n
+        pad_mask = length_to_mask(torch.tensor([num_frames], device=device))
+        first_heading_angle = torch.zeros(1, device=device) if history is None else None
+        seg_mask = None
+        seg_observed = None
+        if mask_full is not None:
+            call_start = seg_start - history_len
+            seg_mask = mask_full[:, call_start : call_start + num_frames].clone()
+            seg_observed = observed_full[:, call_start : call_start + num_frames].clone()
+            seg_mask[:, :history_len] = 0.0
+            seg_observed[:, :history_len] = 0.0
+            if not bool(seg_mask.any()):
+                # no waypoint lands in this call's generated span; a pure
+                # zero mask would only cost compute
+                seg_mask = None
+                seg_observed = None
+        with torch.no_grad():
+            motion = model(
+                [prompt],
+                num_frames,
+                num_denoising_steps=diffusion_steps,
+                pad_mask=pad_mask,
+                first_heading_angle=first_heading_angle,
+                motion_mask=seg_mask,
+                observed_motion=seg_observed,
+                cfg_weight=cfg_weight,
+                progress_bar=lambda iterable: iterable,
+                init_history_sequence=history,
+                crop_history_length=history_budget if history is None else None,
+            )
+        return motion[:, history_len:]
+
+    def tail_jump(segment_motion: torch.Tensor) -> float:
+        """Max posed-joint jump over the last `patch` transitions of a segment (meters/frame)."""
+        with torch.no_grad():
+            out = model.motion_rep.inverse(segment_motion[:, -(patch + 1) :], is_normalized=True)
+        pj = out["posed_joints"][0].cpu().numpy()
+        return float(_posed_joint_jumps(pj).max())
+
+    # Single continuous rollout: `acc` holds the accumulated NORMALIZED motion
+    # tensor (1, F, D); each segment conditions on its tail.
+    acc = None
+    gate_worst_jump = None
+    gate_exhausted = False
+    for seg_idx, seg in enumerate(seg_bounds):
+        n = seg["end_frame"] - seg["start_frame"] + 1
+        is_final = seg_idx == len(seg_bounds) - 1
+        best, best_jump = None, None
+        attempts = 1 if is_final else args.max_segment_attempts
+        for attempt in range(attempts):
+            if args.seed is not None:
+                seed_everything(args.seed + _SEED_STRIDE * seg_idx + attempt)
+            if acc is None:
+                history = None
+            else:
+                # Clamp to the largest token multiple available so a short
+                # first segment never feeds a non-token-multiple history.
+                hist_len = min(transition_frames, (acc.shape[1] // patch) * patch)
+                history = acc[:, -hist_len:]
+            segment_motion = generate_segment(seg["prompt"], n, history, seg["start_frame"])
+            if is_final:
+                best = segment_motion
+                break
+            jump = tail_jump(segment_motion)
+            if best_jump is None or jump < best_jump:
+                best, best_jump = segment_motion, jump
+            if jump <= args.max_boundary_jump:
+                break
+        if not is_final:
+            print(
+                f"[segment {seg['start_frame']}-{seg['end_frame']}] tail jump "
+                f"{best_jump:.4f} m/frame after {attempt + 1} attempt(s)"
+            )
+            gate_worst_jump = best_jump if gate_worst_jump is None else max(gate_worst_jump, best_jump)
+            if best_jump > args.max_boundary_jump:
+                gate_exhausted = True
+        acc = best if acc is None else torch.cat([acc, best], dim=1)
+
+    # Unnormalize + inverse ONCE over the full rollout, then the same tail as
+    # scripts/generate.py: postprocess once, numpy, save.
+    with torch.no_grad():
+        output = model.motion_rep.inverse(acc, is_normalized=True)
+
+    use_postprocess = "g1" not in resolved_model.lower() and not args.no_postprocess
+    if use_postprocess:
+        corrected = post_process_motion(
+            output["local_rot_mats"],
+            output["root_positions"],
+            output["foot_contacts"],
+            model.skeleton,
+            # With waypoints, ARDY's own postprocess enforces the pinned root
+            # contacts instead of skating them away (as the one-shot path).
+            constraint_lst=constraint_lst if constraint_lst else None,
+        )
+        output.update(corrected)
+
+    if isinstance(model.skeleton, SOMASkeleton30):
+        output = model.skeleton.output_to_SOMASkeleton77(output)
+
+    output = to_numpy(output)
+    # Single sample: drop the batch dim (generate.py _select_sample with index 0).
+    motion_dict = {
+        k: (v[0] if hasattr(v, "shape") and len(v.shape) > 0 and v.shape[0] == 1 else v)
+        for k, v in output.items()
+    }
+
+    text = " then ".join(seg["prompt"] for seg in seg_bounds)
+    output_base = _resolve_output_base(args.output)
+    npz_path = _single_file_path(output_base, ".npz")
+    print(f"Saving the npz output to {npz_path}")
+    save_motion_npz(npz_path, motion_dict, fps, text)
+
+    boundaries = [seg["start_frame"] for seg in seg_bounds[1:]]
+    frames = int(motion_dict["posed_joints"].shape[0])
+    if frames != total_frames:
+        raise RuntimeError(
+            f"generated frame count {frames} does not match the planned segment table "
+            f"({total_frames}); segment boundaries would be mislabeled."
+        )
+    result = {
+        "frames": frames,
+        "fps": int(fps),
+        "model": resolved_model,
+        "segments": seg_bounds,
+        "continuity": _continuity_metrics(np.asarray(motion_dict["posed_joints"]), boundaries),
+        "boundary_gate": {
+            "threshold_m": args.max_boundary_jump,
+            "max_attempts": args.max_segment_attempts,
+            "worst_tail_jump_m": gate_worst_jump,
+            "exhausted": gate_exhausted,
+        },
+        # Measured, never asserted: horizontal error at every pinned frame.
+        "waypoints": measure_waypoints(
+            waypoints, np.asarray(motion_dict["posed_joints"]), model.skeleton
+        ),
+    }
+    print(json.dumps(result))
+
+
+if __name__ == "__main__":
+    main()