fdanyone_app.py

"""Gradio and ZeroGPU wiring for the 4DAnyone Rerun streaming runtime.

The four-link event chain is ``begin -> prepare_cpu -> run_gpu -> publish_cpu``.
ZeroGPU runs decorated callbacks in forked children, so each link creates its
own recording after ``rr.cleanup_if_forked_child()`` and uses the shared token
as its recording ID. The merged ``run_gpu`` remains one allocation because its
motion result contains live state that cannot cross a process boundary.
"""

from __future__ import annotations

import logging
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import gradio as gr
import rerun as rr
import spaces
from gradio_rerun import Rerun

from fdanyone.errors import FourDAnyoneError
from fdanyone.rerun_streaming import (
    APPLICATION_ID,
    EXAMPLE_DIR,
    GPU_DURATION,
    STREAM_SMOKE,
    RunSpec,
    Session,
    begin_phase,
    generate_phase,
    motion_phase,
    new_spec,
    publish_phase,
    request_stop,
    smoke_generate_phase,
    smoke_motion_phase,
    smoke_publish_phase,
    source_phase,
)


@dataclass(frozen=True)
class Link:
    """One chain link's own recording, and the sink its bytes are drained from."""

    recording: rr.RecordingStream
    """Explicit stream this link and its pipeline hooks log through."""
    stream: rr.BinaryStream
    """Binary sink feeding the Gradio viewer."""

    def read(self) -> bytes | None:
        """Return bytes buffered since the last read."""

        return self.stream.read()


def open_link(token: str) -> Link:
    """Open this callback's own recording, in whatever process the callback runs in.

    Called first thing in every link, GPU ones included, where "this callback"
    means the forked child. A stream inherited across the fork cannot be used:
    the SDK compares pids on flush and raises "Fork detected during flush". A
    stream built here belongs to this process, and sharing ``token`` as the
    recording id is what makes the viewer treat all four links as one recording.

    ``cleanup_if_forked_child`` drops whatever streams were inherited. The SDK
    registers it with ``os.register_at_fork`` already, and ``rr.init`` calls it
    outright for the same reason, so this is belt and braces and a no-op outside
    a child. It is called unconditionally: guarding it on the name existing
    would turn a rename in the SDK into a silent return of the very bug this
    function exists to prevent.
    """

    rr.cleanup_if_forked_child()
    recording: rr.RecordingStream = rr.RecordingStream(APPLICATION_ID, recording_id=token)
    return Link(recording=recording, stream=recording.binary_stream())


def begin(
    video: str | None, start_time: float, seed: int
) -> Iterator[tuple[RunSpec, bytes | None, str, Any]]:
    """Validate the input on CPU, open a recording, and switch to the outputs."""

    if video is None:
        raise gr.Error("Upload a video, or pick the bundled example.")
    try:
        spec: RunSpec = new_spec(Path(video), float(start_time), int(seed))
    except FourDAnyoneError as exc:
        raise gr.Error(str(exc)) from None
    link: Link = open_link(spec.token)
    begin_phase(link.recording, spec)
    yield spec, link.read(), "Preparing the source clip.", gr.Tabs(selected="outputs")


def prepare_cpu(spec: RunSpec) -> Iterator[tuple[bytes | None, str]]:
    """Put the source clip on the frame timeline before any GPU work starts."""

    link: Link = open_link(spec.token)
    label: str = source_phase(link.recording, spec)
    yield link.read(), label


@spaces.GPU(duration=GPU_DURATION)
def run_gpu(spec: RunSpec) -> Iterator[tuple[dict | None, bytes | None, str]]:
    """Recover the motion and generate the six views, inside one allocation.

    The two phases have to share a process. ``PreparedRun`` carries the decoded
    clip, the open conditioning artifacts, and a completion barrier for the
    skeleton renderer, so it cannot be pickled from one ZeroGPU worker into
    another — and every ``@spaces.GPU`` function gets a worker of its own.

    Everything in here runs in the forked child, including ``open_link``: this
    is the only process that may build the recording it logs through.
    """

    link: Link = open_link(spec.token)
    session: Session = Session(spec)
    try:
        motion: Iterator[str] = (
            smoke_motion_phase(link.recording, session)
            if STREAM_SMOKE
            else motion_phase(link.recording, session)
        )
        for label in motion:
            yield None, link.read(), label
        generate: Iterator[str] = (
            smoke_generate_phase(link.recording, session)
            if STREAM_SMOKE
            else generate_phase(link.recording, session)
        )
        for label in generate:
            yield None, link.read(), label
    except FourDAnyoneError as exc:
        raise gr.Error(str(exc)) from None
    yield session.summary, link.read(), "Publishing the result."


