From 2b974e6173b265ad6f443bb02e9ac171f3126892 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 2 Jul 2026 01:32:44 -0700 Subject: [PATCH] Add multi-control transfer inference smoke test Add `test_nano_inference_multi_control_transfer` to the Cosmos3-Nano generator inference smoke test, mirroring the existing single-control transfer run (latency preset, 4 ranks -> cfgp=2/cp=2). The generated spec sets two control hints (edge + blur) with per-hint weights and no control_path, so both controls are derived on the fly from a single `vision_path` (the pinned public robot_pouring.mp4 clip) and blended by `multi_control_two_way_attention` (N maskless SDPA passes summed by weight). This exercises the multi-control path added in the multi-control transfer inference feature. Asserts: exactly one output; both edge and blur hints active with a per-hint weight and no control_path; vision_path set; control_guidance and guidance > 1.0; and a non-degenerate vision.mp4 (via the existing _assert_video_has_content). Smoke-level (output validity + path executed under the framework's multi-control runtime asserts), not numeric goldens. Verified end to end on a 4-rank latency run (exit 0, both controls computed on the fly, all assertions pass). Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/nano_inference_smoke_test.py | 115 +++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/tests/nano_inference_smoke_test.py b/tests/nano_inference_smoke_test.py index 14defa15..098d8510 100644 --- a/tests/nano_inference_smoke_test.py +++ b/tests/nano_inference_smoke_test.py @@ -101,6 +101,54 @@ "edge": {"control_path": _TRANSFER_CONTROL_URL, "preset_edge_threshold": "medium"}, } +# Multi-control transfer (video2video, edge + blur) input, written to a temp file +# at run time. Mirrors the cookbook +# ``cookbooks/cosmos3/generator/transfer/specs/multi_control.json`` — two control +# hints (edge + blur) computed on the fly from a single source video (``vision_path``) +# and blended by ``multi_control_two_way_attention`` (N independent maskless SDPA +# passes, one per control, summed by the per-hint ``weight``) — but downscaled +# (480p / 10 steps / single 29-frame chunk) for a fast smoke run. The source clip is +# the exact one the cookbook uses (a robot arm pouring into a glass), pinned to a +# public raw URL; the prompt is a compact caption of it. Unlike ``_TRANSFER_SPEC`` +# (a single pre-computed ``control_path``), both controls here are derived on the +# fly, so this exercises the transfer control augmentor in addition to the weighted +# multi-control aggregation. ``guidance`` + ``control_guidance`` > 1.0 also keep the +# text-CFG and control-CFG branches active. +_MULTI_CONTROL_VISION_URL = ( + "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/" + "2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/vision/robot_pouring.mp4" +) +_MULTI_CONTROL_SPEC = { + "name": "transfer_multi_control", + "model_mode": "video2video", + "resolution": "480", + "aspect_ratio": "16,9", + "num_frames": 29, + "fps": 30, + "shift": 10.0, + "num_steps": 10, + "seed": 2026, + "num_video_frames_per_chunk": 29, + "max_frames": 29, + "num_conditional_frames": 1, + "num_first_chunk_conditional_frames": 0, + "share_vision_temporal_positions": True, + "guidance": 3.0, + "control_guidance": 1.5, + "vision_path": _MULTI_CONTROL_VISION_URL, + "prompt": ( + "A white robotic arm with black joints and cables carefully pours a clear liquid from a " + "small light-green pitcher into a glass on a white tabletop, in a clean, brightly lit " + "modern indoor setting." + ), + "negative_prompt": "blurry, distorted, deformed, low quality, flickering, artifacts", + # Two hints, no control_path -> both derived on the fly from vision_path; the + # per-hint weights drive the weighted multi-control attention aggregation. + "edge": {"weight": 0.5, "preset_edge_threshold": "medium"}, + "blur": {"weight": 0.5, "preset_blur_strength": "medium"}, + "emphasize_control_in_prompt": False, +} + # Audio sanity thresholds for the muxed sound track. _RMS_SILENCE_FLOOR = 1e-4 # below this the track is effectively silence _PEAK_SANITY_CEIL = 1.5 # decoded float audio should sit within ~[-1, 1] @@ -355,3 +403,70 @@ def test_nano_inference_omni(tmp_path: Path) -> None: transfer_video = so.parent / "vision.mp4" assert transfer_video.is_file(), f"transfer run produced no vision.mp4 ({so})" _assert_video_has_content(transfer_video) + + @pytest.mark.level(2) + @pytest.mark.gpus(8) + def test_nano_inference_multi_control_transfer(tmp_path: Path) -> None: + """Multi-control transfer: edge + blur derived on the fly from ONE source + video, blended by ``multi_control_two_way_attention``. + + Mirrors ``test_nano_inference_omni``'s single-control transfer run (same + ``latency`` preset, 4 ranks -> cfgp=2, cp=2 -- the cookbook Cosmos3-Super + transfer layout), but the generated spec sets TWO control hints (edge + + blur) each with a per-hint ``weight`` and no ``control_path``, so both + controls are computed on the fly from ``vision_path`` and aggregated by the + weighted multi-control attention path (``multi_control_two_way_attention``: + N maskless SDPA passes summed by weight). A non-degenerate clip confirms + that path ran end to end -- a broken multi-control route would raise + mid-sampling, and a numerically broken one would collapse the output + (caught by ``_assert_video_has_content``). The on-the-fly derivation also + exercises the transfer control augmentor (opencv), unlike the single-control + run above which loads a pre-computed control_path.""" + spec_file = tmp_path / "transfer_multi_control.json" + spec_file.write_text(json.dumps(_MULTI_CONTROL_SPEC)) + out_dir = tmp_path / "out_multi_control" + cmd = [ + "torchrun", + "--nproc_per_node=4", + f"--master_port={_free_port()}", + "-m", + "cosmos_framework.scripts.inference", + "--parallelism-preset=latency", + "-i", + str(spec_file), + "-o", + str(out_dir), + "--checkpoint-path", + "Cosmos3-Nano", + "--seed=0", + ] + _run(cmd, tmp_path / "inference_multi_control.log") + + results = sorted(out_dir.rglob("sample_outputs.json")) + assert len(results) == 1, ( + f"expected 1 multi-control sample_outputs.json, found {[str(p) for p in results]}" + ) + so = results[0] + args = json.loads(so.read_text()).get("args", {}) + # Multi-control-specific: BOTH edge and blur hints are active (2 controls -> + # the weighted multi_control_two_way_attention path), each carries a weight, + # and neither has a control_path (both derived on the fly from vision_path). + edge = args.get("edge") or {} + blur = args.get("blur") or {} + assert edge and blur, f"expected both edge and blur hints active ({so}); edge={edge} blur={blur}" + assert edge.get("weight") is not None and blur.get("weight") is not None, ( + f"expected a per-hint weight on both controls ({so}); edge={edge} blur={blur}" + ) + assert not edge.get("control_path") and not blur.get("control_path"), ( + f"expected on-the-fly controls (no control_path) ({so}); edge={edge} blur={blur}" + ) + assert args.get("vision_path"), f"multi-control run missing vision_path ({so})" + assert args.get("control_guidance", 1.0) > 1.0, ( + f"expected control-CFG (control_guidance > 1.0), got {args.get('control_guidance')} ({so})" + ) + assert (args.get("guidance") or 1.0) > 1.0, ( + f"expected text-CFG (guidance > 1.0), got {args.get('guidance')} ({so})" + ) + video = so.parent / "vision.mp4" + assert video.is_file(), f"multi-control run produced no vision.mp4 ({so})" + _assert_video_has_content(video)