Matija Kecman
84950fa371
This change introduces parallel chunked downloading for large model files (.mjb) directly into WASM linear memory, enables model loading without full page reloads and reduces memory overhead. Key changes: - Implement a chunked model endpoint in the Python web server to support range requests. - Add parallel chunked fetching in the frontend with retry logic and a single-fetch fallback. - Increase initial WASM memory to 3 GB to accommodate large models and prevent heap fragmentation. - Support in-place model reloading in the C++ client, including texture cache invalidation when the Filament context is recreated. - Display a model download progress bar and model parsing/loading banner to the UI. - Fixes model drag and drop (caused by typo in sessionId, corrected to session_id). PiperOrigin-RevId: 960249236 Change-Id: Icac89e6a4ca099882aaf9b111744c1b6ab0cc6c1
369 lines
15 KiB
HTML
369 lines
15 KiB
HTML
<!doctype html>
|
|
<!--
|
|
Copyright 2026 DeepMind Technologies Limited
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
https://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
-->
|
|
<html>
|
|
<head>
|
|
<meta charset="utf-8" />
|
|
<title>MuJoCo Web Viewer</title>
|
|
<link rel="icon" href="favicon.ico" sizes="16x16 32x32 48x48" />
|
|
<style>
|
|
body {
|
|
margin: 0;
|
|
overflow: hidden;
|
|
background: #1a1a2e;
|
|
}
|
|
canvas {
|
|
width: 100vw;
|
|
height: 100vh;
|
|
display: block;
|
|
}
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<canvas id="canvas" oncontextmenu="event.preventDefault()"></canvas>
|
|
<script>
|
|
// Standard WebSocket close codes.
|
|
const WS_CLOSE_POLICY_VIOLATION = 1008;
|
|
const WS_CLOSE_MESSAGE_TOO_BIG = 1009;
|
|
|
|
// Our custom WebSocket close codes in range 4xxx.
|
|
const WS_CLOSE_CONTROLLER_TAKEN = 4001;
|
|
const WS_CLOSE_SESSION_FULL = 4002;
|
|
const WS_CLOSE_INACTIVE = 4003;
|
|
const WS_CLOSE_NOT_CONTROLLER = 4004;
|
|
|
|
// Parallel chunked model download.
|
|
//
|
|
// Fetches /model in up to PARALLEL concurrent 64 MiB requests
|
|
// and streams the chunks directly into a WASM linear heap buffer
|
|
// allocated via allocModelBuffer(totalSize). Each chunk is retried up
|
|
// to MAX_RETRIES times on transient errors (proxy resets, HTTP2
|
|
// protocol errors, etc.). Returns { ptr, size } on success, or null
|
|
// on failure.
|
|
async function fetchModelChunked() {
|
|
const CHUNK = 64 * 1024 * 1024; // 64 MiB per request
|
|
const PARALLEL = 6; // concurrent fetches
|
|
const MAX_RETRIES = 3; // per-chunk retry limit
|
|
const RETRY_DELAY = 500; // ms between retries
|
|
let ptr = 0;
|
|
try {
|
|
const totalSizeResp = await fetch("/model?total_bytes");
|
|
if (!totalSizeResp.ok) return null;
|
|
const totalSizeData = await totalSizeResp.json();
|
|
const totalSize = totalSizeData.total_bytes;
|
|
if (!totalSize || totalSize <= 0) return null;
|
|
|
|
ptr = Module.allocModelBuffer(totalSize);
|
|
if (!ptr) {
|
|
console.error("[model] failed to allocate WASM buffer of size:", totalSize);
|
|
return null;
|
|
}
|
|
const chunks = [];
|
|
for (let offset = 0; offset < totalSize; offset += CHUNK) {
|
|
chunks.push({ offset: offset, size: Math.min(CHUNK, totalSize - offset) });
|
|
}
|
|
|
|
if (Module.updateModelDownloadProgress) {
|
|
Module.updateModelDownloadProgress(0, totalSize, 0);
|
|
}
|
|
let bytesDownloaded = 0;
|
|
let maxRetry = 0;
|
|
// Bounded-parallelism worker pool with per-chunk retries.
|
|
let nextIndex = 0;
|
|
let errors = 0;
|
|
async function worker() {
|
|
while (nextIndex < chunks.length) {
|
|
const index = nextIndex++;
|
|
const chunk = chunks[index];
|
|
const url = "/model?offset_bytes=" + chunk.offset + "&size_bytes=" + chunk.size;
|
|
let ok = false;
|
|
for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
|
|
try {
|
|
if (attempt > 0) {
|
|
if (attempt > maxRetry) maxRetry = attempt;
|
|
if (Module.updateModelDownloadProgress) {
|
|
Module.updateModelDownloadProgress(bytesDownloaded, totalSize, maxRetry);
|
|
}
|
|
console.log("[model] chunk", index, "retry", attempt);
|
|
await new Promise((r) => setTimeout(r, RETRY_DELAY * attempt));
|
|
}
|
|
const resp = await fetch(url);
|
|
if (!resp.ok) continue;
|
|
const buffer = await resp.arrayBuffer();
|
|
HEAPU8.set(new Uint8Array(buffer), ptr + chunk.offset);
|
|
bytesDownloaded += chunk.size;
|
|
if (Module.updateModelDownloadProgress) {
|
|
Module.updateModelDownloadProgress(bytesDownloaded, totalSize, maxRetry);
|
|
}
|
|
ok = true;
|
|
break;
|
|
} catch (e) {
|
|
console.warn("[model] chunk", index, "attempt", attempt, e);
|
|
}
|
|
}
|
|
if (!ok) {
|
|
console.error("[model] chunk", index, "failed after", MAX_RETRIES, "retries");
|
|
errors++;
|
|
}
|
|
}
|
|
}
|
|
const workers = [];
|
|
for (let i = 0; i < Math.min(PARALLEL, chunks.length); i++) {
|
|
workers.push(worker());
|
|
}
|
|
await Promise.all(workers);
|
|
if (errors > 0) return null; // finally frees ptr
|
|
if (Module.updateModelDownloadProgress) {
|
|
Module.updateModelDownloadProgress(totalSize, totalSize, 0);
|
|
}
|
|
// Allow ~4 frames (60ms) to elapse so the 100% progress bar paints
|
|
// before C++ synchronous parsing blocks the thread.
|
|
await new Promise((resolve) => setTimeout(resolve, 60));
|
|
const result = { ptr: ptr, size: totalSize };
|
|
ptr = 0; // transfer ownership to caller
|
|
return result;
|
|
} catch (e) {
|
|
console.error("[model] chunked fetch failed:", e);
|
|
return null; // finally frees ptr
|
|
} finally {
|
|
if (ptr) Module.freeModelBuffer(ptr);
|
|
}
|
|
}
|
|
|
|
// Single-fetch fallback if chunked download using fetchModelChunked() fails.
|
|
async function fetchModelSingle() {
|
|
let ptr = 0;
|
|
try {
|
|
const resp = await fetch("/model");
|
|
if (!resp.ok) return null;
|
|
const buffer = await resp.arrayBuffer();
|
|
const totalSize = buffer.byteLength;
|
|
if (Module.updateModelDownloadProgress) {
|
|
Module.updateModelDownloadProgress(0, totalSize, 0);
|
|
}
|
|
ptr = Module.allocModelBuffer(totalSize);
|
|
if (!ptr) return null;
|
|
HEAPU8.set(new Uint8Array(buffer), ptr);
|
|
if (Module.updateModelDownloadProgress) {
|
|
Module.updateModelDownloadProgress(totalSize, totalSize, 0);
|
|
}
|
|
// Allow ~4 frames (60ms) to elapse so the 100% progress bar paints
|
|
// before C++ synchronous parsing blocks the thread.
|
|
await new Promise((resolve) => setTimeout(resolve, 60));
|
|
const result = { ptr: ptr, size: totalSize };
|
|
ptr = 0; // transfer ownership to caller
|
|
return result;
|
|
} catch (e) {
|
|
console.error("[model] single fetch failed:", e);
|
|
return null; // finally frees ptr
|
|
} finally {
|
|
if (ptr) Module.freeModelBuffer(ptr);
|
|
}
|
|
}
|
|
|
|
// Called from C++ (via EM_ASM) when the Python side swaps the model.
|
|
async function reloadModel() {
|
|
// Retry the chunked download up to 3 times with increasing delay.
|
|
// Falling back to a single-fetch for large models hits the same
|
|
// proxy errors, so retrying chunks is the only viable path.
|
|
for (let attempt = 0; attempt < 3; attempt++) {
|
|
if (attempt > 0) {
|
|
console.log("[model] retrying chunked download, attempt", attempt + 1);
|
|
await new Promise((r) => setTimeout(r, 1000 * attempt));
|
|
}
|
|
const modelRes = await fetchModelChunked();
|
|
if (modelRes) {
|
|
Module.parseModelBuffer(modelRes.ptr, modelRes.size);
|
|
Module.freeModelBuffer(modelRes.ptr);
|
|
return;
|
|
}
|
|
}
|
|
console.error("[model] all chunked download attempts failed; trying single-fetch");
|
|
const modelRes = await fetchModelSingle();
|
|
if (modelRes) {
|
|
Module.parseModelBuffer(modelRes.ptr, modelRes.size);
|
|
Module.freeModelBuffer(modelRes.ptr);
|
|
}
|
|
}
|
|
|
|
var Module = {
|
|
canvas: (() => {
|
|
const canvas = document.getElementById("canvas");
|
|
return canvas;
|
|
})(),
|
|
locateFile: function (path, prefix) {
|
|
if (path.endsWith(".data")) {
|
|
return "web_client.data";
|
|
}
|
|
if (path.includes("assets/")) {
|
|
const filename = path.substring(path.lastIndexOf("/") + 1);
|
|
return "assets/" + filename;
|
|
}
|
|
return prefix + path;
|
|
},
|
|
onRuntimeInitialized: () => {
|
|
const assetsToPrefetch = [
|
|
"AtkinsonHyperlegibleNext[wght].ttf",
|
|
"AtkinsonHyperlegibleMono-Regular.ttf",
|
|
"fontawesome-webfont.ttf",
|
|
"ibl.ktx",
|
|
"pbr.filamat",
|
|
"pbr_transparent.filamat",
|
|
"pbr_packed.filamat",
|
|
"pbr_packed_transparent.filamat",
|
|
"phong_2d.filamat",
|
|
"phong_2d_fade.filamat",
|
|
"phong_2d_reflect.filamat",
|
|
"phong_2d_uv.filamat",
|
|
"phong_2d_uv_fade.filamat",
|
|
"phong_2d_uv_reflect.filamat",
|
|
"phong_color.filamat",
|
|
"phong_color_fade.filamat",
|
|
"phong_color_reflect.filamat",
|
|
"phong_cube.filamat",
|
|
"phong_cube_fade.filamat",
|
|
"phong_cube_reflect.filamat",
|
|
"outline_composite.filamat",
|
|
"outline_flatten.filamat",
|
|
"outline_jumpflood.filamat",
|
|
"decor.filamat",
|
|
"unlit_depth.filamat",
|
|
"unlit_segmentation.filamat",
|
|
"unlit_ui.filamat",
|
|
];
|
|
|
|
const assetPromises = assetsToPrefetch.map(async (filename) => {
|
|
for (let attempt = 0; attempt < 3; attempt++) {
|
|
try {
|
|
if (attempt > 0) {
|
|
await new Promise((r) => setTimeout(r, 500 * attempt));
|
|
}
|
|
const response = await fetch("assets/" + filename);
|
|
if (!response.ok) continue;
|
|
const buffer = await response.arrayBuffer();
|
|
Module.registerAsset(filename, new Uint8Array(buffer));
|
|
return;
|
|
} catch (error) {
|
|
console.warn(`Attempt ${attempt + 1} failed for ${filename}:`, error);
|
|
}
|
|
}
|
|
console.error(`Error prefetching asset ${filename} after 3 attempts`);
|
|
});
|
|
|
|
Promise.all(assetPromises).then(() => {
|
|
Module.startApp();
|
|
reloadModel();
|
|
});
|
|
},
|
|
};
|
|
// Set only the CSS size; SDL owns canvas.width/height. Use whole
|
|
// pixels: a fractional size (100vw/100vh under display scaling) makes
|
|
// the streamed UI relayout — and visibly flicker — every frame.
|
|
function fitCanvas() {
|
|
Module.canvas.style.width = Math.floor(window.innerWidth) + "px";
|
|
Module.canvas.style.height = Math.floor(window.innerHeight) + "px";
|
|
}
|
|
fitCanvas();
|
|
window.addEventListener("resize", fitCanvas);
|
|
|
|
// Drag & drop: upload dropped model files to the viewer over the
|
|
// /drop WebSocket. Each file is one binary frame ([u32 path length]
|
|
// [relative path utf-8][file bytes], little-endian), an empty frame
|
|
// marks the end of the drop, and the server closes once it has
|
|
// everything (parsed by web_server.py's drop_handler). Folders are
|
|
// walked recursively so models with separate asset files work; the
|
|
// viewer picks the root model file (see web_viewer._pick_drop_root).
|
|
// Loading a model changes the session for everyone, so only the
|
|
// controller may drop; the server rejects drops from other pages and the
|
|
// spectator check below just skips the pointless upload.
|
|
window.addEventListener("dragover", (e) => e.preventDefault());
|
|
window.addEventListener("drop", async (e) => {
|
|
e.preventDefault();
|
|
if (window.Module && Module.isSpectator) {
|
|
console.warn("Model drop ignored: only the controlling page can load models.");
|
|
return;
|
|
}
|
|
// Capture entries synchronously: DataTransferItems are invalidated
|
|
// as soon as this handler awaits.
|
|
const items = e.dataTransfer ? [...e.dataTransfer.items] : [];
|
|
const entries = items
|
|
.map((i) => i.webkitGetAsEntry && i.webkitGetAsEntry())
|
|
.filter(Boolean);
|
|
const files = []; // {path, file}
|
|
async function walk(entry, prefix) {
|
|
if (entry.isFile) {
|
|
const file = await new Promise((ok, err) => entry.file(ok, err));
|
|
files.push({ path: prefix + entry.name, file });
|
|
} else if (entry.isDirectory) {
|
|
const reader = entry.createReader();
|
|
for (;;) {
|
|
// readEntries returns batches (<=100); loop until empty.
|
|
const batch = await new Promise((ok, err) => reader.readEntries(ok, err));
|
|
if (!batch.length) break;
|
|
for (const child of batch) {
|
|
await walk(child, prefix + entry.name + "/");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
for (const entry of entries) await walk(entry, "");
|
|
if (!files.length) return;
|
|
// The server limits each WebSocket message size. If any file exceeds the limit, the server
|
|
// closes the connection with WS_CLOSE_MESSAGE_TOO_BIG mid-upload and the entire drop is
|
|
// discarded. Check client-side first to give a clear message, not a silent failure.
|
|
const MAX_DROP_FILE_BYTES = 64 * 1024 * 1024;
|
|
const tooBig = files.find((f) => f.file.size > MAX_DROP_FILE_BYTES);
|
|
if (tooBig) {
|
|
console.warn(
|
|
`Model drop rejected: "${tooBig.path}" is larger than ` +
|
|
`${MAX_DROP_FILE_BYTES / (1024 * 1024)} MiB.`,
|
|
);
|
|
return;
|
|
}
|
|
const frames = [];
|
|
for (const { path, file } of files) {
|
|
const bytes = new Uint8Array(await file.arrayBuffer());
|
|
const name = new TextEncoder().encode(path);
|
|
const frame = new Uint8Array(4 + name.length + bytes.length);
|
|
new DataView(frame.buffer).setUint32(0, name.length, true);
|
|
frame.set(name, 4);
|
|
frame.set(bytes, 4 + name.length);
|
|
frames.push(frame);
|
|
}
|
|
const proto = location.protocol === "https:" ? "wss://" : "ws://";
|
|
// The session id tells the server which page is dropping; it only
|
|
// accepts drops from the controller.
|
|
const sid = (window.Module && Module.session_id) || "";
|
|
const dropUrl = proto + location.host + "/drop?sid=" + encodeURIComponent(sid);
|
|
const ws = new WebSocket(dropUrl);
|
|
ws.binaryType = "arraybuffer";
|
|
ws.onopen = () => {
|
|
for (const frame of frames) ws.send(frame);
|
|
ws.send(new Uint8Array(0)); // End marker; the server closes.
|
|
};
|
|
ws.onclose = (ev) => {
|
|
if (ev.code === WS_CLOSE_NOT_CONTROLLER) {
|
|
console.warn("Model drop rejected: only the controlling page can load models.");
|
|
} else if (ev.code === WS_CLOSE_MESSAGE_TOO_BIG) {
|
|
console.warn("Model drop rejected: a file exceeded the server's size limit.");
|
|
}
|
|
};
|
|
});
|
|
</script>
|
|
<script async src="web_client.js"></script>
|
|
</body>
|
|
</html>
|