Spaces:
Running
Running
| 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.") | |