def publish_cpu(spec: RunSpec, summary: dict | None) -> Iterator[tuple[bytes | None, str]]:
    """Attach the finished rig and its six videos, off the GPU allocation."""

    if summary is None:
        raise gr.Error("The generation phase did not publish a result.")
    link: Link = open_link(spec.token)
    label: str = (
        smoke_publish_phase(link.recording, spec)
        if STREAM_SMOKE
        else publish_phase(link.recording, spec, summary)
    )
    yield link.read(), label


DESCRIPTION: str = """
# 4DAnyone × Rerun

One monocular clip of a person becomes six synchronized novel views. GVHMR
recovers the SMPL-X motion, a Wan 2.2 diffusion transformer generates every
view, and each phase streams into the Rerun viewer as it happens.
"""

SOURCE_CLIP_HEIGHT: int = 360
"""Display height of the source preview. A portrait clip scaled to the column
width is taller than the fold, which pushes every control below it out of sight."""

STATUS_CSS: str = """
#run-status {
    display: flex;
    flex-direction: column;
    justify-content: center;
    /* Tall enough for the progress overlay Gradio draws over a running event. */
    min-height: 5rem;
    /* Gradio writes `overflow: auto` inline on every block. */
    overflow: visible !important;
    padding: 0.75rem 1rem;
    border-radius: var(--radius-lg);
    background: var(--background-fill-secondary);
}
#run-status p {
    font-size: 1.15rem;
    line-height: 1.5;
    margin: 0;
}
"""
"""Banner styling for the run status, the only readout of a multi-minute run.
Gradio otherwise gives a Markdown block the body font and a height its own text
overflows, which turns the line into a scrolling sliver. Gradio 6 takes ``css``
on ``launch``, not on the ``Blocks`` constructor."""


def build_demo() -> gr.Blocks:
    """Assemble the persistent viewer, the inputs, and the run chain.

    Controls sit in a narrow column and the viewer fills a wide one beside it,
    the same split every other Rerun Space here uses. Stacking them instead
    pushes the viewer off the fold on a laptop, and the viewer is the demo. It
    stays outside the tabs for the same reason it is created once: a tab switch
    would tear down the component and drop the stream mid-run.

    The chain is joined by ``success`` rather than ``then``: each link's inputs
    are the previous link's outputs, so a link that failed leaves the next one
    nothing to work with, and running it anyway would bury the real error under
    a second, meaningless one.
    """

    # Serve the example clips in place; without this Gradio copies all 209 MB
    # of them into its temp cache at construction.
    gr.set_static_paths([EXAMPLE_DIR])
    with gr.Blocks(title="4DAnyone × Rerun") as demo:
        gr.Markdown(DESCRIPTION)
        spec: gr.State = gr.State(None)
        summary: gr.State = gr.State(None)
        with gr.Row():
            with gr.Column(scale=1):
                with gr.Tabs() as tabs:
                    with gr.Tab("Input", id="input"):
                        video: gr.Video = gr.Video(
                            label="Source clip", sources=["upload"], height=SOURCE_CLIP_HEIGHT
                        )
                        start_time: gr.Number = gr.Number(
                            value=0.0, label="Start time (seconds)", minimum=0.0
                        )
                        seed: gr.Number = gr.Number(value=0, label="Seed", precision=0, minimum=0)
                        # A single input keeps the compact thumbnail gallery; a
                        # second column would turn it into a labeled table.
                        gr.Examples(
                            examples=[[str(clip)] for clip in sorted(EXAMPLE_DIR.glob("*.mp4"))],
                            inputs=[video],
                            cache_examples=False,
                            examples_per_page=4,
                        )
                    with gr.Tab("Outputs", id="outputs"):
                        gr.Markdown(
                            "The viewer beside this switches layout with the run: detections on the "
                            "source clip, then a grid of per-step previews, then the finished rig "
                            "playing on a loop."
                        )
                with gr.Row():
                    run_button: gr.Button = gr.Button("Run", variant="primary")
                    stop_button: gr.Button = gr.Button("Stop", variant="stop")
                status: gr.Markdown = gr.Markdown(
                    "Upload a clip, or pick the example, then press Run.", elem_id="run-status"
                )
            with gr.Column(scale=3):
                viewer: Rerun = Rerun(
                    streaming=True,
                    height=760,
                    panel_states={"time": "collapsed", "blueprint": "hidden", "selection": "hidden"},
                )

        started = run_button.click(begin, [video, start_time, seed], [spec, viewer, status, tabs])
        prepared = started.success(prepare_cpu, spec, [viewer, status])
        generated = prepared.success(run_gpu, spec, [summary, viewer, status])
        published = generated.success(publish_cpu, [spec, summary], [viewer, status])
        # `cancels` detaches the UI at once, but the pipeline itself only halts
        # when it reads the sentinel `request_stop` writes: the worker thread
        # (and on ZeroGPU, the forked child) never sees a Gradio cancellation.
        stop_button.click(
            request_stop, spec, status, cancels=[started, prepared, generated, published]
        )
    return demo


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
    build_demo().queue(default_concurrency_limit=1).launch(ssr_mode=False, css=STATUS_CSS)