[TRTLLM-14268][feat] Cosmos3 Transfer (control-video conditioning) - #16394
Conversation
9319387 to
7cc0a5e
Compare
568a264 to
e6fb891
Compare
e6fb891 to
c5bb75c
Compare
ishovkun
changed the title
|
/bot run |
|
PR_Github #63660 [ run ] triggered by Bot. Commit: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughCosmos3 now supports Transfer control-video conditioning. The change adds validated hints, encoded media handling, generated controls, GPU kernels, chunked inference, control-latent transformer support, CLI examples, and tests. ChangesCosmos3 Transfer conditioning
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant parse_visual_gen_params
participant Cosmos3Transfer
participant Cosmos3OmniMoTPipeline
participant Cosmos3VFMTransformer
Client->>parse_visual_gen_params: Submit Transfer parameters and encoded media
parse_visual_gen_params->>Cosmos3Transfer: Decode and validate controls
Cosmos3Transfer->>Cosmos3OmniMoTPipeline: Provide transfer configuration and control frames
Cosmos3OmniMoTPipeline->>Cosmos3VFMTransformer: Denoise with control latents
Cosmos3OmniMoTPipeline->>Client: Return assembled video
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
🧹 Nitpick comments (12)
tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py (1)
1700-1714: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueGuard the prompt log with a rank check.
Lines 1586 and 1598 in this same function already wrap their
logger.infocalls withif self.rank == 0. Line 1704 and Line 1707 do not. On a multi-rank job every rank emits the full user prompt, which duplicates user content across all rank logs.♻️ Proposed rank guard
prompt = [prompt] if isinstance(prompt, str) else list(prompt) prompt = prompt[0] - logger.info(f"Transfer prompt: '{prompt}'") - - # 1. Tokenize prompts (no separate text encoder — transformer embeds internally) - logger.info("Tokenizing prompts...") + if self.rank == 0: + logger.info(f"Transfer prompt: '{prompt}'") + logger.info("Tokenizing prompts...")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py` around lines 1700 - 1714, Guard the transfer prompt logging in the current function with self.rank == 0, matching the existing rank checks around the other logger.info calls. Ensure nonzero ranks do not emit the user prompt while preserving the prompt normalization and tokenization flow.tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py (2)
805-860: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd cases for the new transfer rejection guards.
This PR adds two validation branches in
Cosmos3OmniMoTPipeline.forwardat Lines 933-939 oftensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py: transfer withoutput_type='image'and transfer withenable_audio=True.TestCosmos3TransferRoutingcovers only theuse_system_promptdefault. Neither rejection branch has a test.The existing stub setup in this test already supports both cases, so each addition is a few lines.
💚 Proposed additional cases
def test_transfer_rejects_image_output_and_audio(self): from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import resolve_transfer_config pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) pipeline.action_gen = False pipeline.audio_gen = True pipeline._apply_flow_shift = lambda *args, **kwargs: None class FakeSampling: is_distilled = False checkpoint_flow_shift = 1.0 def validate_request(self, num_inference_steps, guidance_scale): return None def generation_default_overrides(self): return {} pipeline.sampling = FakeSampling() pipeline._forward_transfer = lambda **kwargs: None cfg = resolve_transfer_config( {"edge": True}, SimpleNamespace(num_frames=93, guidance_scale=None), None ) common = dict( prompt="bounce", transfer_config=cfg, height=16, width=16, num_frames=5, num_inference_steps=1, guidance_scale=1.0, seed=1, max_sequence_length=8, frame_rate=8.0, use_guardrails=False, ) with pytest.raises(ValueError, match="only for video outputs"): pipeline.forward(output_type="image", **common) with pytest.raises(ValueError, match="cannot be combined with sound generation"): pipeline.forward(enable_audio=True, **common)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py` around lines 805 - 860, Add coverage to TestCosmos3TransferRouting for both transfer validation guards in Cosmos3OmniMoTPipeline.forward: assert a ValueError when output_type is "image" and when enable_audio is true. Reuse the existing stub pipeline and transfer configuration setup, and verify each error with the expected message pattern.
805-862: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd tests for the remaining Cosmos3 branches.
Test coverage summary
- Added
TestCosmos3TransferRouting.test_transfer_use_system_prompt_defaults_off. No tests were modified or removed.- The test is covered by
tests/integration/test_lists/test-db/l0_b200.yml. No matching QA entry exists.- Transfer routing and helper paths are covered. The transfer output and audio rejection paths, control/audio exclusivity, control-latent shape validation, and multi-hint mRoPE paths remain uncovered. Verdict: insufficient.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py` around lines 805 - 862, Expand TestCosmos3TransferRouting with tests for the uncovered Cosmos3 branches: transfer output handling, audio-input rejection, control/audio exclusivity, control-latent shape validation, and multi-hint mRoPE behavior. Reuse the existing Cosmos3 pipeline test fixtures and assert each branch’s expected result or error, without changing the existing system-prompt default test.Source: Path instructions
tensorrt_llm/serve/visual_gen_utils.py (1)
105-105: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the
specsparameter.
_decode_inline_mediaannotatesextra_paramsbut leavesspecsbare. Use the declared spec type so the signature is complete.-def _decode_inline_media(extra_params: dict | None, specs) -> None: +def _decode_inline_media( + extra_params: dict | None, specs: Optional[Dict[str, Any]] +) -> None:As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/serve/visual_gen_utils.py` at line 105, Update the _decode_inline_media function signature to annotate the specs parameter with its declared specification type, while preserving the existing extra_params annotation and None return annotation.Source: Coding guidelines
tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py (4)
424-424: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the return type of
_import_cv2.Use
types.ModuleTypeso the function matches the annotation rule applied to the rest of this module.+from types import ModuleType ... -def _import_cv2(hint_key: str): +def _import_cv2(hint_key: str) -> ModuleType:As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py` at line 424, Update the _import_cv2 function signature to annotate its return type as types.ModuleType, ensuring the required types import is available and preserving the function’s existing behavior.Source: Coding guidelines
436-441: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueImport cv2 after the
noneshort-circuit.
make_blur_controlimports cv2 before it checks the preset. Thenonepreset returns a clone and uses no cv2 function, but it still raisesImportErrorwhen opencv-python is absent. Move the import below the short-circuit.♻️ Proposed change
def make_blur_control(frames: torch.Tensor, preset: str) -> torch.Tensor: - cv2 = _import_cv2("blur") preset = preset.lower() if preset not in BLUR_PRESETS: raise ValueError(f"Unsupported Cosmos3 blur preset: {preset!r}.") if preset == "none": return frames.clone() + cv2 = _import_cv2("blur")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py` around lines 436 - 441, Update make_blur_control to normalize and validate preset, then handle the "none" preset by returning the cloned frames before calling _import_cv2("blur"). Import cv2 only after that short-circuit so the no-blur path does not require OpenCV.
177-195: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueConsider treating
Falseas "hint not selected".
raw is Trueenables the hint, butFalsefalls through to theTypeErrorbranch. A client that sends{"edge": false}to disable the hint receives an error instead of a request without the hint. The error text lists the accepted forms, so the failure is actionable, but skippingFalseis the symmetric behavior.♻️ Proposed change
for key in TRANSFER_HINT_KEYS: raw = extra_params.get(key, None) if raw is None: continue + if raw is False: + continue if raw is True: raw = {}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py` around lines 177 - 195, Update the TRANSFER_HINT_KEYS normalization loop to treat raw is False the same as raw is None: skip that transfer hint without raising an error or including it in the request. Preserve the existing raw is True enablement and validation behavior for all other values.
129-144: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDerive the chunking defaults from
TRANSFER_SAMPLE_DEFAULTS.The values 93, 1, 5000, 0, and
Trueare declared twice: once inTRANSFER_SAMPLE_DEFAULTSand once as dataclass field defaults. The two copies can drift, and a directCosmos3TransferConfig()construction would then disagree withresolve_transfer_config.♻️ Proposed refactor
- num_video_frames_per_chunk: int = 93 - num_conditional_frames: int = 1 - max_frames: int = 5000 - show_control_condition: bool = False - show_input: bool = False - num_first_chunk_conditional_frames: int = 0 - share_vision_temporal_positions: bool = True + num_video_frames_per_chunk: int = TRANSFER_SAMPLE_DEFAULTS["num_video_frames_per_chunk"] + num_conditional_frames: int = TRANSFER_SAMPLE_DEFAULTS["num_conditional_frames"] + max_frames: int = TRANSFER_SAMPLE_DEFAULTS["max_frames"] + show_control_condition: bool = TRANSFER_SAMPLE_DEFAULTS["show_control_condition"] + show_input: bool = TRANSFER_SAMPLE_DEFAULTS["show_input"] + num_first_chunk_conditional_frames: int = TRANSFER_SAMPLE_DEFAULTS[ + "num_first_chunk_conditional_frames" + ] + share_vision_temporal_positions: bool = TRANSFER_SAMPLE_DEFAULTS[ + "share_vision_temporal_positions" + ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py` around lines 129 - 144, The Cosmos3TransferConfig field defaults duplicate chunking values from TRANSFER_SAMPLE_DEFAULTS and can diverge. Update the affected fields—num_video_frames_per_chunk, num_conditional_frames, max_frames, num_first_chunk_conditional_frames, and share_vision_temporal_positions—to derive their defaults from TRANSFER_SAMPLE_DEFAULTS, keeping direct Cosmos3TransferConfig construction consistent with resolve_transfer_config.tests/unittest/_torch/visual_gen/test_visual_gen_utils.py (1)
642-669: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a pass-through test for values that already arrived as bytes.
The docstring of
_decode_inline_mediastates that values that already arrived as bytes pass through. No test pins that contract, so a future change that base64-decodes bytes would not fail here. The nestedcontrolbytes case has the same gap.💚 Suggested additional test
def test_existing_bytes_pass_through_unchanged(self): request = VideoGenerationRequest( prompt="storm", extra_params={"video": b"\x00mp4", "edge": {"control": b"\x00ctrl"}}, ) params = parse_visual_gen_params(request, "id-bytes", self._generator()) assert params.extra_params["video"] == b"\x00mp4" assert params.extra_params["edge"]["control"] == b"\x00ctrl"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_visual_gen_utils.py` around lines 642 - 669, The tests cover base64 strings but not media values already represented as bytes. Add a test alongside test_base64_extra_param_reaches_the_pipeline_as_bytes and test_nested_control_reaches_the_pipeline_as_bytes that passes byte values for both video and nested edge.control through parse_visual_gen_params, then asserts both remain unchanged.tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py (3)
99-186: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove
calls_num_train_timestepsaboveforward.The class attribute is defined at line 134, after the method that reads it. The code works because the class body runs before instantiation, but a reader meets the name before its definition.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py` around lines 99 - 186, Move the StubTransformer class attribute calls_num_train_timesteps above forward so the normalization constant is declared before the method that references it; leave the value and all other behavior unchanged.
715-721: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSimplify the
check_video_safetystub.The
andexpression relies onsetdefaultreturning a truthy value to passvideothrough. A plain function states the intent and cannot return the sentinel by accident.♻️ Proposed change
+ def fake_check(video, checker): + seen["called"] = True + return video + monkeypatch.setattr( pipeline_module, "check_video_safety", - lambda video, checker: seen.setdefault("called", True) and video, + fake_check, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py` around lines 715 - 721, Replace the lambda used to monkeypatch check_video_safety in the guardrails test with a simple named or inline function that records the call and explicitly returns video, avoiding reliance on setdefault’s return value. Keep the existing seen assertion and _run invocation unchanged.
193-255: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the config validation errors.
resolve_transfer_configraises for a non-positivenum_video_frames_per_chunk, a negativenum_conditional_frames, a non-positivemax_frames, a negativenum_first_chunk_conditional_frames, and an unsupported edge or blur preset._as_intervalraises for a wrong-length value and forlo > hi, and it accepts the"lo,hi"string form. None of these paths are exercised. They are all client-facing 400 responses over serving.💚 Suggested additional tests
`@pytest.mark.parametrize`( "extra, match", [ ({"num_video_frames_per_chunk": 0}, "num_video_frames_per_chunk must be positive"), ({"num_conditional_frames": -1}, "num_conditional_frames must be non-negative"), ({"max_frames": 0}, "max_frames must be positive"), ( {"num_first_chunk_conditional_frames": -1}, "num_first_chunk_conditional_frames must be non-negative", ), ({"edge": {"preset_edge_threshold": "nope"}}, "Unsupported Cosmos3 edge preset"), ], ) def test_invalid_settings_are_rejected(self, extra, match): with pytest.raises(ValueError, match=match): resolve_transfer_config({"edge": True, **extra}, _req()) def test_control_guidance_interval_forms(self): cfg = resolve_transfer_config( {"edge": True, "control_guidance_interval": "0.2, 0.8"}, _req() ) assert cfg.control_guidance_interval == (0.2, 0.8) with pytest.raises(ValueError, match="exactly two values"): resolve_transfer_config({"edge": True, "control_guidance_interval": [0.2]}, _req()) with pytest.raises(ValueError, match=r"\[lo, hi\]"): resolve_transfer_config( {"edge": True, "control_guidance_interval": [0.8, 0.2]}, _req() )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py` around lines 193 - 255, Add pytest coverage in TestTransferConfig for resolve_transfer_config validation: parameterize non-positive or negative frame settings and unsupported edge/blur presets, asserting the expected ValueError messages. Add interval coverage for the control_guidance_interval setting, verifying "lo,hi" parsing and ValueError handling for incorrect length and lo greater than hi.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/visual_gen/models/cosmos3/cosmos3.py`:
- Around line 222-234: The --extra_params parser currently accepts any JSON
value, but downstream code expects a mapping. Add an argparse type helper near
the parser setup that decodes the input with json.loads, validates the result is
a dict, and raises argparse.ArgumentTypeError for other JSON types; use this
helper in the --extra_params argument instead of json.loads.
- Around line 49-101: Annotate both helpers using built-in generic types and
union syntax, replacing Dict, Optional, and Any. Give
_fit_output_to_source.params the concrete parameter/configuration type used by
its caller and annotate _load_transfer_controls.extra_params with an appropriate
value type that supports the accessed hint data. Preserve the existing return
annotations and behavior while ensuring every function parameter is precisely
typed.
In `@examples/visual_gen/serve/README.md`:
- Line 319: Update the Cosmos3 controls documentation near the existing
transfer-hint description to document precomputed controls as JSON under
extra_params.<hint>.control, using a base64-encoded MP4/AVI string and the depth
example specified by the reviewer. Document the optional preset_edge_threshold
and preset_blur_strength fields for edge and blur, and add a complete request
example showing the structure.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`:
- Line 1787: Rename the comprehension target in _forward_transfer that iterates
over control_norms.values() from video to a distinct name, and update its use in
_encode_video_tensor accordingly; preserve the existing control_latents
construction.
- Around line 933-939: Update the transfer_config validation block in the
pipeline forward path to reject requests that provide both transfer_config and
an image input, before routing to _forward_transfer. Preserve the existing
is_t2i/output-type and enable_audio validation, and raise a clear ValueError for
the conflicting image request.
- Around line 1565-1585: Update the `_forward_transfer` signature annotations to
match its callers: use `list[str]` for `prompt`, `bool` for
`use_duration_template` and `use_resolution_template`, add concrete annotations
for `max_sequence_length`, `use_system_prompt`, and `seed` consistent with
values passed by `forward`, and replace `video: Any` with `bytes | None`;
preserve the existing `PipelineOutput` return annotation and annotate every
parameter without using `Any`.
- Around line 1401-1411: Rename the static helper positive_float to
_positive_float and update its call in _forward_transfer to use the new private
name, preserving the existing validation behavior.
- Line 334: Update classify_worker_error() to classify TypeError alongside
ValueError as a "client" error, covering malformed transfer hints raised by
resolve_transfer_config(). Preserve the existing mappings for other exception
types.
- Around line 1413-1422: Document that control_guidance_interval is interpreted
in raw scheduler-timestep units, matching the direct comparison performed by
_transfer_active_at and the float conversion in _as_interval. Clarify that
fractional values such as [0.0, 0.8] are not normalized to schedule progress;
preserve the existing comparison behavior.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py`:
- Around line 1053-1054: Update the docstring for the surrounding transformer
method: remove the inline `#TODO` from the control_latents description so it
clearly explains the clean vision context, and add
transfer_share_vision_temporal_positions to the Google-style Args section with
an accurate description of its behavior.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py`:
- Around line 1-45: Register the test_cosmos3_transfer.py module in an
appropriate test-list YAML, such as the l0_b200 list, so all 35 transfer test
functions are discoverable by the test infrastructure. Do not modify the test
implementations or add unrelated coverage changes.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`:
- Around line 1700-1714: Guard the transfer prompt logging in the current
function with self.rank == 0, matching the existing rank checks around the other
logger.info calls. Ensure nonzero ranks do not emit the user prompt while
preserving the prompt normalization and tokenization flow.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py`:
- Line 424: Update the _import_cv2 function signature to annotate its return
type as types.ModuleType, ensuring the required types import is available and
preserving the function’s existing behavior.
- Around line 436-441: Update make_blur_control to normalize and validate
preset, then handle the "none" preset by returning the cloned frames before
calling _import_cv2("blur"). Import cv2 only after that short-circuit so the
no-blur path does not require OpenCV.
- Around line 177-195: Update the TRANSFER_HINT_KEYS normalization loop to treat
raw is False the same as raw is None: skip that transfer hint without raising an
error or including it in the request. Preserve the existing raw is True
enablement and validation behavior for all other values.
- Around line 129-144: The Cosmos3TransferConfig field defaults duplicate
chunking values from TRANSFER_SAMPLE_DEFAULTS and can diverge. Update the
affected fields—num_video_frames_per_chunk, num_conditional_frames, max_frames,
num_first_chunk_conditional_frames, and share_vision_temporal_positions—to
derive their defaults from TRANSFER_SAMPLE_DEFAULTS, keeping direct
Cosmos3TransferConfig construction consistent with resolve_transfer_config.
In `@tensorrt_llm/serve/visual_gen_utils.py`:
- Line 105: Update the _decode_inline_media function signature to annotate the
specs parameter with its declared specification type, while preserving the
existing extra_params annotation and None return annotation.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py`:
- Around line 805-860: Add coverage to TestCosmos3TransferRouting for both
transfer validation guards in Cosmos3OmniMoTPipeline.forward: assert a
ValueError when output_type is "image" and when enable_audio is true. Reuse the
existing stub pipeline and transfer configuration setup, and verify each error
with the expected message pattern.
- Around line 805-862: Expand TestCosmos3TransferRouting with tests for the
uncovered Cosmos3 branches: transfer output handling, audio-input rejection,
control/audio exclusivity, control-latent shape validation, and multi-hint mRoPE
behavior. Reuse the existing Cosmos3 pipeline test fixtures and assert each
branch’s expected result or error, without changing the existing system-prompt
default test.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py`:
- Around line 99-186: Move the StubTransformer class attribute
calls_num_train_timesteps above forward so the normalization constant is
declared before the method that references it; leave the value and all other
behavior unchanged.
- Around line 715-721: Replace the lambda used to monkeypatch check_video_safety
in the guardrails test with a simple named or inline function that records the
call and explicitly returns video, avoiding reliance on setdefault’s return
value. Keep the existing seen assertion and _run invocation unchanged.
- Around line 193-255: Add pytest coverage in TestTransferConfig for
resolve_transfer_config validation: parameterize non-positive or negative frame
settings and unsupported edge/blur presets, asserting the expected ValueError
messages. Add interval coverage for the control_guidance_interval setting,
verifying "lo,hi" parsing and ValueError handling for incorrect length and lo
greater than hi.
In `@tests/unittest/_torch/visual_gen/test_visual_gen_utils.py`:
- Around line 642-669: The tests cover base64 strings but not media values
already represented as bytes. Add a test alongside
test_base64_extra_param_reaches_the_pipeline_as_bytes and
test_nested_control_reaches_the_pipeline_as_bytes that passes byte values for
both video and nested edge.control through parse_visual_gen_params, then asserts
both remain unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: efa2821e-6682-428e-ad58-c767586ae0c9
📒 Files selected for processing (13)
examples/visual_gen/models/cosmos3/README.mdexamples/visual_gen/models/cosmos3/cosmos3.pyexamples/visual_gen/models/cosmos3/generate_bouncing_ball_control.pyexamples/visual_gen/serve/README.mdtensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.pytensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.pytensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.pytensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.pytensorrt_llm/serve/visual_gen_utils.pytensorrt_llm/visual_gen/params.pytests/unittest/_torch/visual_gen/test_cosmos3_pipeline.pytests/unittest/_torch/visual_gen/test_cosmos3_transfer.pytests/unittest/_torch/visual_gen/test_visual_gen_utils.py
|
PR_Github #63660 [ run ] completed with state
|
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (4)
tests/unittest/_torch/visual_gen/test_control_kernels.py (1)
204-209: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename
test_blur_reduces_local_variance.The test compares
blurred.var()againstsharp.var(). Both are global variances over the whole clip, not local variances. Rename totest_blur_reduces_varianceso the name states what the assertion checks.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_control_kernels.py` around lines 204 - 209, Rename the test method test_blur_reduces_local_variance to test_blur_reduces_variance, leaving its implementation and assertions unchanged.tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py (1)
151-242: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the validator parameters.
The new validators declare return types but leave
valueunannotated. The coding guidelines require every function to be annotated. The accepted shapes arebool | bytes | Mapping[str, Any]for the hint validators andAnyfor the interval/frame validators.♻️ Proposed annotations
-def _transfer_hint_payload(value) -> Mapping: +def _transfer_hint_payload(value: bool | bytes | Mapping[str, Any]) -> Mapping[str, Any]:-def _validate_edge_hint(value) -> None: +def _validate_edge_hint(value: bool | bytes | Mapping[str, Any]) -> None:Apply the same change to
_validate_blur_hint,_validate_precomputed_control_hint,_validate_control_guidance_interval, and the two frame validators.As per coding guidelines: "Annotate every function, use
Nonefor procedures".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py` around lines 151 - 242, Annotate the validator parameters: use bool | bytes | Mapping[str, Any] for _validate_edge_hint, _validate_blur_hint, and _validate_precomputed_control_hint, and Any for _validate_control_guidance_interval, _validate_positive_frames, and _validate_non_negative_frames. Preserve their existing -> None return annotations and ensure Any/Mapping imports are available.Source: Coding guidelines
tensorrt_llm/_torch/visual_gen/triton_kernels/canny.py (2)
178-180: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDerive the NMS grid from a single row-per-program constant.
The launch repeats the literal
4twice: once intriton.cdiv(H, 4)and once inR=4. The grid and thestatic_range(R)loop must agree. If a later change updates only one of them, the kernel writes the wrong row count with no error. Bind both to one name.♻️ Proposed fix
_BLOCK = 256 +_NMS_ROWS = 4 # output rows per program; the grid and _nms_kernel's R must match _TG22 = 13573 # round(tan(22.5deg) * 2**15)- _nms_kernel[(triton.cdiv(W, _BLOCK), triton.cdiv(H, 4), T)]( - bdx, bdy, mag, cmap, low, high, H, W, R=4, BLOCK=_BLOCK - ) + _nms_kernel[(triton.cdiv(W, _BLOCK), triton.cdiv(H, _NMS_ROWS), T)]( + bdx, bdy, mag, cmap, low, high, H, W, R=_NMS_ROWS, BLOCK=_BLOCK + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/triton_kernels/canny.py` around lines 178 - 180, Update the _nms_kernel launch to define one row-per-program constant and reuse it for both the H grid dimension and the kernel’s R argument, ensuring the launch grid and static_range(R) remain synchronized.
122-122: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReuse the module constant for the Q15 tangent.
Line 122 hardcodes
13573although_TG22holds the same value at line 28. Triton captures module-level Python constants at trace time, so the kernel can reference_TG22directly. This removes the duplicated literal and its duplicated comment.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/triton_kernels/canny.py` at line 122, Update the Canny kernel calculation in the function containing tg22 to use the existing module constant _TG22 instead of the hardcoded 13573 literal, removing the duplicated inline comment while preserving the Q15 tangent behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/visual_gen/models/cosmos3/cosmos3.py`:
- Around line 59-62: Validate control_path in the hint-processing flow before
calling _resolve_path or Path, ensuring it is a non-empty string. For malformed
values such as numbers, reject the request at the CLI boundary with an
actionable error instead of allowing a TypeError, while preserving the existing
behavior for missing paths.
In `@examples/visual_gen/models/cosmos3/README.md`:
- Around line 75-77: Update the description near the example command to label
the output as V2V instead of T2V, keeping the remaining aspect-ratio and
cropping details unchanged.
In `@tensorrt_llm/_torch/visual_gen/triton_kernels/resize.py`:
- Around line 40-43: Bound the module-level _tbl_cache to prevent unbounded
GPU-resident growth across caller-supplied dimensions. Add a small maximum-entry
LRU mechanism with _tbl_get and _tbl_put, then update all four axis-table
builders to retrieve and store entries through those helpers while preserving
the existing cache key and reuse behavior.
- Around line 267-421: Guard the `_area3_kernel` selection in `resize_area_u8`
with `frames.data_ptr() % 4 == 0`, in addition to the existing C=3,
divisibility, and contiguity checks. Route contiguous but unaligned inputs
through the existing `_area_kernel` fallback.
In `@tests/unittest/_torch/visual_gen/test_control_kernels.py`:
- Around line 1-233: Register the CUDA test module in the CI test database by
adding tests/unittest/_torch/visual_gen/test_control_kernels.py to l0_b200.yml
alongside the existing Cosmos3 tests. Do not add a QA entry unless this test is
explicitly owned by manual QA.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py`:
- Around line 151-242: Annotate the validator parameters: use bool | bytes |
Mapping[str, Any] for _validate_edge_hint, _validate_blur_hint, and
_validate_precomputed_control_hint, and Any for
_validate_control_guidance_interval, _validate_positive_frames, and
_validate_non_negative_frames. Preserve their existing -> None return
annotations and ensure Any/Mapping imports are available.
In `@tensorrt_llm/_torch/visual_gen/triton_kernels/canny.py`:
- Around line 178-180: Update the _nms_kernel launch to define one
row-per-program constant and reuse it for both the H grid dimension and the
kernel’s R argument, ensuring the launch grid and static_range(R) remain
synchronized.
- Line 122: Update the Canny kernel calculation in the function containing tg22
to use the existing module constant _TG22 instead of the hardcoded 13573
literal, removing the duplicated inline comment while preserving the Q15 tangent
behavior.
In `@tests/unittest/_torch/visual_gen/test_control_kernels.py`:
- Around line 204-209: Rename the test method test_blur_reduces_local_variance
to test_blur_reduces_variance, leaving its implementation and assertions
unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: cf3202b6-9dc6-4d7a-8bd9-5f4ac9d122c3
📒 Files selected for processing (16)
examples/visual_gen/models/cosmos3/README.mdexamples/visual_gen/models/cosmos3/cosmos3.pytensorrt_llm/_torch/visual_gen/executor.pytensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.pytensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.pytensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.pytensorrt_llm/_torch/visual_gen/triton_kernels/__init__.pytensorrt_llm/_torch/visual_gen/triton_kernels/bilateral.pytensorrt_llm/_torch/visual_gen/triton_kernels/canny.pytensorrt_llm/_torch/visual_gen/triton_kernels/reference.pytensorrt_llm/_torch/visual_gen/triton_kernels/resize.pytensorrt_llm/media/decoding.pytensorrt_llm/visual_gen/visual_gen.pytests/unittest/_torch/visual_gen/test_control_kernels.pytests/unittest/_torch/visual_gen/test_cosmos3_distilled.pytests/unittest/_torch/visual_gen/test_cosmos3_transfer.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py
|
PR_Github #67693 [ run ] triggered by Bot. Commit: |
|
PR_Github #67662 [ run ] completed with state |
|
PR_Github #67693 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #67875 [ run ] triggered by Bot. Commit: |
|
PR_Github #67875 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #67939 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #67959 [ run ] triggered by Bot. Commit: |
|
PR_Github #67939 [ run ] completed with state |
|
PR_Github #67959 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68138 [ run ] triggered by Bot. Commit: |
Signed-off-by: Igor Shovkun <igshov@gmail.com>
|
/bot run --disable-fail-fast |
|
PR_Github #68138 [ run ] completed with state
|
|
PR_Github #68178 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #68178 [ run ] completed with state
|
|
PR_Github #68300 [ run ] triggered by Bot. Commit: |
|
PR_Github #68300 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68398 [ run ] triggered by Bot. Commit: |
|
PR_Github #68398 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68423 [ run ] triggered by Bot. Commit: |
|
PR_Github #68423 [ run ] completed with state |
Transfer landed on main as a squash (cbc3784), so the conflict set is the same action-vs-transfer overlap staged earlier on cosmos3-action-staging; the resolution is taken from there (1e4c1bb286) with the two hunks transfer gained after that staging base folded in via a file-level three-way merge (the Cosmos3CrossAttention backend-fallback warning_once from NVIDIA#17698 and the guardrail early-return fix from NVIDIA#17510). Resolution summary, argued in full in the staging commit: - Token layout control|video|action(|audio): transfer prepends control tokens, action appends, so the velocity decode reads video at T_control..T_control+T_vid_tokens and the extra span starts after it. Control latents are mutually exclusive with both audio (transfer's guard) and action (added here). - Kept the per-request cached_freqs_gen_combined mRoPE table for the audio branch rather than transfer's per-step rebuild. - Dropped action's probe_video_dimensions for transfer's video_stream_info (a strict superset: same header read also returns the frame rate); action_reference_size raises on an unreadable container instead of propagating None. - Action skips the source-following parameter resolution: its frame rate is a trained property of the embodiment, not of the reference footage. - Action's scheduler stream joins transfer's `transfer_config is None` rebuild guard; the two never arrive together. Signed-off-by: Igor Shovkun <igshov@gmail.com>
11 participants
Summary
edge,blur,depth,seg, andwsmcontrol hints.Dev Engineer Review
defaults.py.extra_paramsobjects, and invalid base64 data.extra_paramsAPI consistency.tests/integration/test_lists/test-db/l0_b200.yml.QA Engineer Review
test_cosmos3_pipeline.py.test_cosmos3_transfer.py.test_visual_gen_utils.py.test_control_kernels.py.test_cosmos3_distilled.py.tests/integration/test_lists/test-db/l0_b200.yml.Description
Adds Transfer (ControlNet-style control-video conditioning) for Cosmos3,
offline via
--extra_paramsand over the serving API. A control videoconstrains output structure frame by frame; the prompt supplies appearance.
Builds on Cosmos3 V2V (#16155, merged): transfer's chunk stitching reuses V2V
frame-pinning, and its control media uses V2V's media transport unchanged.
edge,blur,depth,seg,wsm.edge/blurare derived on the GPU from the
videoreference; any hint also accepts aprecomputed control clip as encoded MP4/AVI bytes, the same contract the
videoreference uses. Multiple hints compose.that share the target's mRoPE positions, so control and output patches align
at zero displacement — no extra cross-attention wiring.
cond_full/cond_no_control/uncond_full), combinedu + gs·((nc + cg·(f − nc)) − u); idle branchesskipped, gated over the denoise schedule by
control_guidance_interval.wsm),stitched by V2V-pinning the previous chunk's tail.
--extra_paramsJSON flag (control paths are readclient-side and sent as bytes), a synthetic edge-map control generator
(
generate_bouncing_ball_control.py, no media assets), README usage examples.tensorrt_llm.media.decoding.decode_video_reference_window, resized duringdecode and returned on device. The decode window is bounded by the output
length rather than
max_frames, whose 5000-frame default would otherwisereserve ~14 GB at 720p. Media prepare converges all ranks through
synchronize_media_prepare_statusbefore any model collective.edge(Canny) andblur(bilateralfilter plus a resize chain) are Triton kernels under
tensorrt_llm/_torch/visual_gen/triton_kernels/, consuming the decoded frameswhere they already are instead of round-tripping them through the host.
Measured on 1× B300, 8 frames at 1280×704, wall clock end to end (median of
3 runs on an otherwise idle GPU):
edge195 ms → 1.17 ms (167×),blur medium1.32 s → 1.88 ms (700×),blur high19.7 s → 18.4 ms(1070×). A 93-frame chunk on the
highpreset goes from ~3.8 min to~0.21 s.
control kernels use Triton, itself already a TensorRT-LLM dependency. No
request-schema changes (api-compatible).
Validated on Cosmos3-Nano (1× B200 pre-rebase): edge-fidelity F1 0.971,
multi-hint edge+blur F1 0.956, 189-frame chunked output stitched seamlessly,
synthetic bouncing-ball edge control tracked to ~8 px. Re-confirmed post-rebase
on 1× B300 (multi-hint edge+blur) — RESULT PENDING.
Behavior change for existing V2V requests
Transfer needs the output to follow the control clip, so a request that leaves
height/width/frame_rateunset now derives them from the source video —the closest supported aspect bucket, and the source's frame rate — instead of
falling back to the fixed 1280x720 at the mode default.
That resolution path is shared with V2V, so it changes V2V output for
clients that did not change anything: an unchanged request against a 9:16 or
square source now produces video in that shape rather than 1280x720, and at the
source's fps. Requests that set
height/width/frame_rateexplicitly areunaffected, and pinning them restores the previous behavior exactly.
Called out here rather than only in the README because it is the one part of
this PR that reaches users who are not using Transfer.
Test Coverage
tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py— CPU tests(stubbed scheduler/transformer/decoder, no GPU/model): hint-config resolution;
control-payload contract (bytes accepted bare or in an object, paths and
control_pathrejected, non-bytes rejected, decode window and[3, T, H, W]layout, non-positive window rejected); media helpers (resolution-scaled
bilateral parameters, temporal reflect-pad, generated hints requiring a video
input); 3-branch CFG
arithmetic (applies control+text guidance, skips idle branches, interval
switches branch counts); chunk-count arithmetic, multichunk overlap stitching,
and the decode bound;
find_closest_target_sizebucketing (exact, nearest,order, error paths). Verified against vllm-omni PR fix: fix accuracy and illegal memory access issues when using mtp + attention dp #4379.
tests/unittest/_torch/visual_gen/test_control_kernels.py— GPU tests for thecontrol kernels (83 cases): every kernel asserted bitwise against a torch
reference implementation across shapes, scale ratios, channel counts and
threshold presets, plus input validation (CUDA-only, uint8-only, unsupported
resize geometry) and the transfer entry points that compose them (shape,
RGB broadcast, device, per-preset behaviour, error paths).
tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py::TestCosmos3TransferRouting— transfer requests default
use_system_promptoff.Known limitations / follow-ups
chunk. Control-generation scratch is now O(window): a 189-frame 704p clip
peaked at 7.5 GiB (edge) / 10.2 GiB (blur) and now holds under 3.4 GiB. But
decode and control retention still scale with the clip — 0.49 GiB per 720p
clip at the 189-frame default (1.46 GiB for input + edge + blur), rising to
12.87 GiB per clip if a request asks for the 5000-frame maximum. Bounding
that means decoding and generating per chunk;
show_control_condition/show_inputneed the full clips regardless.are checked on the coordinator and 400 before enqueue. Two combinations are
not: transfer options with no hint selected, and
edge/blurset toauto-compute with no
videoreference — no per-key validator can see anotherkey. Both still fail in the worker before any decode, as a
ValueErrorclassified as a client error, so the cost is a 202-then-error rather than a
400. Closing it needs a whole-request hook on the pipeline base class, which
is a framework change that should stand on its own merits rather than ride
in here.
LPIPS-or-equivalent reference comparison; Cosmos3 already carries several and
reviewers have pushed back on their CI cost, so this PR does not add another.
Correctness is pinned instead by bitwise kernel tests against a torch
reference (83 GPU cases) plus the validation runs above.
Unrelated CI waiver carried in this PR
eccc039touchestests/integration/test_lists/waives.txtto restore twowaivers that are unrelated to Cosmos3. This PR cannot reach green without it,
so it rides here rather than in a separate PR that would need its own full
pipeline before this one could rerun.
Both were waived until #17632 removed them as fixed. They are not fixed — they
still fail on
DGX_B200-4_GPUs-PyTorch-Ray-1with "Disaggregated server failedto start within 5 minutes". The
tp1variants on H100 are unaffected.#17632 was green on its own pre-merge CI (L0 #54328 ran all four cases and
passed them), but its tested head was 79 commits behind main
(
status=diverged, ahead_by=1, behind_by=79), so that result does not cover thetree it merged into.
Scanning every
L0_Test-x86_64-Multi-GPUbuild that reported these cases: runsbefore #17632 merged show
SKIPPED, and failures begin immediately after.Other PRs blocked so far:
_pythonFAILED_python[tp2]failed in all seven;[tp2]passed once, in #204. The blastradius grows as PRs pull in main, since a branch keeps the old
waives.txtuntil it merges.
Restoring the waivers under their original nvbug IDs unblocks CI while the real
fix is worked in #17632. Note also that the test creates the server subprocess
with
stdout=PIPE, stderr=PIPEand nothing ever drains or prints them, so theserver's own error is discarded on assertion failure — which is why the CI
artifacts carry no usable diagnostic.
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.