lazy loading webgpu content in docs

ChrLackner · ChrLackner · commit d5e54fd1d19e · 2026-05-25T18:29:50.000+02:00
diff --git a/docs/conf.py b/docs/conf.py
@@ -7,6 +7,7 @@
 import zipfile
 
 os.environ["WEBGPU_EXPORTING"] = "1"
+os.environ["WEBGPU_LAZY_LOAD"] = "1"
 master_doc = "index"
 source_suffix = [".rst", ".md"]
 
diff --git a/webgpu/engine/engine.js b/webgpu/engine/engine.js
@@ -123,6 +123,7 @@ class RenderEngine {
       device: this.device,
       format: this.canvasFormat,
       alphaMode: 'premultiplied',
+      usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.COPY_SRC,
     });
 
     // --- Create GPU resources ---
diff --git a/webgpu/export/screenshot.py b/webgpu/export/screenshot.py
@@ -0,0 +1,152 @@
+"""Persistent screenshot worker for exported scene blobs.
+
+Launches one Chrome instance on a virtual display (Xvfb) and handles multiple
+screenshot requests over stdin/stdout.
+
+Protocol (line-based):
+    Request:  <width> <height> <blob_b64>\n
+    Response: <png_b64>\n
+    Shutdown: (close stdin)
+
+Requires: Xvfb, playwright, chrome
+"""
+
+import sys
+import os
+import base64
+import subprocess
+import tempfile
+import threading
+from pathlib import Path
+from http.server import HTTPServer, SimpleHTTPRequestHandler
+
+
+def main():
+    # Use Xvfb for a private display so Chrome stays invisible
+    for disp_num in range(99, 120):
+        if not os.path.exists(f'/tmp/.X11-unix/X{disp_num}'):
+            break
+    xvfb_proc = subprocess.Popen(
+        ['Xvfb', f':{disp_num}', '-screen', '0', '1280x1024x24', '-ac'],
+        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+    )
+    os.environ['DISPLAY'] = f':{disp_num}'
+    os.environ.pop('WAYLAND_DISPLAY', None)
+    import time
+    time.sleep(0.3)
+
+    try:
+        _run_worker()
+    finally:
+        xvfb_proc.terminate()
+        xvfb_proc.wait()
+
+
+def _run_worker():
+    from playwright.sync_api import sync_playwright
+
+    ARGS = [
+        "--no-sandbox",
+        "--enable-unsafe-webgpu",
+        "--enable-features=Vulkan,UnsafeWebGPU",
+        "--use-vulkan=native",
+        "--ignore-gpu-blocklist",
+        "--disable-dev-shm-usage",
+        "--enable-dawn-features=allow_unsafe_apis,disable_adapter_blocklist",
+        "--ozone-platform=x11",
+    ]
+
+    # Load engine JS once
+    engine_js_path = Path(__file__).parent.parent / "engine"
+    js_files = ["format.js", "compute.js", "camera.js", "input.js", "interactions.js", "engine.js"]
+    engine_js = "\n".join(
+        (engine_js_path / f).read_text().replace("export ", "") for f in js_files
+    )
+    engine_js += "\nif (typeof window !== 'undefined') { window.RenderEngine = RenderEngine; }\n"
+
+    # Start HTTP server for serving pages to Chrome
+    tmpdir = Path(tempfile.mkdtemp(prefix="webgpu_ss_"))
+
+    class Quiet(SimpleHTTPRequestHandler):
+        def __init__(self, *a, **kw):
+            super().__init__(*a, directory=str(tmpdir), **kw)
+        def log_message(self, *a):
+            pass
+
+    server = HTTPServer(("127.0.0.1", 0), Quiet)
+    port = server.server_address[1]
+    threading.Thread(target=server.serve_forever, daemon=True).start()
+
+    # Launch browser once (not headless — uses the Xvfb virtual display)
+    pw = sync_playwright().start()
+    browser = pw.chromium.launch(
+        channel="chrome", headless=False,
+        args=ARGS,
+    )
+    page = browser.new_page()
+
+    # Signal ready
+    sys.stdout.write("READY\n")
+    sys.stdout.flush()
+
+    # Process requests from stdin
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+
+        parts = line.split(' ', 2)
+        if len(parts) != 3:
+            sys.stdout.write("\n")
+            sys.stdout.flush()
+            continue
+
+        width, height, blob_b64 = int(parts[0]), int(parts[1]), parts[2]
+
+        html = f"""<!DOCTYPE html><html><body style="margin:0;padding:0;overflow:hidden;">
+<canvas id="c" width="{width}" height="{height}" style="display:block;"></canvas>
+<script>
+{engine_js}
+(async () => {{
+    try {{
+        const engine = await RenderEngine.create('c', `{blob_b64}`);
+        document.title = 'READY';
+    }} catch(e) {{
+        document.title = 'ERROR:' + (e.stack || e.message || e);
+    }}
+}})();
+</script></body></html>"""
+
+        (tmpdir / "index.html").write_text(html)
+        page.goto(f"http://127.0.0.1:{port}/index.html")
+
+        try:
+            page.wait_for_function(
+                "document.title === 'READY' || document.title.startsWith('ERROR:')",
+                timeout=30000,
+            )
+            title = page.title()
+            if title.startswith('ERROR:'):
+                print(f"[screenshot] JS error: {title}", file=sys.stderr)
+                sys.stdout.write("\n")
+                sys.stdout.flush()
+                continue
+
+            # Wait for rAF compositing then screenshot the canvas element
+            page.wait_for_timeout(100)
+            el = page.query_selector('#c')
+            png_bytes = el.screenshot() if el else b''
+            sys.stdout.write(base64.b64encode(png_bytes).decode() + "\n")
+            sys.stdout.flush()
+        except Exception as e:
+            print(f"[screenshot] error: {e}", file=sys.stderr)
+            sys.stdout.write("\n")
+            sys.stdout.flush()
+
+    browser.close()
+    pw.stop()
+    server.shutdown()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/webgpu/jupyter.py b/webgpu/jupyter.py
@@ -126,6 +126,163 @@ def _DrawHTML(scene, width=640, height=600):
     return scene
 
 
