[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:
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 |
|
PR_Github #62974 [ run ] triggered by Bot. Commit: |
|
PR_Github #62938 [ run ] completed with state |
|
PR_Github #62974 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63020 [ run ] triggered by Bot. Commit: |
|
PR_Github #63020 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
Brings in the DMD2-distilled Cosmos3 4-step I2V checkpoint (NVIDIA#16690), which lands in the same call paths as V2V. Non-obvious resolutions: - forward(): drop the guard rejecting image conditioning on distilled checkpoints. NVIDIA#16690 implements exactly that, so the guard would have rejected the feature it just merged. - denoise(): both sides added a post_step_fn argument to the same call and git took both, which does not compile. V2V pins the conditioning latents and distilled I2V re-anchors the conditioning frame; a request carries an image or a video, never both, so one selects the other. - use_system_prompt: upstream made it checkpoint-declared, this branch made it mode-dependent. Unset now resolves to "V2V uses it, otherwise the checkpoint decides". The extra-param spec keeps its None default because the executor materializes spec defaults into every request -- publishing the checkpoint's boolean there would destroy "unset" before forward() could see it, silently dropping the system prompt from served V2V. The checkpoint value stays discoverable as default_use_system_prompt. - T2I + audio: keep upstream's force-disable. This branch raised instead, left over from the action work already removed from this PR. - The LPIPS golden archive is unioned: ours carries the V2V reference and upstream's the I2V-4Step one, and taking either side whole would have deleted the other PR's baseline. Also scopes failure classification back to Cosmos3. classify_worker_error declared every ValueError from every pipeline a client error, changing the public failure type for Wan, LTX2, Qwen-Image and Flux -- 130 raise sites this PR never touched. It is now an opt-in BasePipeline hook, classify_request_failure, unclassified by default. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #63120 [ run ] triggered by Bot. Commit: |
|
PR_Github #63120 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63156 [ run ] triggered by Bot. Commit: |
|
PR_Github #63156 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63185 [ run ] triggered by Bot. Commit: |
|
PR_Github #63185 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63188 [ run ] triggered by Bot. Commit: |
|
PR_Github #63188 [ run ] completed with state |
|
✅ LFS objects already in storage (1 file) — no sync needed. These LFS-tracked files are already present in this repository's LFS storage:
|
Brings in Cosmos3 video-to-video (NVIDIA#16155) and 52 other commits. Five conflicts; the substantive one is the scheduler. V2V made flow_shift a public request extra-param, which is incompatible with prebuilding one scheduler per output mode at load time: a scheduler built at load cannot honor a value that arrives with the request. The prebuild is replaced by a lazy cache keyed on the resolved (flow_shift, use_karras_sigmas, stream) triple, so a configuration is built on first use and never rebuilt, a mode that is never served never builds one, and modes that resolve to the same pair share an instance (Edge's video and image tables both declare shift 3.0). set_flow_shift keeps main's name and signature so V2V's callers are untouched; Edge's native flow schedule becomes an implicit use_karras_sigmas=False rather than a second code path. Each stream caches from its own basis, so the audio scheduler is derived from the audio instance rather than the video one - main's test_v2v_rebuilds_the_audio_scheduler_too pins that, and an earlier resolution that derived both from the video basis would have changed behavior for audio-capable checkpoints under V2V. Other resolutions: unions of imports and module-level helpers; the generation-defaults resolution keeps the per-(family, mode) table path and drops main's now-redundant distilled pre-resolution; the LFS golden zip takes main's 17 entries plus Edge's 3. Two of main's V2V fixtures construct a pipeline via __new__ and now set `family`, which forward() reads its mode table through. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
## Summary - extend the existing TensorRT-LLM audiovisual notebook with Nano and Super text-to-video and image-to-video audio examples - add Nano and Super video-to-video examples using the merged multipart `input_reference` API - document the `ffmpeg` requirement for MP4 encoding and synchronized audio muxing - update the shared setup, audiovisual quickstart, and root example index TensorRT-LLM runtime support is available through NVIDIA/TensorRT-LLM#14827 (audio) and NVIDIA/TensorRT-LLM#16155 (V2V). ## Validation - validated notebook JSON, unique cell IDs, and empty execution outputs - compiled all notebook code cells and README Python examples - generated every Nano and Super payload and validated referenced assets - validated audio JSON and image/video multipart curl construction with mocked server responses - `git diff --check` Full GPU inference was not run as part of this documentation-only change. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
15 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.