workflow_app.py

# SPDX-License-Identifier: Apache-2.0
"""Viggle-Animate as a gr.Workflow: the same prepare -> render -> finish pipeline as the Blocks
app in ../space_viggle/app.py, exposed as bound function nodes on gradio's node canvas.

Written by @akhaliq (https://huggingface.co/akhaliq) as PR #1 on Viggle/viggle-animate
(https://huggingface.co/spaces/Viggle/viggle-animate/discussions/1). Changes from that PR:

  - `prepare` declares `request` and forwards it. Without it app.py's `gate` got None and keyed
    every visitor to "?", i.e. one daily repaint counter for the whole Space.
  - A character image wins over the pre-painted frame. That slot opens preloaded with the prop
    example, and app.py uses a painted frame first, so an uploaded character was silently
    ignored unless the visitor also cleared the preload.
  - The job lives server-side in JOBS; the canvas only carries its id. Values on the canvas
    round-trip through the browser, so a job dict there let any client choose `frames` and
    the file paths `render` and `finish` read. The Blocks app keeps them in gr.State.
  - No AOTI build node and no extra OAuth scopes: the packages are built by the main Space and
    shared through the weights repo, so nothing here needs to write anywhere.

The pipeline itself is app.py, imported rather than reimplemented: ZeroGPU admission/billing,
the audio pin, the daily repaint gate and AOTI wiring all live there.
"""

import json
import uuid
from pathlib import Path

import gradio as gr
from gradio_client.utils import encode_file_path

import app as va

GRAPH = Path(__file__).parent / "workflow.json"
JOBS = {}  # job id -> what prepare/render produced; tiny, and a restart empties it anyway


def _path(v):
    """Media ports arrive as {"path": ...}/{"url": ...} dicts on the canvas and as bare
    filepaths through the API executor; accept all three. A canvas file value often carries
    NO path at all -- only a percent-encoded /gradio_api/file= URL -- so the prefix is
    stripped and unquoted back to a local path."""
    if isinstance(v, dict):
        v = v.get("path") or v.get("url")
    if isinstance(v, str) and v.startswith("/gradio_api/file="):
        from urllib.parse import unquote
        v = unquote(v.removeprefix("/gradio_api/file="))
    return v


def _media(p):
    """The {path, url, is_file} shape the canvas needs to display media; a bare /tmp path has
    no URL the browser can load. Workflow.launch adds the tempdir to allowed_paths."""
    return ({"path": p, "url": f"/gradio_api/file={encode_file_path(p)}", "is_file": True}
            if p else None)


def _job(job):
    if job not in JOBS:
        raise gr.Error("this run predates a Space restart -- run Prepare again")
    return JOBS[job]


def prepare(video, character=None, painted=None, seed: float = 42,
            request: gr.Request = None):
    """CPU stage: retime/crop the driving clip onto the ladder canvas, repaint frame 0 with
    gpt-image-2 (unless only a pre-painted frame is supplied), extract the soundtrack for the
    pin. `request` is injected by gradio from the annotation; app.py's daily gate keys on it."""
    character, painted = _path(character), _path(painted)
    if character:
        painted = None
    ref, cond, keep, frames, apath, _shown, _out, msg = va.prepare(
        _path(video), character, painted, request)
    job = uuid.uuid4().hex
    JOBS[job] = {"ref": ref, "cond": cond, "keep": keep, "frames": frames, "apath": apath,
                 "seed": int(seed if seed is not None else 42)}
    return _media(ref), job, msg


def render(job: str):
    """GPU stage: 4-step ref2va render on ZeroGPU xlarge (va.render carries the
    @spaces.GPU decorator and its admission budget)."""
    j = _job(job)
    (j["raw"], _keep, _apath), msg = va.render(j["ref"], j["cond"], j["keep"], j["frames"],
                                               j["apath"], j["seed"])
    return job, msg


def finish(job: str):
    """CPU stage: trim the frozen 17k+5 tail, watermark, remux the original soundtrack."""
    j = _job(job)
    if "raw" not in j:
        raise gr.Error("run Render on this job first")
    out = va.finish((j["raw"], j["keep"], j["apath"]))
    return _media(out), f"trimmed to {j['keep']:.1f}s, watermarked, remuxed"


