Spaces:
Running
Running
| 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 | |