Files
AIGame/Tools/run_littlebear_tokenrig.py

64 lines
3.1 KiB
Python

"""Run the LittleBear fixed-skeleton skin-only job in the existing DGXA container."""
from pathlib import Path
import hashlib
import json
import subprocess
ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "Tools/output/littlebear_studio/rig_mixamo41"
JOB_NAME = "littlebear_mixamo41_20261026_a1"
REMOTE = "/home/spark/SkinTokens/tmp_jobs/" + JOB_NAME
CONTAINER = "/workspace/SkinTokens/tmp_jobs/" + JOB_NAME
def run(args, capture=False):
print("RUN", args[0], flush=True)
return subprocess.run(args, check=True, capture_output=capture, text=capture,
encoding="utf-8" if capture else None)
def sha(path):
return hashlib.sha256(path.read_bytes()).hexdigest()
def main():
source = OUT / "source.glb"
if not source.is_file():
raise SystemExit("Prepare the fixed Mixamo skeleton and source.glb first")
script = OUT / "run_tokenrig.sh"
script.write_text("\n".join([
"#!/usr/bin/env bash", "set -euo pipefail",
"export PATH=/opt/py310/bin:$PATH",
"export PYTHONPATH=/workspace/SkinTokens:/workspace/blender_dev/build_bpy/bin",
"export BLENDER_SYSTEM_SCRIPTS=/workspace/blender_dev/blender/scripts",
"export BLENDER_SYSTEM_DATAFILES=/workspace/blender_dev/blender/release/datafiles",
"export PYTHONUNBUFFERED=1",
"python3 demo.py --input " + CONTAINER + "/source.glb --output " + CONTAINER + "/rigged_raw.glb --use_skeleton --use_transfer --skeleton_template original --rig_cls vroid",
"" ]), encoding="utf-8", newline="\n")
ssh = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=8", "DGXA"]
run(ssh + ["mkdir -p " + REMOTE])
existing = run(ssh + ["if test -s " + REMOTE + "/rigged_raw.glb; then sha256sum " + REMOTE + "/source.glb; fi"], capture=True).stdout.split()
if existing and existing[0] != sha(source):
raise RuntimeError("Existing job belongs to a different input; use a new job name")
if not existing:
run(["scp", str(source), str(script), "DGXA:" + REMOTE + "/"])
remote_hash = run(ssh + ["sha256sum " + REMOTE + "/source.glb"], capture=True).stdout.split()[0]
if remote_hash != sha(source):
raise RuntimeError("Remote input checksum mismatch")
if not existing:
run(ssh + ["docker exec -w /workspace/SkinTokens tokenrig bash " + CONTAINER + "/run_tokenrig.sh > " + REMOTE + "/tokenrig.log 2>&1"])
for filename in ["rigged_raw.glb", "tokenrig.log"]:
run(["scp", "DGXA:" + REMOTE + "/" + filename, str(OUT / filename)])
remote_output = run(ssh + ["sha256sum " + REMOTE + "/rigged_raw.glb"], capture=True).stdout.split()[0]
if remote_output != sha(OUT / "rigged_raw.glb"):
raise RuntimeError("Remote output checksum mismatch")
(OUT / "remote_receipt.json").write_text(json.dumps(dict(job=JOB_NAME,
input_sha256=remote_hash, output_sha256=remote_output,
flags=["--use_skeleton", "--use_transfer", "--skeleton_template", "original", "--rig_cls", "vroid"],
reference_weights_used=False), indent=2), encoding="utf-8")
print("TOKENRIG_WEIGHTS_READY", flush=True)
if __name__ == "__main__":
main()