def _media_ref(name):
    """Preload value for a media reference node: a bundled example file, so the canvas opens
    with a ready-to-run pipeline. Absolute paths make this generation-time only."""
    return _media(str(va.EX / name))


def _port(pid, label, ptype, **kw):
    return {"id": pid, "label": label, "type": ptype, **kw}


def _node(nid, label, role, asset_type, port_label, **kw):
    return {"id": nid, "label": label, "role": role, "asset_type": asset_type,
            "inputs": [_port("in", port_label, asset_type)],
            "outputs": [_port("out", port_label, asset_type)], **kw}


def _edge(eid, src, sport, dst, dport, ptype):
    return {"id": eid, "from_node_id": src, "from_port_id": sport,
            "to_node_id": dst, "to_port_id": dport, "type": ptype}


def _default_graph():
    return {
        "schema_version": "2",
        "name": "Viggle-Animate — workflow by @akhaliq",
        "references": [
            # prop: deterministic at seed 42 and needs no OpenAI call (frame pre-painted).
            _node("ref_video", "Driving Video", "reference", "video", "Video",
                  data={"out": _media_ref("prop_drive.mp4")}),
            _node("ref_character", "Character Image (repainted into frame 0)", "reference",
                  "image", "Image"),
            _node("ref_painted", "Pre-painted Frame 0 (ignored if a character is given)",
                  "reference", "image", "Image", data={"out": _media_ref("prop_ref.png")}),
            _node("ref_seed", "Seed", "reference", "number", "Seed", data={"out": 42}),
        ],
        "operators": [
            {"id": "op_prepare", "label": "Prepare — retime, repaint frame 0",
             "role": "operator", "kind": "fn", "fn": "prepare",
             "inputs": [_port("video", "video", "video", required=True),
                        _port("character", "character", "image"),
                        _port("painted", "painted", "image"),
                        _port("seed", "seed", "number")],
             "outputs": [_port("ref", "Repainted frame", "image", output_index=0),
                         _port("job", "job", "text", output_index=1),
                         _port("log", "log", "text", output_index=2)]},
            {"id": "op_render", "label": "Render — 4 steps on ZeroGPU xlarge",
             "role": "operator", "kind": "fn", "fn": "render",
             "inputs": [_port("job", "job", "text", required=True)],
             "outputs": [_port("job", "job", "text", output_index=0),
                         _port("log", "log", "text", output_index=1)]},
            {"id": "op_finish", "label": "Finish — trim, watermark, remux audio",
             "role": "operator", "kind": "fn", "fn": "finish",
             "inputs": [_port("job", "job", "text", required=True)],
             "outputs": [_port("out_0", "Video", "video", output_index=0),
                         _port("log", "log", "text", output_index=1)]},
        ],
        "subjects": [
            _node("sub_frame", "Repainted Frame", "subject", "image", "Image"),
            _node("sub_video", "Output Video", "subject", "video", "Video"),
            _node("sub_plog", "Prepare Log", "subject", "text", "Text"),
            _node("sub_rlog", "Render Log", "subject", "text", "Text"),
            _node("sub_flog", "Finish Log", "subject", "text", "Text"),
        ],
        "edges": [
            _edge("e1", "ref_video", "out", "op_prepare", "video", "video"),
            _edge("e2", "ref_character", "out", "op_prepare", "character", "image"),
            _edge("e3", "ref_painted", "out", "op_prepare", "painted", "image"),
            _edge("e4", "ref_seed", "out", "op_prepare", "seed", "number"),
            _edge("e5", "op_prepare", "ref", "sub_frame", "in", "image"),
            _edge("e6", "op_prepare", "job", "op_render", "job", "text"),
            _edge("e7", "op_render", "job", "op_finish", "job", "text"),
            _edge("e8", "op_finish", "out_0", "sub_video", "in", "video"),
            _edge("e9", "op_prepare", "log", "sub_plog", "in", "text"),
            _edge("e10", "op_render", "log", "sub_rlog", "in", "text"),
            _edge("e11", "op_finish", "log", "sub_flog", "in", "text"),
        ],
    }


if not GRAPH.exists():
    GRAPH.write_text(json.dumps(_default_graph(), indent=2))

gr.Workflow(
    graph=str(GRAPH),
    bind={"prepare": prepare, "render": render, "finish": finish},
).launch(show_error=True, max_file_size="200mb", allowed_paths=[str(va.EX)])