Dream Engines
Open in Colab

Bulk augmentation — GR-1

This example takes the public kingJulio/dream-engine-example-frames dataset — 16 rows of synthetic GR-1-shaped data — and runs DreamDojo against every row using client.predict_many. The output is a local directory of mp4 files + a Parquet metadata table you can inspect immediately.

The source data is synthetic (not real teleop data). The point of this example is the SDK shape and the end-to-end plumbing, not the visual fidelity of the predictions.

Try it

Open the Colab notebook: notebooks/bulk-augment-gr1.ipynb

Or run the script below directly.

Install

BASH
pip install "dream-engine[io,decode]"
export DREAM_API_KEY="dre_..."

Code

PYTHON
"""
bulk_augment_gr1.py — predict every row in kingJulio/dream-engine-example-frames,
write results to ./out/, inspect the first rollout.
"""
import os
import dream
import pyarrow.parquet as pq
from dream.io import frames_from_hf, RolloutSink
 
# ── 1. Client ──────────────────────────────────────────────────────
client = dream.Client()
 
# ── 2. Source ──────────────────────────────────────────────────────
# kingJulio/dream-engine-example-frames is public — no token needed.
# 16 rows × (start_frame: PIL.Image, action_sequence: (48, 384) float32)
src = frames_from_hf(
    "kingJulio/dream-engine-example-frames",
    frame_field="start_frame",
    actions_field="action_sequence",
)
print(f"Source rows: {src.rows_hint}")  # 16
 
# ── 3. Estimate cost ───────────────────────────────────────────────
estimate = client.estimate_cost(src, spec="dreamdojo-2b-gr1")
print(
    f"Estimate: {estimate.rows} rows × {estimate.frames_per_row} frames "
    f"= {estimate.total_frames} total frames "
    f"≈ ${estimate.total_usd:.4f}"
)
# Estimate: 16 rows × 49 frames = 784 total frames ≈ $0.3920
 
# ── 4. Sink ────────────────────────────────────────────────────────
# RolloutSink.dir writes mp4s + a Parquet table to a local directory.
# Swap for RolloutSink.hf("my-org/gr1-predictions") to push to HF.
sink = RolloutSink.dir("./out")
 
# ── 5. Predict ─────────────────────────────────────────────────────
result = client.predict_many(
    src,
    sink,
    spec="dreamdojo-2b-gr1",
    concurrency=4,    # up to 4 in-flight requests
    on_error="skip",  # skip failures; use "halt" to stop on first error
    progress=True,    # tqdm bar
)
 
print(f"Done: {result.ok} ok / {result.failed} failed")
print(f"Output: {result.output_uri}")
 
# ── 6. Inspect ─────────────────────────────────────────────────────
table = pq.read_table("./out/metadata.parquet")
df = table.to_pandas()
print(df[["row_id", "status", "engine_wall_ms", "cost_usd"]].head())
 
# Load and display the first rollout
first_row = df.iloc[0]
mp4_path  = f"./out/videos/{first_row['row_id']}.mp4"
 
try:
    import mediapy
    frames = mediapy.read_video(mp4_path)
    print(f"Video shape: {frames.shape}")  # (49, 480, 640, 3)
    # In a Jupyter notebook:
    # mediapy.show_video(frames[:4])
except ImportError:
    from PIL import Image
    from dream._decode import decode_mp4_to_ndarray
    # Decode via dream's built-in frame access (requires [decode] extra)
    with open(mp4_path, "rb") as f:
        mp4_bytes = f.read()
    frames = decode_mp4_to_ndarray(mp4_bytes)   # (T+1, 480, 640, 3) uint8
    Image.fromarray(frames[0]).save("first_frame.png")
    print("First frame saved to first_frame.png")

What this script does

  1. Streams kingJulio/dream-engine-example-frames via the HF datasets library (no full download — rows arrive one at a time).
  2. Calls client.estimate_cost to preview the spend before any GPU work runs.
  3. Runs client.predict_many — an async concurrency pool (default 4) posts each row to POST /v1/predict, retrying on 429/5xx.
  4. Writes each mp4 to ./out/videos/<row_id>.mp4 and accumulates a Parquet metadata table at ./out/metadata.parquet.
  5. Reads the Parquet table with pyarrow and decodes the first rollout for visual inspection.

Adapting this to your own data

HF dataset with PIL.Image columns

PYTHON
src = frames_from_hf(
    "my-org/my-episodes",
    split="train",
    frame_field="start_frame",      # PIL.Image column
    actions_field="action_sequence", # (T, A) float32 ndarray column
    token=os.environ["HF_TOKEN"],
)

HF dataset with encoded image dicts

Some HF datasets store images as {"bytes": b"...", "path": "img.png"}:

PYTHON
src = frames_from_hf(
    "my-org/my-episodes",
    frame_field="image",      # {"bytes":…, "path":…} dict column
    actions_field="actions",
)

The loader handles both PIL images and encoded dicts automatically.

Local directory of .png + .npy pairs

PYTHON
from dream.io import frames_from_dir
 
# Expects: /data/ep_001.png + /data/ep_001.npy, etc.
src = frames_from_dir("/data/episodes")

No [io] extra required for the local loader.

Cost

The public fixture has 16 rows × 49 frames × $0.0005/frame:

16 × 49 × $0.0005 = $0.392

Small enough to run a few times while iterating on your pipeline. For production runs (thousands of rows), call client.estimate_cost first and confirm the balance in your dashboard.

Next