[TRTLLM-14268][feat]Cosmos3 Transfer (control-video conditioning) - #16394
Draft
ishovkun wants to merge 42 commits into
Draft
[TRTLLM-14268][feat]Cosmos3 Transfer (control-video conditioning)#16394ishovkun wants to merge 42 commits into
ishovkun wants to merge 42 commits into
Conversation
ishovkun
force-pushed
the
cosmos3_control
branch
from
July 15, 2026 22:05
9319387 to
7cc0a5e
Compare
ishovkun
force-pushed
the
cosmos3_control
branch
from
July 17, 2026 04:38
7cc0a5e to
568a264
Compare
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
Previously `input_reference` was always treated as an image and stored as `<id>_reference.png`, routing unconditionally to `params.image`. Introduce `_reference_is_image` and `_reference_is_video` helpers that probe decoded content (PIL and PyAV respectively) rather than relying on file extension or content-type. Classification order is image-first: FFmpeg demuxes still images as single-frame video streams, so the PIL probe must gate the PyAV probe. On a positive video classification the reference is stored as `<id>_reference.mp4` and routed to `params.extra_params["video"]` instead of `params.image`. Undecodable content removes the temporary `.part` file and raises `ValueError`. The field docstring is updated to document the content-based dispatch contract. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
- Document V2V mode in README and cosmos3.py docstring alongside existing T2V/T2I/I2V/T2AV modes - Update `--video_path` help text to reflect V2V as the primary use case - Add `condition_video_keep="last"` smoke test and audio+V2V combined smoke test to the pipeline unit tests - Add serve endpoint tests: multipart video reference routes to `extra_params["video"]` (V2V), undecodable reference returns HTTP 400 - Document Cosmos3 `extra_params` knobs and V2V multipart curl example in the serve README; clarify `input_reference` classification behavior Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
- Add missing `prompts/v2v.json` example for video-to-video mode - Add `av` (PyAV) to `requirements.txt` for video decode support - Add "Media I/O dependencies" section to README covering ffmpeg (mp4 output) and av (V2V reference video decode) - Update V2V bullet to cross-reference the new section instead of inline pip-install note Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
- Remove leftover conflict markers from transformer_cosmos3.py - Consolidate multi-line f-strings and error messages onto single lines - Reorder imports alphabetically in test_cosmos3_pipeline.py - Add blank line after `import av` per formatting conventions Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Local path handling remains unchanged; only HTTP(S), data:, and file: URIs are routed through `load_image` to match the existing I2V image branch behavior. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Wrap the input_reference materialization block in a try/except that removes the `.part` temp file before re-raising. Previously, any failure after the file was opened (bad base64, broken upload stream, unrecognized media type) would leave the partial file on disk. Also add explicit validation for malformed base64 input, raising a descriptive ValueError instead of propagating the raw decode error. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Action generation (policy / forward_dynamics / inverse_dynamics), embodiment domain presets, action mRoPE, the DomainAwareLinear projections, and all action output fields, CLI flags, prompts, and tests are removed. V2V, audio, and the shared video encode/decode infrastructure they were built on are retained. Action generation will land in a separate PR. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Remove unused `step_extra_stream_latents` parameter from `post_step_fn` and update its return value accordingly. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Remove the two-argument `post_step_fn(latents, extra_stream_latents)` overload and the `inspect.signature` dispatch logic. The callable now always takes and returns `latents` only. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
opt-in explicitly Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Switch the V2V video reference decoder from `av` (PyAV) to OpenCV, aligning it with the shared multimodal video path. This removes the `TRTLLM_ENABLE_PYAV=1` gate, drops the `av` dependency, and updates docs, tests, and the serve-side `_reference_is_video` probe accordingly. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Reference files are now stored without a `.png`/`.mp4` suffix. Content classification uses `is_image_file`/`is_video_file` (PIL and OpenCV header probes) instead of file extension, matching what the serve path already produces when writing raw bytes. The two helpers are lifted from `visual_gen_utils.py` into `media_io.py` for reuse by the Cosmos3 `normalize_video_input_path` utility. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Replace the `extra_params["video"]` path-based approach with decoded `VideoData` under `multi_modal_data["video"]`, aligning V2V with the framework's multimodal convention. The worker now receives pre-decoded frames and VAE-encodes them directly. Also rename `condition_frame_indexes_vision` to `condition_video_latent_indexes` to clarify that the indexes refer to output latent frames, not source frame selection. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
…ia decode Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
…or, docs Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Introduces `VIDEO_RES_SIZE_INFO` resolution bucket table and `find_closest_target_size` helper to map source frame dimensions to the nearest supported Cosmos3 output size by aspect ratio. Also removes the restriction that prevented combining transfer inference with action generation. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
- Move `VIDEO_RES_SIZE_INFO` from `transfer.py` to `defaults.py` to eliminate the circular dependency and give defaults a single source of truth for resolution buckets - Rename ambiguous helpers for clarity: - `media_hw` → `media_height_width` - `media_to_uint8_cthw` → `media_to_uint8_frames` - `uint8_cthw_to_normalized_5d` → `uint8_frames_to_model_input` - `_path_media_to_uint8_cthw` → `_path_media_to_uint8_frames` - Merge `BLUR_DOWNUP_PRESETS` and `BLUR_PRE_BLUR_DOWNSCALE_PRESETS` into a single `BLUR_PRESETS` dict keyed by preset name - Add inline comments to Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
Precomputed transfer controls (depth/seg/wsm, or any hint's control_path) now ride multi_modal_data["control"] as decoded frames, produced client-side by the example; extra_params carries only configuration. The worker attaches control media to hints in resolve_transfer_config and rejects media forms (control_path/control) inside extra_params, so it never decodes or reads paths — completing the knobs-vs-media split used for the V2V reference. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>
ishovkun
force-pushed
the
cosmos3_control
branch
from
July 17, 2026 05:37
568a264 to
e6fb891
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
Adds Transfer (ControlNet-style control-video conditioning) for Cosmos3,
offline via
--extra_paramsand over the serving API. A control videoconstrains output structure frame by frame; the prompt supplies appearance.
Built on the Cosmos3 V2V PR (#16155) — merge that first.
edge,blur,depth,seg,wsm.edge/blurauto-compute from
--video_path(OpenCV); any hint accepts a precomputed{"control_path": <server-local media>}. Multiple hints compose.that share the target's mRoPE positions, so control and output patches align
at zero displacement — no extra cross-attention wiring.
cond_full/cond_no_control/uncond_full), combinedu + gs·((nc + cg·(f − nc)) − u); idle branchesskipped, gated over the denoise schedule by
control_guidance_interval.wsm),stitched by V2V-pinning the previous chunk's tail.
--extra_paramsJSON flag, a synthetic edge-map controlgenerator (
generate_bouncing_ball_control.py, no media assets), READMEusage examples.
multimodal decoder), used by default but not pinned — users
pip install opencv-python-headless(documented). No PyAV, no opt-in gate.No request-schema changes (api-compatible).
Validated on Cosmos3-Nano (1× B200): edge-fidelity F1 0.971, multi-hint
edge+blur F1 0.956, 189-frame chunked output stitched seamlessly, synthetic
bouncing-ball edge control tracked to ~8 px.
Test Coverage
tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py— CPU tests(stubbed scheduler/transformer, no GPU/model): hint-config resolution and
precomputed-control passthrough; media helpers (RGB Canny
edge, scaledbilateral
blur, temporal reflect-pad,cv2-missing error, OpenCV decoderound-trip); 3-branch CFG arithmetic (applies control+text guidance, skips
idle branches, interval switches branch counts); chunk-count arithmetic and
multichunk overlap stitching;
find_closest_target_sizebucketing (exact,nearest, order, error paths). Verified against vllm-omni PR fix: fix accuracy and illegal memory access issues when using mtp + attention dp #4379.
tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py::TestCosmos3TransferRouting— transfer requests default
use_system_promptoff.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.