Spaces:
Running
Running
File size: 4,472 Bytes
8f65673 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | 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.")
|