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", | |
| "app_blender.py", | |
| "util/blender_utils.py", | |
| ): | |
| pristine = subprocess.check_output( | |
| ["git", "-c", "safe.directory=/data/app", "-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.") | |
| # STATICAI Standard animation library integration | |
| def _replace_once(path, old, new, label): | |
| source = path.read_text() | |
| count = source.count(old) | |
| if count != 1: | |
| raise RuntimeError(f"{label}: expected exactly one patch target, found {count}") | |
| path.write_text(source.replace(old, new, 1)) | |
| blender_utils_path = Path("app/util/blender_utils.py") | |
| _replace_once( | |
| blender_utils_path, | |
| "def load_mixamo_anim(char_file: str, anim_file: str, do_retarget=False, inplace=False, to_tris=False):", | |
| "def load_mixamo_anim(char_file: str, anim_file: str, do_retarget=False, inplace=False, to_tris=False, action_name=None):", | |
| "blender_utils signature", | |
| ) | |
| _replace_once( | |
| blender_utils_path, | |
| ' assert anim_armature.animation_data is not None and len(bpy.data.actions) > 0, f"Animation not found in {anim_file}"\n\n set_action(char_armature, anim_armature.animation_data.action)', | |
| ''' assert anim_armature.animation_data is not None and len(bpy.data.actions) > 0, f"Animation not found in {anim_file}" | |
| if action_name: | |
| candidates = [ | |
| action | |
| for action in bpy.data.actions | |
| if action.name == action_name | |
| or action.name.endswith("|" + action_name) | |
| or action.name.startswith(action_name + ".") | |
| ] | |
| if not candidates: | |
| available = ", ".join(sorted(action.name for action in bpy.data.actions)) | |
| raise ValueError(f"Animation action {action_name!r} not found. Available: {available}") | |
| anim_armature.animation_data.action = candidates[0] | |
| print(f"Selected animation action: {candidates[0].name}") | |
| set_action(char_armature, anim_armature.animation_data.action)''', | |
| "blender_utils action selection", | |
| ) | |
| app_blender_path = Path("app/app_blender.py") | |
| _replace_once( | |
| app_blender_path, | |
| ''' blender_utils.load_mixamo_anim( | |
| [armature_obj, mesh_obj], args.animation_path, do_retarget=args.retarget, inplace=args.inplace | |
| )''', | |
| ''' blender_utils.load_mixamo_anim( | |
| [armature_obj, mesh_obj], | |
| args.animation_path, | |
| do_retarget=args.retarget, | |
| inplace=args.inplace, | |
| action_name=args.animation_action, | |
| )''', | |
| "app_blender load_mixamo_anim", | |
| ) | |
| _replace_once( | |
| app_blender_path, | |
| ' parser.add_argument("--animation_path", type=str, default=None)\n', | |
| ' parser.add_argument("--animation_path", type=str, default=None)\n parser.add_argument("--animation_action", type=str, default=None)\n', | |
| "app_blender parser action", | |
| ) | |
| app_path = Path("app/app.py") | |
| _replace_once(app_path, "import gc\nimport os\n", "import gc\nimport json\nimport os\n", "app json import") | |
| _replace_once( | |
| app_path, | |
| ''' animation_file: str, | |
| retarget: bool,''', | |
| ''' animation_file: str, | |
| animation_action: str, | |
| retarget: bool,''', | |
| "vis_blender action argument", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ''' animation_path=animation_file, | |
| retarget=retarget,''', | |
| ''' animation_path=animation_file, | |
| animation_action=animation_action, | |
| retarget=retarget,''', | |
| "Namespace action argument", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ''' if animation_file is not None: | |
| cmd += f" --animation_path '{os.path.abspath(animation_file)}'" | |
| if retarget: | |
| cmd += " --retarget" | |
| if inplace: | |
| cmd += " --inplace"''', | |
| ''' if animation_file is not None: | |
| cmd += f" --animation_path '{os.path.abspath(animation_file)}'" | |
| if animation_action: | |
| safe_action = animation_action.replace("'", "") | |
| cmd += f" --animation_action '{safe_action}'" | |
| if retarget: | |
| cmd += " --retarget" | |
| if inplace: | |
| cmd += " --inplace"''', | |
| "subprocess action argument", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ''' animation_file: str = None, | |
| retarget=True,''', | |
| ''' animation_file: str = None, | |
| animation_action: str = None, | |
| retarget=True,''', | |
| "pipeline action argument", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ''' yield vis_blender( | |
| reset_to_rest, no_fingers, rest_pose_type, ignore_pose_parts, animation_file, retarget, inplace, db | |
| )''', | |
| ''' yield vis_blender( | |
| reset_to_rest, | |
| no_fingers, | |
| rest_pose_type, | |
| ignore_pose_parts, | |
| animation_file, | |
| animation_action, | |
| retarget, | |
| inplace, | |
| db, | |
| )''', | |
| "pipeline vis_blender action", | |
| ) | |
| _replace_once(app_path, ' title = "Make-It-Animatable"\n', ' title = "Motion Rig"\n', "Motion Rig title") | |
| _replace_once( | |
| app_path, | |
| ' animation_dir = "data/Mixamo/animation"\n', | |
| ''' animation_dir = "data/Mixamo/animation" | |
| standard_library_root = "/data/animations/Standard" | |
| standard_manifest_path = os.path.join(standard_library_root, "ANIMATION_MANIFEST.json") | |
| standard_animation_lookup = {} | |
| standard_animation_choices = [] | |
| if os.path.isfile(standard_manifest_path): | |
| with open(standard_manifest_path, "r", encoding="utf-8") as fh: | |
| manifest = json.load(fh) | |
| for entry_index, entry in enumerate(manifest.get("entries", [])): | |
| motion_file = os.path.join(standard_library_root, entry["file"]) | |
| variant = "Root Motion" if entry.get("root_motion") else "In Place" | |
| for action in entry.get("actions", []): | |
| if not action or action == "A_TPose": | |
| continue | |
| key = f"{entry_index}:{action}" | |
| label = f"Standard {entry.get('set')} | {action} | {variant}" | |
| standard_animation_lookup[key] = {"path": motion_file, "action": action} | |
| standard_animation_choices.append((label, key)) | |
| ''', | |
| "Standard library manifest load", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ' sorted(glob("*.fbx", root_dir=animation_dir)),\n', | |
| ' standard_animation_choices,\n', | |
| "Standard animation choices", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ' label="Select Animation File",\n', | |
| ' label="Standard Animation Library",\n', | |
| "Standard animation label", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ''' input_animation_file = gr.File( | |
| label="Animation File", | |
| file_types=[".fbx"], | |
| value=lambda: "./data/Standard Run.fbx", | |
| interactive=True, | |
| )''', | |
| ''' input_animation_file = gr.File( | |
| label="Animation File", | |
| file_types=[".fbx"], | |
| value=lambda: "./data/Standard Run.fbx", | |
| interactive=True, | |
| ) | |
| animation_action_state = gr.State(value=None)''', | |
| "animation action state", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ''' input_animation_file, | |
| input_retarget,''', | |
| ''' input_animation_file, | |
| animation_action_state, | |
| input_retarget,''', | |
| "inputs tuple action state", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ''' animation_file=inputs[input_animation_file], | |
| retarget=inputs[input_retarget],''', | |
| ''' animation_file=inputs[input_animation_file], | |
| animation_action=inputs[animation_action_state], | |
| retarget=inputs[input_retarget],''', | |
| "pipeline UI action", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ''' input_animation_file, | |
| input_retarget,''', | |
| ''' input_animation_file, | |
| animation_action_state, | |
| input_retarget,''', | |
| "animate inputs action state", | |
| ) | |
| _replace_once( | |
| app_path, | |
| ''' def select2file(selected: str): | |
| return None if selected is None else os.path.join(animation_dir, selected) | |
| select_animation_file.input( | |
| lambda x: (select2file(x), gr.Checkbox(value=True)), | |
| inputs=select_animation_file, | |
| outputs=[input_animation_file, input_reset_to_rest], | |
| preprocess=False, | |
| ) | |
| input_animation_file.upload( | |
| lambda: (gr.Dropdown(value=None), gr.Checkbox(value=True)), | |
| outputs=[select_animation_file, input_reset_to_rest], | |
| ) | |
| input_animation_file.clear( | |
| lambda: gr.Dropdown(value=None), outputs=select_animation_file, show_progress="hidden" | |
| )''', | |
| ''' def select2file(selected: str): | |
| if selected is None: | |
| return None, None | |
| item = standard_animation_lookup.get(selected) | |
| if item is None: | |
| raise gr.Error(f"Unknown Standard animation selection: {selected}") | |
| return item["path"], item["action"] | |
| def select_library_animation(selected: str): | |
| path, action = select2file(selected) | |
| return path, action, gr.Checkbox(value=True) | |
| select_animation_file.input( | |
| select_library_animation, | |
| inputs=select_animation_file, | |
| outputs=[input_animation_file, animation_action_state, input_reset_to_rest], | |
| preprocess=False, | |
| ) | |
| input_animation_file.upload( | |
| lambda: (gr.Dropdown(value=None), None, gr.Checkbox(value=True)), | |
| outputs=[select_animation_file, animation_action_state, input_reset_to_rest], | |
| ) | |
| input_animation_file.clear( | |
| lambda: (gr.Dropdown(value=None), None), | |
| outputs=[select_animation_file, animation_action_state], | |
| show_progress="hidden", | |
| )''', | |
| "Standard selection events", | |
| ) | |
| print("STATICAI Standard animation library integration patched.") | |