+def _DrawHTMLLazy(scene, width=640, height=600):
+    """Export scene as a lazy-loading HTML snippet with a pre-baked screenshot.
+
+    The screenshot is rendered in a separate subprocess (own Chrome + GPU device)
+    so it doesn't interfere with the main build process's mapAsync calls.
+    The scene blob and engine JS are saved as separate static files and only
+    fetched when the user clicks to interact.
+    """
+    global _engine_emitted
+    from IPython.display import HTML, Javascript, display
+
+    from .engine import engine_js
+
+    if isinstance(scene, Renderer):
+        scene = [scene]
+    if isinstance(scene, list):
+        scene = Scene(scene)
+
+    id_ = f"__webgpu_{next(_id_counter)}_"
+    canvas_id = f"{id_}canvas"
+
+    # Create a canvas in the headless Chrome DOM for GPU initialization.
+    js = platform.js
+    html_canvas = js.document.createElement("canvas")
+    html_canvas.width = width
+    html_canvas.height = height
+    html_canvas.id = canvas_id
+    js.document.body.appendChild(html_canvas)
+
+    canvas = Canvas(init_device_sync(), html_canvas)
+    scene.init(canvas)
+
+    # Disable render() — patchedRequestAnimationFrame hangs in headless Chrome.
+    scene.render = lambda *a, **kw: None
+    scene._on_camera_changed = lambda *a, **kw: None
+
+    # Export scene to blob (scene.init already filled all GPU buffers)
+    blob = scene.export()
+    blob_b64 = base64.b64encode(blob).decode()
+
+    # Save blob and engine JS as static files for deferred loading.
+    # Use a unique hash to avoid collisions across notebooks.
+    import hashlib
+    blob_hash = hashlib.md5(blob).hexdigest()[:10]
+    static_dir = _get_static_dir()
+
+    # Save blob as a JS file (script src works on file:// unlike fetch)
+    scene_filename = f"scene_{blob_hash}.js"
+    scene_var = f"__webgpu_blob_{blob_hash}"
+    (static_dir / scene_filename).write_text(
+        f"window.{scene_var} = \"{blob_b64}\";"
+    )
+    scene_url = f"_static/webgpu_scenes/{scene_filename}"
+
+    if not _engine_emitted:
+        (static_dir / "engine.js").write_text(engine_js)
+        _engine_emitted = True
+    engine_url = "_static/webgpu_scenes/engine.js"
+
+    # Capture screenshot in a separate subprocess
+    screenshot_b64 = _capture_screenshot_subprocess(blob_b64, width, height)
+    screenshot_filename = f"screenshot_{blob_hash}.png"
+    if screenshot_b64:
+        (static_dir / screenshot_filename).write_bytes(base64.b64decode(screenshot_b64))
+    screenshot_url = f"_static/webgpu_scenes/{screenshot_filename}"
+
+    # Emit the lazy-load HTML: only screenshot + overlay, everything else loaded on click
+    lazy_html = f"""
+    <div id='{id_}root'
+         style="position: relative; width: {width}px; max-width: 100%; overflow: hidden;"
+    >
+        <img id='{id_}img'
+             src='{screenshot_url}'
+             style='width: {width}px; height: {height}px; max-width: 100%; display: block;'
+        />
+        <div id='{id_}overlay'
+             style='position: absolute; top: 0; left: 0; width: 100%; height: 100%;
+                    display: flex; align-items: center; justify-content: center;
+                    background: rgba(0,0,0,0); cursor: pointer; transition: background 0.2s;'
+             onmouseover="this.style.background='rgba(0,0,0,0.18)'; this.querySelector('span').style.opacity='1'"
+             onmouseout="this.style.background='rgba(0,0,0,0)'; this.querySelector('span').style.opacity='0'"
+             onclick="(function() {{ var r = document.getElementById('{id_}root'); if (r.__activated) return; r.__activated = true; function activate() {{ document.getElementById('{id_}img').style.display = 'none'; document.getElementById('{id_}overlay').style.display = 'none'; document.getElementById('{canvas_id}').style.display = 'block'; RenderEngine.create('{canvas_id}', window.{scene_var}); }} function loadBlob() {{ var s = document.createElement('script'); s.src = '{scene_url}'; s.onload = activate; document.head.appendChild(s); }} if (typeof RenderEngine === 'undefined') {{ var s = document.createElement('script'); s.src = '{engine_url}'; s.onload = loadBlob; document.head.appendChild(s); }} else {{ loadBlob(); }} }})()"
+        >
+            <span style='color: white; font-size: 1.3em; font-weight: bold;
+                         text-shadow: 0 1px 4px rgba(0,0,0,0.7); pointer-events: none;
+                         opacity: 0; transition: opacity 0.2s;'
+            >&#9654; Click to interact</span>
+        </div>
+        <canvas
+            id='{canvas_id}'
+            style='background-color: white; width: {width}px; height: {height}px;
+                   touch-action: none; max-width: 100%; display: none;'
+        ></canvas>
+        <div id='{id_}lilgui'
+             style='position: absolute; top: 0; right: 0; z-index: 10;'
+        ></div>
+    </div>
+    """
+
+    display(HTML(lazy_html))
+    return scene
+
+
+def _get_static_dir():
+    """Get or create the static directory for webgpu scene assets."""
+    from pathlib import Path
+    # Try to find the docs source _static directory
+    for candidate in [
+        Path("docs/_static/webgpu_scenes"),
+        Path("_static/webgpu_scenes"),
+        Path("webgpu_scenes"),
+    ]:
+        # Check if the parent _static dir exists (we're in a docs build)
+        if candidate.parent.exists():
+            candidate.mkdir(parents=True, exist_ok=True)
+            return candidate
+    # Fallback: create in current directory
+    fallback = Path("_static/webgpu_scenes")
+    fallback.mkdir(parents=True, exist_ok=True)
+    return fallback
+
+
+# Persistent screenshot worker process (launched once, handles all screenshots)
+_screenshot_worker = None
+
+
+def _capture_screenshot_subprocess(blob_b64, width, height):
+    """Send a screenshot request to the persistent worker process."""
+    import subprocess
+    import sys
+
+    global _screenshot_worker
+    if _screenshot_worker is None or _screenshot_worker.poll() is not None:
+        _screenshot_worker = subprocess.Popen(
+            [sys.executable, "-m", "webgpu.export.screenshot"],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        # Wait for READY signal
+        ready = _screenshot_worker.stdout.readline().strip()
+        if ready != "READY":
+            print(f"Warning: screenshot worker failed to start: {ready}")
+            _screenshot_worker = None
+            return ""
+
+    try:
+        _screenshot_worker.stdin.write(f"{width} {height} {blob_b64}\n")
+        _screenshot_worker.stdin.flush()
+        result = _screenshot_worker.stdout.readline().strip()
+        return result
+    except Exception as e:
+        print(f"Warning: screenshot capture failed: {e}")
+        return ""
+
+
 def _init_export_gpu():
     """Start headless Chrome with WebGPU for buffer export during nbconvert.
 
@@ -265,7 +422,11 @@ def Draw(
     if is_exporting:
         # Launch headless Chrome for GPU access during export
         _init_export_gpu()
-        Draw = _DrawHTML
+        is_lazy_load = "WEBGPU_LAZY_LOAD" in os.environ
+        if is_lazy_load:
+            Draw = _DrawHTMLLazy
+        else:
+            Draw = _DrawHTML
     else:
         # Not exporting and not running in pyodide -> Start a websocket server
         # and wait for the client to connect.
diff --git a/webgpu/notebook_to_html.py b/webgpu/notebook_to_html.py
@@ -36,6 +36,8 @@ def main():
     parser.add_argument("-o", "--output", help="Output HTML path (default: same name .html)")
     parser.add_argument("--timeout", type=int, default=120, help="Cell execution timeout in seconds")
     parser.add_argument("--kernel", default=None, help="Jupyter kernel name")
+    parser.add_argument("--lazy-load", action="store_true", default=False,
+                        help="Embed screenshot previews that load the WebGPU scene on click")
     args = parser.parse_args()
 
     if not os.path.exists(args.notebook):
@@ -45,6 +47,8 @@ def main():
     # Build environment
     env = os.environ.copy()
     env["WEBGPU_EXPORTING"] = "1"
+    if args.lazy_load:
+        env["WEBGPU_LAZY_LOAD"] = "1"
 
     if not _has_real_gpu():
         icd = _find_lavapipe_icd()
