Restore runtime.py from Make-It-Animatable

runtime.py

@@ -0,0 +1,107 @@
+import hashlib
+import importlib.metadata
+import json
+import os
+import platform
+import sys
+import time
+import uuid
+from pathlib import Path
+
+ROOT = Path("/data")
+UPSTREAM = Path("/opt/space/upstream")
+MODEL_ROOT = ROOT / "models/make-it-animatable"
+WEIGHTS = {"bw.pth":424693698, "bw_normal.pth":429008352,
+           "joints.pth":525833771, "joints_coarse.pth":424706373,
+           "pose.pth":525848433}
+
+def configure_storage():
+    if not ROOT.is_dir():
+        raise RuntimeError("/data is not mounted. Mount the existing persistent bucket at /data.")
+    for name in ("cache/huggingface","cache/torch","cache/xdg","cache/pip",
+                 "models","inputs","outputs","temp","logs"):
+        (ROOT/name).mkdir(parents=True,exist_ok=True)
+    os.environ.update(HF_HOME=str(ROOT/"cache/huggingface"),
+                      HF_HUB_CACHE=str(ROOT/"cache/huggingface/hub"),
+                      TORCH_HOME=str(ROOT/"cache/torch"),
+                      XDG_CACHE_HOME=str(ROOT/"cache/xdg"),
+                      PIP_CACHE_DIR=str(ROOT/"cache/pip"),
+                      TMPDIR=str(ROOT/"temp"),
+                      GRADIO_TEMP_DIR=str(ROOT/"temp/gradio"))
+    probe = ROOT/"temp"/("write-check-"+uuid.uuid4().hex)
+    with probe.open("x") as stream:
+        stream.write("ok")
+    if probe.read_text() != "ok":
+        raise RuntimeError("/data read/write verification failed")
+    probe.unlink()
+    import tempfile
+    tempfile.tempdir = str(ROOT/"temp")
+    missing = [str(MODEL_ROOT/"output/best/new"/name) for name,size in WEIGHTS.items()
+               if not (MODEL_ROOT/"output/best/new"/name).is_file()
+               or (MODEL_ROOT/"output/best/new"/name).stat().st_size != size]
+    bones = MODEL_ROOT/"data/Mixamo/bones.fbx"
+    if not bones.is_file() or bones.stat().st_size == 0:
+        missing.append(str(bones))
+    if missing:
+        raise RuntimeError("Required cached assets missing or incomplete; startup stopped without "
+                           "overwriting or redownloading assets: "+", ".join(missing))
+    for name in ("data","output"):
+        target = UPSTREAM/name
+        expected = MODEL_ROOT/name
+        if target.is_symlink():
+            if target.resolve() != expected.resolve():
+                raise RuntimeError("Unexpected asset link: "+str(target))
+        elif target.exists():
+            raise RuntimeError("Asset path must be an image-managed symlink: "+str(target))
+        else:
+            target.symlink_to(expected, target_is_directory=True)
+    print("STORAGE_READY /data writable; cached weights reused; downloads=0; installs=0",flush=True)
+
+def startup_checks():
+    import torch
+    import torch_cluster
+    import pytorch3d
+    from pytorch3d.ops import sample_farthest_points
+    versions = {}
+    for name in ("torch","torchvision","torch-cluster","torch-geometric",
+                 "torch-scatter","torch-sparse","torch-spline-conv","pytorch3d","numpy","gradio"):
+        try:
+            versions[name] = importlib.metadata.version(name)
+        except importlib.metadata.PackageNotFoundError:
+            versions[name] = "not required / not installed"
+    report = dict(python=platform.python_version(),packages=versions,cuda=torch.version.cuda,
+                  cuda_available=torch.cuda.is_available(),
+                  gpu=torch.cuda.get_device_name(0) if torch.cuda.is_available() else None)
+    print("RUNTIME_VERSIONS "+json.dumps(report),flush=True)
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA GPU unavailable. Dedicated GPU required; CPU inference is disabled.")
+    if torch.__version__ != "2.1.2+cu121" or torch.version.cuda != "12.1":
+        raise RuntimeError("Runtime Torch/CUDA differs from the pinned build")
+    x = torch.rand((128,3),device="cuda")
+    out = torch_cluster.fps(x,ratio=0.25)
+    _, indices = sample_farthest_points(x[None],K=32)
+    y = torch.ones((32,32),device="cuda") @ torch.ones((32,32),device="cuda")
+    torch.cuda.synchronize()
+    if len(out) != 32 or indices.shape != (1,32) or y[0,0].item() != 32:
+        raise RuntimeError("CUDA binary smoke test failed")
+    report["cuda_execution"] = "PASS: matmul, torch_cluster.fps, pytorch3d FPS"
+    persistence = {}
+    for name in ("cache/huggingface","cache/torch","cache/xdg","cache/pip","inputs","outputs"):
+        marker = ROOT/name/".persistence-check"
+        existed = marker.is_file()
+        if not existed:
+            with marker.open("x") as stream:
+                stream.write(uuid.uuid4().hex)
+        persistence[name] = {"marker":marker.read_text(),"survived_previous_start":existed}
+    report["persistence"] = persistence
+    report["cached_weights"] = WEIGHTS
+    report["startup_time"] = time.time()
+    report["startup_id"] = uuid.uuid4().hex
+    report["outputs_present"] = sorted(str(p.relative_to(ROOT)) for p in (ROOT/"outputs").glob("*"))
+    previous = ROOT/"logs/last-startup.json"
+    if previous.is_file():
+        report["previous_startup"] = json.loads(previous.read_text()).get("startup_id")
+    (ROOT/"logs"/(report["startup_id"]+".json")).write_text(json.dumps(report,indent=2))
+    previous.write_text(json.dumps(report,indent=2))
+    print("CUDA_EXECUTION_PASS "+json.dumps(report),flush=True)
+    return report