diff --git a/README.md b/README.md index ed06368..c78c474 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,7 @@ The orchestrator is a **transparent reverse proxy**: every endpoint you expose i - **HTTP** request/response — the common case. (`hello-world`, `tiles`) - **HTTP + SSE** — streamed / token responses. (`vllm`) -- **Trickle** — continuous realtime video in/out. (`echo`) +- **Trickle** — continuous realtime video in/out. (`echo`, `flux-klein`) - **WebSocket** — long-lived bidirectional sessions. (external: `scope`) Need a schema that isn't here? [Open an issue](https://github.com/livepeer/live-runner-example-apps/issues). @@ -40,12 +40,13 @@ Need a schema that isn't here? [Open an issue](https://github.com/livepeer/live- ## Examples -| Example | Goal | Registration | Mode | Transport | -| ------------------------------ | ----------------------------------------------- | ------------ | ---------------------------------- | ----------------- | -| [`hello-world`](./hello-world) | The simplest app: one request, one response | dynamic | persistent (single-shot by nature) | HTTP (JSON) | -| [`tiles`](./tiles) | Capacity fan-out — one session per tile | dynamic | persistent (single-shot by nature) | HTTP (base64 PNG) | -| [`echo`](./echo) | Realtime video, transformed and echoed back | dynamic | persistent | trickle | -| [`vllm`](./vllm) | Drop-in OpenAI API; the client stays unmodified | static | persistent (single-shot by nature) | HTTP + SSE | +| Example | Goal | Registration | Mode | Transport | +| ------------------------------ | -------------------------------------------------- | ------------ | ---------------------------------- | ----------------- | +| [`hello-world`](./hello-world) | The simplest app: one request, one response | dynamic | persistent (single-shot by nature) | HTTP (JSON) | +| [`tiles`](./tiles) | Capacity fan-out — one session per tile | dynamic | persistent (single-shot by nature) | HTTP (base64 PNG) | +| [`echo`](./echo) | Realtime video, transformed and echoed back | dynamic | persistent | trickle | +| [`vllm`](./vllm) | Drop-in OpenAI API; the client stays unmodified | static | persistent (single-shot by nature) | HTTP + SSE | +| [`flux-klein`](./flux-klein) | Realtime video diffusion, steered by a live prompt | static | persistent | trickle | Start with `hello-world` (the smallest end-to-end path); the others each layer on one new idea. More will follow, including a full example that exercises every feature. Each is self-contained and runs **offchain** (free, no wallet); most also run **on-chain** (paid) — see each README. @@ -54,7 +55,7 @@ Start with `hello-world` (the smallest end-to-end path); the others each layer o How the app attaches to the orchestrator: - **Dynamic** — the app self-registers via the SDK (`register_runner`) and heartbeats; the orchestrator drops it when heartbeats stop. Best for apps that come and go. (`hello-world`, `echo`) -- **Static** — the orchestrator is configured with the app's URL in a `runners.json` and health-polls it; the app needs no SDK. Best for fixed, long-running deployments. (`vllm`) +- **Static** — the orchestrator is configured with the app's URL in a `runners.json` and health-polls it; the app needs no SDK. Best for fixed, long-running deployments. (`vllm`, `flux-klein`) The arrow flips — dynamic, the app announces itself; static, the orchestrator is told about a passive app: @@ -74,7 +75,7 @@ flowchart LR Chosen _at_ registration (above); **defaults to `persistent`** — set on both `register_runner(...)` and in `runners.json`. The examples set it explicitly. -- **Persistent** — a held-open session billed per second of wall-clock. Best for realtime / streaming. (`echo`) +- **Persistent** — a held-open session billed per second of wall-clock. Best for realtime / streaming. (`echo`, `flux-klein`) - **Single-shot** — one request in, one response out. Best for batch / request-response. (`hello-world`, `vllm` are single-shot by nature.) > [!IMPORTANT] diff --git a/flux-klein/.env.example b/flux-klein/.env.example new file mode 100644 index 0000000..56d870c --- /dev/null +++ b/flux-klein/.env.example @@ -0,0 +1,24 @@ +# Copy to .env (gitignored) and fill in. Never commit secrets. +# Keystore dirs: absolute paths OUTSIDE this repo, mounted read-only. +# Only needed for the on-chain (paid) path (docker-compose.onchain.yml). + +NETWORK=arbitrum-one-mainnet +ETH_RPC_URL=https://arb1.arbitrum.io/rpc + +# Signer (payer): needs an on-chain deposit + reserve. +SIGNER_KEYSTORE_DIR=/absolute/path/to/signer-keystore +SIGNER_ETH_ACCT=0xYourSignerAddress +SIGNER_ETH_PASSWORD=your-signer-keystore-password + +# Orchestrator operating key (split-key): needs ETH for gas to redeem tickets. +ORCH_KEYSTORE_DIR=/absolute/path/to/operator-keystore +ORCH_ETH_ACCT=0xYourOperatorAddress +ORCH_ETH_PASSWORD=your-operator-keystore-password +# Registered orch = ticket recipient (-ethOrchAddr); empty = use the operating key. +ORCH_ONCHAIN_ADDR=0xYourRegisteredOrchestrator + +# Pricing (on-chain), in USD, converted to wei via the price feed. Keep +# PIXELS_PER_UNIT small — too large and the per-unit price floors to 0 wei. +PRICE_PER_UNIT=1 +PIXELS_PER_UNIT=1000 +MAX_PRICE_PER_UNIT=0.10USD diff --git a/flux-klein/.gitignore b/flux-klein/.gitignore new file mode 100644 index 0000000..fe63024 --- /dev/null +++ b/flux-klein/.gitignore @@ -0,0 +1,5 @@ +# Output written by client.py (MPEG-TS stream) +flux-klein-out.ts + +# HF weights cache mounted into the container (~15 GB, downloaded on first boot) +models/ diff --git a/flux-klein/Dockerfile b/flux-klein/Dockerfile new file mode 100644 index 0000000..bb58419 --- /dev/null +++ b/flux-klein/Dockerfile @@ -0,0 +1,47 @@ +# Self-contained realtime FLUX Klein app for the new Livepeer runner. +# +# Depends only on pip packages + our flux_klein.py/runner.py — NOT the Daydream +# Scope platform. FLUX.2-klein-4B needs a diffusers build that ships +# Flux2KleinPipeline (not yet in a stable release), so diffusers is installed +# from git. CUDA 12.8 / torch 2.7.1 to match the GPU wheels. +FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive PYTHONUNBUFFERED=1 +ENV HF_HUB_ENABLE_HF_TRANSFER=1 + +# Python 3.12 (the SDK declares >=3.12; FLUX/diffusers run fine there) + git for +# the pip-from-git installs. +RUN apt-get update && apt-get install -y --no-install-recommends \ + software-properties-common ca-certificates curl git \ + && add-apt-repository ppa:deadsnakes/ppa \ + && apt-get update && apt-get install -y --no-install-recommends \ + python3.12 python3.12-venv python3.12-dev \ + && ln -sf /usr/bin/python3.12 /usr/local/bin/python \ + && curl -sS https://bootstrap.pypa.io/get-pip.py | python \ + && python -m pip --version \ + && apt-get clean && rm -rf /var/lib/apt/lists/* + +# Torch (CUDA 12.8 wheels) first so the diffusers install below doesn't pull a +# CPU build over the top. +RUN python -m pip install --no-cache-dir \ + torch==2.7.1+cu128 torchvision==0.22.1+cu128 \ + --index-url https://download.pytorch.org/whl/cu128 + +# diffusers from git: Flux2KleinPipeline (and the flux2 internals refine_frame +# reaches into) is not in a stable diffusers release yet. transformers/accelerate +# for the Qwen3 text encoder + device management. +RUN python -m pip install --no-cache-dir \ + "diffusers @ git+https://github.com/huggingface/diffusers.git" \ + transformers accelerate \ + av aiohttp pillow numpy hf_transfer \ + "livepeer-gateway @ git+https://github.com/livepeer/livepeer-python-gateway@rs/live-runner-session-payments" + +WORKDIR /app +COPY flux_klein.py runner.py ./ + +# HF model weights are cached under /models/hf; mount a host dir so they persist +# across restarts and don't need to be re-downloaded (~15 GB for klein-4B). +ENV HUGGINGFACE_HUB_CACHE=/models/hf +EXPOSE 8720 + +ENTRYPOINT ["python", "runner.py"] diff --git a/flux-klein/README.md b/flux-klein/README.md new file mode 100644 index 0000000..a670f01 --- /dev/null +++ b/flux-klein/README.md @@ -0,0 +1,151 @@ +# FLUX Klein app (realtime trickle video-to-video) + +A realtime **FLUX.2-klein-4B** image-to-image app on the Livepeer network. It receives a live video stream over **trickle**, transforms every frame against a text prompt using a **Krea-style feedback loop**, and streams the result back. Prompt, noise seed and camera anchoring are all adjustable live over the control path. + +Like [`vllm`](../vllm), it's attached as a **static runner** — the best fit for a fixed, long-running GPU deployment: the orchestrator is configured with the app's URL in `runners.json` and health-polls `/health`, so there's no SDK registrar or heartbeat in the app. + +| | | +| ------------ | ------------------------------------------- | +| App id | `livepeer-example/flux-klein` | +| Runner mode | persistent (held-open session) | +| Registration | static (orchestrator config + health poll) | +| Transport | trickle (realtime video in/out + `/update`) | +| Port | 8720 | + +**Requires an NVIDIA GPU** with ~13 GB VRAM. Prerequisites (Docker, `uv`, the not-yet-released SDK) and the shared on-chain/payment setup are in the [repo README](../README.md). + +> [!NOTE] +> This example pins the SDK's `rs/live-runner-session-payments` branch rather than the `ja/live-runner` the other examples use, because a long-lived streaming session needs `session.start_payments()`. See the note in `pyproject.toml`. + +| File | Role | +| ---- | ---- | +| `flux_klein.py` | FLUX Klein driver (load / process / update prompt, seed, blend) + the feedback loop | +| `runner.py` | the app: warm-up load, `/health` `/status` `/stats`, trickle `in`/`out`, latest-frame inference worker | +| `client.py` | publishes a file or webcam (Linux/macOS), reads the output, updates prompt/seed/blend live | +| `view.sh` | turnkey showcase: webcam → FLUX Klein → an mpv/ffplay window | +| `runners.json` | static config the orchestrator loads (`-liveRunnerConfig`) | +| `FEEDBACK.md` | developer-experience notes collected while building this example | + +## How it's wired + +**Static registration.** The orchestrator reads `runners.json` via `-liveRunnerConfig`, learns the app at `http://app:8720`, and health-polls `/health` — no registrar, no heartbeat, no SDK registration in the app. The SDK is used only for trickle plumbing once a session starts. + +**Self-contained inference.** `flux_klein.py` wraps `diffusers` (`Flux2KleinPipeline`) directly. It does **not** run [Daydream Scope](https://github.com/daydreamlive/scope) or any plugin runtime. The reference implementation, [`hthillman/scope-flux-klein`](https://github.com/hthillman/scope-flux-klein), is a *Scope plugin*; this example borrows its recipe — the Krea-style feedback loop and the truncated-schedule `refine_frame` — and reimplements it standalone. + +**Warm-up gates routing.** On boot `runner.py` loads the FLUX pipeline, downloading ~15 GB of weights on first run into the mounted `./models`. Until it's ready `/health` returns 503, so the orchestrator routes nothing; once ready it returns 200 and sessions flow. A failed load exits non-zero. Watch `curl localhost:8720/status` → `{state: building|ready|error}`. + +**Per frame — latest-frame worker.** On `/stream` the app opens trickle `in`/`out` channels with `create_trickle_channels`. Decoded frames are *not* queued: the frame callback stashes only the **newest** frame, and a worker processes whatever is newest whenever the GPU is free. Frames arriving while Klein is busy are skipped — essential for a model slower than the input frame rate, otherwise the backlog and end-to-end latency grow without bound. `/update` changes prompt/seed/blend mid-stream. The client drives it with `reserve_session` → `MediaPublish`/`MediaOutput`, routed by **app id**, the same flow as [`echo`](../echo). + +**Session stats.** `GET /stats` reports live throughput while a session runs: trickle input (`video_frames_decoded`, `input_fps`, decoder queue and subscriber health), output (per-track `frames_in`, `encode_fps`, publisher throughput) and inference (`frames_processed`, `avg/last_inference_s`, `inference_fps`, `frames_skipped`), plus the active `seed`/`input_blend`. + +## The feedback loop + +FLUX Klein is a 4-step model — full inference on every frame would be far too slow for video. The [Krea](https://www.krea.ai/)-style feedback loop avoids that: + +- **First frame** → full inference (`image_to_image`), cached as the previous output. +- **Every frame after** → `refine_frame`: the incoming frame is blended with the previous output (so the model refines something that already carries the prompt, not the raw camera feed), then **partially denoised**, running only `feedback_strength × steps` steps. `0.5` runs ~2 of 4 steps. + +Because `Flux2KleinPipeline` has no `strength` parameter (it conditions on the image via joint attention), `refine_frame` hand-rolls a truncated flow-matching schedule against diffusers internals. Prompt embeddings, VAE batch-norm stats and latent position IDs are cached across frames to keep per-frame cost down. + +### Tuning: shimmer vs. drift + +Two loop parameters shape the look, both adjustable live: + +- **`seed`** — `-1` (default) injects fresh random noise every frame: organic but shimmery. A **fixed seed** reuses the same noise each frame: much steadier, but the constant bias compounds, so the image drifts harder toward the prompt and seed-specific artifacts can "burn in". +- **`input_blend`** — the camera's weight (0..1, default 0.5) in the prev-output/camera blend. Higher anchors the stream to reality and washes burned-in artifacts out; lower lets the prompt take over. + +A good starting combination is **`--seed 50 --input-blend 0.85`** (steady *and* anchored). Raise `input_blend` if fixed-seed artifacts appear; lower it if the prompt is too weak. Seeds are not interchangeable — each one settles into its own attractor, so if a fixed seed keeps drawing the same stray shapes, try another value before reaching for the blend. + +## Run offchain (free) + +```sh +docker compose up -d --build # first boot downloads ~15 GB of weights +curl -s localhost:8720/status # building -> ready +curl -sk https://localhost:8935/discovery | jq '.[].runners[].app' # confirm livepeer-example/flux-klein + +uv sync +./view.sh "a cyberpunk portrait, neon lighting" # webcam -> FLUX Klein -> live window + +docker compose down +``` + +`view.sh` is just `client.py` piped to a player (prefers `mpv`, falls back to `ffplay`), auto-detecting macOS (AVFoundation) vs Linux (v4l2). Point it at a remote orchestrator with `DISCOVERY=https://host:8935/discovery ./view.sh "..."`. + +Full control via the client: + +```sh +uv run client.py --webcam-macos 0 --fps 10 \ + --prompt "a watercolor painting" \ + --seed 50 --input-blend 0.85 \ + --reprompt "10=Change to a 90-year-old man" \ + --reprompt "20=seed:7" \ + --reprompt "30=blend:0.6" \ + --output - | mpv --profile=low-latency --no-cache - +``` + +- **Webcam by platform:** Linux uses `--webcam /dev/videoN` (v4l2); macOS uses `--webcam-macos INDEX` (AVFoundation — list cameras with `ffmpeg -f avfoundation -list_devices true -i ""`, and grant your terminal camera access under System Settings → Privacy & Security → Camera, or capture times out). On macOS capture always runs at `--capture-fps` (default 30, since AVFoundation rejects unsupported rates) and `--fps` is honoured by client-side decimation. +- **Live updates:** `--reprompt "SECONDS=..."` (repeatable) fires `/update` mid-stream — plain text changes the prompt, `seed:N` the noise seed (`seed:-1` = random), `blend:X` the camera weight. Seed/blend changes never touch the prompt. +- **Publish less, not slower:** `--fps 10` is plenty; Klein only processes ~19 fps anyway, and less input means less decode work everywhere. +- **File instead of webcam:** `uv run client.py ~/samples/clip.mp4 --output out.ts`. +- **No window, and the runner logs `Trickle ... channel does not exist`?** The player never opened, so it never drained the pipe; the client's stdout then blocks and the output channel times out. That's a dead viewer, not a pipeline bug. `ffplay` (SDL/X11) can fail under Wayland — `mpv` is more reliable there. If in doubt, record then play: `uv run client.py --webcam-macos 0 --max-frames 300 --output live.ts`. +- **`mpegts: Packet corrupt` warnings** appear at every trickle segment joint (independent TS chunks concatenated). Cosmetic; silence with `--msg-level=ffmpeg/demuxer=error`. + +## Run on-chain (paid) + +Layer `docker-compose.onchain.yml` to add the shared remote signer and re-point the orchestrator on-chain. The price is already in `runners.json`, so unlike the dynamic examples there's no `--price` on the app. This needs an Ethereum RPC, a funded signer wallet and an orchestrator wallet — see [On-chain (paid) setup](../README.md#on-chain-paid-setup) in the repo README. + +```sh +cp .env.example .env # fill in RPC, network, keystore paths, accounts, pricing +docker compose -f docker-compose.yml -f docker-compose.onchain.yml up -d --build +uv run client.py --webcam /dev/video0 \ + --discovery https://localhost:8935/discovery \ + --signer http://localhost:7936 --output out.ts +docker compose -f docker-compose.yml -f docker-compose.onchain.yml down +``` + +Keep `PIXELS_PER_UNIT` small — if it's too large the per-unit price floors to 0 wei and calls are effectively free despite being on-chain. + +## Parameters + +| Runner env | Default | Meaning | +| ---------- | ------- | ------- | +| `FLUX_MODEL` | `black-forest-labs/FLUX.2-klein-4B` | HF model id | +| `FLUX_WIDTH` / `FLUX_HEIGHT` | `384` × `384` | inference resolution (snapped to /16) | +| `FLUX_STEPS` | `4` | full-inference steps; refine runs a fraction | +| `FLUX_FEEDBACK` | `0.5` | fraction of steps per refine frame; `int(FLUX_STEPS × FLUX_FEEDBACK)` must be ≥ 2, or the refine starts at ~zero noise and stops transforming | +| `FLUX_GUIDANCE` | `1.0` | prompt adherence (Klein's recommended value) | +| `FLUX_MAX_TEXT_TOKENS` | `128` | prompt token budget. The DiT attends over text+image tokens jointly, so padding to the pipeline default of 512 makes it carry ~512 dead tokens next to only 576 image tokens. Raise it only if long prompts get truncated — the runner logs the embed shape on every prompt change | +| `FLUX_SEED` | `-1` | noise seed; `-1` random per frame, fixed = steady | +| `FLUX_INPUT_BLEND` | `0.5` | camera weight in the feedback blend | +| `FLUX_CPU_OFFLOAD` | off | set `1` for < 16 GB VRAM (slower per frame) | +| `FLUX_COMPILE` | off (unset) | `torch.compile` mode for the transformer + VAE, passed through verbatim: `default` (fast compile), `max-autotune` (longest warm-up, fastest kernels), `reduce-overhead` (CUDA graphs). Adds minutes to warm-up | +| `FLUX_BATCH` | `1` | frames refined per GPU pass; `2` trades latency and feedback fidelity for throughput (see below) | + +Client flags: `--prompt`, `--seed`, `--input-blend` set the session at `/stream` (omit to keep the runner defaults); `--reprompt` changes any of them live. Plus `--fps` (publish rate), `--capture-fps` (macOS capture), `--video-size`, `--max-frames`, `--output`, `--discovery`, `--signer`. + +### Throughput + +Measured on an NVIDIA RTX 6000 Pro at 384×384, 4 steps, feedback 0.5: + +| Config | fps | +| ------ | --- | +| Defaults | **~19 fps** | +| `FLUX_BATCH=2 FLUX_COMPILE=max-autotune` | **~22.7 fps** (best measured) | + +### `FLUX_BATCH` — frame batching (opt-in) + +At 384×384 with 2 refine steps the GPU is not saturated at batch 1. `FLUX_BATCH=N` makes the runner collect the newest N frames and refine them **in one pass** from the same previous output. Because the pass is bandwidth-bound rather than compute-bound at these shapes, `N=2` costs only ~1.25–1.35× the wall time of `N=1` for 2× the frames. + +Two costs come with it: + +- **Latency.** You can't refine a frame that hasn't arrived yet, so the batch adds the collection wait — about one frame interval (~34 ms at a 29 fps input) per additional frame. At `N=2` that's one extra frame of delay. +- **Stride-N feedback.** Every frame in a batch descends from the *same* parent output, so the feedback loop advances once per batch instead of once per frame. At `N=2` this is subtle (a faint pairing in the temporal texture); at `N=4` and above it is clearly visible as stepping. + +`N=2` is the only value worth using. `FLUX_BATCH=1` (the default) is byte-identical to the non-batched path. Batching is orthogonal to `FLUX_COMPILE`, but a batched run compiles **two** shapes (N, plus 1 for a partial batch at a stream tail); warm-up compiles both before `/health` reports ready, so neither stalls a live session. `/stats` reports the configured `batch` plus `inference.last_batch_size`, which drops below `FLUX_BATCH` when the input can't fill the buffer. + +## Notes + +- **One GPU, one session.** The app holds a single active session (`capacity: 1` in `runners.json`). The pipeline loads at boot, so the first session starts immediately once `/status` is `ready`. +- **Fixed model/resolution per container.** Model, resolution, steps and feedback are load-time env vars; prompt, seed and blend are per-session and live-updatable. +- **Framerate expectations.** FLUX Klein is not a TensorRT-class realtime model — expect ~19 fps at the defaults, not a locked 30. The latest-frame worker plus segment skipping keeps latency bounded at roughly one inference interval. Lower resolution trades quality for fps; check `GET /stats` for live numbers. +- **diffusers version.** `Flux2KleinPipeline` and the flux2 internals that `refine_frame` reaches into ship from **diffusers git, not a stable release** — the Dockerfile installs diffusers from GitHub. If those private APIs change, `refine_frame` is where it would break; the full-inference paths don't touch internals. diff --git a/flux-klein/client.py b/flux-klein/client.py new file mode 100644 index 0000000..1134e77 --- /dev/null +++ b/flux-klein/client.py @@ -0,0 +1,367 @@ +#!/usr/bin/env python3 +# Client for the self-hosted realtime FLUX Klein app (new runner). +# +# Routed by app id like the `echo` example, so it uses the same flow +# (reserve_session + /stream), NOT the lv2v capability flow. Offchain/free +# against your local orchestrator; pass --signer for the on-chain (paid) path. +# +# Publishes a video file or a Linux v4l2 webcam into the trickle `in` channel, +# reads the FLUX-transformed output from `out`, and can re-prompt live via /update. +from __future__ import annotations + +import argparse +import asyncio +import sys +import time +from contextlib import nullcontext, suppress +from pathlib import Path + +import av + +from livepeer_gateway.errors import LivepeerGatewayError +from livepeer_gateway.http import post_json +from livepeer_gateway.media_output import MediaOutput +from livepeer_gateway.media_publish import MediaPublish, MediaPublishConfig, VideoOutputConfig +from livepeer_gateway.selection import reserve_session + +DEFAULT_DISCOVERY = "https://localhost:8935/discovery" +APP_ID = "livepeer-example/flux-klein" +DEFAULT_OUTPUT = "flux-klein-out.ts" +DEFAULT_PROMPT = "a psychedelic landscape, vivid colors, intricate details" +DEFAULT_VIDEO_SIZE = "640x480" + + +def _log(*args: object) -> None: + print(*args, file=sys.stderr) + + +def _lowlatency_publish_config(fps: float = 30.0) -> MediaPublishConfig: + # Trickle segments rotate on keyframes; the 2s default GOP makes each segment + # ~2s of video, which dominates end-to-end latency. Short GOP + segment floor + # flush segments ~4x/sec for near-realtime (at the cost of more keyframes). + return MediaPublishConfig( + tracks=[VideoOutputConfig(fps=fps, keyframe_interval_s=0.25)], + min_segment_wallclock_s=0.25, + ) + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Run video through the self-hosted FLUX Klein app.") + p.add_argument("input", nargs="?", default=None, help="Local mp4/mov. Omit and pass --webcam for the camera.") + p.add_argument("--webcam", nargs="?", const="/dev/video0", default=None, metavar="DEVICE", + help="Capture from a Linux v4l2 webcam (default /dev/video0).") + p.add_argument("--webcam-macos", nargs="?", const="0", default=None, metavar="INDEX", + help="Capture from a macOS AVFoundation camera by index (default 0). " + "List devices with: ffmpeg -f avfoundation -list_devices true -i \"\"") + p.add_argument("--fps", type=float, default=30.0, + help="Publish rate. Frames are decimated client-side to hit it (macOS capture " + "runs at --capture-fps; files keep their own pacing). Setting this near the " + "runner's inference fps avoids sending frames it would only drop.") + p.add_argument("--capture-fps", type=float, default=30.0, + help="macOS capture framerate — must be a mode the camera supports (almost always 30). " + "Don't confuse with --fps: AVFoundation rejects unsupported rates outright.") + p.add_argument("--video-size", default=DEFAULT_VIDEO_SIZE, help="Webcam capture size (default 640x480, a near-universal native size; the app resizes to its engine size). On macOS this is only applied if you override it (newer Macs don't offer 640x480); otherwise the camera's native size is used.") + p.add_argument("--discovery", default=DEFAULT_DISCOVERY) + p.add_argument("--prompt", default=DEFAULT_PROMPT) + p.add_argument("--output", default=DEFAULT_OUTPUT, help="Output file, or '-' for stdout (pipe to ffplay).") + p.add_argument("--max-frames", type=int, default=0) + p.add_argument("--reprompt", action="append", default=[], metavar="SECONDS=PROMPT", + help="Live prompt change, e.g. --reprompt 5=an oil painting. Repeatable. " + "Also accepts seed changes (--reprompt \"15=seed:42\"; seed:-1 = random) " + "and camera-blend changes (--reprompt \"20=blend:0.7\").") + p.add_argument("--seed", type=int, default=None, + help="Noise seed. -1 = fresh random noise every frame (shimmery); any fixed " + "value reuses the same noise each frame (steadier, reproducible). " + "Omit to keep the runner's default.") + p.add_argument("--input-blend", type=float, default=None, + help="Camera weight 0..1 in the feedback blend. Higher = sticks closer to the " + "real feed (less fantasy drift), lower = prompt takes over. " + "Omit to keep the runner's default (0.5).") + p.add_argument("--signer", default="", help="Remote signer URL; omit for the offchain (free) local path.") + p.add_argument("--payment-interval", type=float, default=3.0) + return p.parse_args() + + +def _channel_url(resp: dict[str, object], name: str) -> str: + url = resp.get(name) + if not isinstance(url, str) or not url: + raise LivepeerGatewayError(f"/stream response missing {name!r} url") + return url + + +def _parse_reprompts(values: list[str]) -> list[tuple[float, str]]: + out: list[tuple[float, str]] = [] + for v in values: + at, _, prompt = v.partition("=") + if not prompt: + raise SystemExit(f"--reprompt expects SECONDS=PROMPT, got {v!r}") + out.append((float(at), prompt)) + return sorted(out, key=lambda x: x[0]) + + +async def _reprompt(app_url: str, schedule: list[tuple[float, str]]) -> None: + start = time.monotonic() + for at, prompt in schedule: + delay = at - (time.monotonic() - start) + if delay > 0: + await asyncio.sleep(delay) + # "seed:N" / "blend:X" entries change parameters instead of the prompt. + if prompt.startswith("seed:"): + payload: dict = {"seed": int(prompt.partition(":")[2])} + elif prompt.startswith("blend:"): + payload = {"input_blend": float(prompt.partition(":")[2])} + else: + payload = {"prompt": prompt} + with suppress(Exception): + await post_json(f"{app_url.rstrip('/')}/update", payload) + _log(f"updated at {at:.1f}s: {payload}") + + +async def _publish_file(input_path: Path, publish_url: str, *, fps: float = 0.0, + max_frames: int = 0) -> None: + input_ = av.open(str(input_path)) + try: + if not input_.streams.video: + raise LivepeerGatewayError(f"No video stream in {input_path}") + src_fps = float(input_.streams.video[0].average_rate or 30.0) + # Publishing faster than the runner can infer just fills the latest-frame + # slot with frames it will drop. Decimating to ~the inference rate means + # (almost) every published frame is actually transformed. Playback stays + # paced to the source timestamps either way — this changes which frames are + # sent, not how long the file takes. + stride = max(1, round(src_fps / fps)) if fps > 0 else 1 + track_fps = fps if fps > 0 else src_fps + if stride > 1: + _log(f"source {src_fps:.1f} fps -> publishing every {stride} frames (~{track_fps:.1f} fps)") + publisher = MediaPublish(publish_url, config=_lowlatency_publish_config(track_fps)) + prev_pts_time = prev_wall = None + sent = 0 + try: + for index, frame in enumerate(input_.decode(video=0), start=1): + if (index - 1) % stride: + continue + sent += 1 + if max_frames > 0 and sent > max_frames: + break + cur = float(frame.pts * frame.time_base) if frame.pts is not None and frame.time_base else None + if prev_pts_time is not None and prev_wall is not None and cur is not None: + sleep_s = max(0.0, (cur - prev_pts_time) - (time.monotonic() - prev_wall)) + if sleep_s > 0: + await asyncio.sleep(sleep_s) + if cur is not None: + prev_pts_time, prev_wall = cur, time.monotonic() + await publisher.write_frame(frame) + finally: + await publisher.close() + finally: + input_.close() + + +async def _publish_webcam(device: str, publish_url: str, *, fps: float, video_size: str, max_frames: int = 0) -> None: + input_ = av.open(device, format="v4l2", + container_options={"framerate": str(fps), "video_size": video_size, "input_format": "mjpeg"}) + try: + if not input_.streams.video: + raise LivepeerGatewayError(f"No video stream on {device}") + publisher = MediaPublish(publish_url, config=_lowlatency_publish_config(fps)) + try: + index = 0 + decode_errors = 0 + done = False + # Read packets in a thread: v4l2 demux() blocks (~1/fps per packet) in C, + # which would freeze the asyncio loop and starve MediaPublish's background + # encode/POST task. We also skip the occasional malformed packet v4l2 emits + # (notably the first) that the decoder rejects. + demux = input_.demux(video=0) + while not done: + try: + packet = await asyncio.to_thread(next, demux) + except StopIteration: + break + try: + frames = packet.decode() + except av.FFmpegError: + decode_errors += 1 + if decode_errors > 300: + raise LivepeerGatewayError(f"webcam {device}: too many decode errors, giving up") + continue + for frame in frames: + index += 1 + if max_frames > 0 and index > max_frames: + done = True + break + frame.pts = None # let MediaPublish stamp wall-clock pts + await publisher.write_frame(frame) + finally: + await publisher.close() + finally: + input_.close() + + +async def _publish_webcam_macos(device: str, publish_url: str, *, fps: float, video_size: str, + capture_fps: float = 30.0, max_frames: int = 0) -> None: + # macOS AVFoundation capture. `device` is a video device index (see + # `ffmpeg -f avfoundation -list_devices true -i ""`), e.g. "0" for the + # built-in FaceTime camera. Video only (no audio). + # + # AVFoundation only accepts capture framerates the camera natively supports: + # omitting the option makes ffmpeg request 29.97 (NTSC), which Mac cameras + # reject, and arbitrary values like 10 fail the same way — both with + # "[Errno 5] Input/output error". So capture at a supported rate (30 nearly + # everywhere; override with --capture-fps) and decimate to --fps client-side. + options = {"framerate": f"{capture_fps:g}"} + # Newer Macs' cameras don't offer 640x480; only force a size if the user + # overrode the default, else let AVFoundation pick the native mode. + if video_size and video_size != DEFAULT_VIDEO_SIZE: + options["video_size"] = video_size + try: + input_ = av.open(device, format="avfoundation", container_options=options) + except OSError as exc: + raise LivepeerGatewayError( + f"AVFoundation rejected the capture config for device {device!r} ({exc}). " + "The camera only accepts specific modes: try --capture-fps 30 and omit " + "--video-size; check the device index with " + "`ffmpeg -f avfoundation -list_devices true -i \"\"`" + ) from exc + try: + if not input_.streams.video: + raise LivepeerGatewayError(f"No video stream on AVFoundation device {device}") + publisher = MediaPublish(publish_url, config=_lowlatency_publish_config(fps)) + try: + index = 0 + done = False + eagain = 0 + min_interval = 1.0 / fps if fps > 0 else 0.0 + last_sent = 0.0 + # Same threaded-demux reasoning as the v4l2 path: the capture read + # blocks in C, so offload it to keep MediaPublish's encode/POST task alive. + demux = input_.demux(video=0) + while not done: + try: + packet = await asyncio.to_thread(next, demux) + eagain = 0 + except StopIteration: + break + except BlockingIOError: + # AVFoundation returns EAGAIN until the camera warms up, and + # PyAV kills the demux generator on it — back off and recreate. + eagain += 1 + if eagain > 500: # ~5s of nothing + raise LivepeerGatewayError( + f"AVFoundation device {device!r} produced no frames in ~5s — check " + "camera permission (System Settings > Privacy & Security > Camera) " + "and that the device index is correct") + await asyncio.sleep(0.01) + demux = input_.demux(video=0) + continue + try: + frames = packet.decode() + except av.FFmpegError: + continue + for frame in frames: + now = time.monotonic() + if min_interval and now - last_sent < min_interval: + continue # decimate down to the target --fps + last_sent = now + index += 1 + if max_frames > 0 and index > max_frames: + done = True + break + frame.pts = None # let MediaPublish stamp wall-clock pts + await publisher.write_frame(frame) + finally: + await publisher.close() + finally: + input_.close() + + +async def main() -> None: + args = _parse_args() + use_webcam = args.webcam is not None or args.webcam_macos is not None + if not use_webcam and not args.input: + raise SystemExit("provide an input file, --webcam (Linux), or --webcam-macos") + input_path = None + if not use_webcam: + input_path = Path(args.input).expanduser() + if not input_path.exists(): + raise SystemExit(f"input file does not exist: {input_path}") + reprompts = _parse_reprompts(args.reprompt) + + output_stdout = args.output.strip().lower() in {"-", "stdout"} + output_path = None if output_stdout else Path(args.output).expanduser() + signer_url = args.signer.strip() or None + + session = reprompt_task = None + try: + session = await reserve_session( + discovery_url=args.discovery, app=APP_ID, signer_url=signer_url, + payment_interval=args.payment_interval, + ) + session.start_payments() # on-chain: keep the session funded for the whole stream; no-op offchain + _log("session_id:", session.session_id, "app_url:", session.app_url) + + stream_body: dict = {"prompt": args.prompt} + if args.seed is not None: + stream_body["seed"] = args.seed + if args.input_blend is not None: + stream_body["input_blend"] = args.input_blend + resp = await post_json(f"{session.app_url.rstrip('/')}/stream", stream_body) + in_url, out_url = _channel_url(resp, "in"), _channel_url(resp, "out") + _log("in:", in_url, "out:", out_url) + if reprompts: + reprompt_task = asyncio.create_task(_reprompt(session.app_url, reprompts)) + + with nullcontext(sys.stdout.buffer) if output_stdout else output_path.open("wb") as fh: + viewer_gone = False + + def _write_chunk(chunk: bytes) -> None: + # The viewer (mpv/ffplay) quitting breaks the stdout pipe; stop + # writing quietly instead of crashing the whole client. + nonlocal viewer_gone + if viewer_gone: + return + try: + fh.write(chunk) + if output_stdout: + fh.flush() + except BrokenPipeError: + viewer_gone = True + _log("viewer closed the pipe; dropping further output (ctrl-c to stop)") + + # Small window + skip-to-latest so the viewer stays on the live edge + # instead of buffering a growing backlog (keeps latency from creeping up). + async with MediaOutput(out_url, on_bytes=_write_chunk, max_segments=2): + if args.webcam_macos is not None: + _log(f"streaming macOS camera {args.webcam_macos}; ctrl-c to stop") + await _publish_webcam_macos(args.webcam_macos, in_url, fps=args.fps, + video_size=args.video_size, + capture_fps=args.capture_fps, + max_frames=max(0, args.max_frames)) + elif args.webcam is not None: + _log(f"streaming webcam {args.webcam} ({args.video_size}); ctrl-c to stop") + await _publish_webcam(args.webcam, in_url, fps=args.fps, + video_size=args.video_size, max_frames=max(0, args.max_frames)) + else: + await _publish_file(input_path, in_url, fps=args.fps, + max_frames=max(0, args.max_frames)) + _log("publish complete; waiting for output to drain...") + with suppress(BrokenPipeError): + fh.flush() + except LivepeerGatewayError as exc: + raise SystemExit(f"ERROR: {exc}") from exc + finally: + if reprompt_task is not None: + reprompt_task.cancel() + with suppress(asyncio.CancelledError, Exception): + await reprompt_task + if session is not None: + # aclose() cancels the payment loop (if any) and stops the runner session. + with suppress(Exception): + await session.aclose() + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + _log("interrupted; session closed") diff --git a/flux-klein/docker-compose.onchain.yml b/flux-klein/docker-compose.onchain.yml new file mode 100644 index 0000000..822c8cf --- /dev/null +++ b/flux-klein/docker-compose.onchain.yml @@ -0,0 +1,48 @@ +# On-chain (paid) overlay for flux-klein. Layer it on the offchain base: +# docker compose -f docker-compose.yml -f docker-compose.onchain.yml up -d --build +# +# Adds the shared remote signer and re-points the orchestrator on-chain (see +# ../compose.onchain.yml). The price is already in runners.json (static runner), +# so unlike the dynamic hello-world example there's no --price on the app. +# Requires a local .env (gitignored); copy .env.example and fill it in. Pay +# through the signer: +# uv run client.py --webcam /dev/video0 \ +# --discovery https://localhost:8935/discovery \ +# --signer http://localhost:7936 --output out.ts + +services: + signer: + extends: + file: ../compose.onchain.yml + service: signer + ports: + - "7936:7936" + + orchestrator: + extends: + file: ../compose.onchain.yml + service: orchestrator + # extends copies the on-chain command; compose replaces command lists (can't + # append), so restate it here WITH -liveRunnerConfig, and re-mount runners.json + # next to the keystore (a static runner needs the config in the on-chain path too). + command: + - -orchestrator + - -useLiveRunners + - -serviceAddr=127.0.0.1:8935 + - -httpAddr=0.0.0.0:8935 + - -liveRunnerAddr=https://orchestrator:8935 + - -orchSecret=abcdef + - -network=${NETWORK} + - -ethUrl=${ETH_RPC_URL} + - -ethKeystorePath=/keystore + - -ethAcctAddr=${ORCH_ETH_ACCT} + - -ethPassword=${ORCH_ETH_PASSWORD} + - -ethOrchAddr=${ORCH_ONCHAIN_ADDR} + - -pricePerUnit=0 + - -ticketEV=1000000000000 + - -liveRunnerConfig=/config/runners.json + - -monitor=false + - -v=6 + volumes: + - ${ORCH_KEYSTORE_DIR}:/keystore:ro + - ./runners.json:/config/runners.json:ro diff --git a/flux-klein/docker-compose.yml b/flux-klein/docker-compose.yml new file mode 100644 index 0000000..23b0c1c --- /dev/null +++ b/flux-klein/docker-compose.yml @@ -0,0 +1,45 @@ +# Offchain demo: orchestrator + the self-hosted realtime FLUX Klein app, +# attached as a STATIC runner via runners.json (-liveRunnerConfig), like vllm. +# The orchestrator health-polls the app's /health and routes only once the FLUX +# pipeline is loaded (warm-up-on-start). Requires an NVIDIA GPU (~13 GB VRAM). +# +# docker compose up -d --build # orchestrator + flux-klein app +# uv run client.py --webcam /dev/video0 --output - | ffplay - + +services: + orchestrator: + extends: + file: ../compose.orchestrator.yml + service: orchestrator + # `extends` can't append a single flag to `command`, so the command is + # restated here with -liveRunnerConfig added and runners.json mounted in. + command: + - -orchestrator + - -useLiveRunners + - -serviceAddr=127.0.0.1:8935 + - -httpAddr=0.0.0.0:8935 + - -liveRunnerAddr=https://orchestrator:8935 + - -orchSecret=abcdef + - -network=offchain + - -monitor=false + - -liveRunnerConfig=/config/runners.json + - -v=6 + volumes: + - ./runners.json:/config/runners.json:ro + + app: + build: . + container_name: example_apps_flux_klein + # FLUX Klein needs an NVIDIA GPU; the HF weights cache lives in ./models so + # the ~15 GB download survives restarts. + gpus: all + volumes: + - ./models:/models + depends_on: + orchestrator: + condition: service_healthy + # Static runner: no registration flags — only the orchestrator URL, used for + # create_trickle_channels once a session starts. + command: + - --host=0.0.0.0 + - --orchestrator=https://orchestrator:8935 diff --git a/flux-klein/flux_klein.py b/flux-klein/flux_klein.py new file mode 100644 index 0000000..13fafbf --- /dev/null +++ b/flux-klein/flux_klein.py @@ -0,0 +1,535 @@ +#!/usr/bin/env python3 +# Realtime FLUX.2-klein-4B wrapper with a Krea-style feedback loop. +# +# SELF-CONTAINED: this is our own inference code (diffusers + Flux2KleinPipeline), +# NOT the Daydream Scope platform. The feedback-loop recipe is adapted from the +# reference plugin hthillman/scope-flux-klein — we borrow the *idea*, not the +# runtime. The runner (runner.py) calls process() once per incoming video frame. +# +# The trick (why Klein is usable in realtime): +# frame 0 -> full inference (text_to_image / image_to_image), cached. +# frame 1..N -> refine_frame(): partial-denoise the PREVIOUS output instead +# of generating from scratch. In video mode the previous output +# is blended 50/50 with the incoming frame first, so the model +# refines something that already carries the prompt rather than +# denoising the raw camera feed. strength controls how many of +# the N steps actually run (0.5 -> ~2 of 4 steps). +# +# Flux2KleinPipeline has NO `strength` parameter (it conditions on the image via +# joint attention), so refine_frame() hand-rolls a truncated flow-matching +# schedule against diffusers internals. That ties us to a diffusers build that +# ships Flux2KleinPipeline (>= 0.37.0.dev0 / install from git). +from __future__ import annotations + +import logging +import os +import time + +import numpy as np +import torch +from PIL import Image + +log = logging.getLogger("flux-klein") + +DEFAULT_MODEL = os.environ.get("FLUX_MODEL", "black-forest-labs/FLUX.2-klein-4B") +DEFAULT_PROMPT = "a psychedelic landscape, vivid colors, intricate details" +DEFAULT_WIDTH = 384 +DEFAULT_HEIGHT = 384 +DEFAULT_STEPS = 4 # full-inference steps; refine uses a fraction of these +DEFAULT_GUIDANCE = 1.0 # Klein works best at 1.0 (CFG effectively off) +DEFAULT_FEEDBACK = 0.5 # 0 = regenerate fully each frame; 0.3 ~1 step, 0.5 ~2 of 4 +DEFAULT_SEED = int(os.environ.get("FLUX_SEED", "-1")) # -1 = fresh random noise every +# frame (organic shimmer); any fixed value reuses the SAME noise each frame, which +# visibly steadies the feedback loop (less flicker) and makes output reproducible. +DEFAULT_INPUT_BLEND = float(os.environ.get("FLUX_INPUT_BLEND", "0.5")) # camera weight in +# the prev-output/camera blend (0..1). Higher = anchors to the real feed (less "fantasy" +# drift, weaker prompt); lower = prompt influence compounds harder across frames. +DEFAULT_BATCH = int(os.environ.get("FLUX_BATCH", "1")) # frames refined per GPU pass. +# 1 (default) = today's behaviour exactly. >1 collects N frames and refines them all from +# the SAME previous output in one transformer/VAE pass: at 384x384 the GPU is far from +# saturated at batch 1, so B=2 costs ~1.3x the wall time of B=1 for 2x the frames. Costs +# latency (you must wait for N frames to arrive) and makes feedback stride-N — see README. +DEFAULT_MAX_TEXT_TOKENS = int(os.environ.get("FLUX_MAX_TEXT_TOKENS", "128")) # prompt token +# budget. The DiT attends over text+image tokens jointly, so padding the text to the +# pipeline default of 512 makes the transformer carry ~512 dead tokens next to only 576 +# image tokens (384x384). Short prompts fit in 128 easily; raise it if a long prompt +# gets truncated. The per-prompt log line in _refine shows the actual embed length. + + +def _snap(value: int, multiple: int = 16) -> int: + """FLUX requires spatial dims to be multiples of 16.""" + return max(multiple, (value // multiple) * multiple) + + +# --- boundary conversions (numpy uint8 RGB <-> tensors) --------------------- + +def _numpy_to_thwc( + frame_rgb: np.ndarray, device: torch.device, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + """HxWx3 uint8 RGB -> (1, H, W, 3) float [0, 1] on `device`. + Uploads uint8 (4x smaller than fp32) and converts on the GPU. `dtype` defaults + to fp32 but the feedback path passes the model dtype so the blend never has to + round-trip through fp32.""" + t = torch.from_numpy(np.ascontiguousarray(frame_rgb)).to(device) + return t.to(dtype).div_(255.0).unsqueeze(0) + + +class FluxKleinModel: + """Container-wide FLUX Klein pipeline. Loaded once at warm-up; process() is + called per frame from the trickle loop. Resolution/model are fixed per + container; only the prompt changes at runtime (via update_prompt).""" + + def __init__(self) -> None: + self.pipe = None + self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self._width = DEFAULT_WIDTH + self._height = DEFAULT_HEIGHT + self._steps = DEFAULT_STEPS + self._guidance = DEFAULT_GUIDANCE + self._feedback = DEFAULT_FEEDBACK + self._prompt = DEFAULT_PROMPT + self._seed = DEFAULT_SEED + self._input_blend = DEFAULT_INPUT_BLEND + self._batch = DEFAULT_BATCH + + # Krea feedback loop: last emitted frame, kept as THWC [0,1] on device. + self._prev: torch.Tensor | None = None + + # Per-frame caches (recomputing these every frame would dominate latency). + self._cached_prompt_str: str | None = None + self._cached_prompt_embeds: torch.Tensor | None = None + self._cached_text_ids: torch.Tensor | None = None + self._bn_mean: torch.Tensor | None = None + self._bn_std: torch.Tensor | None = None + self._bn_inv_std: torch.Tensor | None = None + self._cached_latent_shape: tuple | None = None + self._cached_latent_ids: torch.Tensor | None = None + self._cached_timesteps_key: tuple | None = None + self._cached_mu: float | None = None + self._noise_key: tuple | None = None + self._noise: torch.Tensor | None = None + self._dtype: torch.dtype = torch.float16 + # Pinned staging buffer for the per-frame device->host copy (see _thwc_to_numpy). + self._pinned: torch.Tensor | None = None + # One-shot flag so the prompt-embed shape log fires on prompt change, not per frame. + self._log_embed_shapes = False + + # --- lifecycle ---------------------------------------------------------- + def load( + self, + model: str = DEFAULT_MODEL, + prompt: str = DEFAULT_PROMPT, + width: int = DEFAULT_WIDTH, + height: int = DEFAULT_HEIGHT, + steps: int = DEFAULT_STEPS, + guidance: float = DEFAULT_GUIDANCE, + feedback_strength: float = DEFAULT_FEEDBACK, + seed: int = DEFAULT_SEED, + input_blend: float = DEFAULT_INPUT_BLEND, + batch: int = DEFAULT_BATCH, + enable_cpu_offload: bool = False, + compile_mode: "str | None" = None, + ) -> None: + self._width = _snap(width) + self._height = _snap(height) + self._steps = steps + self._guidance = guidance + self._feedback = feedback_strength + self._prompt = prompt + self._seed = seed + self._input_blend = input_blend + self._batch = max(1, int(batch)) + + try: + from diffusers import Flux2KleinPipeline + except ImportError as exc: + raise ImportError( + "Flux2KleinPipeline not found — needs a diffusers build with FLUX.2 " + "support (>= 0.37.0.dev0). Install from git: " + "pip install git+https://github.com/huggingface/diffusers.git" + ) from exc + + # Shapes are static here (fixed resolution, fixed step count), so cudnn's + # autotuner pays for itself on the first frames instead of thrashing; TF32 + # matmuls apply to whatever fp32 work remains (VAE bits, compile guards). + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + + log.info("loading %s (%dx%d, %d steps)", model, self._width, self._height, self._steps) + self.pipe = Flux2KleinPipeline.from_pretrained(model, torch_dtype=torch.float16) + if enable_cpu_offload: + self.pipe.enable_model_cpu_offload() + else: + self.pipe = self.pipe.to(self._device) + # Cache once: after torch.compile this would resolve through + # OptimizedModule.__getattr__ on every frame. + self._dtype = self.pipe.transformer.dtype + + if compile_mode: + # torch.compile the transformer (the per-step cost). Compilation is + # triggered by the dummy passes below so it lands during warm-up + # (health still 503) instead of stalling the first live session. + # Two passes: full-inference and refine call the transformer with + # different shapes, and each shape compiles separately. + # compile_mode is passed to torch.compile verbatim: "default" (fast + # compile), "max-autotune" (much longer warm-up, fastest kernels), or + # "reduce-overhead" (CUDA graphs — shapes here are static so it applies). + log.info("compiling transformer + VAE (mode=%s; adds minutes to warm-up)", compile_mode) + self.pipe.transformer = torch.compile(self.pipe.transformer, mode=compile_mode) + # The VAE runs a full encode AND decode every refine frame — with only + # ~2 transformer steps at feedback 0.5, it's a large share of frame time + # and was the biggest remaining uncompiled block. Compile the encoder / + # decoder submodules, not pipe.vae: encode()/decode() are wrappers that + # do tiling/slicing and would graph-break. + self.pipe.vae.to(memory_format=torch.channels_last) + self.pipe.vae.encoder = torch.compile(self.pipe.vae.encoder, mode=compile_mode) + self.pipe.vae.decoder = torch.compile(self.pipe.vae.decoder, mode=compile_mode) + dummy = np.zeros((self._height, self._width, 3), dtype=np.uint8) + self.process(dummy) # compiles the full-inference path (first frame) + self.process(dummy) # compiles the refine path (every frame after) + self.process(dummy) # second refine pass: settles CUDA-graph capture + if self._batch > 1: + # A batched refine is a *different* shape to the compiler, so it would + # otherwise trigger a fresh (multi-minute under max-autotune) compile on + # the first live batch — i.e. after /health already said ready. Compile + # it here. Note both shapes stay live: a stream tail can hand the worker + # a partial batch of 1, which uses the B=1 graph compiled above. + self.process_batch([dummy] * self._batch) + self.process_batch([dummy] * self._batch) # settles CUDA-graph capture + self._prev = None # don't leak warm-up state into the first session + log.info("compile warm-up done") + + log.info("FLUX Klein ready") + + def reset(self) -> None: + """Drop the feedback state. MUST be called between sessions — otherwise the + next session's first frame refines the previous client's last output instead + of doing a fresh full inference.""" + self._prev = None + + def update_prompt(self, prompt: str) -> None: + self._prompt = prompt # prompt-embed cache invalidates itself on change + + def update_seed(self, seed: int) -> None: + self._seed = int(seed) + self._noise_key = None # force regeneration for the new seed + + @property + def seed(self) -> int: + return self._seed + + def update_input_blend(self, input_blend: float) -> None: + self._input_blend = min(1.0, max(0.0, float(input_blend))) + + @property + def input_blend(self) -> float: + return self._input_blend + + @property + def batch(self) -> int: + return self._batch + + def _make_generator(self) -> "torch.Generator | None": + """Fresh generator with the SAME seed every call: a fixed seed therefore + reuses identical noise each frame (steady, reproducible), while -1 means + fully random noise per frame (organic shimmer).""" + if self._seed < 0: + return None + gen = torch.Generator(device="cpu") + gen.manual_seed(self._seed) + return gen + + # --- per-frame entry point --------------------------------------------- + def process(self, frame_rgb: np.ndarray | None) -> np.ndarray: + """Run one feedback-loop step. frame_rgb is HxWx3 uint8 RGB (the incoming + video frame) or None for pure text mode. Returns HxWx3 uint8 RGB.""" + if self.pipe is None: + raise RuntimeError("pipeline not loaded — call load() first") + + prompt = self._prompt if self._prompt.strip() else "high quality image" + + if frame_rgb is not None: + if self._prev is None or self._feedback <= 0: + out = self._image_to_image(prompt, frame_rgb) # first frame: full + else: + blended = self._blend_prev_with_frame(frame_rgb) # steer toward the feed + out = self._refine(prompt, blended, self._feedback) + else: + if self._prev is None or self._feedback <= 0: + out = self._text_to_image(prompt) # first frame: full + else: + out = self._refine(prompt, self._prev, self._feedback) + + self._prev = out.detach() + return self._thwc_to_numpy(out) + + def process_batch(self, frames_rgb: "list[np.ndarray]") -> "list[np.ndarray]": + """Refine N frames in ONE GPU pass and return N HxWx3 uint8 RGB frames, in + input order. Every frame in the batch is blended against the SAME self._prev, + so feedback advances once per batch (stride-N) rather than once per frame — + that's the behavioural price of the throughput win. The newest output seeds + the next batch. N == 1 is exactly process().""" + if not frames_rgb: + return [] + if self.pipe is None: + raise RuntimeError("pipeline not loaded — call load() first") + if len(frames_rgb) == 1: + return [self.process(frames_rgb[0])] + + prompt = self._prompt if self._prompt.strip() else "high quality image" + + # Cold start: the full-inference path takes a PIL image and has no batch form, + # so the first frame runs alone at B=1 to seed _prev; the rest batch normally. + head: list[np.ndarray] = [] + if self._prev is None or self._feedback <= 0: + head.append(self.process(frames_rgb[0])) + frames_rgb = frames_rgb[1:] + if not frames_rgb: + return head + if self._feedback <= 0: + # Feedback disabled: every frame is a full inference, one at a time. + return head + [self.process(f) for f in frames_rgb] + + blended = torch.cat([self._blend_prev_with_frame(f) for f in frames_rgb], dim=0) + out = self._refine(prompt, blended, self._feedback) + self._prev = out[-1:].detach() # newest output seeds the next batch + return head + [self._thwc_to_numpy(out[i : i + 1]) for i in range(out.shape[0])] + + def _thwc_to_numpy(self, tensor: torch.Tensor) -> np.ndarray: + """(1, H, W, 3) float [0, 1] (any device) -> HxWx3 uint8 RGB. Takes ONE frame: + batched callers slice (out[i:i+1]) and call this per element, which keeps the + pinned buffer below a single fixed shape. + Converts to uint8 on the GPU so only 1/4 of the bytes cross PCIe, then copies + through a reusable pinned host buffer — a plain .cpu() on a pageable + destination makes the driver stage through its own bounce buffer and blocks + the whole copy on the critical path.""" + t = tensor[0].clamp(0, 1).mul(255.0).round().to(torch.uint8) + if t.device.type != "cuda": # CPU-only dev boxes: nothing to pin + return t.numpy() + if self._pinned is None or self._pinned.shape != t.shape: + self._pinned = torch.empty(t.shape, dtype=torch.uint8, device="cpu", pin_memory=True) + self._pinned.copy_(t, non_blocking=True) + torch.cuda.current_stream().synchronize() + # .copy() is load-bearing: the caller hands this array to PyAV + # (VideoFrame.from_ndarray), which keeps a reference — returning a view of + # the reusable buffer would let the next frame overwrite a frame still in + # flight. The copy is a ~440 KB memcpy at 384x384, far cheaper than the + # pageable transfer it replaces. + return self._pinned.numpy().copy() + + def _blend_prev_with_frame(self, frame_rgb: np.ndarray) -> torch.Tensor: + """Blend the previous output with the new frame at engine resolution + (THWC [0,1]). Refining this (rather than the raw feed) keeps prompt + influence alive; input_blend is the camera's weight — raise it to anchor + the stream to reality, lower it to let the prompt take over.""" + # Match _prev's dtype (model dtype, fp16): the whole feedback path stays in + # fp16 so nothing pays a full-frame fp32 pass. interpolate() is fine on fp16 CUDA. + frame = _numpy_to_thwc(frame_rgb, self._prev.device, self._prev.dtype) + # Resize the incoming frame to engine resolution so we blend against prev. + if frame.shape[1] != self._height or frame.shape[2] != self._width: + frame = torch.nn.functional.interpolate( + frame.permute(0, 3, 1, 2), + size=(self._height, self._width), + mode="bilinear", align_corners=False, + ).permute(0, 2, 3, 1) + w = self._input_blend + return (1.0 - w) * self._prev + w * frame + + # --- full inference (first frame / feedback disabled) ------------------- + @torch.no_grad() + def _text_to_image(self, prompt: str) -> torch.Tensor: + result = self.pipe( + prompt=prompt, + width=self._width, + height=self._height, + guidance_scale=self._guidance, + num_inference_steps=self._steps, + generator=self._make_generator(), + output_type="pt", # tensor on-device; skip the PIL round-trip + ) + # images: (B, C, H, W) float in [0, 1] -> THWC on device. Left in model dtype: + # _prev is fp16 throughout so the feedback path never upcasts. + return result.images[0].permute(1, 2, 0).unsqueeze(0) + + @torch.no_grad() + def _image_to_image(self, prompt: str, frame_rgb: np.ndarray) -> torch.Tensor: + image = Image.fromarray(frame_rgb, mode="RGB").resize( + (self._width, self._height), Image.LANCZOS + ) + result = self.pipe( + prompt=prompt, + image=image, + width=self._width, + height=self._height, + guidance_scale=self._guidance, + num_inference_steps=self._steps, + generator=self._make_generator(), + output_type="pt", # tensor on-device; skip the PIL round-trip + ) + return result.images[0].permute(1, 2, 0).unsqueeze(0) + + # --- feedback refine (the Krea trick) ----------------------------------- + def _get_bn_stats(self, device, dtype): + """VAE batch-norm stats — constant after load. Returns (mean, std, inv_std): + normalize multiplies by inv_std, denormalize multiplies by std, so neither + hot path divides.""" + if self._bn_mean is None: + self._bn_mean = self.pipe.vae.bn.running_mean.view(1, -1, 1, 1).to(device, dtype) + self._bn_std = torch.sqrt( + self.pipe.vae.bn.running_var.view(1, -1, 1, 1).to(device, dtype) + + self.pipe.vae.config.batch_norm_eps + ) + self._bn_inv_std = 1.0 / self._bn_std + return self._bn_mean, self._bn_std, self._bn_inv_std + + def _get_latent_ids(self, latent_shape: tuple, device): + # Keyed on the spatial dims only, so a future batch dim doesn't thrash it. + key = tuple(latent_shape[1:]) + if self._cached_latent_shape != key: + dummy = torch.empty(latent_shape) + self._cached_latent_ids = self.pipe._prepare_latent_ids(dummy).to(device) + self._cached_latent_shape = key + return self._cached_latent_ids + + def _get_noise(self, shape: tuple, device, dtype) -> torch.Tensor: + """Sampling noise for the refine step. + + Fixed seed: the same seed re-generates bit-identical noise every frame (that's + what makes output steady), so generate ONCE and reuse — the previous code paid + CPU RNG plus a full host->device copy per frame to recompute a constant. + seed -1: fresh noise per frame, generated directly on the device. + + Batched (shape[0] > 1): the fixed-seed noise is generated at B=1 and expanded, so + every element of the batch gets the SAME noise. Distinct per-element noise would + silently break the "fixed seed = steady output" property — the batch elements + would shimmer against each other even though the seed never changed. For seed -1 + per-element random noise is the correct behaviour (that path is the shimmer).""" + from diffusers.utils.torch_utils import randn_tensor + if self._seed < 0: + return randn_tensor(shape, device=device, dtype=dtype) + base_shape = (1,) + tuple(shape[1:]) + key = (self._seed, base_shape) + if self._noise_key != key: + gen = torch.Generator(device=device) + gen.manual_seed(self._seed) + self._noise = randn_tensor(base_shape, generator=gen, device=device, dtype=dtype) + self._noise_key = key + # expand() is a free view; it's only ever read (the interpolation below is + # non-inplace), so the broadcast stride is safe. + return self._noise.expand(shape) if shape[0] != 1 else self._noise + + @torch.no_grad() + def _refine(self, prompt: str, prev_thwc: torch.Tensor, strength: float) -> torch.Tensor: + """Partial-denoise prev_thwc (THWC [0,1]) toward `prompt`. Adapted from + scope-flux-klein's refine_frame: VAE-encode -> add noise at a partial + timestep -> denoise only the remaining steps -> decode.""" + device = self.pipe._execution_device + dtype = self._dtype + t0 = time.perf_counter() + + # THWC [0,1] -> BCHW [-1,1] for the VAE, fused in model dtype (one pass, no + # fp32 intermediate). Non-inplace on purpose: prev_thwc IS self._prev, which + # the next frame's blend still reads. + image_tensor = prev_thwc.to(device).permute(0, 3, 1, 2).to(dtype).mul(2.0).sub(1.0) + + # 1. Prompt embeddings (re-encode only when the prompt changes). + if prompt != self._cached_prompt_str: + self._cached_prompt_embeds, self._cached_text_ids = self.pipe.encode_prompt( + prompt=prompt, device=device, num_images_per_prompt=1, + max_sequence_length=DEFAULT_MAX_TEXT_TOKENS, + ) + self._cached_prompt_str = prompt + self._log_embed_shapes = True # logged below, where image_seq_len is known + prompt_embeds, text_ids = self._cached_prompt_embeds, self._cached_text_ids + # Batched refine: the embeds are cached at B=1, so broadcast them across the + # batch (a view, free — all elements share one prompt anyway). text_ids/img_ids + # are NOT expanded: in the FLUX family they are unbatched (seq, 3) position ids + # shared by the whole batch — img_ids comes from _get_latent_ids, whose cache key + # is deliberately spatial-only for exactly that reason. + batch_size = prev_thwc.shape[0] + if batch_size > 1: + prompt_embeds = prompt_embeds.expand(batch_size, *prompt_embeds.shape[1:]) + + # 2. Encode -> patchify -> BN normalize -> pack. + image_latents = self.pipe.vae.encode(image_tensor).latent_dist.mode() + image_latents = self.pipe._patchify_latents(image_latents) + bn_mean, bn_std, bn_inv_std = self._get_bn_stats(device, dtype) + image_latents = (image_latents - bn_mean) * bn_inv_std + latent_ids = self._get_latent_ids(image_latents.shape, device) + clean_latents = self.pipe._pack_latents(image_latents) + + # 3. Truncated schedule: run only the last `strength * steps` steps. + # set_timesteps MUST run every frame: besides rebuilding the (tiny) sigma + # array it resets the scheduler's internal _step_index, which otherwise + # keeps advancing across frames and indexes past the end of sigmas. + # Only `mu` is cached — it depends solely on (seq_len, steps), both fixed. + image_seq_len = clean_latents.shape[1] + if self._log_embed_shapes: + # The DiT attends over text+image jointly, so text padding is real work. + # If the embed length here is much larger than the prompt, lower + # FLUX_MAX_TEXT_TOKENS. + log.info("prompt embeds=%s text_ids=%s image_tokens=%d", + tuple(self._cached_prompt_embeds.shape), + tuple(self._cached_text_ids.shape), image_seq_len) + self._log_embed_shapes = False + ts_key = (image_seq_len, self._steps) + if self._cached_timesteps_key != ts_key: + # Imported lazily so load()'s friendly ImportError still fires first. + from diffusers.pipelines.flux2.pipeline_flux2_klein import compute_empirical_mu + self._cached_mu = compute_empirical_mu(image_seq_len, self._steps) + self._cached_timesteps_key = ts_key + self.pipe.scheduler.set_timesteps(self._steps, device=device, mu=self._cached_mu) + all_timesteps = self.pipe.scheduler.timesteps + + init_timestep = min(int(self._steps * strength), self._steps) + t_start = max(self._steps - init_timestep, 0) + timesteps = all_timesteps[t_start:] + # Load-bearing: without it the scheduler falls back to a per-step + # (timesteps == t).nonzero().item() lookup, i.e. a GPU->CPU sync every step. + self.pipe.scheduler.set_begin_index(t_start) + + if len(timesteps) == 0: # strength ~0: nothing to denoise + return prev_thwc + + # 4. Add noise at the starting sigma (flow-matching interpolation). + start_sigma = (timesteps[0].float() / self.pipe.scheduler.config.num_train_timesteps) + start_sigma = start_sigma.view(-1, 1, 1).to(device=device, dtype=dtype) + noise = self._get_noise(clean_latents.shape, device, dtype) + latents = start_sigma * noise + (1.0 - start_sigma) * clean_latents + + # 5. Denoise the remaining steps (transformer called directly; CFG off at guidance=1). + # Iterate on CPU: indexing a CUDA tensor in a Python loop materializes a 0-dim + # GPU tensor (and a launch) per step. The scheduler only uses `t` for its + # sigma bookkeeping — and with set_begin_index above it doesn't even look `t` + # up — so a CPU scalar is fine and costs no sync. + timesteps_cpu = timesteps.to("cpu") + for t in timesteps_cpu: + # t is a CPU scalar now, so move the expanded timestep to the device here. + timestep = t.expand(latents.shape[0]).to(latents.device, latents.dtype) + noise_pred = self.pipe.transformer( + hidden_states=latents.to(dtype), + timestep=timestep / 1000, + guidance=None, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_ids, + return_dict=False, + )[0] + noise_pred = noise_pred[:, :image_seq_len] + latents = self.pipe.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + # 6. Unpack -> BN denormalize -> unpatchify -> VAE decode. + latents = self.pipe._unpack_latents_with_ids(latents, latent_ids) + latents = latents * bn_std + bn_mean # inverse of the normalize above + latents = self.pipe._unpatchify_latents(latents) + decoded = self.pipe.vae.decode(latents.to(dtype), return_dict=False)[0] + + decoded = (decoded.clamp(-1, 1) + 1) / 2 # [-1,1] -> [0,1] + # Stays in model dtype (fp16): _prev, the blend and the VAE input are all fp16, + # so upcasting here only bought three extra full-frame passes per frame. + out = decoded.permute(0, 2, 3, 1) # BCHW -> (B,H,W,C) + + if log.isEnabledFor(logging.DEBUG): # args eval eagerly; don't pay per frame + dt = time.perf_counter() - t0 + log.debug("refine: %d steps in %.3fs (%.1f fps potential)", len(timesteps), dt, 1.0 / dt) + return out diff --git a/flux-klein/pyproject.toml b/flux-klein/pyproject.toml new file mode 100644 index 0000000..3ea726e --- /dev/null +++ b/flux-klein/pyproject.toml @@ -0,0 +1,32 @@ +[project] +name = "livepeer-flux-klein" +version = "0.1.0" +description = "Realtime FLUX.2-klein-4B app (new runner, trickle realtime video) for the Livepeer network." +requires-python = ">=3.12" +# Host-side deps only — enough to run client.py (and to import runner.py) with +# `uv run`. The GPU inference stack that flux_klein.py needs (torch, diffusers, +# transformers, accelerate, pillow, numpy) is deliberately NOT listed here: they +# are multi-GB CUDA wheels pinned to a specific CUDA/driver combo, so `uv sync` +# on a laptop would be enormous and would resolve to the wrong platform build. +# The app only ever runs in the container, so they are installed in the +# Dockerfile instead. +dependencies = [ + "av", # client: decode the input file / webcam, publish frames + "aiohttp", # runner: the /health /status /stats /stream /update app + "livepeer-gateway", +] + +# livepeer-gateway is not on PyPI yet; pull it from a branch. +# +# NOTE: this example pins rs/live-runner-session-payments rather than the +# ja/live-runner branch the other examples use. client.py calls +# `session.start_payments()` and passes `payment_interval=` to `reserve_session` +# to keep a long-lived streaming session funded for its whole duration; both +# only exist on the payments branch (on ja/live-runner `LiveRunnerSession` is a +# plain dataclass with no such method). Everything else this example uses +# (`create_trickle_channels`, `MediaPublish`/`MediaPublishConfig`, +# `VideoOutputConfig`, `MediaOutput`, `reserve_session`, `post_json`) exists on +# both branches, so this pin can move to ja/live-runner once session payments +# land there. +[tool.uv.sources] +livepeer-gateway = { git = "https://github.com/livepeer/livepeer-python-gateway", branch = "rs/live-runner-session-payments" } diff --git a/flux-klein/runner.py b/flux-klein/runner.py new file mode 100644 index 0000000..c829ded --- /dev/null +++ b/flux-klein/runner.py @@ -0,0 +1,451 @@ +#!/usr/bin/env python3 +# Realtime FLUX Klein app for the new Livepeer runner, attached as a STATIC +# runner (like vllm): the orchestrator reads runners.json via +# -liveRunnerConfig and health-polls /health — no SDK registrar or heartbeat in +# the app. The SDK is used only for the trickle channel plumbing once a session +# starts. +# +# Same trickle loop shape as the `echo` example, with the per-frame transform +# swapped for FLUX.2-klein-4B img2img + a Krea-style feedback loop +# (flux_klein.py). SELF- +# CONTAINED: inference is our own flux_klein.py (diffusers only); NOT the +# Daydream Scope platform (scope-flux-klein is only a reference for the recipe). +# +# Operator lifecycle (warm-up-on-start, gated by /health): +# boot -> state "building": load the FLUX pipeline + download weights. +# /health returns 503, so the orchestrator won't route here. +# ready -> /health returns 200; the orchestrator (which knows this runner from +# runners.json) starts routing sessions. +# error -> state "error", log, exit non-zero; health never passes. +from __future__ import annotations + +import argparse +import asyncio +import dataclasses +import json +import logging +import os +import time +from contextlib import suppress +from dataclasses import dataclass, field +from typing import Any + +import av +from aiohttp import web + +from livepeer_gateway.live_runner import create_trickle_channels +from livepeer_gateway.media_output import MediaOutput +from livepeer_gateway.media_publish import MediaPublish, MediaPublishConfig, VideoOutputConfig + +from flux_klein import ( + FluxKleinModel, + DEFAULT_MAX_TEXT_TOKENS, + DEFAULT_MODEL, + DEFAULT_PROMPT, + DEFAULT_WIDTH, + DEFAULT_HEIGHT, + DEFAULT_STEPS, + DEFAULT_GUIDANCE, + DEFAULT_FEEDBACK, +) + +DEFAULT_HOST = "127.0.0.1" +DEFAULT_PORT = 8720 +APP_ID = "livepeer-example/flux-klein" + +# Model/resolution are fixed per container; the client only changes the prompt. +FLUX_MODEL = os.environ.get("FLUX_MODEL", DEFAULT_MODEL) +FLUX_WIDTH = int(os.environ.get("FLUX_WIDTH", str(DEFAULT_WIDTH))) +FLUX_HEIGHT = int(os.environ.get("FLUX_HEIGHT", str(DEFAULT_HEIGHT))) +FLUX_STEPS = int(os.environ.get("FLUX_STEPS", str(DEFAULT_STEPS))) +FLUX_GUIDANCE = float(os.environ.get("FLUX_GUIDANCE", str(DEFAULT_GUIDANCE))) +FLUX_FEEDBACK = float(os.environ.get("FLUX_FEEDBACK", str(DEFAULT_FEEDBACK))) +FLUX_SEED = int(os.environ.get("FLUX_SEED", "-1")) # -1 random; fixed = steadier output +FLUX_INPUT_BLEND = float(os.environ.get("FLUX_INPUT_BLEND", "0.5")) # camera weight 0..1; +# higher = anchors to the real feed (less "fantasy" drift), lower = prompt takes over +FLUX_CPU_OFFLOAD = os.environ.get("FLUX_CPU_OFFLOAD", "").lower() in {"1", "true", "yes"} +FLUX_BATCH = max(1, int(os.environ.get("FLUX_BATCH", "1"))) # frames per GPU pass; 1 = +# unchanged behaviour. 2 buys throughput (the GPU is under-occupied at 384x384) at the +# cost of one extra frame of latency and stride-2 feedback. See README. +# torch.compile mode for the transformer + VAE; unset disables it. Takes a +# torch.compile mode verbatim ("default", "max-autotune", "reduce-overhead"). +# Any mode trades minutes of warm-up for faster frames; "max-autotune" is the +# fastest and what the README's throughput numbers were measured with. +FLUX_COMPILE = os.environ.get("FLUX_COMPILE", "").strip() or None + +log = logging.getLogger("flux-klein-runner") + +STATE = {"state": "building", "model": FLUX_MODEL, "resolution": f"{FLUX_WIDTH}x{FLUX_HEIGHT}", "error": None} +MODEL = FluxKleinModel() # single, container-wide pipeline (loaded once at warm-up) +_ORCHESTRATOR_URL = "http://localhost:8935" +session: "StreamSession | None" = None +# Guards the whole /stream handler: the session check and the assignment are +# separated by awaits, so two concurrent requests could otherwise both create a +# session and orphan the first one's worker/publisher. +_session_lock = asyncio.Lock() + + +@dataclass +class StreamSession: + session_id: str + in_url: str + out_url: str + output: MediaOutput + publisher: MediaPublish + prompt: str + worker: "asyncio.Task | None" = None + inference: dict[str, Any] = field(default_factory=dict) + + def to_json(self) -> dict[str, Any]: + return {"session": self.session_id, "in": self.in_url, "out": self.out_url, "prompt": self.prompt} + + +def _session_id(request: web.Request) -> str: + sid = request.headers.get("Livepeer-Session-Id", "").strip() + if not sid: + raise web.HTTPBadRequest(text="missing Livepeer-Session-Id header") + return sid + + +# --- warm-up: load the pipeline, then flip /health to ready ----------------- +async def _warmup(app: web.Application) -> None: + try: + log.info("warming up: model=%s %dx%d steps=%d feedback=%.2f batch=%d compile=%s " + "(may download weights)", + FLUX_MODEL, FLUX_WIDTH, FLUX_HEIGHT, FLUX_STEPS, FLUX_FEEDBACK, + FLUX_BATCH, FLUX_COMPILE) + # Blocking + GPU-bound; keep it off the event loop so /status stays responsive. + await asyncio.to_thread( + MODEL.load, + model=FLUX_MODEL, prompt=DEFAULT_PROMPT, + width=FLUX_WIDTH, height=FLUX_HEIGHT, steps=FLUX_STEPS, + guidance=FLUX_GUIDANCE, feedback_strength=FLUX_FEEDBACK, + seed=FLUX_SEED, input_blend=FLUX_INPUT_BLEND, batch=FLUX_BATCH, + enable_cpu_offload=FLUX_CPU_OFFLOAD, + compile_mode=FLUX_COMPILE, + ) + except Exception as exc: + STATE.update(state="error", error=f"{type(exc).__name__}: {exc}") + log.error("model load FAILED: %s", exc, exc_info=True) + os._exit(1) + + STATE["state"] = "ready" + log.info("FLUX Klein ready; /health now returns 200 — orchestrator may route sessions") + + +# --- status / health (static-runner readiness signal) ----------------------- +async def _handle_status(request: web.Request) -> web.Response: + return web.json_response(STATE) + + +async def _handle_health(request: web.Request) -> web.Response: + if STATE["state"] != "ready": + raise web.HTTPServiceUnavailable(text=STATE["state"]) + return web.Response(text="ok") + + +# --- session --------------------------------------------------------------- +async def _close_session() -> None: + global session + if session is None: + return + current, session = session, None + if current.worker is not None: + current.worker.cancel() + try: + await current.worker + except asyncio.CancelledError: + pass + except Exception as exc: # don't discard the reason the worker died + log.error("worker task ended with error: %s", exc, exc_info=True) + with suppress(Exception): + await current.publisher.close() + with suppress(Exception): + await current.output.close() + # Drop the feedback state so the next session starts from a fresh full + # inference instead of continuing this client's imagery. + MODEL.reset() + log.info("closed session %s", current.session_id) + + +async def _handle_stream(request: web.Request) -> web.Response: + # Serialized: the "is there a session?" check and the assignment below are + # separated by several awaits. + async with _session_lock: + return await _start_session(request) + + +async def _start_session(request: web.Request) -> web.Response: + global session + if STATE["state"] != "ready": + raise web.HTTPServiceUnavailable(text=f"runner not ready ({STATE['state']})") + session_id = _session_id(request) + if session is not None: + if session.session_id != session_id: + raise web.HTTPConflict(text="runner already has an active session") + return web.json_response(session.to_json()) + + MODEL.reset() # defensive: never inherit a previous session's feedback state + + channels = await create_trickle_channels( + session_id, + [{"name": "in", "mime_type": "video/mp2t"}, {"name": "out", "mime_type": "video/mp2t"}], + orchestrator_url=_ORCHESTRATOR_URL, + runner_id=request.headers.get("Livepeer-Runner-Route", "").strip(), + session_token=request.headers.get("Livepeer-Session-Token", "").strip(), + ) + by_name = {c["name"]: c for c in channels} + if "in" not in by_name or "out" not in by_name: + raise web.HTTPInternalServerError(text="orchestrator did not return in/out channels") + + body = json.loads(await request.read() or "{}") + prompt = str(body.get("prompt", DEFAULT_PROMPT)) + MODEL.update_prompt(prompt) # model/res fixed per container; this is one assignment + if "seed" in body: + MODEL.update_seed(int(body["seed"])) # -1 random; fixed = steadier output + if "input_blend" in body: + MODEL.update_input_blend(float(body["input_blend"])) + + # Short GOP/segment so output trickle segments flush frequently -> lower latency. + # (FLUX Klein is a few fps, not 30, so segments carry fewer frames each.) + publisher = MediaPublish( + by_name["out"].get("internal_url") or by_name["out"]["url"], + config=MediaPublishConfig( + tracks=[VideoOutputConfig(fps=30.0, keyframe_interval_s=0.25)], + min_segment_wallclock_s=0.25, + ), + ) + + # FLUX Klein runs at a few fps while input arrives at ~30fps. Decouple decode + # from inference with a latest-frame slot: _on_frame only stashes the newest + # frame(s) (never blocks the decode loop), and the worker below processes + # whatever is newest whenever the GPU is free. If inference were awaited + # inside _on_frame, decode would back up behind it and end-to-end latency + # would grow without bound (frames queue instead of being skipped). + # The slot holds the newest FLUX_BATCH frames (just one at the default FLUX_BATCH=1, + # i.e. the original single-frame slot). Anything older is dropped outright. + latest: list[Any] = [] + frame_ready = asyncio.Event() + inference_stats: dict[str, Any] = { + "frames_processed": 0, "total_inference_s": 0.0, "last_inference_s": None, + "last_batch_size": None, "errors": 0, "started_at": time.monotonic(), + } + + async def _on_frame(decoded) -> None: + if decoded.kind != "video": + return + latest.append(decoded.frame) + del latest[:-FLUX_BATCH] # keep only the newest N; older frames are skipped + # Signalled on the FIRST frame, not once the batch is full: if the input stalls + # (or the stream ends) with a partial batch buffered, waiting for N more frames + # would stall output forever. + frame_ready.set() + + def _transform(frames: "list[Any]") -> "list[av.VideoFrame]": + """Decode -> model -> encode for a whole batch. Runs entirely in a worker + thread: the colour conversions are full-frame sws_scale work and would + otherwise stall the event loop (starving MediaPublish's POST task) between + inferences. Each output carries ITS OWN source pts/time_base — reusing one + frame's timestamps for the batch would wreck playback timing.""" + rgbs = [f.to_ndarray(format="rgb24") for f in frames] + outs = [] + for src, out_rgb in zip(frames, MODEL.process_batch(rgbs)): + out = av.VideoFrame.from_ndarray(out_rgb, format="rgb24") + out.pts = src.pts + out.time_base = src.time_base + outs.append(out) + return outs + + async def _process_latest() -> None: + while True: + await frame_ready.wait() + frame_ready.clear() + frames, latest[:] = list(latest), [] # take whatever is buffered (1..N) + if not frames: + continue + t0 = time.perf_counter() + try: + outs = await asyncio.to_thread(_transform, frames) + except asyncio.CancelledError: + raise + except Exception as exc: + # Without this the task dies silently and the client only sees the + # channel vanish — which looks exactly like a dead-viewer bug. + STATE["error"] = f"{type(exc).__name__}: {exc}" + inference_stats["errors"] += 1 + log.error("frame processing failed: %s", exc, exc_info=True) + raise + dt = time.perf_counter() - t0 + inference_stats["frames_processed"] += len(frames) + inference_stats["total_inference_s"] += dt + # Per-FRAME time, so avg/last stay comparable across batch sizes (and keep + # matching /stats' inference_fps) rather than jumping with the batch. + inference_stats["last_inference_s"] = round(dt / len(frames), 3) + inference_stats["last_batch_size"] = len(frames) + for out in outs: # publish in source order + await publisher.write_frame(out) + + # Segment-level skipping too: small window + default LagPolicy.LATEST keeps + # the reader on the live edge if decode itself ever falls behind. + output = MediaOutput( + by_name["in"].get("internal_url") or by_name["in"]["url"], + on_frame=_on_frame, + max_segments=2, + ) + worker = asyncio.create_task(_process_latest()) + session = StreamSession( + session_id=session_id, in_url=by_name["in"]["url"], out_url=by_name["out"]["url"], + output=output, publisher=publisher, prompt=prompt, worker=worker, + inference=inference_stats, + ) + worker.add_done_callback(lambda _t: asyncio.create_task(_close_session())) + for task in output.callback_tasks(): + task.add_done_callback(lambda _t: asyncio.create_task(_close_session())) + log.info("started flux-klein session %s", session_id) + return web.json_response(session.to_json()) + + +def _stats_dict(obj: Any) -> Any: + """Best-effort conversion of SDK stats objects (dataclasses / plain objects) + into JSON-serializable dicts, recursing into nested stats.""" + if obj is None or isinstance(obj, (int, float, str, bool)): + return obj + if dataclasses.is_dataclass(obj) and not isinstance(obj, type): + return {k: _stats_dict(v) for k, v in dataclasses.asdict(obj).items()} + if isinstance(obj, dict): + return {k: _stats_dict(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_stats_dict(v) for v in obj] + if hasattr(obj, "__dict__"): + return {k: _stats_dict(v) for k, v in vars(obj).items() if not k.startswith("_")} + return str(obj) + + +async def _handle_stats(request: web.Request) -> web.Response: + """Session statistics: trickle input/output stats from the SDK plus our own + inference metrics. input fps = video_frames_decoded / elapsed_s; + encode fps = track frames_in / elapsed_s; frames_skipped = decoded - processed + (what the latest-frame slot dropped because the GPU was busy).""" + if session is None: + return web.json_response({"state": STATE["state"], "session": None}) + + in_stats = session.output.get_stats() + if asyncio.iscoroutine(in_stats): + in_stats = await in_stats + out_stats = session.publisher.get_stats() + if asyncio.iscoroutine(out_stats): + out_stats = await out_stats + + input_ = _stats_dict(in_stats) or {} + output_ = _stats_dict(out_stats) or {} + + # Computed rates on top of the raw dumps. + in_elapsed = float(getattr(in_stats, "elapsed_s", 0) or 0) + decoded = int(getattr(in_stats, "video_frames_decoded", 0) or 0) + if in_elapsed > 0: + input_["input_fps"] = round(decoded / in_elapsed, 2) + + out_elapsed = float(getattr(out_stats, "elapsed_s", 0) or 0) + for i, track in enumerate(getattr(out_stats, "track_queue_stats", None) or []): + frames_in = int(getattr(track, "frames_in", 0) or 0) + if out_elapsed > 0 and isinstance(output_.get("track_queue_stats"), list): + output_["track_queue_stats"][i]["encode_fps"] = round(frames_in / out_elapsed, 2) + + inference = {k: v for k, v in session.inference.items() if k != "started_at"} + processed = inference.get("frames_processed", 0) + session_elapsed = time.monotonic() - session.inference.get("started_at", time.monotonic()) + if processed: + inference["avg_inference_s"] = round(inference.pop("total_inference_s", 0.0) / processed, 3) + if session_elapsed > 0: + inference["inference_fps"] = round(processed / session_elapsed, 2) + if decoded: + inference["frames_skipped"] = max(0, decoded - processed) + + return web.json_response( + { + "state": STATE["state"], + "session": session.session_id, + "prompt": session.prompt, + "model": FLUX_MODEL, + "resolution": f"{FLUX_WIDTH}x{FLUX_HEIGHT}", + "steps": FLUX_STEPS, + "feedback_strength": FLUX_FEEDBACK, + # int(steps * feedback) is what actually runs per frame — the number that + # drives fps. Surfaced so tuning doesn't require doing the arithmetic. + "refine_steps": min(int(FLUX_STEPS * FLUX_FEEDBACK), FLUX_STEPS), + "max_text_tokens": DEFAULT_MAX_TEXT_TOKENS, + # Configured batch; inference.last_batch_size shows what the worker actually + # got last pass (smaller when the input can't keep the buffer full). + "batch": FLUX_BATCH, + "seed": MODEL.seed, + "input_blend": MODEL.input_blend, + "compile": FLUX_COMPILE, + "input": input_, + "output": output_, + "inference": inference, + }, + dumps=lambda o: json.dumps(o, default=str), + ) + + +async def _handle_update(request: web.Request) -> web.Response: + if session is None: + raise web.HTTPNotFound(text="session not started") + if session.session_id != _session_id(request): + raise web.HTTPConflict(text="runner has a different active session") + body = json.loads(await request.read() or "{}") + prompt = str(body.get("prompt", session.prompt)) + MODEL.update_prompt(prompt) + session.prompt = prompt + if "seed" in body: + MODEL.update_seed(int(body["seed"])) + if "input_blend" in body: + MODEL.update_input_blend(float(body["input_blend"])) + return web.json_response(session.to_json() | {"seed": MODEL.seed, "input_blend": MODEL.input_blend}) + + +# --- lifecycle ------------------------------------------------------------- +async def _on_startup(app: web.Application) -> None: + global _ORCHESTRATOR_URL + _ORCHESTRATOR_URL = app["args"].orchestrator + # Warm up in the background so /status + /health serve immediately while the + # pipeline loads (/health stays 503 until ready). + app["warmup"] = asyncio.create_task(_warmup(app)) + + +async def _on_cleanup(app: web.Application) -> None: + task = app.get("warmup") + if task is not None: + task.cancel() + with suppress(asyncio.CancelledError, Exception): + await task + await _close_session() + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Realtime FLUX Klein static app on the new Livepeer runner.") + # Orchestrator URL is needed only for create_trickle_channels; there is no + # registration (static runner — the orchestrator knows us via runners.json). + p.add_argument("--orchestrator", default="http://localhost:8935") + p.add_argument("--host", default=DEFAULT_HOST) + p.add_argument("--port", type=int, default=DEFAULT_PORT) + return p.parse_args() + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + args = _parse_args() + app = web.Application() + app["args"] = args + app.router.add_get("/status", _handle_status) + app.router.add_get("/health", _handle_health) + app.router.add_get("/stats", _handle_stats) + app.router.add_post("/stream", _handle_stream) + app.router.add_post("/update", _handle_update) + app.on_startup.append(_on_startup) + app.on_cleanup.append(_on_cleanup) + web.run_app(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/flux-klein/runners.json b/flux-klein/runners.json new file mode 100644 index 0000000..66ba290 --- /dev/null +++ b/flux-klein/runners.json @@ -0,0 +1,13 @@ +{ + "runners": [ + { + "label": "flux-klein", + "app": "livepeer-example/flux-klein", + "runner_url": "http://app:8720", + "health_url": "/health", + "mode": "persistent", + "capacity": 1, + "price_info": { "price_per_unit": 1, "pixels_per_unit": 1000, "unit": "USD" } + } + ] +} diff --git a/flux-klein/view.sh b/flux-klein/view.sh new file mode 100755 index 0000000..303b6b5 --- /dev/null +++ b/flux-klein/view.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# Turnkey live showcase: webcam -> FLUX Klein -> a low-latency player window. +# Under the hood it's just client.py piping MPEG-TS to a player (prefers mpv, +# falls back to ffplay). Auto-detects macOS (AVFoundation) vs Linux (v4l2). +# +# ./view.sh "a cyberpunk portrait, neon lighting" +# ./view.sh "an oil painting" 1 # camera index (macOS) ... +# ./view.sh "an oil painting" /dev/video2 # ... or device (Linux) +# DISCOVERY=https://1.2.3.4:8935/discovery ./view.sh "..." # remote orchestrator +set -euo pipefail + +PROMPT="${1:-a psychedelic landscape, vivid colors, intricate details}" +DISCOVERY="${DISCOVERY:-https://localhost:8935/discovery}" + +# macOS -> AVFoundation camera index (default 0); Linux -> v4l2 device. +if [[ "$(uname)" == "Darwin" ]]; then + DEVICE="${2:-0}" + WEBCAM=(--webcam-macos "$DEVICE") +else + DEVICE="${2:-/dev/video0}" + WEBCAM=(--webcam "$DEVICE") +fi + +# Prefer uv if available, else the venv's python. +if command -v uv >/dev/null 2>&1; then + RUN=(uv run client.py) +else + RUN=(python client.py) +fi + +if command -v mpv >/dev/null 2>&1; then + PLAYER=(mpv --profile=low-latency --no-cache -) +elif command -v ffplay >/dev/null 2>&1; then + PLAYER=(ffplay -fflags nobuffer -flags low_delay -) +else + echo "need mpv or ffplay installed to view the stream" >&2 + exit 1 +fi + +"${RUN[@]}" "${WEBCAM[@]}" --prompt "$PROMPT" --discovery "$DISCOVERY" --output - | "${PLAYER[@]}"