cozyclay 2.1.1 → 3.0.1

package/tools/ardy/cclay_worker.py

@@ -1,304 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: GPL-3.0-or-later
-# SPDX-FileCopyrightText: CozyClay contributors
-#
-# cclay-owned script. Source of truth lives in the CozyClay repo at
-# tools/ardy/cclay_worker.py.
-"""Persistent local generation worker for the CozyClay ARDY bridge.
-
-Loading the motion model and compiling its GPU kernels costs minutes on MPS,
-while the sampling itself takes seconds. Spawning a fresh python per request
-(the remote-box contract run-local.mjs mirrors) pays that cost every time, so
-this worker pays it ONCE: it loads the model, runs a tiny no-text warmup to
-force kernel compilation, and then serves generation jobs from the warm
-process. run-local.mjs prefers the worker and silently falls back to a direct
-spawn when it is not running, so the worker is an accelerator, never a
-dependency.
-
-Protocol (loopback TCP, one connection per job, one job at a time):
-
-    client -> worker   one JSON line
-                       {"mode": "ping"}
-                       {"mode": "single"|"sequence"|"edit", "argv": [...]}
-    worker -> client   the generator's stdout lines, verbatim, then
-                       "worker: pong"  (ping)  or  "worker: exit <code>"
-
-`argv` is exactly the CLI argument vector of the corresponding
-cclay_*_generate/edit script; the worker calls that script's main() in
-process with the cached model injected, so worker and direct-spawn runs share
-one code path and cannot drift.
-
-env:
-    CCLAY_ARDY_WORKER_PORT  listen port (default 9552, loopback only)
-    CHECKPOINTS_DIR / TEXT_ENCODER_URL / TEXT_ENCODERS_DIR as usual
-"""
-
-import importlib
-import json
-import os
-import socket
-import sys
-import traceback
-
-HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, HERE)
-
-MODULES = {
-    "single": "cclay_constrained_generate",
-    "sequence": "cclay_sequence_generate",
-    "edit": "cclay_motion_edit",
-}
-
-
-class _Tee:
-    """Mirror generator stdout to the job's socket AND the worker console.
-
-    A vanished client must not kill the job: the generation keeps running,
-    the npz still lands on disk, and only the socket copy is dropped.
-    """
-
-    def __init__(self, sock_file):
-        self._sock_file = sock_file
-
-    def write(self, data):
-        sys.__stdout__.write(data)
-        try:
-            self._sock_file.write(data)
-            self._sock_file.flush()
-        except OSError:
-            pass
-
-    def flush(self):
-        sys.__stdout__.flush()
-
-
-def requested_model_name(mode, argv, resolve_model_name, default_model, checkpoints_dir):
-    """The model a job's argv asks for (edit has no --model flag)."""
-    name = default_model
-    if mode != "edit":
-        for index, arg in enumerate(argv):
-            if arg == "--model" and index + 1 < len(argv):
-                name = argv[index + 1]
-    return resolve_model_name(name, checkpoints_dir=checkpoints_dir)
-
-
-def fix_text_encoder_device(model, device):
-    """Move an in-process LLM2Vec encoder onto the generation device, fast.
-
-    ARDY's load_text_encoder ends with `.to(device or "cuda"/"cpu")`, so on a
-    Mac the fallback encoder lands on CPU even when TEXT_ENCODER_DEVICE says
-    mps. On MPS the dtype matters as much as the device: bf16 matmuls are
-    emulated (a single prompt costs ~30-50 s) while fp16 is native, so fp16
-    is tried first and a canary encode guards against fp16 overflow — NaNs
-    demote back to bf16. Only the real in-process encoder has a `.model`;
-    the API client stub is left alone.
-    """
-    encoder = getattr(model, "text_encoder", None)
-    if encoder is None or not hasattr(encoder, "model") or device == "cpu":
-        return
-    import torch
-
-    for dtype in (torch.float16, torch.bfloat16):
-        try:
-            encoder.to(device=device, dtype=dtype)
-            canary, _ = encoder(["a person walks forward"])
-            if bool(torch.isfinite(canary.float()).all()):
-                print(f"worker: text encoder on {device} as {dtype}", flush=True)
-                return
-            print(f"worker: text encoder {dtype} produced non-finite values; trying next dtype", flush=True)
-        except Exception as exc:
-            print(f"worker: text encoder {dtype} on {device} failed ({exc}); trying next dtype", flush=True)
-    try:  # both GPU dtypes failed: back to the slow-but-correct CPU encoder
-        encoder.to(device="cpu", dtype=torch.bfloat16)
-        print("worker: text encoder stays on cpu", flush=True)
-    except Exception as exc:
-        print(f"worker: text encoder left as loaded ({exc})", flush=True)
-
-
-class _CachingEncoder:
-    """Per-prompt embedding cache in front of the in-process encoder.
-
-    Authoring iterates on ONE prompt (new seed, new pose pin, new waypoints),
-    so most requests re-encode text the worker has already seen. Embeddings
-    are deterministic per (model, dtype) — exactly the cache key's tag — and
-    a few KB each, so they live in memory and under ~/.cozyclay/embed-cache
-    across worker restarts. Misses are timed on the console; hits are free.
-    """
-
-    def __init__(self, inner, cache_dir, tag):
-        import os
-
-        self._inner = inner
-        self._dir = cache_dir
-        self._tag = tag
-        self._mem = {}
-        os.makedirs(cache_dir, exist_ok=True)
-
-    def _key(self, text):
-        import hashlib
-
-        return hashlib.sha256(f"{self._tag}\x00{text}".encode("utf-8")).hexdigest()
-
-    def _encode_one(self, text):
-        import os
-        import time
-
-        import numpy as np
-
-        key = self._key(text)
-        if key in self._mem:
-            return self._mem[key]
-        path = os.path.join(self._dir, f"{key}.npy")
-        if os.path.isfile(path):
-            try:
-                embedding = np.load(path)
-                self._mem[key] = embedding
-                return embedding
-            except Exception:
-                pass  # unreadable cache entry: re-encode below
-        started = time.time()
-        tensor, _ = self._inner([text])
-        embedding = tensor[0].detach().float().cpu().numpy()
-        print(f"worker: text encode took {time.time() - started:.1f}s (cache miss)", flush=True)
-        self._mem[key] = embedding
-        try:
-            np.save(path, embedding)
-        except Exception:
-            pass  # cache write failure must never fail the job
-        return embedding
-
-    def __call__(self, texts):
-        import numpy as np
-        import torch
-
-        is_string = isinstance(texts, str)
-        text_list = [texts] if is_string else list(texts)
-        stacked = np.stack([self._encode_one(text) for text in text_list])
-        device = getattr(self._inner, "_device", "cpu")
-        encoded = torch.from_numpy(stacked).to(device)
-        lengths = [1] * len(text_list)
-        if is_string:
-            return encoded[0], 1
-        return encoded, lengths
-
-    def __getattr__(self, name):
-        return getattr(self._inner, name)
-
-
-def warmup(model, device):
-    """One throwaway sampling call so kernel compilation happens off the clock.
-
-    Zero text features skip the text encoder entirely — the warmup must not
-    drag the 16 GB Llama encoder into RAM just to compile Metal shaders.
-    """
-    import torch
-    from ardy.motion_rep.tools import length_to_mask
-
-    try:
-        frames = int(model.gen_horizon_len)
-        lengths = torch.tensor([frames], device=device)
-        with torch.no_grad():
-            model(
-                [""],
-                frames,
-                num_denoising_steps=1,
-                pad_mask=length_to_mask(lengths),
-                first_heading_angle=torch.zeros(1, device=device),
-                motion_mask=None,
-                observed_motion=None,
-                cfg_weight=2.0,
-                text_feat=torch.zeros(1, 1, 4096, device=device),
-                text_pad_mask=torch.ones(1, 1, device=device, dtype=torch.bool),
-                progress_bar=lambda iterable: iterable,
-            )
-        print("worker: warmup complete", flush=True)
-    except Exception as exc:  # non-fatal: first real job just pays the compile
-        print(f"worker: warmup skipped ({exc})", flush=True)
-
-
-def main():
-    port = int(os.environ.get("CCLAY_ARDY_WORKER_PORT", "9552"))
-
-    single = importlib.import_module(MODULES["single"])
-    device = single.cclay_pick_device()
-    single.cclay_mps_compat(device)
-
-    from ardy.model import DEFAULT_MODEL, load_model
-    from ardy.model.loading import get_env_var
-    from ardy.model.registry import resolve_model_name
-
-    checkpoints_dir = get_env_var("CHECKPOINTS_DIR", None)
-    cached_name = resolve_model_name(DEFAULT_MODEL, checkpoints_dir=checkpoints_dir)
-    print(f"worker: loading {cached_name} on {device} ...", flush=True)
-    cached_model = load_model(cached_name, device=device, checkpoints_dir=checkpoints_dir)
-    cached_model.eval()
-    fix_text_encoder_device(cached_model, device)
-    encoder = getattr(cached_model, "text_encoder", None)
-    if encoder is not None and hasattr(encoder, "model"):
-        try:
-            param = next(encoder.model.parameters())
-            tag = f"llm2vec:{param.dtype}:{param.device.type}"
-        except Exception:
-            tag = "llm2vec:unknown"
-        cache_dir = os.path.join(os.path.expanduser("~"), ".cozyclay", "embed-cache")
-        cached_model.text_encoder = _CachingEncoder(encoder, cache_dir, tag)
-        print(f"worker: embedding cache at {cache_dir} ({tag})", flush=True)
-    warmup(cached_model, device)
-
-    server = socket.create_server(("127.0.0.1", port))
-    print(f"worker: ready on 127.0.0.1:{port} (model {cached_name}, device {device})", flush=True)
-
-    while True:
-        conn, _ = server.accept()
-        with conn:
-            rfile = conn.makefile("r", encoding="utf-8")
-            wfile = conn.makefile("w", encoding="utf-8")
-            try:
-                job = json.loads(rfile.readline())
-            except (ValueError, OSError):
-                continue
-            mode = job.get("mode") if isinstance(job, dict) else None
-            if mode == "ping":
-                try:
-                    wfile.write("worker: pong\n")
-                    wfile.flush()
-                except OSError:
-                    pass
-                continue
-            argv = job.get("argv") if isinstance(job, dict) else None
-            if mode not in MODULES or not isinstance(argv, list) or not all(isinstance(a, str) for a in argv):
-                try:
-                    wfile.write("worker: exit 2\n")
-                    wfile.flush()
-                except OSError:
-                    pass
-                continue
-
-            module = importlib.import_module(MODULES[mode])
-            code = 0
-            old_stdout = sys.stdout
-            sys.stdout = _Tee(wfile)
-            try:
-                wanted = requested_model_name(mode, argv, resolve_model_name, DEFAULT_MODEL, checkpoints_dir)
-                if wanted != cached_name:
-                    print(f"worker: switching model {cached_name} -> {wanted}", flush=True)
-                    cached_model = load_model(wanted, device=device, checkpoints_dir=checkpoints_dir)
-                    cached_model.eval()
-                    cached_name = wanted
-                module.main(argv, preloaded_model=cached_model)
-            except SystemExit as exit_exc:
-                code = exit_exc.code if isinstance(exit_exc.code, int) else 1
-            except Exception:
-                traceback.print_exc(file=sys.stdout)
-                code = 1
-            finally:
-                sys.stdout = old_stdout
-            try:
-                wfile.write(f"worker: exit {code}\n")
-                wfile.flush()
-            except OSError:
-                pass
-
-
-if __name__ == "__main__":
-    main()

