From 00ccd57ed54c8b64d1c3153ed46deafd0755b2a0 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Mon, 24 Aug 2026 07:23:11 +0100 Subject: [PATCH] Capture upstream multi-prompt fixtures --- CMakeLists.txt | 5 + .../dump_kimodo_multiprompt_reference.py | 173 ++++++++++++++++++ src/encode.cpp | 32 ++++ 3 files changed, 210 insertions(+) create mode 100644 reference/dump_kimodo_multiprompt_reference.py create mode 100644 src/encode.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 0950898..8b075d8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -53,6 +53,11 @@ if(EXISTS "${KIMODO_GGML_SOURCE_DIR}/CMakeLists.txt") target_compile_definitions(kmd-sample-fixture PRIVATE KIMODO_HAVE_GGML=1 KIMODO_HAVE_GGML_VULKAN=1) add_executable(kmd-generate src/generate.cpp) target_link_libraries(kmd-generate PRIVATE kimodo) + add_executable(kmd-encode src/encode.cpp) + target_sources(kmd-encode PRIVATE src/llm_text_encoder.cpp src/llm_tokenizer.cpp) + target_include_directories(kmd-encode PRIVATE src) + target_link_libraries(kmd-encode PRIVATE ggml ggml-vulkan) + target_compile_definitions(kmd-encode PRIVATE KIMODO_HAVE_GGML=1 KIMODO_HAVE_GGML_VULKAN=1) add_executable(kimodo-llm-layer-parity tests/llm_layer_parity.cpp) target_link_libraries(kimodo-llm-layer-parity PRIVATE ggml ggml-vulkan) target_include_directories(kimodo-llm-layer-parity PRIVATE src) diff --git a/reference/dump_kimodo_multiprompt_reference.py b/reference/dump_kimodo_multiprompt_reference.py new file mode 100644 index 0000000..1a99926 --- /dev/null +++ b/reference/dump_kimodo_multiprompt_reference.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""Capture an upstream Kimodo multi-prompt transition fixture. + +The script deliberately accepts precomputed F32 LLM2Vec embeddings. This +keeps the 8B text model out of the PyTorch process while the upstream Kimodo +motion code remains the authority for the conditioned transition and DDIM +trajectory. Record the provenance of every embedding in ``--embedding-note``. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import platform +import subprocess +import sys +from pathlib import Path +from typing import Any + +import numpy as np +import torch + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def options() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--upstream", required=True, type=Path) + parser.add_argument("--checkpoint-dir", required=True, type=Path) + parser.add_argument("--prompt", action="append", required=True) + parser.add_argument("--frames", action="append", required=True, type=int) + parser.add_argument("--embedding", action="append", required=True, type=Path) + parser.add_argument("--embedding-note", action="append", default=[]) + parser.add_argument("--transition-frames", default=5, type=int) + parser.add_argument("--steps", default=2, type=int) + parser.add_argument("--seed", default=42, type=int) + parser.add_argument("--device", default="cpu") + parser.add_argument("--output", required=True, type=Path) + return parser.parse_args() + + +def as_f32(value: torch.Tensor) -> np.ndarray: + return value.detach().to(device="cpu", dtype=torch.float32).contiguous().numpy() + + +class CapturedEmbeddings: + """Text-encoder-compatible source of one already-captured vector per prompt.""" + + def __init__(self, prompts: list[str], paths: list[Path]) -> None: + self.values: dict[str, torch.Tensor] = {} + for prompt, path in zip(prompts, paths): + raw = np.fromfile(path, dtype=" None: + args = options() + if len(args.prompt) != len(args.frames) or len(args.prompt) != len(args.embedding): + raise SystemExit("--prompt, --frames, and --embedding must occur equally often") + if len(args.prompt) < 2: + raise SystemExit("a multi-prompt fixture needs at least two prompts") + if any(frame <= args.transition_frames for frame in args.frames): + raise SystemExit("every segment must be longer than --transition-frames") + upstream, checkpoint = args.upstream.resolve(), args.checkpoint_dir.resolve() + if not (upstream / "kimodo").is_dir() or not (checkpoint / "Kimodo-SMPLX-RP-v1" / "config.yaml").is_file(): + raise SystemExit("expected an upstream checkout and local Kimodo-SMPLX-RP-v1 checkpoint") + import os + os.environ["CHECKPOINT_DIR"] = str(checkpoint) + sys.path.insert(0, str(upstream)) + from kimodo import load_model # pylint: disable=import-outside-toplevel + + torch.manual_seed(args.seed) + encoder = CapturedEmbeddings(args.prompt, args.embedding) + model, resolved = load_model("kimodo-smplx-rp", device=args.device, + text_encoder=encoder, return_resolved_name=True) + calls: list[dict[str, list[torch.Tensor] | torch.Tensor]] = [] + inverse_inputs: list[torch.Tensor] = [] + original_step = model.denoising_step + + def capture_step(*values: Any, **kwargs: Any) -> torch.Tensor: + if not calls or len(calls[-1]["input"]) == args.steps: + calls.append({ + "input": [], "output": [], + "pad_mask": values[1].detach().clone(), + "text_features": values[2].detach().clone(), + "text_pad_mask": values[3].detach().clone(), + "first_heading_angle": values[5].detach().clone(), + "motion_mask": values[6].detach().clone(), + "observed_motion": values[7].detach().clone(), + }) + call = calls[-1] + call["input"].append(values[0].detach().clone()) # type: ignore[index] + result = original_step(*values, **kwargs) + call["output"].append(result.detach().clone()) # type: ignore[index] + return result + + original_inverse = model.motion_rep.inverse + + def capture_inverse(motion: torch.Tensor, *values: Any, **kwargs: Any): + inverse_inputs.append(motion.detach().clone()) + return original_inverse(motion, *values, **kwargs) + + model.motion_rep.inverse = capture_inverse + model.denoising_step = capture_step + try: + output = model( + args.prompt, num_frames=args.frames, multi_prompt=True, + num_denoising_steps=args.steps, num_samples=1, + cfg_type="separated", cfg_weight=[2.0, 2.0], + num_transition_frames=args.transition_frames, post_processing=False, + return_numpy=True, progress_bar=lambda values: values, + ) + finally: + model.denoising_step = original_step + model.motion_rep.inverse = original_inverse + if len(calls) != len(args.prompt) or not inverse_inputs: + raise RuntimeError(f"expected {len(args.prompt)} segment trajectories, got {len(calls)}") + + arrays: dict[str, np.ndarray] = {"stitched_motion_rep": as_f32(inverse_inputs[-1])} + for index, call in enumerate(calls): + prefix = f"segment_{index:02d}_" + for name in ("pad_mask", "text_features", "text_pad_mask", "first_heading_angle", "motion_mask", "observed_motion"): + arrays[prefix + name] = as_f32(call[name]) # type: ignore[arg-type,index] + for step, (sample_in, sample_out) in enumerate(zip(call["input"], call["output"])): # type: ignore[arg-type,index] + arrays[f"{prefix}sampling_input_{step:03d}"] = as_f32(sample_in) + arrays[f"{prefix}sampling_output_{step:03d}"] = as_f32(sample_out) + for key, value in output.items(): + if isinstance(value, np.ndarray): + arrays["motion_" + key] = value + args.output.parent.mkdir(parents=True, exist_ok=True) + archive = args.output.with_suffix(".npz") + np.savez_compressed(archive, **arrays) + try: + revision = subprocess.check_output( + ["git", "-c", "safe.directory=*", "-C", str(upstream), "rev-parse", "HEAD"], text=True).strip() + except subprocess.CalledProcessError: + revision = "unknown" + notes = args.embedding_note + ["unspecified"] * (len(args.prompt) - len(args.embedding_note)) + metadata = { + "fixture_format": 1, "upstream_revision": revision, "resolved_model": resolved, + "prompts": args.prompt, "frames_per_segment": args.frames, + "transition_frames": args.transition_frames, "diffusion_steps": args.steps, + "seed": args.seed, "cfg_type": "separated", "cfg_weight": [2.0, 2.0], + "post_processing": False, "device": args.device, "torch": torch.__version__, + "python": platform.python_version(), "embedding_sources": [ + {"path": str(path), "sha256": sha256(path), "note": note} + for path, note in zip(args.embedding, notes) + ], "checkpoint": { + "revision": (checkpoint / "Kimodo-SMPLX-RP-v1" / "REVISION").read_text(encoding="utf-8").strip(), + "model_safetensors_sha256": sha256(checkpoint / "Kimodo-SMPLX-RP-v1" / "model.safetensors"), + }, "npz_sha256": sha256(archive), + } + args.output.with_suffix(".json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/src/encode.cpp b/src/encode.cpp new file mode 100644 index 0000000..07b6171 --- /dev/null +++ b/src/encode.cpp @@ -0,0 +1,32 @@ +// Small developer utility for capturing a portable F32 LLM2Vec embedding. +// It uses the same serial GGML text session as the public prompt API. +#include "llm_text_encoder.hpp" + +#include +#include +#include +#include +#include + +int main(int argc, char **argv) try { + if (argc != 4) { + std::cerr << "usage: " << argv[0] << " TEXT_BUNDLE PROMPT.txt OUTPUT.f32\n"; + return 2; + } + std::ifstream prompt_file(argv[2]); + const std::string prompt{std::istreambuf_iterator(prompt_file), {}}; + if (!prompt_file && prompt.empty()) throw std::runtime_error("cannot read prompt"); + auto encoder = kimodo::detail::llm_text_encoder::load(argv[1]); + if (!encoder) throw std::runtime_error(encoder.error()); + auto embedding = (*encoder)->encode(prompt); + if (!embedding) throw std::runtime_error(embedding.error()); + std::ofstream output(argv[3], std::ios::binary | std::ios::trunc); + if (!output) throw std::runtime_error("cannot open output"); + output.write(reinterpret_cast(embedding->data()), + static_cast(embedding->size() * sizeof(float))); + if (!output) throw std::runtime_error("cannot write output"); + return 0; +} catch (const std::exception &error) { + std::cerr << error.what() << '\n'; + return 1; +}