PyTrickle reference

pip install git+https://github.com/livepeer/pytrickle.git

from pytrickle import FrameProcessor
from pytrickle.frames import VideoFrame, AudioFrame
from typing import Optional, List
import torch

class MyProcessor(FrameProcessor):

    async def initialize(self):
        """Called once on startup. Load your model here."""
        self.model = load_model()  # your model loading logic

    async def process_video_async(self, frame: VideoFrame) -> Optional[VideoFrame]:
        """
        Called once per video frame.

        Args:
            frame: VideoFrame containing a PyTorch tensor (H, W, C) and metadata

        Returns:
            Processed VideoFrame, or None to drop the frame
        """
        tensor = frame.tensor  # torch.Tensor, shape (H, W, C), dtype uint8

        with torch.no_grad():
            processed = self.model(tensor)

        return frame.replace_tensor(processed)

    async def process_audio_async(self, frame: AudioFrame) -> Optional[List[AudioFrame]]:
        """
        Called once per audio frame.

        Returns:
            List of AudioFrames to output, or None to drop
        """
        return [frame]  # pass through

    def update_params(self, params: dict):
        """
        Called when the gateway or client sends updated parameters mid-stream.
        Implement to support dynamic model configuration.
        """
        pass

# Replace the tensor while keeping metadata
new_frame = frame.replace_tensor(processed_tensor)

# Move tensor to GPU
frame_gpu = frame.to('cuda')

# Move tensor to CPU
frame_cpu = frame.to('cpu')

from pytrickle import StreamServer

server = StreamServer(
    frame_processor=MyProcessor(),
    port=8000,
    capability_name='live-video-to-video',  # pipeline type identifier
    host='0.0.0.0',
    target_fps=24,        # output frame rate (1-60)
    max_queue_size=30,    # frames to buffer before dropping
)

# Run the server (blocks until stopped)
import asyncio
asyncio.run(server.run_forever())

from pytrickle import TrickleClient

async def stream_frames():
    async with TrickleClient(
        subscribe_url='http://trickle-server/input-stream',
        publish_url='http://trickle-server/output-stream',
    ) as client:
        async for frame in client.video_frames():
            processed = await my_model(frame)
            await client.publish_video_frame(processed)

metrics = processor.get_metrics()
# {
#   'fps': 24.1,
#   'latency_ms': 42.3,
#   'gpu_memory_mb': 4096,
#   'frames_processed': 10234,
#   'frames_dropped': 3,
#   'error_count': 0,
# }