From bc48e9403614abb7ffd4ade9874e18b1f6ef30b1 Mon Sep 17 00:00:00 2001 From: rwagwani Date: Thu, 29 Oct 2026 13:32:59 +0000 Subject: [PATCH] fix(inference): name the cause in Transfer rejection messages Three Transfer rejections were correct but did not say what the request did wrong: - The direct causal Transfer check listed all four rules ("exactly one control hint, num_first_chunk_conditional_frames=0, control_guidance=1.0, and share_vision_temporal_positions=true"). It now names only the rules this request breaks, with the values passed. - On the 480-only SF-DMD checkpoint, 256/704/768 were told "480 and 720 tiers only", so a user tried 720 and only then learned "480 only". The direct-resize check now raises the checkpoint's own limit, using the same predicate as its resolution contract, and both messages include the rejected value. - The direct_resize aspect error ("causes N% aspect distortion") did not mention a forced aspect_ratio. It now names it, says direct_resize does not crop, and names the bucket that would fit, if any. The Transfer cookbook's multi_control.json stated fps 30 / num_frames 121 for a 16 fps, 93-frame video. Neither value is read for a video input, so the spec now states the real values without changing its output, and the README field reference says fps/num_frames should match the video. Signed-off-by: rwagwani Co-Authored-By: Claude Opus 5.5 --- .../cosmos3/generator/transfer/README.md | 6 +- .../transfer/specs/multi_control.json | 4 +- cosmos_framework/inference/args.py | 76 +++++++++++++------ .../inference/sim_transfer_args_test.py | 44 +++++++++++ .../inference/transfer_cookbook_specs_test.py | 37 +++++++++ .../inference/transfer_direct_resize_test.py | 36 ++++++++- .../inference/transfer_preprocessing.py | 21 ++++- .../inference/transfer_preprocessing_test.py | 26 +++++++ 8 files changed, 214 insertions(+), 36 deletions(-) create mode 100644 cosmos_framework/inference/transfer_cookbook_specs_test.py diff --git a/.agents/skills/cookbooks/cosmos3/generator/transfer/README.md b/.agents/skills/cookbooks/cosmos3/generator/transfer/README.md index ba7399c7d..f3179ee20 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/transfer/README.md +++ b/.agents/skills/cookbooks/cosmos3/generator/transfer/README.md @@ -47,7 +47,7 @@ for the negative caption. | Depth | `assets/depth/` | `control_depth.mp4` + `prompt.json` | 121 frames @ 30 FPS | | Segmentation | `assets/seg/` | `control_seg.mp4` + `prompt.json` | 121 frames @ 30 FPS | | World scenario (WSM) | `assets/wsm/` | `control_wsm.mp4` + `prompt.json` | 100 frames @ 10 FPS | -| Multi-control | `assets/multi_control/` | `vision_path` + multiple hints (Framework example) | 121 frames @ 30 FPS | +| Multi-control | `assets/multi_control/` | `vision_path` + multiple hints (Framework example) | 93 frames @ 16 FPS | Transfer inference is selected automatically when any hint key is present in the Framework spec or in vLLM-Omni `extra_params`. @@ -398,9 +398,9 @@ Key fields: - **`aspect_ratio`** — aspect ratio of the control video; together with `resolution` determines the spatial dimensions (e.g. `720` + `16,9` → 1280 × 720). -- **`fps`** — model conditioning signal and playback rate of the saved output video. Should match the native fps of the control video. +- **`fps`** — frame rate of the input video. Keep it equal to the native fps of the control/input video. Cosmos Framework Transfer takes the output frame rate from the video itself and uses `fps` only for image inputs; the Diffusers notebook reads it from the spec. -- **`num_frames`** — number of video frames. +- **`num_frames`** — number of video frames. Keep it equal to the control/input video's frame count. Cosmos Framework Transfer takes the output length from the video (capped by `max_frames`); the Diffusers notebook reads it from the spec. ### Cookbook entrypoints diff --git a/.agents/skills/cookbooks/cosmos3/generator/transfer/specs/multi_control.json b/.agents/skills/cookbooks/cosmos3/generator/transfer/specs/multi_control.json index d022effc8..68d8aa9b0 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/transfer/specs/multi_control.json +++ b/.agents/skills/cookbooks/cosmos3/generator/transfer/specs/multi_control.json @@ -3,8 +3,8 @@ "model_mode": "video2video", "resolution": "720", "aspect_ratio": "16,9", - "num_frames": 121, - "fps": 30, + "num_frames": 93, + "fps": 16, "shift": 10.0, "num_steps": 50, "seed": 2026, diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 0b1f78c98..c4e432439 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -1559,16 +1559,26 @@ def _build_transfer_data( "checkpoint with a fixed sampler or Stage 3 SF-DMD checkpoint with a fixed sampler, " "with one latent frame per chunk" ) - if ( - len(hints) != 1 - or self.num_first_chunk_conditional_frames != 0 - or getattr(self, "control_guidance", None) != 1.0 - or not self.share_vision_temporal_positions - ): + # Name only the requirements this request breaks, with the values it passed. + culprits = [] + if len(hints) != 1: + culprits.append(f"{len(hints)} control hints {sorted(hints)} (must be exactly 1)") + if self.num_first_chunk_conditional_frames != 0: + culprits.append( + f"num_first_chunk_conditional_frames={self.num_first_chunk_conditional_frames} (must be 0)" + ) + control_guidance = getattr(self, "control_guidance", None) + if control_guidance != 1.0: + culprits.append( + f"control_guidance={control_guidance} (must be 1.0: the causal sampler has no " + "control-guidance pass)" + ) + if not self.share_vision_temporal_positions: + culprits.append("share_vision_temporal_positions=false (must be true to match training)") + if culprits: raise ValueError( - "Direct causal Transfer requires exactly one control hint, " - "num_first_chunk_conditional_frames=0, control_guidance=1.0, " - "and share_vision_temporal_positions=true" + "Direct causal Transfer (transfer_resize_mode='direct_resize' on a causal checkpoint) " + "cannot run this request:\n - " + "\n - ".join(culprits) ) if self.multiview is not None or any( key not in {TransferHintKey.EDGE, TransferHintKey.BLUR, TransferHintKey.DEPTH, TransferHintKey.SEG} @@ -1576,7 +1586,18 @@ def _build_transfer_data( ): raise ValueError("Direct Transfer resizing supports single-view edge, blur, depth, and seg only") if self.resolution not in {"480", "720"}: - raise ValueError("Direct Transfer resizing supports the teacher's trained 480 and 720 tiers only") + # A 480-only checkpoint would reject 720 later too; name its real limit instead of the teacher's tiers. + if ( + self._is_sf_dmd_transfer_model(model_config) + and str(getattr(model_config, "resolution", "")) == "480" + ): + raise ValueError( + f"This SF-DMD Transfer checkpoint supports only resolution='480', got resolution='{self.resolution}'." + ) + raise ValueError( + "Direct Transfer resizing supports the teacher's trained 480 and 720 tiers only, " + f"got resolution='{self.resolution}'." + ) for hint_key, config in hints.items(): for field, default in self._TRANSFER_HINT_DEFAULTS.get(hint_key, {}).items(): @@ -1607,6 +1628,23 @@ def _build_transfer_data( if user_fields is None or field not in user_fields: setattr(self, field, value) + @staticmethod + def _is_sf_dmd_transfer_model(model_config: "OmniMoTModelConfig") -> bool: + """Whether the loaded model is the causal SF-DMD Transfer student whose contract applies.""" + video_temporal_causal = bool(getattr(model_config, "video_temporal_causal", False)) + is_self_forcing = getattr(model_config, "simulation_mode", None) == "self_forcing" + is_distilled = getattr(model_config, "fixed_step_sampler_config", None) is not None + if not video_temporal_causal or not (is_self_forcing or is_distilled): + return False + if int(getattr(model_config, "teacher_forcing_frames_per_chunk", 1) or 1) != 1: + return False + replay_policy = getattr(model_config, "teacher_forcing_replay_policy", None) + return ( + getattr(replay_policy, "control_visibility", None) == "causal" + and bool(getattr(replay_policy, "controls_read_strict_past_clean_rgb", False)) + and getattr(replay_policy, "clean_pass_causality", None) == "frame" + ) + def _apply_sf_dmd_transfer_inference_contract( self, *, @@ -1615,26 +1653,16 @@ def _apply_sf_dmd_transfer_inference_contract( ) -> None: """Apply the resolution contract recorded by the SF-DMD Transfer checkpoint.""" self = cast("SampleDataOverrides", self) - video_temporal_causal = bool(getattr(model_config, "video_temporal_causal", False)) - is_self_forcing = getattr(model_config, "simulation_mode", None) == "self_forcing" - is_distilled = getattr(model_config, "fixed_step_sampler_config", None) is not None - if not video_temporal_causal or not (is_self_forcing or is_distilled): - return - if int(getattr(model_config, "teacher_forcing_frames_per_chunk", 1) or 1) != 1: - return - replay_policy = getattr(model_config, "teacher_forcing_replay_policy", None) - if ( - getattr(replay_policy, "control_visibility", None) != "causal" - or not getattr(replay_policy, "controls_read_strict_past_clean_rgb", False) - or getattr(replay_policy, "clean_pass_causality", None) != "frame" - ): + if not self._is_sf_dmd_transfer_model(model_config): return resolution = self.resolution # The selected 480-only checkpoint must not inherit the legacy 720 preset. # Earlier mixed-resolution students record resolution=720 and retain both tiers. if str(getattr(model_config, "resolution", "")) == "480" and resolution != "480": - raise ValueError("This SF-DMD Transfer checkpoint supports only resolution='480'.") + raise ValueError( + f"This SF-DMD Transfer checkpoint supports only resolution='480', got resolution='{resolution}'." + ) if not self.transfer_hints or not self.autoregressive: return frame_cap = self._SF_DMD_TRANSFER_MAX_FRAMES.get(resolution) diff --git a/cosmos_framework/inference/sim_transfer_args_test.py b/cosmos_framework/inference/sim_transfer_args_test.py index 40212edc2..df6e09ab6 100644 --- a/cosmos_framework/inference/sim_transfer_args_test.py +++ b/cosmos_framework/inference/sim_transfer_args_test.py @@ -375,3 +375,47 @@ def test_public_causal_cache_overrides_preserve_exported_config( assert config["attention_sink_size"] == sink assert json.loads(config_path.read_text()) == source_config assert overrides.build_setup(world_size=1, local_world_size=1, device_memory_bytes=_GB200_MEMORY_BYTES) == setup + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("resolution", ["256", "704", "720", "768"]) +def test_480_only_checkpoint_names_its_own_limit_for_every_rejected_tier(tmp_path: Path, resolution: str) -> None: + control = tmp_path / "depth.mp4" + control.touch() + model_config = SimpleNamespace( + resolution="480", + fixed_step_sampler_config=SimpleNamespace(t_list=[1.0, 0.9375, 5 / 6, 0.625]), + video_temporal_causal=True, + causal_training_strategy="none", + teacher_forcing_frames_per_chunk=1, + teacher_forcing_replay_policy=SimpleNamespace( + control_visibility="causal", controls_read_strict_past_clean_rgb=True, clean_pass_causality="frame" + ), + ) + sample = OmniSampleOverrides( + depth={"control_path": str(control)}, + autoregressive=True, + resolution=resolution, + transfer_resize_mode="direct_resize", + control_guidance=1.0, + num_first_chunk_conditional_frames=0, + share_vision_temporal_positions=True, + ) + with pytest.raises(ValueError, match=f"supports only resolution='480', got resolution='{resolution}'"): + sample._build_transfer_data(model_config, sample.sample_meta, user_fields=frozenset()) + sample._apply_sf_dmd_transfer_inference_contract(model_config=model_config, user_fields=frozenset()) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_bidirectional_480_model_keeps_the_teacher_tier_message(tmp_path: Path) -> None: + # Only the causal SF-DMD student is 480-only; a bidirectional model recording 480 still runs 720. + control = tmp_path / "depth.mp4" + control.touch() + model_config = SimpleNamespace(resolution="480", fixed_step_sampler_config=None, video_temporal_causal=False) + sample = OmniSampleOverrides( + depth={"control_path": str(control)}, resolution="256", transfer_resize_mode="direct_resize" + ) + with pytest.raises(ValueError, match="teacher's trained 480 and 720 tiers only, got resolution='256'"): + sample._build_transfer_data(model_config, sample.sample_meta, user_fields=frozenset()) diff --git a/cosmos_framework/inference/transfer_cookbook_specs_test.py b/cosmos_framework/inference/transfer_cookbook_specs_test.py new file mode 100644 index 000000000..638cf1f73 --- /dev/null +++ b/cosmos_framework/inference/transfer_cookbook_specs_test.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""The Transfer cookbook specs must describe the videos they point at.""" + +import json +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + +_COOKBOOK_SPECS = Path(__file__).resolve().parents[2] / ".agents/skills/cookbooks/cosmos3/generator/transfer/specs" + + +def _is_lfs_pointer(path: Path) -> bool: + return path.stat().st_size < 1024 and path.read_bytes().startswith(b"version https://git-lfs") + + +@pytest.mark.parametrize("spec", sorted(_COOKBOOK_SPECS.glob("*.json")), ids=lambda path: path.stem) +def test_cookbook_spec_fps_and_num_frames_match_its_video(spec: Path) -> None: + """Runs where the Git LFS videos are checked out; CI's unit-test checkout does not fetch them.""" + import av + + raw = json.loads(spec.read_text()) + sources = [raw.get("vision_path")] + [hint.get("control_path") for hint in raw.values() if isinstance(hint, dict)] + source = next(path for path in sources if path) + if source.startswith("https://"): + pytest.skip("spec downloads its input") + path = (spec.parent / source).resolve() + if _is_lfs_pointer(path): + pytest.skip("Git LFS assets are not checked out") + with av.open(str(path)) as container: + stream = container.streams.video[0] + assert round(float(stream.average_rate)) == raw["fps"] + # Compared on the 4n+1 latent grid, as generation does. + assert ((stream.frames - 1 + 3) // 4) * 4 + 1 == raw["num_frames"] diff --git a/cosmos_framework/inference/transfer_direct_resize_test.py b/cosmos_framework/inference/transfer_direct_resize_test.py index 0478b60bb..024df63b5 100644 --- a/cosmos_framework/inference/transfer_direct_resize_test.py +++ b/cosmos_framework/inference/transfer_direct_resize_test.py @@ -142,10 +142,10 @@ def test_direct_stage1_transfer_enables_framewise_ar( ), ({"teacher_forcing_frames_per_chunk": 2}, {}, "one latent frame"), ({}, {"autoregressive": False}, "Stage 1 teacher_forcing"), - ({}, {"num_first_chunk_conditional_frames": 1}, "exactly one control hint"), - ({}, {"share_vision_temporal_positions": False}, "share_vision_temporal_positions"), - ({}, {"control_guidance": 2.0}, "control_guidance=1.0"), - ({}, {"edge": {}}, "exactly one control hint"), + ({}, {"num_first_chunk_conditional_frames": 1}, r"num_first_chunk_conditional_frames=1 \(must be 0\)"), + ({}, {"share_vision_temporal_positions": False}, r"share_vision_temporal_positions=false \(must be true"), + ({}, {"control_guidance": 2.0}, r"control_guidance=2.0 \(must be 1.0"), + ({}, {"edge": {}}, r"2 control hints .* \(must be exactly 1\)"), ], ) def test_direct_causal_transfer_rejects_unsupported_contract( @@ -172,6 +172,34 @@ def test_direct_causal_transfer_rejects_unsupported_contract( OmniSampleOverrides(**values).build_sample(model_config=config) +@pytest.mark.L0 +@pytest.mark.CPU +def test_direct_causal_transfer_error_names_only_the_broken_requirements(tmp_path: Path) -> None: + config = _fake_vfm_model_config() + config.video_temporal_causal = True + config.causal_training_strategy = "teacher_forcing" + config.teacher_forcing_frames_per_chunk = 1 + control_path = tmp_path / "depth.mp4" + control_path.touch() + sample = OmniSampleOverrides( + name="two_culprits", + output_dir=tmp_path / "output", + model_mode=ModelMode.VIDEO2VIDEO, + depth={"control_path": str(control_path)}, + transfer_resize_mode="direct_resize", + control_guidance=1.5, + num_first_chunk_conditional_frames=1, + ) + with pytest.raises(ValueError) as excinfo: + sample.build_sample(model_config=config) + + message = str(excinfo.value) + assert "control_guidance=1.5" in message + assert "num_first_chunk_conditional_frames=1" in message + assert "control hints" not in message + assert "share_vision_temporal_positions" not in message + + @pytest.mark.L0 @pytest.mark.CPU @pytest.mark.parametrize("has_rgb", [False, True]) diff --git a/cosmos_framework/inference/transfer_preprocessing.py b/cosmos_framework/inference/transfer_preprocessing.py index 9353b6aaa..c36558372 100644 --- a/cosmos_framework/inference/transfer_preprocessing.py +++ b/cosmos_framework/inference/transfer_preprocessing.py @@ -75,10 +75,25 @@ def select_transfer_geometry( target_hw = (target_h, target_w) ratio = _scale_ratio(source_hw, target_hw) if ratio > 1.0 + _MAX_DISTORTION: + resize = ( + f"{source_hw[1]}x{source_hw[0]} -> {target_hw[1]}x{target_hw[0]} causes {ratio - 1:.2%} " + "aspect distortion, exceeding 3%" + ) + if aspect_ratio is not None: + # A forced bucket is the usual cause: direct_resize stretches the full frame and never crops. + try: + closest = select_transfer_geometry(source_hw, resolution=resolution).aspect_ratio + advice = f"Omit aspect_ratio: the closest trained bucket for this source is {closest!r}." + except ValueError: + advice = "No trained bucket is within 3% of this source either; skip this sample." + raise ValueError( + f"Unsupported Transfer sample: aspect_ratio={aspect_ratio!r} was requested, but " + f"transfer_resize_mode='direct_resize' stretches the full frame without cropping, so {resize}. " + + advice + ) raise ValueError( - f"Unsupported Transfer sample: {source_hw[1]}x{source_hw[0]} -> " - f"{target_hw[1]}x{target_hw[0]} causes {ratio - 1:.2%} aspect distortion, exceeding 3%. " - "Choose a compatible source aspect ratio or skip this sample." + f"Unsupported Transfer sample: {resize}. No trained aspect-ratio bucket is within 3% of the source; " + "skip this sample." ) return TransferGeometry(source_hw, target_hw, selected, ratio - 1.0) diff --git a/cosmos_framework/inference/transfer_preprocessing_test.py b/cosmos_framework/inference/transfer_preprocessing_test.py index 439bbedb7..058042e4a 100644 --- a/cosmos_framework/inference/transfer_preprocessing_test.py +++ b/cosmos_framework/inference/transfer_preprocessing_test.py @@ -56,6 +56,32 @@ def test_exact_three_percent_boundary_and_explicit_override() -> None: select_transfer_geometry((1080, 1920), resolution="480", aspect_ratio="4,3") +def test_forced_aspect_ratio_error_names_the_field_and_the_bucket_that_fits() -> None: + with pytest.raises(ValueError) as excinfo: + select_transfer_geometry((720, 1280), resolution="720", aspect_ratio="3,4") + message = str(excinfo.value) + assert "aspect_ratio='3,4' was requested" in message + assert "without cropping" in message + assert "Omit aspect_ratio: the closest trained bucket for this source is '16,9'." in message + + +def test_forced_aspect_ratio_on_a_source_no_bucket_fits_does_not_suggest_omitting_it() -> None: + with pytest.raises(ValueError) as excinfo: + select_transfer_geometry((1080, 2560), resolution="480", aspect_ratio="16,9") + message = str(excinfo.value) + assert "aspect_ratio='16,9' was requested" in message + assert "Omit aspect_ratio" not in message + assert "No trained bucket is within 3% of this source either" in message + + +def test_auto_bucket_error_says_no_bucket_fits() -> None: + with pytest.raises(ValueError) as excinfo: + select_transfer_geometry((1080, 2560), resolution="480") + message = str(excinfo.value) + assert "No trained aspect-ratio bucket is within 3%" in message + assert "aspect_ratio=" not in message + + @pytest.mark.parametrize("source_hw", [(1080, 2560), (720, 1080), (0, 1920)]) def test_unsupported_source_is_rejected_without_extra_buckets(source_hw: tuple[int, int]) -> None: # The global size table also has a 2:3 robotics bucket, but this teacher was