Efficient VLM Data Loading with Ray Data and kornia-rs

In this tutorial we build a high-performance image loading pipeline for Vision-Language Models using Ray Data for streaming and kornia-rs for Rust-based image decoding and resizing.

In this tutorial, we demonstrate how to build a high-performance image loading pipeline for Vision-Language Models (VLMs) like LLaVA or Qwen-VL.

Efficient data loading is critical for VLM training. We leverage:

• Ray Data: For scalable streaming of large datasets.
• kornia-rs: For blazing fast, Rust-based image decoding and resizing.
• PyTorch: For final tensor collation.

By performing image decoding and resizing in Rust (before converting to PyTorch tensors), we significantly reduce overhead and memory usage.

%%capture
!pip install kornia kornia-rs "ray[data]" pydantic

import kornia.utils
import kornia_rs
import matplotlib.pyplot as plt
import numpy as np
import ray
import requests

# Configuration
TARGET_SIZE = (336, 336)  # Standard size for models like LLaVA / SigLIP
BATCH_SIZE = 4

Build a small VLM-style dataset

A VLM training sample pairs an image with a conversation. To keep this tutorial fast and self-contained we build a tiny dataset from images hosted on kornia/data, but each row has the exact shape you would get from a real instruction-tuning dataset.

In a real setup you would stream the metadata from Hugging Face instead, for example:

from datasets import load_dataset


def stream_hf_dataset():
    # The 'id' column in this dataset mixes integers and strings; Ray's Arrow
    # schema inference is strict, so we force conversion to keep it stable.
    dataset = load_dataset("liuhaotian/LLaVA-Instruct-150K", split="train", streaming=True)
    for row in dataset:
        yield {"id": str(row["id"]), "image_filename": row["image"], "conversations": row["conversations"]}


# True streaming (lazy loading): never uses more RAM than needed,
# even for terabytes of data.
ds = ray.data.from_generator(stream_hf_dataset)

BASE_URL = "https://github.com/kornia/data/raw/main"
SAMPLE_IMAGES = [
    "panda.jpg",
    "crowd.jpg",
    "soccer.jpg",
    "arturito.jpg",
    "drslump.jpg",
    "squirrel.jpg",
    "mountains.jpg",
    "ninja_turtles.jpg",
]

data_items = [
    {
        "id": str(i),
        "image_url": f"{BASE_URL}/{filename}",
        "conversations": [
            {"from": "human", "value": "What is in this image?"},
            {"from": "gpt", "value": f"A caption describing {filename}."},
        ],
    }
    for i, filename in enumerate(SAMPLE_IMAGES)
]

Decode and resize in Rust

The hot path of the pipeline is decoding the JPEG bytes and resizing to the model resolution. We do both with kornia_rs, which runs in Rust and releases the GIL, so Ray can parallelize it across CPU cores.

def process_batch_rust(batch: dict) -> dict:
    """Decode and resize images using the kornia-rs Rust backend, and format them for model input."""
    processed_images = []

    for url in batch["image_url"]:
        try:
            # 1. Fetch image bytes
            response = requests.get(url, timeout=10)
            response.raise_for_status()

            # 2. Decode & resize (Rust backend): accelerated JPEG decoding (in RGB) and resizing
            decoded_img = kornia_rs.io.decode_image_jpeg(response.content)
            resized_img = kornia_rs.resize(decoded_img, TARGET_SIZE, interpolation="bilinear")

            # 3. Format for the model: HWC -> CHW
            processed_images.append(resized_img.transpose(2, 0, 1))

        except Exception:
            # Fallback: return a black image if download or processing fails
            processed_images.append(np.zeros((3, *TARGET_SIZE), dtype=np.uint8))

    batch["pixel_values"] = processed_images
    return batch

Run the Ray Data pipeline

# 1. Initialize Ray
if not ray.is_initialized():
    ray.init(ignore_reinit_error=True)

# 2. Create the Ray Dataset
ds = ray.data.from_items(data_items)

# 3. Apply the Rust transformation
processed_ds = ds.map_batches(process_batch_rust, batch_size=BATCH_SIZE)

Consume torch batches

iter_torch_batches automatically collates the numpy arrays into torch tensors — this is exactly the batch format a VLM vision encoder expects.

# We select only the image column for the tensor batch
image_ds = processed_ds.select_columns(["pixel_values"])

for batch in image_ds.iter_torch_batches(batch_size=BATCH_SIZE):
    imgs = batch["pixel_values"]

    print(f"Batch tensor shape: {imgs.shape}")
    print(f"Dtype: {imgs.dtype}")

    # Visualize the batch: (C, H, W) tensor -> (H, W, C) numpy array
    fig, axs = plt.subplots(1, len(imgs), figsize=(4 * len(imgs), 4))
    for ax, img in zip(axs, imgs):
        ax.imshow(kornia.utils.tensor_to_image(img))
        ax.axis("off")
    plt.show()

    break  # Only show one batch

ray.shutdown()

Takeaways

• Ray Data gives you streaming, backpressure and parallelism for free — swap from_items for from_generator over a Hugging Face streaming dataset and the same pipeline scales to 150k+ samples.
• kornia_rs decodes and resizes in Rust, keeping the Python workers light.
• iter_torch_batches hands you collated torch.Tensor batches ready for a VLM vision encoder.