---
title: "Use your model"
description: "Run PyTorch, ONNX Runtime, MLX or any other framework on live frames, with the frame handed to the model on the GPU and the results published for the rest of the stream."
type: "docs-page"
url: "https://tatolab.com/docs/build/models"
site: "Tatolab"
section: "Build"
---

# Use your model

> Run PyTorch, ONNX Runtime, MLX or any other framework on live frames, with the frame handed to the model on the GPU and the results published for the rest of the stream.

Tatolab doesn't ship an inference stack. Your model runs in a node, with whatever framework you already use, installed in your project like any other package. This page shows how to hand it frames on the GPU, and publish its results for the rest of the stream and other streams to read.

## Where the model goes

Put the model in its own node, as one more reader of the camera's output, with a `newest` input:

```txt
CameraSource ─┬─ DisplayWindow
              └─ ObjectDetector → detections → (logger, dashboard, agent, …)
```

- The node runs in its own process, so a slow or crashing model can't stall the picture or anything else in the stream.
- `newest` means the model always gets the latest frame. If inference takes 50 ms and the camera delivers a frame every 33 ms, the model skips frames instead of falling further and further behind.
- Load the model once, in `setup`, and run it in `process`.

Install your framework into the project:

```sh
uv add torch
```

## Get the frame as a tensor

