Skip to content

Commit 378275f

Browse files
rickstaaclaude
andcommitted
feat(live): add @pipeline decorator for simplified pipeline creation
Add a `@pipeline` decorator that turns a plain function or class into a fully working live Pipeline subclass — handling frame queues, lifecycle management, parameter validation, and thread safety automatically. Includes two example pipelines: - green_shift: minimal function form (4 lines of user code) - depth_midas: class form with MiDaS Small depth estimation, demonstrating prepare_models, on_ready, transform, on_update, and on_stop Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 50a742c commit 378275f

6 files changed

Lines changed: 428 additions & 22 deletions

File tree

docs/custom-pipeline.md

Lines changed: 148 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -74,63 +74,177 @@ touch src/my_pipeline/pipeline/params.py
7474

7575
---
7676

77-
## Step 2: Implement the Pipeline Interface
77+
## Step 2: Implement the Pipeline
7878

79-
### 2.1 Define Parameters (Optional)
79+
You have two options: the `@pipeline` decorator (recommended) or the raw `Pipeline` interface (full control over frame queuing, batching, and threading).
8080

81-
Implement `src/my_pipeline/pipeline/params.py`:
81+
### Option A: `@pipeline` Decorator (Recommended)
82+
83+
The `@pipeline` decorator handles frame queues, lifecycle, threading, and parameter validation automatically.
84+
85+
**Function form** — simplest possible pipeline:
8286

8387
```python
84-
from runner.live.pipelines import BaseParams
88+
# src/my_pipeline/pipeline/pipeline.py
89+
import torch
90+
from runner.live.pipelines import pipeline, BaseParams
91+
from runner.live.trickle import VideoFrame
92+
93+
@pipeline(name="green-shift")
94+
async def green_shift(frame: VideoFrame, params: BaseParams) -> torch.Tensor:
95+
# Process frame tensor and return modified tensor
96+
tensor = frame.tensor.clone()
97+
tensor[:, :, :, 1] = torch.clamp(tensor[:, :, :, 1] + 0.3, -1.0, 1.0)
98+
return tensor
99+
100+
GreenShiftPipeline = green_shift
101+
```
85102

86-
class MyPipelineParams(BaseParams):
87-
# Define your custom fields here
103+
**Class form** — for model loading, state, and mid-stream parameter updates:
104+
105+
```python
106+
# src/my_pipeline/pipeline/pipeline.py
107+
import logging
108+
import torch
109+
from pydantic import Field
110+
from runner.live.pipelines import pipeline, BaseParams
111+
from runner.live.trickle import VideoFrame
112+
113+
class DepthParams(BaseParams):
114+
colormap: bool = Field(default=True, description="Apply colormap to depth output.")
115+
near_clip: float = Field(default=0.0, ge=0.0, le=1.0)
116+
far_clip: float = Field(default=1.0, ge=0.0, le=1.0)
117+
118+
@pipeline(name="depth-midas", params=DepthParams)
119+
class DepthMidas:
120+
121+
def on_ready(self, **params):
122+
# Load model and set up device
123+
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
124+
self.model = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", pretrained=True)
125+
self.model.to(self.device).eval()
126+
self.transforms = torch.hub.load("intel-isl/MiDaS", "transforms").small_transform
127+
128+
def transform(self, frame: VideoFrame, params: DepthParams) -> torch.Tensor:
129+
# Run inference and return output tensor
130+
tensor = frame.tensor[0]
131+
h, w = tensor.shape[:2]
132+
img = ((tensor + 1.0) / 2.0 * 255).byte().cpu().numpy()
133+
input_batch = self.transforms(img).to(self.device)
134+
with torch.no_grad():
135+
depth = self.model(input_batch)
136+
depth = torch.nn.functional.interpolate(
137+
depth.unsqueeze(1), size=(h, w), mode="bilinear", align_corners=False
138+
).squeeze()
139+
depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8)
140+
out = depth.unsqueeze(-1).expand(-1, -1, 3)
141+
return (out * 2.0 - 1.0).unsqueeze(0).to(frame.tensor.device)
142+
143+
def on_update(self, **params):
144+
# Handle mid-stream parameter changes
145+
logging.info(f"Params updated: {params}")
146+
147+
def on_stop(self):
148+
logging.info("DepthMidas stopped")
149+
150+
@classmethod
151+
def prepare_models(cls):
152+
# Download model weights at build time
153+
torch.hub.load("intel-isl/MiDaS", "MiDaS_small", pretrained=True)
154+
155+
DepthMidasPipeline = DepthMidas
88156
```
89157

