Datasets I/O — dream.io

pip install "dream-engine[io]"

from dream.io import frames_from_hf
 
src = frames_from_hf(
    repo,
    *,
    split="train",
    frame_field,
    actions_field,
    id_field=None,
    revision=None,
    token=None,
) -> IterableSource

import dream
from dream.io import frames_from_hf, RolloutSink
 
src = frames_from_hf(
    "kingJulio/dream-engine-example-frames",
    frame_field="start_frame",
    actions_field="action_sequence",
)
 
for row in src:
    print(row.row_id, len(row.frame_bytes), len(row.actions_bytes))
    break  # 16 rows total in the public fixture

from dream.io import frames_from_s3
 
src = frames_from_s3(
    uri,
    *,
    frame_glob="*.png",
    actions_glob="*.npy",
    id_pattern=r"(.*)\.png$",
    boto3_kwargs=None,
) -> IterableSource

from dream.io import frames_from_s3
 
src = frames_from_s3(
    "s3://my-robot-data/episodes/",
    frame_glob="*.png",
    actions_glob="*.npy",
)

from dream.io import frames_from_dir
 
src = frames_from_dir(
    path,
    *,
    frame_glob="*.png",
    actions_glob="*.npy",
) -> IterableSource

from dream.io import frames_from_dir
 
# Directory layout:
#   /data/episodes/ep_001.png  /data/episodes/ep_001.npy
#   /data/episodes/ep_002.png  /data/episodes/ep_002.npy
src = frames_from_dir("/data/episodes")

sink = RolloutSink.dir(
    path,
    *,
    videos_subdir="videos",
    metadata_filename="metadata.parquet",
) -> RolloutSink

from dream.io import RolloutSink
 
sink = RolloutSink.dir("./out")
 
result = client.predict_many(src, sink, spec="dreamdojo-2b-gr1")
# ./out/videos/row_00000000.mp4
# ./out/videos/row_00000001.mp4
# …
# ./out/metadata.parquet
print(result.output_uri)  # /absolute/path/to/out

sink = RolloutSink.s3(
    uri,
    *,
    boto3_kwargs=None,
    videos_subdir="videos",
    metadata_filename="metadata.parquet",
) -> RolloutSink

sink = RolloutSink.s3("s3://my-bucket/rollouts/run-2026-05-06/")
result = client.predict_many(src, sink, spec="dreamdojo-2b-gr1")
print(result.output_uri)  # s3://my-bucket/rollouts/run-2026-05-06

sink = RolloutSink.hf(
    repo,
    *,
    token=None,
    private=False,
    shard_rows=500,
    license="cc-by-4.0",
    pretty_name=None,
) -> RolloutSink

import os
from dream.io import frames_from_hf, RolloutSink
 
src  = frames_from_hf(
    "kingJulio/dream-engine-example-frames",
    frame_field="start_frame",
    actions_field="action_sequence",
)
sink = RolloutSink.hf(
    "my-org/gr1-predictions",
    token=os.environ["HF_WRITE_TOKEN"],
    shard_rows=500,
)
result = client.predict_many(src, sink, spec="dreamdojo-2b-gr1")
print(result.output_uri)  # hf://datasets/my-org/gr1-predictions

import numpy as np
from io import BytesIO
from PIL import Image
from dream.io import IterableSource, SourceRow
 
class MyDatabaseSource(IterableSource):
    """Loads episodes from a custom database."""
 
    def __init__(self, connection_string: str):
        self._conn_str = connection_string
        self.rows_hint = self._count_rows()  # optional; enables estimate_cost
 
    def _count_rows(self) -> int:
        # query your database for the total row count
        ...
 
    def __iter__(self):
        for episode in self._fetch_episodes():
            # encode frame to PNG bytes
            buf = BytesIO()
            Image.fromarray(episode["start_frame"]).save(buf, format="PNG")
            frame_bytes = buf.getvalue()
 
            # encode actions to .npy bytes
            actions_buf = BytesIO()
            np.save(actions_buf, episode["actions"].astype(np.float32))
            actions_bytes = actions_buf.getvalue()
 
            yield SourceRow(
                row_id=str(episode["id"]),
                frame_bytes=frame_bytes,
                actions_bytes=actions_bytes,
                metadata={"source": "my-database"},
            )

src = MyDatabaseSource("postgresql://...")
result = client.predict_many(src, sink, spec="dreamdojo-2b-gr1")

from dream.io import RolloutSink, RolloutRecord
 
class MyWarehouseSink(RolloutSink):
    """Writes rollouts to a proprietary video warehouse."""
 
    def __init__(self, warehouse_url: str):
        self._url = warehouse_url
        self._written = 0
 
    def write(self, record: RolloutRecord) -> None:
        # record.row_id     — matches source row id
        # record.mp4_bytes  — raw mp4 from the engine
        # record.metadata   — provenance dict from Client.predict_many
        self._upload(record.row_id, record.mp4_bytes)
        self._written += 1
 
    def finalize(self) -> str:
        return f"{self._url} ({self._written} rollouts)"
 
    def _upload(self, row_id: str, mp4_bytes: bytes) -> None:
        ...