A frame read with `into=VideoFrame` hands itself to any framework that speaks [DLPack](https://dmlc.github.io/dlpack/latest/), without copying it to the CPU:

```python
import torch

frame = ctx.inputs.read("video", into=VideoFrame)
if frame is not None:
    pixels = torch.from_dlpack(frame)  # height × width × 4, uint8, on the GPU
    rgb = pixels[..., :3].permute(2, 0, 1).unsqueeze(0).float() / 255.0  # 1 × 3 × H × W
```

Where the tensor lives depends on the machine:

| Machine | `torch.from_dlpack(frame)` gives |
|---|---|
| Linux with an NVIDIA GPU and CUDA | A CUDA tensor |
| Apple silicon Mac | A Metal tensor (`mps`), with PyTorch 2.10 or later. MLX 0.32 or later imports it as an MLX array. |
| No GPU framework available | The frame in CPU memory |

The tensor is valid while the `frame` object exists, and the runtime won't reuse the frame's memory until you let it go. With a GPU framework, getting the tensor costs at most one copy on the GPU, never a trip through the CPU.

Treat the tensor as read-only. Other nodes can be reading the same frame, so writing into it would change their picture too. Operations that produce new tensors, like the `.float()` above, are fine.

## A complete detector

This node runs [YOLO](https://docs.ultralytics.com) on every frame it gets and publishes what it finds. It prepares the model's input on the GPU with `ModelInputTensorKernel`, which converts the frame to RGB, scales it to 0–1 and lays it out channels-first in one GPU pass:

```python title="nodes/object_detector.py"
"""Finds objects in the camera's frames with YOLO, on the GPU."""

from dataclasses import dataclass
from typing import Annotated

import torch
from tatolab.stream import (
    ModelInputTensorKernel,
    RuntimeContextFullAccess,
    RuntimeContextLimitedAccess,
    VideoFrame,
    log,
    node,
)
from ultralytics import YOLO


@dataclass
class ObjectDetectorConfig:
    weights: Annotated[str, "YOLO weights file, in the project folder."] = "yolov8n.pt"
    confidence: Annotated[float, "Ignore detections scored below this."] = 0.35


@node
class ObjectDetector:
    """Finds objects in each frame with YOLO and publishes them."""

    @node.input(delivery_profile="newest")
    def video(self) -> VideoFrame: ...

    @node.output()
    def detections(self) -> dict: ...

    def __init__(self, config: ObjectDetectorConfig) -> None:
        self.weights = config.weights
        self.confidence = config.confidence

    def setup(self, ctx: RuntimeContextFullAccess) -> None:
        # The frame arrives on the GPU, so the model goes on the same device.
        device = torch.accelerator.current_accelerator(check_available=True)
        if device is None:
            raise RuntimeError("torch sees no GPU it can use")
        self.model = YOLO(self.weights)
        self.model.to(device)
        # RGB, channels first, scaled to 0-1, and padded at the bottom and right
        # to a multiple of 32 (YOLO's stride), without resizing.
        self.prepare = ModelInputTensorKernel.compile(
            ctx.gpu_full_access, fit="pad_bottom_right", pad_to_multiple_of=32
        )
        log.info("detector ready", weights=self.weights, device=str(device))

    def process(self, ctx: RuntimeContextLimitedAccess) -> None:
        frame = ctx.inputs.read("video", into=VideoFrame)
        if frame is None:
            return
        gpu = ctx.gpu_limited_access
        with gpu.resolve_surface(frame.surface_id) as surface:
            model_input = self.prepare.apply_to_surface(gpu, surface)
            with model_input.tensor_surface as tensor:
                result = self.model.predict(
                    source=torch.from_dlpack(tensor), conf=self.confidence, verbose=False
                )[0]

        objects = []
        if result.boxes is not None:
            boxes = model_input.geometry.boxes_to_source(result.boxes.xyxy).round().int().tolist()
            for box, label, score in zip(boxes, result.boxes.cls.tolist(), result.boxes.conf.tolist()):
                objects.append({"label": result.names[int(label)], "score": round(score, 3), "box": box})
        ctx.outputs.write("detections", {"objects": objects, "timestamp_ns": frame.timestamp_ns}, timestamp_ns=frame.timestamp_ns)
```

And the stream:

```python title="stream.py"
from tatolab.stream import CameraSource, DisplayWindow, StreamBuilder, stream

from nodes.object_detector import ObjectDetector


@stream
def detector(stream_builder: StreamBuilder) -> None:
    """The camera in a window, with YOLO finding objects in every frame."""
    camera = stream_builder.add(CameraSource, config={"max_width": 1280, "max_height": 720})
    window = stream_builder.add(DisplayWindow)
    finder = stream_builder.add(ObjectDetector, config={"confidence": 0.4})
    stream_builder.connect(camera.output("video"), window.input("video"))
    stream_builder.connect(camera.output("video"), finder.input("video"))
    stream_builder.expose(finder.output("detections"))
```

Install its dependencies with `uv add torch ultralytics`, and fetch the weights once before you run it, so the node's `setup` doesn't spend its start-up time downloading:

```sh
uv run python -c "from ultralytics import YOLO; YOLO('yolov8n.pt')"
```

That leaves `yolov8n.pt` in the project folder, where the node's process starts and looks for it.

### How the input is prepared

`ModelInputTensorKernel.compile` takes the model's input format once, in `setup`:

- `fit`: `"stretch"` resizes to `width` × `height`; `"letterbox"` resizes keeping the aspect ratio and pads the rest; `"pad_bottom_right"` doesn't resize, and pads to a multiple of `pad_to_multiple_of`.
- `channel_order`: `"rgb"` or `"bgr"`. `layout`: `"nchw"` or `"nhwc"`. `dtype`: `"float32"` or `"float16"`.
- `scale`, `mean` and `std`: each pixel value `x` (0–255) becomes `(x * scale - mean) / std`. The defaults scale to 0–1. For a model trained with ImageNet normalization, pass `mean=(0.485, 0.456, 0.406)` and `std=(0.229, 0.224, 0.225)`.

`apply_to_surface` returns the tensor and its `geometry`. Use `geometry.boxes_to_source(boxes)` to map boxes from the model's coordinates back to the frame's, whatever fit you chose.

## Publish the results

Write results as a plain dict, one bag per frame you processed. Carry the frame's `timestamp_ns`, in the bag and as the write's `timestamp_ns=`, so anything downstream can line the result up with the picture it came from:

```python
{"objects": [{"label": "person", "score": 0.91, "box": [412, 128, 640, 590]}], "timestamp_ns": 893417260215}
```

From there:

- Another node can log it, count it, raise an alarm, or steer a [GPU effect](https://tatolab.com/docs/build/gpu-effects) through a dial.
- `stream_builder.expose(...)` makes it readable by other streams on the machine and by agents. Once something reads it, a node in this stream or another stream, an agent can [tap](https://tatolab.com/docs/run/inspect) the detections port and answer questions like "how many people are in view?"
- To draw boxes on the video, the node needs an output frame of its own to draw into. The [fisheye-object-detection](https://github.com/tatolab/streamlib/tree/main/examples/fisheye-object-detection) example does exactly that: it draws into a texture it owns and publishes it, with the detections in the same bag.

## Other frameworks and the CPU

- **ONNX Runtime, TensorRT, CuPy, JAX:** anything that accepts a DLPack tensor can take `frame` or the prepared tensor directly.
- **MLX on a Mac:** MLX 0.32 or later imports the frame as an array.
- **CPU-only models:** read the pixels with `frame.cpu()`, which gives you a numpy array. The array is valid only inside the `with` block, so take what you need there:

  ```python
  with frame.cpu() as pixels:
      rgb = pixels[..., :3].copy()  # height × width × 3, uint8
  ```

  That's fine a few times a second. For every frame at full resolution, keep the work on the GPU.

## Keeping up

- If inference can't keep up with the camera, the node processes the frames it can and skips the rest, because of `newest`. To see how long inference takes, log it: take `monotonic_now_ns()` before and after `predict`.
- Lower the camera's resolution with `max_width` and `max_height` on `CameraSource` if the model doesn't need full resolution.
- To run only a few times a second, compare `ctx.time` against your next due time and return early, as the [tutorial's motion detector](https://tatolab.com/docs/tutorial/first-node) does.
- Each model node is its own process with its own copy of the model. Two detectors on one GPU means two models in GPU memory.