90-
### 2.2 Implement the Pipeline
158+
**Decorator class methods:**
159+
160+
| Method | Required | When it runs |
161+
|---|---|---|
162+
| `transform(self, frame, params)` | Yes | Every frame |
163+
| `on_ready(self, **params)` | No | Once on startup |
164+
| `on_update(self, **params)` | No | When params change mid-stream |
165+
| `on_stop(self)` | No | On shutdown |
166+
| `prepare_models(cls)` | No | At build time (model download) |
167+
168+
Both `async def` and `def` work. Sync functions automatically run in a thread pool.
169+
170+
See [`examples/pipelines/`](../examples/pipelines/) for complete working examples.
171+
172+
### Option B: Raw `Pipeline` Interface
91173

92-
Implement `src/my_pipeline/pipeline/pipeline.py`:
174+
For advanced use cases where you need full control over frame queuing, custom batching, or complex threading. Here's the same depth-midas pipeline implemented with the raw interface:
93175

94176
```python
177+
# src/my_pipeline/pipeline/pipeline.py
95178
import asyncio
96179
import logging
97-
180+
import torch
98181
from runner.live.pipelines import Pipeline
99182
from runner.live.trickle import VideoFrame, VideoOutput
100183

101-
class MyPipeline(Pipeline):
184+
class DepthMidasPipeline(Pipeline):
102185
def __init__(self):
103186
self.frame_queue: asyncio.Queue[VideoOutput] = asyncio.Queue()
104187

105188
async def initialize(self, **params):
106-
logging.info(f"Initializing with params: {params}")
107-
# Load your model here (use asyncio.to_thread for blocking operations)
108-
# self.model = await asyncio.to_thread(load_model, params)
189+
# Load model (use asyncio.to_thread for blocking operations)
190+
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
191+
self.model = await asyncio.to_thread(
192+
torch.hub.load, "intel-isl/MiDaS", "MiDaS_small", pretrained=True
193+
)
194+
self.model.to(self.device).eval()
195+
self.transforms = torch.hub.load("intel-isl/MiDaS", "transforms").small_transform
109196

110197
async def put_video_frame(self, frame: VideoFrame, request_id: str):
111-
# Process frame here (use asyncio.to_thread for blocking inference)
112-
# result = await asyncio.to_thread(self.model.predict, frame.tensor)
113-
await self.frame_queue.put(VideoOutput(frame, request_id))
198+
# Run inference in thread pool to avoid blocking the event loop
199+
result = await asyncio.to_thread(self._run_inference, frame)
200+
await self.frame_queue.put(
201+
VideoOutput(frame, request_id).replace_tensor(result)
202+
)
203+
204+
def _run_inference(self, frame: VideoFrame) -> torch.Tensor:
205+
tensor = frame.tensor[0]
206+
h, w = tensor.shape[:2]
207+
img = ((tensor + 1.0) / 2.0 * 255).byte().cpu().numpy()
208+
input_batch = self.transforms(img).to(self.device)
209+
with torch.no_grad():
210+
depth = self.model(input_batch)
211+
depth = torch.nn.functional.interpolate(
212+
depth.unsqueeze(1), size=(h, w), mode="bilinear", align_corners=False
213+
).squeeze()
214+
depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8)
215+
out = depth.unsqueeze(-1).expand(-1, -1, 3)
216+
return (out * 2.0 - 1.0).unsqueeze(0).to(frame.tensor.device)
114217

115218
async def get_processed_video_frame(self) -> VideoOutput:
116219
return await self.frame_queue.get()
117220

118221
async def update_params(self, **params):
222+
# Handle mid-stream parameter changes
223+
# Return asyncio.create_task(...) for slow reloads (shows loading overlay)
119224
logging.info(f"Updating params: {params}")
120-
# Return asyncio.create_task(...) if reload needed (shows loading overlay)
121225

122226
async def stop(self):
123227
logging.info("Stopping pipeline")
124228

125229
@classmethod
126230
def prepare_models(cls):
127-
logging.info("Preparing models")
128-
# Download models, compile TensorRT engines, etc.
231+
# Download model weights at build time
232+
torch.hub.load("intel-isl/MiDaS", "MiDaS_small", pretrained=True)
129233
```
130234

