Restore patch_runtime.py from Make-It-Animatable

patch_runtime.py

@@ -0,0 +1,100 @@
+from pathlib import Path
+import os
+import subprocess
+
+# The upstream checkout lives on persistent storage and therefore retains edits
+# across container restarts. Restore only the files managed by this patcher from
+# its pinned Git commit before applying changes, making this script idempotent.
+for relative_path in (
+    "util/dataset_mixamo_additional.py",
+    "model.py",
+    "models_ae.py",
+    "app.py",
+):
+    pristine = subprocess.check_output(
+        ["git", "-C", "app", "show", f"HEAD:{relative_path}"]
+    )
+    Path("app", relative_path).write_bytes(pristine)
+
+# ── 1. dataset_mixamo_additional.py — optional extended skeleton guard ─────────
+path = Path("app/util/dataset_mixamo_additional.py")
+if path.exists():
+    source = path.read_text()
+    old = """TEMPLATE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "../data/Mixamo/bones_vroid.fbx"))
+KINEMATIC_TREE = get_kinematic_tree(TEMPLATE_PATH, BONES_IDX_DICT)
+assert len(KINEMATIC_TREE) == len(MIXAMO_JOINTS)
+"""
+    new = """TEMPLATE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "../data/Mixamo/bones_vroid.fbx"))
+if os.path.isfile(TEMPLATE_PATH):
+    KINEMATIC_TREE = get_kinematic_tree(TEMPLATE_PATH, BONES_IDX_DICT)
+    assert len(KINEMATIC_TREE) == len(MIXAMO_JOINTS)
+else:
+    # The optional extended skeleton lives in the restricted Mixamo dataset.
+    # Standard humanoid inference does not use it (joints_additional=False).
+    KINEMATIC_TREE = dataset_mixamo.KINEMATIC_TREE
+"""
+    if old in source:
+        path.write_text(source.replace(old, new))
+        print("[dataset_mixamo_additional.py] extended-skeleton guard patched")
+    else:
+        print("[dataset_mixamo_additional.py] already patched or block not found")
+
+# ── fps fallback: replaces torch_cluster.fps with pytorch3d equivalent ─────────
+FPS_FALLBACK = """try:
+    from torch_cluster import fps
+except ImportError:
+    import numpy as np
+    from pytorch3d.ops import sample_farthest_points
+
+    def fps(x, batch, ratio):
+        import torch
+        sampled = []
+        for batch_id in torch.unique(batch, sorted=True):
+            global_idx = torch.nonzero(batch == batch_id, as_tuple=False).flatten()
+            points = x[global_idx].unsqueeze(0)
+            count = max(1, int(np.ceil(points.shape[1] * ratio)))
+            _, local_idx = sample_farthest_points(points, K=count, random_start_point=True)
+            sampled.append(global_idx[local_idx[0]])
+        return torch.cat(sampled)
+"""
+CLUSTER_IMPORT = "from torch_cluster import fps"
+
+# ── 2. model.py ───────────────────────────────────────────────────────────────
+model_path = Path("app/model.py")
+if model_path.exists():
+    src = model_path.read_text()
+    if CLUSTER_IMPORT in src:
+        model_path.write_text(src.replace(CLUSTER_IMPORT, FPS_FALLBACK, 1))
+        print("[model.py] torch_cluster fps patched")
+    else:
+        print("[model.py] already patched")
+else:
+    print("[model.py] not found")
+
+# ── 3. models_ae.py — THE CRITICAL MISSING PATCH ─────────────────────────────
+# The original patch_runtime.py only patched model.py but NOT models_ae.py,
+# which also imports torch_cluster. This caused the ModuleNotFoundError.
+models_ae_path = Path("app/models_ae.py")
+if models_ae_path.exists():
+    src = models_ae_path.read_text()
+    if CLUSTER_IMPORT in src:
+        models_ae_path.write_text(src.replace(CLUSTER_IMPORT, FPS_FALLBACK, 1))
+        print("[models_ae.py] torch_cluster fps patched")
+    else:
+        print("[models_ae.py] already patched")
+else:
+    print("[models_ae.py] not found")
+
+# ── 4. app/app.py — @spaces.GPU duration ─────────────────────────────────────
+app_path = Path("app/app.py")
+if app_path.exists():
+    src = app_path.read_text()
+    bare_gpu_decorators = src.count("@spaces.GPU\n")
+    if bare_gpu_decorators != 2:
+        raise RuntimeError(f"Expected 2 bare GPU decorators, found {bare_gpu_decorators}")
+    app_path.write_text(src.replace("@spaces.GPU\n", "@spaces.GPU(duration=300)\n"))
+    print(f"[app/app.py] {bare_gpu_decorators} @spaces.GPU -> duration=300 patched")
+else:
+    print("[app/app.py] not found")
+
+print("patch_runtime.py done.")