64 lines
3.1 KiB
Python
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()
|