131-
For a real-world implementation, see [scope-runner's pipeline](https://github.com/daydreamlive/scope-runner/blob/dec9ecf7e306892df9cfae21759c23fdf15b0510/src/scope_runner/pipeline/pipeline.py#L22).
235+
For another real-world example, see [scope-runner's pipeline](https://github.com/daydreamlive/scope-runner/blob/dec9ecf7e306892df9cfae21759c23fdf15b0510/src/scope_runner/pipeline/pipeline.py#L22).
236+
237+
### Define Parameters (Optional)
238+
239+
```python
240+
# src/my_pipeline/pipeline/params.py
241+
from runner.live.pipelines import BaseParams
242+
243+
class MyPipelineParams(BaseParams):
244+
# Define your custom fields here
245+
```
132246

133-
### 2.3 Keep Module Exports Minimal
247+
### Keep Module Exports Minimal
134248

135249
> **⚠️ Important**: Do **not** export `Pipeline` or `Params` classes from `__init__.py`. The loader imports these by their full path (`module.path:ClassName`), and re-exporting from `__init__.py` would trigger expensive imports (torch, etc.) when only loading the params class.
136250
@@ -405,6 +519,19 @@ The orchestrator will route requests to your local runner at `http://localhost:8
405519

406520
---
407521

522+
## Publishing to the Marketplace
523+
524+
Publish your pipeline to make it discoverable on the marketplace. The CLI extracts the parameter schema from your `@pipeline` decorator and registers it automatically.
525+
526+
```bash
527+
livepeer publish examples/pipelines/depth_midas.py
528+
livepeer teardown depth-midas
529+
livepeer list
530+
livepeer inspect examples/pipelines/depth_midas.py
531+
```
532+
533+
---
534+
408535
## Reference Implementation
409536

410537
For a complete working example, see [scope-runner](https://github.com/livepeer/scope-runner).

examples/live-video-to-video/__init__.py

Whitespace-only changes.
Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
"""Depth Estimation (MiDaS) -- @pipeline example with model loading.
2+
3+
Uses MiDaS Small to produce a real-time depth map from video frames.
4+
Demonstrates prepare_models (download + load) and GPU inference in transform.
5+
6+
Usage:
7+
python -m runner.live.infer \
8+
--pipeline '{
9+
"pipeline_cls": "examples.pipelines.depth_midas:DepthMidasPipeline"
10+
}'
11+
"""
12+
13+
import logging
14+
15+
import torch
16+
from pydantic import Field
17+
18+
from runner.live.pipelines import pipeline, BaseParams
19+
from runner.live.trickle import VideoFrame
20+
21+
22+
class DepthParams(BaseParams):
23+
"""Parameters for depth estimation, adjustable mid-stream."""
24+
25+
colormap: bool = Field(
26+
default=True,
27+
description="Apply colormap to depth output. False = grayscale.",
28+
)
29+
near_clip: float = Field(
30+
default=0.0, ge=0.0, le=1.0,
31+
description="Clip depth values below this threshold (0.0 = no clip).",
32+
)
33+
far_clip: float = Field(
34+
default=1.0, ge=0.0, le=1.0,
35+
description="Clip depth values above this threshold (1.0 = no clip).",
36+
)
37+
38+
39+
@pipeline(name="depth-midas", params=DepthParams)
40+
class DepthMidas:
41+
"""Real-time depth estimation using MiDaS Small.
42+
43+
Demonstrates:
44+
- prepare_models: downloads the model at startup
45+
- on_ready: loads model to GPU
46+
- transform: runs inference per frame
47+
- on_update: recomputes clip range when params change mid-stream
48+
"""
49+
50+
def on_ready(self, **params):
51+
self._near = params.get("near_clip", 0.0)
52+
self._far = params.get("far_clip", 1.0)
53+
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
54+
self.model = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", pretrained=True)
55+
self.model.to(self.device).eval()
56+
57+
self.transforms = torch.hub.load("intel-isl/MiDaS", "transforms").small_transform
58+
logging.info(f"MiDaS Small loaded on {self.device}")
59+
60+
def transform(self, frame: VideoFrame, params: DepthParams) -> torch.Tensor:
61+
"""Run MiDaS depth estimation on each frame."""
62+
# frame.tensor is (B, H, W, C) in [-1.0, 1.0]
63+
tensor = frame.tensor[0] # (H, W, C), single batch
64+
h, w = tensor.shape[:2]
65+
66+
# Convert to uint8 RGB for MiDaS transforms
67+
img = ((tensor + 1.0) / 2.0 * 255).byte().cpu().numpy()
68+
69+
# MiDaS expects (B, C, H, W) float32
70+
input_batch = self.transforms(img).to(self.device)
71+
72+
with torch.no_grad():
73+
depth = self.model(input_batch) # (1, h', w')
74+
75+
# Resize back to original resolution
76+
depth = torch.nn.functional.interpolate(
77+
depth.unsqueeze(1), size=(h, w), mode="bilinear", align_corners=False
78+
).squeeze() # (H, W)
79+
80+
# Normalize to [0, 1] and apply depth clipping
81+
depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8)
82+
depth = torch.clamp(depth, self._near, self._far)
83+
depth = (depth - self._near) / (self._far - self._near + 1e-8)
84+
85+
if params.colormap:
86+
# Simple colormap: near=warm, far=cool
87+
r = depth
88+
g = 1.0 - torch.abs(depth - 0.5) * 2.0
89+
b = 1.0 - depth
90+
out = torch.stack([r, g, b], dim=-1) # (H, W, 3)
91+
else:
92+
# Grayscale
93+
out = depth.unsqueeze(-1).expand(-1, -1, 3) # (H, W, 3)
94+
95+
# Convert to [-1, 1] and add batch dim → (1, H, W, C)
96+
out = (out * 2.0 - 1.0).unsqueeze(0)
97+
return out.to(frame.tensor.device)
98+
99+
def on_update(self, **params):
100+
"""Update clip range when params change mid-stream via control_url."""
101+
self._near = params.get("near_clip", 0.0)
102+
self._far = params.get("far_clip", 1.0)
103+
logging.info(f"Depth clip range updated: [{self._near}, {self._far}]")
104+
105+
def on_stop(self):
106+
logging.info("DepthMidas stopped")
107+
108+
@classmethod
109+
def prepare_models(cls):
110+
"""Download MiDaS Small weights ahead of time."""
111+
torch.hub.load("intel-isl/MiDaS", "MiDaS_small", pretrained=True)
112+
logging.info("MiDaS Small weights downloaded")
113+
114+
115+
DepthMidasPipeline = DepthMidas
Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
"""Green Shift -- minimal @pipeline example (function form).
2+
3+
Boosts the green channel of every video frame. No model download, no GPU
4+
inference -- just a tensor op. Ideal for verifying the pipeline infrastructure.
5+
6+
Usage:
7+
python -m runner.live.infer \
8+
--pipeline '{"pipeline_cls":"examples.pipelines.green_shift:GreenShiftPipeline"}'
9+
"""
10+
11+
import torch
12+
13+
from runner.live.pipelines import pipeline, BaseParams
14+
from runner.live.trickle import VideoFrame
15+
16+
17+
@pipeline(name="green-shift")
18+
async def green_shift(frame: VideoFrame, params: BaseParams) -> torch.Tensor:
19+
"""Boost the green channel of every frame.
20+
21+
Frame tensor layout: (B, H, W, C), values in [-1.0, 1.0].
22+
Channel order: R=0, G=1, B=2.
23+
"""
24+
tensor = frame.tensor.clone()
25+
tensor[:, :, :, 1] = torch.clamp(tensor[:, :, :, 1] + 0.3, -1.0, 1.0)
26+
return tensor
27+
28+
29+
GreenShiftPipeline = green_shift
Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from .interface import Pipeline, BaseParams
22
from .spec import PipelineSpec, builtin_pipeline_spec
3+
from .create import pipeline
34

4-
__all__ = ["Pipeline", "BaseParams", "PipelineSpec", "builtin_pipeline_spec"]
5+
__all__ = ["Pipeline", "BaseParams", "PipelineSpec", "builtin_pipeline_spec", "pipeline"]

0 commit comments

Comments
 (0)