window_viewer.py

"""Local live-window viewer. Shared rendering assets remain in ``space/``."""

from pathlib import Path

import gradio as gr

from space.viewer import CSS as BASE_CSS


HERE = Path(__file__).resolve().parent
ASSETS = HERE / "space"
HTML = """
<div class="motion-viewer">
  <div class="viewer-heading">
    <span class="viewer-action">A person walks forward.</span>
    <span class="viewer-status" role="status">Ready</span>
  </div>
  <div class="viewer-instructions" aria-label="Controls">
    <span><strong>Start</strong>, then click the scene</span>
    <span><kbd>WASD</kbd> · Move</span>
    <!-- action-key-help -->
  </div>
  <div class="viewer-stage">
    <div class="viewer-focus-hint">Click to control</div>
    <div class="viewer-body-notice"><span class="viewer-body-status" role="status"></span><button class="viewer-body-retry" type="button" hidden>Retry</button></div>
    <button class="viewer-joystick" type="button" aria-label="Movement joystick"><span class="viewer-joystick-thumb" aria-hidden="true"></span></button>
  </div>
  <div class="viewer-controls">
    <!-- action-slots -->
    <span><kbd data-key="KeyW">W</kbd><kbd data-key="KeyA">A</kbd><kbd data-key="KeyS">S</kbd><kbd data-key="KeyD">D</kbd> / Left stick · Move</span>
    <span>Drag / Right stick · Camera</span><span>Scroll · Zoom</span>
    <span class="viewer-gamepad-status" role="status"></span>
  </div>
  <div class="viewer-stats">
    <div><span>Inference</span><strong data-stat="generation">— fps</strong></div>
    <div><span>Playback</span><strong data-stat="playback">0 fps</strong></div>
  </div>
</div>
"""
ACTION_SLOTS_HTML = """
<span class="viewer-action-slots" role="group" aria-label="Actions">
  <span class="viewer-action-shortcut">Switch action · keyboard 1–4 or mouse click:</span>
  <button class="viewer-action-slot" data-action-slot="0" type="button" style="--slot-color: #397ad5" aria-label="Action 1">1</button>
  <button class="viewer-action-slot" data-action-slot="1" type="button" style="--slot-color: #8b5cdd" aria-label="Action 2">2</button>
  <button class="viewer-action-slot" data-action-slot="2" type="button" style="--slot-color: #2d9e72" aria-label="Action 3">3</button>
  <button class="viewer-action-slot" data-action-slot="3" type="button" style="--slot-color: #d55a9a" aria-label="Action 4">4</button>
</span>
"""
CSS = BASE_CSS + """
.viewer-instructions { display: flex; flex-wrap: wrap; gap: 6px 20px; padding: 0 18px 14px; color: #415563; font-size: 13px; line-height: 1.6; }
.viewer-instructions strong { color: #243b4c; font-weight: 600; }
.viewer-action-help { flex-basis: 100%; }
.viewer-instructions kbd { display: inline-block; padding: 1px 6px; border: 1px solid #cbd6df; border-radius: 4px; background: #f2f6f9; color: #243b4c; font: inherit; font-weight: 600; }
.viewer-action-shortcut { margin-right: 3px; color: #415563; font-size: 12px; }
.viewer-controls { padding-bottom: 12px; color: #667782; }
.viewer-action-slots { display: inline-flex; flex-wrap: wrap; align-items: center; gap: 5px; }
.viewer-action-slot { width: 27px; height: 27px; padding: 0; border: 1px solid var(--slot-color); border-radius: 6px; background: transparent; color: var(--slot-color); font: inherit; font-weight: 600; cursor: pointer; }
.viewer-action-slot.is-selected { background: var(--slot-color); color: white; }
.viewer-action-slot:focus-visible { outline: 2px solid var(--slot-color); outline-offset: 2px; }
.viewer-gamepad-status { margin-left: auto; font-size: 11px; }
.viewer-gamepad-status:empty { display: none; }
.viewer-stats { grid-template-columns: repeat(2, minmax(0, 1fr)); }
.viewer-stage { height: clamp(280px, 40vw, 420px); }
.viewer-joystick { position: absolute; bottom: 18px; left: 18px; z-index: 2; display: block; box-sizing: border-box; width: 96px; height: 96px; margin: 0; padding: 0; border: 1px solid #c8d4dd; border-radius: 50%; background: rgba(239, 244, 248, .82); box-shadow: 0 2px 8px rgba(35, 51, 64, .08); appearance: none; touch-action: none; user-select: none; -webkit-user-select: none; -webkit-tap-highlight-color: transparent; cursor: grab; }
.viewer-joystick:focus-visible { outline: 2px solid #397ad5; outline-offset: 3px; }
.viewer-joystick.is-active { background: rgba(229, 239, 248, .9); border-color: #8baac7; cursor: grabbing; }
.viewer-joystick-thumb { position: absolute; top: 50%; left: 50%; width: 28%; height: 28%; border-radius: 50%; background: #738da3; box-shadow: 0 1px 3px rgba(35, 51, 64, .16); transform: translate(-50%, -50%); pointer-events: none; }
.viewer-joystick.is-active .viewer-joystick-thumb { background: #397ad5; }
@media (max-width: 600px) { .viewer-joystick { bottom: 14px; left: 14px; width: 88px; height: 88px; } }
"""


