Cache assets locally

Pipfile

@@ -9,6 +9,8 @@ black = "*"
 isort = "*"
 mypy = "*"
 gradio = "==3.16.2"
+types-requests = "*"
+requests = "*"
 
 [dev-packages]
 

install.py

@@ -0,0 +1,44 @@
+import pathlib
+
+import requests
+
+try:
+    root_path = pathlib.Path(__file__).resolve().parents[0]
+except NameError:
+    import inspect
+
+    root_path = pathlib.Path(inspect.getfile(lambda: None)).resolve().parents[0]
+
+
+def download(url: str, dest: pathlib.Path):
+    try:
+        with requests.get(url, stream=True) as r:
+            with open(dest, "wb") as f:
+                for chunk in r.iter_content(chunk_size=1024):
+                    f.write(chunk)
+    except requests.exceptions.RequestException as e:
+        print(e)
+
+
+def main():
+    MEDIAPIPE_POSE_VERSION = "0.5.1675469404"
+    mediapipe_dir = root_path / "downloads" / "pose" / MEDIAPIPE_POSE_VERSION
+    mediapipe_dir.mkdir(mode=0o755, parents=True, exist_ok=True)
+
+    for file_name in [
+        "pose_landmark_full.tflite",
+        "pose_web.binarypb",
+        "pose_solution_packed_assets.data",
+        "pose_solution_simd_wasm_bin.wasm",
+        "pose_solution_packed_assets_loader.js",
+        "pose_solution_simd_wasm_bin.js",
+    ]:
+        file_path = mediapipe_dir / file_name
+        if file_path.exists():
+            continue
+        url = f"https://cdn.jsdelivr.net/npm/@mediapipe/pose@{MEDIAPIPE_POSE_VERSION}/{file_name}"
+        print(f"Downloading {file_name}...")
+        download(url, file_path)
+
+
+main()

scripts/openpose_editor.py

@@ -26,12 +26,28 @@ def create_ui():
     except (ImportError, AttributeError):
         cn_max = 0
 
-    consts = {
-        "handFbxPath": str(root_path / "models" / "hand.fbx"),
-        "footFbxPath": str(root_path / "models" / "foot.fbx"),
-        "posesPath": str(root_path / "src" / "poses" / "data.bin"),
+    assets = {
+        "models/hand.fbx": "/file=" + str(root_path / "models" / "hand.fbx"),
+        "models/foot.fbx": "/file=" + str(root_path / "models" / "foot.fbx"),
+        "src/poses/data.bin": "/file=" + str(root_path / "src" / "poses" / "data.bin"),
     }
 
+    MEDIAPIPE_POSE_VERSION = "0.5.1675469404"
+    mediapipe_dir = root_path / "downloads" / "pose" / MEDIAPIPE_POSE_VERSION
+    for file_name in [
+        "pose_landmark_full.tflite",
+        "pose_web.binarypb",
+        "pose_solution_packed_assets.data",
+        "pose_solution_simd_wasm_bin.wasm",
+        "pose_solution_packed_assets_loader.js",
+        "pose_solution_simd_wasm_bin.js",
+    ]:
+        file_path = mediapipe_dir / file_name
+        if not file_path.exists():
+            continue
+        assets[file_name] = "/file=" + str(file_path.absolute())
+
+    consts = {"assets": assets}
     gr.HTML(
         f"""
         <div id="openpose3d_consts">{html.escape(json.dumps(consts))}</div>

src/environments/extension/assets.ts

@@ -5,8 +5,9 @@ const footFBXFileUrl =
 const PosesLibraryUrl =
     '/file=extensions/sd-webui-3d-open-pose-editor/src/poses/data.bin'
 
-export default {
+const files: Record<string, string> = {
     'models/hand.fbx': handFBXFileUrl,
     'models/foot.fbx': footFBXFileUrl,
     'src/poses/data.bin': PosesLibraryUrl,
 }
+export default files

src/environments/extension/init.ts

@@ -6,9 +6,7 @@ import { gradioAppElem, updateGradioImage } from './internal/gradio'
 const consts = JSON.parse(
     gradioAppElem.querySelector('#openpose3d_consts')!.textContent!
 )
-assets['models/hand.fbx'] = '/file=' + consts.handFbxPath
-assets['models/foot.fbx'] = '/file=' + consts.footFbxPath
-assets['src/poses/data.bin'] = '/file=' + consts.posesPath
+Object.assign(assets, consts['assets'])
 
 options['Width'] = 512
 options['Height'] = 512