Make SFT metadata filters configurable for short-task datasets

cosmos_framework/data/generator/local_datasets/sft_dataset.py

@@ -34,6 +34,8 @@
 from cosmos_framework.utils.lazy_config import instantiate as lazy_instantiate
 
 _MAX_CAPTION_TOKENS = 1024
+_MAX_VIDEO_DURATION_S = 61.0
+_MIN_WINDOW_FRAMES = 61
 _DURATION_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS."
 _RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution."
 
@@ -478,19 +480,24 @@ def _flatten_metadata_by_window(metadata_list: list[dict]) -> list[dict]:
 def _load_sft_metadata_from_s3(
     s3_client,
     jsonl_url: str,
-    min_frames: int,
+    min_frames: int | None,
+    max_duration_s: float | None = _MAX_VIDEO_DURATION_S,
     uuid_prefix: str = "",
     min_short_edge: int = 0,
 ) -> list[dict]:
     """Load SFT metadata from a single JSONL file on S3.
 
     Returns one entry per video.  Each entry keeps only the windows whose frame
-    span is at least *min_frames*; videos with no qualifying windows are dropped.
+    span is at least *min_frames* when that filter is enabled; videos with no
+    qualifying windows are dropped.
 
     Args:
         s3_client: Boto3 S3 client
         jsonl_url: S3 URL to the JSONL metadata file
-        min_frames: Minimum number of frames required per window
+        min_frames: Minimum number of frames required per window.  None disables
+            the metadata window-length filter.
+        max_duration_s: Drop videos whose duration exceeds this value.  None
+            disables the metadata duration filter.
         uuid_prefix: Prefix prepended to each uuid for disambiguation when
             multiple JSONL files are loaded
         min_short_edge: Drop videos whose shortest spatial edge (min of width,
@@ -520,8 +527,8 @@ def _load_sft_metadata_from_s3(
             num_raw_records += 1
             record = json.loads(line.decode("utf-8"))
             uuid = f"{uuid_prefix}{record['uuid']}" if uuid_prefix else record["uuid"]
-            if record["duration"] > 61.0:
-                print(f"Skipping video with too long duration: {uuid}, {record['duration']} > 61.0")
+            if max_duration_s is not None and record["duration"] > max_duration_s:
+                print(f"Skipping video with too long duration: {uuid}, {record['duration']} > {max_duration_s}")
                 num_filtered_duration += 1
                 continue
             if min_short_edge > 0 and min(record["width"], record["height"]) < min_short_edge:
@@ -536,7 +543,7 @@ def _load_sft_metadata_from_s3(
             for window in windows:
                 num_raw_windows += 1
                 frames_in_window = window["end_frame"] - window["start_frame"] + 1
-                if frames_in_window < min_frames:
+                if min_frames is not None and frames_in_window < min_frames:
                     num_filtered_windows += 1
                 else:
                     kept_windows.append(window)
@@ -563,12 +570,14 @@ def _load_sft_metadata_from_s3(
                 }
             )
 
+    duration_filter = f"> {max_duration_s}s" if max_duration_s is not None else "disabled"
+    window_filter = f"< {min_frames}f" if min_frames is not None else "disabled"
     log.info(
         f"Finished decoding SFT metadata from {jsonl_url}. "
         f"Records: {num_raw_records}, "
-        f"Duration > 61s: {num_filtered_duration}, "
+        f"Duration filter {duration_filter}: {num_filtered_duration}, "
         f"Short edge < {min_short_edge}: {num_filtered_short_edge}, "
-        f"Windows: {num_raw_windows}, Windows < {min_frames}f: {num_filtered_windows}, "
+        f"Windows: {num_raw_windows}, Window filter {window_filter}: {num_filtered_windows}, "
         f"Videos kept: {len(metadata_list)}"
     )
     return metadata_list
@@ -589,6 +598,8 @@ def get_sft_dataset(
     cfg_dropout_keep_metadata: bool = False,
     sample_by_window: bool = False,
     min_short_edge: int = 0,
+    min_window_frames: int | None = _MIN_WINDOW_FRAMES,
+    max_duration_s: float | None = _MAX_VIDEO_DURATION_S,
     caption_suffix: str = "",
     conditioning_fps: float = 24,
     conditioning_fps_noise_std: float = 0.0,
@@ -629,6 +640,12 @@ def get_sft_dataset(
             window is chosen on every access.
         min_short_edge: Drop videos whose shortest spatial edge (min of width,
             height) is below this value.  0 (default) disables the filter.
+        min_window_frames: Metadata pre-filter for t2w_window length.  Default
+            61 matches the historical loader behavior.  Set to None for
+            short-task datasets whose valid samples may be shorter.
+        max_duration_s: Metadata pre-filter for video duration.  Default 61.0
+            matches the historical loader behavior.  Set to None for short-task
+            datasets.
         caption_suffix: Text appended to every caption before the
             duration/FPS/resolution templates, e.g.
             ``"Overall, the video is of poor quality."``.  Empty string
@@ -671,7 +688,8 @@ def get_sft_dataset(
             _load_sft_metadata_from_s3(
                 s3_client,
                 jsonl_url,
-                min_frames=61,
+                min_frames=min_window_frames,
+                max_duration_s=max_duration_s,
                 uuid_prefix=prefix,
                 min_short_edge=min_short_edge,
             )

cosmos_framework/data/generator/local_datasets/sft_dataset_caption_test.py

@@ -4,7 +4,7 @@
 
 import json
 
-from cosmos_framework.data.generator.local_datasets.sft_dataset import _select_caption
+from cosmos_framework.data.generator.local_datasets.sft_dataset import _load_sft_metadata_from_s3, _select_caption
 from cosmos_framework.inference.structured_caption import CAPTION_JSON_KEY
 
 
@@ -53,3 +53,31 @@ def test_weighted_caption_types_fallback():
 
 def test_no_known_caption_key_returns_none():
     assert _select_caption({"start_frame": 0, "end_frame": 84}) is None
+
+
+def test_load_sft_metadata_can_disable_duration_and_window_prefilters(tmp_path):
+    jsonl = tmp_path / "short_tasks.jsonl"
+    record = {
+        "uuid": "short-task",
+        "duration": 120.0,
+        "width": 256,
+        "height": 256,
+        "vision_path": "clips/short-task.mp4",
+        "t2w_windows": [
+            {
+                "start_frame": 0,
+                "end_frame": 9,
+                "temporal_interval": 1,
+                "caption": "short task",
+            }
+        ],
+    }
+    jsonl.write_text(json.dumps(record) + "\n")
+
+    default_filtered = _load_sft_metadata_from_s3(None, str(jsonl), min_frames=61)
+    assert default_filtered == []
+
+    kept = _load_sft_metadata_from_s3(None, str(jsonl), min_frames=None, max_duration_s=None)
+    assert len(kept) == 1
+    assert kept[0]["uuid"] == "short-task"
+    assert kept[0]["t2w_windows"][0]["caption"] == "short task"

cosmos_framework/data/generator/local_datasets/transfer_sft_dataset.py

@@ -66,6 +66,8 @@
 
 _ON_THE_FLY_CONTROLS = frozenset({"edge", "blur"})
 _PRECOMPUTED_CONTROLS = frozenset({"depth", "seg"})
+_MAX_VIDEO_DURATION_S = 61.0
+_MIN_WINDOW_FRAMES = 61
 
 # Canny threshold presets matching AddControlInputEdge.
 _EDGE_THRESHOLDS = [
@@ -235,7 +237,8 @@ def process_one_sample(self, metadata: dict) -> dict | None:
 def _load_transfer_metadata(
     s3_client: Any,
     jsonl_url: str,
-    min_frames: int,
+    min_frames: int | None,
+    max_duration_s: float | None = _MAX_VIDEO_DURATION_S,
     uuid_prefix: str = "",
     min_short_edge: int = 0,
     control_type: str = "edge",
@@ -261,15 +264,17 @@ def _load_transfer_metadata(
             record = json.loads(line.decode("utf-8"))
             uuid = f"{uuid_prefix}{record['uuid']}" if uuid_prefix else record["uuid"]
 
-            if record["duration"] > 61.0:
+            if max_duration_s is not None and record["duration"] > max_duration_s:
                 continue
             if min_short_edge > 0 and min(record["width"], record["height"]) < min_short_edge:
                 continue
             if control_type in _PRECOMPUTED_CONTROLS and not record.get("control_path"):
                 continue
 
             kept_windows = [
-                w for w in (record.get("t2w_windows") or []) if w["end_frame"] - w["start_frame"] + 1 >= min_frames
+                w
+                for w in (record.get("t2w_windows") or [])
+                if min_frames is None or w["end_frame"] - w["start_frame"] + 1 >= min_frames
             ]
             if not kept_windows:
                 continue
@@ -315,6 +320,8 @@ def get_transfer_sft_dataset(
     append_duration_fps_timestamps: bool = True,
     append_resolution_info: bool = True,
     min_short_edge: int = 0,
+    min_window_frames: int | None = _MIN_WINDOW_FRAMES,
+    max_duration_s: float | None = _MAX_VIDEO_DURATION_S,
     conditioning_config: dict[int, float] | None = None,
     temporal_compression_factor: int = 4,
     **kwargs,
@@ -338,6 +345,10 @@ def get_transfer_sft_dataset(
         append_duration_fps_timestamps: Append duration/FPS text to captions.
         append_resolution_info: Append resolution text to captions.
         min_short_edge: Drop videos whose shortest edge is below this value.
+        min_window_frames: Metadata pre-filter for t2w_window length.  Default
+            61 matches historical behavior.  Set to None for short-task datasets.
+        max_duration_s: Metadata pre-filter for video duration.  Default 61.0
+            matches historical behavior.  Set to None for short-task datasets.
         conditioning_config: I2V conditioning frame distribution, e.g.
             ``{0: 0.7, 1: 0.2, 2: 0.1}`` (70% unconditioned).
         temporal_compression_factor: VAE temporal compression factor (default 4).
@@ -363,7 +374,8 @@ def get_transfer_sft_dataset(
             _load_transfer_metadata(
                 s3_client,
                 jsonl_url,
-                min_frames=61,
+                min_frames=min_window_frames,
+                max_duration_s=max_duration_s,
                 uuid_prefix=prefix,
                 min_short_edge=min_short_edge,
                 control_type=control_type,

cosmos_framework/data/generator/local_datasets/transfer_sft_dataset_test.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+"""Tests for transfer SFT metadata filtering."""
+
+import json
+
+from cosmos_framework.data.generator.local_datasets.transfer_sft_dataset import _load_transfer_metadata
+
+
+def test_load_transfer_metadata_can_disable_duration_and_window_prefilters(tmp_path):
+    jsonl = tmp_path / "short_transfer_tasks.jsonl"
+    record = {
+        "uuid": "short-task",
+        "duration": 120.0,
+        "width": 256,
+        "height": 256,
+        "vision_path": "clips/short-task.mp4",
+        "control_type": "edge",
+        "t2w_windows": [
+            {
+                "start_frame": 0,
+                "end_frame": 9,
+                "temporal_interval": 1,
+                "caption": "short task",
+            }
+        ],
+    }
+    jsonl.write_text(json.dumps(record) + "\n")
+
+    assert _load_transfer_metadata(None, str(jsonl), min_frames=61) == []
+
+    kept = _load_transfer_metadata(None, str(jsonl), min_frames=None, max_duration_s=None)
+    assert len(kept) == 1
+    assert kept[0]["uuid"] == "short-task"
+    assert kept[0]["t2w_windows"][0]["caption"] == "short task"

cosmos_framework/scripts/captions_to_sft_jsonl.py

@@ -16,13 +16,15 @@
 If a clip has no ``caption.json`` (e.g. produced by an older captioner), the row is
 written dense-only, exactly as before.
 
-Filters mirror what training actually consumes so dataset counts match:
+By default, filters mirror what training actually consumes so dataset counts match:
 
 * clips longer than 61 s are dropped (matches the loader's hard cap);
 * windows shorter than ``max(61, num_video_frames)`` frames are dropped.  Pass
   ``--num-video-frames`` to match your training recipe.  The default (-1) applies
   only the loader's metadata minimum of 61 frames, matching the example recipe
   (``num_video_frames=-1``) so short example clips (~85 frames) are kept.
+* for short-task datasets, pass ``--max-duration-s None --min-window-frames None``
+  and use ``--num-video-frames -1`` so metadata conversion does not discard short clips.
 
 A sibling ``<output>.summary.json`` records kept/dropped counts per reason.
 
@@ -60,8 +62,8 @@
 from cosmos_framework.inference.structured_caption import CAPTION_JSON_KEY
 from cosmos_framework.scripts.video_metadata import probe_video_metadata
 
-_MAX_DURATION = 61.0  # seconds; matches hard-coded limit in sft_dataset.py
-_MIN_FRAMES = 61  # matches the metadata min_frames=61 in get_sft_dataset()
+_MAX_DURATION = 61.0  # seconds; matches default in sft_dataset.py
+_MIN_FRAMES = 61  # matches default metadata min_frames in get_sft_dataset()
 _VIDEO_EXTENSIONS = (".mp4", ".mov", ".avi", ".mkv", ".webm")
 
 
@@ -96,10 +98,28 @@ def main(
         int,
         tyro.conf.arg(
             help="Decoded frames per window in your training recipe; windows shorter than "
-            "max(61, this) are dropped. -1 (default) applies only the 61-frame metadata "
-            "minimum, matching the example recipe."
+            "max(min_window_frames, this) are dropped when min_window_frames is set. "
+            "-1 (default) applies only the metadata minimum."
         ),
     ] = -1,
+    min_window_frames: Annotated[
+        int | None,
+        tyro.conf.arg(
+            help=(
+                f"Drop windows shorter than this metadata floor. Default {_MIN_FRAMES} matches "
+                "get_sft_dataset(). Set to None for short-task datasets."
+            ),
+        ),
+    ] = _MIN_FRAMES,
+    max_duration_s: Annotated[
+        float | None,
+        tyro.conf.arg(
+            help=(
+                f"Drop clips longer than this. Default {_MAX_DURATION} matches get_sft_dataset(). "
+                "Set to None for short-task datasets."
+            ),
+        ),
+    ] = _MAX_DURATION,
     min_short_edge: Annotated[
         int, tyro.conf.arg(help="Drop clips whose shortest spatial edge is below this value. 0 disables.")
     ] = 0,
@@ -115,7 +135,12 @@ def main(
         print(f"No caption.txt / caption.json files found under {captions_dir}", file=sys.stderr)
         sys.exit(1)
 
-    effective_min_frames = _MIN_FRAMES if num_video_frames <= 0 else max(_MIN_FRAMES, num_video_frames)
+    if min_window_frames is None:
+        effective_min_frames = None if num_video_frames <= 0 else num_video_frames
+    elif num_video_frames <= 0:
+        effective_min_frames = min_window_frames
+    else:
+        effective_min_frames = max(min_window_frames, num_video_frames)
 
     records = []
     drops: Counter[str] = Counter()
@@ -153,12 +178,12 @@ def main(
             drops["ffprobe_error"] += 1
             continue
 
-        if meta["duration"] > _MAX_DURATION:
-            print(f"  SKIP {name}: duration {meta['duration']:.1f}s > {_MAX_DURATION}s")
+        if max_duration_s is not None and meta["duration"] > max_duration_s:
+            print(f"  SKIP {name}: duration {meta['duration']:.1f}s > {max_duration_s}s")
             drops["duration_too_long"] += 1
             continue
 
-        if meta["total_frames"] < effective_min_frames:
+        if effective_min_frames is not None and meta["total_frames"] < effective_min_frames:
             print(f"  SKIP {name}: only {meta['total_frames']} frames < {effective_min_frames}")
             drops["too_few_frames"] += 1
             continue
@@ -204,10 +229,11 @@ def main(
         "records_dropped": sum(drops.values()),
         "drops_by_reason": dict(drops),
         "filters": {
-            "max_duration_s": _MAX_DURATION,
+            "max_duration_s": max_duration_s,
             "min_window_frames": effective_min_frames,
             "min_short_edge": min_short_edge,
             "num_video_frames": num_video_frames,
+            "metadata_min_window_frames": min_window_frames,
         },
     }
     summary_path = output.with_suffix(output.suffix + ".summary.json")

cosmos_framework/scripts/captions_to_sft_jsonl_test.py

@@ -96,6 +96,29 @@ def fake_probe(path):
     assert summary["drops_by_reason"]["too_few_frames"] == 1
 
 
+def test_can_disable_duration_and_min_frame_filters_for_short_tasks(dirs, tmp_path, monkeypatch):
+    captions_dir, videos_dir = dirs
+    _make_clip(captions_dir, videos_dir, "short", caption_json={"x": 1})
+    monkeypatch.setattr(mod, "probe_video_metadata", lambda p: _meta(duration=0.5, total_frames=12))
+
+    out = tmp_path / "ds.jsonl"
+    mod.main(
+        captions_dir=captions_dir,
+        videos_dir=videos_dir,
+        output=out,
+        min_window_frames=None,
+        max_duration_s=None,
+        num_video_frames=-1,
+    )
+
+    rows = _read_jsonl(out)
+    assert [r["uuid"] for r in rows] == ["short"]
+    summary = json.loads((tmp_path / "ds.jsonl.summary.json").read_text())
+    assert summary["filters"]["max_duration_s"] is None
+    assert summary["filters"]["min_window_frames"] is None
+    assert summary["drops_by_reason"] == {}
+
+
 def test_num_video_frames_filter_drops_85_frame_clip(dirs, tmp_path, monkeypatch):
     captions_dir, videos_dir = dirs
     _make_clip(captions_dir, videos_dir, "ep0", caption_json={"x": 1})