Upload data.py

data.py

@@ -0,0 +1,792 @@
+"""
+WebDataset-based data loader for foveated VLM training.
+
+Reads tar shards produced by video2dataset / the CPU precompute pipeline.
+Each sample in a shard contains EITHER:
+  A) Pre-extracted frames:
+     - {key}.jpg or {key}_000.jpg, {key}_001.jpg, ... -- JPEG frames (224x224)
+     - {key}.json -- metadata: {caption, token_ids, loss_mask, ...}
+  B) Raw MP4 from video2dataset:
+     - {key}.mp4 -- raw video file
+     - {key}.txt -- caption text
+     - {key}.json -- metadata: {videoid, duration, url, ...}
+
+On-the-fly tokenization: if token_ids/loss_mask are missing from JSON,
+the sample is tokenized at load time using the provided tokenizer.
+
+Returns dicts with:
+  frames:     [T, 3, 224, 224]  float32, ImageNet-normalized for DINO
+  input_ids:  [S]               long, token IDs
+  loss_mask:  [S]               float32, 1.0 for answer tokens, 0.0 otherwise
+  num_frames: int               actual frame count before any padding
+"""
+
+import io
+import json
+import os
+import re
+import subprocess
+import tempfile
+from typing import Optional
+
+import torch
+import torchvision.transforms.functional as TF
+import webdataset as wds
+
+# ImageNet normalization for DINOv2 (same constants as src/data/llava_video_dataset.py)
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+# Regex to detect multi-frame filenames like "sample_003.jpg"
+_FRAME_INDEX_RE = re.compile(r"^(.+)_(\d{3})\.(jpg|jpeg|png)$")
+
+# Regex to detect single-frame filenames like "sample.jpg"
+_SINGLE_FRAME_RE = re.compile(r"^(.+)\.(jpg|jpeg|png)$")
+
+
+_NORM_MEAN = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
+_NORM_STD = torch.tensor(IMAGENET_STD).view(3, 1, 1)
+
+
+def _load_image_tensor(data: bytes) -> torch.Tensor:
+    """Decode JPEG/PNG bytes to a [3, 224, 224] float32 tensor, ImageNet-normalized."""
+    try:
+        # Fast path: torchvision decode_jpeg — avoids PIL/numpy overhead
+        from torchvision.io import decode_jpeg
+        raw = torch.frombuffer(bytearray(data), dtype=torch.uint8)
+        tensor = decode_jpeg(raw).float().div_(255.0)  # [3, H, W]
+        tensor.sub_(_NORM_MEAN).div_(_NORM_STD)
+        return tensor
+    except Exception:
+        # Fallback: PIL (handles PNG and edge cases)
+        from PIL import Image
+        img = Image.open(io.BytesIO(data)).convert("RGB")
+        tensor = TF.to_tensor(img)  # [3, H, W] float32 in [0, 1]
+        tensor = TF.normalize(tensor, mean=IMAGENET_MEAN, std=IMAGENET_STD)
+        return tensor
+
+
+def _decode_mp4_frames(mp4_bytes: bytes, max_frames: int = 64) -> list[torch.Tensor]:
+    """Decode MP4 bytes to a list of [3, 224, 224] tensors at 1 FPS."""
+    try:
+        import decord
+        decord.bridge.set_bridge("torch")
+        vr = decord.VideoReader(io.BytesIO(mp4_bytes), width=224, height=224)
+        fps = vr.get_avg_fps()
+        total = len(vr)
+        # Sample at 1 FPS
+        step = max(1, int(fps))
+        indices = list(range(0, total, step))[:max_frames]
+        if not indices:
+            return []
+        batch = vr.get_batch(indices)  # [T, H, W, C] uint8
+        frames = []
+        for i in range(batch.shape[0]):
+            t = batch[i].permute(2, 0, 1).float() / 255.0  # [3, 224, 224]
+            t = TF.normalize(t, mean=IMAGENET_MEAN, std=IMAGENET_STD)
+            frames.append(t)
+        return frames
+    except ImportError:
+        pass
+
+    # Fallback: ffmpeg subprocess
+    with tempfile.NamedTemporaryFile(suffix=".mp4", dir="/workspace/tmp", delete=True) as f:
+        f.write(mp4_bytes)
+        f.flush()
+        frames_dir = f.name + "_frames"
+        os.makedirs(frames_dir, exist_ok=True)
+        try:
+            subprocess.run(
+                ["ffmpeg", "-y", "-i", f.name,
+                 "-vf", "fps=1,scale=224:224:force_original_aspect_ratio=increase,crop=224:224",
+                 "-frames:v", str(max_frames), "-q:v", "2",
+                 os.path.join(frames_dir, "frame_%03d.jpg")],
+                capture_output=True, timeout=30,
+            )
+            from PIL import Image
+            frame_files = sorted(os.listdir(frames_dir))
+            frames = []
+            for fname in frame_files[:max_frames]:
+                fp = os.path.join(frames_dir, fname)
+                img = Image.open(fp).convert("RGB")
+                t = TF.to_tensor(img)
+                t = TF.normalize(t, mean=IMAGENET_MEAN, std=IMAGENET_STD)
+                frames.append(t)
+            return frames
+        except Exception:
+            return []
+        finally:
+            import shutil
+            shutil.rmtree(frames_dir, ignore_errors=True)
+
+
+def decode_sample(sample: dict, max_frames: int = 64,
+                  tokenizer=None, stage: int = 1,
+                  replicate_image_frames: int = 1) -> Optional[dict]:
+    """
+    Decode a single webdataset sample dict into training tensors.
+
+    The sample dict has keys like:
+      "jpg" or "jpeg" or "png" -- single frame bytes
+      "000.jpg", "001.jpg", ... -- multi-frame bytes
+      "json" -- metadata JSON bytes or dict
+
+    Returns None if the sample is malformed (caller should filter).
+    """
+    # ------------------------------------------------------------------
+    # 1. Parse metadata JSON
+    # ------------------------------------------------------------------
+    meta_raw = sample.get("json")
+    if meta_raw is None:
+        return None
+
+    if isinstance(meta_raw, bytes):
+        try:
+            meta = json.loads(meta_raw.decode("utf-8"))
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            return None
+    elif isinstance(meta_raw, str):
+        try:
+            meta = json.loads(meta_raw)
+        except json.JSONDecodeError:
+            return None
+    elif isinstance(meta_raw, dict):
+        meta = meta_raw
+    else:
+        return None
+
+    token_ids = meta.get("token_ids")
+    loss_mask = meta.get("loss_mask")
+
+    # On-the-fly tokenization if pre-tokenized data is missing
+    if token_ids is None or loss_mask is None:
+        from tokenization import (
+            tokenize_stage1, tokenize_sft, SOURCE_PROMPTS, DEFAULT_VISUAL_PROMPT,
+        )
+
+        # Unified format: user/assistant keys
+        user_text = meta.get("user", "")
+        assistant_text = meta.get("assistant", "")
+        source = meta.get("source", "")
+
+        if user_text or assistant_text:
+            # Has structured user/assistant format
+            is_text_only = meta.get("frame_count", 0) == 0
+            if stage == 1 and not is_text_only:
+                # Stage 1 visual data: per-source conditioning prompt
+                # Use shard's user field if non-empty, else per-source default
+                user_prompt = user_text if user_text else SOURCE_PROMPTS.get(source, DEFAULT_VISUAL_PROMPT)
+                tok = tokenize_stage1(assistant_text, tokenizer=tokenizer, user_prompt=user_prompt)
+            elif stage == 1 and is_text_only:
+                # Stage 1 text retention: keep proper chat format, all-text loss
+                tok = tokenize_sft(
+                    user_text,
+                    assistant_text,
+                    stage=stage,
+                    tokenizer=tokenizer,
+                )
+                tok["loss_mask"] = [1] * len(tok["token_ids"])
+            else:
+                # Stage 2-3: answer-only loss on assistant portion
+                # Use shard's user field if non-empty, else per-source default
+                effective_user = user_text if user_text else SOURCE_PROMPTS.get(source, DEFAULT_VISUAL_PROMPT)
+                tok = tokenize_sft(
+                    effective_user,
+                    assistant_text,
+                    stage=stage,
+                    tokenizer=tokenizer,
+                )
+        else:
+            # Legacy format: caption key or .txt file
+            caption = meta.get("caption", "")
+            if not caption:
+                txt_raw = sample.get("txt")
+                if isinstance(txt_raw, bytes):
+                    caption = txt_raw.decode("utf-8", errors="replace").strip()
+                elif isinstance(txt_raw, str):
+                    caption = txt_raw.strip()
+
+            if not caption or tokenizer is None:
+                return None
+
+            user_prompt = SOURCE_PROMPTS.get(source, DEFAULT_VISUAL_PROMPT)
+            if stage == 1:
+                tok = tokenize_stage1(caption, tokenizer=tokenizer, user_prompt=user_prompt)
+            else:
+                tok = tokenize_sft(user_prompt, caption, stage=stage, tokenizer=tokenizer)
+
+        if tokenizer is None:
+            return None
+
+        token_ids = tok["token_ids"]
+        loss_mask = tok["loss_mask"]
+
+    # ------------------------------------------------------------------
+    # 2. Collect frames (JPEG bytes or decode from MP4)
+    # ------------------------------------------------------------------
+    frames: list[torch.Tensor] = []
+
+    # Try MP4 first (video2dataset raw output)
+    mp4_data = sample.get("mp4")
+    if isinstance(mp4_data, bytes) and len(mp4_data) > 100:
+        frames = _decode_mp4_frames(mp4_data, max_frames=max_frames)
+    else:
+        # Try numbered JPEG frames (000.jpg, 001.jpg, ...)
+        numbered_keys: list[tuple[int, str]] = []
+        for key in sample:
+            m = re.match(r"^(\d{3})\.(jpg|jpeg|png)$", key)
+            if m:
+                numbered_keys.append((int(m.group(1)), key))
+
+        if numbered_keys:
+            numbered_keys.sort(key=lambda x: x[0])
+            for _, key in numbered_keys:
+                raw = sample[key]
+                if isinstance(raw, bytes):
+                    try:
+                        frames.append(_load_image_tensor(raw))
+                    except Exception:
+                        continue
+        else:
+            # Single frame: look for jpg / jpeg / png key
+            for ext in ("jpg", "jpeg", "png"):
+                if ext in sample and isinstance(sample[ext], bytes):
+                    try:
+                        frames.append(_load_image_tensor(sample[ext]))
+                    except Exception:
+                        pass
+                    break
+
+    if not frames:
+        return None
+
+    # Truncate to max_frames
+    if len(frames) > max_frames:
+        frames = frames[:max_frames]
+
+    # Replicate single-frame images to N frames (A8 ablation: static video)
+    if replicate_image_frames > 1 and len(frames) == 1:
+        frames = frames * replicate_image_frames
+
+    num_frames = len(frames)
+    frames_tensor = torch.stack(frames, dim=0)  # [T, 3, 224, 224]
+
+    # ------------------------------------------------------------------
+    # 3. Build text tensors
+    # ------------------------------------------------------------------
+    input_ids = torch.tensor(token_ids, dtype=torch.long)
+    loss_mask_t = torch.tensor(loss_mask, dtype=torch.float32)
+
+    # Ensure consistent lengths
+    min_len = min(len(input_ids), len(loss_mask_t))
+    input_ids = input_ids[:min_len]
+    loss_mask_t = loss_mask_t[:min_len]
+
+    return {
+        "frames": frames_tensor,       # [T, 3, 224, 224]
+        "input_ids": input_ids,         # [S]
+        "loss_mask": loss_mask_t,       # [S]
+        "num_frames": num_frames,       # int
+    }
+
+
+def decode_dpo_sample(sample: dict, max_frames: int = 64,
+                      tokenizer=None, replicate_image_frames: int = 1) -> Optional[dict]:
+    """
+    Decode a single DPO webdataset sample into training tensors.
+
+    DPO samples have JSON with keys:
+      user:               user prompt
+      chosen_assistant:   preferred response
+      rejected_assistant: dispreferred response
+      source:             dataset source (e.g. "rlaif_v")
+      frame_count:        number of frames (1 for images)
+
+    Returns None if the sample is malformed (caller should filter).
+
+    Returns dict with:
+      frames:             [T, 3, 224, 224]  shared visual input
+      chosen_input_ids:   [S_c]             tokenized user+chosen
+      chosen_loss_mask:   [S_c]             answer-only mask for chosen
+      rejected_input_ids: [S_r]             tokenized user+rejected
+      rejected_loss_mask: [S_r]             answer-only mask for rejected
+      num_frames:         int               actual frame count
+    """
+    # ------------------------------------------------------------------
+    # 1. Parse metadata JSON
+    # ------------------------------------------------------------------
+    meta_raw = sample.get("json")
+    if meta_raw is None:
+        return None
+
+    if isinstance(meta_raw, bytes):
+        try:
+            meta = json.loads(meta_raw.decode("utf-8"))
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            return None
+    elif isinstance(meta_raw, str):
+        try:
+            meta = json.loads(meta_raw)
+        except json.JSONDecodeError:
+            return None
+    elif isinstance(meta_raw, dict):
+        meta = meta_raw
+    else:
+        return None
+
+    user_text = meta.get("user", "")
+    chosen_text = meta.get("chosen_assistant", "")
+    rejected_text = meta.get("rejected_assistant", "")
+
+    if not chosen_text or not rejected_text:
+        return None
+    if tokenizer is None:
+        return None
+
+    # ------------------------------------------------------------------
+    # 2. Tokenize chosen and rejected with answer-only loss masks
+    # ------------------------------------------------------------------
+    from tokenization import tokenize_sft, SOURCE_PROMPTS, DEFAULT_VISUAL_PROMPT
+
+    source = meta.get("source", "")
+    effective_user = user_text if user_text else SOURCE_PROMPTS.get(source, DEFAULT_VISUAL_PROMPT)
+
+    chosen_tok = tokenize_sft(effective_user, chosen_text, stage=3, tokenizer=tokenizer)
+    rejected_tok = tokenize_sft(effective_user, rejected_text, stage=3, tokenizer=tokenizer)
+
+    # ------------------------------------------------------------------
+    # 3. Collect frames (same logic as decode_sample)
+    # ------------------------------------------------------------------
+    frames: list[torch.Tensor] = []
+
+    mp4_data = sample.get("mp4")
+    if isinstance(mp4_data, bytes) and len(mp4_data) > 100:
+        frames = _decode_mp4_frames(mp4_data, max_frames=max_frames)
+    else:
+        numbered_keys: list[tuple[int, str]] = []
+        for key in sample:
+            m = re.match(r"^(\d{3})\.(jpg|jpeg|png)$", key)
+            if m:
+                numbered_keys.append((int(m.group(1)), key))
+
+        if numbered_keys:
+            numbered_keys.sort(key=lambda x: x[0])
+            for _, key in numbered_keys:
+                raw = sample[key]
+                if isinstance(raw, bytes):
+                    try:
+                        frames.append(_load_image_tensor(raw))
+                    except Exception:
+                        continue
+        else:
+            for ext in ("jpg", "jpeg", "png"):
+                if ext in sample and isinstance(sample[ext], bytes):
+                    try:
+                        frames.append(_load_image_tensor(sample[ext]))
+                    except Exception:
+                        pass
+                    break
+
+    if not frames:
+        return None
+
+    if len(frames) > max_frames:
+        frames = frames[:max_frames]
+
+    if replicate_image_frames > 1 and len(frames) == 1:
+        frames = frames * replicate_image_frames
+
+    num_frames = len(frames)
+    frames_tensor = torch.stack(frames, dim=0)  # [T, 3, 224, 224]
+
+    # ------------------------------------------------------------------
+    # 4. Build text tensors
+    # ------------------------------------------------------------------
+    chosen_ids = torch.tensor(chosen_tok["token_ids"], dtype=torch.long)
+    chosen_mask = torch.tensor(chosen_tok["loss_mask"], dtype=torch.float32)
+    rejected_ids = torch.tensor(rejected_tok["token_ids"], dtype=torch.long)
+    rejected_mask = torch.tensor(rejected_tok["loss_mask"], dtype=torch.float32)
+
+    # Ensure consistent lengths within each pair
+    c_len = min(len(chosen_ids), len(chosen_mask))
+    chosen_ids = chosen_ids[:c_len]
+    chosen_mask = chosen_mask[:c_len]
+
+    r_len = min(len(rejected_ids), len(rejected_mask))
+    rejected_ids = rejected_ids[:r_len]
+    rejected_mask = rejected_mask[:r_len]
+
+    return {
+        "frames": frames_tensor,             # [T, 3, 224, 224]
+        "chosen_input_ids": chosen_ids,       # [S_c]
+        "chosen_loss_mask": chosen_mask,       # [S_c]
+        "rejected_input_ids": rejected_ids,   # [S_r]
+        "rejected_loss_mask": rejected_mask,   # [S_r]
+        "num_frames": num_frames,             # int
+    }
+
+
+def _sample_decoder(max_frames: int, tokenizer=None, stage: int = 1,
+                    replicate_image_frames: int = 1):
+    """Return a map function for use in a webdataset pipeline."""
+    def _decode(sample):
+        result = decode_sample(sample, max_frames=max_frames,
+                               tokenizer=tokenizer, stage=stage,
+                               replicate_image_frames=replicate_image_frames)
+        if result is None:
+            return None
+        return result
+    return _decode
+
+
+def _dpo_sample_decoder(max_frames: int, tokenizer=None,
+                        replicate_image_frames: int = 1):
+    """Return a map function for DPO samples in a webdataset pipeline."""
+    def _decode(sample):
+        result = decode_dpo_sample(sample, max_frames=max_frames,
+                                   tokenizer=tokenizer,
+                                   replicate_image_frames=replicate_image_frames)
+        if result is None:
+            return None
+        return result
+    return _decode
+
+
+def _is_valid(sample) -> bool:
+    """Filter predicate: keep only successfully decoded samples."""
+    return sample is not None
+
+
+def _min_frames_filter(min_frames: int):
+    """Filter predicate: keep only samples with >= min_frames frames."""
+    def _filter(sample):
+        return sample is not None and sample["frames"].shape[0] >= min_frames
+    return _filter
+
+
+def _length_sort_buffer(buffer_size: int = 1000):
+    """
+    Sort samples by frame count within a rolling buffer.
+
+    When the DataLoader forms batches from consecutive samples, this ensures
+    samples with similar frame counts end up in the same batch — dramatically
+    reducing padding waste.  A buffer of 1000 samples (default) gives good
+    grouping while maintaining enough randomization.
+    """
+    def _sort(src):
+        buf = []
+        for sample in src:
+            buf.append(sample)
+            if len(buf) >= buffer_size:
+                buf.sort(key=lambda s: s["frames"].shape[0])
+                yield from buf
+                buf = []
+        if buf:
+            buf.sort(key=lambda s: s["frames"].shape[0])
+            yield from buf
+    return _sort
+
+
+def create_webdataset(
+    shard_pattern: str,
+    tokenizer=None,
+    stage: int = 1,
+    max_frames: int = 64,
+    min_frames: int = 0,
+    shuffle: bool = True,
+    seed: int = 42,
+    epoch: int = 0,
+    num_workers: int = 4,
+    batch_size: Optional[int] = None,
+    shardshuffle: int = 1000,
+    replicate_image_frames: int = 1,
+) -> wds.WebDataset:
+    """
+    Create a webdataset pipeline that streams tar shards.
+
+    Parameters
+    ----------
+    shard_pattern : str
+        Brace-expansion pattern for tar shards, e.g.
+        "/workspace/webvid_frames/{00000..02999}.tar"
+    tokenizer : optional
+        Tokenizer for on-the-fly tokenization of raw captions.
+        If None, samples must have pre-tokenized token_ids in JSON.
+    max_frames : int
+        Maximum number of frames per sample (extras truncated). Default 64,
+        matching SmolVLM2's frame cap.
+    shuffle : bool
+        Whether to shuffle shards and samples.  Disable for deterministic
+        evaluation.
+    seed : int
+        Random seed for reproducible shard + sample shuffling.
+    epoch : int
+        Epoch counter — combined with seed for per-epoch shuffling so that
+        each epoch sees a different order without losing reproducibility.
+    num_workers : int
+        Hint for shard splitting across DataLoader workers.  webdataset
+        handles the splitting internally via its nodesplitter.
+    batch_size : int, optional
+        If provided, the pipeline batches internally (rare — usually the
+        external DataLoader + collate_foveated handles batching).
+    shardshuffle : int
+        Buffer size for shard-level shuffle.  Larger = better randomisation
+        at the cost of memory.  1000 shards ~= 1M samples for our shard
+        size of 1000 samples/shard.
+
+    Returns
+    -------
+    wds.WebDataset
+        An iterable dataset that yields dicts:
+          frames:     [T, 3, 224, 224]
+          input_ids:  [S]
+          loss_mask:  [S]
+          num_frames: int
+    """
+    effective_seed = seed + epoch
+
+    # Resolve shard_pattern: can be a string glob, brace-expansion, or a list of globs.
+    # webdataset handles brace-expansion ({0000..0999}.tar) but NOT shell globs (*.tar).
+    import glob as globmod
+    if isinstance(shard_pattern, list):
+        urls = []
+        for pat in shard_pattern:
+            urls.extend(sorted(globmod.glob(pat)))
+        if not urls:
+            raise ValueError(f"No shards found for patterns: {shard_pattern}")
+    elif '*' in shard_pattern or '?' in shard_pattern:
+        urls = sorted(globmod.glob(shard_pattern))
+        if not urls:
+            raise ValueError(f"No shards found for pattern: {shard_pattern}")
+    else:
+        urls = shard_pattern
+
+    # Build the pipeline.
+    dataset = wds.WebDataset(
+        urls,
+        nodesplitter=wds.split_by_worker,
+        shardshuffle=shardshuffle if shuffle else False,
+        seed=effective_seed if shuffle else None,
+        empty_check=False,  # avoid crash when workers get no valid samples
+        handler=wds.warn_and_continue,  # skip corrupted shards instead of crashing
+    )
+
+    if shuffle:
+        # Shuffle within a buffer of samples (after shard-level shuffle).
+        dataset = dataset.shuffle(size=5000, seed=effective_seed)
+
+    # Decode: we do NOT use wds.decode() because we need custom multi-frame
+    # logic.  Instead we pass raw bytes and decode in _sample_decoder.
+    dataset = dataset.map(_sample_decoder(max_frames, tokenizer=tokenizer, stage=stage,
+                                          replicate_image_frames=replicate_image_frames))
+    dataset = dataset.select(_is_valid)
+
+    if min_frames > 0:
+        dataset = dataset.select(_min_frames_filter(min_frames))
+
+    # Length-sort buffer DISABLED: grouping long videos into same batch causes
+    # (1) GPU OOM cascades (n_real > 700), (2) RAM growth from worker backlog
+    # during OOM retry loops, (3) system OOM crashes.  Random batching with
+    # bucketed padding is safer and only ~10-15% less efficient.
+    # if shuffle:
+    #     dataset = dataset.compose(_length_sort_buffer(128))
+
+    if batch_size is not None:
+        dataset = dataset.batched(batch_size)
+
+    return dataset
+
+
+def create_dpo_webdataset(
+    shard_pattern: str,
+    tokenizer=None,
+    max_frames: int = 64,
+    shuffle: bool = True,
+    seed: int = 42,
+    epoch: int = 0,
+    num_workers: int = 4,
+    batch_size: Optional[int] = None,
+    shardshuffle: int = 1000,
+    replicate_image_frames: int = 1,
+) -> wds.WebDataset:
+    """
+    Create a webdataset pipeline for DPO (preference) data.
+
+    Each sample contains chosen and rejected responses for the same visual input.
+    Returns dicts with:
+      frames:             [T, 3, 224, 224]
+      chosen_input_ids:   [S_c]
+      chosen_loss_mask:   [S_c]
+      rejected_input_ids: [S_r]
+      rejected_loss_mask: [S_r]
+      num_frames:         int
+
+    Parameters
+    ----------
+    shard_pattern : str
+        Brace-expansion pattern for tar shards.
+    tokenizer : optional
+        Tokenizer for on-the-fly tokenization.
+    max_frames : int
+        Maximum number of frames per sample.
+    shuffle : bool
+        Whether to shuffle shards and samples.
+    seed : int
+        Random seed for shuffling.
+    epoch : int
+        Epoch counter for per-epoch shuffling.
+    num_workers : int
+        Hint for shard splitting.
+    batch_size : int, optional
+        If provided, batch internally (rare).
+    shardshuffle : int
+        Buffer size for shard-level shuffle.
+    replicate_image_frames : int
+        Replicate single-frame images to N frames.
+    """
+    effective_seed = seed + epoch
+
+    import glob as globmod
+    if isinstance(shard_pattern, list):
+        urls = []
+        for pat in shard_pattern:
+            urls.extend(sorted(globmod.glob(pat)))
+        if not urls:
+            raise ValueError(f"No shards found for patterns: {shard_pattern}")
+    elif '*' in shard_pattern or '?' in shard_pattern:
+        urls = sorted(globmod.glob(shard_pattern))
+        if not urls:
+            raise ValueError(f"No shards found for pattern: {shard_pattern}")
+    else:
+        urls = shard_pattern
+
+    dataset = wds.WebDataset(
+        urls,
+        nodesplitter=wds.split_by_worker,
+        shardshuffle=shardshuffle if shuffle else False,
+        seed=effective_seed if shuffle else None,
+        empty_check=False,
+        handler=wds.warn_and_continue,
+    )
+
+    if shuffle:
+        dataset = dataset.shuffle(size=5000, seed=effective_seed)
+
+    dataset = dataset.map(_dpo_sample_decoder(max_frames, tokenizer=tokenizer,
+                                               replicate_image_frames=replicate_image_frames))
+    dataset = dataset.select(_is_valid)
+
+    if batch_size is not None:
+        dataset = dataset.batched(batch_size)
+
+    return dataset
+
+
+def make_dynamic_dataloader(
+    shard_pattern: str,
+    max_total_frames: int = 512,
+    max_batch_size: int = 64,
+    max_frames: int = 64,
+    min_frames: int = 0,
+    shuffle: bool = True,
+    seed: int = 42,
+    epoch: int = 0,
+    num_workers: int = 4,
+    pin_memory: bool = True,
+    prefetch_factor: int = 4,
+    tokenizer=None,
+    stage: int = 1,
+    replicate_image_frames: int = 1,
+) -> torch.utils.data.DataLoader:
+    """
+    Dynamic-batch dataloader: batch size varies per batch based on total
+    frame count.  Short-video batches get more samples; long-video batches
+    get fewer.  Total frames per batch is capped at max_total_frames.
+
+    This keeps GPU work roughly constant across batches and eliminates the
+    pathological case where one T=64 sample forces the entire batch to pad
+    to 64 frames.
+    """
+    from collate import token_budget_batcher
+
+    dataset = create_webdataset(
+        shard_pattern=shard_pattern,
+        tokenizer=tokenizer,
+        stage=stage,
+        max_frames=max_frames,
+        min_frames=min_frames,
+        shuffle=shuffle,
+        seed=seed,
+        epoch=epoch,
+        num_workers=num_workers,
+        replicate_image_frames=replicate_image_frames,
+    )
+
+    # The batcher forms variable-size batches and collates them internally.
+    # length_bucket=True sorts by total length within a buffer to reduce padding waste.
+    dataset = dataset.compose(token_budget_batcher(
+        max_total_frames, max_batch_size,
+        length_bucket=True, bucket_buffer=max_batch_size * 4,
+    ))
+
+    # batch_size=None: each dataset item is already a collated batch dict
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=None,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None,
+        persistent_workers=num_workers > 0,
+    )
+    return loader
+
+
+def make_dataloader(
+    shard_pattern: str,
+    batch_size: int,
+    max_frames: int = 64,
+    min_frames: int = 0,
+    shuffle: bool = True,
+    seed: int = 42,
+    epoch: int = 0,
+    num_workers: int = 4,
+    collate_fn=None,
+    pin_memory: bool = True,
+    prefetch_factor: int = 4,
+    tokenizer=None,
+    stage: int = 1,
+    replicate_image_frames: int = 1,
+) -> torch.utils.data.DataLoader:
+    """
+    Convenience wrapper: creates the webdataset pipeline and wraps it in a
+    standard PyTorch DataLoader with the given collate function.
+
+    If collate_fn is None, use collate.collate_foveated.
+    """
+    if collate_fn is None:
+        from collate import collate_foveated
+        collate_fn = collate_foveated
+
+    dataset = create_webdataset(
+        shard_pattern=shard_pattern,
+        tokenizer=tokenizer,
+        stage=stage,
+        max_frames=max_frames,
+        min_frames=min_frames,
+        shuffle=shuffle,
+        seed=seed,
+        epoch=epoch,
+        num_workers=num_workers,
+        replicate_image_frames=replicate_image_frames,
+    )
+
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        collate_fn=collate_fn,
+        pin_memory=pin_memory,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None,
+        persistent_workers=num_workers > 0,
+    )
+    return loader