[None][feat] Cosmos3 Action Support - #15890
Conversation
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (9)
tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py (1)
335-338: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExpose
action_modeas a literal schema value.This is a fixed set of accepted values, and
output_typealready uses the literal form in the same schema. As per coding guidelines, useLiteral[...]instead ofstrwhen a fixed set of values is expected.Proposed schema polish
"action_mode": ExtraParamSchema( - type="str", + type="Literal['policy', 'forward_dynamics', 'inverse_dynamics']", default=None, description="Action generation mode: policy, forward_dynamics, or inverse_dynamics.", ),🤖 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 335 - 338, The `action_mode` field in `defaults.py` is using a generic string schema even though it only accepts a fixed set of values. Update the `ExtraParamSchema` for `action_mode` to use a `Literal[...]` type, matching the existing pattern used by `output_type`, so the accepted values are explicitly constrained to the supported modes.Source: Coding guidelines
tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py (3)
1133-1165: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new action forward arguments.
The Args block still stops at
audio_latents; please add shapes/semantics foraction_latents,action_domain_ids,action_noisy_mask,action_start_frame_offset, andaction_fps. As per coding guidelines, externally called functions should have docstrings and their arguments should be documented.🤖 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/transformer_cosmos3.py` around lines 1133 - 1165, The forward docstring in TransformerCosmos3.forward is missing the newly added action inputs, so update the Args block to document action_latents, action_domain_ids, action_noisy_mask, action_start_frame_offset, and action_fps with their shapes and semantics. Keep the descriptions consistent with the existing audio_latents documentation and ensure the method’s returned action behavior is reflected in the docstring.Source: Coding guidelines
155-175: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winWire or remove
base_temporal_compression_factor.
compute_mrope_position_ids_action()accepts this value and_compute_action_rope_freqs()passesself.temporal_compression_factor, but the helper ignores it and hard-codestemporal_compression_factor=1. If action tokens intentionally use uncompressed frame positions, remove the parameter/call argument; otherwise pass it through so config changes are honored.Also applies to: 1089-1091
🤖 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/transformer_cosmos3.py` around lines 155 - 175, The `compute_mrope_position_ids_action()` helper is ignoring `base_temporal_compression_factor` even though `_compute_action_rope_freqs()` passes `self.temporal_compression_factor`, so config changes are not honored. Either remove the unused parameter from `compute_mrope_position_ids_action()` and its caller if action tokens must always use uncompressed positions, or thread the value through to the `compute_mrope_position_ids_vision()` call instead of hard-coding `temporal_compression_factor=1`.
201-201: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the missing return annotation.
DomainAwareLinear.forward()should declare-> torch.Tensor. As per coding guidelines, “Always annotate functions.”Proposed fix
- def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: + def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor:🤖 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/transformer_cosmos3.py` at line 201, The DomainAwareLinear.forward method is missing its explicit return type annotation. Update the forward signature in transformer_cosmos3.py so that DomainAwareLinear.forward declares a torch.Tensor return type, matching the project’s “Always annotate functions” guideline and keeping the method signature consistent with the surrounding Tensor-returning APIs.Source: Coding guidelines
tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py (1)
392-529: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCoverage gap: transformer-level guard clauses for action are untested.
Cosmos3VFMTransformer.forwardraisesValueErrorfor (1)action_latents+audio_latentssupplied together and (2)action_latentssupplied whenaction_gen=False. Neither is exercised inTestCosmos3Action; only the pipeline-level equivalents are covered intest_cosmos3_pipeline.py. Since the transformer implements its own redundant checks, add two small unit tests here (e.g.,test_forward_action_and_audio_rejected,test_forward_action_latents_without_action_gen_rejected) usingcosmos3_model_config_noactionfor the second case.As per path instructions, tests/** reviewers should flag coverage sufficiency; this is a targeted gap on a guard clause with cross-file evidence in
transformer_cosmos3.py. Want me to draft these two tests?🤖 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 392 - 529, Add two transformer-level coverage tests in TestCosmos3Action for the guard clauses in Cosmos3VFMTransformer.forward: one should pass both action_latents and audio_latents and assert ValueError, and the other should use cosmos3_model_config_noaction and assert ValueError when action_latents is provided with action_gen disabled. Reuse the existing action_model_config/cosmos3_model_config_noaction fixtures and _build_random_weight_model/_cosmos3_inputs helpers so the tests directly exercise Cosmos3VFMTransformer.forward rather than the pipeline.Source: Path instructions
tests/unittest/_torch/visual_gen/test_cosmos3_action.py (1)
85-152: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsolidate repeated inline imports.
resolve_domain_action_config,get_domain_preset, andnormalize_action_resolutionare re-imported inside 6 separate test methods (Lines 87-89, 100-102, 110-112, 123-125, 132-134, 148). Moving these to the top-level import block alongside the existingactionanddefaultsimports would remove the duplication.🤖 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 85 - 152, The test methods repeat inline imports for resolve_domain_action_config, get_domain_preset, and normalize_action_resolution, so consolidate them into the module-level import section used by TestDomainActionPresets. Update the imports near the top of the test file to include the symbols from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults and tensorrt_llm._torch.visual_gen.models.cosmos3.action, then remove the repeated local imports from each test while keeping the assertions in test_bridge_preset_fills_missing_fields, test_av_preset_uses_longer_chunk, test_mismatch_emits_warning, test_action_fps_defaults_to_frame_rate, test_explicit_action_fps_overrides_default, test_alias_maps_to_canonical_preset, and test_unknown_resolution_raises unchanged.tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py (2)
1090-1176: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
post_step_fn'simage_latentbranch is currently unreachable.
post_step_fnis only passed todenoise()whendo_actionis true (line 1205), but everydo_actioncode path explicitly setsimage_latent = None(lines 1005, 1020). This means theelif velocity_mask is not None and image_latent is not None:branch (lines 1159-1163) can never execute under the current wiring — it silently does nothing for the classic I2V path (which doesn't usepost_step_fnat all) and never triggers for the action path either.If this is forward-looking scaffolding for the planned action+audio/I2V merge (per the
FUTURE(action+audio)comment at line 1194), consider a short comment clarifying that intent; otherwise this branch can be removed until it's wired up.🤖 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 1090 - 1176, The image_latent branch inside post_step_fn is unreachable with the current denoise wiring. Review the forward_fn/post_step_fn flow in pipeline_cosmos3.py and either remove the unused elif velocity_mask is not None and image_latent is not None path, or add a brief comment near post_step_fn/denoise() explaining it is reserved for the future action+audio/I2V merge. Keep the logic aligned with the do_action path where image_latent is always None.
739-761: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
action/videoparams use bareAnyunlike the siblingimageparam.
imageis precisely typed asOptional[Union[PIL.Image.Image, torch.Tensor, str]], but the newaction/videoparams fall back toAny. Givenaction.py'snormalize_action_video_input/load_action_tensoronly accept a bounded set of types (str/Path/list/PIL image/torch.Tensor), aUnionalias would be both more precise and consistent with the existingimageparameter's style.♻️ Suggested typing
- action: Any = None, + action: Optional[Union[torch.Tensor, list]] = None, action_resolution: Optional[int] = None, action_fps: Optional[float] = None, - video: Any = None, + video: Optional[Union[PIL.Image.Image, torch.Tensor, str, list]] = None,As per coding guidelines, "Do not use
typing.Anyif avoidable" and "Prefer built-in types... use|syntax instead oftyping.Union."Also applies to: 757-757, 760-760
🤖 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 739 - 761, The new action/video parameters in the Cosmos3 pipeline are typed as bare Any, which is too loose and inconsistent with the existing image parameter. Update the pipeline signature in pipeline_cosmos3.py so action and video use a precise bounded union type matching the accepted inputs from normalize_action_video_input and load_action_tensor (for example str/Path, list, PIL image, and torch.Tensor), using the same modern built-in/| style as image. Keep the typing consistent across action, video, and any related helper signatures so callers get proper type safety and editor support.Source: Coding guidelines
tensorrt_llm/_torch/visual_gen/pipeline.py (1)
1183-1188: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winHoist
inspect.signature()out of the per-step loop.
inspect.signature(post_step_fn)is recomputed on every denoising step even though the callable's arity never changes across the loop. This re-runs reflection unnecessarily on a hot path and is explicitly discouraged by the project's coding guidelines ("Avoid using reflection when functionality can be easily achieved without reflection").⚡ Suggested fix
+ post_step_fn_takes_streams = ( + 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_fn_takes_streams: 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 per-step denoising loop in pipeline.py recomputes inspect.signature(post_step_fn) on every iteration, which is unnecessary reflection on a hot path. Hoist the arity check for post_step_fn out of the loop in the relevant pipeline method, cache whether it accepts 1 or 2 arguments once before stepping begins, and then reuse that decision inside the loop when calling post_step_fn.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 `@examples/visual_gen/models/cosmos3/cosmos3.py`:
- Around line 24-34: The local ACTION_MODES definition in cosmos3.py duplicates
the canonical set from action.py and can drift from the pipeline contract.
Remove the inline frozenset and import ACTION_MODES from
tensorrt_llm._torch.visual_gen.models.cosmos3.action alongside
normalize_action_mode-related usage, then update any CLI choices or validation
in this module to reference the imported constant so the command-line behavior
stays consistent with the canonical action definitions.
- Around line 370-388: The `params.extra_params["action_resolution"]` assignment
in `cosmos3.py` is unconditional and can overwrite the domain-resolved value
from `_apply_action_generation_params` or inject a `None` default when
`action_mode` is absent. Update the `action_resolution` wiring to follow the
same guarded pattern as the sibling fields in `build_args`/parameter assembly:
only set it when `args.action_mode is not None` and `args.action_resolution is
not None`, and preserve the already resolved `cfg["action_resolution"]`
otherwise. Make sure `params.extra_params` stays consistent with the
`width`/`height` derived from the resolved action settings.
- Around line 143-164: The action parameter setup in
_apply_action_generation_params is hard-coding the canvas to the 16:9 entry from
VIDEO_RES_SIZE_INFO, which bypasses aspect-aware sizing. Update this path to
derive width and height from the reference frame or select the matching aspect
bucket based on the chosen action_resolution so portrait inputs are handled
correctly while keeping the rest of the cfg-driven settings unchanged.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py`:
- Around line 402-411: Forward-dynamics action inputs are currently being
silently padded or truncated when their feature dimension does not match the
configured raw action contract. In the ACTION_MODE_FORWARD_DYNAMICS path in
action.py, add an explicit validation around load_action_tensor,
pad_action_to_dim, and the raw_action_dim handling so that if raw_action_dim is
already set and action.shape[-1] differs from it, the code fails fast instead of
treating missing columns as conditioned dimensions or zeroing extra columns
later. Apply the same check in the related forward-dynamics branches referenced
by the surrounding logic so the contract is enforced consistently before any
padding or slicing.
- Around line 333-364: The action_reference_image() helper only treats string
paths as file inputs, so Path objects fall through to the TypeError even though
normalize_action_video_input() already accepts them. Update
action_reference_image() to handle Path alongside str for both the image and
video fallback paths, preserving the existing PIL.Image.Image behavior and
file-extension check before opening or normalizing the source.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py`:
- Around line 285-298: The default resolution logic in the Cosmos3 defaults
builder should validate timing/length fields before returning them. In the
function that computes resolved values and returns the config dict, add
positive-value checks for action_chunk_size, num_frames, frame_rate, and
action_fps after _resolve_field / fallback resolution and before the final
return, so zero or negative overrides are rejected early with a clear error
instead of propagating into downstream latent/mRoPE setup.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py`:
- Around line 817-830: The action projection layers are being initialized and
preserved in a hardcoded BF16 dtype, which can cause dtype mismatches in
DomainAwareLinear.forward when the model runs in fp16 or fp32. Update the
initialization in the transformer_cosmos3 model to use the configured model
dtype (model_config.torch_dtype) for action_proj_in, action_proj_out, and
action_modality_embed, and make post_load_weights() retain that same dtype
instead of forcing BF16.
In `@tensorrt_llm/visual_gen/output.py`:
- Around line 99-104: The tensor payload serialization currently omits the new
action tensor, so save/load round-trips drop it. Update the serialization path
in tensor_payload.py by adding action to _modalities() and ensuring the
batch-size and slicing logic in the tensor payload handling includes action
alongside image/video/audio. Also verify the output dataclass in output.py and
any payload assembly code use the action field consistently so save(...,
fmt="safetensors"|"pt") preserves it.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py`:
- Around line 335-338: The `action_mode` field in `defaults.py` is using a
generic string schema even though it only accepts a fixed set of values. Update
the `ExtraParamSchema` for `action_mode` to use a `Literal[...]` type, matching
the existing pattern used by `output_type`, so the accepted values are
explicitly constrained to the supported modes.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`:
- Around line 1090-1176: The image_latent branch inside post_step_fn is
unreachable with the current denoise wiring. Review the forward_fn/post_step_fn
flow in pipeline_cosmos3.py and either remove the unused elif velocity_mask is
not None and image_latent is not None path, or add a brief comment near
post_step_fn/denoise() explaining it is reserved for the future action+audio/I2V
merge. Keep the logic aligned with the do_action path where image_latent is
always None.
- Around line 739-761: The new action/video parameters in the Cosmos3 pipeline
are typed as bare Any, which is too loose and inconsistent with the existing
image parameter. Update the pipeline signature in pipeline_cosmos3.py so action
and video use a precise bounded union type matching the accepted inputs from
normalize_action_video_input and load_action_tensor (for example str/Path, list,
PIL image, and torch.Tensor), using the same modern built-in/| style as image.
Keep the typing consistent across action, video, and any related helper
signatures so callers get proper type safety and editor support.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py`:
- Around line 1133-1165: The forward docstring in TransformerCosmos3.forward is
missing the newly added action inputs, so update the Args block to document
action_latents, action_domain_ids, action_noisy_mask, action_start_frame_offset,
and action_fps with their shapes and semantics. Keep the descriptions consistent
with the existing audio_latents documentation and ensure the method’s returned
action behavior is reflected in the docstring.
- Around line 155-175: The `compute_mrope_position_ids_action()` helper is
ignoring `base_temporal_compression_factor` even though
`_compute_action_rope_freqs()` passes `self.temporal_compression_factor`, so
config changes are not honored. Either remove the unused parameter from
`compute_mrope_position_ids_action()` and its caller if action tokens must
always use uncompressed positions, or thread the value through to the
`compute_mrope_position_ids_vision()` call instead of hard-coding
`temporal_compression_factor=1`.
- Line 201: The DomainAwareLinear.forward method is missing its explicit return
type annotation. Update the forward signature in transformer_cosmos3.py so that
DomainAwareLinear.forward declares a torch.Tensor return type, matching the
project’s “Always annotate functions” guideline and keeping the method signature
consistent with the surrounding Tensor-returning APIs.
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 1183-1188: The per-step denoising loop in pipeline.py recomputes
inspect.signature(post_step_fn) on every iteration, which is unnecessary
reflection on a hot path. Hoist the arity check for post_step_fn out of the loop
in the relevant pipeline method, cache whether it accepts 1 or 2 arguments once
before stepping begins, and then reuse that decision inside the loop when
calling post_step_fn.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_action.py`:
- Around line 85-152: The test methods repeat inline imports for
resolve_domain_action_config, get_domain_preset, and
normalize_action_resolution, so consolidate them into the module-level import
section used by TestDomainActionPresets. Update the imports near the top of the
test file to include the symbols from
tensorrt_llm._torch.visual_gen.models.cosmos3.defaults and
tensorrt_llm._torch.visual_gen.models.cosmos3.action, then remove the repeated
local imports from each test while keeping the assertions in
test_bridge_preset_fills_missing_fields, test_av_preset_uses_longer_chunk,
test_mismatch_emits_warning, test_action_fps_defaults_to_frame_rate,
test_explicit_action_fps_overrides_default, test_alias_maps_to_canonical_preset,
and test_unknown_resolution_raises unchanged.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py`:
- Around line 392-529: Add two transformer-level coverage tests in
TestCosmos3Action for the guard clauses in Cosmos3VFMTransformer.forward: one
should pass both action_latents and audio_latents and assert ValueError, and the
other should use cosmos3_model_config_noaction and assert ValueError when
action_latents is provided with action_gen disabled. Reuse the existing
action_model_config/cosmos3_model_config_noaction fixtures and
_build_random_weight_model/_cosmos3_inputs helpers so the tests directly
exercise Cosmos3VFMTransformer.forward rather than the pipeline.
🪄 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: 4e52edce-bddb-4a9d-8e25-06d48d0d0a40
📒 Files selected for processing (16)
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.jsontensorrt_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/output.pytensorrt_llm/_torch/visual_gen/pipeline.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.py
📝 WalkthroughWalkthroughAdds action generation support (policy, forward_dynamics, inverse_dynamics modes) to the Cosmos3 visual generation model, including new domain-aware action preprocessing utilities, transformer action-token injection, pipeline scheduler/denoise wiring, output dataclass extensions, tensor payload serialization support, a CLI example, and unit/multi-GPU tests. ChangesCosmos3 Action Generation
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant CLI as cosmos3.py
participant Pipeline as Cosmos3OmniMoTPipeline
participant Transformer as Cosmos3VFMTransformer
participant Output as PipelineOutput
CLI->>Pipeline: infer(action_mode, domain, action_json, video_path)
Pipeline->>Pipeline: prepare_action_latents(action_input)
Pipeline->>Transformer: forward(action_latents, action_domain_ids)
Transformer-->>Pipeline: TransformerOutput(action=action_noise_pred)
Pipeline->>Pipeline: denoise loop updates action_latents via post_step_fn
Pipeline->>Output: PipelineOutput(action, raw_action_dim, action_mode, domain_id)
Output-->>CLI: VisualGenOutput.action
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/visual_gen/output.py (1)
95-106: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd an explicit
action-only branch tosave()'s non-tensor dispatch.New
actionfield means an output can be action-only (no image/video/audio). The existing dispatch insave()(further down, ~lines 211-217) doesn't checkself.action, so an action-only output saved with a non-tensor format falls through toValueError("... carries no media (image/video/audio are all None)"), which is misleading — it does carry media, just not one the encoder path supports.🐛 Proposed fix
if self.audio is not None: raise NotImplementedError("Saving audio-only outputs is not supported in this release.") + if self.action is not None: + raise NotImplementedError( + "Saving action-only outputs requires a tensor format " + "(safetensors/pt); non-tensor image/video encoders are not supported." + ) + raise ValueError( f"Cannot save output: request {self.request_id} carries no media " - "(image/video/audio are all None)." + "(image/video/audio/action are all None)." )🤖 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/output.py` around lines 95 - 106, The non-tensor dispatch in VisualGenOutput.save() currently treats action-only outputs as having no media, which triggers the wrong ValueError. Update the save() branch logic to explicitly check self.action alongside image/video/audio, and add an action-only path or a clear unsupported-media branch before the existing “no media” error. Use the VisualGenOutput.save method and the self.action field to locate and adjust the dispatch.tensorrt_llm/media/tensor_payload.py (1)
63-69: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAction metadata (
raw_action_dim/action_mode/domain_id) isn't preserved in the tensor payload.
_modalities()now includes theactiontensor, so the tensor values round-trip correctly (confirmed by the new tests). But_collect_tensors_and_metadata(further down, ~lines 102-150) still only forwardsoutput.frame_rate/output.audio_sample_rateintometadata.output.raw_action_dim,output.action_mode, andoutput.domain_idare silently dropped on save, so a loadedsafetensors/ptpayload has the action trajectory but no way to know which mode/domain it was generated for.Note:
action_modeis a string, so it can't go through the safetensors branch'storch.as_tensor(v)scalar conversion unchanged — it needs to land only in the stringmetadata=header (or a dedicated non-tensor field), whileraw_action_dim/domain_id(ints) can follow the existing numeric pattern.🐛 Proposed fix (in `_collect_tensors_and_metadata`)
if frame_rate is not None: metadata["frame_rate"] = float(frame_rate) if audio_sample_rate is not None: metadata["audio_sample_rate"] = int(audio_sample_rate) + if output.raw_action_dim is not None: + metadata["raw_action_dim"] = int(output.raw_action_dim) + if output.domain_id is not None: + metadata["domain_id"] = int(output.domain_id) + if output.action_mode is not None: + metadata["action_mode"] = str(output.action_mode) return tensors, metadataAnd in
serialize_visual_gen_output's safetensors branch, exclude string-valued metadata from thescalar_tensorsconversion so it's only stored in the string header:- scalar_tensors = {k: torch.as_tensor(v) for k, v in metadata.items()} + scalar_tensors = { + k: torch.as_tensor(v) for k, v in metadata.items() if not isinstance(v, str) + } return safetensors_save( {**tensors, **scalar_tensors}, metadata={k: str(v) for k, v in metadata.items()}, )Also applies to: 72-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/media/tensor_payload.py` around lines 63 - 69, The tensor payload serialization is dropping action-related metadata, so `_collect_tensors_and_metadata` should be updated to preserve `raw_action_dim`, `action_mode`, and `domain_id` alongside the existing frame/audio fields. In `serialize_visual_gen_output`, make sure `action_mode` is written only to the safetensors string metadata header, while numeric fields like `raw_action_dim` and `domain_id` continue through the existing tensor/scalar metadata path. Keep the fix localized around `_modalities`, `_collect_tensors_and_metadata`, and the safetensors branch so loaded payloads retain both the action tensor and its generation context.
🧹 Nitpick comments (10)
tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py (2)
335-339: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
Literal[...]type string foraction_mode, consistent withoutput_type.
output_typeabove (Line 331) documents its fixed value set viatype="Literal['video', 'image']".action_modeaccepts only the three values inACTION_MODESbut is declared as genericstr, losing that self-documenting/tooling-friendly contract.♻️ Proposed fix
"action_mode": ExtraParamSchema( - type="str", + type=f"Literal[{', '.join(sorted(repr(m) for m in ACTION_MODES))}]", default=None, description="Action generation mode: policy, forward_dynamics, or inverse_dynamics.", ),🤖 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 335 - 339, The action_mode schema is too loosely typed as a generic string in the defaults definition, even though it only supports a fixed set of values. Update the ExtraParamSchema for action_mode in the Cosmos3 defaults to use a Literal[...] type string matching ACTION_MODES, following the same pattern used by output_type so the contract is explicit and tooling-friendly.
29-35: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPrefer builtin generics/
|overtyping.Dict/List/Optionalper coding guidelines.New code imports and uses
Dict,List,Optionalthroughout (e.g.Dict[str, Cosmos3DomainPreset],List[str],Optional[int]) instead ofdict,list,X | None. Also several params (domain_name: Any,domain_id: Any) could be narrowed tostr | int | Nonerather thanAny.♻️ Example fix
-from typing import Any, Dict, List, Optional, TypedDict +from typing import Any, TypedDict ... -COSMOS3_DOMAIN_PRESETS: Dict[str, Cosmos3DomainPreset] = { +COSMOS3_DOMAIN_PRESETS: dict[str, Cosmos3DomainPreset] = { ... -COSMOS3_DOMAIN_PRESET_ALIASES: Dict[str, str] = { +COSMOS3_DOMAIN_PRESET_ALIASES: dict[str, str] = { ... def resolve_domain_action_config( *, domain_name: Any = None, domain_id: Any = None, - raw_action_dim: Optional[int] = None, - action_chunk_size: Optional[int] = None, - action_resolution: Optional[int] = None, - frame_rate: Optional[float] = None, - action_fps: Optional[float] = None, - num_frames: Optional[int] = None, -) -> Dict[str, Any]: + raw_action_dim: int | None = None, + action_chunk_size: int | None = None, + action_resolution: int | None = None, + frame_rate: float | None = None, + action_fps: float | None = None, + num_frames: int | None = None, +) -> dict[str, Any]:As per coding guidelines: "Prefer built-in types
list,dict,tupleovertyping.List,typing.Dict,typing.Tuple; use|syntax instead oftyping.Union" and "do not usetyping.Anyif avoidable."Also applies to: 192-301
🤖 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 29 - 35, Update the Cosmos3 defaults typing to follow the coding guidelines by replacing `typing.Dict/List/Optional` with builtin generics and `| None` throughout `defaults.py`. In the affected constants, type aliases, and function signatures, use `dict`, `list`, and `str | int | None` instead of `Any` where the domain parameters are constrained, and adjust any related annotations in `Cosmos3DomainPreset`/preset helpers so the public API stays consistent. Focus on the imports and the definitions around `normalize_action_resolution`, `COSMOS3_ACTION_RESOLUTIONS`, and the preset mappings where `Dict[str, Cosmos3DomainPreset]`, `List[str]`, and `Optional[int]` are currently used.Source: Coding guidelines
tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py (2)
391-401: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a docstring to
prepare_action_latents.This is a non-trivial public function (imported by pipeline/transformer layers) returning a 4-tuple (
action_latents,action_velocity_mask,clean_action,raw_action_dim) whose meaning isn't obvious from the signature alone. A short Google-style docstring documenting args, return values, and per-mode conditioning behavior would help downstream maintainers.As per coding guidelines: "Externally called functions should have docstrings, and their arguments should be documented."
🤖 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 391 - 401, `prepare_action_latents` is a public helper without documentation; add a short Google-style docstring to this function. Document the arguments (`mode`, `action_chunk_size`, `raw_action_dim`, `action_dim`, `generator`, `device`, `dtype`, `action_input`) and clearly describe the returned 4-tuple (`action_latents`, `action_velocity_mask`, `clean_action`, `raw_action_dim`). Also note the per-mode conditioning behavior so callers of `prepare_action_latents` can understand how inputs are interpreted.Source: Coding guidelines
8-9: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMixed
typing.List/Optionaland builtin generics; prefer consistent builtin style.The file imports
List/Optionalfromtyping(Line 9) and uses them in several signatures (e.g. Lines 261, 296, 315-318), while other annotations already use builtin generics (dict[str, int]Line 25,list[int]Line 138). Per coding guidelines, preferlist/dict/tuple/X | Noneuniformly.As per coding guidelines: "Prefer built-in types
list,dict, andtupleto the legacytyping.List,typing.Dict, andtyping.Tuple."Also applies to: 25-25, 45-45, 138-138
🤖 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 8 - 9, Mixed typing style is used in action.py, with legacy List/Optional imported from typing alongside builtin generics. Update the relevant annotations in the affected module, especially in the action-related signatures and any dict/list/tuple annotations, to use builtin generics consistently (list, dict, tuple, and X | None) and remove the typing.List/Optional import if it becomes unused. Use the existing symbols in this file, including the affected action methods and the module-level type annotations, to make the changes consistently across the file.Source: Coding guidelines
tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py (1)
496-551: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winInconsistent post-forward assertions across action-mode tests.
test_policy_smokeassertsresult.action_mode == "policy"andresult.domain_id == 7(Lines 493-494), buttest_forward_dynamics_smokeandtest_inverse_dynamics_smokeomit these checks even though the same output fields should be populated for these modes. Adding the equivalent assertions would strengthen regression coverage foraction_mode/domain_idpropagation across all three modes.🤖 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 496 - 551, The action-mode smoke tests are missing the same post-forward metadata checks used in test_policy_smoke. Update test_forward_dynamics_smoke and test_inverse_dynamics_smoke in test_cosmos3_pipeline to assert that the returned result includes the expected action_mode for each case and that result.domain_id is set to 7, so coverage matches the other action-mode path and verifies propagation through _run_forward.Source: Path instructions
tests/unittest/_torch/visual_gen/test_cosmos3_action.py (2)
85-151: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsolidate repeated local imports.
resolve_domain_action_config,get_domain_preset, andnormalize_action_resolutionare each re-imported inside individual test methods (Lines 87-89, 100-102, 110-112, 123-125, 132-134, 141, 148) instead of once at module scope alongside the other imports at Lines 14-21. This is repetitive and inconsistent with the rest of the file's import style.♻️ Proposed fix
from tensorrt_llm._torch.visual_gen.models.cosmos3.action import ( VIDEO_RES_SIZE_INFO, action_reference_image, find_closest_target_size, normalize_action_video_input, + normalize_action_resolution, resolve_action_size, ) -from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_EXTRA_SPECS, + get_domain_preset, + resolve_domain_action_config, +)Then remove the per-test
from ... import ...lines.🤖 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 85 - 151, The test methods are redundantly re-importing resolve_domain_action_config, get_domain_preset, and normalize_action_resolution inside each test, which is inconsistent with the module’s import style. Move these imports to the top-level import section of TestDomainActionPresets’s module alongside the existing imports, then remove the local import statements from each affected test method. Keep the tests themselves using the imported symbols directly.
186-254: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCoverage gap: untested
normalize_action_video_inputbranches.Per the upstream contract,
normalize_action_video_inputreturns[]forNoneand raisesValueErrorfor an empty list, but neither branch is exercised here. Suggest adding:
test_none_returns_empty_list— assertsnormalize_action_video_input(None) == []test_empty_list_raises— assertsValueErrorfornormalize_action_video_input([])Coverage is otherwise good (image path, frame directory, extension validation, decode + max_frames), so this is the only notable gap.
🤖 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 186 - 254, Add tests for the missing branches in normalize_action_video_input: create one test that verifies passing None returns an empty list, and another that verifies passing an empty list raises ValueError. Place these alongside the existing TestNormalizeActionVideoInput cases so the behavior of normalize_action_video_input is fully covered, including the None and empty-sequence inputs.Source: Path instructions
tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py (2)
438-528: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCoverage: add negative/boundary-path tests for action forward.
The forward-pass tests only cover the happy path (valid
action_latentsshape, validdomain_ids). Consider adding:
- a case where
action_latentslast-dim doesn't matchmodel.action_dim, asserting theValueErrorraised bypack_action(per the transformer's contract).- a case with
action_domain_idsout of range (e.g.,>= num_embodiment_domains) to confirm expected error behavior rather than silent out-of-bounds embedding lookup.🤖 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 438 - 528, The action forward tests in test_cosmos3_transformer.py only cover valid inputs; add negative-path coverage around Cosmos3 transformer action handling. Extend the existing action tests to verify that pack_action raises a ValueError when action_latents has a last dimension different from model.action_dim, and add a separate test that passes out-of-range action_domain_ids (at or above num_embodiment_domains) to confirm the expected error instead of allowing a silent embedding lookup failure.Source: Path instructions
392-437: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCoverage: broaden negative structural check.
test_video_only_model_has_no_action_headsonly checksaction_proj_inis absent (Line 436). Sinceaction_proj_in,action_proj_out, andaction_modality_embedare all initialized together whenaction_genis enabled, asserting their absence together would better guard against a partial-initialization regression.🤖 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 392 - 437, The negative structural check in TestCosmos3Action.test_video_only_model_has_no_action_heads is too narrow because it only verifies action_proj_in is missing. Update this test to also assert the absence of action_proj_out and action_modality_embed on Cosmos3VFMTransformer when action_gen is disabled, matching the grouped initialization used in test_action_model_structure.Source: Path instructions
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 is fixed for the whole call todenoise(), but the signature is currently re-inspected on every denoising step. Compute it once before the loop.♻️ Proposed refactor
+ post_step_fn_takes_extra_streams = ( + 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_fn_takes_extra_streams: latents, extra_stream_latents = post_step_fn(latents, extra_stream_latents) else: latents = post_step_fn(latents)🤖 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 `denoise()` path in `pipeline.py` is re-checking `inspect.signature(post_step_fn)` on every denoising iteration even though `post_step_fn`’s arity is constant for the whole call; compute the signature once before the loop and reuse it inside the step body. Update the logic around `post_step_fn` in `denoise()` so the parameter count is determined up front, then keep the existing branch that calls `post_step_fn(latents, extra_stream_latents)` or `post_step_fn(latents)` based on that cached result.
🤖 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/defaults.py`:
- Around line 233-247: Unresolved domain_name/domain_id currently falls back to
generic defaults without surfacing any diagnostic. In
resolve_domain_action_config(), after canonical_domain_preset_key() and preset
lookup, add a warning path when no preset is found so typos or unknown
embodiments are reported through the existing warnings list. Keep the behavior
of using generic defaults, but make the fallback explicit by appending a clear
message before returning the config, alongside the existing mismatch warnings
logic.
- Around line 362-383: `action_chunk_size` and `action_resolution` are being
prefilled as explicit values, which blocks Cosmos3 preset inference in
`resolve_domain_action_config()`. Update the `ExtraParamSchema` defaults in the
Cosmos3 defaults spec so these two fields default to `None`, and keep their
descriptions indicating they are inferred from the domain preset when omitted.
Make sure the surrounding `extra_param_specs` entries for `action_chunk_size`
and `action_resolution` preserve their validation/range metadata while allowing
preset values to populate them later during `infer()`.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py`:
- Around line 856-863: The action-reference path in forward() is using
action_reference_image() before tensor inputs are handled, so torch.Tensor
images can fail when do_action is enabled. Update the Cosmos3 pipeline logic to
either convert tensor image/video inputs into a supported format before calling
action_reference_image(), or explicitly reject tensor-based action requests with
a clear error; keep the fix near forward() and the
action_reference_image()/resolve_action_size() flow.
- Around line 995-999: In the ACTION_MODE_INVERSE_DYNAMICS path of
`pipeline_cosmos3.py`, `normalize_action_video_input(...,
max_frames=num_frames)` can still return too-short clips, so add an explicit
length check before calling `_preprocess_action_video` or
`_prepare_latents_action_video`. Use the `inverse_video`/`video` handling in
`normalized_action_mode == ACTION_MODE_INVERSE_DYNAMICS` to verify the
normalized video has at least `num_frames` frames and raise an error if it does
not.
---
Outside diff comments:
In `@tensorrt_llm/media/tensor_payload.py`:
- Around line 63-69: The tensor payload serialization is dropping action-related
metadata, so `_collect_tensors_and_metadata` should be updated to preserve
`raw_action_dim`, `action_mode`, and `domain_id` alongside the existing
frame/audio fields. In `serialize_visual_gen_output`, make sure `action_mode` is
written only to the safetensors string metadata header, while numeric fields
like `raw_action_dim` and `domain_id` continue through the existing
tensor/scalar metadata path. Keep the fix localized around `_modalities`,
`_collect_tensors_and_metadata`, and the safetensors branch so loaded payloads
retain both the action tensor and its generation context.
In `@tensorrt_llm/visual_gen/output.py`:
- Around line 95-106: The non-tensor dispatch in VisualGenOutput.save()
currently treats action-only outputs as having no media, which triggers the
wrong ValueError. Update the save() branch logic to explicitly check self.action
alongside image/video/audio, and add an action-only path or a clear
unsupported-media branch before the existing “no media” error. Use the
VisualGenOutput.save method and the self.action field to locate and adjust the
dispatch.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/action.py`:
- Around line 391-401: `prepare_action_latents` is a public helper without
documentation; add a short Google-style docstring to this function. Document the
arguments (`mode`, `action_chunk_size`, `raw_action_dim`, `action_dim`,
`generator`, `device`, `dtype`, `action_input`) and clearly describe the
returned 4-tuple (`action_latents`, `action_velocity_mask`, `clean_action`,
`raw_action_dim`). Also note the per-mode conditioning behavior so callers of
`prepare_action_latents` can understand how inputs are interpreted.
- Around line 8-9: Mixed typing style is used in action.py, with legacy
List/Optional imported from typing alongside builtin generics. Update the
relevant annotations in the affected module, especially in the action-related
signatures and any dict/list/tuple annotations, to use builtin generics
consistently (list, dict, tuple, and X | None) and remove the
typing.List/Optional import if it becomes unused. Use the existing symbols in
this file, including the affected action methods and the module-level type
annotations, to make the changes consistently across the file.
In `@tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py`:
- Around line 335-339: The action_mode schema is too loosely typed as a generic
string in the defaults definition, even though it only supports a fixed set of
values. Update the ExtraParamSchema for action_mode in the Cosmos3 defaults to
use a Literal[...] type string matching ACTION_MODES, following the same pattern
used by output_type so the contract is explicit and tooling-friendly.
- Around line 29-35: Update the Cosmos3 defaults typing to follow the coding
guidelines by replacing `typing.Dict/List/Optional` with builtin generics and `|
None` throughout `defaults.py`. In the affected constants, type aliases, and
function signatures, use `dict`, `list`, and `str | int | None` instead of `Any`
where the domain parameters are constrained, and adjust any related annotations
in `Cosmos3DomainPreset`/preset helpers so the public API stays consistent.
Focus on the imports and the definitions around `normalize_action_resolution`,
`COSMOS3_ACTION_RESOLUTIONS`, and the preset mappings where `Dict[str,
Cosmos3DomainPreset]`, `List[str]`, and `Optional[int]` are currently used.
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 1183-1188: The `denoise()` path in `pipeline.py` is re-checking
`inspect.signature(post_step_fn)` on every denoising iteration even though
`post_step_fn`’s arity is constant for the whole call; compute the signature
once before the loop and reuse it inside the step body. Update the logic around
`post_step_fn` in `denoise()` so the parameter count is determined up front,
then keep the existing branch that calls `post_step_fn(latents,
extra_stream_latents)` or `post_step_fn(latents)` based on that cached result.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_action.py`:
- Around line 85-151: The test methods are redundantly re-importing
resolve_domain_action_config, get_domain_preset, and normalize_action_resolution
inside each test, which is inconsistent with the module’s import style. Move
these imports to the top-level import section of TestDomainActionPresets’s
module alongside the existing imports, then remove the local import statements
from each affected test method. Keep the tests themselves using the imported
symbols directly.
- Around line 186-254: Add tests for the missing branches in
normalize_action_video_input: create one test that verifies passing None returns
an empty list, and another that verifies passing an empty list raises
ValueError. Place these alongside the existing TestNormalizeActionVideoInput
cases so the behavior of normalize_action_video_input is fully covered,
including the None and empty-sequence inputs.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py`:
- Around line 496-551: The action-mode smoke tests are missing the same
post-forward metadata checks used in test_policy_smoke. Update
test_forward_dynamics_smoke and test_inverse_dynamics_smoke in
test_cosmos3_pipeline to assert that the returned result includes the expected
action_mode for each case and that result.domain_id is set to 7, so coverage
matches the other action-mode path and verifies propagation through
_run_forward.
In `@tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py`:
- Around line 438-528: The action forward tests in test_cosmos3_transformer.py
only cover valid inputs; add negative-path coverage around Cosmos3 transformer
action handling. Extend the existing action tests to verify that pack_action
raises a ValueError when action_latents has a last dimension different from
model.action_dim, and add a separate test that passes out-of-range
action_domain_ids (at or above num_embodiment_domains) to confirm the expected
error instead of allowing a silent embedding lookup failure.
- Around line 392-437: The negative structural check in
TestCosmos3Action.test_video_only_model_has_no_action_heads is too narrow
because it only verifies action_proj_in is missing. Update this test to also
assert the absence of action_proj_out and action_modality_embed on
Cosmos3VFMTransformer when action_gen is disabled, matching the grouped
initialization used in test_action_model_structure.
🪄 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: 40aa992a-ef70-4c1b-9575-df862ca13634
📒 Files selected for processing (18)
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.jsontensorrt_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/output.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/media/tensor_payload.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_tensor_payload.py
7e26d00 to
6519278
Compare
|
/bot run |
|
PR_Github #59007 [ run ] triggered by Bot. Commit: |
|
PR_Github #59007 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59085 [ run ] triggered by Bot. Commit: |
|
PR_Github #59085 [ run ] completed with state
|
|
/bot run |
|
PR_Github #59243 [ run ] triggered by Bot. Commit: |
|
PR_Github #59243 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59500 [ run ] triggered by Bot. Commit: |
|
PR_Github #59500 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59577 [ run ] triggered by Bot. Commit: |
| from tensorrt_llm._torch.visual_gen.models.cosmos3.action import VIDEO_RES_SIZE_INFO | ||
| from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( |
There was a problem hiding this comment.
Examples should avoid importing internal components (_torch.visual_gen.*), otherwise we don't have a bar of what's committed to be stable. Can the domain-preset resolution move behind the public API so this example only needs tensorrt_llm.VisualGen?
There was a problem hiding this comment.
Fixed: the example only imports public VisualGen/VisualGenArgs now; domain-preset resolution stays in the pipeline.
| "action_chunk_size": ExtraParamSchema( | ||
| type="int", | ||
| default=COSMOS3_ACTION_PARAMS["action_chunk_size"], | ||
| description=( | ||
| "Number of action tokens to generate (16 for most robots, 60 for av/camera_pose). " | ||
| "Inferred from domain_name preset when omitted." | ||
| ), | ||
| ), | ||
| "action": ExtraParamSchema( | ||
| type="list", | ||
| default=None, | ||
| description="Action trajectory [T, D] for forward_dynamics mode.", | ||
| ), | ||
| "action_resolution": ExtraParamSchema( | ||
| type="int", | ||
| default=480, |
There was a problem hiding this comment.
These materialized defaults defeat the domain presets: action_chunk_size=16 / action_resolution=480 reach resolve_domain_action_config as non-None, so domain_name="av" keeps chunk 16 (warning only) instead of the preset's 60. Should these default to None and let the resolver fill preset then generic fallback, as the module docstring describes?
There was a problem hiding this comment.
Fixed: both defaults are None now; resolver fills domain preset values first, then generic fallback.
| image: Optional[torch.Tensor] = None | ||
| video: Optional[torch.Tensor] = None | ||
| audio: Optional[torch.Tensor] = None | ||
| action: Optional[torch.Tensor] = None |
There was a problem hiding this comment.
An action request over /v1/videos still takes the video save path, so the default serve flow reports success while silently dropping action (only safetensors/pt formats carry it). Should action requests fail loudly on non-tensor formats instead?
There was a problem hiding this comment.
Fixed: VisualGenOutput.save() now raises for action-bearing outputs unless the format is safetensors/pt.
| raw_action_dim: Optional[int] = None | ||
| action_mode: Optional[str] = None | ||
| domain_id: Optional[int] = None |
There was a problem hiding this comment.
raw_action_dim duplicates action.shape[-1], and action_mode/domain_id echo request inputs in Cosmos3 vocabulary — do these need to be top-level VisualGenOutput fields? The tensor payload also only serializes frame_rate/audio_sample_rate, so they don't round-trip through serve anyway. Public-surface changes here should go through a team sync per ENGINEERING_CRITERIA §2.
And, should we document the "action" semantic? It seems to me it's not a well known representation.
There was a problem hiding this comment.
Kept the fields for action context, and tensor payloads now preserve raw_action_dim/domain_id plus action_mode in the safetensors header.
| "video": ExtraParamSchema( | ||
| type="path_or_list", |
There was a problem hiding this comment.
path_or_list isn't in the params validator's _TYPE_MAP, so this skips preflight type checking and only fails in the background worker (same for action_mode as free str, and the discrete resolution set expressed as a continuous range). Also, over trtllm-serve this accepts a raw server-local path from the client, while input_reference routes through media storage — should the serve path do the same here?
There was a problem hiding this comment.
Fixed: Literal validation covers action_mode/action_resolution, path_or_list is typed, and serve rejects extra_params["video"] raw paths.
| sig = inspect.signature(post_step_fn) | ||
| if len(sig.parameters) >= 2: |
There was a problem hiding this comment.
Arity-sniffing via inspect.signature breaks on functools.partial / **kwargs callables. Does it make sense to keep a single hook signature (always take and return the extra-streams dict)?
There was a problem hiding this comment.
Fixed: denoise post_step_fn now has one fixed signature: (latents, extra_stream_latents) -> (latents, extra_stream_latents).
| @@ -0,0 +1,254 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
New test file isn't referenced in any test list — should it be added to tests/integration/test_lists/test-db/l0_*.yml?
There was a problem hiding this comment.
Fixed: added unittest/_torch/visual_gen/test_cosmos3_action.py to l0_b200.yml.
|
PR_Github #59577 [ run ] completed with state
|
| "audio_sample_rate", | ||
| "pre_denoise", | ||
| "denoise", | ||
| "post_denoise", |
There was a problem hiding this comment.
This assertion will fail. PipelineOutput (in _torch/visual_gen/output.py) is a plain dataclass with no base class, and its actual fields are:
{image, video, audio, action, frame_rate, audio_sample_rate, raw_action_dim, action_mode, domain_id, pre_denoise, denoise, post_denoise}.The asserted set here adds request_id, error, metrics (these three belong to VisualGenOutput, not PipelineOutput) and drops pre_denoise, denoise, post_denoise. So field_names == {...} is always False and the test raises AssertionError.It's also self-contradictory with test_pipeline_output_default_construction below, which asserts p.pre_denoise == 0.0 / denoise / post_denoise — confirming those three fields do exist.
Please assert the real PipelineOutput fields (keep the three timing fields, drop request_id/error/metrics) and update the test name and docstring from "eight" to the actual count (12).
There was a problem hiding this comment.
Fixed: the test now asserts PipelineOutput's 12 real fields and the name/docstring match that count.
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
LGTM on the code. Verified the one correctness question (per-domain action projections): the checkpoint uses action_proj_*.fc.weight Embedding-style keys, which match the DomainAwareLinear (input_size, output_size) layout — so no transpose is needed and load_weights is correct. VisualGenOutput additions are additive/backward-compatible; the default_generation_params→resolve-in-forward() refactor looks consistent for existing T2V/T2I.
One thing before merge: the DCO check is red — please re-sign your commits and force-push:
git rebase --exec 'git commit --amend -s --no-edit' origin/main
git push -f
(or git commit --amend -s if it's a single commit).
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
This reverts commit c115ea6. Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
818f5c0 to
6bf9ba9
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #60389 [ run ] triggered by Bot. Commit: |
|
PR_Github #60389 [ run ] completed with state
|
mikeiovine
left a comment
mikeiovine
left a comment
There was a problem hiding this comment.
Stamp for runtime devs, did not review visualgen stuff
|
you can close this one now. replaced by PR17325. |
8 participants
Summary by CodeRabbit
Description
Adds Cosmos3 action-generation support for policy, forward-dynamics, and inverse-dynamics workflows. The change wires domain-aware action presets through the Cosmos3 pipeline, injects/decodes action tokens in the transformer, preserves action tensors through VisualGen outputs and tensor payloads, and updates the Cosmos3 example to use tensor payloads for action runs so action data is not dropped.
Reviewer follow-up fixes in this revision tighten preset inference, extra-param validation, inverse-video length checks, action metadata serialization, non-tensor save behavior, and DCO coverage.
Test Coverage
tests/unittest/_torch/visual_gen/test_cosmos3_action.pytotests/integration/test_lists/test-db/l0_b200.yml.ruff check,ruff format --check, and targetedpy_compileon edited files.mpi4py/tensorrtimports.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.