def create_viewer(*, require_body_ready: bool = False, max_buffer_frames: int | None = None, skeleton_only: bool = False, action_slots: bool = False) -> gr.HTML:
    """Create the local viewer with actual-motion-only playback.

    Existing mesh, mesh_pose, skeleton, input and legacy path modes are retained.
    For a new path revision send a strictly increasing sequence, the absolute
    control_replace_from_frame, and control_points containing the complete new
    suffix. Omitted future points are removed; earlier accepted history is fixed.
    Replace controls before adding new generated frames in the same packet: its
    start_frame is the correct replacement boundary when a packet commits frames.
    An empty frames list supports startup/replanning without speculative motion.
    Each point's slot/prompt must be its actual conditioning metadata. Replanning
    may update pending-frame text and color; committed history stays fixed.
    The viewer never recolors committed frames from current selection. Send explicit
    target_point=None or spring_point=None to clear those markers.

    Send no more than three generated frames per packet. By default playback and
    buffering are identical to the existing viewer for a clean A/B comparison.
    max_buffer_frames is an explicit optional received-pose buffer limit. During
    an underrun the last actual pose is held; controls never synthesize motion.
    Session changes reset both path and playback; retired sessions and stale
    sequence numbers are ignored.

    skeleton_only skips SMPL body assets and renders the supplied 22-joint poses
    as bones and joints. It signals body_ready immediately, including when the
    canvas fallback is used. Playback always follows the packet's fps.
    Otherwise, packets may set representation="mesh" or "skeleton", or model
    to "SEED" or "HumanML3D". The explicit representation takes priority. Send
    this on a fresh preparation packet to remove the old mesh before new motion
    arrives. skeleton_only=True always overrides packet rendering requests.
    action_slots adds four action buttons and page-level 1-4 shortcuts; it is
    disabled by default so the single-prompt viewer retains its existing input.

    The on-screen joystick accepts mouse or touch with camera-relative analog
    movement and returns to center on release. WASD takes priority over an active
    joystick, which takes priority over a connected gamepad's left stick.
    Standard-mapped browser gamepads use left-stick movement and right-stick
    orbit. Prompt and speed are changed in the form
    between runs. The controller works across the foreground
    page except while an editor has focus, including inside shadow DOM. Editing,
    window focus loss, tab hiding or disconnect clears controller input; neutral
    sticks rearm it. The joystick clears on Stop, focus loss and session changes.
    Keyboard uses the viewer canvas; interacting with the joystick focuses it.
    Keyboard and mouse remain available. Controller visibility
    and mapping depend on browser/OS support and may need an initial button press.
    """
    if max_buffer_frames is not None and (not isinstance(max_buffer_frames, int) or isinstance(max_buffer_frames, bool) or max_buffer_frames < 1):
        raise ValueError("max_buffer_frames must be None or a positive integer")
    vendor = (ASSETS / "three.min.js").read_text(encoding="utf-8")
    body_script = ""
    body_props = {}
    if not skeleton_only:
        body_script = (ASSETS / "body.js").read_text(encoding="utf-8")
        body_manifest = ASSETS / "assets" / "body.json"
        body_data = ASSETS / "assets" / "body.bin"
        gr.set_static_paths(paths=[body_manifest, body_data])
        body_props = {
            "body_manifest_url": "/gradio_api/file=" + body_manifest.resolve().as_posix(),
            "body_data_url": "/gradio_api/file=" + body_data.resolve().as_posix(),
        }
    script = (HERE / "window_viewer.js").read_text(encoding="utf-8")
    js = "if (!globalThis.THREE) {\n" + vendor + "\n}\n" + body_script + "\n" + script
    return gr.HTML(
        value={"session_id": "", "sequence": -1, "frames": [], "status": "idle"},
        html_template=HTML.replace("<!-- action-slots -->", ACTION_SLOTS_HTML if action_slots else "").replace(
            "<!-- action-key-help -->",
            '<span class="viewer-action-help"><strong>Switch actions:</strong> press <kbd>1–4</kbd> on your <strong>keyboard</strong>, or <strong>click</strong> the colored 1–4 buttons below.</span>' if action_slots else ""),
        css_template=CSS,
        js_on_load=js,
        apply_default_css=False,
        container=False,
        padding=False,
        elem_id="live-motion-viewer",
        skeleton_only=skeleton_only,
        action_slots=action_slots,
        require_body_ready=require_body_ready,
        max_buffer_frames=max_buffer_frames,
        **body_props,
    )