diff --git a/scripts/fixtures/stub-worker.mjs b/scripts/fixtures/stub-worker.mjs index 1ad2ccc..ce777b6 100644 --- a/scripts/fixtures/stub-worker.mjs +++ b/scripts/fixtures/stub-worker.mjs @@ -1,7 +1,7 @@ // Stands in for kimodo-resident-worker.py in the daemon's tests. It binds port 0, prints the ready line the // daemon waits for, and answers /health and /generate with what the daemon gave it. The --model value picks a // misbehaviour: crash-before-ready, hang-before-ready, sick, slow-health, stubborn, fail-generate, -// refuse-generate, drop-generate, slow-generate, writer (appends to .writes every 50 ms, never answers, ignores SIGTERM), slow-to-stop (no answer; prints a line and exits 1.5 s after SIGTERM), chatty (prints a line the moment it is healthy), encoder (a grandchild, then no answer), graceful (exits 0.5 s after SIGTERM), die-after-ready, bvh-outside (lists a BVH outside the job's directory), bvh-link (lists a symlink as its BVH), bvh-fifo (lists a FIFO as its BVH), plant-link (a normal answer with a file symlink and a directory symlink planted beside the BVH). A request with bvh writes .bvh (motion/motion_NN.bvh for several samples) and lists it. Any other value is just the model the worker serves. +// refuse-generate, drop-generate, slow-generate, writer (appends to .writes every 50 ms, never answers, ignores SIGTERM), slow-to-stop (no answer; prints a line and exits 1.5 s after SIGTERM), chatty (prints a line the moment it is healthy), encoder (a grandchild, then no answer), graceful (exits 0.5 s after SIGTERM), die-after-ready, bvh-outside (lists a BVH outside the job's directory), bvh-link (lists a symlink as its BVH), bvh-fifo (lists a FIFO as its BVH), plant-link (a normal answer with a file symlink and a directory symlink planted beside the BVH), no-loops (a normal answer that leaves out the loops a loop request asked for), bad-loops (loops whose window is not a frame range), short-loops (one window fewer than samples). A loop request is answered with loops, sample k as {first 3 + k, last 223 + k}. A request with bvh writes .bvh (motion/motion_NN.bvh for several samples) and lists it. Any other value is just the model the worker serves. import http from "node:http"; import os from "node:os"; import path from "node:path"; @@ -117,6 +117,11 @@ if (model === "hang-before-ready") { seed: payload.seed, metaSeed: payload.meta?.seed, bvh: payload.bvh, + ...(payload.loop && model !== "no-loops" ? { + loops: Array.from({ length: payload.numSamples - (model === "short-loops" ? 1 : 0) }, (_, k) => (model === "bad-loops" + ? { first: 5.5, last: 9, poseCost: 0, velocityCost: 0 } + : { first: 3 + k, last: 223 + k, poseCost: 0.03, velocityCost: 0.2 })), + } : {}), bvhStandardTpose: payload.bvhStandardTpose, model: model || "stub-model", pid: process.pid, diff --git a/scripts/kimodo-resident-worker.py b/scripts/kimodo-resident-worker.py index e855a32..27d45f6 100644 --- a/scripts/kimodo-resident-worker.py +++ b/scripts/kimodo-resident-worker.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import argparse import json +import math import os import shutil import sys @@ -14,7 +15,7 @@ import torch from huggingface_hub.constants import HF_HOME from kimodo import DEFAULT_MODEL, load_model -from kimodo.constraints import load_constraints_lst +from kimodo.constraints import compute_global_heading, 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 @@ -28,6 +29,7 @@ from kimodo.scripts.generate import ( ) from kimodo.tools import save_json, seed_everything +import loop_window import text_encoding @@ -184,6 +186,20 @@ 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) @@ -319,6 +335,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 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 @@ -347,6 +364,7 @@ def generate(payload): "constraints": constraints_path, "seed": seed, "files": files, + **({"loops": loops} if loops is not None else {}), "encoder": encoder_record, "worker": worker_identity, "timing": { diff --git a/scripts/test-vili-daemon.mjs b/scripts/test-vili-daemon.mjs index 66f9819..0e827c9 100644 --- a/scripts/test-vili-daemon.mjs +++ b/scripts/test-vili-daemon.mjs @@ -1441,6 +1441,81 @@ test("fbx refusals happen before anything is written or spawned", async () => { } }); +test("loop asks for the seamless window, and the window is the FBX's frame range", async () => { + const daemon = await startDaemon(); + try { + const constraints = [{ type: "root2d", frame_indices: [0, 10], smooth_root_2d: [[0, 0], [0, 0.5]] }]; + const answer = await generate(daemon, exportRequest({ clipName: "walk", seed: 7, loop: true, constraints })); + assert.equal(answer.status, 200); + const job = answer.body; + assert.equal(job.request.loop, true); + assert.deepEqual(job.worker.loops, [{ first: 3, last: 223, poseCost: 0.03, velocityCost: 0.2 }]); + assert.deepEqual(job.request.constraints, constraints, "the record keeps inline constraints as sent"); + assert.deepEqual(JSON.parse(readFileSync(path.join(job.artifactDir, "constraints.json"), "utf8")), constraints); + assert.equal(JSON.parse(readFileSync(path.join(job.artifactDir, "request.json"), "utf8")).loop, true, "the worker is asked for the loop"); + assert.equal(job.worker.constraints, path.join(job.artifactDir, "constraints.json"), "the worker still gets the path"); + const [run] = exportStarts(daemon); + assert.deepEqual(run.argv.slice(run.argv.indexOf("--") + 3), ["walk", "--frames", "4", "224"], "Blender's frames count from 1"); + assert.deepEqual(job.export.clips.map((clip) => clip.frames), [[1, 60]], "the clip's frames are what Blender reports"); + + const several = await generate(daemon, exportRequest({ clipName: "pair", numSamples: 2, loop: true })); + assert.equal(several.status, 200); + assert.deepEqual(several.body.worker.loops.map((loop) => [loop.first, loop.last]), [[3, 223], [4, 224]]); + const frames = Object.fromEntries(exportStarts(daemon).slice(1).map((start) => { + const [, , clip, ...rest] = start.argv.slice(start.argv.indexOf("--") + 1); + return [clip, rest]; + })); + assert.deepEqual(frames, { pair_00: ["--frames", "4", "224"], pair_01: ["--frames", "5", "225"] }, "each sample is exported over its own window"); + + const plain = await generate(daemon, exportRequest({ clipName: "plain" })); + assert.equal(plain.status, 200); + assert.equal(plain.body.worker.loops, undefined); + assert.equal(plain.body.request.loop, undefined); + assert.equal(JSON.parse(readFileSync(path.join(plain.body.artifactDir, "request.json"), "utf8")).loop, undefined); + assert.equal(exportStarts(daemon).at(-1).argv.at(-1), "plain", "no window, no --frames"); + const kept = await generate(daemon, exportRequest({ constraints_path: "/some/where.json", constraints: undefined })); + assert.equal(kept.body.request.constraints, "/some/where.json", "a path stays a path"); + } finally { + await daemon.stop(); + } +}); + +test("loop is a boolean and needs fbx, refused before anything is written or spawned", async () => { + const daemon = await startDaemon(); + try { + for (const bad of [{ fbx: false, loop: true }, { fbx: undefined, loop: true }, { loop: "yes canary" }, { loop: 1 }, { loop: null }, { loop: "true" }]) { + const answer = await generate(daemon, exportRequest(bad)); + assert.equal(answer.status, 400, JSON.stringify(bad)); + assert.deepEqual(answer.body, { ok: false, error: "Invalid fbx request." }); + assert.ok(!answer.raw.includes("canary"), "the refusal repeated the input"); + } + assert.equal(existsSync(path.join(daemon.stateRoot, "artifacts")), false, "no artifact directory was made"); + assert.deepEqual(jobFiles(daemon), []); + assert.deepEqual(daemon.spawns(), [], "no worker was started"); + assert.equal(exportStarts(daemon).length, 0); + assert.equal((await generate(daemon, exportRequest({ loop: false }))).status, 200, "loop false is no loop"); + assert.equal((await generate(daemon, { prompt: "a person waves", duration: 2, loop: false })).status, 200); + } finally { + await daemon.stop(); + } +}); + +test("a loop request whose answer has no usable window for a sample fails without exporting that sample", async () => { + for (const [model, exported] of [["no-loops", 0], ["bad-loops", 0], ["short-loops", 1]]) { + const daemon = await startDaemon(); + try { + const answer = await generate(daemon, exportRequest({ model, loop: true, numSamples: 2 })); + assert.equal(answer.status, 502, model); + assert.equal(answer.body.ok, false); + assert.equal(answer.body.error, "fbx-export-failed (no loop window)"); + assert.equal(exportStarts(daemon).length, exported, `${model}: only a sample with a window is exported`); + assert.equal(answer.body.export.clips.length, exported); + } finally { + await daemon.stop(); + } + } +}); + test("the BVH is never given to Blender unless it is a regular file inside the job's directory", async () => { const daemon = await startDaemon(); try { diff --git a/scripts/vili-daemon.mjs b/scripts/vili-daemon.mjs index 908ca98..77e4966 100644 --- a/scripts/vili-daemon.mjs +++ b/scripts/vili-daemon.mjs @@ -592,8 +592,9 @@ async function readInside(file, directory) { // the BVH and writes a private FBX in a directory made for this export alone, and the FBX is then created beside the // BVH exclusively, so nothing planted at either path is read, written through or replaced. The FBX counts only if // Blender exits 0, prints its result line and its file exists. A failure names the exit or the signal and never -// Blender's output. -async function exportClip(bvh, fbx, clipName, timeoutMs, artifactDir) { +// Blender's output. A window {first, last} (BVH frame indices, last inclusive) is handed to Blender as the frame +// range to export, in Blender's frames, which count from 1; the daemon never chooses one. +async function exportClip(bvh, fbx, clipName, timeoutMs, artifactDir, window = null) { const input = await readInside(bvh, artifactDir); if (input === null) return { error: "fbx-export-refused" }; guard(); @@ -606,6 +607,7 @@ async function exportClip(bvh, fbx, clipName, timeoutMs, artifactDir) { const run = await runBlender([ "-b", "--factory-startup", "--python-use-system-env", "--python-exit-code", "1", "--python", exportScript, "--", privateBvh, privateFbx, clipName, + ...(window ? ["--frames", String(window.first + 1), String(window.last + 1)] : []), ], timeoutMs); let info = null; if (run.code === 0 && run.result !== null) { @@ -648,8 +650,9 @@ async function sweepWork() { // Every BVH the worker listed becomes an FBX, or the job fails. A listed BVH must be a regular file inside the job's // artifact directory, and that is decided when its export begins (exportClip), never earlier. Sample N of several is -// _NN. -async function exportListed(listed, artifactDir, clipName, timeoutMs) { +// _NN. With loops (the worker's answer, an array when the request asked for a loop, else undefined), sample k +// is exported over loops[k], and a job whose answer has no usable window for a sample fails. +async function exportListed(listed, artifactDir, clipName, timeoutMs, loops = undefined) { const bvhs = []; for (const file of Array.isArray(listed) ? listed : []) { if (typeof file !== "string" || !file.endsWith(".bvh")) continue; @@ -660,7 +663,11 @@ async function exportListed(listed, artifactDir, clipName, timeoutMs) { const clips = []; for (const [index, bvh] of bvhs.entries()) { const name = bvhs.length === 1 ? clipName : `${clipName}_${String(index).padStart(2, "0")}`; - const done = await exportClip(bvh, `${bvh.slice(0, -".bvh".length)}.fbx`, name, timeoutMs, artifactDir); + const window = loops === undefined ? null : loops[index]; + if (window !== null && !(Number.isInteger(window?.first) && Number.isInteger(window?.last) && window.first >= 0 && window.last > window.first)) { + return { clips, error: "fbx-export-failed (no loop window)" }; + } + const done = await exportClip(bvh, `${bvh.slice(0, -".bvh".length)}.fbx`, name, timeoutMs, artifactDir, window); if (done.error) return { clips, error: done.error }; clips.push({ ...done.clip, file: path.relative(artifactDir, done.clip.file).split(path.sep).join("/") }); } @@ -864,8 +871,10 @@ async function generate(request, response) { // One export plan: fbx implies the BVH and its standard T-pose, which the FBX is made from. const fbx = boolValue(body.fbx); const askedTpose = firstDefined(body.bvhStandardTpose, body.bvh_standard_tpose); - if (fbx && ((askedTpose !== undefined && !boolValue(askedTpose)) - || (body.clipName !== undefined && !(typeof body.clipName === "string" && CLIP_NAME.test(body.clipName))))) { + const loop = body.loop === true; + if ((body.loop !== undefined && typeof body.loop !== "boolean") || (loop && !fbx) + || (fbx && ((askedTpose !== undefined && !boolValue(askedTpose)) + || (body.clipName !== undefined && !(typeof body.clipName === "string" && CLIP_NAME.test(body.clipName)))))) { json(response, 400, { ok: false, error: "Invalid fbx request." }); return; } @@ -934,9 +943,10 @@ async function generate(request, response) { numSamples, diffusion_steps: diffusionSteps, numTransitionFrames: firstDefined(body.numTransitionFrames, body.num_transition_frames), - constraints: constraintsPath, + constraints: body.constraints !== undefined && typeof body.constraints !== "string" ? body.constraints : constraintsPath, seed, fbx: fbx || undefined, + loop: loop || undefined, clipName, cfg: body.cfg, cfgType: firstDefined(body.cfgType, body.cfg_type), @@ -967,6 +977,7 @@ async function generate(request, response) { numTransitionFrames: firstDefined(body.numTransitionFrames, body.num_transition_frames), constraints: constraintsPath, seed, + loop: loop || undefined, cfg: body.cfg, cfgType: firstDefined(body.cfgType, body.cfg_type), cfgWeight: firstDefined(body.cfgWeight, body.cfg_weight), @@ -993,7 +1004,7 @@ async function generate(request, response) { let exported = null; if (fbx && result.ok) { const timeoutMs = Math.min(EXPORT_TIMEOUT_MS, Math.max(1, Number(firstDefined(body.exportTimeoutMs, EXPORT_TIMEOUT_MS)) || EXPORT_TIMEOUT_MS)); - exported = await exportListed(result.body?.files, artifactDir, clipName, timeoutMs); + exported = await exportListed(result.body?.files, artifactDir, clipName, timeoutMs, loop ? (Array.isArray(result.body?.loops) ? result.body.loops : []) : undefined); if (exported.error) result = { ok: false, error: exported.error, body: result.body }; }