Spaces:
Running
Running
| // Runs the MotionBricks G1 ONNX model with onnxruntime-web (WASM) off the main thread. | |
| /* global ort */ | |
| importScripts('https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/ort.wasm.min.js'); | |
| let session = null; | |
| async function cachedFetch(url, onProgress) { | |
| let cache = null; | |
| try { cache = await caches.open('motionbricks-models-v1'); } catch (e) { cache = null; } | |
| if (cache) { | |
| const hit = await cache.match(url); | |
| if (hit) { | |
| onProgress(1, 1, true); | |
| return new Uint8Array(await hit.arrayBuffer()); | |
| } | |
| } | |
| const response = await fetch(url, { mode: 'cors' }); | |
| if (!response.ok) throw new Error(`HTTP ${response.status} for ${url}`); | |
| const total = Number(response.headers.get('content-length')) || 0; | |
| const reader = response.body.getReader(); | |
| const chunks = []; | |
| let received = 0; | |
| for (;;) { | |
| const { done, value } = await reader.read(); | |
| if (done) break; | |
| chunks.push(value); | |
| received += value.length; | |
| onProgress(received, total, false); | |
| } | |
| const bytes = new Uint8Array(received); | |
| let offset = 0; | |
| for (const c of chunks) { bytes.set(c, offset); offset += c.length; } | |
| if (cache) { | |
| try { await cache.put(url, new Response(bytes, { headers: { 'content-type': 'application/octet-stream' } })); } catch (e) { /* quota */ } | |
| } | |
| return bytes; | |
| } | |
| self.onmessage = async (event) => { | |
| const msg = event.data; | |
| if (msg.type === 'load') { | |
| try { | |
| const threads = self.crossOriginIsolated ? Math.max(1, Math.min(8, (navigator.hardwareConcurrency || 4) - 1)) : 1; | |
| ort.env.wasm.numThreads = threads; | |
| ort.env.wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/'; | |
| let bytes = null; | |
| let lastError = null; | |
| for (const url of msg.urls) { | |
| try { | |
| bytes = await cachedFetch(url, (r, t, cached) => self.postMessage({ type: 'progress', received: r, total: t, cached })); | |
| break; | |
| } catch (e) { lastError = e; } | |
| } | |
| if (!bytes) throw lastError || new Error('Model download failed'); | |
| self.postMessage({ type: 'compiling' }); | |
| const t0 = performance.now(); | |
| session = await ort.InferenceSession.create(bytes, { executionProviders: ['wasm'], graphOptimizationLevel: 'all' }); | |
| self.postMessage({ type: 'ready', threads, compileMs: performance.now() - t0 }); | |
| } catch (e) { | |
| self.postMessage({ type: 'error', message: String(e && e.message ? e.message : e) }); | |
| } | |
| return; | |
| } | |
| if (msg.type === 'run') { | |
| const f = msg.feeds; | |
| const n = f.n || 1; | |
| const t0 = performance.now(); | |
| try { | |
| const out = await session.run({ | |
| context_qpos: new ort.Tensor('float32', f.context_qpos, [n, 4, 36]), | |
| target_qpos: new ort.Tensor('float32', f.target_qpos, [n, 4, 36]), | |
| has_global_root_values: new ort.Tensor('int32', f.has_global_root_values, [n, 8]), | |
| has_local_root_values: new ort.Tensor('int32', f.has_local_root_values, [n, 8]), | |
| has_local_poses: new ort.Tensor('int32', f.has_local_poses, [n, 8]), | |
| allowed_pred_num_tokens: new ort.Tensor('int32', f.allowed_pred_num_tokens, [n, 11]), | |
| }); | |
| const pred_qpos = out.pred_qpos.data, pred_foot_contacts = out.pred_foot_contacts.data; | |
| const ntok = Array.from(out.pred_num_tokens.data, Number); | |
| self.postMessage({ type: 'result', id: msg.id, pred_qpos, pred_foot_contacts, ntok, ms: performance.now() - t0 }, | |
| [pred_qpos.buffer, pred_foot_contacts.buffer]); | |
| } catch (e) { | |
| self.postMessage({ type: 'result', id: msg.id, error: String(e && e.message ? e.message : e) }); | |
| } | |
| } | |
| }; | |