Spaces:
Paused
Paused
| import os, shutil, subprocess, tempfile | |
| SUPPORTED = (".glb", ".gltf", ".fbx", ".dae") | |
| _SCRIPT = os.path.join(os.path.dirname(__file__), "tools", "fbx2glb.mjs") | |
| _NODE_DIR = os.environ.get("VID2RIG_NODE_DIR", "/tmp/vid2rig_node") | |
| def _find(cmd): | |
| """Locate a binary even if it's not on the Python process PATH (HF runtime PATH != build PATH).""" | |
| return shutil.which(cmd) or next( | |
| (p for p in (f"/usr/local/bin/{cmd}", f"/usr/bin/{cmd}", f"/opt/homebrew/bin/{cmd}") | |
| if os.path.exists(p)), None) | |
| def _node_convert(src: str, out: str) -> bool: | |
| """FBX -> GLB via three.js FBXLoader + GLTFExporter (the path the mule scenes use). Correctly | |
| collapses Mixamo $AssimpFbx$ pivots, fixes the up-axis, and writes TRS bones so the baked | |
| animation is valid. Installs `three` into a work dir on first use. Returns True on success.""" | |
| node, npm = _find("node"), _find("npm") | |
| if not node or not npm or not os.path.exists(_SCRIPT): | |
| print(f"[rig_io] node convert unavailable (node={node} npm={npm} script={os.path.exists(_SCRIPT)})", flush=True) | |
| return False | |
| src, out = os.path.abspath(src), os.path.abspath(out) # node runs with cwd=_NODE_DIR | |
| os.makedirs(_NODE_DIR, exist_ok=True) | |
| shutil.copy(_SCRIPT, os.path.join(_NODE_DIR, "fbx2glb.mjs")) | |
| if not os.path.isdir(os.path.join(_NODE_DIR, "node_modules", "three")): | |
| r = subprocess.run([npm, "install", "--no-audit", "--no-fund", "three@0.169.0"], | |
| cwd=_NODE_DIR, capture_output=True, text=True) | |
| if r.returncode != 0: | |
| print("[rig_io] npm install three FAILED:", (r.stderr or r.stdout)[-500:], flush=True) | |
| return False | |
| r = subprocess.run([node, os.path.join(_NODE_DIR, "fbx2glb.mjs"), src, out], | |
| cwd=_NODE_DIR, capture_output=True, text=True) | |
| ok = r.returncode == 0 and os.path.exists(out) | |
| if ok: | |
| print("[rig_io] three.js conversion OK ->", out, flush=True) | |
| else: | |
| print("[rig_io] three.js conversion FAILED:", (r.stderr or r.stdout)[-500:], flush=True) | |
| return ok | |
| def _embedded_png(fbx: str): | |
| """(bytes, mime) of the FBX's first embedded texture via assimp, or None.""" | |
| assimp = _find("assimp") | |
| if not assimp: | |
| return None | |
| try: | |
| import pygltflib as gl | |
| tmp = tempfile.mktemp(suffix=".glb") | |
| subprocess.run([assimp, "export", fbx, tmp], check=True, capture_output=True) | |
| ga = gl.GLTF2().load(tmp) | |
| if not ga.images: | |
| return None | |
| im = ga.images[0]; bv = ga.bufferViews[im.bufferView] | |
| png = (ga.binary_blob() or b"")[(bv.byteOffset or 0):(bv.byteOffset or 0) + bv.byteLength] | |
| return png, (im.mimeType or "image/png") | |
| except Exception: | |
| return None | |
| def _finalize_rig(glb: str, fbx: str) -> None: | |
| """Post-process the clean three.js GLB in place: (1) point skins[0].skeleton at Hips | |
| (fixes SKIN_SKELETON_INVALID), (2) transplant the FBX's embedded texture onto the material | |
| (three.js export drops it but keeps the matching UV0). Best-effort; failures leave a valid rig.""" | |
| try: | |
| import pygltflib as gl | |
| g = gl.GLTF2().load(glb) | |
| changed = False | |
| if g.skins: | |
| hips = next((i for i, n in enumerate(g.nodes) | |
| if n.name and n.name.replace(":", "").replace("_", "") == "mixamorigHips"), None) | |
| if hips is not None and g.skins[0].skeleton != hips: | |
| g.skins[0].skeleton = hips; changed = True | |
| tex = _embedded_png(fbx) | |
| prims = [p for m in (g.meshes or []) for p in m.primitives | |
| if getattr(p.attributes, "TEXCOORD_0", None) is not None] | |
| if tex and prims: | |
| png, mime = tex | |
| blob = bytearray(g.binary_blob() or b"") | |
| while len(blob) % 4: | |
| blob.append(0) | |
| off = len(blob); blob += png | |
| g.bufferViews.append(gl.BufferView(buffer=0, byteOffset=off, byteLength=len(png))) | |
| g.images = (g.images or []) + [gl.Image(bufferView=len(g.bufferViews) - 1, mimeType=mime)] | |
| g.samplers = (g.samplers or []) + [gl.Sampler()] | |
| g.textures = (g.textures or []) + [gl.Texture(source=len(g.images) - 1, sampler=len(g.samplers) - 1)] | |
| tex_i = len(g.textures) - 1 | |
| for prim in prims: | |
| if prim.material is None: | |
| g.materials = (g.materials or []) + [gl.Material()]; prim.material = len(g.materials) - 1 | |
| mat = g.materials[prim.material] | |
| if mat.pbrMetallicRoughness is None: | |
| mat.pbrMetallicRoughness = gl.PbrMetallicRoughness() | |
| mat.pbrMetallicRoughness.baseColorTexture = gl.TextureInfo(index=tex_i, texCoord=0) | |
| mat.pbrMetallicRoughness.baseColorFactor = [1, 1, 1, 1] | |
| g.set_binary_blob(bytes(blob)); g.buffers[0].byteLength = len(blob) | |
| changed = True | |
| print("[rig_io] texture transplanted onto rig", flush=True) | |
| if changed: | |
| g.save(glb) | |
| except Exception as e: | |
| print("[rig_io] _finalize_rig error:", repr(e), flush=True) | |
| return | |
| def to_glb(path: str) -> str: | |
| """Return a .glb path for the rig. .glb/.gltf pass through; .fbx is converted via three.js | |
| (preferred — handles Mixamo rigs correctly); assimp is a last-resort fallback (it mangles | |
| Mixamo pivots, so only used when node/three is unavailable, and for .dae).""" | |
| ext = os.path.splitext(path)[1].lower() | |
| if ext in (".glb", ".gltf"): | |
| return path | |
| if ext not in SUPPORTED: | |
| raise ValueError(f"unsupported rig format {ext!r}; use one of {SUPPORTED}") | |
| out = tempfile.mktemp(suffix=".glb") | |
| if ext == ".fbx" and _node_convert(path, out): | |
| _finalize_rig(out, path) # fix skin root + borrow embedded texture | |
| return out | |
| print("[rig_io] FALLING BACK to assimp (mangles Mixamo rig + no texture)", flush=True) | |
| assimp = _find("assimp") | |
| if assimp: | |
| subprocess.run([assimp, "export", path, out], check=True, capture_output=True) | |
| return out | |
| fbx2gltf = shutil.which("FBX2glTF") or shutil.which("fbx2gltf") | |
| if fbx2gltf and ext == ".fbx": | |
| subprocess.run([fbx2gltf, "-i", path, "-o", out, "--binary"], check=True, capture_output=True) | |
| return out if os.path.exists(out) else out + ".glb" | |
| raise RuntimeError( | |
| "No FBX/DAE converter available. Need Node.js + three (preferred) or `brew install assimp`.") | |