package/tools/ardy/dump-npz.py

@@ -1,205 +0,0 @@
-#!/usr/bin/env python3
-"""Dump one frame of an ARDY motion npz as a JSON document on stdout.
-
-Runs under the ARDY box python (``~/ardy/.venv/bin/python``); it needs numpy
-but no torch. Usage::
-
-    python dump-npz.py <motion.npz> <frame>
-
-Emitted document (schema "ardy.frame.v1"):
-
-    {
-      "schema": "ardy.frame.v1",
-      "source_npz": "<basename of the npz>",
-      "frame": <requested frame index>,
-      "frames_total": <frame count of the motion>,
-      "members": {"<member name>": {"shape": [...], "dtype": "..."}, ...},
-      "local_rot_mats": [ 27 x [[3],[3],[3]] row-major ],
-      "posed_joints":   [ 27 x [x, y, z] ]
-    }
-
-Member names and axis order match the ARDY archives exactly as consumed by
-CozyClay (blender-addon/cclay): ``local_rot_mats[frame, joint]`` is the
-per-joint LOCAL rotation matrix in cskel27 order (Hips = joint 0), and
-``posed_joints[frame, joint]`` the joint center position. ``foot_contacts``
-is optional and reported only when present. ``members`` lists every array in
-the archive (sorted), including carried members like ``fps``, so a consumer
-can see the exact contents without re-opening the npz.
-
-Before emitting, every one of the 27 rotation matrices of the requested frame
-is checked for finite entries, orthonormal rows/columns, and determinant +1,
-with the same tolerance CozyClay's motion_retarget uses
-(ROTATION_MATRIX_TOLERANCE = 1e-3 on squared norms, pairwise dots, and
-determinant). On the first violation the matrix index and the failing
-quantity are printed to stderr, the exit status is 1, and NO JSON is emitted:
-a garbage rotation must not silently reach a downstream consumer.
-"""
-
-import json
-import math
-import sys
-
-import numpy as np
-
-SCHEMA = "ardy.frame.v1"
-JOINTS = 27
-# Mirror of CozyClay motion_retarget.ROTATION_MATRIX_TOLERANCE: applies to
-# squared row/column norms, pairwise dot products, and determinant, in
-# float64. float32 ARDY serialization noise stays far below 1e-3.
-TOLERANCE = 1e-3
-
-
-def _fail(message):
-    """Report a fatal error on stderr and exit 1 without emitting JSON."""
-    print(f"dump-npz: {message}", file=sys.stderr)
-    sys.exit(1)
-
-
-def _check_rotation(matrix, frame, joint):
-    """Return a list of violation strings for one 3x3 matrix, or [] if clean."""
-    violations = []
-    rows = [[float(matrix[i][j]) for j in range(3)] for i in range(3)]
-    columns = [[rows[i][j] for i in range(3)] for j in range(3)]
-    prefix = f"frame {frame} joint {joint}"
-    for label, vectors in (("row", rows), ("column", columns)):
-        for index, vector in enumerate(vectors):
-            squared_norm = sum(c * c for c in vector)
-            if not math.isfinite(squared_norm):
-                violations.append(
-                    f"{prefix} {label} {index} has non-finite entries"
-                )
-            elif abs(squared_norm - 1.0) > TOLERANCE:
-                violations.append(
-                    f"{prefix} {label} {index} squared norm {squared_norm:.6f} "
-                    f"differs from 1.0 by more than {TOLERANCE}"
-                )
-        for first, second in ((0, 1), (0, 2), (1, 2)):
-            dot = sum(
-                vectors[first][i] * vectors[second][i] for i in range(3)
-            )
-            if not math.isfinite(dot):
-                violations.append(
-                    f"{prefix} {label} {first}/{second} dot is non-finite"
-                )
-            elif abs(dot) > TOLERANCE:
-                violations.append(
-                    f"{prefix} {label} {first}/{second} dot {dot:.6f} "
-                    f"exceeds {TOLERANCE} in magnitude"
-                )
-    determinant = (
-        rows[0][0] * (rows[1][1] * rows[2][2] - rows[1][2] * rows[2][1])
-        - rows[0][1] * (rows[1][0] * rows[2][2] - rows[1][2] * rows[2][0])
-        + rows[0][2] * (rows[1][0] * rows[2][1] - rows[1][1] * rows[2][0])
-    )
-    if not math.isfinite(determinant):
-        violations.append(f"{prefix} determinant is non-finite")
-    elif abs(determinant - 1.0) > TOLERANCE:
-        violations.append(
-            f"{prefix} determinant {determinant:.6f} differs from 1.0 "
-            f"by more than {TOLERANCE}"
-        )
-    return violations
-
-
-def main(argv):
-    if len(argv) != 3:
-        _fail(f"usage: {argv[0]} <motion.npz> <frame>")
-    npz_path, frame_text = argv[1], argv[2]
-    try:
-        frame = int(frame_text)
-    except ValueError:
-        _fail(f"frame must be an integer, got {frame_text!r}")
-    if frame < 0:
-        _fail(f"frame must be >= 0, got {frame}")
-
-    try:
-        data = np.load(npz_path, allow_pickle=False)
-    except Exception as exc:
-        _fail(f"cannot load {npz_path}: {exc}")
-
-    try:
-        if "local_rot_mats" not in data.files:
-            _fail(f"{npz_path} has no 'local_rot_mats' member")
-        if "posed_joints" not in data.files:
-            _fail(f"{npz_path} has no 'posed_joints' member")
-        rotations = data["local_rot_mats"]
-        joints = data["posed_joints"]
-        if rotations.shape[1:] != (JOINTS, 3, 3):
-            _fail(
-                f"local_rot_mats must have shape (F, {JOINTS}, 3, 3), "
-                f"got {rotations.shape}"
-            )
-        if joints.shape[1:] != (JOINTS, 3):
-            _fail(
-                f"posed_joints must have shape (F, {JOINTS}, 3), "
-                f"got {joints.shape}"
-            )
-        if joints.shape[0] != rotations.shape[0]:
-            _fail(
-                "local_rot_mats and posed_joints frame counts differ "
-                f"({rotations.shape[0]} vs {joints.shape[0]})"
-            )
-        frames_total = int(rotations.shape[0])
-        if frames_total < 1:
-            _fail(f"motion has no frames ({frames_total})")
-        if frame >= frames_total:
-            _fail(
-                f"frame {frame} out of range: motion has {frames_total} frames"
-            )
-
-        # Validate every rotation matrix of the requested frame BEFORE any
-        # JSON is written; a violation aborts with a stderr report and a
-        # non-zero exit status instead of emitting garbage.
-        violations = []
-        for joint in range(JOINTS):
-            violations.extend(_check_rotation(rotations[frame][joint], frame, joint))
-        for joint in range(JOINTS):
-            for axis in range(3):
-                if not math.isfinite(float(joints[frame][joint][axis])):
-                    violations.append(
-                        f"frame {frame} joint {joint} axis {axis} "
-                        "position is not finite"
-                    )
-        if violations:
-            for violation in violations[:20]:
-                print(f"dump-npz: VIOLATION {violation}", file=sys.stderr)
-            if len(violations) > 20:
-                print(
-                    f"dump-npz: ... and {len(violations) - 20} more",
-                    file=sys.stderr,
-                )
-            _fail(
-                f"refusing to emit {npz_path} frame {frame}: "
-                f"{len(violations)} rotation/position check(s) failed"
-            )
-
-        members = {}
-        for name in sorted(data.files):
-            member = data[name]
-            members[name] = {
-                "shape": [int(size) for size in member.shape],
-                "dtype": str(member.dtype),
-            }
-
-        document = {
-            "schema": SCHEMA,
-            "source_npz": npz_path.rsplit("/", 1)[-1],
-            "frame": frame,
-            "frames_total": frames_total,
-            "members": members,
-            "local_rot_mats": [
-                [[float(value) for value in row] for row in rotations[frame][joint]]
-                for joint in range(JOINTS)
-            ],
-            "posed_joints": [
-                [float(value) for value in joints[frame][joint]]
-                for joint in range(JOINTS)
-            ],
-        }
-        print(json.dumps(document, indent=2))
-    finally:
-        data.close()
-
-
-if __name__ == "__main__":
-    main(sys.argv)