From e1ad7f2460d8a019d4abecceb72fedc606de6c3a Mon Sep 17 00:00:00 2001 From: Hans Yang Date: Thu, 16 Jul 2026 06:53:30 -0700 Subject: [PATCH] feat(data): make SFT metadata filters configurable --- .../generator/local_datasets/sft_dataset.py | 36 +++++++++++---- .../sft_dataset_caption_test.py | 30 +++++++++++- .../local_datasets/transfer_sft_dataset.py | 20 ++++++-- .../transfer_sft_dataset_test.py | 35 ++++++++++++++ .../scripts/captions_to_sft_jsonl.py | 46 +++++++++++++++---- .../scripts/captions_to_sft_jsonl_test.py | 23 ++++++++++ .../scripts/curator_to_sft_jsonl.py | 31 ++++++++----- .../scripts/curator_to_sft_jsonl_test.py | 30 ++++++++++++ 8 files changed, 215 insertions(+), 36 deletions(-) create mode 100644 cosmos_framework/data/generator/local_datasets/transfer_sft_dataset_test.py diff --git a/cosmos_framework/data/generator/local_datasets/sft_dataset.py b/cosmos_framework/data/generator/local_datasets/sft_dataset.py index 1a9fcdbc..11f6710f 100644 --- a/cosmos_framework/data/generator/local_datasets/sft_dataset.py +++ b/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, ) diff --git a/cosmos_framework/data/generator/local_datasets/sft_dataset_caption_test.py b/cosmos_framework/data/generator/local_datasets/sft_dataset_caption_test.py index 04560883..500dd6fb 100644 --- a/cosmos_framework/data/generator/local_datasets/sft_dataset_caption_test.py +++ b/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" diff --git a/cosmos_framework/data/generator/local_datasets/transfer_sft_dataset.py b/cosmos_framework/data/generator/local_datasets/transfer_sft_dataset.py index 20f926c7..9f3b23ea 100644 --- a/cosmos_framework/data/generator/local_datasets/transfer_sft_dataset.py +++ b/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,7 +264,7 @@ 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 @@ -269,7 +272,9 @@ def _load_transfer_metadata( 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, diff --git a/cosmos_framework/data/generator/local_datasets/transfer_sft_dataset_test.py b/cosmos_framework/data/generator/local_datasets/transfer_sft_dataset_test.py new file mode 100644 index 00000000..af652252 --- /dev/null +++ b/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" diff --git a/cosmos_framework/scripts/captions_to_sft_jsonl.py b/cosmos_framework/scripts/captions_to_sft_jsonl.py index 1d2e9b9d..452d4602 100644 --- a/cosmos_framework/scripts/captions_to_sft_jsonl.py +++ b/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 ``.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") diff --git a/cosmos_framework/scripts/captions_to_sft_jsonl_test.py b/cosmos_framework/scripts/captions_to_sft_jsonl_test.py index b00f5307..4d0802c0 100644 --- a/cosmos_framework/scripts/captions_to_sft_jsonl_test.py +++ b/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}) diff --git a/cosmos_framework/scripts/curator_to_sft_jsonl.py b/cosmos_framework/scripts/curator_to_sft_jsonl.py index 0ae07573..dc5d1bab 100644 --- a/cosmos_framework/scripts/curator_to_sft_jsonl.py +++ b/cosmos_framework/scripts/curator_to_sft_jsonl.py @@ -10,8 +10,8 @@ Curator writes a richer schema per clip at ``/metas_jsonl/v0/*.jsonl``. This script renames and trims those -rows into the loader's format, applies the same hard filters the loader applies -silently at train time (so dataset counts match), and writes a sidecar +rows into the loader's format, applies the same default filters the loader applies +at train time (so dataset counts match), and writes a sidecar ``.summary.json`` with per-reason drop counts. Transfer (control-conditioned) training @@ -58,8 +58,9 @@ import tyro -# Hard filters mirror sft_dataset.py defaults so the converter drops the same -# rows the loader would silently drop at train time. +# Default filters mirror sft_dataset.py defaults so the converter drops the same +# rows the loader would drop at train time. Pass None through the CLI/config to +# disable a filter for short-task datasets. MAX_VIDEO_DURATION_S: float = 61.0 MIN_WINDOW_FRAMES: int = 61 DEFAULT_TEMPORAL_INTERVAL: int = 1 @@ -213,8 +214,8 @@ def _build_sft_row( caption_model: str | None, enhanced_caption_model: str | None, min_short_edge: int, - min_window_frames: int, - max_duration_s: float, + min_window_frames: int | None, + max_duration_s: float | None, temporal_interval: int, control_type: str | None = None, control_path_root: Path | None = None, @@ -240,7 +241,7 @@ def _build_sft_row( return None, "missing_identity" if width is None or height is None or num_frames is None or framerate is None or duration_s is None: return None, "missing_clip_metadata" - if duration_s > max_duration_s: + if max_duration_s is not None and duration_s > max_duration_s: return None, "duration_too_long" if min_short_edge > 0 and min(int(width), int(height)) < min_short_edge: return None, "short_edge_too_small" @@ -260,7 +261,7 @@ def _build_sft_row( if not isinstance(start_frame, int) or not isinstance(end_frame, int): continue frames_in_window = end_frame - start_frame + 1 - if frames_in_window < min_window_frames: + if min_window_frames is not None and frames_in_window < min_window_frames: continue caption_text = _resolve_window_caption( window, @@ -323,15 +324,21 @@ def main( # noqa: PLR0913 tyro.conf.arg(help="Drop clips whose shortest spatial edge is below this value. 0 disables."), ] = 0, min_window_frames: Annotated[ - int, + int | None, tyro.conf.arg( - help=f"Drop windows shorter than this. Default {MIN_WINDOW_FRAMES} matches sft_dataset.py.", + help=( + f"Drop windows shorter than this. Default {MIN_WINDOW_FRAMES} matches sft_dataset.py. " + "Set to None for short-task datasets." + ), ), ] = MIN_WINDOW_FRAMES, max_duration_s: Annotated[ - float, + float | None, tyro.conf.arg( - help=f"Drop clips longer than this. Default {MAX_VIDEO_DURATION_S} matches sft_dataset.py.", + help=( + f"Drop clips longer than this. Default {MAX_VIDEO_DURATION_S} matches sft_dataset.py. " + "Set to None for short-task datasets." + ), ), ] = MAX_VIDEO_DURATION_S, temporal_interval: Annotated[ diff --git a/cosmos_framework/scripts/curator_to_sft_jsonl_test.py b/cosmos_framework/scripts/curator_to_sft_jsonl_test.py index 0c0c4fc6..d69a29f5 100644 --- a/cosmos_framework/scripts/curator_to_sft_jsonl_test.py +++ b/cosmos_framework/scripts/curator_to_sft_jsonl_test.py @@ -135,6 +135,18 @@ def test_build_sft_row_drops_clip_longer_than_max_duration() -> None: assert reason == "duration_too_long" +@pytest.mark.L0 +def test_build_sft_row_can_disable_max_duration_filter() -> None: + # Short-task configs can disable the metadata duration cap without + # changing the default behavior for long-video SFT datasets. + record = _make_record(num_frames=120, framerate=1.0) + kwargs = _default_kwargs() | {"max_duration_s": None} + row, reason = _build_sft_row(record, **kwargs) + assert reason is None + assert row is not None + assert row["duration"] == pytest.approx(120.0) + + @pytest.mark.L0 def test_build_sft_row_keeps_clip_at_exactly_max_duration() -> None: # 61 frames at 1 fps = 61.0 s, must pass because loader uses strict >. @@ -180,6 +192,24 @@ def test_build_sft_row_drops_when_all_windows_too_short() -> None: assert reason == "no_valid_window" +@pytest.mark.L0 +def test_build_sft_row_can_disable_min_window_frames_filter() -> None: + record = _make_record( + windows=[ + { + "start_frame": _SHORT_WINDOW[0], + "end_frame": _SHORT_WINDOW[1], + "qwen_caption": "short window", + } + ], + ) + kwargs = _default_kwargs() | {"min_window_frames": None} + row, reason = _build_sft_row(record, **kwargs) + assert reason is None + assert row is not None + assert row["t2w_windows"][0]["caption"] == "short window" + + @pytest.mark.L0 def test_build_sft_row_drops_when_no_window_has_caption() -> None: record = _make_record(