[None][feat] Cosmos3 video-to-video (V2V) generation - #16155
Conversation
|
/bot run |
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughCosmos3 gains V2V reference-video conditioning, action-modality support, NVDEC decoding, content-based media routing, classified failures, expanded outputs, CLI examples, documentation, and unit/integration coverage. ChangesCosmos3 V2V and action generation
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Client
participant VisualGenServe
participant Cosmos3Pipeline
participant NVDEC
participant Transformer
Client->>VisualGenServe: submit image or video reference
VisualGenServe->>VisualGenServe: classify and validate payload
VisualGenServe->>Cosmos3Pipeline: pass V2V video bytes
Cosmos3Pipeline->>NVDEC: decode conditioning window
NVDEC-->>Cosmos3Pipeline: resized RGB frames
Cosmos3Pipeline->>Transformer: denoise video and optional action streams
Transformer-->>Cosmos3Pipeline: video, audio, or action outputs
Cosmos3Pipeline-->>Client: visual generation response
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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: 3
🧹 Nitpick comments (8)
tensorrt_llm/_torch/visual_gen/pipeline.py (1)
1183-1188: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winHoist
inspect.signature()out of the per-step loop.
post_step_fn's arity never changes across denoising steps, so recomputinginspect.signature()on every iteration is unnecessary reflection overhead in the hot path. Compute it once before the loop.♻️ Proposed fix
+ post_step_takes_extra = ( + post_step_fn is not None and len(inspect.signature(post_step_fn).parameters) >= 2 + ) + for i, t in enumerate(timesteps): ... if post_step_fn is not None: - sig = inspect.signature(post_step_fn) - if len(sig.parameters) >= 2: + if post_step_takes_extra: latents, extra_stream_latents = post_step_fn(latents, extra_stream_latents) else: latents = post_step_fn(latents)As per coding guidelines, "Avoid using reflection when functionality can be easily achieved without reflection."
🤖 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/pipeline.py` around lines 1183 - 1188, The denoising loop in the pipeline hot path recomputes post_step_fn arity with inspect.signature() on every step, which is unnecessary reflection overhead. Move the signature inspection for post_step_fn out of the per-step loop in the pipeline logic, cache the parameter count once before iterating, and reuse that result inside the loop when deciding whether to call post_step_fn with one or two arguments.Source: Coding guidelines
tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py (2)
17-57: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd docstrings to the public helper functions.
pil_to_rgb,decode_video_file, andnormalize_video_input_pathhave no docstrings, whilenormalize_video_inputdoes. Since none of these are underscore-prefixed, they're part of this module's public interface and should be documented (Google style), similar tonormalize_video_input.As per coding guidelines, "For interfaces that may be used outside a file, prefer docstrings over comments."
🤖 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/utils.py` around lines 17 - 57, The public helper functions pil_to_rgb, decode_video_file, and normalize_video_input_path are missing docstrings, unlike normalize_video_input. Add Google-style docstrings to each function that describe its purpose, arguments, return values, and raised errors as applicable, keeping the documentation consistent with the module’s existing public API style.Source: Coding guidelines
8-9: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse modern typing style (
list[...],X | None) instead oftyping.List/typing.Optional.This is a new file, so it's a good opportunity to follow the guideline directly rather than mixing legacy typing conventions.
♻️ Proposed fix
-from typing import Any, List, Optional +from typing import Any -def decode_video_file(path: Path, max_frames: Optional[int] = None) -> List[PIL.Image.Image]: +def decode_video_file(path: Path, max_frames: int | None = None) -> list[PIL.Image.Image]: -def normalize_video_input_path(path: Path, max_frames: Optional[int] = None) -> List[Any]: +def normalize_video_input_path(path: Path, max_frames: int | None = None) -> list[Any]: -def normalize_video_input(video: Any, max_frames: Optional[int] = None) -> List[Any]: +def normalize_video_input(video: Any, max_frames: int | None = None) -> list[Any]:As per coding guidelines, "Prefer built-in types list, dict, and tuple to the legacy typing.List, typing.Dict, and typing.Tuple... use the | syntax instead of typing.Union."
Also applies to: 17-76
🤖 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/utils.py` around lines 8 - 9, The type hints in this new module still use legacy typing imports, so update the imports and annotations in utils.py to the modern built-in style. Replace typing.List and typing.Optional with list[...] and ... | None throughout the file, and adjust any affected function signatures or class attributes so symbols like the module imports and any typed helpers in this file consistently follow the new guideline.Source: Coding guidelines
tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py (1)
400-536: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGood happy-path coverage; consider adding negative-path tests for the action modality.
Structural and forward-pass coverage (presence/absence of action heads, noisy mask, multiframe) is solid. Coverage is currently insufficient for boundary/error conditions on the new action inputs — e.g.
action_domain_idsvalues>= NUM_DOMAINS, oraction_latentswith a mismatchedaction_dim. If these error paths aren't already covered intest_cosmos3_action.py(not included in this review batch), consider adding them there or here.As per path instructions, "Keep feedback actionable: suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR."
🤖 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_transformer.py` around lines 400 - 536, Coverage for the new action modality is insufficient because there are no negative-path tests for invalid action inputs. Add focused assertions in TestCosmos3Action (or follow up in test_cosmos3_action.py if that is the intended home) for cases like action_domain_ids at or above NUM_DOMAINS and action_latents with the wrong action_dim, using Cosmos3VFMTransformer and the existing action_model_config fixture to locate the behavior quickly.Source: Path instructions
tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py (1)
776-909: 🚀 Performance & Scalability | 🔵 TrivialCoverage gap:
resolve_domain_action_config/canonical_domain_preset_keyare only exercised via GPU-gated integration tests.
TestCosmos3Actionis marked@pytest.mark.integration+@pytest.mark.high_cuda_memory, so it won't run in typical CI. The branching logic indefaults.py(alias resolution, ambiguousdomain_id→Nonefallback, mismatch-warning generation,num_frames = action_chunk_size + 1derivation) is pure Python with no GPU dependency, yet has no equivalent fast unit test path here (unlikeTestCosmos3V2VConditioningParams, which does cover its pure-Python counterparts without GPU marks).Coverage assessment: insufficient for non-GPU CI on this specific logic. Suggest adding a
tests/unittest/_torch/visual_gen/test_cosmos3_defaults.py(or a similarly-marker-free test class in this file) that directly unit-testscanonical_domain_preset_key(including the ambiguous-domain_id-returns-Nonecase and alias mapping) andresolve_domain_action_config(mismatch-warning content,num_framesfallback derivation), independent ofcosmos3_pipeline/checkpoint fixtures.Do you want me to draft this unit test module?
🤖 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 776 - 909, The pure-Python domain action config logic is only covered by GPU-gated integration tests, so add a fast unit test path for the defaults helpers. Create a marker-free test module (or a non-integration class in this area) that directly exercises canonical_domain_preset_key and resolve_domain_action_config, including alias mapping, ambiguous domain_id returning None, mismatch-warning text, and the num_frames fallback derived from action_chunk_size + 1. Keep the tests independent of cosmos3_pipeline and checkpoint fixtures so they run in normal CI.tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py (1)
845-902: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd coverage for base64-JSON video
input_reference.New tests cover multipart video upload (routes to
extra_params["video"]) and undecodable multipart content (400), but the PR objectives call out base64-JSON video references as a supported path, which isn't exercised here (only multipart is tested for video). Recommend adding atest_sync_video_generation_json_with_video_reference(or similar) that postsinput_referenceas a base64-encoded string in the JSON body and asserts the sameextra_params["video"]routing.🤖 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_trtllm_serve_endpoints.py` around lines 845 - 902, The current video reference tests only cover multipart upload and the rejection path, but they miss the supported base64-JSON `input_reference` flow. Add a new test alongside `test_sync_video_generation_multipart_with_video_reference` that posts `input_reference` as a base64-encoded video in the JSON body, then verify the request is accepted and that `video_client.mock_gen.last_params.extra_params["video"]` is populated (matching the same routing asserted in the multipart test).Source: Path instructions
tests/unittest/_torch/visual_gen/test_cosmos3_action.py (1)
154-183: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd coverage for
inverse_dynamicsand the missing-source error path inaction_reference_image.Current tests only cover
forward_dynamicsandpolicyselection order.action_reference_image'sinverse_dynamicsbranch (video preferred over image) and itsValueErrorguard when neither image nor video is provided are untested. Coverage is insufficient for this function; recommend adding these two cases totest_cosmos3_action.py.🤖 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_action.py` around lines 154 - 183, Add tests for the missing `action_reference_image` paths in `TestActionReferenceImage`: cover `inverse_dynamics` to verify it prefers `video` over `image`, and add a case asserting the `ValueError` raised when both sources are absent. Keep the new tests alongside the existing `test_forward_dynamics_accepts_mp4_on_image_path` and `test_policy_prefers_image_path_over_video` cases so the branch behavior in `action_reference_image` is fully covered.Source: Path instructions
tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py (1)
82-109: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd Google-style docstrings to public helper functions, especially
prepare_action_latents.Most module-level functions here (e.g.
normalize_action_resolution,normalize_action_mode,prepare_action_latents) lack docstrings despite being imported and used outside this file (tests, presumablypipeline_cosmos3.py).prepare_action_latentsin particular has non-obvious parameters (raw_action_dimvsaction_dim,action_chunk_size) that would benefit from documentation.Based on coding guidelines: "For interfaces used outside a file, prefer docstrings over comments... Externally called functions should have docstrings, and their arguments should be documented."
Also applies to: 325-333
🤖 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/action.py` around lines 82 - 109, Add Google-style docstrings to the public helper functions in this module, including normalize_action_resolution, normalize_action_mode, and especially prepare_action_latents. Document each function’s purpose, arguments, and return values so externally used interfaces are clear, with special attention to the non-obvious parameters in prepare_action_latents such as raw_action_dim, action_dim, and action_chunk_size.Source: Coding guidelines
🤖 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 `@tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py`:
- Around line 325-376: The in-place zeroing in prepare_action_latents can mutate
caller-owned input when clean_action still shares storage with action_input.
Update the prepare_action_latents flow so the tensor is cloned before any
in-place writes, especially before clean_action[:, :, raw_action_dim:] = 0,
while preserving the existing shape/device/dtype handling and the raw_action_dim
logic.
- Around line 267-298: action_reference_image() does not handle URI-based image
references, so http(s):// and data: strings are incorrectly treated as local
paths. Update the string branch in action_reference_image() to detect URI inputs
and pass them through the same URI-aware loading path used elsewhere before
falling back to filesystem/path or normalize_video_input(), while preserving the
existing PIL.Image and inverse_dynamics behavior.
In `@tensorrt_llm/serve/visual_gen_utils.py`:
- Around line 196-202: The temporary reference-file write path in
visual_gen_utils.py needs explicit failure handling and cleanup; wrap the base64
decode and file write logic used for request.input_reference in the same
try/finally flow that handles tmp_path so malformed data or write errors are
converted into a clean validation-style error instead of bubbling up, and ensure
tmp_path is deleted on every failure path. Locate the reference handling block
around request.input_reference, base64.b64decode, shutil.copyfileobj, and the
cleanup logic for tmp_path, and make sure both string and file-object inputs
share the same safe cleanup behavior.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py`:
- Around line 82-109: Add Google-style docstrings to the public helper functions
in this module, including normalize_action_resolution, normalize_action_mode,
and especially prepare_action_latents. Document each function’s purpose,
arguments, and return values so externally used interfaces are clear, with
special attention to the non-obvious parameters in prepare_action_latents such
as raw_action_dim, action_dim, and action_chunk_size.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/utils.py`:
- Around line 17-57: The public helper functions pil_to_rgb, decode_video_file,
and normalize_video_input_path are missing docstrings, unlike
normalize_video_input. Add Google-style docstrings to each function that
describe its purpose, arguments, return values, and raised errors as applicable,
keeping the documentation consistent with the module’s existing public API
style.
- Around line 8-9: The type hints in this new module still use legacy typing
imports, so update the imports and annotations in utils.py to the modern
built-in style. Replace typing.List and typing.Optional with list[...] and ... |
None throughout the file, and adjust any affected function signatures or class
attributes so symbols like the module imports and any typed helpers in this file
consistently follow the new guideline.
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 1183-1188: The denoising loop in the pipeline hot path recomputes
post_step_fn arity with inspect.signature() on every step, which is unnecessary
reflection overhead. Move the signature inspection for post_step_fn out of the
per-step loop in the pipeline logic, cache the parameter count once before
iterating, and reuse that result inside the loop when deciding whether to call
post_step_fn with one or two arguments.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_action.py`:
- Around line 154-183: Add tests for the missing `action_reference_image` paths
in `TestActionReferenceImage`: cover `inverse_dynamics` to verify it prefers
`video` over `image`, and add a case asserting the `ValueError` raised when both
sources are absent. Keep the new tests alongside the existing
`test_forward_dynamics_accepts_mp4_on_image_path` and
`test_policy_prefers_image_path_over_video` cases so the branch behavior in
`action_reference_image` is fully covered.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py`:
- Around line 776-909: The pure-Python domain action config logic is only
covered by GPU-gated integration tests, so add a fast unit test path for the
defaults helpers. Create a marker-free test module (or a non-integration class
in this area) that directly exercises canonical_domain_preset_key and
resolve_domain_action_config, including alias mapping, ambiguous domain_id
returning None, mismatch-warning text, and the num_frames fallback derived from
action_chunk_size + 1. Keep the tests independent of cosmos3_pipeline and
checkpoint fixtures so they run in normal CI.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py`:
- Around line 400-536: Coverage for the new action modality is insufficient
because there are no negative-path tests for invalid action inputs. Add focused
assertions in TestCosmos3Action (or follow up in test_cosmos3_action.py if that
is the intended home) for cases like action_domain_ids at or above NUM_DOMAINS
and action_latents with the wrong action_dim, using Cosmos3VFMTransformer and
the existing action_model_config fixture to locate the behavior quickly.
In `@tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py`:
- Around line 845-902: The current video reference tests only cover multipart
upload and the rejection path, but they miss the supported base64-JSON
`input_reference` flow. Add a new test alongside
`test_sync_video_generation_multipart_with_video_reference` that posts
`input_reference` as a base64-encoded video in the JSON body, then verify the
request is accepted and that
`video_client.mock_gen.last_params.extra_params["video"]` is populated (matching
the same routing asserted in the multipart test).
🪄 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: fffe0e58-286c-43b2-a067-78b9ec05bb3c
📒 Files selected for processing (24)
examples/visual_gen/models/cosmos3/README.mdexamples/visual_gen/models/cosmos3/cosmos3.pyexamples/visual_gen/models/cosmos3/prompts/action_forward_dynamics.jsonexamples/visual_gen/models/cosmos3/prompts/action_inverse_dynamics.jsonexamples/visual_gen/models/cosmos3/prompts/action_policy.jsonexamples/visual_gen/models/cosmos3/prompts/v2v.jsonexamples/visual_gen/serve/README.mdrequirements.txttensorrt_llm/_torch/visual_gen/models/cosmos3/action.pytensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.pytensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.pytensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.pytensorrt_llm/_torch/visual_gen/models/cosmos3/utils.pytensorrt_llm/_torch/visual_gen/output.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/serve/openai_protocol.pytensorrt_llm/serve/visual_gen_utils.pytensorrt_llm/visual_gen/output.pytests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.pytests/unittest/_torch/visual_gen/test_cosmos3_action.pytests/unittest/_torch/visual_gen/test_cosmos3_pipeline.pytests/unittest/_torch/visual_gen/test_cosmos3_transformer.pytests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.pytests/unittest/_torch/visual_gen/test_visual_gen_utils.py
|
/bot run |
|
PR_Github #58510 [ run ] triggered by Bot. Commit: |
|
PR_Github #58510 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58699 [ run ] triggered by Bot. Commit: |
|
PR_Github #58699 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58706 [ run ] triggered by Bot. Commit: |
|
PR_Github #58706 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58736 [ run ] triggered by Bot. Commit: |
|
PR_Github #58736 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58996 [ run ] triggered by Bot. Commit: |
|
PR_Github #58996 [ run ] completed with state
|
|
/bot run |
- Make `is_decodable_image_file` strict (calls `image.load()`) so truncated files are rejected at classification, matching the bytes probe behavior - Add `TestMediaFileProbes` with a truncated-PNG test covering the file-path probe - Add AVI/MJPEG as a documented second container/codec pair: new `_avi_bytes` helper, `test_multipart_avi_reference_routes_to_video`, and `test_avi_mjpeg_bytes_decode` - Update `VideoGenerationRequest.input_reference` docstring and `examples/visual_gen/serve/README.md` to enumerate tested formats Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Each worker rank demuxes encoded video bytes from memory and decodes the V2V conditioning window on NVDEC: forward-only feed, preallocated ring at the request's target resolution (keep=first stops early, keep=last rings to EOS), emitted-frame work limit (TRTLLM_MAX_REFERENCE_DECODE_FRAMES), local-tap Lanczos resize with bounded PIL parity, client/capacity error classes, and an all-rank CPU/gloo status protocol so a rank-local decode failure cannot hang healthy ranks in model collectives. PyNvVideoCodec becomes a declared, pinned dependency (compliance approval recorded, covering the bundled FFmpeg demux libraries); the import stays function-local so importing tensorrt_llm never loads the driver-linked module. Checked-in H.264 fixtures (MP4 + AVI, forced B-frames, per-frame index patterns, provenance in README) drive the direct tests. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
The serve boundary now classifies input_reference by magic bytes (PNG, JPEG, MP4 ftyp, AVI RIFF) — routing only, never acceptance: images still require a full PIL decode at the boundary, while video payloads pass through as untouched encoded bytes for the workers' NVDEC to decode. This removes every OpenCV import and probe from the VisualGen path (cv2 cannot be declared in requirements) along with the boundary video decoders and the decoded-size budget. Base64 decodes strictly; payload size is deliberately not part of the request-validity contract (body limits are deployment policy). Format contract documented accordingly: PNG/JPEG images; H.264 in MP4/AVI decode-tested; the rest best-effort per the GPU's decoder capabilities. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
DiffusionResponse carries an error_type so worker-side failures keep their class over the wire: client errors (unusable reference content, conditioning bounds, decode-work limit) surface as HTTP 400 and as ValueError from VisualGen.generate() — uniform with coordinator preflight — while capacity failures (allocation/OOM for a valid request) surface as HTTP 503 and RuntimeError. Ambiguous NVDEC decoder-init failures and everything unclassified stay 500. No new public exception class. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
extra_params['video'] now carries encoded MP4/AVI bytes only, decoded per rank inside forward(): conditioning indexes are bound-checked against the output latent length before any window math, the decoded frames are freed as soon as the VAE has encoded them, and per-rank prepare outcomes converge through the status protocol before any model collective. The interim tensor form and everything that existed to ship it are removed: the tensor/tensor_or_bytes type-map entries, the coordinator crop reducer, the ExtraParamSchema.reducer hook and reduce_visual_gen_params plumbing, and cosmos3/utils.py (orphaned). Smokes drive the production bytes path via the checked-in fixture — keep=last is now exercised through the real NVDEC decode. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
--video_path reads the file's bytes straight into extra_params['video']; each worker decodes the conditioning window on NVDEC. Docs updated for the NVDEC path (no OpenCV install step, MP4/AVI files). Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
The synthesized-tensor reference died with the tensor intake; the gate now conditions on a 2.7 KB checked-in H.264 clip (same moving-block content, provenance in test_data/README.md — decoded YUV is normative per spec, and the YUV->RGB conversion is this stack's, so the decode is stable for this path). Golden frame regenerated once for the codec-lossy conditioning pixels. Registers test_media_decode.py in l0_b200. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #62004 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/inputs/media_io.py (1)
358-372: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing parameter type annotations on new media-classification helpers. All three new functions omit the
bytestype hint on their sole input parameter, while sibling new functions added in this same PR (e.g.media_decode.py'sdecode_video_reference_window,_read) are fully annotated.
tensorrt_llm/inputs/media_io.py#L358-L372: annotatesniff_media_kind(data: bytes) -> Optional[str].tensorrt_llm/inputs/media_io.py#L374-L389: annotateis_decodable_image_bytes(data: bytes) -> bool.tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py#L93-L101: annotate_validate_video_reference(video: bytes) -> None.As per coding guidelines, "Annotate every function, use
Nonefor non-returning functions, avoidAnyand unnecessary type ignores... use preciseCallabletypes."🤖 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/inputs/media_io.py` around lines 358 - 372, Annotate the three media helpers with precise parameter and return types: update sniff_media_kind in tensorrt_llm/inputs/media_io.py lines 358-372 to accept bytes, update is_decodable_image_bytes in tensorrt_llm/inputs/media_io.py lines 374-389 to accept bytes and return bool, and update _validate_video_reference in tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py lines 93-101 to accept bytes and return None.Source: Coding guidelines
tensorrt_llm/serve/visual_gen_utils.py (1)
87-101: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the
referenceparameter.
referenceis untyped; the two supported shapes are a base64strand a multipart upload exposing.file. As per coding guidelines ("Annotate every function ... prefer built-in generic types and|"), give it an explicit union (e.g.str | UploadFile).🤖 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` around lines 87 - 101, Update the `_read_reference_payload` parameter annotation to explicitly accept the two supported input shapes: a base64 `str` or the project’s multipart upload type exposing `.file` (for example, `UploadFile`). Keep the return type and existing decoding/read behavior unchanged.Source: Coding guidelines
🤖 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 `@tensorrt_llm/visual_gen/visual_gen.py`:
- Around line 166-185: Update the class and aresult() docstrings to document
ValueError for client failures, VisualGenCapacityError for capacity failures,
and RuntimeError for other single-prompt failures, matching _resolved_value().
Add regression tests covering each of these error types.
---
Nitpick comments:
In `@tensorrt_llm/inputs/media_io.py`:
- Around line 358-372: Annotate the three media helpers with precise parameter
and return types: update sniff_media_kind in tensorrt_llm/inputs/media_io.py
lines 358-372 to accept bytes, update is_decodable_image_bytes in
tensorrt_llm/inputs/media_io.py lines 374-389 to accept bytes and return bool,
and update _validate_video_reference in
tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py lines 93-101 to accept
bytes and return None.
In `@tensorrt_llm/serve/visual_gen_utils.py`:
- Around line 87-101: Update the `_read_reference_payload` parameter annotation
to explicitly accept the two supported input shapes: a base64 `str` or the
project’s multipart upload type exposing `.file` (for example, `UploadFile`).
Keep the return type and existing decoding/read behavior 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: 1535fd51-79dd-47db-8842-73e68483c197
⛔ Files ignored due to path filters (4)
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zipis excluded by!**/*.ziptests/integration/defs/examples/visual_gen/test_data/cosmos3_v2v_lpips_reference.mp4is excluded by!**/*.mp4tests/unittest/_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.aviis excluded by!**/*.avitests/unittest/_torch/visual_gen/test_data/cosmos3_v2v_ref_9f_bframes.mp4is excluded by!**/*.mp4
📒 Files selected for processing (24)
examples/visual_gen/models/cosmos3/README.mdexamples/visual_gen/models/cosmos3/cosmos3.pyexamples/visual_gen/serve/README.mdrequirements.txttensorrt_llm/_torch/visual_gen/executor.pytensorrt_llm/_torch/visual_gen/media_decode.pytensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.pytensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/inputs/media_io.pytensorrt_llm/serve/openai_protocol.pytensorrt_llm/serve/openai_video_routes.pytensorrt_llm/serve/visual_gen_utils.pytensorrt_llm/visual_gen/params.pytensorrt_llm/visual_gen/visual_gen.pytests/integration/defs/examples/visual_gen/test_data/README.mdtests/integration/defs/examples/visual_gen/test_visual_gen.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/visual_gen/test_cosmos3_pipeline.pytests/unittest/_torch/visual_gen/test_data/README.mdtests/unittest/_torch/visual_gen/test_media_decode.pytests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.pytests/unittest/_torch/visual_gen/test_visual_gen_params.pytests/unittest/_torch/visual_gen/test_visual_gen_utils.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/serve/openai_protocol.py
| def _resolved_value(self): | ||
| # For single prompts, surface engine-side failure as | ||
| # ``RuntimeError``. Request-parameter validation is enforced | ||
| # synchronously at :meth:`VisualGen.generate_async` entry, so | ||
| # anything reaching this point is by definition a runtime | ||
| # failure from ``pipeline.infer()``. For batches, return the | ||
| # list as-is so callers iterate per-item ``error``. | ||
| # For single prompts, surface engine-side failure typed by class: | ||
| # worker-classified client errors (unusable reference content, | ||
| # conditioning bounds) raise ``ValueError`` — uniform with the | ||
| # synchronous parameter validation at ``generate_async`` entry — | ||
| # and capacity failures raise ``RuntimeError`` | ||
| # (``VisualGenCapacityError``), like any unclassified runtime | ||
| # failure. For batches, return the list as-is so callers iterate | ||
| # per-item ``error``. | ||
| if self._batch_size is None and isinstance(self._resolved, VisualGenOutput): | ||
| if self._resolved.error is not None: | ||
| raise RuntimeError(f"Generation failed: {self._resolved.error}") | ||
| message = f"Generation failed: {self._resolved.error}" | ||
| error_type = getattr(self, "_error_type", None) | ||
| if error_type == "client": | ||
| raise ValueError(message) | ||
| if error_type == "capacity": | ||
| from tensorrt_llm._torch.visual_gen.media_decode import VisualGenCapacityError | ||
|
|
||
| raise VisualGenCapacityError(message) | ||
| raise RuntimeError(message) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the documented exception contract.
This now raises ValueError for client failures and VisualGenCapacityError for capacity failures, while the class and aresult() docstrings still promise RuntimeError for all single-prompt failures (Lines [54] and [94]). Update those docstrings and add regression coverage for each error type.
🤖 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/visual_gen/visual_gen.py` around lines 166 - 185, Update the
class and aresult() docstrings to document ValueError for client failures,
VisualGenCapacityError for capacity failures, and RuntimeError for other
single-prompt failures, matching _resolved_value(). Add regression tests
covering each of these error types.
The class and aresult() docstrings promised RuntimeError for any single-prompt failure, but _resolved_value() raises by failure class: ValueError for client errors, VisualGenCapacityError for capacity, and RuntimeError only when unclassified. ValueError is not a RuntimeError subclass, so a caller following the docs would miss client failures entirely. Adds regression tests over all three classes plus the batch path (which keeps Option B semantics and never raises). They assert the exact type: VisualGenCapacityError subclasses RuntimeError, so a plain pytest.raises(RuntimeError) would pass for a capacity failure and leave the 400-vs-503 distinction the routes depend on untested. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/unittest/visual_gen/test_output.py`:
- Around line 830-842: Parameterize
test_batch_failure_never_raises_regardless_of_error_class over error_type values
"client", "capacity", and None. Pass each value into DiffusionResponse while
preserving the existing assertions that batch results contain per-item errors
and do not raise.
🪄 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: 729a412e-5bc2-4b42-a56c-4c08827d1edc
📒 Files selected for processing (2)
tensorrt_llm/visual_gen/visual_gen.pytests/unittest/visual_gen/test_output.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/visual_gen/visual_gen.py
|
PR_Github #62009 [ run ] triggered by Bot. Commit: |
The test asserted Option B semantics 'regardless of error class' while only exercising the client class; run it over client, capacity, and unclassified so the name matches what it covers. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #62004 [ run ] completed with state |
|
/bot run |
|
PR_Github #62015 [ run ] triggered by Bot. Commit: |
|
PR_Github #62016 [ run ] triggered by Bot. Commit: |
|
PR_Github #62009 [ run ] completed with state |
|
PR_Github #62015 [ run ] completed with state |
|
PR_Github #62016 [ run ] completed with state
|
|
/bot -h |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot run --disable-fail-fast |
|
PR_Github #62106 [ run ] triggered by Bot. Commit: |
|
PR_Github #62106 [ run ] completed with state
|
|
CI status for The failure: Why it is unrelated:
This PR's tests passed. Also note the previous pipeline's failure — |
9 participants
Dev Engineer Review
video(reference bytes)condition_video_latent_indexescondition_video_keepflow_shiftuse_system_promptis now tri-state (None= model decides; CLI supports--no-use_system_prompt)condition_video_latent_indexesnormalization/validation (integer, non-empty, non-negative)condition_video_keeprestricted tofirst/lastoutput_typevalidated, and preflight rejects invalid modality combinations (e.g., image+video together; V2V only for video outputs)PyNvVideoCodecwith a fixed conditioning windowTRTLLM_MAX_REFERENCE_DECODE_FRAMES/DEFAULT_MAX_REFERENCE_DECODE_FRAMES(default 7200;0disables)input_referencenow supports image or video payloads using content/container sniffing (not filename/MIME)VisualGenCapacityErrornum_framesupper bound now uses sharedMAX_VIDEO_FRAMESDomainAwareLinear, action injection/packing into transformer forward, and new optional action inputs/outputsQA Engineer Review
Test code changes (files outside
tests/integration/test_lists/)tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.pyuse_karras_sigmas), V2V conditioning defaults/normalization tests (incl. error cases), flow-shift propagation override, modality rejection (image+video; V2V withoutput_type="image"), and prompt/tokenization role behavior (new system-prompt-default variant).tests/integration/test_lists/: not listed in the updated CI test-db entry provided (needs follow-up).tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.pyinput_referencesuccess routing toparams.extra_params["video"]; invalid non-mediainput_referencereturns HTTP 400; adjusted multipart image reference naming expectation.tests/integration/test_lists/: not listed in provided update (needs follow-up).tests/unittest/_torch/visual_gen/test_visual_gen_utils.pyparse_visual_gen_paramsrouting for base64/multipart MP4+AVI→extra_params["video"], JPEG→params.image, undecodable/malformed base64 rejection with cleanup, and media probing/sniffing unit tests.tests/integration/test_lists/: not listed in provided update (needs follow-up).tests/unittest/_torch/visual_gen/test_media_decode.pyPyNvVideoCodecnot eagerly loaded; distributed synchronization convergence tests; decode limit parsing (0disables) and GPU-gated correctness/perf for decoding window behavior.tests/integration/test_lists/: Yes (unittest/_torch/visual_gen/test_media_decode.pyadded).tests/unittest/_torch/visual_gen/test_visual_gen_params.pyCOSMOS3_EXTRA_SPECS.tests/integration/test_lists/: not listed in provided update (needs follow-up).tests/unittest/visual_gen/test_output.pyerror_type, including mapping"capacity"→VisualGenCapacityError, plus batch behavior ensuring errors populate per-item without raising.tests/integration/test_lists/: not listed in provided update (needs follow-up).Integration test code changes
tests/integration/defs/examples/visual_gen/test_visual_gen.pytest_cosmos3_nano_v2v_lpips_against_golden(CUDA-gated) covering V2V LPIPS comparison using an inputvideopayload.tests/integration/test_lists/: Yes (examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_v2v_lpips_against_goldenadded).Test-list changes
tests/integration/test_lists/test-db/l0_b200.ymlunittest/_torch/visual_gen/test_media_decode.pyexamples/visual_gen/test_visual_gen.py::test_cosmos3_nano_v2v_lpips_against_golden(TIMEOUT (10))</final_summary>
Description
Adds video-conditioned generation (V2V) for Cosmos3, offline and via the
OpenAI-style videos API, plus the action-generation groundwork it builds on.
[0,1]= first 5input frames), velocity masking + post-step re-injection of clean latents,
flow_shift=10.0with Karras sigmas forced off (restored per request),condition_frame_indexes_vision/condition_video_keeprequest params,shared video decode utils (
.mp4/.avifile, frame directory, image, PILlist). Mode selection is implicit: a video input without an action mode runs
V2V. Faithful to the vllm-omni reference implementation.
input_referenceon the videos API is now classified bydecoding its content (PIL-readable → image/I2V, PyAV-readable video stream →
video/V2V); filename and content-type are ignored. Video uploads route to
the same pipeline entry as the offline path; undecodable content returns
HTTP 400; base64-JSON video references work. Image-reference behavior is
byte-identical to before. No request-schema changes (
input_referencetype/default untouched — api-compatible).
embodiment domain presets,
action_fps— included because V2V structurallyreuses this machinery (VAE video encode,
--video_pathCLI, videonormalize helpers extracted into shared
utils.py).README (
input_referencesemantics, V2V curl, Cosmos3extra_params),prompts/v2v.json. No media assets ship with the repo — examples point atthe checkpoint's own
assets/; test videos are synthesized with PyAV/PIL.av(PyAV) inrequirements.txt— video referencedecode (torchvision
read_videobackend) and serve-side classification.Validated on Cosmos3-Nano (1x B200): offline and served V2V outputs pin the
conditioning frames (frame-0 pixel MAE ≈ 2 vs ≈ 18 different-frame scale) and
diverge afterwards; I2V upload path regression-checked; garbage upload returns
a clean 400.
The default system prompt intentionally says "a give prompt" — it matches the
model's training data; please do not "fix" it in review.
Test Coverage
tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py -k v2v—checkpoint-gated GPU suite (skips without
DIFFUSION_MODEL_PATH_COSMOS3/LLM_MODELS_ROOT): V2V smoke,condition_video_keep="last"behavioral test(dark-head/bright-tail input → output must track the tail), audio+V2V smoke,
conditioning-param normalization, flow-shift request path,
_prepare_latents_v2vvalues.tests/unittest/_torch/visual_gen/test_visual_gen_utils.py— CPU:content-classifier and routing units (mp4 upload →
extra_params["video"],base64 mp4, JPEG → image, undecodable → 400 + temp-file cleanup).
tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py— CPU,mock generator: multipart video reference → V2V routing with
.mp4suffix;undecodable reference → HTTP 400; existing image-reference tests cover the
I2V regression.
tests/unittest/_torch/visual_gen/test_cosmos3_action.pyandmulti_gpu/test_cosmos3_transformer_parallel.py— action-generationcoverage from the underlying commits.
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.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.