From 1f2fbaa730e908776b8a56b11fac8aceac06c3d7 Mon Sep 17 00:00:00 2001 From: Xiangyu Lu <169013972+xlu451@users.noreply.github.com> Date: Wed, 1 Jul 2026 19:58:50 -0700 Subject: [PATCH 01/26] Add Docker build CI workflow (#73) Co-authored-by: Xiangyu Lu --- .github/workflows/docker-build.yml | 71 ++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 .github/workflows/docker-build.yml diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml new file mode 100644 index 00000000..063076c4 --- /dev/null +++ b/.github/workflows/docker-build.yml @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +name: Docker Build + +on: + pull_request: + branches: [main] + paths: + - .github/workflows/docker-build.yml + - Dockerfile + - .dockerignore + - .python-version + - pyproject.toml + - uv.lock + - docker/** + - packages/*/pyproject.toml + - packages/*/uv.lock + push: + branches: [main] + paths: + - .github/workflows/docker-build.yml + - Dockerfile + - .dockerignore + - .python-version + - pyproject.toml + - uv.lock + - docker/** + - packages/*/pyproject.toml + - packages/*/uv.lock + workflow_dispatch: + +concurrency: + group: docker-build-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + runtime-image: + runs-on: [self-hosted, docker-build, h200, cosmos-framework] + timeout-minutes: 120 + env: + DOCKER_BUILDKIT: "1" + BUILDER_NAME: cosmos-docker-build-${{ github.run_id }}-${{ github.run_attempt }} + steps: + - uses: actions/checkout@v6 + + - name: Show Docker disk usage before build + run: | + docker system df + df -h "$(docker info --format '{{ .DockerRootDir }}')" + + - name: Create temporary BuildKit builder + run: | + docker buildx rm -f "$BUILDER_NAME" >/dev/null 2>&1 || true + docker buildx create --driver docker-container --name "$BUILDER_NAME" --use + + - name: Build Docker image + run: docker buildx build --builder "$BUILDER_NAME" --progress=plain . + + - name: Remove temporary BuildKit builder + if: always() + run: docker buildx rm -f "$BUILDER_NAME" || true + + - name: Show Docker disk usage after cleanup + if: always() + run: | + docker system df + df -h "$(docker info --format '{{ .DockerRootDir }}')" From 61460b629de2f0b08153c3b623ec114d76821bdb Mon Sep 17 00:00:00 2001 From: yy-code-nv Date: Thu, 2 Jul 2026 11:51:21 +0800 Subject: [PATCH 02/26] =?UTF-8?q?Release:=20sync=20from=20i4=20(vlm=20?= =?UTF-8?q?=E2=86=92=20reasoner=20rename)=20+=20inference/config=20fixes?= =?UTF-8?q?=20(#70)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Upstream i4 refactor `4190136a09` renamed `projects/cosmos3/vfm/` to `projects/cosmos3/cosmos3/` and moved the `vlm/` subtree to `reasoner/` under configs/base, models, datasets, and utils. This release re-runs the release pipeline against that new source layout and cleans up CF-side knock-on damage. Highlights: * Release pipeline output: 459 files mapped from the new i4 layout. Files added under `cosmos_framework/**/reasoner/`; the old `cosmos_framework/**/vlm/` counterparts are removed (orphan cleanup). Also ships `callbacks/tokens_per_sec.py`, `model/tokenizer/utils/vlm_prompt_format.py`, a new `configs/base/defaults/experimental/`-excluded layout, and refreshed `data/vfm/augmentors/reasoner/*` etc. * CF-owned config move: `configs/base/vlm/{defaults,experiment}/*.py` moved to `configs/base/reasoner/{defaults,experiment}/` (5 files), with imports rewritten `cosmos_framework.configs.base.vlm.` → `cosmos_framework.configs.base.reasoner.` so `task="vlm"` still finds a live experiment tree. * `configs/toml_config/toml_config_helper.py`: retarget `task="vlm"` → `cosmos_framework/configs/base/reasoner/config.py`. * Inference config rewriters (`inference/common/config.py`, `inference/common/public_model_config.py`): add `vlm` → `reasoner` rewrite rules ahead of the general vfm rules so old checkpoint JSONs (which still ship `cosmos3._src.vfm.configs.base.defaults.vlm.*` targets) resolve to the new module paths at load time. * `configs/base/experiment/sft/models/{nano,super}_model_config.py` and the `inference/configs/model/Cosmos3-{Nano,Super}.yaml` + `examples/checkpoints/Cosmos3-Nano/model/config.json`: drop the removed `DiffusionExpertConfig` fields (`position_embedding_type`, `rope_{h,t,w}_extrapolation_ratio`) and retarget the shipped Qwen3-VL JSON path to the `reasoner/` layout. * `inference/common/public_model_config_test.py`: update stale `vlm` paths in the round-trip test fixture. * `model/vfm/mot/unified_mot_test.py`: removed (tested private helpers `_PACKAGE_ROOT` / `_resolve_packaged_config_path` that were refactored out of `unified_mot.py`). Testing (4 × NVIDIA GB200 node): * `test_launch_regression[vision_sft_nano]`: 10-iter losses match gb200 goldens exactly. * `test_launch_regression[llava_ov]`: trains 10 iters end to end; loss drifts vs. 2026-05-18 gb200 goldens (goldens need recapture). * `nano_reasoner_inference_smoke_test`: inference runs cleanly; argmax hard-gate passes; tight allclose (rtol=atol=1e-3) fails by bf16 noise vs. presumably-H100 goldens. * Unit suite (`--num-gpus=0 --levels=0`): 327 passed / 1 failed (the failure is the `convert_model_to_dcp` script test, which needs `HF_TOKEN` for the private `nvidia/Cosmos3-Experimental` snapshot). --------- Co-authored-by: lfengad --- cosmos_framework/callbacks/expert_heatmap.py | 2 +- .../callbacks/moe_specialization_callback.py | 2 +- .../callbacks/moe_stability_callback.py | 2 +- cosmos_framework/callbacks/norm_monitor.py | 4 +- cosmos_framework/callbacks/tokens_per_sec.py | 415 ++++++++++++++++++ .../callbacks/wandb_log_simple.py | 155 ------- cosmos_framework/checkpoint/dcp.py | 390 ++++++++++++++-- cosmos_framework/configs/base/config.py | 2 +- .../configs/base/defaults/checkpointer.py | 27 +- .../configs/base/defaults/llm_model.py | 64 --- .../configs/base/defaults/model_config.py | 19 +- .../base/defaults/open_source_dataloader.py | 2 +- .../base/defaults/{vlm.py => reasoner.py} | 138 +++++- .../sft/models/nano_model_config.py | 8 +- .../sft/models/super_model_config.py | 8 +- .../base/{vlm => reasoner}/__init__.py | 0 .../configs/base/{vlm => reasoner}/config.py | 12 +- .../{vlm => reasoner}/defaults/__init__.py | 0 .../{vlm => reasoner}/defaults/callbacks.py | 5 + .../defaults/checkpointer.py | 0 .../base/{vlm => reasoner}/defaults/config.py | 8 +- .../base/{vlm => reasoner}/defaults/model.py | 2 +- .../{vlm => reasoner}/defaults/optimizer.py | 2 +- .../defaults/policy_config.py | 11 +- .../{vlm => reasoner}/defaults/vlm_policy.py | 6 +- .../{vlm => reasoner}/experiment/__init__.py | 0 .../experiment/dataflow_roles.py | 0 .../experiment/llava_ov_vlm.py | 2 +- .../{vlm => reasoner}/experiment/utils.py | 0 .../experiment/videophy2_dataflow_roles.py | 0 .../experiment/videophy2_sft_nano.py | 4 +- .../base/{vlm => reasoner}/freeze_config.py | 0 .../configs/toml_config/toml_config_helper.py | 2 +- .../augmentors/{vlm => reasoner}/__init__.py | 0 .../{vlm => reasoner}/bytes_to_media.py | 133 +++++- .../{vlm => reasoner}/filter_output_key.py | 0 .../{vlm => reasoner}/filter_seq_length.py | 0 .../floating_number_format.py | 0 .../format_describe_anything.py | 2 +- .../nvlm_data_to_conversation.py | 0 .../{vlm => reasoner}/prompt_format.py | 0 .../shuffle_text_media_order.py | 0 .../augmentors/{vlm => reasoner}/timestamp.py | 2 +- .../timestamp_with_subject_tracking.py | 2 +- .../timestamp_without_augment_message.py | 2 +- .../timestamp_without_end_time.py | 2 +- .../{vlm => reasoner}/tokenize_data.py | 58 ++- .../user_prompt_caption_general.json | 0 .../{vlm => reasoner}/user_prompt_ocr.json | 0 .../data/vfm/augmentors/text_tokenizer.py | 3 + .../augmentors/text_transforms_for_image.py | 7 +- cosmos_framework/data/vfm/joint_dataloader.py | 3 +- .../data/vfm/local_datasets/sft_dataset.py | 2 +- .../data/vfm/packing_iterable_dataset.py | 2 +- .../data/vfm/processors/__init__.py | 2 +- cosmos_framework/data/vfm/processors/base.py | 4 +- .../processors/nemotron3densevl_processor.py | 26 +- .../data/vfm/processors/qwen3vl_processor.py | 14 +- .../{vlm => reasoner}/video_decoder_qwen.py | 0 .../data/vfm/sequence_packing/modalities.py | 163 +++---- .../data/vfm/sequence_packing/packers.py | 39 +- .../vfm/sequence_packing/temporal_causal.py | 178 ++++---- .../data/vfm/sequence_packing/types.py | 29 +- cosmos_framework/inference/common/config.py | 14 + .../inference/common/public_model_config.py | 16 + .../common/public_model_config_test.py | 6 +- .../inference/configs/model/Cosmos3-Nano.yaml | 4 - .../configs/model/Cosmos3-Super.yaml | 4 - cosmos_framework/inference/inference.py | 2 +- cosmos_framework/inference/transfer.py | 2 +- .../benchmark_fmha_head_sharding.py | 398 +++++++++++++++++ .../evaluation/reconstruction_metrics.py | 4 + .../tokenizer/models/sparse_autoencoder.py | 10 +- .../model/tokenizer/models/text_decoder.py | 272 ++++++++---- .../tokenizer/utils/vlm_prompt_format.py | 23 + .../model/vfm/algorithm/loss/__init__.py | 4 +- .../model/vfm/algorithm/loss/cross_entropy.py | 68 ++- .../vfm/algorithm/loss/load_balancing.py | 2 +- cosmos_framework/model/vfm/hf_model.py | 13 +- cosmos_framework/model/vfm/mot/attention.py | 47 +- .../model/vfm/mot/context_parallel_utils.py | 21 + .../model/vfm/mot/cosmos3_vfm_network.py | 120 +---- .../model/vfm/mot/modeling_utils.py | 305 ------------- .../model/vfm/mot/parallelize_unified_mot.py | 12 + .../model/vfm/mot/parallelize_vfm_network.py | 2 - .../model/vfm/mot/und_k_norm_example_test.py | 324 ++++++++++++++ cosmos_framework/model/vfm/mot/unified_mot.py | 131 ++++-- .../model/vfm/mot/unified_mot_test.py | 68 --- cosmos_framework/model/vfm/omni_mot_model.py | 175 ++++---- cosmos_framework/model/vfm/parallelize_vlm.py | 160 ++++++- .../model/vfm/{vlm => reasoner}/__init__.py | 0 .../nemotron_3_dense_vl/__init__.py | 0 .../configs/Nemotron-2B-Dense-VL.json | 0 .../configuration_nemotron_3_dense_vl.py | 0 .../nemotron_3_dense_vl.py | 2 +- .../{vlm => reasoner}/qwen3_vl/__init__.py | 0 .../configs/Qwen3-VL-2B-Instruct.json | 0 .../configs/Qwen3-VL-32B-Instruct.json | 0 .../configs/Qwen3-VL-4B-Instruct.json | 0 .../configs/Qwen3-VL-8B-Instruct.json | 0 .../qwen3_vl/configs/__init__.py | 0 .../qwen3_vl/configuration_qwen3_vl.py | 0 .../{vlm => reasoner}/qwen3_vl/qwen3_vl.py | 8 +- .../vfm/{vlm => reasoner}/qwen3_vl/utils.py | 65 ++- .../qwen3_vl/video_processing_qwen3_vl.py | 0 .../qwen3_vl_moe/__init__.py | 0 .../configs/Qwen3-VL-235B-A22B-Instruct.json | 0 .../configs/Qwen3-VL-30B-A3B-Instruct.json | 0 .../qwen3_vl_moe/configs/__init__.py | 0 .../configuration_qwen3_vl_moe.py | 0 .../vfm/{vlm => reasoner}/qwen3_vl_moe/moe.py | 4 +- .../qwen3_vl_moe/moe_bench.py | 20 +- .../qwen3_vl_moe/moe_kernels.py | 0 .../qwen3_vl_moe/moe_test.py | 4 +- .../qwen3_vl_moe/qwen3_vl_moe.py | 10 +- .../vfm/tokenizers/uniae/noncausal_4x16x16.py | 4 +- .../model/vfm/utils/dcp_loader.py | 120 ----- .../model/vfm/utils/safetensors_loader.py | 13 +- cosmos_framework/model/vfm/vlm_model.py | 41 +- cosmos_framework/utils/config.py | 24 +- .../utils/env_parsers/cred_env_parser.py | 4 +- cosmos_framework/utils/vfm/flash_attn.py | 2 +- cosmos_framework/utils/vfm/fused_adam.py | 7 +- cosmos_framework/utils/vfm/model_loader.py | 3 + cosmos_framework/utils/vfm/monkey_patch.py | 104 +++-- .../utils/vfm/{vlm => reasoner}/__init__.py | 0 .../utils/vfm/{vlm => reasoner}/constant.py | 0 .../{vlm => reasoner}/create_position_ids.py | 0 .../vfm/{vlm => reasoner}/flop_calculator.py | 0 .../pretrained_models_downloader.py | 0 .../vlm/configs_defaults/checkpointer.py | 22 +- examples/integration/net_level.py | 2 +- .../integration/trainer_level_inference.py | 4 +- .../integration/trainer_level_training.py | 4 +- 134 files changed, 3032 insertions(+), 1607 deletions(-) create mode 100644 cosmos_framework/callbacks/tokens_per_sec.py delete mode 100644 cosmos_framework/callbacks/wandb_log_simple.py delete mode 100644 cosmos_framework/configs/base/defaults/llm_model.py rename cosmos_framework/configs/base/defaults/{vlm.py => reasoner.py} (83%) rename cosmos_framework/configs/base/{vlm => reasoner}/__init__.py (100%) rename cosmos_framework/configs/base/{vlm => reasoner}/config.py (77%) rename cosmos_framework/configs/base/{vlm => reasoner}/defaults/__init__.py (100%) rename cosmos_framework/configs/base/{vlm => reasoner}/defaults/callbacks.py (93%) rename cosmos_framework/configs/base/{vlm => reasoner}/defaults/checkpointer.py (100%) rename cosmos_framework/configs/base/{vlm => reasoner}/defaults/config.py (83%) rename cosmos_framework/configs/base/{vlm => reasoner}/defaults/model.py (90%) rename cosmos_framework/configs/base/{vlm => reasoner}/defaults/optimizer.py (98%) rename cosmos_framework/configs/base/{vlm => reasoner}/defaults/policy_config.py (79%) rename cosmos_framework/configs/base/{vlm => reasoner}/defaults/vlm_policy.py (95%) rename cosmos_framework/configs/base/{vlm => reasoner}/experiment/__init__.py (100%) rename cosmos_framework/configs/base/{vlm => reasoner}/experiment/dataflow_roles.py (100%) rename cosmos_framework/configs/base/{vlm => reasoner}/experiment/llava_ov_vlm.py (99%) rename cosmos_framework/configs/base/{vlm => reasoner}/experiment/utils.py (100%) rename cosmos_framework/configs/base/{vlm => reasoner}/experiment/videophy2_dataflow_roles.py (100%) rename cosmos_framework/configs/base/{vlm => reasoner}/experiment/videophy2_sft_nano.py (97%) rename cosmos_framework/configs/base/{vlm => reasoner}/freeze_config.py (100%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/__init__.py (100%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/bytes_to_media.py (62%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/filter_output_key.py (100%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/filter_seq_length.py (100%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/floating_number_format.py (100%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/format_describe_anything.py (99%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/nvlm_data_to_conversation.py (100%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/prompt_format.py (100%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/shuffle_text_media_order.py (100%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/timestamp.py (99%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/timestamp_with_subject_tracking.py (99%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/timestamp_without_augment_message.py (99%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/timestamp_without_end_time.py (99%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/tokenize_data.py (88%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/user_prompt_caption_general.json (100%) rename cosmos_framework/data/vfm/augmentors/{vlm => reasoner}/user_prompt_ocr.json (100%) rename cosmos_framework/data/vfm/{vlm => reasoner}/video_decoder_qwen.py (100%) create mode 100644 cosmos_framework/model/attention/benchmarks/benchmark_fmha_head_sharding.py create mode 100644 cosmos_framework/model/tokenizer/utils/vlm_prompt_format.py create mode 100644 cosmos_framework/model/vfm/mot/und_k_norm_example_test.py delete mode 100644 cosmos_framework/model/vfm/mot/unified_mot_test.py rename cosmos_framework/model/vfm/{vlm => reasoner}/__init__.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/nemotron_3_dense_vl/__init__.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/nemotron_3_dense_vl/nemotron_3_dense_vl.py (98%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl/__init__.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl/configs/__init__.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl/configuration_qwen3_vl.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl/qwen3_vl.py (99%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl/utils.py (90%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl/video_processing_qwen3_vl.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl_moe/__init__.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl_moe/configs/__init__.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl_moe/configuration_qwen3_vl_moe.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl_moe/moe.py (98%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl_moe/moe_bench.py (94%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl_moe/moe_kernels.py (100%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl_moe/moe_test.py (93%) rename cosmos_framework/model/vfm/{vlm => reasoner}/qwen3_vl_moe/qwen3_vl_moe.py (99%) delete mode 100644 cosmos_framework/model/vfm/utils/dcp_loader.py rename cosmos_framework/utils/vfm/{vlm => reasoner}/__init__.py (100%) rename cosmos_framework/utils/vfm/{vlm => reasoner}/constant.py (100%) rename cosmos_framework/utils/vfm/{vlm => reasoner}/create_position_ids.py (100%) rename cosmos_framework/utils/vfm/{vlm => reasoner}/flop_calculator.py (100%) rename cosmos_framework/utils/vfm/{vlm => reasoner}/pretrained_models_downloader.py (100%) diff --git a/cosmos_framework/callbacks/expert_heatmap.py b/cosmos_framework/callbacks/expert_heatmap.py index 8fb745dd..f3262fc7 100644 --- a/cosmos_framework/callbacks/expert_heatmap.py +++ b/cosmos_framework/callbacks/expert_heatmap.py @@ -10,7 +10,7 @@ from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.trainer import ImaginaireTrainer from cosmos_framework.utils import distributed -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock def compute_expert_heatmap(vfm: torch.nn.Module) -> dict[str, torch.Tensor]: diff --git a/cosmos_framework/callbacks/moe_specialization_callback.py b/cosmos_framework/callbacks/moe_specialization_callback.py index 5f1f32c6..4d5d4233 100644 --- a/cosmos_framework/callbacks/moe_specialization_callback.py +++ b/cosmos_framework/callbacks/moe_specialization_callback.py @@ -39,7 +39,7 @@ from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.trainer import ImaginaireTrainer from cosmos_framework.utils import distributed -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock def _get_device_mesh(vfm: torch.nn.Module): diff --git a/cosmos_framework/callbacks/moe_stability_callback.py b/cosmos_framework/callbacks/moe_stability_callback.py index e223e95d..8e241071 100644 --- a/cosmos_framework/callbacks/moe_stability_callback.py +++ b/cosmos_framework/callbacks/moe_stability_callback.py @@ -107,7 +107,7 @@ from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.trainer import ImaginaireTrainer from cosmos_framework.utils import distributed -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock def _effective_experts( diff --git a/cosmos_framework/callbacks/norm_monitor.py b/cosmos_framework/callbacks/norm_monitor.py index 3c683038..460804cf 100644 --- a/cosmos_framework/callbacks/norm_monitor.py +++ b/cosmos_framework/callbacks/norm_monitor.py @@ -212,8 +212,8 @@ def _get_named_parameters(self, model: nn.Module) -> dict[str, nn.Parameter]: def _should_track_param(self, param_name: str) -> bool: """Check if parameter should be tracked based on naming conventions.""" - # Track only generation tower params, exclude EMA params - return "moe_gen" in param_name and "net_ema" not in param_name + # Track generation tower params and und→gen cross-attention norms; exclude EMA params + return ("moe_gen" in param_name or "k_norm_und_for_gen" in param_name) and "net_ema" not in param_name def _compute_l2_stats(self, tensor: torch.Tensor, detach: bool = True) -> dict[str, torch.Tensor]: """Compute statistics (squared sum and max) for a tensor. diff --git a/cosmos_framework/callbacks/tokens_per_sec.py b/cosmos_framework/callbacks/tokens_per_sec.py new file mode 100644 index 00000000..1634db95 --- /dev/null +++ b/cosmos_framework/callbacks/tokens_per_sec.py @@ -0,0 +1,415 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Tokens-per-second throughput callback for VLM training. + +Logs a single, directly-comparable throughput number per logging window so that +optimization ablations (compile, sharding strategy, comm dtype, packing budget) +can be A/B'd on equal footing. + +Why this exists instead of ``MFUCallback`` (``cosmos_framework/callbacks/mfu.py``): +``MFUCallback`` is written for the OmniMoT/VFM network -- it reads +``model.net.language_model.config`` and per-modality token counts +(``output_batch["und_token_length"]`` etc.). ``VLMModel`` +(``cosmos_framework/model/vfm/vlm_model.py``) is a plain HF wrapper: it exposes +``self.model`` (no ``.net``) and its ``training_step`` returns only +``{"loss", "loss_avg", "labels"}``. ``MFUCallback`` would therefore silently +no-op (token length is ``None``) or fail on ``model.net``. This callback instead +counts the real input tokens off ``data_batch["input_ids"]`` -- the one tensor +the VLM forward always consumes -- so it is correct for the VLM path and +self-normalizes for any residual per-step token variation. + +Measurement notes: + - Reports rank-0 *per-GPU* throughput. With token-based packing (each rank + packs to the same ``max_tokens`` budget) this is representative of every + rank; no cross-rank collective is taken, so the callback adds no + communication and does not contaminate timing. + - ``input_ids.numel()`` is a metadata read (tensor shape), so it triggers no + device sync. + - Content (non-pad) tokens are read from the loader-emitted CPU int + ``data_batch["content_tokens"]`` (set in ``collate_fn.custom_collate``), so the + packing-efficiency metric also triggers no per-step device sync. The legacy + ``(input_ids != pad_id).sum()`` device reduction remains only as a fallback when + that key is absent. + - ``hit_thres`` warm-up steps are skipped so compilation / allocator warm-up + don't pollute the first window (mirrors ``IterSpeed`` / ``MFUCallback``). + +Packing-aware telemetry. On top of the headline +throughput rate, this callback reports the "iteration speed per sample at a given packing +density" view, in four groups: + - per-sample / per-step: ``sec_per_sample``, ``samples_per_step``, ``tokens_per_step``. + - density: ``useful_util`` (= U/padded), supervision density ``rho_sup`` (= U*/U), + ``attention_quadratic_waste`` (the O(L^2) work true packing would remove), and the + unpadded ``seq_max_len`` (l_max) mean/max. + - utilization / comm amortization: ``useful_mfu`` (+ ``compute_bound`` flag) tells us when + raising the batch budget stops paying; ``comm_bytes_per_useful_token`` makes the + fixed-per-step FSDP collective volume amortized over useful work explicit. + - cost-model calibration: ``realized`` vs packer-``predicted`` step time (FLOP path only). + +All inputs are loader-emitted CPU ints/floats (``content_tokens``, ``supervised_tokens``, +``sum_len_sq``, ``seq_max_len``, ``predicted_runtime_ms`` from ``collate_fn.custom_collate``) +plus tensor-shape metadata and a one-time parameter count, so the callback adds NO per-step +device sync. ``COSMOS_VLM_TELEMETRY=0`` turns the whole callback into a no-op for the +telemetry-overhead A/B. We intentionally do not log ``sec_per_useful_token`` (it is the exact +reciprocal of ``useful_tokens_per_sec_per_gpu``). +""" + +import os +import time + +import torch +import wandb +from torch import Tensor + +from cosmos_framework.callbacks.every_n import EveryN +from cosmos_framework.model._base import ImaginaireModel +from cosmos_framework.trainer import ImaginaireTrainer +from cosmos_framework.utils import log +from cosmos_framework.utils.distributed import is_rank0, rank0_only + + +class VLMTokensPerSec(EveryN): + """Log per-GPU tokens/sec over each logging window. + + Args: + hit_thres: Number of warm-up steps to skip before timing begins. + length_key: Key in ``data_batch`` whose tensor holds the packed token + ids for this rank (``input_ids`` for the VLM path). + """ + + def __init__( + self, + *args, + hit_thres: int = 50, + length_key: str = "input_ids", + peak_flops_per_gpu: float | None = None, + comm_factor: float = 3.0, + comm_dtype_bytes: float = 2.0, + compute_bound_mfu_thresh: float = 0.5, + fwd_bwd_flop_coeff: float | None = None, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + self.hit_thres = hit_thres + self.length_key = length_key + # A/B toggle for the telemetry-overhead validation: + # COSMOS_VLM_TELEMETRY=0 makes the whole callback a no-op so a "telemetry off" run + # measures the true cost of the callback (base + extended) against a "telemetry on" run. + # Defaults to enabled when the env var is unset (normal training). + self._enabled: bool = os.environ.get("COSMOS_VLM_TELEMETRY", "1") != "0" + # Hardware/structural constants for the comm-amortization and utilization metrics. + # peak_flops_per_gpu: dense BF16 peak; None => resolve from torch.cuda device name. + # comm_factor: per-step FSDP collective volume in units of the (global) parameter + # bytes -- 2 all-gathers (fwd + activation-ckpt bwd recompute) + 1 grad + # reduce-scatter ~= 3x. comm_dtype_bytes: 2 for BF16 collectives. + self._peak_flops: float | None = peak_flops_per_gpu + self._comm_factor: float = comm_factor + self._comm_dtype_bytes: float = comm_dtype_bytes + self._compute_bound_mfu_thresh: float = compute_bound_mfu_thresh + # fwd+bwd FLOPs per parameter per token for useful_mfu: 6 without activation + # recomputation, 8 with full activation checkpointing (the +2 is the recompute + # forward run during backward). None => auto-resolve from the model's AC mode. + self._fwd_bwd_flop_coeff: float | None = fwd_bwd_flop_coeff + self._flop_coeff_resolved: float | None = None # cached auto-resolved coeff + self._num_params: int | None = None # lazily filled from model.parameters() (global/DTensor) + self._hit_counter: int = 0 + self._tokens_in_window: int = 0 + self._useful_tokens_in_window: int = 0 + self._samples_in_window: int = 0 + self._singleton_steps_in_window: int = 0 + self._steps_in_window: int = 0 + self._window_start_time: float | None = None + # Extended packing telemetry (all from loader-emitted CPU ints -> sync-free). + self._supervised_tokens_in_window: int = 0 # U* = #(labels != ignore_index) + self._padded_attn_in_window: int = 0 # Σ_step k * padded_len^2 (attention work paid) + self._content_attn_in_window: int = 0 # Σ_step Σ_i L_i^2 (attention work needed) + self._seq_max_len_sum: int = 0 # for mean l_max over the window + self._seq_max_len_max: int = 0 # max l_max over the window + self._seq_max_len_steps: int = 0 # #steps that carried seq_max_len (for the mean) + # Cost-model calibration: packer-predicted per-step runtime (FLOP path only). + self._predicted_runtime_ms_in_window: float = 0.0 + self._predicted_steps_in_window: int = 0 + + def on_training_step_end( + self, + model: ImaginaireModel, + data_batch: dict[str, torch.Tensor], + output_batch: dict[str, torch.Tensor], + loss: torch.Tensor, + iteration: int = 0, + ) -> None: + # Telemetry-off A/B arm: complete no-op (no accumulation, no periodic report). + if not self._enabled: + return + # Skip warm-up steps entirely (no timing, no accumulation). + if self._hit_counter < self.hit_thres: + self._hit_counter += 1 + return + + # Only rank 0 accumulates AND reports. every_n_impl (which resets the window) is + # @rank0_only, so accumulating on the other ranks would never reset -- their window + # counters (notably _padded_attn_in_window ~ Σ k·l_max²) would grow for the entire + # run. This callback takes no cross-rank collective (it reports rank-0 per-GPU + # throughput), so scoping accumulation to rank 0 matches the report/reset scope. + # Non-rank0 still delegates to EveryN below so its periodic distributed barrier + # (barrier_after_run) stays collective and in lockstep. + if not is_rank0(): + super().on_training_step_end(model, data_batch, output_batch, loss, iteration) + return + + # Open the timing window on the first post-warm-up step. + if self._window_start_time is None: + self._window_start_time = time.time() + + ids = data_batch.get(self.length_key) + if ids is not None: + # numel() is the PADDED token count (dynamic batching pads each batch to + # its longest member); shape[0] is the number of packed samples this step. + n_samples = int(ids.shape[0]) if ids.dim() > 1 else 1 + self._tokens_in_window += int(ids.numel()) + self._samples_in_window += n_samples + if n_samples == 1: + self._singleton_steps_in_window += 1 + # USEFUL (non-pad) tokens. Prefer the loader-emitted CPU int + # data_batch["content_tokens"] (= Σ len(input_ids) captured PRE-pad in + # collate_fn.custom_collate). Reading it is sync-free, so this throughput + # callback no longer perturbs the very step timing it measures. Only if the + # key is absent (e.g. a non-VLM collate) do we fall back to a device + # reduction off pad_token_id -- correct, but it forces a per-step D2H sync. + # useful/padded ratio = packing efficiency -> whether the padded budget is + # real content or padding (i.e. is true no-pad packing worth it). + content = data_batch.get("content_tokens") + if content is not None: + self._useful_tokens_in_window += int(content) + else: + pad = data_batch.get("pad_token_id") + if pad is not None: + pad_id = int(pad.flatten()[0]) + self._useful_tokens_in_window += int((ids != pad_id).sum()) + + # Extended packing telemetry. All reads are loader-emitted CPU ints (or tensor + # shape metadata), so this adds NO device sync to the timed step. Each is guarded + # so a non-VLM collate that omits a key simply skips that derived metric. + sup = data_batch.get("supervised_tokens") + if sup is not None: + self._supervised_tokens_in_window += int(sup) + slq = data_batch.get("sum_len_sq") + if slq is not None: + self._content_attn_in_window += int(slq) + smax = data_batch.get("seq_max_len") + if smax is not None: + smax = int(smax) + self._seq_max_len_sum += smax + self._seq_max_len_max = max(self._seq_max_len_max, smax) + self._seq_max_len_steps += 1 + # Padded attention work this step = n_samples * padded_len^2. padded_len is the + # row length the model actually attends over (shape metadata -> no sync). + padded_len = int(ids.shape[1]) if ids.dim() > 1 else int(ids.shape[0]) + self._padded_attn_in_window += n_samples * padded_len * padded_len + # Packer-predicted per-step runtime (FLOP cost model), a CPU float emitted by + # collate_fn only on the FLOP-batching path -> sync-free; absent => skipped. + pred = data_batch.get("predicted_runtime_ms") + if pred is not None: + self._predicted_runtime_ms_in_window += float(pred) + self._predicted_steps_in_window += 1 + self._steps_in_window += 1 + + # Delegate to EveryN for the periodic reporting cadence. + super().on_training_step_end(model, data_batch, output_batch, loss, iteration) + + def _get_num_params(self, model: ImaginaireModel) -> int | None: + """Global parameter count, cached. For FSDP2/DTensor params ``.numel()`` returns the + logical (unsharded) size, so this is the true model size regardless of shard degree.""" + if self._num_params is None: + try: + self._num_params = int(sum(p.numel() for p in model.parameters())) + except Exception: # pragma: no cover - defensive; MFU/comm metrics just skip + self._num_params = None + return self._num_params + + def _resolve_peak_flops(self) -> float | None: + """Dense BF16 peak FLOP/s per GPU. Explicit override wins; else map the device name. + Unknown device -> None so MFU is skipped rather than reported against a wrong peak.""" + if self._peak_flops is not None: + return self._peak_flops + if not torch.cuda.is_available(): + return None + name = torch.cuda.get_device_name() + if any(k in name for k in ("B200", "GB200", "Blackwell")): + return 2.45e15 # GB200 NVL72 dense BF16 ~2.45 PFLOP/s/GPU (no sparsity) + if "H200" in name or "H100" in name: + return 989e12 # SXM dense BF16 + if "A100" in name: + return 312e12 + log.warning(f"VLMTokensPerSec: unknown device '{name}', MFU not reported (set peak_flops_per_gpu)") + return None + + def _resolve_flop_coeff(self, model: ImaginaireModel) -> float: + """fwd+bwd FLOPs per parameter per token for useful_mfu. 6 without activation + recomputation; 8 with full activation checkpointing (the +2 is the recompute + forward during backward). Explicit override (``fwd_bwd_flop_coeff``) wins; else + read the model's AC mode. On the VLM/HF path ONLY ``mode == "full"`` actually + recomputes -- ``"selective"`` degrades to a no-op because the HF backbone has no + per-op SAC (see vlm_model.py / activation_checkpointing.py), so we key on "full" + rather than the looser ``!= "none"`` MFUCallback uses. Cached (config is static).""" + if self._fwd_bwd_flop_coeff is not None: + return self._fwd_bwd_flop_coeff + if self._flop_coeff_resolved is None: + ac_cfg = getattr(getattr(model, "config", None), "activation_checkpointing", None) + ac_mode = getattr(ac_cfg, "mode", "none") + self._flop_coeff_resolved = 8.0 if ac_mode == "full" else 6.0 + return self._flop_coeff_resolved + + @rank0_only + def every_n_impl( + self, + trainer: ImaginaireTrainer, + model: ImaginaireModel, + data_batch: dict[str, Tensor], + output_batch: dict[str, Tensor], + loss: Tensor, + iteration: int, + ) -> None: + if self._window_start_time is None or self._steps_in_window == 0: + return + + elapsed = time.time() - self._window_start_time + if elapsed <= 0 or self._tokens_in_window == 0: + return + + tokens_per_sec_per_gpu = self._tokens_in_window / elapsed + tokens_per_step = self._tokens_in_window / self._steps_in_window + samples_per_step = self._samples_in_window / self._steps_in_window + singleton_rate = self._singleton_steps_in_window / self._steps_in_window + + # Packing efficiency: useful (non-pad) tokens vs the padded budget actually + # paid for. useful_util ~1.0 => packing is dense (true no-pad packing would + # buy little); useful_util << 1 => padding waste (pool_size / scoring / true + # sequence packing are the levers). useful_tokens_per_sec is the REAL training + # throughput (what loss-vs-tokens should be measured against). + useful_tokens_per_sec_per_gpu = self._useful_tokens_in_window / elapsed + useful_util = self._useful_tokens_in_window / self._tokens_in_window if self._tokens_in_window else 0.0 + + # --- Extended packing telemetry (the "iteration speed per sample at a density" view) --- + # sec_per_sample is only fair alongside the density terms below, so they are logged + # together. NOTE: we intentionally do NOT log a + # sec_per_useful_token -- it is the exact reciprocal of useful_tokens_per_sec_per_gpu. + sec_per_sample = elapsed / self._samples_in_window if self._samples_in_window else 0.0 + # rho_sup = U*/U : fraction of CONTENT tokens that carry a loss signal. + rho_sup = ( + self._supervised_tokens_in_window / self._useful_tokens_in_window if self._useful_tokens_in_window else 0.0 + ) + useful_supervised_tokens_per_sec_per_gpu = self._supervised_tokens_in_window / elapsed + # attention-quadratic waste = 1 - Σ L_i^2 / Σ k*l_max^2 : the part of attention cost + # that true sequence packing removes but the linear useful_util cannot see. + attention_quadratic_waste = ( + 1.0 - self._content_attn_in_window / self._padded_attn_in_window if self._padded_attn_in_window else 0.0 + ) + seq_max_len_mean = self._seq_max_len_sum / self._seq_max_len_steps if self._seq_max_len_steps else 0.0 + seq_max_len_max = self._seq_max_len_max + + # --- Amortized communication per useful token --- + # Per-step FSDP collective volume per rank is a FIXED structural constant (independent of + # batch size): comm_factor * N_params * comm_dtype_bytes. Dividing the fixed per-step bytes + # by useful tokens/step makes the amortization explicit -- this is exactly why denser packing + # lowers comm-per-sample (more useful work under the same collectives). Structural estimate, + # not a measured NCCL byte count. + num_params = self._get_num_params(model) + useful_tokens_per_step = self._useful_tokens_in_window / self._steps_in_window + comm_bytes_per_useful_token = 0.0 + if num_params and useful_tokens_per_step > 0: + comm_bytes_per_step = self._comm_factor * num_params * self._comm_dtype_bytes + comm_bytes_per_useful_token = comm_bytes_per_step / useful_tokens_per_step + + # --- Useful-MFU + compute-bound flag --- + # useful_mfu = coeff * N * useful_tokens/s / peak (conventional linear-term MFU, but on + # USEFUL tokens so padding is not credited). coeff is 6 (no recompute) or 8 (full activation + # checkpointing), resolved from the model's AC mode so recompute runs are not underestimated + # by ~25%. It is a conservative lower bound (ignores the attention quadratic term -- which is + # separately visible as attn_quadratic_waste). It answers "when to stop raising the budget": + # useful_mfu near peak => compute-bound => denser packing buys little more throughput; far + # from peak => overhead/comm-bound => packing still pays. + peak_flops = self._resolve_peak_flops() + useful_mfu = 0.0 + compute_bound = 0 + if num_params and peak_flops: + flop_coeff = self._resolve_flop_coeff(model) + useful_mfu = (flop_coeff * num_params * useful_tokens_per_sec_per_gpu) / peak_flops + compute_bound = 1 if useful_mfu >= self._compute_bound_mfu_thresh else 0 + + # --- Cost-model calibration: realized vs packer-predicted step time --- + # The FLOP packer sizes each step to ~target_runtime using estimate_runtime_ms; logging the + # realized/predicted ratio calibrates that cost model. ~1.x is expected (predicted models + # compute only, realized includes optimizer/comm/dataloader); a drifting or large ratio means + # the packer is mis-sizing steps (oversized => comm under-amortized; undersized => throughput + # left on the table). Populated only on the FLOP-batching path. + realized_step_ms = elapsed * 1000.0 / self._steps_in_window + predicted_step_ms = 0.0 + realized_over_predicted = 0.0 + if self._predicted_steps_in_window > 0: + predicted_step_ms = self._predicted_runtime_ms_in_window / self._predicted_steps_in_window + if predicted_step_ms > 0: + realized_over_predicted = realized_step_ms / predicted_step_ms + + # Cumulative device peak (NOT reset) -> OOM-headroom gate for Nmax/Tmax sweeps. + peak_mem_gb = torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0 + + log.info( + f"{iteration} : tokens_per_sec_per_gpu {tokens_per_sec_per_gpu:.1f} | " + f"useful_tokens_per_sec_per_gpu {useful_tokens_per_sec_per_gpu:.1f} | " + f"useful_supervised_tokens_per_sec_per_gpu {useful_supervised_tokens_per_sec_per_gpu:.1f} | " + f"tokens_per_step {tokens_per_step:.1f} | samples_per_step {samples_per_step:.2f} | " + f"sec_per_sample {sec_per_sample:.4f} | " + f"useful_util {useful_util:.3f} | rho_sup {rho_sup:.3f} | " + f"attn_quadratic_waste {attention_quadratic_waste:.3f} | " + f"seq_max_len_mean {seq_max_len_mean:.0f} | seq_max_len_max {seq_max_len_max} | " + f"useful_mfu {useful_mfu:.3f} | compute_bound {compute_bound} | " + f"comm_bytes_per_useful_token {comm_bytes_per_useful_token:.1f} | " + f"realized_step_ms {realized_step_ms:.0f} | predicted_step_ms {predicted_step_ms:.0f} | " + f"realized_over_predicted {realized_over_predicted:.2f} | " + f"singleton_rate {singleton_rate:.3f} | " + f"peak_mem_gb {peak_mem_gb:.1f} | steps_in_window {self._steps_in_window}", + rank0_only=False, + ) + + if wandb.run: + wandb.log( + { + "throughput/tokens_per_sec_per_gpu": tokens_per_sec_per_gpu, + "throughput/useful_tokens_per_sec_per_gpu": useful_tokens_per_sec_per_gpu, + "throughput/useful_supervised_tokens_per_sec_per_gpu": useful_supervised_tokens_per_sec_per_gpu, + "throughput/tokens_per_step": tokens_per_step, + "throughput/samples_per_step": samples_per_step, + "throughput/sec_per_sample": sec_per_sample, + "throughput/useful_mfu": useful_mfu, + "throughput/compute_bound": compute_bound, + "throughput/singleton_rate": singleton_rate, + "throughput/peak_mem_gb": peak_mem_gb, + "packing/useful_util": useful_util, + "packing/rho_sup": rho_sup, + "packing/attention_quadratic_waste": attention_quadratic_waste, + "packing/seq_max_len_mean": seq_max_len_mean, + "packing/seq_max_len_max": seq_max_len_max, + "comm/comm_bytes_per_useful_token": comm_bytes_per_useful_token, + "cost_model/predicted_step_ms": predicted_step_ms, + "cost_model/realized_step_ms": realized_step_ms, + "cost_model/realized_over_predicted": realized_over_predicted, + }, + step=iteration, + ) + + # Reset the window (peak mem is intentionally left cumulative). + self._tokens_in_window = 0 + self._useful_tokens_in_window = 0 + self._samples_in_window = 0 + self._singleton_steps_in_window = 0 + self._steps_in_window = 0 + self._supervised_tokens_in_window = 0 + self._padded_attn_in_window = 0 + self._content_attn_in_window = 0 + self._seq_max_len_sum = 0 + self._seq_max_len_max = 0 + self._seq_max_len_steps = 0 + self._predicted_runtime_ms_in_window = 0.0 + self._predicted_steps_in_window = 0 + self._window_start_time = time.time() diff --git a/cosmos_framework/callbacks/wandb_log_simple.py b/cosmos_framework/callbacks/wandb_log_simple.py deleted file mode 100644 index 6087a337..00000000 --- a/cosmos_framework/callbacks/wandb_log_simple.py +++ /dev/null @@ -1,155 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Tuple - -import torch -import torch.distributed as dist -import torch.utils.data -import wandb - -from cosmos_framework.model._base import ImaginaireModel -from cosmos_framework.utils import distributed, log -from cosmos_framework.utils.callback import Callback -from cosmos_framework.utils.easy_io import easy_io - - -@dataclass -class _LossRecord: - loss: float = 0 - iter_count: int = 0 - name: str = None - - def reset(self) -> None: - self.loss = 0 - self.iter_count = 0 - - def get_stat(self, return_valid_mask_sum: bool = False) -> Tuple[float, float]: - if self.iter_count == 0: - self.loss = torch.tensor([float("nan")], device="cuda") # [1] - self.iter_count = 1 - msg_str = f"{self.name}: sum_loss={self.loss.item()}/iter_count={self.iter_count}=" - avg_loss_tensor = self.loss / self.iter_count - # Create a mask (1 if valid, 0 if NaN or Inf) - valid_mask = torch.tensor([torch.isfinite(avg_loss_tensor).float()], device="cuda") # [1] - msg_str += f"avg_loss={avg_loss_tensor.item()}, valid_mask={valid_mask.item()}, " - - # Replace NaN/Inf with 0 to avoid affecting sum - avg_loss_tensor = torch.where( - torch.isfinite(avg_loss_tensor), - avg_loss_tensor, - torch.tensor([0.0], device="cuda"), # [1] - ) - - # Reduce across all ranks - dist.all_reduce(avg_loss_tensor, op=dist.ReduceOp.SUM) # Sum of valid losses - dist.all_reduce(valid_mask, op=dist.ReduceOp.SUM) # Count of valid losses - msg_str += f" | all_reduce: avg_loss={avg_loss_tensor.item()}, valid_mask={valid_mask.item()}" - # Compute final average, avoiding division by zero - if valid_mask.item() > 0: - final_avg_loss = (avg_loss_tensor / valid_mask).item() - valid_mask_sum = valid_mask.item() - else: - final_avg_loss = 0.0 # Default to zero if all values were invalid - valid_mask_sum = 0 - - avg_loss = final_avg_loss - msg_str += f" | final: avg_loss={final_avg_loss}" - if self.name is not None: - log.debug(msg_str, rank0_only=False) - self.reset() - if return_valid_mask_sum: - return avg_loss, valid_mask_sum - else: - return avg_loss - - -class WandbCallback(Callback): - def __init__( - self, - logging_iter_multipler: int = 1, - save_logging_iter_multipler: int = 1, - save_s3: bool = False, - ) -> None: - super().__init__() - self.final_loss_log = _LossRecord() - self.final_all_loss_log = {} - self.logging_iter_multipler = logging_iter_multipler - self.save_logging_iter_multipler = save_logging_iter_multipler - assert self.logging_iter_multipler > 0, "logging_iter_multipler should be greater than 0" - self.save_s3 = save_s3 - self.wandb_extra_tag = f"@{logging_iter_multipler}" if logging_iter_multipler > 1 else "" - self.name = "wandb_loss_log" + self.wandb_extra_tag - self.unstable_count = torch.zeros(1, device="cuda") # [1] - - def on_training_step_end( - self, - model: ImaginaireModel, - data_batch: dict[str, torch.Tensor], - output_batch: dict[str, torch.Tensor], - loss: torch.Tensor, - iteration: int = 0, - ) -> None: - if torch.isnan(loss) or torch.isinf(loss): - log.critical( - f"Unstable loss {loss} at iteration {iteration}", - rank0_only=False, - ) - self.unstable_count += 1 - - self.final_loss_log.loss += loss.detach().float() - self.final_loss_log.iter_count += 1 - - for key in output_batch.keys(): - if "loss" in key: - if key not in self.final_all_loss_log: - self.final_all_loss_log[key] = _LossRecord() - self.final_all_loss_log[key].loss += output_batch[key].detach().float() - self.final_all_loss_log[key].iter_count += 1 - - if iteration % (self.config.trainer.logging_iter * self.logging_iter_multipler) == 0: - avg_final_loss = self.final_loss_log.get_stat() - - avg_final_all_loss = {} - for key in self.final_all_loss_log.keys(): - avg_final_all_loss[key] = self.final_all_loss_log[key].get_stat() - - dist.all_reduce(self.unstable_count, op=dist.ReduceOp.SUM) - - if distributed.is_rank0() and wandb.run is not None: - info = {} - info.update( - { - f"train{self.wandb_extra_tag}/loss": avg_final_loss, - f"train{self.wandb_extra_tag}/unstable_count": self.unstable_count.item(), - "iteration": iteration, - } - ) - for key, loss in avg_final_all_loss.items(): - info.update( - { - f"train{self.wandb_extra_tag}_detail/{key}": loss, - } - ) - if self.save_s3: - if ( - iteration - % ( - self.config.trainer.logging_iter - * self.logging_iter_multipler - * self.save_logging_iter_multipler - ) - == 0 - ): - easy_io.dump( - info, - f"s3://rundir/{self.name}/Train_Iter{iteration:09d}.json", - ) - - if wandb: - wandb.log(info, step=iteration) - # reset unstable count - self.unstable_count.zero_() diff --git a/cosmos_framework/checkpoint/dcp.py b/cosmos_framework/checkpoint/dcp.py index 7318e4a7..d3cad0ea 100644 --- a/cosmos_framework/checkpoint/dcp.py +++ b/cosmos_framework/checkpoint/dcp.py @@ -7,8 +7,8 @@ The checkpointer saves model state in a sharded format across multiple processes: self.save_dirname/ -├── iter_000000005/ # Checkpoint at iteration 5 -│ ├── model/ # Model state shards +├── iter_000000005/ # Checkpoint at iteration 5 +│ ├── model/ # Model state shards │ │ ├── __0_0.distcp # Shard 0 from rank 0 │ │ └── __1_0.distcp # Shard 1 from rank 1 │ ├── optim/ # Optimizer state shards @@ -39,30 +39,34 @@ 3. Supporting both local and cloud storage backends """ +import dataclasses import enum import multiprocessing import os import re import time from multiprocessing import get_context -from typing import Any, Dict, List, Optional, Protocol, Tuple, Union, runtime_checkable +from typing import Any, Dict, Optional, Protocol, Tuple, Union, runtime_checkable import torch import torch.distributed as dist import torch.distributed.checkpoint as dcp from torch import nn +from torch.distributed.checkpoint.default_planner import create_default_local_load_plan from torch.distributed.checkpoint.filesystem import FileSystemReader, FileSystemWriter from torch.distributed.checkpoint.metadata import ( STATE_DICT_TYPE, Metadata, StorageMeta, ) +from torch.distributed.checkpoint.planner import LoadPlan from torch.distributed.checkpoint.state_dict import ( StateDictOptions, get_model_state_dict, set_model_state_dict, ) from torch.distributed.checkpoint.stateful import Stateful +from torch.distributed.tensor import DTensor, Replicate from torch.nn.modules.module import _IncompatibleKeys from cosmos_framework.checkpoint.base import AbstractCheckpointer @@ -269,6 +273,62 @@ class CustomLoadPlanner(dcp.DefaultLoadPlanner): CustomLoadPlanner that supports ignoring keys during checkpoint load. This is useful when the checkpoint is saved with a different component architecture, e.g. different RoPE embeddings than the current model. + + Setting ``keys_to_skip_loading``: do not construct this planner directly. Set the experiment's + checkpoint config field ``checkpoint.keys_to_skip_loading`` (a ``list[str]``; see + :class:`cosmos_framework.utils.config.CheckpointConfig`) and the checkpointer forwards it here. Each entry is + matched as a *substring* against every flattened fully-qualified name in the model/optimizer state + dict, and any leaf whose fqn contains one of the entries is skipped. The list is only honored on a + warm start (loading from a *different* run via ``checkpoint.load_path``); when resuming the latest + checkpoint of the same run it is forced to ``[]``, since there is nothing to skip. Examples: + + 1. Skip a reshaped positional embedding (different sequence length) for reg + EMA copies. + checkpoint.keys_to_skip_loading=["net.latent_pos_embed.seq", "net_ema.latent_pos_embed.seq"] + 2. Skip action-head layers when warm-starting an action model from a non-action checkpoint. + checkpoint.keys_to_skip_loading=["action2llm", "llm2action", "action_modality_embed", "action_pos_embed"] + + When ``dedup=True`` it additionally elects a single reader per replicated leaf to kill + redundant storage reads. Two replication patterns waste reads at large world sizes: + + 1. Fully-replicated leaves (e.g. optimizer scalar ``step`` tensors, replicated buffers) are + saved on global rank 0 but appear in *every* rank's state dict, so all ``world_size`` ranks + read the same object — a single-object hotspot. + 2. HSDP shards are identical across the ``dp_replicate`` dim, so each shard file is read by + ``dp_replicate`` ranks. + + With ``dedup=True``, ``create_local_plan`` drops, from the read plan it builds, the items this + rank is *not* the designated reader for, so each leaf is fetched from storage exactly once. + + ``create_local_plan`` always builds the read plan by calling ``create_default_local_load_plan`` + directly on ``self._skip_keys_if_found(self.state_dict)`` rather than delegating to + ``super().create_local_plan()``. This serves two ends: + + 1. It drops ``keys_to_skip_loading`` *before* plan construction, so skipped keys never reach the + base helper's strict missing-key validation or its *unconditional* size-mismatch check. + Skipping therefore works with ``strict_resume=True``, with a skipped key absent from the + checkpoint, and with a skipped key present-but-reshaped (e.g. different RoPE embeddings). + 2. It avoids the base ``create_local_plan``'s pre-2.4 "missing keys" fallback, which would + re-derive ``self.state_dict`` from the full ``original_state_dict`` and reinstate any skipped + key still present in the checkpoint. Dropping that fallback is acceptable -- this repo does + not load pre-2.4 checkpoints. + + ``self.state_dict`` is never mutated (``_skip_keys_if_found`` returns a pruned *copy*, or the dict + unchanged when nothing matches), so the loader still writes non-tensor leaves back into the caller's + state dict in place. + + Dedup removes reads at a different point: ``create_local_plan`` filters the already-built plan's + read items via ``_drop_non_reader_items``, looking each leaf up in the full ``self.state_dict``. + + - ``DTensor`` leaf -> reader is local rank 0 along the tensor's replicate mesh dim (one reader + per replicate group; a fully-sharded tensor has none, so every rank reads its own shard); + - non-``DTensor`` leaf (plain tensor / python scalar) -> fully replicated across the world, so + the reader is global rank 0. + + Reading the replicate dim from each tensor's *own* mesh (rather than one global ``dp_replicate`` + group) generalizes to tensors on different meshes, e.g. MoE expert weights under expert + parallelism. The dropped data is filled in afterwards by :func:`_broadcast_state_dict`. + Because each rank then reads a disjoint subset, dedup requires ``no_dist=True`` (there is no + global plan to coordinate). """ def __init__( @@ -276,18 +336,26 @@ def __init__( flatten_state_dict: bool = True, flatten_sharded_tensors: bool = True, allow_partial_load: bool = False, - keys_to_skip_loading: List[str] = [], + keys_to_skip_loading: list[str] | None = None, load_ema_to_reg: bool = False, + dedup: bool = False, + global_rank: int = 0, ) -> None: super().__init__( flatten_state_dict=flatten_state_dict, flatten_sharded_tensors=flatten_sharded_tensors, allow_partial_load=allow_partial_load, ) - self.keys_to_skip_loading = keys_to_skip_loading + + # Default to [] without a mutable default argument (which would be shared across instances). + self.keys_to_skip_loading = keys_to_skip_loading or [] self.load_ema_to_reg = load_ema_to_reg - if len(keys_to_skip_loading) > 0: - log.info(f"Skipping loading of keys that match the following patterns: {keys_to_skip_loading}") + # When set, prune non-reader leaves so each replicated leaf is read by exactly one rank. + self.dedup = dedup + self._global_rank = global_rank + + if len(self.keys_to_skip_loading) > 0: + log.info(f"Skipping loading of keys that match the following patterns: {self.keys_to_skip_loading}") def set_up_planner( self, @@ -295,8 +363,6 @@ def set_up_planner( metadata: Metadata | None = None, is_coordinator: bool = False, ) -> None: - state_dict = self._skip_keys_if_found(state_dict) - if self.load_ema_to_reg: state_dict = _replace_keys_with_ema_keys(state_dict) @@ -306,24 +372,251 @@ def set_up_planner( is_coordinator=is_coordinator, ) - def _skip_keys_if_found( - self, - state_dict: STATE_DICT_TYPE, - ) -> Dict[str, Any]: + def create_local_plan(self) -> LoadPlan: + # Build the read plan from a pruned *copy* of self.state_dict so skipped keys never + # reach create_default_local_load_plan, which would otherwise (a) raise a missing-key + # error for a skipped key absent from the checkpoint under strict load + # (allow_partial_load=False), and (b) raise an *unconditional* size-mismatch ValueError + # for a skipped key that is present in the checkpoint but reshaped. + if self.metadata is None: + raise AssertionError("metadata must be set (via set_up_planner) before create_local_plan") + + plan = create_default_local_load_plan( + state_dict=self._skip_keys_if_found(self.state_dict), + metadata=self.metadata, + strict=not self.allow_partial_load, + ) + + return self._drop_non_reader_items(plan) + + def _drop_non_reader_items(self, plan: LoadPlan) -> LoadPlan: + """Under dedup load, drop the read items this rank is not the elected reader for. + + ``self.state_dict`` is the full flattened fqn -> leaf mapping at this point, so we look + each item's leaf up by fqn to decide readership. The dropped reads are filled afterward by + :func:`_broadcast_state_dict`. No-op unless ``self.dedup`` is set. """ - While loading the checkpoint, skip the weight loading for the keys - that contain any element of `self.keys_to_skip_loading` as a substring. + if not self.dedup: + return plan + + kept_items = [ + item + for item in plan.items + if _is_assigned_reader(self.state_dict.get(item.dest_index.fqn), self._global_rank) + ] + dropped = len(plan.items) - len(kept_items) + log.info( + f"[DCP-LOAD-DEDUP] kept_read_items={len(kept_items)} dropped_read_items={dropped}", + rank0_only=False, + ) + return dataclasses.replace(plan, items=kept_items) + + def _skip_keys_if_found(self, state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE: + """Return a *copy* of ``state_dict`` without keys whose fqn contains any ``keys_to_skip_loading`` substring. + + ``create_local_plan`` feeds this pruned copy to ``create_default_local_load_plan`` so skipped + keys are kept out of the read plan, the base planner's strict missing-key validation, and its + unconditional size-mismatch check (skipping therefore works with ``strict_resume=True``, with a + skipped key absent from the checkpoint, and with a skipped key present-but-reshaped). The caller + never mutates ``self.state_dict`` with this result, so the loader still writes non-tensor leaves + back into the caller's state dict in place. No-op when no skip patterns are configured. """ if len(self.keys_to_skip_loading) == 0: return state_dict - new_state_dict = {} + kept = {} for fqn, obj in state_dict.items(): if any(skip_key in fqn for skip_key in self.keys_to_skip_loading): log.warning(f"Skipping loading of key: {fqn}") continue - new_state_dict[fqn] = obj - return new_state_dict + kept[fqn] = obj + log.info( + f"[DCP-LOAD-SKIP-KEYS] kept_keys={len(kept)} dropped_keys={len(state_dict) - len(kept)}", + rank0_only=False, + ) + return kept + + +def _is_assigned_reader(value: Any, global_rank: int) -> bool: + """Whether this rank is the single elected reader for ``value`` under dedup load. + + Used by :meth:`CustomLoadPlanner.create_local_plan` to drop the read items a rank is not + responsible for, and mirrored by :func:`_broadcast_state_dict` to fill those dropped leaves back + in. The election rule depends on how the leaf is replicated: + + - ``DTensor`` -> reader iff this rank sits at local rank 0 along EVERY (size>1) replicate mesh + dim (see :func:`_replicate_mesh_dims`). A fully-sharded tensor has no replicate dims, so the + ``all(...)`` over an empty list is ``True`` and every rank reads its own (unique) shard. A + fully-replicated leaf (e.g. an optimizer ``step`` with ``[Replicate(), Replicate()]``) elects + exactly one reader instead of a whole ``dp_shard`` line. + - any other leaf (plain replicated tensor or python scalar) -> fully replicated across the world, + so global rank 0 is the sole reader. + """ + if isinstance(value, DTensor): + dims = _replicate_mesh_dims(value) + return all(value.device_mesh.get_local_rank(dim) == 0 for dim in dims) + return global_rank == 0 + + +def _replicate_mesh_dims(value: DTensor) -> list[int]: + """Mesh-dim indices (ascending) along which a ``DTensor``'s local shard is replicated. + + A ``DTensor``'s local shard is byte-identical across every mesh dim whose placement is + ``Replicate()``; it differs along ``Shard``/``Partial`` dims. Reading those dims' data once and + broadcasting along them is therefore lossless. Deriving this from the tensor's *own* mesh (not a + single global ``dp_replicate`` group) is what lets the dedup generalize to tensors on other + meshes — e.g. MoE expert weights on a future expert-parallel mesh whose replicate dim differs. + + Mesh dims of **size 1** are excluded: they carry no replication, and — critically — every rank + is local rank 0 along a size-1 dim, so treating one as a replicate dim would make the + single-reader election (see :func:`_is_assigned_reader`) pass on *every* rank + and silently disable dedup along the real replicate dim. This is exactly the failure for a + fully-replicated leaf (e.g. an optimizer ``step`` with ``[Replicate(), Replicate()]``) on the + 2-D ``(dp_replicate=1, dp_shard=N)`` FSDP mesh: without this filter the size-1 ``dp_replicate`` + dim is picked, so all ``N`` ranks read the single object in ``__0_0.distcp`` instead of just + global rank 0. + + Returns multiple dims for a leaf replicated across several mesh dims (e.g. a fully-replicated + tensor on an HSDP ``(dp_replicate>1, dp_shard)`` mesh); empty when the tensor is fully sharded + (no replication, so every rank legitimately reads its own shard). + """ + mesh = value.device_mesh + return [i for i, placement in enumerate(value.placements) if isinstance(placement, Replicate) and mesh.size(i) > 1] + + +def _broadcast_tensor_leaf(value: torch.Tensor) -> None: + """Broadcast a tensor leaf from its single reader to the ranks that share it. + + Handles the two tensor leaf kinds the dedup planner drops (see + :func:`_is_assigned_reader`): + + - ``DTensor`` -> broadcast the local shard across the tensor's replicate mesh dims. The reader + is the rank at local rank 0 along every (size>1) replicate dim; the local shard is + byte-identical across those dims, so broadcasting from that rank is lossless. Groups are taken + from the tensor's *own* mesh so this works regardless of which mesh the tensor lives on (e.g. a + future expert-parallel mesh). Dims are broadcast in ascending order so each step's source + already holds valid data: after broadcasting along dim ``d``, every rank that is local rank 0 + along the remaining (higher) replicate dims is populated, which is exactly the source set the + next broadcast needs. A fully-sharded ``DTensor`` (no replicate dims) needs no broadcast — every + rank already read its own shard. + - plain (non-``DTensor``) tensor -> fully replicated across the world and read by global rank 0 + only, so broadcast it over the default (world) group with ``src=0``. This covers replicated + ``param_groups`` tensors such as a capturable optimizer ``step``. + """ + if not torch.is_tensor(value): + raise ValueError(f"Unsupported type: {type(value)}") + if isinstance(value, DTensor): + dims = _replicate_mesh_dims(value) + if not dims: + # Fully sharded across its mesh: every rank already read its own shard. + return + local = value.to_local() + for dim in dims: + group = value.device_mesh.get_group(dim) + dist.broadcast(local, group=group, group_src=0) + else: + dist.broadcast(value, src=0) + + +def _broadcast_state_dict( + state_dict: STATE_DICT_TYPE, + global_rank: int, +) -> None: + """Broadcast every leaf from its single reader (see :class:`CustomLoadPlanner` dedup mode) to + the ranks that share it, filling in the reads that the planner dropped. + + Must be called by ALL ranks in lockstep, *after* a dedup ``dcp.load`` and *before* the + ``load_state_dict`` that consumes ``state_dict``. Each leaf's replication is a + structural property (identical on every rank), so all ranks walk the (identically ordered) tree + and issue the same sequence of collectives: + + - ``DTensor`` leaf -> broadcast along its mesh's replicate dims (src = local rank 0 there); a + fully-sharded tensor needs no broadcast; + - non-``DTensor`` tensor leaf (e.g. a replicated ``param_groups`` ``step`` tensor) -> replicated + across the world, broadcast (src = global rank 0); + - non-tensor leaf (python scalars, e.g. optimizer ``param_groups`` ``lr``/``betas``/``eps``/ + ``step``) -> a single world ``broadcast_object_list`` of the tensor-stripped skeleton, merged + back in. + + Reader election here mirrors :func:`_is_assigned_reader` exactly so that the + set of leaves broadcast is precisely the set the planner dropped: DTensors elect a per-mesh + reader; every other leaf (plain tensor or scalar) is read by global rank 0 only. + """ + + # 1) Broadcast tensor leaves in place. Keys are sorted by repr() so the traversal order is + # identical across ranks regardless of key type (int param ids vs str fqns). DTensors go over + # their own mesh's replicate dims; plain (non-DTensor) tensors are world-replicated and read by + # global rank 0 only, so they broadcast from global src 0. Non-tensor leaves are skipped here + # and handled by the object broadcast in step 2. + def _walk_tensors(obj: Any) -> None: + if torch.is_tensor(obj): + _broadcast_tensor_leaf(obj) + elif isinstance(obj, dict): + for k in sorted(obj.keys(), key=repr): + _walk_tensors(obj[k]) + elif isinstance(obj, (list, tuple)): + for v in obj: + _walk_tensors(v) + # else: non-tensor scalar leaf -> handled in step 2. + + start_time_tensor_broadcast = time.monotonic() + _walk_tensors(state_dict) + end_time_tensor_broadcast = time.monotonic() + + # 2) Broadcast all non-tensor leaves as one tensor-stripped skeleton (world, src = rank 0). + # Tensors (already broadcast in step 1) are stripped to None so they are not pickled; every + # other value (python scalars, lists, None, ...) rides along in the object broadcast. + def _strip(obj: Any) -> Any: + if torch.is_tensor(obj): + return None + if isinstance(obj, dict): + return {k: _strip(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return type(obj)(_strip(v) for v in obj) + return obj + + start_time_object_broadcast = time.monotonic() + box = [_strip(state_dict) if global_rank == 0 else None] + dist.broadcast_object_list(box, src=0) + skeleton = box[0] + end_time_object_broadcast = time.monotonic() + + # 3) Rebuild: keep the (already-broadcast) tensors in place, take non-tensors from rank 0. + def _rebuild(local_obj: Any, skeleton_obj: Any) -> Any: + if torch.is_tensor(local_obj): + return local_obj + if isinstance(local_obj, dict): + return {k: _rebuild(local_obj[k], skeleton_obj[k]) for k in local_obj} + if isinstance(local_obj, list): + return [_rebuild(lv, sv) for lv, sv in zip(local_obj, skeleton_obj)] + if isinstance(local_obj, tuple): + return tuple(_rebuild(lv, sv) for lv, sv in zip(local_obj, skeleton_obj)) + return skeleton_obj + + start_time_rebuild = time.monotonic() + rebuilt = _rebuild(state_dict, skeleton) + state_dict.clear() + state_dict.update(rebuilt) + end_time_rebuild = time.monotonic() + + log.info( + "[DCP-LOAD-TIMING] broadcast: " + f"tensors={end_time_tensor_broadcast - start_time_tensor_broadcast:.1f}s, " + f"objects={end_time_object_broadcast - start_time_object_broadcast:.1f}s, " + f"rebuild={end_time_rebuild - start_time_rebuild:.1f}s", + rank0_only=False, + ) + + +def _copy_ema_weights_to_reg(state_dict: STATE_DICT_TYPE) -> None: + """Copy EMA weights to regular weights.""" + for sd_key in list(state_dict.keys()): + if sd_key.startswith("net."): + key_ema = "net_ema." + sd_key.removeprefix("net.") + assert key_ema in state_dict, ( + f"EMA key {key_ema} not found in state_dict. Ensure the model has net_ema submodule." + ) + state_dict[sd_key] = state_dict[key_ema] class CustomSavePlanner(dcp.DefaultSavePlanner): @@ -383,6 +676,13 @@ def __init__( ): super().__init__(config_checkpoint, config_job, callbacks) self.config_checkpoint = config_checkpoint + if config_checkpoint.load_ema_to_reg and config_checkpoint.load_training_state: + raise ValueError( + "load_ema_to_reg=True requires load_training_state=False. " + "Loading optimizer/EMA state after copying EMA->reg weights produces an " + "inconsistent checkpoint: the optimizer moments would track the original " + "reg-weight trajectory, not the EMA weights just copied in." + ) if config_checkpoint.dcp_async_mode_enabled and not disable_async: self.async_mode = AsyncMode.ASYNC_WITH_PINNED_MEM else: @@ -497,6 +797,7 @@ def load( log.critical(f"Resuming ckpt {checkpoint_path} with keys: {resume_keys}") iteration = 0 + global_rank = dist.get_rank() if dist.is_initialized() else 0 if checkpoint_path is not None: self._check_checkpoint_exists(checkpoint_path) @@ -514,10 +815,21 @@ def load( # the latest checkpoint of the same model, then we don't need to skip any keys. keys_to_skip_loading = self.config_checkpoint.keys_to_skip_loading if warm_start else [] - load_planner = CustomLoadPlanner( - allow_partial_load=not strict_resume, - keys_to_skip_loading=keys_to_skip_loading, - ) + # Dedup-load context: when enabled, replicated state is read by a single rank and + # broadcast over the device mesh instead of being read redundantly by every rank. The + # per-tensor replicate group is derived from each DTensor's own mesh (see + # _broadcast_state_dict). Applied only to the "model" and "optim" components + # (the bulk of the data); the per-rank-unique "dataloader" state and + # "trainer" RNG and the tiny "scheduler" state stay on the normal read-everywhere path. + if key in ("model", "optim"): + load_planner = CustomLoadPlanner( + allow_partial_load=not strict_resume, + keys_to_skip_loading=keys_to_skip_loading, + dedup=self.config_checkpoint.dcp_load_dedup, + global_rank=global_rank, + ) + else: + load_planner = dcp.DefaultLoadPlanner(allow_partial_load=not strict_resume) if key == "model": log.info("- Loading the model...") @@ -527,19 +839,26 @@ def load( _state_dict, storage_reader=storage_reader, planner=load_planner, + no_dist=True, ) - if self.config_checkpoint.load_ema_to_reg: + if self.config_checkpoint.dcp_load_dedup: + # Fill in the reads the planner dropped by broadcasting from each leaf's + # single reader. Must run before the EMA copy and load_state_dict below. + _broadcast_state_dict(_state_dict, global_rank) + if self.config_checkpoint.load_ema_to_reg and warm_start: # The model has both net.* and net_ema.* submodules, so _state_dict # contains both sets of keys after dcp.load(). Copy EMA weights into - # regular model weights so we can resume from EMA and reset EMA. - for sd_key in list(_state_dict.keys()): - if sd_key.startswith("net."): - key_ema = "net_ema." + sd_key.removeprefix("net.") - assert key_ema in _state_dict, ( - f"EMA key {key_ema} not found in state_dict. " - "Ensure the model has net_ema submodule." - ) - _state_dict[sd_key] = _state_dict[key_ema] + # regular model weights so we can warm-start from EMA (and reset EMA + # when load_training_state=False). + # + # This must only run on warm start. During regular auto-resume + # (warm_start=False) the full training state is always reloaded, so + # copying EMA->reg here would silently overwrite the resumed regular + # weights with the (lagging) EMA snapshot on every restart while the + # optimizer state still tracks the pre-copy trajectory -- corrupting + # training in a way that depends on how often the job is preempted. + _copy_ema_weights_to_reg(_state_dict) + results = _model_wrapper.load_state_dict(_state_dict) if len(results.missing_keys) > 0: raise ValueError(f"Missing keys (not found in checkpoint): {results.missing_keys}") @@ -555,7 +874,12 @@ def load( _state_dict, storage_reader=storage_reader, planner=load_planner, + no_dist=True, ) + if self.config_checkpoint.dcp_load_dedup: + # Fill in the reads the planner dropped by broadcasting from each leaf's + # single reader. Must run before load_state_dict below. + _broadcast_state_dict(_state_dict, global_rank) optimizer.load_state_dict(_state_dict) elif key == "scheduler": @@ -565,6 +889,7 @@ def load( _state_dict, storage_reader=storage_reader, planner=load_planner, + no_dist=True, ) scheduler.load_state_dict(_state_dict) @@ -590,6 +915,7 @@ def load( _state_dict, storage_reader=storage_reader, planner=load_planner, + no_dist=True, ) grad_scaler.load_state_dict(_state_dict["grad_scaler"]) iteration = _state_dict["iteration"] diff --git a/cosmos_framework/configs/base/config.py b/cosmos_framework/configs/base/config.py index e766c5c9..99fbd808 100644 --- a/cosmos_framework/configs/base/config.py +++ b/cosmos_framework/configs/base/config.py @@ -77,7 +77,7 @@ def make_config() -> Config: from cosmos_framework.configs.base.defaults.model import register_model from cosmos_framework.configs.base.defaults.optimizer import register_optimizer, register_scheduler from cosmos_framework.configs.base.defaults.tokenizer import register_sound_tokenizer, register_tokenizer - from cosmos_framework.configs.base.defaults.vlm import register_vlm + from cosmos_framework.configs.base.defaults.reasoner import register_vlm # Call this function to register config groups for advanced overriding. the order follows the default config groups # register_data() diff --git a/cosmos_framework/configs/base/defaults/checkpointer.py b/cosmos_framework/configs/base/defaults/checkpointer.py index 94d86902..9502070e 100644 --- a/cosmos_framework/configs/base/defaults/checkpointer.py +++ b/cosmos_framework/configs/base/defaults/checkpointer.py @@ -25,7 +25,7 @@ s3_object_store = config.ObjectStoreConfig( enabled=True, - credentials="credentials/s3_training.secret", + credentials="credentials/s3_checkpoint.secret", bucket="bucket4", ) @@ -35,6 +35,12 @@ bucket="checkpoints-eu-west-3", ) +s3_east2_object_store = config.ObjectStoreConfig( + enabled=True, + credentials="credentials/s3_east2_checkpoint.secret", + bucket="bucket", +) + gcp_object_store = config.ObjectStoreConfig( enabled=True, credentials="credentials/gcp_checkpoint.secret", @@ -53,6 +59,7 @@ save_iter=5000, broadcast_via_filesystem=True, dcp_async_mode_enabled=True, + dcp_load_dedup=True, ) CHECKPOINT_PDX = CheckpointConfig( @@ -61,6 +68,7 @@ save_iter=5000, broadcast_via_filesystem=True, dcp_async_mode_enabled=True, + dcp_load_dedup=True, ) CHECKPOINT_S3 = CheckpointConfig( @@ -69,6 +77,7 @@ save_iter=5000, broadcast_via_filesystem=True, dcp_async_mode_enabled=True, + dcp_load_dedup=True, ) CHECKPOINT_S3_EU = CheckpointConfig( @@ -77,6 +86,7 @@ save_iter=5000, broadcast_via_filesystem=True, dcp_async_mode_enabled=True, + dcp_load_dedup=True, ) CHECKPOINT_GCP = CheckpointConfig( @@ -90,11 +100,25 @@ dcp_async_mode_enabled=True, ) +CHECKPOINT_S3_EAST2 = CheckpointConfig( + save_to_object_store=s3_east2_object_store, + save_iter=1000, + load_from_object_store=s3_east2_object_store, + load_path="", + load_training_state=False, + strict_resume=True, + enable_gcs_patch_in_boto3=False, + dcp_async_mode_enabled=True, + dcp_load_dedup=True, +) + CHECKPOINT_NEB_EU = CheckpointConfig( save_to_object_store=neb_eu_object_store, load_from_object_store=neb_eu_object_store, save_iter=2000, broadcast_via_filesystem=True, + dcp_async_mode_enabled=True, + dcp_load_dedup=True, ) @@ -106,6 +130,7 @@ def register_checkpoint(): cs.store(group="checkpoint", package="checkpoint", name="s3_eu", node=CHECKPOINT_S3_EU) cs.store(group="checkpoint", package="checkpoint", name="gcp", node=CHECKPOINT_GCP) cs.store(group="checkpoint", package="checkpoint", name="neb_eu", node=CHECKPOINT_NEB_EU) + cs.store(group="checkpoint", package="checkpoint", name="s3_east2", node=CHECKPOINT_S3_EAST2) DUMMY_CHECKPOINTER: Dict[str, str] = L(DummyCheckpointer)() diff --git a/cosmos_framework/configs/base/defaults/llm_model.py b/cosmos_framework/configs/base/defaults/llm_model.py deleted file mode 100644 index a4e55a9d..00000000 --- a/cosmos_framework/configs/base/defaults/llm_model.py +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -"""Register LLM and DiT model configs alongside the existing ``mot_fsdp``.""" - -from hydra.core.config_store import ConfigStore - -from cosmos_framework.utils.lazy_config import LazyCall as L -from cosmos_framework.model.vfm.llm.dit.image_dit_model import DiTPretrainModel, DiTPretrainModelConfig -from cosmos_framework.model.vfm.llm.llm_pretrain_model import LLMPretrainModel, LLMPretrainModelConfig - -# ── FSDP config (production, multi-GPU) ────────────────────────────────────── - -LLM_FSDP_CONFIG = dict( - trainer=dict( - distributed_parallelism="fsdp", - ), - model=L(LLMPretrainModel)( - config=LLMPretrainModelConfig(), - _recursive_=False, - ), -) - -# ── DDP config (debug, single-node) ───────────────────────────────────────── - -LLM_DDP_CONFIG = dict( - trainer=dict( - distributed_parallelism="ddp", - ), - model=L(LLMPretrainModel)( - config=LLMPretrainModelConfig(), - _recursive_=False, - ), -) - -# ── Image DiT configs ─────────────────────────────────────────────────────── - -DIT_FSDP_CONFIG = dict( - trainer=dict( - distributed_parallelism="fsdp", - ), - model=L(DiTPretrainModel)( - config=DiTPretrainModelConfig(), - _recursive_=False, - ), -) - -DIT_DDP_CONFIG = dict( - trainer=dict( - distributed_parallelism="ddp", - ), - model=L(DiTPretrainModel)( - config=DiTPretrainModelConfig(), - _recursive_=False, - ), -) - - -def register_llm_model(): - cs = ConfigStore.instance() - cs.store(group="model", package="_global_", name="llm_fsdp", node=LLM_FSDP_CONFIG) - cs.store(group="model", package="_global_", name="llm_ddp", node=LLM_DDP_CONFIG) - cs.store(group="model", package="_global_", name="dit_fsdp", node=DIT_FSDP_CONFIG) - cs.store(group="model", package="_global_", name="dit_ddp", node=DIT_DDP_CONFIG) diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py index 0695aab1..d0d47ca0 100644 --- a/cosmos_framework/configs/base/defaults/model_config.py +++ b/cosmos_framework/configs/base/defaults/model_config.py @@ -10,7 +10,7 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.ema import EMAConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig -from cosmos_framework.configs.base.defaults.vlm import VLMConfig +from cosmos_framework.configs.base.defaults.reasoner import VLMConfig @attrs.define(slots=False) @@ -24,21 +24,8 @@ class DiffusionExpertConfig: max_vae_latent_side_after_patchify: int = ( 20 # Max dimension (h or w) of the VAE latent after patchification (320/(8*2)) ) - # Position embedding type for vision tokens: - # - "3d_rope": Additive 3D RoPE embeddings (VideoRopePosition3DEmb) + 1D position IDs for attention - # - "flattened_sin_cos": Additive flattened sin/cos embeddings + 1D position IDs for attention - # - "unified_3d_mrope": No additive embedding + 3D position IDs for Qwen3VL-style mRoPE attention - position_embedding_type: str = "3d_rope" - # When finetuning from lower resolution to higher resolution, the spatial resolution of videos increase. - # So, we need to adjust the position embedding. - # We use NTK based RoPE extrapolation to adjust the position embedding. - # Reference: (https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/) - # Design adapted from Cosmos2.5 (https://arxiv.org/pdf/2511.00062) - # extrapolation_ratio here is how the base of the RoPE is scaled - # b' = b * extrapolation_ratio^(dim / (dim - 2)) - rope_h_extrapolation_ratio: float = 1.0 - rope_w_extrapolation_ratio: float = 1.0 - rope_t_extrapolation_ratio: float = 1.0 + # Vision/action/sound position information is always provided through + # Qwen3VL-style 3D mRoPE attention IDs. enable_fps_modulation: bool = False base_fps: int = 24 # Base temporal compression factor for SOUND m-RoPE. None = current behavior diff --git a/cosmos_framework/configs/base/defaults/open_source_dataloader.py b/cosmos_framework/configs/base/defaults/open_source_dataloader.py index af0d4852..4dd42b33 100644 --- a/cosmos_framework/configs/base/defaults/open_source_dataloader.py +++ b/cosmos_framework/configs/base/defaults/open_source_dataloader.py @@ -26,7 +26,7 @@ from hydra.core.config_store import ConfigStore -from cosmos_framework.configs.base.defaults.vlm import create_qwen2_tokenizer_with_download +from cosmos_framework.configs.base.defaults.reasoner import create_qwen2_tokenizer_with_download from cosmos_framework.data.vfm.joint_dataloader import ( PackingDataLoader, RankPartitionedDataLoader, diff --git a/cosmos_framework/configs/base/defaults/vlm.py b/cosmos_framework/configs/base/defaults/reasoner.py similarity index 83% rename from cosmos_framework/configs/base/defaults/vlm.py rename to cosmos_framework/configs/base/defaults/reasoner.py index 32718f43..7293b6e6 100644 --- a/cosmos_framework/configs/base/defaults/vlm.py +++ b/cosmos_framework/configs/base/defaults/reasoner.py @@ -52,6 +52,9 @@ def download_tokenizer_files(model_name: str, config_variant: str) -> str: if config_variant == "s3": ckpt_bucket = "bucket4" credentials = "credentials/s3_checkpoint.secret" + elif config_variant == "s3_east2": + ckpt_bucket = "nv-cosmos-checkpoint-us-east-2" + credentials = "credentials/s3_east2_checkpoint.secret" elif config_variant == "gcp": ckpt_bucket = "bucket0" credentials = "credentials/gcp_checkpoint.secret" @@ -214,7 +217,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -240,7 +243,7 @@ class VLMConfig: model_instance=L(Qwen3VLMoeTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoeMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json" ), layer_module="Qwen3VLMoeTextMoTDecoderLayer", qk_norm_for_text=True, @@ -263,7 +266,7 @@ class VLMConfig: model_instance=L(Qwen3VLMoeTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoeMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json" ), layer_module="Qwen3VLMoeTextMoTDecoderLayer", qk_norm_for_text=True, @@ -288,7 +291,7 @@ class VLMConfig: model_instance=L(Qwen3VLMoeTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoeMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json" ), layer_module="Qwen3VLMoeTextMoTDecoderLayer", qk_norm_for_text=True, @@ -314,7 +317,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" ), qk_norm_for_text=True, ), @@ -334,7 +337,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" ), qk_norm_for_text=True, ), @@ -355,7 +358,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" ), qk_norm_for_text=True, ), @@ -371,7 +374,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -393,7 +396,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -415,7 +418,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" ), qk_norm_for_text=True, ), @@ -438,7 +441,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json" ), qk_norm_for_text=True, ), @@ -458,7 +461,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json" ), qk_norm_for_text=True, ), @@ -481,7 +484,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -501,7 +504,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -522,7 +525,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -543,7 +546,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -564,7 +567,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -580,6 +583,26 @@ class VLMConfig: ), ) +Cosmos3NanoReasoner_VLM_S3_EAST2_Config_0517: VLMConfig = VLMConfig( + model_name="nvidia/Cosmos3-Nano-Reasoner", + model_instance=L(Qwen3VLTextForCausalLM)( + config=L(create_vlm_config)( + base_config=L(Qwen3VLMoTConfig.from_json_file)( + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + ), + qk_norm_for_text=True, + ), + ), + tokenizer=L(create_qwen2_tokenizer_with_download)( + pretrained_model_name="Qwen/Qwen3-VL-8B-Instruct", + config_variant="s3_east2", + ), + pretrained_weights=PretrainedWeightsConfig( + backbone_path="s3://nv-cosmos-checkpoint-us-east-2/cosmos3/pretrained/huggingface/Cosmos-Reason/Cosmos3-Nano-Reasoner-bb9c6f5/", + credentials_path="credentials/s3_east2_checkpoint.secret", + ), +) + # Config for Qwen3VL 32B Instruct model # Qwen3VL uses Qwen2Tokenizer Qwen3VLMoT_VLM_32b_Instruct_Config: VLMConfig = VLMConfig( @@ -587,7 +610,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" ), qk_norm_for_text=True, ), @@ -607,7 +630,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" ), qk_norm_for_text=True, ), @@ -628,7 +651,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" ), qk_norm_for_text=True, ), @@ -649,7 +672,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" ), qk_norm_for_text=True, ), @@ -670,7 +693,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" + json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" ), qk_norm_for_text=True, ), @@ -697,7 +720,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -721,7 +744,31 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + ), + qk_norm_for_text=False, + ), + ), + tokenizer=L(build_processor_lazy)( + tokenizer_type="nvidia/Cosmos3-Edge-Reasoner", + config_variant="gcp", + ), + pretrained_weights=PretrainedWeightsConfig( + backbone_path="s3://bucket0/cosmos3/pretrained/huggingface/nvidia/Cosmos3-Edge-Reasoner-9b4c028/", + credentials_path="credentials/gcp_checkpoint.secret", + enable_gcs_patch_in_boto3=True, + checkpoint_format="nemotron_3_dense_vl", + ), +) + +# Cosmos3-Edge-Reasoner at commit 590c1c0 (2026-06-28). +# Updated weights uploaded 2026-06-28. +Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0: VLMConfig = VLMConfig( + model_name="nvidia/Cosmos3-Edge-Reasoner", + model_instance=L(Nemotron3DenseVLTextForCausalLM)( + config=L(create_vlm_config)( + base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( + json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -730,6 +777,31 @@ class VLMConfig: tokenizer_type="nvidia/Cosmos3-Edge-Reasoner", config_variant="gcp", ), + pretrained_weights=PretrainedWeightsConfig( + backbone_path="s3://bucket0/cosmos3/pretrained/huggingface/nvidia/Cosmos3-Edge-Reasoner-590c1c0/", + credentials_path="credentials/gcp_checkpoint.secret", + enable_gcs_patch_in_boto3=True, + checkpoint_format="nemotron_3_dense_vl", + ), +) + +# Same as 9b4c028 but with use_und_k_norm_for_gen=True: normalises K_und before +# it is used as a key in the gen→und cross-attention path. +Cosmos3EdgeReasoner_VLM_GCP_Config_9b4c028_UndKNorm: VLMConfig = VLMConfig( + model_name="nvidia/Cosmos3-Edge-Reasoner", + model_instance=L(Nemotron3DenseVLTextForCausalLM)( + config=L(create_vlm_config)( + base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( + json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + ), + qk_norm_for_text=False, + use_und_k_norm_for_gen=True, + ), + ), + tokenizer=L(build_processor_lazy)( + tokenizer_type="nvidia/Cosmos3-Edge-Reasoner", + config_variant="gcp", + ), pretrained_weights=PretrainedWeightsConfig( backbone_path="s3://bucket0/cosmos3/pretrained/huggingface/nvidia/Cosmos3-Edge-Reasoner-9b4c028/", credentials_path="credentials/gcp_checkpoint.secret", @@ -831,6 +903,12 @@ def register_vlm(): name="cosmos3_nano_reasoner_vlm_gcp_0517", node=Cosmos3NanoReasoner_VLM_GCP_Config_0517, ) + cs.store( + group="vlm_config", + package="model.config.vlm_config", + name="cosmos3_nano_reasoner_vlm_s3_east2_0517", + node=Cosmos3NanoReasoner_VLM_S3_EAST2_Config_0517, + ) cs.store( group="vlm_config", package="model.config.vlm_config", @@ -897,3 +975,15 @@ def register_vlm(): name="cosmos3_edge_reasoner_vlm_gcp_9b4c028", node=Cosmos3EdgeReasoner_VLM_GCP_Config_9b4c028, ) + cs.store( + group="vlm_config", + package="model.config.vlm_config", + name="cosmos3_edge_reasoner_vlm_gcp_9b4c028_und_k_norm", + node=Cosmos3EdgeReasoner_VLM_GCP_Config_9b4c028_UndKNorm, + ) + cs.store( + group="vlm_config", + package="model.config.vlm_config", + name="cosmos3_edge_reasoner_vlm_gcp_590c1c0", + node=Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0, + ) diff --git a/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py b/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py index 7340737e..cd7b6a3f 100644 --- a/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py +++ b/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py @@ -8,7 +8,7 @@ paths, video-style loss scales, ``load_weights_from_pretrained=True``). """ -from cosmos_framework.configs.base.defaults.vlm import ( +from cosmos_framework.configs.base.defaults.reasoner import ( create_qwen2_tokenizer_with_download, create_vlm_config, ) @@ -40,10 +40,6 @@ load_weights_from_pretrained=True, max_vae_latent_side_after_patchify=20, patch_spatial=2, - position_embedding_type="unified_3d_mrope", - rope_h_extrapolation_ratio=1.0, - rope_t_extrapolation_ratio=1.0, - rope_w_extrapolation_ratio=1.0, timestep_range=1.0, unified_3d_mrope_reset_spatial_ids=True, unified_3d_mrope_temporal_modality_margin=15000, @@ -128,7 +124,7 @@ config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( json_file=( - "cosmos_framework/model/vfm/vlm/qwen3_vl/configs/" + "cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/" "Qwen3-VL-8B-Instruct.json" ), ), diff --git a/cosmos_framework/configs/base/experiment/sft/models/super_model_config.py b/cosmos_framework/configs/base/experiment/sft/models/super_model_config.py index 8c7278ae..a10cee1a 100644 --- a/cosmos_framework/configs/base/experiment/sft/models/super_model_config.py +++ b/cosmos_framework/configs/base/experiment/sft/models/super_model_config.py @@ -26,7 +26,7 @@ from cosmos_framework.utils.lazy_config import LazyCall as L -from cosmos_framework.configs.base.defaults.vlm import ( +from cosmos_framework.configs.base.defaults.reasoner import ( create_qwen2_tokenizer_with_download, create_vlm_config, ) @@ -62,10 +62,6 @@ load_weights_from_pretrained=True, max_vae_latent_side_after_patchify=20, patch_spatial=2, - position_embedding_type="unified_3d_mrope", - rope_h_extrapolation_ratio=1.0, - rope_t_extrapolation_ratio=1.0, - rope_w_extrapolation_ratio=1.0, timestep_range=1.0, unified_3d_mrope_reset_spatial_ids=True, unified_3d_mrope_temporal_modality_margin=15000, @@ -150,7 +146,7 @@ config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( json_file=( - "cosmos_framework/model/vfm/vlm/qwen3_vl/configs/" + "cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/" "Qwen3-VL-32B-Instruct.json" ), ), diff --git a/cosmos_framework/configs/base/vlm/__init__.py b/cosmos_framework/configs/base/reasoner/__init__.py similarity index 100% rename from cosmos_framework/configs/base/vlm/__init__.py rename to cosmos_framework/configs/base/reasoner/__init__.py diff --git a/cosmos_framework/configs/base/vlm/config.py b/cosmos_framework/configs/base/reasoner/config.py similarity index 77% rename from cosmos_framework/configs/base/vlm/config.py rename to cosmos_framework/configs/base/reasoner/config.py index 4d1de856..1a1455d6 100644 --- a/cosmos_framework/configs/base/vlm/config.py +++ b/cosmos_framework/configs/base/reasoner/config.py @@ -5,11 +5,11 @@ from cosmos_framework.utils import log from cosmos_framework.utils.config_helper import import_all_modules_from_package from cosmos_framework.configs.base.defaults.checkpointer import register_checkpoint, register_ckpt_type -from cosmos_framework.configs.base.vlm.defaults.callbacks import register_callbacks -from cosmos_framework.configs.base.vlm.defaults.config import Config -from cosmos_framework.configs.base.vlm.defaults.model import register_model -from cosmos_framework.configs.base.vlm.defaults.optimizer import register_optimizer, register_scheduler -from cosmos_framework.configs.base.vlm.defaults.vlm_policy import register_vlm_policy +from cosmos_framework.configs.base.reasoner.defaults.callbacks import register_callbacks +from cosmos_framework.configs.base.reasoner.defaults.config import Config +from cosmos_framework.configs.base.reasoner.defaults.model import register_model +from cosmos_framework.configs.base.reasoner.defaults.optimizer import register_optimizer, register_scheduler +from cosmos_framework.configs.base.reasoner.defaults.vlm_policy import register_vlm_policy def make_config() -> Config: @@ -47,5 +47,5 @@ def make_config() -> Config: register_checkpoint() register_ckpt_type() register_callbacks() - import_all_modules_from_package("cosmos_framework.configs.base.vlm.experiment", reload=True) + import_all_modules_from_package("cosmos_framework.configs.base.reasoner.experiment", reload=True) return c diff --git a/cosmos_framework/configs/base/vlm/defaults/__init__.py b/cosmos_framework/configs/base/reasoner/defaults/__init__.py similarity index 100% rename from cosmos_framework/configs/base/vlm/defaults/__init__.py rename to cosmos_framework/configs/base/reasoner/defaults/__init__.py diff --git a/cosmos_framework/configs/base/vlm/defaults/callbacks.py b/cosmos_framework/configs/base/reasoner/defaults/callbacks.py similarity index 93% rename from cosmos_framework/configs/base/vlm/defaults/callbacks.py rename to cosmos_framework/configs/base/reasoner/defaults/callbacks.py index bb1cd4b3..7d1911f5 100644 --- a/cosmos_framework/configs/base/vlm/defaults/callbacks.py +++ b/cosmos_framework/configs/base/reasoner/defaults/callbacks.py @@ -19,6 +19,7 @@ from cosmos_framework.callbacks.learning_rate_logger import LearningRateLogger from cosmos_framework.callbacks.log_tensor_shape import LogTensorShapeCallback from cosmos_framework.callbacks.param_count import ParamCount +from cosmos_framework.callbacks.tokens_per_sec import VLMTokensPerSec from cosmos_framework.callbacks.wandb_log import WandbCallback as WandBCallbackMultiplier from cosmos_framework.callbacks.wandb_vis import VisualizationLoggingCallback from cosmos_framework.configs.base.defaults.callbacks import JOB_MONITOR_CALLBACKS @@ -35,6 +36,10 @@ def register_callbacks(): save_s3_every_log_n=500, hit_thres=50, ), + vlm_tokens_per_sec=L(VLMTokensPerSec)( # per-GPU tokens/sec + packing efficiency + every_n="${trainer.logging_iter}", + hit_thres=50, + ), manual_gc=L(ManualGarbageCollection)(every_n=5), # does not use model or optimizer wandb=L(WandBCallback)(), param_count=L(ParamCount)( # use model diff --git a/cosmos_framework/configs/base/vlm/defaults/checkpointer.py b/cosmos_framework/configs/base/reasoner/defaults/checkpointer.py similarity index 100% rename from cosmos_framework/configs/base/vlm/defaults/checkpointer.py rename to cosmos_framework/configs/base/reasoner/defaults/checkpointer.py diff --git a/cosmos_framework/configs/base/vlm/defaults/config.py b/cosmos_framework/configs/base/reasoner/defaults/config.py similarity index 83% rename from cosmos_framework/configs/base/vlm/defaults/config.py rename to cosmos_framework/configs/base/reasoner/defaults/config.py index 21280721..6da4659a 100644 --- a/cosmos_framework/configs/base/vlm/defaults/config.py +++ b/cosmos_framework/configs/base/reasoner/defaults/config.py @@ -6,7 +6,7 @@ import attrs from cosmos_framework.utils import config -from cosmos_framework.configs.base.vlm.defaults.policy_config import PolicyConfig +from cosmos_framework.configs.base.reasoner.defaults.policy_config import PolicyConfig @attrs.define(slots=False) @@ -19,6 +19,10 @@ class DataSetting: text_chat_order: Order of text items in user messages. distributor_type: "with_replace" (WeightedShardlistBasic) or "no_replace" (NoReplaceShardlistBasic). distributor_seed: Seed for the distributor. + max_batch_size: Hard cap on the number of samples in each dynamic batch. + max_tokens: Per-sample token limit used by filtering. + max_tokens_in_batch: Padded-token budget for dynamic batching. + long_threshold: Single seed-sample length threshold that emits the sample alone. """ qwen_max_video_token_length: int = 8192 @@ -36,6 +40,8 @@ class DataSetting: # For packed dataset max_batch_size: int = 1 max_tokens: int = 16000 + max_tokens_in_batch: int = 16000 + long_threshold: int = 6400 # "with_replace" (WeightedShardlistBasic) or "no_replace" (NoReplaceShardlistBasic). distributor_type: str = attrs.field( default="with_replace", diff --git a/cosmos_framework/configs/base/vlm/defaults/model.py b/cosmos_framework/configs/base/reasoner/defaults/model.py similarity index 90% rename from cosmos_framework/configs/base/vlm/defaults/model.py rename to cosmos_framework/configs/base/reasoner/defaults/model.py index e618d11c..0c991f08 100644 --- a/cosmos_framework/configs/base/vlm/defaults/model.py +++ b/cosmos_framework/configs/base/reasoner/defaults/model.py @@ -5,7 +5,7 @@ from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig -from cosmos_framework.configs.base.vlm.defaults.policy_config import VLMModelConfig +from cosmos_framework.configs.base.reasoner.defaults.policy_config import VLMModelConfig from cosmos_framework.model.vfm.vlm_model import VLMModel VLM_FSDP_CONFIG = dict( diff --git a/cosmos_framework/configs/base/vlm/defaults/optimizer.py b/cosmos_framework/configs/base/reasoner/defaults/optimizer.py similarity index 98% rename from cosmos_framework/configs/base/vlm/defaults/optimizer.py rename to cosmos_framework/configs/base/reasoner/defaults/optimizer.py index 538ab7d4..dbd5829c 100644 --- a/cosmos_framework/configs/base/vlm/defaults/optimizer.py +++ b/cosmos_framework/configs/base/reasoner/defaults/optimizer.py @@ -30,7 +30,7 @@ VLM_LAMBDACOSINE_KWARGS: dict[str, Any] = dict( warm_up_steps=[1000], cycle_lengths=["${trainer.max_iter}"], - f_start=[0.05], + f_start=[0.01], f_max=[1.0], f_min=[0.5], ) diff --git a/cosmos_framework/configs/base/vlm/defaults/policy_config.py b/cosmos_framework/configs/base/reasoner/defaults/policy_config.py similarity index 79% rename from cosmos_framework/configs/base/vlm/defaults/policy_config.py rename to cosmos_framework/configs/base/reasoner/defaults/policy_config.py index 7df99877..1cabfce6 100644 --- a/cosmos_framework/configs/base/vlm/defaults/policy_config.py +++ b/cosmos_framework/configs/base/reasoner/defaults/policy_config.py @@ -9,8 +9,8 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.ema import EMAConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig -from cosmos_framework.configs.base.defaults.vlm import VLMConfig -from cosmos_framework.configs.base.vlm.freeze_config import VLMFreezeConfig +from cosmos_framework.configs.base.defaults.reasoner import VLMConfig +from cosmos_framework.configs.base.reasoner.freeze_config import VLMFreezeConfig @attrs.define(slots=False) @@ -23,6 +23,13 @@ class PolicyConfig: # The maximum length for video tokens, only applied to qwen model qwen_max_video_token_length: int = 8000 + use_weighted_ce: bool = False + # Controls the interpolation between per-token (0) and per-sample (1) loss: + # exponent=1 -> per-sample loss: every sample contributes equally to the global loss + # exponent=0 -> per-token loss: every token contributes equally to the global loss + # 0 < exponent < 1 -> interpolation; e.g. exponent=0.5 gives square-root per-token loss (Qwen3-VL) + weighted_ce_exponent: float = 1.0 + # Extra model config lora: Union[str, None] = None enable_liger_kernel: bool = False diff --git a/cosmos_framework/configs/base/vlm/defaults/vlm_policy.py b/cosmos_framework/configs/base/reasoner/defaults/vlm_policy.py similarity index 95% rename from cosmos_framework/configs/base/vlm/defaults/vlm_policy.py rename to cosmos_framework/configs/base/reasoner/defaults/vlm_policy.py index 9fe54a14..5c119ab6 100644 --- a/cosmos_framework/configs/base/vlm/defaults/vlm_policy.py +++ b/cosmos_framework/configs/base/reasoner/defaults/vlm_policy.py @@ -3,12 +3,12 @@ from hydra.core.config_store import ConfigStore -from cosmos_framework.configs.base.defaults.vlm import VLMConfig -from cosmos_framework.configs.base.vlm.defaults.policy_config import PolicyConfig +from cosmos_framework.configs.base.defaults.reasoner import VLMConfig +from cosmos_framework.configs.base.reasoner.defaults.policy_config import PolicyConfig # Each entry replaces cfg.model.config.policy via package="model.config.policy". # Sibling to the VFM vlm_config group at -# cosmos_framework/configs/base/defaults/vlm.py: that group binds VLMConfig +# cosmos_framework/configs/base/defaults/reasoner.py: that group binds VLMConfig # SKUs onto OmniMoTModelConfig.vlm_config; this group binds PolicyConfig SKUs # onto VLMModelConfig.policy. Both groups compose the same VLMConfig: VFM as # vlm_config: VLMConfig, VLM as policy.backbone: VLMConfig. diff --git a/cosmos_framework/configs/base/vlm/experiment/__init__.py b/cosmos_framework/configs/base/reasoner/experiment/__init__.py similarity index 100% rename from cosmos_framework/configs/base/vlm/experiment/__init__.py rename to cosmos_framework/configs/base/reasoner/experiment/__init__.py diff --git a/cosmos_framework/configs/base/vlm/experiment/dataflow_roles.py b/cosmos_framework/configs/base/reasoner/experiment/dataflow_roles.py similarity index 100% rename from cosmos_framework/configs/base/vlm/experiment/dataflow_roles.py rename to cosmos_framework/configs/base/reasoner/experiment/dataflow_roles.py diff --git a/cosmos_framework/configs/base/vlm/experiment/llava_ov_vlm.py b/cosmos_framework/configs/base/reasoner/experiment/llava_ov_vlm.py similarity index 99% rename from cosmos_framework/configs/base/vlm/experiment/llava_ov_vlm.py rename to cosmos_framework/configs/base/reasoner/experiment/llava_ov_vlm.py index 2d80828c..1510ea1b 100644 --- a/cosmos_framework/configs/base/vlm/experiment/llava_ov_vlm.py +++ b/cosmos_framework/configs/base/reasoner/experiment/llava_ov_vlm.py @@ -61,7 +61,7 @@ ) from cosmos_framework.data.vfm.processors import build_processor from cosmos_framework.utils.vlm.constant import IGNORE_INDEX -from cosmos_framework.configs.base.vlm.experiment.dataflow_roles import VLMProcessor, VLMCollator +from cosmos_framework.configs.base.reasoner.experiment.dataflow_roles import VLMProcessor, VLMCollator from cosmos_framework.callbacks.cosmos_dataloader_state import CosmosDataLoaderStateCallback cs = ConfigStore.instance() diff --git a/cosmos_framework/configs/base/vlm/experiment/utils.py b/cosmos_framework/configs/base/reasoner/experiment/utils.py similarity index 100% rename from cosmos_framework/configs/base/vlm/experiment/utils.py rename to cosmos_framework/configs/base/reasoner/experiment/utils.py diff --git a/cosmos_framework/configs/base/vlm/experiment/videophy2_dataflow_roles.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_dataflow_roles.py similarity index 100% rename from cosmos_framework/configs/base/vlm/experiment/videophy2_dataflow_roles.py rename to cosmos_framework/configs/base/reasoner/experiment/videophy2_dataflow_roles.py diff --git a/cosmos_framework/configs/base/vlm/experiment/videophy2_sft_nano.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py similarity index 97% rename from cosmos_framework/configs/base/vlm/experiment/videophy2_sft_nano.py rename to cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py index 358b2c69..baa438af 100644 --- a/cosmos_framework/configs/base/vlm/experiment/videophy2_sft_nano.py +++ b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py @@ -25,8 +25,8 @@ from cosmos_framework.data.vlm.data_sources_videophy2.videophy2 import DATAINFO from cosmos_framework.utils import log from cosmos_framework.utils.vlm.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD -from cosmos_framework.configs.base.vlm.experiment.dataflow_roles import VLMCollator -from cosmos_framework.configs.base.vlm.experiment.videophy2_dataflow_roles import VideoPhy2Processor +from cosmos_framework.configs.base.reasoner.experiment.dataflow_roles import VLMCollator +from cosmos_framework.configs.base.reasoner.experiment.videophy2_dataflow_roles import VideoPhy2Processor cs = ConfigStore.instance() diff --git a/cosmos_framework/configs/base/vlm/freeze_config.py b/cosmos_framework/configs/base/reasoner/freeze_config.py similarity index 100% rename from cosmos_framework/configs/base/vlm/freeze_config.py rename to cosmos_framework/configs/base/reasoner/freeze_config.py diff --git a/cosmos_framework/configs/toml_config/toml_config_helper.py b/cosmos_framework/configs/toml_config/toml_config_helper.py index 82333cf3..211e9d3b 100644 --- a/cosmos_framework/configs/toml_config/toml_config_helper.py +++ b/cosmos_framework/configs/toml_config/toml_config_helper.py @@ -22,7 +22,7 @@ # Maps ``job.task`` to the base Hydra config that ``make_config()`` lives in. TASK_TO_BASE_CONFIG: dict[str, str] = { "vfm": "cosmos_framework/configs/base/config.py", - "vlm": "cosmos_framework/configs/base/vlm/config.py", + "vlm": "cosmos_framework/configs/base/reasoner/config.py", } diff --git a/cosmos_framework/data/vfm/augmentors/vlm/__init__.py b/cosmos_framework/data/vfm/augmentors/reasoner/__init__.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/vlm/__init__.py rename to cosmos_framework/data/vfm/augmentors/reasoner/__init__.py diff --git a/cosmos_framework/data/vfm/augmentors/vlm/bytes_to_media.py b/cosmos_framework/data/vfm/augmentors/reasoner/bytes_to_media.py similarity index 62% rename from cosmos_framework/data/vfm/augmentors/vlm/bytes_to_media.py rename to cosmos_framework/data/vfm/augmentors/reasoner/bytes_to_media.py index 3f61341a..db670b30 100644 --- a/cosmos_framework/data/vfm/augmentors/vlm/bytes_to_media.py +++ b/cosmos_framework/data/vfm/augmentors/reasoner/bytes_to_media.py @@ -16,7 +16,7 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.vlm.video_decoder_qwen import _video_decoder_qwen_func +from cosmos_framework.data.vfm.reasoner.video_decoder_qwen import _video_decoder_qwen_func from cosmos_framework.data.vfm.processors.qwen3vl_processor import Qwen3VLProcessor from cosmos_framework.utils.vfm.video_preprocess import tensor_to_pil_images @@ -81,8 +81,101 @@ def __init__( self.use_start_frame_end_frame = use_start_frame_end_frame self.processor = processor + def _is_video_key(self, name: str) -> bool: + """Returns whether the media key will be decoded as video.""" + name_lower = name.lower() + if self.use_start_frame_end_frame: + return "video" in name_lower or ".mp4" in name_lower + control_video_names = ( + "control_input_world_scenario", + "control_input_seg", + "control_input_blur", + "control_input_depth", + "control_input_edge", + ) + return "video" in name_lower or ".mp4" in name_lower or name_lower in control_video_names + + def _probe_video_duration_seconds( + self, + video_bytes: bytes, + identifier: str, + start_frame: int | None = None, + end_frame: int | None = None, + ) -> float | None: + """Probe the effective video duration in seconds used for proportional budget allocation.""" + try: + import decord + + with io.BytesIO(video_bytes) as stream: + video_reader = decord.VideoReader(stream, num_threads=self.video_decoder_params["num_threads"]) + frame_count = ( + max(end_frame - start_frame, 0) + if start_frame is not None and end_frame is not None + else len(video_reader) + ) + video_fps = video_reader.get_avg_fps() + video_reader.seek(0) + del video_reader + if frame_count <= 0 or video_fps <= 0: + return None + return frame_count / video_fps + except Exception as e: + log.warning(f"Could not probe video duration for '{identifier}': {e}") + return None + + def _get_video_durations(self, data: dict, data_dict: Dict) -> dict[str, float]: + """Return per-video durations in seconds, or an empty dict if any probe fails.""" + start_frame = data_dict.get("start_frame") + end_frame = data_dict.get("end_frame") + video_durations: dict[str, float] = {} + video_names = [name for name, item in data.items() if isinstance(item, bytes) and self._is_video_key(name)] + + for name in video_names: + duration_seconds = self._probe_video_duration_seconds( + data[name], + identifier=f"{self.input_key}['{name}']", + start_frame=start_frame, + end_frame=end_frame, + ) + if duration_seconds is None or duration_seconds <= 0: + return {} + video_durations[name] = duration_seconds + + return video_durations + + def _get_decoder_params( + self, + video_count: int, + video_duration: float | None = None, + total_video_duration: float | None = None, + ) -> dict: + """Scale the decoder pixel budget by video count, weighted by video duration when available.""" + decoder_params = dict(self.video_decoder_params) + if video_count <= 1: + return decoder_params + + max_video_token_length = decoder_params["max_video_token_length"] + min_video_token_length = decoder_params["min_video_token_length"] + if video_duration is not None and total_video_duration: + scaled_video_token_length = round(max_video_token_length * video_duration / total_video_duration) + else: + scaled_video_token_length = max_video_token_length // video_count + + decoder_params["max_video_token_length"] = max( + min_video_token_length, + scaled_video_token_length, + ) + return decoder_params + def _bytes_to_video_frames( - self, video_bytes: bytes, identifier: str = "video", start_frame: int = None, end_frame: int = None + self, + video_bytes: bytes, + identifier: str = "video", + start_frame: int | None = None, + end_frame: int | None = None, + video_count: int = 1, + video_duration: float | None = None, + total_video_duration: float | None = None, ) -> Optional[Dict]: """Converts video bytes to video frame tensors using the video decoder.""" try: @@ -92,14 +185,17 @@ def _bytes_to_video_frames( processor=self.processor, start_frame=start_frame, end_frame=end_frame, - **self.video_decoder_params, + **self._get_decoder_params( + video_count, + video_duration=video_duration, + total_video_duration=total_video_duration, + ), ) - result["videos"] = tensor_to_pil_images(result["videos"]) # 3,T,H,W -> list of PIL images - if result is not None: - return result - else: + if result is None: log.warning(f"Skipping item '{identifier}': Video decoder returned None.") return None + result["videos"] = tensor_to_pil_images(result["videos"]) # 3,T,H,W -> list of PIL images + return result except Exception as e: log.warning(f"Skipping item '{identifier}': Error decoding video bytes: {e}") return None @@ -163,22 +259,37 @@ def __call__(self, data_dict: Dict) -> Dict: output_data = {} if isinstance(data, dict): + video_count = sum(1 for name, item in data.items() if isinstance(item, bytes) and self._is_video_key(name)) + video_durations = self._get_video_durations(data, data_dict) if video_count > 1 else {} + total_video_duration = sum(video_durations.values()) if video_durations else None for name, item in data.items(): if isinstance(item, bytes): # Determine if this is video or image based on the key name - if ("video" in name.lower() or ".mp4" in name.lower()) and not self.use_start_frame_end_frame: + if self._is_video_key(name) and not self.use_start_frame_end_frame: # Decode as video - result = self._bytes_to_video_frames(item, identifier=f"{input_key}['{name}']") + result = self._bytes_to_video_frames( + item, + identifier=f"{input_key}['{name}']", + video_count=video_count, + video_duration=video_durations.get(name), + total_video_duration=total_video_duration, + ) if result: output_data[name] = result - elif ("video" in name.lower() or ".mp4" in name.lower()) and self.use_start_frame_end_frame: + elif self._is_video_key(name) and self.use_start_frame_end_frame: assert "start_frame" in data_dict.keys() and "end_frame" in data_dict.keys(), ( f"start_frame and end_frame are not in data_dict.keys(): {data_dict.keys()}" ) start_frame = data_dict["start_frame"] end_frame = data_dict["end_frame"] result = self._bytes_to_video_frames( - item, identifier=f"{input_key}['{name}']", start_frame=start_frame, end_frame=end_frame + item, + identifier=f"{input_key}['{name}']", + start_frame=start_frame, + end_frame=end_frame, + video_count=video_count, + video_duration=video_durations.get(name), + total_video_duration=total_video_duration, ) if result: output_data[name] = result diff --git a/cosmos_framework/data/vfm/augmentors/vlm/filter_output_key.py b/cosmos_framework/data/vfm/augmentors/reasoner/filter_output_key.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/vlm/filter_output_key.py rename to cosmos_framework/data/vfm/augmentors/reasoner/filter_output_key.py diff --git a/cosmos_framework/data/vfm/augmentors/vlm/filter_seq_length.py b/cosmos_framework/data/vfm/augmentors/reasoner/filter_seq_length.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/vlm/filter_seq_length.py rename to cosmos_framework/data/vfm/augmentors/reasoner/filter_seq_length.py diff --git a/cosmos_framework/data/vfm/augmentors/vlm/floating_number_format.py b/cosmos_framework/data/vfm/augmentors/reasoner/floating_number_format.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/vlm/floating_number_format.py rename to cosmos_framework/data/vfm/augmentors/reasoner/floating_number_format.py diff --git a/cosmos_framework/data/vfm/augmentors/vlm/format_describe_anything.py b/cosmos_framework/data/vfm/augmentors/reasoner/format_describe_anything.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/vlm/format_describe_anything.py rename to cosmos_framework/data/vfm/augmentors/reasoner/format_describe_anything.py index 8ca392ec..45028366 100644 --- a/cosmos_framework/data/vfm/augmentors/vlm/format_describe_anything.py +++ b/cosmos_framework/data/vfm/augmentors/reasoner/format_describe_anything.py @@ -13,7 +13,7 @@ from typing import Dict, List, Literal from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor -from cosmos_framework.data.vfm.augmentors.vlm.timestamp import markdown_to_list +from cosmos_framework.data.vfm.augmentors.reasoner.timestamp import markdown_to_list # reorder dict entries diff --git a/cosmos_framework/data/vfm/augmentors/vlm/nvlm_data_to_conversation.py b/cosmos_framework/data/vfm/augmentors/reasoner/nvlm_data_to_conversation.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/vlm/nvlm_data_to_conversation.py rename to cosmos_framework/data/vfm/augmentors/reasoner/nvlm_data_to_conversation.py diff --git a/cosmos_framework/data/vfm/augmentors/vlm/prompt_format.py b/cosmos_framework/data/vfm/augmentors/reasoner/prompt_format.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/vlm/prompt_format.py rename to cosmos_framework/data/vfm/augmentors/reasoner/prompt_format.py diff --git a/cosmos_framework/data/vfm/augmentors/vlm/shuffle_text_media_order.py b/cosmos_framework/data/vfm/augmentors/reasoner/shuffle_text_media_order.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/vlm/shuffle_text_media_order.py rename to cosmos_framework/data/vfm/augmentors/reasoner/shuffle_text_media_order.py diff --git a/cosmos_framework/data/vfm/augmentors/vlm/timestamp.py b/cosmos_framework/data/vfm/augmentors/reasoner/timestamp.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/vlm/timestamp.py rename to cosmos_framework/data/vfm/augmentors/reasoner/timestamp.py index 88d3ac5f..6228153e 100644 --- a/cosmos_framework/data/vfm/augmentors/vlm/timestamp.py +++ b/cosmos_framework/data/vfm/augmentors/reasoner/timestamp.py @@ -97,7 +97,7 @@ def overlay_text( return images, [compute_timestamps(i, fps, processor) for i in range(len(images))] # Try to use DejaVu Sans Mono font for better readability - font = ImageFont.truetype("/invalid_dir", font_size) + font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", font_size) # Process each image processed_images = [] diff --git a/cosmos_framework/data/vfm/augmentors/vlm/timestamp_with_subject_tracking.py b/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_with_subject_tracking.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/vlm/timestamp_with_subject_tracking.py rename to cosmos_framework/data/vfm/augmentors/reasoner/timestamp_with_subject_tracking.py index 90cc9ab6..34851186 100644 --- a/cosmos_framework/data/vfm/augmentors/vlm/timestamp_with_subject_tracking.py +++ b/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_with_subject_tracking.py @@ -16,7 +16,7 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.augmentors.vlm.timestamp import ( +from cosmos_framework.data.vfm.augmentors.reasoner.timestamp import ( json_to_markdown, markdown_to_list, overlay_text, diff --git a/cosmos_framework/data/vfm/augmentors/vlm/timestamp_without_augment_message.py b/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_augment_message.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/vlm/timestamp_without_augment_message.py rename to cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_augment_message.py index 7212510a..8ca4984b 100644 --- a/cosmos_framework/data/vfm/augmentors/vlm/timestamp_without_augment_message.py +++ b/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_augment_message.py @@ -13,7 +13,7 @@ from typing import Dict, List, Literal, Tuple from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor -from cosmos_framework.data.vfm.augmentors.vlm.timestamp import overlay_text +from cosmos_framework.data.vfm.augmentors.reasoner.timestamp import overlay_text def list_to_markdown(conversation_data: List[Dict]) -> str: diff --git a/cosmos_framework/data/vfm/augmentors/vlm/timestamp_without_end_time.py b/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_end_time.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/vlm/timestamp_without_end_time.py rename to cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_end_time.py index 812e1137..0c9fb05d 100644 --- a/cosmos_framework/data/vfm/augmentors/vlm/timestamp_without_end_time.py +++ b/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_end_time.py @@ -16,7 +16,7 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.augmentors.vlm.timestamp import ( +from cosmos_framework.data.vfm.augmentors.reasoner.timestamp import ( json_to_markdown, markdown_to_list, overlay_text, diff --git a/cosmos_framework/data/vfm/augmentors/vlm/tokenize_data.py b/cosmos_framework/data/vfm/augmentors/reasoner/tokenize_data.py similarity index 88% rename from cosmos_framework/data/vfm/augmentors/vlm/tokenize_data.py rename to cosmos_framework/data/vfm/augmentors/reasoner/tokenize_data.py index f3a9914c..d3446726 100644 --- a/cosmos_framework/data/vfm/augmentors/vlm/tokenize_data.py +++ b/cosmos_framework/data/vfm/augmentors/reasoner/tokenize_data.py @@ -12,9 +12,9 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.vlm.video_decoder_qwen import token_to_pixels +from cosmos_framework.data.vfm.reasoner.video_decoder_qwen import token_to_pixels from cosmos_framework.data.vfm.processors.qwen3vl_processor import Qwen3VLProcessor as Processor -from cosmos_framework.utils.vfm.vlm.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD +from cosmos_framework.utils.vfm.reasoner.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD def maybe_subsample_frames(model_name_or_path, list_of_pil_image, max_video_token_length, processor): @@ -148,7 +148,8 @@ def __call__(self, data_dict: Dict) -> Dict: processor_kwargs = {} total_images = 0 total_videos = 0 - raw_images = [] + raw_images: list[torch.Tensor] = [] + raw_videos: list[torch.Tensor] = [] # Pre-compute the total_images and total_videos for message in conversation: if not isinstance(message, dict): @@ -158,7 +159,6 @@ def __call__(self, data_dict: Dict) -> Dict: if message["role"] == "user" and isinstance(message["content"], list): total_images += len([content for content in message["content"] if content["type"] == "image"]) total_videos += len([content for content in message["content"] if content["type"] == "video"]) - assert total_videos == 1 or total_videos == 0, "Only one video is supported for now" # url url = data_dict["__url__"].root + "/" + data_dict["__url__"].path @@ -220,7 +220,8 @@ def __call__(self, data_dict: Dict) -> Dict: image = data_dict["media"][content["image"]] content["image"] = image content["max_pixels"] = max_pixels_per_image - raw_images.append(image) + raw_image = np.asarray(image.convert("RGB")) # [H,W,3] + raw_images.append(torch.from_numpy(raw_image).permute(2, 0, 1)[:, None]) # [3,1,H,W] elif content["type"] == "video": # as tokenization will NOT upsample the video, we can use a larger value here at the cost of multiple video having 1.5x token length @@ -250,25 +251,27 @@ def __call__(self, data_dict: Dict) -> Dict: videos = maybe_subsample_frames( self.processor.name, videos, self.max_video_token_length, self.processor ) + if len(videos) == 0: + log.info(f"[TokenizerDataError]video {media_key} has no decoded frames. url: {url}") + return None content["video"] = videos max_pixels_per_image = max_total_pixels // total_videos // len(videos) content["fps"] = fps content["max_pixels"] = max_pixels_per_image - data_dict["raw_video"] = torch.from_numpy(np.array(videos)).permute( - 3, 0, 1, 2 - ) # [3,T,H,W], range [0, 255] + raw_video_frames = np.stack( + [np.asarray(frame.convert("RGB")) for frame in videos], axis=0 + ) # [T,H,W,3] + raw_videos.append(torch.from_numpy(raw_video_frames).permute(3, 0, 1, 2)) # [3,T,H,W] new_content_list.append(content) message["content"] = new_content_list - if len(raw_images) > 1: - # resize the raw_image to the size of the first image - image_size = raw_images[0].size - raw_images = [image.resize(image_size) for image in raw_images] - if len(raw_images) > 0: - data_dict["raw_image"] = torch.from_numpy(np.array(raw_images)).permute(3, 0, 1, 2) # [3,num_images,H,W] + data_dict["raw_image"] = raw_images # each: [3,1,H,W] + + if len(raw_videos) > 0: + data_dict["raw_video"] = raw_videos # each: [3,T,H,W] if conversation[0]["role"] != "system" and self.add_system_prompt_if_missing: conversation.insert(0, {"role": "system", "content": "You are a helpful assistant."}) @@ -307,12 +310,15 @@ def __call__(self, data_dict: Dict) -> Dict: return None input_ids = tokenizer_output["input_ids"] if "image_grid_thw" in tokenizer_output and "raw_image" in data_dict: - # image_grid_thw: (1, t, h, w) - t, h, w = tokenizer_output["image_grid_thw"][0] - # interpolate raw_image to the size of the image grid * 14 - data_dict["raw_image"] = torch.nn.functional.interpolate( - data_dict["raw_image"], size=(h * 14, w * 14), mode="bilinear", align_corners=False - ) # [3,num_images,h*14,w*14] + resized_raw_images: list[torch.Tensor] = [] + for raw_image, image_grid_thw in zip(data_dict["raw_image"], tokenizer_output["image_grid_thw"]): + # image_grid_thw: [t,h,w] + _, h, w = image_grid_thw + raw_image = torch.nn.functional.interpolate( + raw_image, size=(int(h) * 14, int(w) * 14), mode="bilinear", align_corners=False + ) # [3,1,h*14,w*14] + resized_raw_images.append(raw_image) + data_dict["raw_image"] = resized_raw_images # each: [3,1,h*14,w*14] try: # token_mask: True for tokens to compute loss on; False for tokens to ignore @@ -348,9 +354,17 @@ def __call__(self, data_dict: Dict) -> Dict: # For debugging purpose msg = f"input_ids: {input_ids.shape[-1]} | __url__: {data_dict['__url__'].root}, {data_dict['__url__'].path} | __key__: {data_dict['__key__']}" if "raw_video" in data_dict: - msg += f" | raw_video: {data_dict['raw_video'].shape} " + raw_video = data_dict["raw_video"] + if isinstance(raw_video, list): + msg += f" | raw_video: {[video.shape for video in raw_video]} " + else: + msg += f" | raw_video: {raw_video.shape} " if "raw_image" in data_dict: - msg += f" | raw_image: {data_dict['raw_image'].shape} " + raw_image = data_dict["raw_image"] + if isinstance(raw_image, list): + msg += f" | raw_image: {[image.shape for image in raw_image]} " + else: + msg += f" | raw_image: {raw_image.shape} " if "pixel_values" in data_dict: msg += f" | pixel_values: {data_dict['pixel_values'].shape} " diff --git a/cosmos_framework/data/vfm/augmentors/vlm/user_prompt_caption_general.json b/cosmos_framework/data/vfm/augmentors/reasoner/user_prompt_caption_general.json similarity index 100% rename from cosmos_framework/data/vfm/augmentors/vlm/user_prompt_caption_general.json rename to cosmos_framework/data/vfm/augmentors/reasoner/user_prompt_caption_general.json diff --git a/cosmos_framework/data/vfm/augmentors/vlm/user_prompt_ocr.json b/cosmos_framework/data/vfm/augmentors/reasoner/user_prompt_ocr.json similarity index 100% rename from cosmos_framework/data/vfm/augmentors/vlm/user_prompt_ocr.json rename to cosmos_framework/data/vfm/augmentors/reasoner/user_prompt_ocr.json diff --git a/cosmos_framework/data/vfm/augmentors/text_tokenizer.py b/cosmos_framework/data/vfm/augmentors/text_tokenizer.py index 2f8196a2..243c518b 100644 --- a/cosmos_framework/data/vfm/augmentors/text_tokenizer.py +++ b/cosmos_framework/data/vfm/augmentors/text_tokenizer.py @@ -79,6 +79,9 @@ def __init__(self, input_keys: list, output_keys: Optional[list] = None, args: O def __call__(self, data_dict: dict) -> dict | None: input_caption = data_dict.get(self.input_keys[0], "") + if isinstance(input_caption, dict): + input_caption = json.dumps(input_caption) + data_dict[self.input_keys[0]] = input_caption if self.cfg_dropout_rate > 0 and random.random() < self.cfg_dropout_rate: input_caption = "" data_dict[self.input_keys[0]] = input_caption diff --git a/cosmos_framework/data/vfm/augmentors/text_transforms_for_image.py b/cosmos_framework/data/vfm/augmentors/text_transforms_for_image.py index d38fae42..95462393 100644 --- a/cosmos_framework/data/vfm/augmentors/text_transforms_for_image.py +++ b/cosmos_framework/data/vfm/augmentors/text_transforms_for_image.py @@ -9,16 +9,13 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log +# COSMOS-RELEASE-END-IGNORE + # For the qwen captions, we have 3 variants: short, medium, long # In addition, for synthetic data, we create prompt embeddings as well. # There is quite a bit of entropy in the way prompt data is saved. # Captions are saved as "prompts", while the corresponding embeddings are saved as "original_prompt" # This part will be cleaned after synthetic data is cleaned to be in the same format as real data. -_CAPTION_EMBEDDING_KEY_MAPPING_IMAGES = { - "ai_v3p1": "ai_v3p1", - "qwen2p5_7b_v4": "qwen2p5_7b_v4", - "prompts": "qwen2p5_7b_v4", -} _AVAILABLE_QWEN_CAPTIONS = ["qwen2p5_7b_short", "qwen2p5_7b_medium", "qwen2p5_7b_long"] _AVAILABLE_QWEN3_30B_A3B_CAPTIONS = [ "qwen3_30b_a3b_short", diff --git a/cosmos_framework/data/vfm/joint_dataloader.py b/cosmos_framework/data/vfm/joint_dataloader.py index 6ad83eb7..fa27e528 100644 --- a/cosmos_framework/data/vfm/joint_dataloader.py +++ b/cosmos_framework/data/vfm/joint_dataloader.py @@ -923,7 +923,8 @@ def __iter__(self): current_sequence_length += num_tokens_in_current_sample num_samples += 1 - output["dataset_name"] = ds_name + # Allows the dataset name to be overridden by the sample itself. + output["dataset_name"] = output.get("dataset_name", ds_name) self._update_output_batch(output_batch, output) for sample in reversed(skipped_samples): diff --git a/cosmos_framework/data/vfm/local_datasets/sft_dataset.py b/cosmos_framework/data/vfm/local_datasets/sft_dataset.py index 38b965ed..0632c2a1 100644 --- a/cosmos_framework/data/vfm/local_datasets/sft_dataset.py +++ b/cosmos_framework/data/vfm/local_datasets/sft_dataset.py @@ -28,7 +28,7 @@ from cosmos_framework.data.vfm.sequence_packing.modalities import add_special_tokens from cosmos_framework.data.vfm.utils import VIDEO_RES_SIZE_INFO from cosmos_framework.inference.structured_caption import CAPTION_JSON_KEY, caption_json_to_prompt -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption from cosmos_framework.utils import log from cosmos_framework.utils.flags import INTERNAL from cosmos_framework.utils.lazy_config import instantiate as lazy_instantiate diff --git a/cosmos_framework/data/vfm/packing_iterable_dataset.py b/cosmos_framework/data/vfm/packing_iterable_dataset.py index ac30f0d1..8d8d022b 100644 --- a/cosmos_framework/data/vfm/packing_iterable_dataset.py +++ b/cosmos_framework/data/vfm/packing_iterable_dataset.py @@ -4,7 +4,7 @@ """ Abstract base class for pool-based token-budget bin-packing over multiple datasets. -Extracted from ``cosmos_framework.data.vfm.vlm.joint_dataset_dynamic_batch_webloader`` +Extracted from ``cosmos_framework.data.vfm.reasoner.joint_dataset_dynamic_batch_webloader`` so that both the VLM and VFM internal dataloaders can share a single packing implementation. Usage diff --git a/cosmos_framework/data/vfm/processors/__init__.py b/cosmos_framework/data/vfm/processors/__init__.py index ce98401a..73ea3404 100644 --- a/cosmos_framework/data/vfm/processors/__init__.py +++ b/cosmos_framework/data/vfm/processors/__init__.py @@ -14,7 +14,7 @@ from cosmos_framework.data.vfm.processors.nemotronvl_processor import NemotronVLProcessor from cosmos_framework.data.vfm.processors.qwen3vl_processor import Qwen3VLProcessor from cosmos_framework.model.vfm.tokenizers.tokenization_qwen2 import Qwen2Tokenizer -from cosmos_framework.utils.vfm.vlm.pretrained_models_downloader import maybe_download_hf_model_from_s3 +from cosmos_framework.utils.vfm.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 _VARIANT_TO_CREDENTIALS = { "s3": ("credentials/s3_training.secret", "bucket4"), diff --git a/cosmos_framework/data/vfm/processors/base.py b/cosmos_framework/data/vfm/processors/base.py index 297c3a23..7b0aa3e0 100644 --- a/cosmos_framework/data/vfm/processors/base.py +++ b/cosmos_framework/data/vfm/processors/base.py @@ -22,8 +22,8 @@ from transformers.models.auto.processing_auto import AutoProcessor from cosmos_framework.utils import log -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import tokenize_caption -from cosmos_framework.utils.vfm.vlm.pretrained_models_downloader import maybe_download_hf_model_from_s3 +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption +from cosmos_framework.utils.vfm.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 def convert_string_content_to_list_content(messages: List[Dict]) -> List[Dict]: diff --git a/cosmos_framework/data/vfm/processors/nemotron3densevl_processor.py b/cosmos_framework/data/vfm/processors/nemotron3densevl_processor.py index c586686c..8024c98a 100644 --- a/cosmos_framework/data/vfm/processors/nemotron3densevl_processor.py +++ b/cosmos_framework/data/vfm/processors/nemotron3densevl_processor.py @@ -88,16 +88,22 @@ def apply_chat_template( num_video, video_fps, video_total_num_frames, video_frames_indices = maybe_parse_video_content(messages) if num_video > 0: - assert num_video == 1, "only support one video for now" - fps = video_fps[0] - total_num_frames = video_total_num_frames[0] - frames_indices = video_frames_indices[0] - kwargs.update( - { - "do_sample_frames": False, - "video_metadata": dict(fps=fps, total_num_frames=total_num_frames, frames_indices=frames_indices), - } - ) + # Here we add the args to avoid the error: + # File "/usr/local/lib/python3.12/dist-packages/transformers/video_processing_utils.py", line 321, in _decode_and_sample_videos + # raise ValueError( + # ValueError: Sampling frames from a list of images is not supported! Set `do_sample_frames=False`. + video_metadata = [ + dict(fps=fps, total_num_frames=total_num_frames, frames_indices=frames_indices) + for fps, total_num_frames, frames_indices in zip( + video_fps, + video_total_num_frames, + video_frames_indices, + ) + ] + kwargs["videos_kwargs"] = { + "do_sample_frames": False, + "video_metadata": video_metadata[0] if num_video == 1 else video_metadata, + } inputs = self.processor.apply_chat_template( messages, diff --git a/cosmos_framework/data/vfm/processors/qwen3vl_processor.py b/cosmos_framework/data/vfm/processors/qwen3vl_processor.py index dffa47ca..2a8eb59e 100644 --- a/cosmos_framework/data/vfm/processors/qwen3vl_processor.py +++ b/cosmos_framework/data/vfm/processors/qwen3vl_processor.py @@ -75,13 +75,17 @@ def apply_chat_template( # raise ValueError( # ValueError: Sampling frames from a list of images is not supported! Set `do_sample_frames=False`. kwargs["videos_kwargs"] = dict(do_sample_frames=False) - assert num_video == 1, "only support one video for now" - fps = video_fps[0] - total_num_frames = video_total_num_frames[0] - frames_indices = video_frames_indices[0] + video_metadata = [ + dict(fps=fps, total_num_frames=total_num_frames, frames_indices=frames_indices) + for fps, total_num_frames, frames_indices in zip( + video_fps, + video_total_num_frames, + video_frames_indices, + ) + ] kwargs["videos_kwargs"] = { "do_sample_frames": False, - "video_metadata": dict(fps=fps, total_num_frames=total_num_frames, frames_indices=frames_indices), + "video_metadata": video_metadata[0] if num_video == 1 else video_metadata, } inputs = self.processor.apply_chat_template( messages, diff --git a/cosmos_framework/data/vfm/vlm/video_decoder_qwen.py b/cosmos_framework/data/vfm/reasoner/video_decoder_qwen.py similarity index 100% rename from cosmos_framework/data/vfm/vlm/video_decoder_qwen.py rename to cosmos_framework/data/vfm/reasoner/video_decoder_qwen.py diff --git a/cosmos_framework/data/vfm/sequence_packing/modalities.py b/cosmos_framework/data/vfm/sequence_packing/modalities.py index e30b1886..247789a9 100644 --- a/cosmos_framework/data/vfm/sequence_packing/modalities.py +++ b/cosmos_framework/data/vfm/sequence_packing/modalities.py @@ -4,7 +4,6 @@ """Modality-specific append helpers for VFM sequence packing.""" import math -from typing import Dict, List, Tuple import torch @@ -118,7 +117,7 @@ def add_special_tokens(tokenizer): def compute_text_split_length( num_caption_tokens: int, - special_tokens: Dict[str, int], + special_tokens: dict[str, int], has_generation: bool = True, ) -> int: """Compute the total text split length without mutating any state. @@ -145,25 +144,23 @@ def compute_text_split_length( def pack_text_tokens( packed_seq: PackedSequence, - text_ids: List[int], - special_tokens: Dict[str, int], - curr_rope_id: int, + text_ids: list[int], + special_tokens: dict[str, int], has_generation: bool, use_float_positions: bool = False, -) -> Tuple[int, int, int]: +) -> int: """Pack text tokens into the sequence. Args: packed_seq: PackedSequence instance to accumulate data into. text_ids: List of text token IDs (integers). special_tokens: Dictionary of special token IDs. - curr_rope_id: Current RoPE position ID. has_generation: Whether there's media/action after text. use_float_positions: If True, generate float position IDs for 3D mRoPE (for consistency with FPS-modulated vision tokens). Returns: - Tuple of (updated curr_rope_id, split_length, sample_length). + Text sample length. """ # Ensure we're in build mode (fields are lists, not tensors) assert isinstance(packed_seq.text_ids, list), "PackedSequence must be in build mode" @@ -212,20 +209,17 @@ def pack_text_tokens( assert split_len == compute_text_split_length(len(text_ids), special_tokens, has_generation) # Update position IDs and attention mode for text split - if packed_seq._use_mrope: - text_mrope_ids, packed_seq._mrope_temporal_offset = get_3d_mrope_ids_text_tokens( - num_tokens=split_len, - temporal_offset=packed_seq._mrope_temporal_offset, - use_float_positions=use_float_positions, - ) # text_mrope_ids: [3,split_len] - packed_seq.position_ids.append(text_mrope_ids) - else: - packed_seq.position_ids.extend(range(curr_rope_id, curr_rope_id + split_len)) + text_mrope_ids, packed_seq._mrope_temporal_offset = get_3d_mrope_ids_text_tokens( + num_tokens=split_len, + temporal_offset=packed_seq._mrope_temporal_offset, + use_float_positions=use_float_positions, + ) # text_mrope_ids: [3,split_len] + packed_seq.position_ids.append(text_mrope_ids) packed_seq.attn_modes.append("causal") packed_seq.split_lens.append(split_len) packed_seq.curr = curr - return curr_rope_id + split_len, split_len, split_len + return split_len def pack_vision_tokens( @@ -233,7 +227,6 @@ def pack_vision_tokens( input_vision_tokens: torch.Tensor, condition_frame_indexes_vision: list[int], input_timestep: float | torch.Tensor, - curr_rope_id: int, latent_patch_size: int = 1, vision_fps: float | None = None, enable_fps_modulation: bool = False, @@ -250,7 +243,6 @@ def pack_vision_tokens( input_timestep: Diffusion timestep. Either a float (teacher_forcing/none — all frames share the same sigma) or a Tensor(T_max,) (diffusion_forcing — per-frame sigma; indexed as input_timestep[frame_idx] for each noisy frame). - curr_rope_id: Current RoPE position ID. latent_patch_size: Patch size for latent patchification. vision_fps: Frames per second of the video. Used when enable_fps_modulation=True. enable_fps_modulation: If True, scale temporal position IDs based on video FPS. @@ -326,29 +318,24 @@ def pack_vision_tokens( curr += num_vision_tokens vision_split_len += num_vision_tokens - # Update position IDs for image split - if packed_seq._use_mrope: - # Determine FPS for this vision segment (None disables FPS modulation) - effective_fps = vision_fps if enable_fps_modulation else None - if vision_temporal_positions is not None: - vision_temporal_positions = vision_temporal_positions.to(device="cpu", dtype=torch.float32) # [T] - - vision_mrope_ids, packed_seq._mrope_temporal_offset = get_3d_mrope_ids_vae_tokens( - grid_t=latent_t, - grid_h=patch_h, - grid_w=patch_w, - temporal_offset=packed_seq._mrope_temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, - fps=effective_fps, - base_fps=base_fps, - temporal_compression_factor=temporal_compression_factor, - temporal_positions=vision_temporal_positions, - actual_temporal_compression_factor=temporal_compression_factor, - ) # vision_mrope_ids: [3,N_vision_tokens] - packed_seq.position_ids.append(vision_mrope_ids) - else: - # All image tokens share the same RoPE position ID - packed_seq.position_ids.extend([curr_rope_id] * vision_split_len) + # Update position IDs for image split. + effective_fps = vision_fps if enable_fps_modulation else None + if vision_temporal_positions is not None: + vision_temporal_positions = vision_temporal_positions.to(device="cpu", dtype=torch.float32) # [T] + + vision_mrope_ids, packed_seq._mrope_temporal_offset = get_3d_mrope_ids_vae_tokens( + grid_t=latent_t, + grid_h=patch_h, + grid_w=patch_w, + temporal_offset=packed_seq._mrope_temporal_offset, + reset_spatial_indices=packed_seq._mrope_reset_spatial, + fps=effective_fps, + base_fps=base_fps, + temporal_compression_factor=temporal_compression_factor, + temporal_positions=vision_temporal_positions, + actual_temporal_compression_factor=temporal_compression_factor, + ) # vision_mrope_ids: [3,N_vision_tokens] + packed_seq.position_ids.append(vision_mrope_ids) packed_seq.curr = curr return vision_split_len @@ -359,7 +346,6 @@ def pack_action_tokens( input_action_tokens: torch.Tensor, condition_frame_indexes_action: list[int], input_timestep: float, - curr_rope_id: int, action_temporal_offset: int | float = 0, enable_fps_modulation: bool = False, base_fps: float = 24.0, @@ -374,7 +360,6 @@ def pack_action_tokens( input_action_tokens: Action latent tokens (T, D). condition_frame_indexes_action: Indexes of conditioning action steps. input_timestep: Diffusion timestep. - curr_rope_id: Current RoPE position ID. action_temporal_offset: Temporal offset for action mRoPE IDs (typically the vision start offset so action aligns temporally with vision). enable_fps_modulation: If True, scale temporal position IDs based on FPS. @@ -438,30 +423,25 @@ def pack_action_tokens( packed_seq.action.mse_loss_indexes.extend(range(frame_start, frame_end)) packed_seq.action.timesteps.extend([input_timestep] * frame_token_stride) - # Update RoPE position IDs for action tokens. - if packed_seq._use_mrope: - # 3D mRoPE: action tokens use a 1x1 spatial grid with start_frame_offset=1 - # so action[0] (null token) aligns with vision frame 1, not frame 0. - effective_fps = action_fps if enable_fps_modulation else None - - action_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( - grid_t=action_split_len, - grid_h=1, - grid_w=1, - temporal_offset=action_temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, - fps=effective_fps, - base_fps=base_fps, - temporal_compression_factor=1, # Action is at frame rate (no temporal compression) - base_temporal_compression_factor=base_temporal_compression_factor, - start_frame_offset=action_start_frame_offset, # Align action[0] with vision frame action_start_frame_offset - ) # action_mrope_ids: [3,N_action_tokens] - packed_seq.position_ids.append(action_mrope_ids) - # Note: we don't update _mrope_temporal_offset here because action tokens - # share the temporal space with vision tokens (they run in parallel). - else: - # All action tokens share the SAME RoPE position as vision tokens (see docs/sequence_packing.md). - packed_seq.position_ids.extend([curr_rope_id] * action_split_len) + # Action tokens use a 1x1 spatial grid with start_frame_offset=1 by default, + # so action[0] (null token) aligns with vision frame 1, not frame 0. + effective_fps = action_fps if enable_fps_modulation else None + + action_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( + grid_t=action_split_len, + grid_h=1, + grid_w=1, + temporal_offset=action_temporal_offset, + reset_spatial_indices=packed_seq._mrope_reset_spatial, + fps=effective_fps, + base_fps=base_fps, + temporal_compression_factor=1, # Action is at frame rate (no temporal compression) + base_temporal_compression_factor=base_temporal_compression_factor, + start_frame_offset=action_start_frame_offset, # Align action[0] with vision frame action_start_frame_offset + ) # action_mrope_ids: [3,N_action_tokens] + packed_seq.position_ids.append(action_mrope_ids) + # Note: we don't update _mrope_temporal_offset here because action tokens + # share the temporal space with vision tokens (they run in parallel). packed_seq.curr = curr + action_split_len return action_split_len @@ -472,7 +452,6 @@ def pack_sound_tokens( input_sound_tokens: torch.Tensor, condition_frame_indexes_sound: list[int], input_timestep: float, - curr_rope_id: int, sound_temporal_offset: int | float = 0, enable_fps_modulation: bool = False, base_fps: float = 24.0, @@ -492,7 +471,6 @@ def pack_sound_tokens( [] means all frames are noised/supervised. All frames specified means all frames are clean (no MSE supervision). input_timestep: Diffusion timestep. - curr_rope_id: Current RoPE position ID. sound_temporal_offset: Temporal offset for m-RoPE position IDs (aligned with vision start). enable_fps_modulation: If True, scale temporal positions by FPS ratio. base_fps: Base FPS for normalization (default 24.0). @@ -556,30 +534,25 @@ def pack_sound_tokens( packed_seq.sound.mse_loss_indexes.extend(range(frame_start, frame_end)) packed_seq.sound.timesteps.extend([input_timestep]) - # Update RoPE position IDs for sound tokens. - if packed_seq._use_mrope: - # 3D mRoPE: sound tokens use a 1x1 spatial grid, aligned with vision temporal positions. - # sound[0] aligns with vision frame 0 (start_frame_offset=0, unlike action which offsets by 1). - effective_fps = sound_fps if enable_fps_modulation else None - - sound_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( - grid_t=sound_split_len, - grid_h=1, - grid_w=1, - temporal_offset=sound_temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, - fps=effective_fps, - base_fps=base_fps, - temporal_compression_factor=1, # Sound latent is already at sound_latent_fps (no further compression) - base_temporal_compression_factor=sound_base_temporal_compression_factor, - start_frame_offset=0, # Sound[0] aligns with vision frame 0 - ) # sound_mrope_ids: [3,N_sound_tokens] - packed_seq.position_ids.append(sound_mrope_ids) - # Note: we don't update _mrope_temporal_offset here because sound tokens - # share the temporal space with vision tokens (they run in parallel). - else: - # All sound tokens share the SAME RoPE position as vision/action tokens (unified generation split). - packed_seq.position_ids.extend([curr_rope_id] * sound_split_len) + # Sound tokens use a 1x1 spatial grid, aligned with vision temporal positions. + # sound[0] aligns with vision frame 0 (start_frame_offset=0, unlike action which offsets by 1). + effective_fps = sound_fps if enable_fps_modulation else None + + sound_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( + grid_t=sound_split_len, + grid_h=1, + grid_w=1, + temporal_offset=sound_temporal_offset, + reset_spatial_indices=packed_seq._mrope_reset_spatial, + fps=effective_fps, + base_fps=base_fps, + temporal_compression_factor=1, # Sound latent is already at sound_latent_fps (no further compression) + base_temporal_compression_factor=sound_base_temporal_compression_factor, + start_frame_offset=0, # Sound[0] aligns with vision frame 0 + ) # sound_mrope_ids: [3,N_sound_tokens] + packed_seq.position_ids.append(sound_mrope_ids) + # Note: we don't update _mrope_temporal_offset here because sound tokens + # share the temporal space with vision tokens (they run in parallel). packed_seq.curr = curr + sound_split_len return sound_split_len diff --git a/cosmos_framework/data/vfm/sequence_packing/packers.py b/cosmos_framework/data/vfm/sequence_packing/packers.py index a4532fd1..643c944d 100644 --- a/cosmos_framework/data/vfm/sequence_packing/packers.py +++ b/cosmos_framework/data/vfm/sequence_packing/packers.py @@ -32,7 +32,6 @@ def pack_input_sequence( latent_patch_size: int = 1, skip_text_tokens: bool = False, include_end_of_generation_token: bool = False, - position_embedding_type: str = "3d_rope", unified_3d_mrope_reset_spatial_ids: bool = True, unified_3d_mrope_temporal_modality_margin: int = 0, enable_fps_modulation: bool = False, @@ -65,13 +64,9 @@ def pack_input_sequence( latent_patch_size: Patch size used by the network to pack latents skip_text_tokens: If True, skip packing text tokens include_end_of_generation_token: If True, append end-of-generation token - position_embedding_type: Position embedding type for vision tokens: - - "3d_rope": Additive 3D RoPE embeddings + 1D position IDs for attention - - "flattened_sin_cos": Additive flattened sin/cos embeddings + 1D position IDs - - "unified_3d_mrope": No additive embedding + 3D position IDs for Qwen3VL-style mRoPE unified_3d_mrope_reset_spatial_ids: If True (default), spatial (H, W) indices start from 0 for each vision segment. If False, spatial indices are offset - by the temporal offset (Qwen2VL-style). Only used when position_embedding_type="unified_3d_mrope". + by the temporal offset (Qwen2VL-style). enable_fps_modulation: If True, scale temporal position IDs based on video FPS to reflect real time. Requires fps_vision in gen_data_clean. Uses the same flag as diffusion_expert_config.enable_fps_modulation. @@ -105,10 +100,6 @@ def pack_input_sequence( has_any_vision = any(plan.has_vision for plan in sequence_plans) explicit_vision_temporal_positions_active = vision_temporal_position_mode != "latent_index" and has_any_vision if explicit_vision_temporal_positions_active: - if position_embedding_type != "unified_3d_mrope": - raise NotImplementedError( - "Explicit vision temporal positions are only supported with position_embedding_type='unified_3d_mrope'." - ) if gen_data_clean.temporal_positions_vision is None: raise ValueError( f"vision_temporal_position_mode={vision_temporal_position_mode} requires " @@ -137,8 +128,7 @@ def pack_input_sequence( # Initialize packed sequence (acts as builder during packing) packed_seq = PackedSequence() - # Configure 3D mRoPE on the builder (enabled when position_embedding_type is unified_3d_mrope) - packed_seq._use_mrope = position_embedding_type == "unified_3d_mrope" + # Configure 3D mRoPE on the builder. packed_seq._mrope_reset_spatial = unified_3d_mrope_reset_spatial_ids # Maintain separate indices for each modality @@ -156,7 +146,6 @@ def pack_input_sequence( # Pack each sample based on its sequence plan for sample_idx, sequence_plan in enumerate(sequence_plans): - curr_rope_id = 0 sample_len = 0 # mRoPE temporal offset resets per sample. @@ -172,11 +161,10 @@ def pack_input_sequence( idx_text += 1 has_generation_for_sample = sequence_plan.has_vision or sequence_plan.has_action or sequence_plan.has_sound - curr_rope_id, _, text_sample_len = pack_text_tokens( + text_sample_len = pack_text_tokens( packed_seq, text_ids, special_tokens, - curr_rope_id, has_generation=has_generation_for_sample, use_float_positions=use_float_mrope_positions, ) @@ -192,9 +180,6 @@ def pack_input_sequence( if video_temporal_causal and sequence_plan.has_vision: # Temporal causal path: when sequence_plan.has_action=True, interleaved supertokens # [action_t, vision_t]; when False, supertokens are just vision patches. - assert position_embedding_type == "unified_3d_mrope", ( - "video_temporal_causal=True requires position_embedding_type='unified_3d_mrope'" - ) input_vision_tokens = gen_data_clean.x0_tokens_vision[idx_vision] idx_vision += 1 @@ -224,7 +209,6 @@ def pack_input_sequence( input_action_tokens=input_action_tokens_tc, condition_frame_indexes_vision=sequence_plan.condition_frame_indexes_vision, input_timestep=input_timestep, - curr_rope_id=curr_rope_id, latent_patch_size=latent_patch_size, temporal_compression_factor=temporal_compression_factor, action_dim=action_dim, @@ -355,7 +339,6 @@ def pack_input_sequence( input_vision_tokens=input_vision_tokens, condition_frame_indexes_vision=item_condition_frames, input_timestep=input_timestep, - curr_rope_id=curr_rope_id, latent_patch_size=latent_patch_size, vision_fps=vision_fps, enable_fps_modulation=enable_fps_modulation, @@ -393,7 +376,6 @@ def pack_input_sequence( input_action_tokens=input_action_tokens, condition_frame_indexes_action=sequence_plan.condition_frame_indexes_action, input_timestep=input_timestep, - curr_rope_id=curr_rope_id, action_temporal_offset=vision_start_temporal_offset, enable_fps_modulation=enable_fps_modulation, base_fps=base_fps, @@ -425,7 +407,6 @@ def pack_input_sequence( input_sound_tokens=input_sound_tokens, condition_frame_indexes_sound=sequence_plan.condition_frame_indexes_sound, input_timestep=input_timestep, - curr_rope_id=curr_rope_id, sound_temporal_offset=vision_start_temporal_offset, enable_fps_modulation=enable_fps_modulation, base_fps=base_fps, @@ -448,15 +429,11 @@ def pack_input_sequence( packed_seq.text_ids.append(special_tokens["end_of_generation"]) packed_seq.text_indexes.append(packed_seq.curr) - # EOV position IDs: 3D mRoPE or 1D RoPE - if packed_seq._use_mrope: - # Use float dtype when any vision mRoPE positions are fractional. - eov_dtype = torch.float32 if use_float_mrope_positions else torch.long - eov_mrope_ids = torch.full((3, 1), packed_seq._mrope_temporal_offset, dtype=eov_dtype) # [3,1] - packed_seq.position_ids.append(eov_mrope_ids) # type: ignore[arg-type] - packed_seq._mrope_temporal_offset += 1 - else: - packed_seq.position_ids.append(curr_rope_id) # type: ignore[arg-type] + # Use float dtype when any vision mRoPE positions are fractional. + eov_dtype = torch.float32 if use_float_mrope_positions else torch.long + eov_mrope_ids = torch.full((3, 1), packed_seq._mrope_temporal_offset, dtype=eov_dtype) # [3,1] + packed_seq.position_ids.append(eov_mrope_ids) # type: ignore[arg-type] + packed_seq._mrope_temporal_offset += 1 packed_seq.curr += 1 eov_len = 1 diff --git a/cosmos_framework/data/vfm/sequence_packing/temporal_causal.py b/cosmos_framework/data/vfm/sequence_packing/temporal_causal.py index abb7345d..2610c26a 100644 --- a/cosmos_framework/data/vfm/sequence_packing/temporal_causal.py +++ b/cosmos_framework/data/vfm/sequence_packing/temporal_causal.py @@ -17,7 +17,6 @@ def pack_supertokens_temporal_causal( input_action_tokens: torch.Tensor | None, condition_frame_indexes_vision: list[int], input_timestep: float | torch.Tensor, - curr_rope_id: int, latent_patch_size: int, temporal_compression_factor: int, action_dim: int, @@ -167,96 +166,95 @@ def pack_supertokens_temporal_causal( curr = packed_seq.curr total_split_len = 0 - # mRoPE: snapshot offset before this sample, compute IDs - if packed_seq._use_mrope: - temporal_offset = packed_seq._mrope_temporal_offset - effective_vision_fps = vision_fps if enable_fps_modulation else None - - # AR generation (single frame OR chunk) is detected by every frame carrying a - # real action (``real_actions`` has ``latent_t*tcf`` rows). There, vision AND - # action both use start_frame_offset=1 so the last action in each group - # co-locates with its vision frame, mirroring whole-clip training; the caller - # (pack_input_sequence_autoregressive) seeds temporal_offset one frame-stride - # back to compensate. Whole-clip training (frame 0 is the null conditioning - # frame, ``real_actions`` has ``(T-1)*tcf`` rows) keeps vision start_frame_offset=0. - all_frames_have_real_action = ( - pack_action_tokens and input_action_tokens is not None and real_actions.shape[0] == latent_t * tcf - ) - vision_sfo = 1 if all_frames_have_real_action else 0 - - vision_ids_flat, new_offset = get_3d_mrope_ids_vae_tokens( - grid_t=latent_t, - grid_h=patch_h, - grid_w=patch_w, - temporal_offset=temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, - fps=effective_vision_fps, - base_fps=base_fps, - temporal_compression_factor=tcf, - start_frame_offset=vision_sfo, - ) # vision_ids_flat: [3,T*patch_h*patch_w] - - if pack_action_tokens: - effective_action_fps = action_fps if enable_fps_modulation else None - - # Action IDs. Real action tokens use start_frame_offset=1 so the last - # sub-token of a group co-locates with its vision frame. Whole-clip training - # has a null action at frame 0 (the conditioning frame); AR units have a real - # action for every frame. - fps_active = effective_action_fps is not None - t_dtype = torch.float32 if fps_active else torch.long - t_offset = float(temporal_offset) if fps_active else int(temporal_offset) - null_t = torch.full((tcf,), t_offset, dtype=t_dtype) # [tcf] - null_hw = torch.zeros(tcf, dtype=t_dtype) # [tcf] - null_ids = torch.stack([null_t, null_hw, null_hw]) # [3,tcf] - - def _real_action_ids(n_frames: int, start_frame_offset: int) -> torch.Tensor: - flat, _ = get_3d_mrope_ids_vae_tokens( - grid_t=n_frames * tcf, - grid_h=1, - grid_w=1, - temporal_offset=temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, - fps=effective_action_fps, - base_fps=base_fps, - temporal_compression_factor=1, - base_temporal_compression_factor=tcf, - start_frame_offset=start_frame_offset, - ) - return flat.reshape(3, n_frames, tcf) # [3,n_frames,tcf] + # Snapshot the offset before this sample and compute mRoPE IDs. + temporal_offset = packed_seq._mrope_temporal_offset + effective_vision_fps = vision_fps if enable_fps_modulation else None + + # AR generation (single frame OR chunk) is detected by every frame carrying a + # real action (``real_actions`` has ``latent_t*tcf`` rows). There, vision AND + # action both use start_frame_offset=1 so the last action in each group + # co-locates with its vision frame, mirroring whole-clip training; the caller + # (pack_input_sequence_autoregressive) seeds temporal_offset one frame-stride + # back to compensate. Whole-clip training (frame 0 is the null conditioning + # frame, ``real_actions`` has ``(T-1)*tcf`` rows) keeps vision start_frame_offset=0. + all_frames_have_real_action = ( + pack_action_tokens and input_action_tokens is not None and real_actions.shape[0] == latent_t * tcf + ) + vision_sfo = 1 if all_frames_have_real_action else 0 + + vision_ids_flat, new_offset = get_3d_mrope_ids_vae_tokens( + grid_t=latent_t, + grid_h=patch_h, + grid_w=patch_w, + temporal_offset=temporal_offset, + reset_spatial_indices=packed_seq._mrope_reset_spatial, + fps=effective_vision_fps, + base_fps=base_fps, + temporal_compression_factor=tcf, + start_frame_offset=vision_sfo, + ) # vision_ids_flat: [3,T*patch_h*patch_w] - if all_frames_have_real_action: - # AR generation (single frame: tcf == 1*tcf, or chunk: latent_t*tcf): - # every supertoken carries a real action. start_frame_offset=1 puts - # a_{j-1}'s last sub-token on vision frame j -- the whole-clip TF - # training layout. The caller seeds temporal_offset (N-1) frame-strides - # back to compensate. - action_ids_3d = _real_action_ids(latent_t, start_frame_offset=1) # [3,T,tcf] - elif latent_t > 1: - # Whole-clip training: supertoken 0 = null (conditioning frame), frames - # 1..T-1 = real with start_frame_offset=1. Covers real-action training - # (real_actions has (T-1)*tcf rows) and the architectural all-null layout - # (input_action_tokens is None); the tokens differ but the IDs match. - null_ids_3d = null_ids.reshape(3, 1, tcf) # [3,1,tcf] - real_ids_3d = _real_action_ids(latent_t - 1, start_frame_offset=1) # [3,T-1,tcf] - action_ids_3d = torch.cat([null_ids_3d, real_ids_3d], dim=1) # [3,T,tcf] - else: - # AR frame 0 / image2video (latent_t == 1, no action): only null. - action_ids_3d = null_ids.reshape(3, 1, tcf) # [3,1,tcf] + if pack_action_tokens: + effective_action_fps = action_fps if enable_fps_modulation else None + + # Action IDs. Real action tokens use start_frame_offset=1 so the last + # sub-token of a group co-locates with its vision frame. Whole-clip training + # has a null action at frame 0 (the conditioning frame); AR units have a real + # action for every frame. + fps_active = effective_action_fps is not None + t_dtype = torch.float32 if fps_active else torch.long + t_offset = float(temporal_offset) if fps_active else int(temporal_offset) + null_t = torch.full((tcf,), t_offset, dtype=t_dtype) # [tcf] + null_hw = torch.zeros(tcf, dtype=t_dtype) # [tcf] + null_ids = torch.stack([null_t, null_hw, null_hw]) # [3,tcf] + + def _real_action_ids(n_frames: int, start_frame_offset: int) -> torch.Tensor: + flat, _ = get_3d_mrope_ids_vae_tokens( + grid_t=n_frames * tcf, + grid_h=1, + grid_w=1, + temporal_offset=temporal_offset, + reset_spatial_indices=packed_seq._mrope_reset_spatial, + fps=effective_action_fps, + base_fps=base_fps, + temporal_compression_factor=1, + base_temporal_compression_factor=tcf, + start_frame_offset=start_frame_offset, + ) + return flat.reshape(3, n_frames, tcf) # [3,n_frames,tcf] + + if all_frames_have_real_action: + # AR generation (single frame: tcf == 1*tcf, or chunk: latent_t*tcf): + # every supertoken carries a real action. start_frame_offset=1 puts + # a_{j-1}'s last sub-token on vision frame j -- the whole-clip TF + # training layout. The caller seeds temporal_offset (N-1) frame-strides + # back to compensate. + action_ids_3d = _real_action_ids(latent_t, start_frame_offset=1) # [3,T,tcf] + elif latent_t > 1: + # Whole-clip training: supertoken 0 = null (conditioning frame), frames + # 1..T-1 = real with start_frame_offset=1. Covers real-action training + # (real_actions has (T-1)*tcf rows) and the architectural all-null layout + # (input_action_tokens is None); the tokens differ but the IDs match. + null_ids_3d = null_ids.reshape(3, 1, tcf) # [3,1,tcf] + real_ids_3d = _real_action_ids(latent_t - 1, start_frame_offset=1) # [3,T-1,tcf] + action_ids_3d = torch.cat([null_ids_3d, real_ids_3d], dim=1) # [3,T,tcf] + else: + # AR frame 0 / image2video (latent_t == 1, no action): only null. + action_ids_3d = null_ids.reshape(3, 1, tcf) # [3,1,tcf] - # (3, T*H*W) → (3, T, H*W) - vision_ids_3d = vision_ids_flat.reshape(3, latent_t, patches_per_frame) # [3,T,patch_h*patch_w] + # (3, T*H*W) -> (3, T, H*W) + vision_ids_3d = vision_ids_flat.reshape(3, latent_t, patches_per_frame) # [3,T,patch_h*patch_w] - # Interleave per frame: (3, T, tcf+H*W) → (3, T*S) - interleaved_ids = torch.cat([action_ids_3d, vision_ids_3d], dim=2).reshape( - 3, latent_t * supertoken_len - ) # [3,T*S] - packed_seq.position_ids.append(interleaved_ids) - else: - # No action tokens: just vision IDs, already in (3, T*H*W) order. - packed_seq.position_ids.append(vision_ids_flat) + # Interleave per frame: (3, T, tcf+H*W) -> (3, T*S) + interleaved_ids = torch.cat([action_ids_3d, vision_ids_3d], dim=2).reshape( + 3, latent_t * supertoken_len + ) # [3,T*S] + packed_seq.position_ids.append(interleaved_ids) + else: + # No action tokens: just vision IDs, already in (3, T*H*W) order. + packed_seq.position_ids.append(vision_ids_flat) - packed_seq._mrope_temporal_offset = new_offset + packed_seq._mrope_temporal_offset = new_offset for frame_t in range(latent_t): if pack_action_tokens: @@ -267,18 +265,12 @@ def _real_action_ids(n_frames: int, start_frame_offset: int) -> torch.Tensor: curr += tcf total_split_len += tcf - if not packed_seq._use_mrope: - packed_seq.position_ids.extend([curr_rope_id] * tcf) - # Pack vision tokens for this frame frame_indexes = list(range(curr, curr + patches_per_frame)) packed_seq.vision.sequence_indexes.extend(frame_indexes) curr += patches_per_frame total_split_len += patches_per_frame - if not packed_seq._use_mrope: - packed_seq.position_ids.extend([curr_rope_id] * patches_per_frame) - # Vision MSE loss: supervise non-conditioning frames if frame_t not in condition_set_vision: packed_seq.vision.mse_loss_indexes.extend(frame_indexes) diff --git a/cosmos_framework/data/vfm/sequence_packing/types.py b/cosmos_framework/data/vfm/sequence_packing/types.py index 0fdd4ef3..38018bd0 100644 --- a/cosmos_framework/data/vfm/sequence_packing/types.py +++ b/cosmos_framework/data/vfm/sequence_packing/types.py @@ -119,17 +119,16 @@ class PackedSequence: # Text modality (list during build, tensor after finalize) text_ids: list[int] | torch.Tensor = field(default_factory=list) text_indexes: list[int] | torch.Tensor = field(default_factory=list) - position_ids: list[int] | torch.Tensor = field(default_factory=list) + position_ids: list[torch.Tensor] | torch.Tensor = field(default_factory=list) # Loss computation - Cross Entropy (text) label_ids: list[int] | torch.Tensor | None = field(default_factory=list) ce_loss_indexes: list[int] | torch.Tensor | None = field(default_factory=list) ce_loss_weights: list[float] | torch.Tensor | None = field(default_factory=list) - # Build-time mRoPE tracking (used during packing, not after finalize) - # When _use_mrope=True, position_ids accumulates (3, N) tensors instead of ints, - # and finalize() produces a (3, total_seq_len) tensor instead of (total_seq_len,). - _use_mrope: bool = False + # Build-time mRoPE tracking (used during packing, not after finalize). + # position_ids accumulates (3, N) tensors and finalize() produces a + # (3, total_seq_len) tensor. # Running temporal index for mRoPE position ID generation within a single sample. # Reset to 0 at the start of each sample, then advanced by text and vision helpers # as segments are packed. Action reuses the pre-vision snapshot (parallel temporal @@ -201,7 +200,7 @@ def finalize( if self.vision is not None and len(self.vision.sequence_indexes) > 0: vision = ModalityData( sequence_indexes=torch.tensor(self.vision.sequence_indexes, dtype=torch.long), # [N_vision_tokens] - timesteps=torch.tensor(self.vision.timesteps), # [N_vision_noisy_tokens] + timesteps=torch.tensor(self.vision.timesteps, dtype=torch.float32), # [N_vision_noisy_tokens] mse_loss_indexes=torch.tensor( self.vision.mse_loss_indexes, dtype=torch.long ), # [N_vision_noisy_tokens] @@ -216,7 +215,7 @@ def finalize( if self.action is not None and len(self.action.sequence_indexes) > 0: action = ModalityData( sequence_indexes=torch.tensor(self.action.sequence_indexes, dtype=torch.long), # [N_action_tokens] - timesteps=torch.tensor(self.action.timesteps), # [N_action_noisy_tokens] + timesteps=torch.tensor(self.action.timesteps, dtype=torch.float32), # [N_action_noisy_tokens] mse_loss_indexes=torch.tensor( self.action.mse_loss_indexes, dtype=torch.long ), # [N_action_noisy_tokens] @@ -237,7 +236,7 @@ def finalize( if self.sound is not None and len(self.sound.sequence_indexes) > 0: sound = ModalityData( sequence_indexes=torch.tensor(self.sound.sequence_indexes, dtype=torch.long), # [N_sound_tokens] - timesteps=torch.tensor(self.sound.timesteps), # [N_sound_noisy_tokens] + timesteps=torch.tensor(self.sound.timesteps, dtype=torch.float32), # [N_sound_noisy_tokens] mse_loss_indexes=torch.tensor(self.sound.mse_loss_indexes, dtype=torch.long), # [N_sound_noisy_tokens] token_shapes=list(self.sound.token_shapes), tokens=self.sound.tokens, @@ -245,12 +244,12 @@ def finalize( noisy_frame_indexes=list(self.sound.noisy_frame_indexes), ) - # Finalize position IDs: 3D mRoPE (3, seq_len) or 1D RoPE (seq_len,) - if self._use_mrope and len(self.position_ids) > 0 and isinstance(self.position_ids[0], torch.Tensor): - mrope_tensors: list[torch.Tensor] = self.position_ids # type: ignore[assignment] - position_ids = torch.cat(mrope_tensors, dim=1) # [3,actual_seq_len] - else: # Original 1D RoPE from Bagel, where all the media tokens share the same 1D position ID - position_ids = torch.tensor(self.position_ids) # [seq_len] + # Finalize position IDs. + assert isinstance(self.position_ids, list) + if len(self.position_ids) > 0: + position_ids = torch.cat(self.position_ids, dim=1) # [3,actual_seq_len] + else: + position_ids = torch.empty((3, 0), dtype=torch.long) # [3,0] return PackedSequence( # Sequence structure @@ -262,7 +261,7 @@ def finalize( # Text modality (converted to tensors) text_ids=torch.tensor(self.text_ids, dtype=torch.long), # [N_text_tokens] text_indexes=torch.tensor(self.text_indexes, dtype=torch.long), # [N_text_tokens] - position_ids=position_ids, # [seq_len] or [3,seq_len] + position_ids=position_ids, # [3,seq_len] # Loss computation - Cross Entropy label_ids=label_ids, ce_loss_indexes=ce_loss_indexes, diff --git a/cosmos_framework/inference/common/config.py b/cosmos_framework/inference/common/config.py index c376d757..121b1252 100644 --- a/cosmos_framework/inference/common/config.py +++ b/cosmos_framework/inference/common/config.py @@ -442,6 +442,20 @@ def _structure_torch_tensor(data: Any, _cls: Any) -> torch.Tensor | None: import cosmos_framework.model.vfm # noqa: F401 CONFIG_REPLACEMENTS_INVERSE = [ + # vlm → reasoner (upstream rename in i4). MUST precede the general vfm rules + # so the specific vlm→reasoner subtree isn't shadowed by them. + (r"(? bool: def _canonicalize_module_path(path: str) -> str: replacements = ( + # vlm → reasoner rename (upstream i4). Longer/more-specific rules first + # so the reasoner subtree isn't shadowed by the general vfm rules below. + ("cosmos_framework.configs.base.defaults.reasoner.", "projects.cosmos3.vfm.configs.base.defaults.vlm."), + ("cosmos_framework.configs.base.reasoner.", "projects.cosmos3.vfm.configs.base.vlm."), + ("cosmos_framework.model.vfm.reasoner.", "projects.cosmos3.vfm.models.vlm."), + ("cosmos_framework.data.vfm.reasoner.", "projects.cosmos3.vfm.datasets.vlm."), + ("cosmos_framework.data.vfm.augmentors.reasoner.", "projects.cosmos3.vfm.datasets.augmentors.vlm."), + ("cosmos_framework.utils.vfm.reasoner.", "projects.cosmos3.vfm.utils.vlm."), ("cosmos3._src.vfm.", "projects.cosmos3.vfm."), ("cosmos_framework.configs.base.", "projects.cosmos3.vfm.configs.base."), ("cosmos_framework.model.vfm.tokenizers.", "projects.cosmos3.vfm.tokenizers."), @@ -205,6 +213,14 @@ def _runtime_module_path(canonical_path: str) -> str: def _replace_vfm_module_prefix(canonical_path: str, *, package: str) -> str: replacements = ( + # vlm → reasoner (upstream rename). MUST precede the general vfm rules + # so the specific vlm→reasoner subtree isn't shadowed by them. + ("projects.cosmos3.vfm.configs.base.defaults.vlm.", f"{package}.configs.base.defaults.reasoner."), + ("projects.cosmos3.vfm.configs.base.vlm.", f"{package}.configs.base.reasoner."), + ("projects.cosmos3.vfm.models.vlm.", f"{package}.model.vfm.reasoner."), + ("projects.cosmos3.vfm.datasets.vlm.", f"{package}.data.vfm.reasoner."), + ("projects.cosmos3.vfm.datasets.augmentors.vlm.", f"{package}.data.vfm.augmentors.reasoner."), + ("projects.cosmos3.vfm.utils.vlm.", f"{package}.utils.vfm.reasoner."), ("projects.cosmos3.vfm.configs.base.", f"{package}.configs.base."), ("projects.cosmos3.vfm.models.", f"{package}.model.vfm."), ("projects.cosmos3.vfm.tokenizers.", f"{package}.model.vfm.tokenizers."), diff --git a/cosmos_framework/inference/common/public_model_config_test.py b/cosmos_framework/inference/common/public_model_config_test.py index 61ba16cf..feb63ec4 100644 --- a/cosmos_framework/inference/common/public_model_config_test.py +++ b/cosmos_framework/inference/common/public_model_config_test.py @@ -33,14 +33,14 @@ def test_public_model_config_round_trip_removes_internal_metadata(): "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth", }, "vlm_config": { - "_type": "cosmos_framework.configs.base.defaults.vlm.VLMConfig", + "_type": "cosmos_framework.configs.base.defaults.reasoner.VLMConfig", "model_instance": { "_target_": "cosmos_framework.model.vfm.mot.unified_mot.Qwen3VLTextForCausalLM", "config": { - "_target_": "cosmos_framework.configs.base.defaults.vlm.create_vlm_config", + "_target_": "cosmos_framework.configs.base.defaults.reasoner.create_vlm_config", "base_config": { "_target_": "cosmos_framework.model.vfm.mot.unified_mot.Qwen3VLMoTConfig.from_json_file", - "json_file": "cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json", + "json_file": "cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json", }, }, }, diff --git a/cosmos_framework/inference/configs/model/Cosmos3-Nano.yaml b/cosmos_framework/inference/configs/model/Cosmos3-Nano.yaml index c202d876..44f75e59 100644 --- a/cosmos_framework/inference/configs/model/Cosmos3-Nano.yaml +++ b/cosmos_framework/inference/configs/model/Cosmos3-Nano.yaml @@ -36,10 +36,6 @@ model: load_weights_from_pretrained: false max_vae_latent_side_after_patchify: 20 patch_spatial: 2 - position_embedding_type: unified_3d_mrope - rope_h_extrapolation_ratio: 1.0 - rope_t_extrapolation_ratio: 1.0 - rope_w_extrapolation_ratio: 1.0 timestep_range: 1.0 unified_3d_mrope_reset_spatial_ids: true unified_3d_mrope_temporal_modality_margin: 15000 diff --git a/cosmos_framework/inference/configs/model/Cosmos3-Super.yaml b/cosmos_framework/inference/configs/model/Cosmos3-Super.yaml index 9eb38147..f8b188af 100644 --- a/cosmos_framework/inference/configs/model/Cosmos3-Super.yaml +++ b/cosmos_framework/inference/configs/model/Cosmos3-Super.yaml @@ -36,10 +36,6 @@ model: load_weights_from_pretrained: false max_vae_latent_side_after_patchify: 20 patch_spatial: 2 - position_embedding_type: unified_3d_mrope - rope_h_extrapolation_ratio: 1.0 - rope_t_extrapolation_ratio: 1.0 - rope_w_extrapolation_ratio: 1.0 timestep_range: 1.0 unified_3d_mrope_reset_spatial_ids: true unified_3d_mrope_temporal_modality_margin: 15000 diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index 564da97f..a35968e5 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -51,7 +51,7 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import _SYSTEM_PROMPT_IMAGE_EDITING +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import _SYSTEM_PROMPT_IMAGE_EDITING from cosmos_framework.model.vfm.upsampler.prompts import is_upsampled_prompt if TYPE_CHECKING: diff --git a/cosmos_framework/inference/transfer.py b/cosmos_framework/inference/transfer.py index 61a99b49..17aead8e 100644 --- a/cosmos_framework/inference/transfer.py +++ b/cosmos_framework/inference/transfer.py @@ -29,7 +29,7 @@ from cosmos_framework.data.vfm.sequence_packing import SequencePlan from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel from cosmos_framework.model.vfm.utils.data_and_condition import GenerationDataClean -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import _SYSTEM_PROMPT_TRANSFER +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import _SYSTEM_PROMPT_TRANSFER @dataclass diff --git a/cosmos_framework/model/attention/benchmarks/benchmark_fmha_head_sharding.py b/cosmos_framework/model/attention/benchmarks/benchmark_fmha_head_sharding.py new file mode 100644 index 00000000..ce0709a3 --- /dev/null +++ b/cosmos_framework/model/attention/benchmarks/benchmark_fmha_head_sharding.py @@ -0,0 +1,398 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Standalone FMHA microbenchmark for CP attention sharding characteristics. + +This script isolates per-rank attention kernel shapes used by possible context +parallelism strategies. ``--shard-mode head`` matches the current +CP shape where ``--cp-size`` only changes the local head counts: + + CP=1: q [1,S_q,H,D], k/v [1,S_kv,H_kv,D] + CP=4: q [1,S_q,H/4,D], k/v [1,S_kv,H_kv/4,D] + +``--shard-mode sequence`` models the local split-KV kernel shape: + + CP=1: q [1,S_q,H,D], k/v [1,S_kv,H_kv,D] + CP=4: q [1,S_q,H,D], k/v [1,S_kv/4,H_kv,D] + +In sequence mode, Q is intentionally not divided. Each rank would compute a +partial result for the same query tokens over a local KV shard, and a complete +distributed implementation would merge partial outputs using LSEs. This +benchmark times only the local attention kernel. + +There is intentionally no all-to-all, all-gather, or all-reduce in the timed +region since the intent is to study kernel scalability. + +Example production-like one-GPU run: + + torchrun --standalone --nproc_per_node=1 \ + -m cosmos_framework.model.attention.benchmarks.benchmark_fmha_head_sharding \ + --cp-size 1 \ + --q-len 396 \ + --kv-len 177771 \ + --num-q-heads 32 \ + --num-kv-heads 8 \ + --head-dim 128 \ + --warmup-iters 20 \ + --iters 100 + +Example true four-GPU CP run: + + torchrun --standalone --nproc_per_node=4 \ + -m cosmos_framework.model.attention.benchmarks.benchmark_fmha_head_sharding \ + --shard-mode head \ + --cp-size 4 \ + --q-len 396 \ + --kv-len 177771 \ + --num-q-heads 32 \ + --num-kv-heads 8 \ + --head-dim 128 \ + --warmup-iters 20 \ + --iters 100 + +Example one-GPU local sequence-shard mock of a CP=4 split-KV attention shape: + + torchrun --standalone --nproc_per_node=1 \ + -m cosmos_framework.model.attention.benchmarks.benchmark_fmha_head_sharding \ + --shard-mode sequence \ + --cp-size 4 \ + --q-len 396 \ + --kv-len 177771 \ + --num-q-heads 32 \ + --num-kv-heads 8 \ + --head-dim 128 \ + --warmup-iters 20 \ + --iters 100 + +Example one-GPU local-head mock of the CP=4 per-rank attention shape: + + torchrun --standalone --nproc_per_node=1 \ + -m cosmos_framework.model.attention.benchmarks.benchmark_fmha_head_sharding \ + --shard-mode head \ + --cp-size 1 \ + --q-len 396 \ + --kv-len 177771 \ + --num-q-heads 8 \ + --num-kv-heads 2 \ + --head-dim 128 \ + --warmup-iters 20 \ + --iters 100 + +The mock uses one process and sets the global head counts to the local head +counts that a CP=4 rank would see. It creates q [1,396,8,128] and k/v +[1,177771,2,128] without launching four ranks. This isolates whether the +attention kernel scales with the smaller local-head shape, independent of +distributed launch overhead, communication, or multi-GPU contention. +""" + +from __future__ import annotations + +import argparse +import json +import os +from dataclasses import asdict, dataclass +from typing import Any + +import torch +import torch.distributed as dist + +from cosmos_framework.model.attention import attention + + +@dataclass(frozen=True) +class BenchmarkConfig: + q_len: int + kv_len: int + num_q_heads: int + num_kv_heads: int + head_dim: int + cp_size: int + shard_mode: str + batch_size: int + warmup_iters: int + iters: int + dtype: str + backend: str | None + compile: bool + seed: int + + +@dataclass(frozen=True) +class LocalAttentionShape: + q_len: int + kv_len: int + num_q_heads: int + num_kv_heads: int + sequence_shard_index: int + sequence_shard_start: int + + +def _parse_args() -> BenchmarkConfig: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--q-len", type=int, default=512, help="Current-frame query token count.") + parser.add_argument("--kv-len", type=int, default=262144, help="Total key/value token count including history.") + parser.add_argument("--num-q-heads", type=int, default=32, help="Global query head count before head sharding.") + parser.add_argument("--num-kv-heads", type=int, default=8, help="Global KV head count before head sharding.") + parser.add_argument("--head-dim", type=int, default=128, help="Per-head dimension.") + parser.add_argument("--cp-size", type=int, default=1, help="Context-parallel sharding factor to simulate.") + parser.add_argument( + "--shard-mode", + choices=("head", "sequence"), + default="head", + help=( + "head: divide Q/KV heads by cp_size and keep full KV length; " + "sequence: keep heads and split KV length across cp_size shards." + ), + ) + parser.add_argument("--batch-size", type=int, default=1, help="Batch size for the dense attention call.") + parser.add_argument("--warmup-iters", type=int, default=20, help="Untimed warmup iterations.") + parser.add_argument("--iters", type=int, default=100, help="Timed iterations.") + parser.add_argument( + "--dtype", + choices=("bfloat16", "float16", "float32"), + default="bfloat16", + help="Q/K/V dtype.", + ) + parser.add_argument("--backend", type=str, default=None, help="Optional cosmos_framework attention backend override.") + parser.add_argument("--compile", action="store_true", help="Compile the attention call before benchmarking.") + parser.add_argument("--seed", type=int, default=1234, help="Random seed.") + args = parser.parse_args() + return BenchmarkConfig( + q_len=args.q_len, + kv_len=args.kv_len, + num_q_heads=args.num_q_heads, + num_kv_heads=args.num_kv_heads, + head_dim=args.head_dim, + cp_size=args.cp_size, + shard_mode=args.shard_mode, + batch_size=args.batch_size, + warmup_iters=args.warmup_iters, + iters=args.iters, + dtype=args.dtype, + backend=args.backend, + compile=args.compile, + seed=args.seed, + ) + + +def _dtype_from_name(name: str) -> torch.dtype: + if name == "bfloat16": + return torch.bfloat16 + if name == "float16": + return torch.float16 + if name == "float32": + return torch.float32 + raise ValueError(f"Unsupported dtype={name!r}") + + +def _init_rank() -> tuple[int, int, int]: + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + if world_size > 1 and not dist.is_initialized(): + dist.init_process_group("nccl") + return rank, local_rank, world_size + + +def _validate_config(config: BenchmarkConfig, rank: int) -> LocalAttentionShape: + if config.batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {config.batch_size}") + if config.q_len <= 0: + raise ValueError(f"q_len must be positive, got {config.q_len}") + if config.kv_len <= 0: + raise ValueError(f"kv_len must be positive, got {config.kv_len}") + if config.head_dim <= 0: + raise ValueError(f"head_dim must be positive, got {config.head_dim}") + if config.warmup_iters < 0: + raise ValueError(f"warmup_iters must be non-negative, got {config.warmup_iters}") + if config.iters <= 0: + raise ValueError(f"iters must be positive, got {config.iters}") + if config.cp_size <= 0: + raise ValueError(f"cp_size must be positive, got {config.cp_size}") + if config.num_q_heads % config.num_kv_heads != 0: + raise ValueError(f"num_q_heads={config.num_q_heads} must be divisible by num_kv_heads={config.num_kv_heads}") + if config.shard_mode == "head": + if config.num_q_heads % config.cp_size != 0: + raise ValueError(f"num_q_heads={config.num_q_heads} must be divisible by cp_size={config.cp_size}") + if config.num_kv_heads % config.cp_size != 0: + raise ValueError(f"num_kv_heads={config.num_kv_heads} must be divisible by cp_size={config.cp_size}") + local_q_heads = config.num_q_heads // config.cp_size + local_kv_heads = config.num_kv_heads // config.cp_size + if local_q_heads % local_kv_heads != 0: + raise ValueError(f"local_q_heads={local_q_heads} must be divisible by local_kv_heads={local_kv_heads}") + return LocalAttentionShape( + q_len=config.q_len, + kv_len=config.kv_len, + num_q_heads=local_q_heads, + num_kv_heads=local_kv_heads, + sequence_shard_index=0, + sequence_shard_start=0, + ) + if config.shard_mode == "sequence": + shard_index = rank % config.cp_size + base_kv_len = config.kv_len // config.cp_size + remainder = config.kv_len % config.cp_size + local_kv_len = base_kv_len + int(shard_index < remainder) + if local_kv_len <= 0: + raise ValueError( + f"sequence-sharded local_kv_len must be positive, got {local_kv_len}; " + f"kv_len={config.kv_len}, cp_size={config.cp_size}, shard_index={shard_index}" + ) + sequence_shard_start = shard_index * base_kv_len + min(shard_index, remainder) + return LocalAttentionShape( + q_len=config.q_len, + kv_len=local_kv_len, + num_q_heads=config.num_q_heads, + num_kv_heads=config.num_kv_heads, + sequence_shard_index=shard_index, + sequence_shard_start=sequence_shard_start, + ) + raise ValueError(f"Unsupported shard_mode={config.shard_mode!r}") + + +def _make_inputs( + config: BenchmarkConfig, + local_shape: LocalAttentionShape, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + generator = torch.Generator(device=device) + generator.manual_seed(config.seed) + query = torch.randn( + config.batch_size, + local_shape.q_len, + local_shape.num_q_heads, + config.head_dim, + device=device, + dtype=dtype, + generator=generator, + ) # [B,S_q,H_local,D] + key = torch.randn( + config.batch_size, + local_shape.kv_len, + local_shape.num_kv_heads, + config.head_dim, + device=device, + dtype=dtype, + generator=generator, + ) # [B,S_kv,H_kv_local,D] + value = torch.randn( + config.batch_size, + local_shape.kv_len, + local_shape.num_kv_heads, + config.head_dim, + device=device, + dtype=dtype, + generator=generator, + ) # [B,S_kv,H_kv_local,D] + return query, key, value + + +def _run_attention( + query: torch.Tensor, # [B,S_q,H_local,D] + key: torch.Tensor, # [B,S_kv,H_kv_local,D] + value: torch.Tensor, # [B,S_kv,H_kv_local,D] + config: BenchmarkConfig, +) -> torch.Tensor: # [B,S_q,H_local,D] + out = attention( + query=query, + key=key, + value=value, + is_causal=False, + backend=config.backend, + return_lse=False, + ) # [B,S_q,H_local,D] + assert isinstance(out, torch.Tensor) + return out + + +def _benchmark( + query: torch.Tensor, # [B,S_q,H_local,D] + key: torch.Tensor, # [B,S_kv,H_kv_local,D] + value: torch.Tensor, # [B,S_kv,H_kv_local,D] + config: BenchmarkConfig, +) -> dict[str, Any]: + def run_once( + local_query: torch.Tensor, # [B,S_q,H_local,D] + local_key: torch.Tensor, # [B,S_kv,H_kv_local,D] + local_value: torch.Tensor, # [B,S_kv,H_kv_local,D] + ) -> torch.Tensor: # [B,S_q,H_local,D] + return _run_attention(local_query, local_key, local_value, config) # [B,S_q,H_local,D] + + if config.compile: + run_once = torch.compile(run_once, fullgraph=True) + + with torch.inference_mode(): + for _ in range(config.warmup_iters): + _warmup_out = run_once(query, key, value) # [B,S_q,H_local,D] + torch.cuda.synchronize() + + torch.cuda.reset_peak_memory_stats(query.device) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.nvtx.range_push(f"fmha_{config.shard_mode}_sharding.cp{config.cp_size}") + start.record() + for _ in range(config.iters): + out = run_once(query, key, value) # [B,S_q,H_local,D] + end.record() + torch.cuda.nvtx.range_pop() + torch.cuda.synchronize() + + elapsed_ms = start.elapsed_time(end) + tokens = config.batch_size * config.q_len + checksum = float(out.float().mean().item()) # [] + return { + "elapsed_ms": elapsed_ms, + "iters": config.iters, + "avg_ms": elapsed_ms / config.iters, + "max_memory_allocated_bytes": torch.cuda.max_memory_allocated(query.device), + "tokens_per_second": (tokens * config.iters) / (elapsed_ms / 1000.0), + "checksum": checksum, + } + + +def main() -> None: + config = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this FMHA benchmark") + rank, local_rank, world_size = _init_rank() + try: + local_shape = _validate_config(config, rank) + device = torch.device("cuda", local_rank) + dtype = _dtype_from_name(config.dtype) + query, key, value = _make_inputs(config, local_shape, device, dtype) + # query: [B,S_q,H_local,D], key/value: [B,S_kv,H_kv_local,D] + + if dist.is_initialized(): + dist.barrier() + result = _benchmark(query, key, value, config) + if dist.is_initialized(): + dist.barrier() + + payload = { + "rank": rank, + "local_rank": local_rank, + "world_size": world_size, + "device": torch.cuda.get_device_name(device), + "config": asdict(config), + "local_shape": asdict(local_shape), + "local_q_heads": local_shape.num_q_heads, + "local_kv_heads": local_shape.num_kv_heads, + "local_q_len": local_shape.q_len, + "local_kv_len": local_shape.kv_len, + "query_shape": list(query.shape), + "key_shape": list(key.shape), + "value_shape": list(value.shape), + "result": result, + } + if rank == 0: + print(json.dumps(payload, sort_keys=True)) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py b/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py index eca542aa..4eeb0657 100644 --- a/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py +++ b/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py @@ -486,11 +486,15 @@ def reset(self) -> None: if self._fid_metric is not None: self._fid_metric.reset() + + + __all__ = [ "TokenizerMetric", "PSNRMetric", "SSIMMetric", "LPIPSMetric", "Rank0FIDMetric", + "FVDMetric", "calculate_psnr", ] diff --git a/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py b/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py index 24002875..f69884b9 100644 --- a/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py +++ b/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py @@ -1354,9 +1354,13 @@ def __init__( self._logged_decoder_temporal_plan = False # Load SigLIP2 pretrained model (text encoder always needed for text alignment) - # Use HF_HUB_CACHE (set by configure_hf_cache_env to HF_HOME/hub) so the cache_dir - # matches the actual hub layout where models are stored. - hf_cache_dir = os.environ.get("HF_HUB_CACHE") or os.environ.get("HF_HOME") + # Prefer HF_HUB_CACHE (set by configure_hf_cache_env to HF_HOME/hub), then fall back to + # HF_HOME/hub (where transformers stores models), then HF_HOME itself. This handles jobs + # that set HF_HOME but not HF_HUB_CACHE (e.g. VFM training via submit_helper). + _hf_home = os.environ.get("HF_HOME") + hf_cache_dir = ( + os.environ.get("HF_HUB_CACHE") or (os.path.join(_hf_home, "hub") if _hf_home else None) or _hf_home + ) local_files_only = hf_cache_dir is not None pretrained_model = None pretrained_vision_model = None diff --git a/cosmos_framework/model/tokenizer/models/text_decoder.py b/cosmos_framework/model/tokenizer/models/text_decoder.py index 234a5f1f..b9f296c1 100644 --- a/cosmos_framework/model/tokenizer/models/text_decoder.py +++ b/cosmos_framework/model/tokenizer/models/text_decoder.py @@ -38,6 +38,7 @@ prepare_nemotron_tokenizer_snapshot, resolve_hf_snapshot_path, ) +from cosmos_framework.model.tokenizer.utils.vlm_prompt_format import densevl_add_vision_id_text QWEN3_PAD_TOKEN_ID = 151643 QWEN3_IM_START_TOKEN_ID = 151644 @@ -57,6 +58,34 @@ NEMOTRON_2B_THINK_START_TOKEN = "" NEMOTRON_2B_THINK_END_TOKEN = "" TEXT_DECODER_ATTN_IMPLEMENTATION_ENV = "TOKENIZER_TEXT_DECODER_ATTN_IMPLEMENTATION" +VQA_THINKING_MODE_OFF = "off" +VQA_THINKING_MODE_ON = "on" +VQA_THINKING_MODE_RAW = "raw" +VQA_THINKING_MODES = frozenset({VQA_THINKING_MODE_OFF, VQA_THINKING_MODE_ON, VQA_THINKING_MODE_RAW}) +VQA_REASONING_SUFFIX = ( + "\nAnswer the question using the following format:\n\n" + "\nYour reasoning.\n\n\n" + "Write your final answer immediately after the tag." +) + + +def normalize_vqa_thinking_mode(thinking_mode: str | None) -> str: + """Normalize the VQA generation thinking-mode knob.""" + normalized = VQA_THINKING_MODE_OFF if thinking_mode is None else str(thinking_mode).strip().lower() + if normalized not in VQA_THINKING_MODES: + raise ValueError( + f"Unsupported VQA thinking_mode={thinking_mode!r}; expected one of {sorted(VQA_THINKING_MODES)}." + ) + return normalized + + +def _append_vqa_reasoning_suffix(question: str, reasoning_suffix: str) -> str: + """Append the training-compatible reasoning-format suffix once.""" + if not reasoning_suffix: + return question + if reasoning_suffix.strip() in question: + return question + return f"{question.rstrip()}{reasoning_suffix}" @dataclass(frozen=True) @@ -744,87 +773,89 @@ def _build_vqa_user_turn( num_image_tokens: int, ) -> tuple[list[int], int]: """Build one multimodal user turn and return the image-pad start offset.""" + user_turn, image_pad_offsets = self._build_vqa_user_turn_with_visual_blocks( + tokenizer=tokenizer, + question=question, + image_token_counts_by_visual=[int(num_image_tokens)], + ) + return user_turn, image_pad_offsets[0] + + def _build_vqa_user_turn_with_visual_blocks( + self, + tokenizer: Any, + question: str, + image_token_counts_by_visual: list[int], + ) -> tuple[list[int], list[int]]: + """Build one multimodal user turn and return image-pad offsets per visual.""" + if not image_token_counts_by_visual: + raise ValueError("VQA prompt construction requires at least one visual token block.") + if any(int(token_count) <= 0 for token_count in image_token_counts_by_visual): + raise ValueError(f"Visual token counts must be positive, got {image_token_counts_by_visual!r}.") def _encode(text: str) -> list[int]: return tokenizer.encode(text, add_special_tokens=False) image_placeholder = "" - normalized_question = _keep_only_first_image_placeholder(str(question).strip(), image_placeholder) - if self.spec.family == QWEN3_SPEC.family: - user_prefix = [QWEN3_IM_START_TOKEN_ID] + _encode("user\n") - - if image_placeholder in normalized_question: - placeholder_index = normalized_question.find(image_placeholder) - before_text = normalized_question[:placeholder_index] - after_text = normalized_question[placeholder_index + len(image_placeholder) :] - before_ids = _encode(before_text) - after_ids = _encode(after_text) - image_pad_offset = len(user_prefix) + len(before_ids) + 1 - user_turn = ( - user_prefix - + before_ids - + [QWEN3_VISION_START_TOKEN_ID] - + [QWEN3_VISION_PAD_TOKEN_ID] * num_image_tokens - + [QWEN3_VISION_END_TOKEN_ID] - + after_ids - + [QWEN3_IM_END_TOKEN_ID] - + _encode("\n") + num_visuals = len(image_token_counts_by_visual) + normalized_question = str(question).strip() + if num_visuals == 1: + normalized_question = _keep_only_first_image_placeholder(normalized_question, image_placeholder) + else: + placeholder_count = normalized_question.count(image_placeholder) + if placeholder_count == 0: + visual_placeholders = "\n".join(image_placeholder for _ in image_token_counts_by_visual) + normalized_question = f"{visual_placeholders}\n{normalized_question}".strip() + elif placeholder_count != num_visuals: + raise ValueError( + f"VQA prompt has {placeholder_count} image placeholders for {num_visuals} visual inputs." ) - return user_turn, image_pad_offset - - image_pad_offset = len(user_prefix) + 1 - user_turn = ( - user_prefix - + [QWEN3_VISION_START_TOKEN_ID] - + [QWEN3_VISION_PAD_TOKEN_ID] * num_image_tokens - + [QWEN3_VISION_END_TOKEN_ID] - + _encode("\n") - + _encode(normalized_question) - + [QWEN3_IM_END_TOKEN_ID] - + _encode("\n") - ) - return user_turn, image_pad_offset - if self.spec.family == NEMOTRON_2B_SPEC.family: + if self.spec.family == QWEN3_SPEC.family: + im_start_id = QWEN3_IM_START_TOKEN_ID + im_end_id = QWEN3_IM_END_TOKEN_ID + vision_start_id = QWEN3_VISION_START_TOKEN_ID + vision_pad_id = QWEN3_VISION_PAD_TOKEN_ID + vision_end_id = QWEN3_VISION_END_TOKEN_ID + elif self.spec.family == NEMOTRON_2B_SPEC.family: im_start_id = _get_required_token_id(tokenizer, NEMOTRON_2B_IM_START_TOKEN) im_end_id = _get_required_token_id(tokenizer, NEMOTRON_2B_IM_END_TOKEN) - user_prefix = [im_start_id] + _encode("user\n") - - if image_placeholder in normalized_question: - placeholder_index = normalized_question.find(image_placeholder) - before_text = normalized_question[:placeholder_index] - after_text = normalized_question[placeholder_index + len(image_placeholder) :] - before_ids = _encode(before_text) - after_ids = _encode(after_text) - image_pad_offset = len(user_prefix) + len(before_ids) + 1 - user_turn = ( - user_prefix - + before_ids - + [self.spec.vision_start_token_id] - + [self.spec.vision_pad_token_id] * num_image_tokens - + [self.spec.vision_end_token_id] - + after_ids - + [im_end_id] - + _encode("\n") - ) - return user_turn, image_pad_offset - - image_pad_offset = len(user_prefix) + 1 - user_turn = ( - user_prefix - + [self.spec.vision_start_token_id] - + [self.spec.vision_pad_token_id] * num_image_tokens - + [self.spec.vision_end_token_id] - + _encode("\n") - + _encode(normalized_question) - + [im_end_id] - + _encode("\n") + vision_start_id = self.spec.vision_start_token_id + vision_pad_id = self.spec.vision_pad_token_id + vision_end_id = self.spec.vision_end_token_id + else: + raise NotImplementedError( + f"VQA prompt construction is not implemented for text decoder family {self.spec.family!r}." ) - return user_turn, image_pad_offset - raise NotImplementedError( - f"VQA prompt construction is not implemented for text decoder family {self.spec.family!r}." - ) + user_turn = [im_start_id] + _encode("user\n") + image_pad_offsets: list[int] = [] + remaining_text = normalized_question + visual_index = 0 + while image_placeholder in remaining_text: + placeholder_index = remaining_text.find(image_placeholder) + before_text = remaining_text[:placeholder_index] + remaining_text = remaining_text[placeholder_index + len(image_placeholder) :] + user_turn.extend(_encode(before_text)) + if num_visuals > 1: + user_turn.extend(_encode(densevl_add_vision_id_text("image", visual_index + 1))) + image_pad_offsets.append(len(user_turn) + 1) + user_turn.extend([vision_start_id]) + user_turn.extend([vision_pad_id] * int(image_token_counts_by_visual[visual_index])) + user_turn.extend([vision_end_id]) + visual_index += 1 + if image_pad_offsets: + user_turn.extend(_encode(remaining_text)) + else: + image_pad_offsets.append(len(user_turn) + 1) + user_turn.extend([vision_start_id]) + user_turn.extend([vision_pad_id] * int(image_token_counts_by_visual[0])) + user_turn.extend([vision_end_id]) + user_turn.extend(_encode("\n")) + user_turn.extend(_encode(normalized_question)) + + user_turn.extend([im_end_id]) + user_turn.extend(_encode("\n")) + return user_turn, image_pad_offsets def forward( self, @@ -1056,13 +1087,18 @@ def _decode_generation_result( tokenizer: Any, max_new_tokens: int, eos_token_ids: tuple[int, ...], + skip_special_tokens: bool = True, + decode_prefix_ids: list[int] | None = None, ) -> tuple[str, dict[str, bool | int]]: """Decode generated token IDs and expose basic generation metadata.""" - generated_only = generated_ids[0, input_len:] + generated_only = generated_ids[0, input_len:] # [T_gen] generated_token_count = int(generated_only.shape[0]) finished_with_eos = generated_token_count > 0 and int(generated_only[-1].item()) in eos_token_ids truncated = generated_token_count >= max_new_tokens and not finished_with_eos - text = tokenizer.decode(generated_only, skip_special_tokens=True) + decode_ids: torch.Tensor | list[int] = generated_only + if decode_prefix_ids is not None: + decode_ids = list(decode_prefix_ids) + generated_only.tolist() + text = tokenizer.decode(decode_ids, skip_special_tokens=skip_special_tokens) metadata: dict[str, bool | int] = { "generated_tokens": generated_token_count, "finished_with_eos": finished_with_eos, @@ -1079,8 +1115,10 @@ def _generate_from_prefix_embeddings( do_sample: bool, temperature: float, eos_token_ids: tuple[int, ...], + suppress_token_ids: tuple[int, ...] | None = None, ) -> torch.Tensor: """Autoregressively generate from a caller-provided embedding prefix.""" + resolved_suppress_token_ids = self.spec.suppress_token_ids if suppress_token_ids is None else suppress_token_ids if self.spec.supports_inputs_embeds_generate: eos_token_id: int | list[int] = eos_token_ids[0] if len(eos_token_ids) == 1 else list(eos_token_ids) generate_kwargs = dict( @@ -1092,8 +1130,8 @@ def _generate_from_prefix_embeddings( pad_token_id=self.spec.pad_token_id, eos_token_id=eos_token_id, ) - if self.spec.suppress_token_ids: - generate_kwargs["suppress_tokens"] = list(self.spec.suppress_token_ids) + if resolved_suppress_token_ids: + generate_kwargs["suppress_tokens"] = list(resolved_suppress_token_ids) if do_sample: generate_kwargs["temperature"] = temperature generate_kwargs["top_p"] = 0.8 @@ -1128,7 +1166,7 @@ def _generate_from_prefix_embeddings( temperature=temperature, top_p=0.8 if do_sample else 1.0, top_k=20 if do_sample else 0, - suppress_token_ids=self.spec.suppress_token_ids, + suppress_token_ids=resolved_suppress_token_ids, ) generated_tokens.append(next_token.unsqueeze(1)) @@ -1250,13 +1288,17 @@ def generate_answer( max_new_tokens: int = 512, do_sample: bool = True, temperature: float = 0.7, + image_token_counts_by_visual: list[int] | None = None, + thinking_mode: str = VQA_THINKING_MODE_OFF, + reasoning_suffix: str = VQA_REASONING_SUFFIX, + decode_skip_special_tokens: bool | None = None, return_metadata: bool = False, ) -> str | tuple[str, dict[str, bool | int]]: """Generate an answer to a question about an image. Qwen3 uses its native chat template. Nemotron uses its own native - chat template with ``<|im_start|>``, ``<|im_end|>``, and the - ``\n`` no-thinking prefix. + chat template with ``<|im_start|>``, ``<|im_end|>``, and an empty + ```` no-thinking prefix. Args: image_feats_tensor: [N, encoder_dim] features for ONE image. @@ -1265,11 +1307,30 @@ def generate_answer( max_new_tokens: Maximum tokens to generate. do_sample: Whether to sample. temperature: Sampling temperature if do_sample=True. + image_token_counts_by_visual: Merged visual-token counts for each visual block. + thinking_mode: ``off`` forces direct-answer mode, ``on`` requests a + visible thinking block, and ``raw`` leaves the assistant turn unconstrained. + reasoning_suffix: User-turn suffix appended in Qwen ``on`` mode. Nemotron + follows Dense VL generation and uses only the assistant think prefix. + decode_skip_special_tokens: Whether tokenizer special tokens are hidden + in the returned answer text. When unset, ``off`` keeps the + legacy skip-special decode and ``on``/``raw`` preserve tags for + answer postprocessing. return_metadata: Whether to also return generation metadata. Returns: Generated answer string, or ``(answer, metadata)`` when requested. """ + resolved_thinking_mode = normalize_vqa_thinking_mode(thinking_mode) + resolved_decode_skip_special_tokens = ( + resolved_thinking_mode == VQA_THINKING_MODE_OFF + if decode_skip_special_tokens is None + else bool(decode_skip_special_tokens) + ) + prompt_question = question + if resolved_thinking_mode == VQA_THINKING_MODE_ON and self.spec.family != NEMOTRON_2B_SPEC.family: + prompt_question = _append_vqa_reasoning_suffix(question, reasoning_suffix) + # Spatial merge + position embeddings (same as generate_caption) if self.spatial_merger is not None: image_features, coords, _ = self.spatial_merger(image_feats_tensor, image_coords) @@ -1287,14 +1348,24 @@ def generate_answer( image_features = self.image_pos_embed(image_features, coords_2d) num_image_tokens = len(image_features) + if image_token_counts_by_visual is None: + visual_token_counts = [num_image_tokens] + else: + visual_token_counts = [int(token_count) for token_count in image_token_counts_by_visual] + if sum(visual_token_counts) != num_image_tokens: + raise ValueError( + f"Visual token counts {visual_token_counts!r} do not sum to merged feature count " + f"{num_image_tokens}." + ) device = image_features.device # Lazy-load tokenizer tok = self._ensure_caption_tokenizer() - def _encode(s): + def _encode(s: str) -> list[int]: return tok.encode(s, add_special_tokens=False) + answer_decode_prefix_ids: list[int] | None = None if self.spec.family == QWEN3_SPEC.family: system_turn = ( [QWEN3_IM_START_TOKEN_ID] @@ -1302,26 +1373,34 @@ def _encode(s): + [QWEN3_IM_END_TOKEN_ID] + _encode("\n") ) - user_turn, user_image_pad_offset = self._build_vqa_user_turn( + user_turn, user_image_pad_offsets = self._build_vqa_user_turn_with_visual_blocks( tokenizer=tok, - question=question, - num_image_tokens=num_image_tokens, + question=prompt_question, + image_token_counts_by_visual=visual_token_counts, ) - # Empty think block signals Qwen3 non-thinking mode - no_think = [QWEN3_THINK_START_TOKEN_ID] + _encode("\n\n") + [QWEN3_THINK_END_TOKEN_ID] + _encode("\n\n") - asst_prefix = [QWEN3_IM_START_TOKEN_ID] + _encode("assistant\n") + no_think + asst_prefix = [QWEN3_IM_START_TOKEN_ID] + _encode("assistant\n") + if resolved_thinking_mode == VQA_THINKING_MODE_OFF: + # Empty think block signals Qwen3 non-thinking mode. + no_think = [QWEN3_THINK_START_TOKEN_ID] + _encode("\n\n") + [QWEN3_THINK_END_TOKEN_ID] + _encode("\n\n") + asst_prefix += no_think eos_token_ids = self._get_eos_token_ids() elif self.spec.family == NEMOTRON_2B_SPEC.family: im_start_id = _get_required_token_id(tok, NEMOTRON_2B_IM_START_TOKEN) im_end_id = _get_required_token_id(tok, NEMOTRON_2B_IM_END_TOKEN) + think_id = _get_required_token_id(tok, NEMOTRON_2B_THINK_START_TOKEN) end_think_id = _get_required_token_id(tok, NEMOTRON_2B_THINK_END_TOKEN) system_turn = [im_start_id] + _encode("system\nYou are a helpful assistant.") + [im_end_id] + _encode("\n") - user_turn, user_image_pad_offset = self._build_vqa_user_turn( + user_turn, user_image_pad_offsets = self._build_vqa_user_turn_with_visual_blocks( tokenizer=tok, - question=question, - num_image_tokens=num_image_tokens, + question=prompt_question, + image_token_counts_by_visual=visual_token_counts, ) - asst_prefix = [im_start_id] + _encode("assistant\n") + [end_think_id] + _encode("\n") + asst_prefix = [im_start_id] + _encode("assistant\n") + if resolved_thinking_mode == VQA_THINKING_MODE_OFF: + asst_prefix += [think_id, end_think_id] + elif resolved_thinking_mode == VQA_THINKING_MODE_ON: + answer_decode_prefix_ids = [think_id] + _encode("\n") + asst_prefix += answer_decode_prefix_ids eos_token_ids = self._get_eos_token_ids() else: raise NotImplementedError( @@ -1336,9 +1415,15 @@ def _encode(s): vision_mask = (input_ids != self.spec.vision_pad_token_id).to(dtype=text_embeds.dtype) text_embeds = text_embeds * vision_mask[:, :, None] - # Image features start inside the multimodal user turn at the first image_pad token. - ip_start = len(system_turn) + user_image_pad_offset - text_embeds[0, ip_start : ip_start + num_image_tokens] = image_features.to(text_embeds.dtype) + # Image features start inside each multimodal user-turn image_pad block. + feature_start = 0 + for user_image_pad_offset, visual_token_count in zip(user_image_pad_offsets, visual_token_counts, strict=True): + ip_start = len(system_turn) + user_image_pad_offset + feature_end = feature_start + visual_token_count + text_embeds[0, ip_start : ip_start + visual_token_count] = image_features[feature_start:feature_end].to( + text_embeds.dtype + ) + feature_start = feature_end # Generate # Qwen3 best practice: do NOT use greedy decoding — causes repetitions. @@ -1351,6 +1436,9 @@ def _encode(s): text_embeds=text_embeds, temperature=temperature, eos_token_ids=eos_token_ids, + suppress_token_ids=( + self.spec.suppress_token_ids if resolved_thinking_mode == VQA_THINKING_MODE_OFF else () + ), ) answer, metadata = self._decode_generation_result( @@ -1359,6 +1447,8 @@ def _encode(s): tokenizer=tok, max_new_tokens=max_new_tokens, eos_token_ids=eos_token_ids, + skip_special_tokens=resolved_decode_skip_special_tokens, + decode_prefix_ids=answer_decode_prefix_ids, ) if return_metadata: return answer, metadata diff --git a/cosmos_framework/model/tokenizer/utils/vlm_prompt_format.py b/cosmos_framework/model/tokenizer/utils/vlm_prompt_format.py new file mode 100644 index 00000000..1eefda37 --- /dev/null +++ b/cosmos_framework/model/tokenizer/utils/vlm_prompt_format.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Prompt-format helpers shared by tokenizer VLM training and generation.""" + +DENSEVL_ADD_VISION_ID_PREFIX_BY_MEDIA_TYPE: dict[str, str] = { + "image": "Picture", + "video": "Video", +} + + +def densevl_add_vision_id_text(media_type: str, one_based_index: int) -> str: + """Render the text DenseVL/Qwen emits for ``add_vision_id=True``.""" + normalized_media_type = str(media_type).strip().lower() + prefix = DENSEVL_ADD_VISION_ID_PREFIX_BY_MEDIA_TYPE.get(normalized_media_type) + if prefix is None: + raise ValueError( + f"Unsupported DenseVL add_vision_id media_type={media_type!r}; " + f"expected one of {sorted(DENSEVL_ADD_VISION_ID_PREFIX_BY_MEDIA_TYPE)}." + ) + if one_based_index < 1: + raise ValueError(f"DenseVL add_vision_id index must be one-based and positive, got {one_based_index}.") + return f"{prefix} {one_based_index}: " diff --git a/cosmos_framework/model/vfm/algorithm/loss/__init__.py b/cosmos_framework/model/vfm/algorithm/loss/__init__.py index cdad7942..4a67ce95 100644 --- a/cosmos_framework/model/vfm/algorithm/loss/__init__.py +++ b/cosmos_framework/model/vfm/algorithm/loss/__init__.py @@ -5,9 +5,9 @@ __all__: list[str] = [] -from cosmos_framework.model.vfm.algorithm.loss.cross_entropy import cross_entropy_loss +from cosmos_framework.model.vfm.algorithm.loss.cross_entropy import cross_entropy_loss, weighted_cross_entropy_loss -__all__ += ["cross_entropy_loss"] +__all__ += ["cross_entropy_loss", "weighted_cross_entropy_loss"] from cosmos_framework.model.vfm.algorithm.loss.load_balancing import compute_load_balancing_loss diff --git a/cosmos_framework/model/vfm/algorithm/loss/cross_entropy.py b/cosmos_framework/model/vfm/algorithm/loss/cross_entropy.py index dea0bf7d..9e3c9b5b 100644 --- a/cosmos_framework/model/vfm/algorithm/loss/cross_entropy.py +++ b/cosmos_framework/model/vfm/algorithm/loss/cross_entropy.py @@ -28,21 +28,19 @@ from __future__ import annotations -from typing import Optional - import torch import torch.distributed as dist import torch.nn.functional as F -from cosmos_framework.utils.vfm.vlm.constant import IGNORE_INDEX +from cosmos_framework.utils.vfm.reasoner.constant import IGNORE_INDEX def cross_entropy_loss( logits: torch.Tensor, labels: torch.Tensor, loss_scaling_factor: float = 1.0, - dp_group: Optional[dist.ProcessGroup] = None, - cp_group: Optional[dist.ProcessGroup] = None, + dp_group: dist.ProcessGroup | None = None, + cp_group: dist.ProcessGroup | None = None, ignore_index: int = IGNORE_INDEX, ) -> torch.Tensor: """Next-token-prediction CE loss with DP/CP group reduction. @@ -99,3 +97,63 @@ def cross_entropy_loss( loss = per_token_loss.sum() / (n_valid_tokens + 1e-8) * (num_dp_workers * loss_scaling_factor) return loss + + +def weighted_cross_entropy_loss( + logits: torch.Tensor, + labels: torch.Tensor, + exponent: float, + loss_scaling_factor: float = 1.0, + dp_group: dist.ProcessGroup | None = None, + cp_group: dist.ProcessGroup | None = None, + ignore_index: int = IGNORE_INDEX, +) -> torch.Tensor: + """Next-token-prediction CE loss interpolated between per-token and per-sample reductions. + + Matches ``cosmos_rl.policy.trainer.llm_trainer.sft_trainer.async_safe_weighted_ce`` + for the non-packed, non-CP VLM path. + + Args: + logits: [B,T,V] float tensor, raw model output before softmax. + labels: [B,T] long tensor, ground-truth token ids. + exponent: 0 gives per-token loss, 1 gives per-sample loss, values in + between interpolate by valid-token-count weight. + loss_scaling_factor: scalar multiplied into the returned loss. + dp_group: Ignored for weighted CE. Kept for call-site parity with + ``cross_entropy_loss``; normalization uses the default process group + to match cosmos-rl. + cp_group: Context-parallel group. Weighted CE does not support CP. + ignore_index: label value to exclude. + + Returns: + Scalar loss tensor. + """ + if cp_group is not None and cp_group.size() > 1: + raise AssertionError("weighted_cross_entropy_loss does not support CP") + del dp_group + + batch_size = labels.shape[0] + shifted_logits = logits[:, :-1].contiguous().view(-1, logits.size(-1)) # [B*(T-1),V] + shifted_labels = labels[:, 1:].contiguous().view(-1) # [B*(T-1)] + + per_token_loss = F.cross_entropy( + shifted_logits.float(), # [B*(T-1),V] + shifted_labels, # [B*(T-1)] + ignore_index=ignore_index, + reduction="none", + ).view(batch_size, -1) # [B,T-1] + valid_mask = (shifted_labels.view(batch_size, -1) != ignore_index).float() # [B,T-1] + valid_counts = valid_mask.sum(dim=1) # [B] + has_valid = (valid_counts > 0).float() # [B] + + sample_losses = (per_token_loss * valid_mask).sum(dim=1) / valid_counts.clamp(min=1).pow(exponent) # [B] + local_loss_sum = (sample_losses * has_valid).sum() # [] + local_exp_weight_sum = (valid_counts.pow(1 - exponent) * has_valid).sum() # [] + + num_dp_workers = 1 + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(local_exp_weight_sum, op=dist.ReduceOp.SUM) + num_dp_workers = dist.get_world_size() + + loss = local_loss_sum / local_exp_weight_sum.clamp(min=1) * (num_dp_workers * loss_scaling_factor) # [] + return loss diff --git a/cosmos_framework/model/vfm/algorithm/loss/load_balancing.py b/cosmos_framework/model/vfm/algorithm/loss/load_balancing.py index 1853e9c6..206d1a39 100644 --- a/cosmos_framework/model/vfm/algorithm/loss/load_balancing.py +++ b/cosmos_framework/model/vfm/algorithm/loss/load_balancing.py @@ -5,7 +5,7 @@ from torch.distributed.tensor import DTensor, Partial from torch.distributed.tensor.device_mesh import DeviceMesh -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.qwen3_vl_moe import LBLMetadata +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.qwen3_vl_moe import LBLMetadata def compute_load_balancing_loss( diff --git a/cosmos_framework/model/vfm/hf_model.py b/cosmos_framework/model/vfm/hf_model.py index a39bbe48..7c3fa839 100644 --- a/cosmos_framework/model/vfm/hf_model.py +++ b/cosmos_framework/model/vfm/hf_model.py @@ -43,7 +43,7 @@ def _tensor_names_to_skip_for(model_type: str) -> list[str]: matched model keys absent from the checkpoint). Registered VLMs (see - cosmos_framework/configs/base/vlm/defaults/vlm_policy.py): + cosmos_framework/configs/base/reasoner/defaults/vlm_policy.py): - Qwen3-VL dense (2B/4B/8B/32B): no skips needed. - NemotronH_Nano_VL_V2 (nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16): RADIO backbone buffers — initialized by the module, not from ckpt. @@ -314,6 +314,17 @@ def load_weights( "pad_token_id", "ignore_index", "collated", + # content_tokens: non-pad token count emitted by custom_collate for the + # VLMTokensPerSec throughput callback; telemetry only, not a forward arg. + "content_tokens", + # Extended packing telemetry emitted by custom_collate (supervision density, + # l_max, attention-quadratic waste) for VLMTokensPerSec; telemetry only. + "supervised_tokens", + "seq_max_len", + "sum_len_sq", + # predicted_runtime_ms: the FLOP packer's per-step runtime estimate, surfaced by + # custom_collate for the VLMTokensPerSec realized-vs-predicted calibration; telemetry only. + "predicted_runtime_ms", "raw_image", "raw_video", # image_sizes is collected by collate_fn but is NOT a Qwen3-VL forward arg diff --git a/cosmos_framework/model/vfm/mot/attention.py b/cosmos_framework/model/vfm/mot/attention.py index 5c55ddfa..babcdd94 100644 --- a/cosmos_framework/model/vfm/mot/attention.py +++ b/cosmos_framework/model/vfm/mot/attention.py @@ -89,10 +89,22 @@ def two_way_attention( packed_query_states: SequencePack, packed_key_states: SequencePack, packed_value_states: SequencePack, -) -> SequencePack: + packed_key_states_normalized: SequencePack | None = None, +): """ Performs two-way attention with causal and full attention. + + ``packed_key_states_normalized``: optional alternative K pack used for the generator's + full attention (gen→all). When provided, the generator attends to these keys + instead of ``packed_key_states``, allowing the und K tokens to be normalised for + the gen cross-attention path while keeping raw K tokens for the reasoner's own + causal self-attention. If ``None``, ``packed_key_states`` is used for both paths. """ + # For gen full-attention, use normed keys when provided, + # otherwise fall back to the standard packed keys. + packed_key_normalized = ( + packed_key_states_normalized if packed_key_states_normalized is not None else packed_key_states + ) causal_q, causal_q_offsets = get_causal_seq(packed_query_states) causal_k, causal_k_offsets = get_causal_seq(packed_key_states) @@ -121,7 +133,7 @@ def two_way_attention( full_res = attention( full_q.unsqueeze(0), # [1,N_full,heads,head_dim] - get_all_seq(packed_key_states).unsqueeze(0), # [1,N_all,heads,head_dim] + get_all_seq(packed_key_normalized).unsqueeze(0), # [1,N_all,heads,head_dim] normed und K for gen get_all_seq(packed_value_states).unsqueeze(0), # [1,N_all,heads,head_dim] cumulative_seqlen_Q=full_q_offsets, cumulative_seqlen_KV=sample_offsets, @@ -142,7 +154,8 @@ def three_way_attention( packed_value_states: SequencePack, natten_metadata: dict | None, attention_meta: SplitInfo | None = None, -) -> SequencePack: + packed_key_states_normalized: SequencePack | None = None, +): """ Performs three-way attention, with understanding and generations attentions fully decomposed, and allows sparsity / multi-dimensional masking in the generation tower. @@ -155,10 +168,23 @@ def three_way_attention( NOTE: the three-way decomposition is only done so we can handle sparsity in the gen tower, but a KEY assumption is that the "full" tokens all correspond to the same modality! We should be careful when extending this to beyond t2i and t2v. + + ``packed_key_states_normalized``: optional alternative K pack for the gen→und cross-attention + (``full_ca``). When provided, ``get_causal_seq(packed_key_states_normalized)`` supplies the und + K tokens seen by the generator, while ``get_causal_seq(packed_key_states)`` (raw und K) is + still used for the reasoner's own causal self-attention. If ``None``, both paths share + ``packed_key_states``. """ causal_q, causal_q_offsets = get_causal_seq(packed_query_states) causal_k, causal_k_offsets = get_causal_seq(packed_key_states) + + # For gen→und cross-attention use normed keys when provided, + # otherwise fall back to the standard causal keys. + if packed_key_states_normalized is not None: + causal_k_normalized, causal_k_normalized_offsets = get_causal_seq(packed_key_states_normalized) + else: + causal_k_normalized, causal_k_normalized_offsets = causal_k, causal_k_offsets causal_v, _ = get_causal_seq(packed_value_states) full_q, full_q_offsets = get_full_only_seq(packed_query_states) full_k, full_k_offsets = get_full_only_seq(packed_key_states) @@ -220,10 +246,10 @@ def three_way_attention( full_ca, full_ca_lse = attention( full_q.unsqueeze(0), # [1,N_full,heads,head_dim] - causal_k.unsqueeze(0), # [1,N_und,heads,head_dim] + causal_k_normalized.unsqueeze(0), # [1,N_und,heads,head_dim] normed und K for gen→und causal_v.unsqueeze(0), # [1,N_und,heads,head_dim] cumulative_seqlen_Q=full_q_offsets, - cumulative_seqlen_KV=causal_k_offsets, + cumulative_seqlen_KV=causal_k_normalized_offsets, max_seqlen_Q=packed_query_states["max_full_len"], max_seqlen_KV=packed_query_states["max_causal_len"], return_lse=True, @@ -323,7 +349,7 @@ def multi_control_two_way_attention( noisy_v = full_v_v[noisy_s:noisy_e] # [N_noisy, Hkv, D] def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - """Maskless attention using cosmos_framework attention() → [N_q, Hq*D].""" + """Maskless attention using cosmos_framework.model.attention() → [N_q, Hq*D].""" n_q, n_kv = q.shape[0], k.shape[0] seqlens_q = torch.tensor([n_q], dtype=torch.int32, device=q.device) seqlens_kv = torch.tensor([n_kv], dtype=torch.int32, device=k.device) @@ -381,6 +407,7 @@ def dispatch_attention( attention_mask: SplitInfo, natten_metadata: dict | None = None, memory_value: MemoryValue | None = None, + packed_key_states_normalized: SequencePack | None = None, ) -> tuple[SequencePack, KVToStore | None]: assert memory_value is None, "Base dispatch_attention does not handle MemoryValue" if isinstance(attention_mask, SplitInfo) and attention_mask.control_stream_token_ranges is not None: @@ -397,9 +424,15 @@ def dispatch_attention( packed_value_states, natten_metadata=natten_metadata, attention_meta=attention_mask, + packed_key_states_normalized=packed_key_states_normalized, ) elif isinstance(attention_mask, SplitInfo): - output = two_way_attention(packed_query_states, packed_key_states, packed_value_states) + output = two_way_attention( + packed_query_states, + packed_key_states, + packed_value_states, + packed_key_states_normalized=packed_key_states_normalized, + ) else: raise TypeError(f"Unsupported attention metadata: {type(attention_mask)}") return output, None diff --git a/cosmos_framework/model/vfm/mot/context_parallel_utils.py b/cosmos_framework/model/vfm/mot/context_parallel_utils.py index 6a6b3dd5..0195b835 100644 --- a/cosmos_framework/model/vfm/mot/context_parallel_utils.py +++ b/cosmos_framework/model/vfm/mot/context_parallel_utils.py @@ -276,6 +276,7 @@ def context_parallel_attention( attention_function: Callable, natten_metadata: dict | None = None, memory_value: MemoryValue | None = None, + packed_key_states_normalized: SequencePack | None = None, ) -> tuple[SequencePack, KVToStore | None]: """Ulysses-style context parallel attention for packed und+gen sequences. @@ -300,6 +301,12 @@ def context_parallel_attention( attention_function: Callable implementing the actual attention kernel. natten_metadata: Optional neighborhood attention metadata. memory_value: Optional memory value for KV-cache training / AR inference. + packed_key_states_normalized: Optional seq-sharded normed K pack (und tokens RMSNorm-ed + before RoPE) used for the gen→und cross-attention path. When provided, the + und portion is gathered/scattered through the same all-to-all as the regular K + and forwarded to ``attention_function`` as ``packed_key_states_normalized``. + The gen portion is shared with ``packed_key_states`` (no separate all-to-all + needed). Pass ``None`` (default) to skip and use raw K for all paths. Returns: (output_pack, kv_to_store): @@ -392,6 +399,19 @@ def context_parallel_attention( packed_key_states_ = from_mode_splits(k_und_seq, k_gen_seq, meta, is_sharded=False) packed_value_states_ = from_mode_splits(v_und_seq, v_gen_seq, meta, is_sharded=False) + # If a normed K pack is provided (und K-norm for gen→und cross-attn), apply the same + # all-to-all to the und portion. The gen portion is identical to k_gen_seq (already + # gathered above), so no second all-to-all is needed for it. + packed_key_states_normalized_: SequencePack | None = None + if packed_key_states_normalized is not None: + k_und_normalized_seq, _ = get_causal_seq(packed_key_states_normalized) # [text_shard_len,H_kv,head_dim] + if kv_head_repeats > 1: + k_und_normalized_seq = _repeat_kv_heads_for_cp(k_und_normalized_seq, kv_head_repeats) + k_und_normalized_seq = gather_seq_scatter_heads( + k_und_normalized_seq, seq_dim=0, head_dim=1, cp_mesh=cp_mesh + ) # [text_len,H_kv_local,head_dim] + packed_key_states_normalized_ = from_mode_splits(k_und_normalized_seq, k_gen_seq, meta, is_sharded=False) + # dispatch_attention returns (output, kv_to_store | None) attn_output_pack_hp, _inner_kv_to_store = attention_function( packed_query_states_, @@ -400,6 +420,7 @@ def context_parallel_attention( attention_mask, natten_metadata=natten_metadata, memory_value=memory_value, + packed_key_states_normalized=packed_key_states_normalized_, ) attn_output_und_hp = get_und_seq(attn_output_pack_hp) # [text_len,H_local,head_dim] diff --git a/cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py b/cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py index 1592f282..0a326996 100644 --- a/cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py +++ b/cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: OpenMDW-1.1 import math -from typing import List, Optional, Tuple +from typing import List, Tuple import torch from torch import nn @@ -15,11 +15,7 @@ get_context_parallel_sharded_sequence, ) from cosmos_framework.model.vfm.mot.domain_aware_linear import DomainAwareLinear -from cosmos_framework.model.vfm.mot.modeling_utils import ( - FlattenedSinCosPositionEmbedding, - TimestepEmbedder, - VideoRopePosition3DEmb, -) +from cosmos_framework.model.vfm.mot.modeling_utils import TimestepEmbedder from cosmos_framework.model.vfm.utils.memory import MemoryState from cosmos_framework.data.vfm.sequence_packing import ModalityData, PackedSequence from cosmos_framework.data.vfm.sequence_packing.natten import verify_natten_parameter_list @@ -35,13 +31,9 @@ def __init__( latent_patch_size=2, latent_downsample_factor=8, latent_channel_size=16, - position_embedding_type="3d_rope", max_latent_h=32, max_latent_w=32, max_latent_t=32, - rope_h_extrapolation_ratio=1.0, - rope_w_extrapolation_ratio=1.0, - rope_t_extrapolation_ratio=1.0, enable_fps_modulation=False, base_fps=24, vit_max_num_patch_per_side=70, @@ -69,13 +61,9 @@ def __init__( self.latent_patch_size = latent_patch_size self.latent_downsample_factor = latent_downsample_factor self.latent_channel_size = latent_channel_size - self.position_embedding_type = position_embedding_type self.max_latent_h = max_latent_h self.max_latent_w = max_latent_w self.max_latent_t = max_latent_t - self.rope_h_extrapolation_ratio = rope_h_extrapolation_ratio - self.rope_w_extrapolation_ratio = rope_w_extrapolation_ratio - self.rope_t_extrapolation_ratio = rope_t_extrapolation_ratio self.enable_fps_modulation = enable_fps_modulation self.base_fps = base_fps self.vit_max_num_patch_per_side = vit_max_num_patch_per_side @@ -160,57 +148,12 @@ def __init__(self, language_model, config: Cosmos3VFMNetworkConfig): self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size) self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim) - assert config.position_embedding_type in ["3d_rope", "flattened_sin_cos", "unified_3d_mrope"] - if config.position_embedding_type == "3d_rope": - self.latent_pos_embed = VideoRopePosition3DEmb( - head_dim=self.hidden_size, - len_h=self.max_latent_h, - len_w=self.max_latent_w, - len_t=self.max_latent_t, - h_extrapolation_ratio=config.rope_h_extrapolation_ratio, - w_extrapolation_ratio=config.rope_w_extrapolation_ratio, - t_extrapolation_ratio=config.rope_t_extrapolation_ratio, - enable_fps_modulation=config.enable_fps_modulation, # fps_modulation scales RoPE by fps. By default, disable FPS RoPE modulation. - base_fps=config.base_fps, - base_temporal_compression_factor=config.temporal_compression_factor_vision, - temporal_compression_factor=config.temporal_compression_factor_vision, - ) - elif config.position_embedding_type == "flattened_sin_cos": - self.latent_pos_embed = FlattenedSinCosPositionEmbedding( - max_latent_h=self.max_latent_h, max_latent_w=self.max_latent_w, hidden_size=self.hidden_size - ) - elif config.position_embedding_type == "unified_3d_mrope": - # No additive position embedding - position info is in 3D position IDs for attention - self.latent_pos_embed = None - else: - raise ValueError(f"Unknown position_embedding_type: {config.position_embedding_type!r}") - if config.action_gen: self.action_dim = config.action_dim self.num_embodiment_domains = config.num_embodiment_domains self.action2llm = DomainAwareLinear(self.action_dim, self.hidden_size, self.num_embodiment_domains) self.llm2action = DomainAwareLinear(self.hidden_size, self.action_dim, self.num_embodiment_domains) - if config.position_embedding_type == "3d_rope": - self.action_pos_embed = VideoRopePosition3DEmb( - head_dim=self.hidden_size, - len_h=1, - len_w=1, - len_t=self.max_latent_t * config.temporal_compression_factor_vision, - h_extrapolation_ratio=config.rope_h_extrapolation_ratio, - w_extrapolation_ratio=config.rope_w_extrapolation_ratio, - t_extrapolation_ratio=config.rope_t_extrapolation_ratio, - enable_fps_modulation=config.enable_fps_modulation, - base_fps=config.base_fps, - base_temporal_compression_factor=config.temporal_compression_factor_vision, # vision compression factor is used for base tps - temporal_compression_factor=config.temporal_compression_factor_action, # Action is at frame rate (no temporal compression) - ) - elif config.position_embedding_type == "unified_3d_mrope": - # No additive position embedding - position info is in 3D position IDs for attention - self.action_pos_embed = None - else: - raise ValueError(f"Unknown position_embedding_type: {config.position_embedding_type!r}") - self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size)) if config.sound_gen: @@ -235,9 +178,6 @@ def init_weights(self, buffer_device: torch.device | None): torch.nn.init.trunc_normal_(self.llm2vae.weight, std=std, a=-3 * std, b=3 * std) torch.nn.init.zeros_(self.llm2vae.bias) - if self.latent_pos_embed is not None: - self.latent_pos_embed._init_weights() - if self.config.action_gen: # DomainAwareLinear uses embeddings for weights, so we initialize them differently # action2llm: input_size=action_dim, output_size=hidden_size @@ -253,9 +193,6 @@ def init_weights(self, buffer_device: torch.device | None): std = 1.0 / math.sqrt(self.hidden_size) torch.nn.init.trunc_normal_(self.action_modality_embed, std=std, a=-3 * std, b=3 * std) - if self.action_pos_embed is not None: - self.action_pos_embed._init_weights() - if self.config.sound_gen: # sound2llm: input_size=sound_dim, output_size=hidden_size std = 1.0 / math.sqrt(self.sound_dim) @@ -279,6 +216,8 @@ def generate_reasoner_text( *, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, @@ -299,9 +238,11 @@ def generate_reasoner_text( prompts through this single entry point: pass ``pixel_values`` + ``image_grid_thw`` (and optionally ``attention_mask``) for image-conditioned prefill via the Qwen3-VL - visual encoder, or omit them for text-only prefill. Uses the - und-pathway weights (those WITHOUT the ``_moe_gen`` suffix) plus - ``embed_tokens`` / ``norm`` / ``lm_head``; the generation pathway + visual encoder, or omit them for text-only prefill. Video + conditioning is also supported via ``pixel_values_videos`` + + ``video_grid_thw``; the image and video pairs are mutually exclusive. + Uses the und-pathway weights (those WITHOUT the ``_moe_gen`` suffix) + plus ``embed_tokens`` / ``norm`` / ``lm_head``; the generation pathway and all VFM-level multimodal embedders / heads (``vae2llm``, ``llm2vae``, ``sound2llm``, etc.) are bypassed. @@ -330,6 +271,8 @@ def generate_reasoner_text( max_new_tokens=max_new_tokens, pixel_values=pixel_values, image_grid_thw=image_grid_thw, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, attention_mask=attention_mask, eos_token_id=eos_token_id, pad_token_id=pad_token_id, @@ -608,7 +551,6 @@ def _encode_vision( packed_seq: PackedSequence, packed_sequence: torch.Tensor, target_dtype: torch.dtype, - fps: Optional[torch.Tensor] = None, ) -> List[Tuple[int, int, int]] | None: """Project vision tokens and fill into packed_sequence. @@ -616,7 +558,6 @@ def _encode_vision( packed_seq: PackedSequence containing vision tokens and metadata. packed_sequence: The packed sequence tensor to fill vision embeddings into (modified in-place). target_dtype: Target dtype for embeddings (typically from text embedding). - fps: Optional FPS tensor for RoPE modulation. Returns: Original latent shapes before padding (for unpadding during decode), or None if no vision tokens. @@ -642,14 +583,6 @@ def _encode_vision( ) # packed_tokens_vision: [total_vision_patches,patch_latent_dim] packed_tokens_vision = self.vae2llm(packed_tokens_vision.to(target_dtype)) # [total_vision_patches,hidden_size] - # Add absolute position embedding only when NOT using unified 3D mRoPE - # (3D mRoPE provides positional information via rotary embeddings instead) - if self.latent_pos_embed is not None: - latent_token_pos_emb = self.latent_pos_embed(vision.token_shapes, fps=fps).to( - target_dtype - ) # [total_vision_patches,hidden_size] - packed_tokens_vision = packed_tokens_vision + latent_token_pos_emb # [total_vision_patches,hidden_size] - has_noisy_vision = vision.mse_loss_indexes.numel() > 0 if has_noisy_vision: @@ -742,7 +675,6 @@ def _encode_action( packed_seq: PackedSequence, packed_sequence: torch.Tensor, target_dtype: torch.dtype, - fps_action: Optional[torch.Tensor] = None, ) -> None: """Encode action tokens and fill into packed_sequence.""" if packed_seq.action is None or packed_seq.action.tokens is None: @@ -761,17 +693,6 @@ def _encode_action( ) packed_tokens_action = self.action2llm(packed_tokens_action, per_token_domain_id) - # Add additive position embedding only if not using unified_3d_mrope - if self.action_pos_embed is not None: - # VideoRopePosition3DEmb expects shapes as (t, h, w). For actions we use a 1x1 spatial grid. - action_shapes_3d = [(ts[0], 1, 1) for ts in action.token_shapes] - action_token_pos_emb = self.action_pos_embed( - action_shapes_3d, - fps=fps_action, - start_frame_offset=1, - ).to(target_dtype) # [B_action*T_action,hidden_size] - packed_tokens_action = packed_tokens_action + action_token_pos_emb # [B_action*T_action,hidden_size] - packed_tokens_action = packed_tokens_action + self.action_modality_embed.view( 1, -1 ) # [B_action*T_action,hidden_size] @@ -862,7 +783,6 @@ def _encode_sound( packed_seq: PackedSequence, packed_sequence: torch.Tensor, target_dtype: torch.dtype, - fps_sound: Optional[torch.Tensor] = None, ) -> None: """Encode sound tokens and fill into packed_sequence. @@ -870,7 +790,6 @@ def _encode_sound( packed_seq: PackedSequence containing sound tokens and metadata. packed_sequence: The packed sequence tensor to fill sound embeddings into (modified in-place). target_dtype: Target dtype for embeddings (typically from text embedding). - fps_sound: FPS tensor for RoPE modulation. Should be the sound latent rate (e.g., 25 Hz). """ if packed_seq.sound is None or packed_seq.sound.tokens is None: # No sound tokens in this batch @@ -888,9 +807,8 @@ def _encode_sound( ) # [total_sound_tokens,sound_dim] packed_tokens_sound = packed_tokens_sound.to(target_dtype) # [total_sound_tokens,sound_dim] - # Project sound tokens + modality embedding - # NOTE: Sound position info comes from m-RoPE position IDs in the attention layers. - # No additive position embedding is used (unlike legacy video which keeps one for backward compat). + # Project sound tokens + modality embedding. Position info comes from + # mRoPE position IDs in the attention layers. packed_tokens_sound = ( self.sound2llm(packed_tokens_sound) + self.sound_modality_embed ) # [total_sound_tokens,hidden_size] @@ -964,9 +882,6 @@ def _decode_sound( def forward( self, packed_seq: PackedSequence, - fps_vision: Optional[torch.Tensor] = None, - fps_action: Optional[torch.Tensor] = None, - fps_sound: Optional[torch.Tensor] = None, memory: MemoryState | None = None, ) -> dict: """ @@ -975,9 +890,6 @@ def forward( Args: packed_seq: PackedSequence containing all packed tensors and metadata. See PackedSequence dataclass for field details. - fps_vision: Optional FPS tensor for vision RoPE modulation. - fps_action: Optional FPS tensor for action RoPE modulation. - fps_sound: Optional FPS tensor for sound RoPE modulation (e.g., sound_latent_fps=25). memory: Optional MemoryState for persistent KV-cache memory (AR inference or rolling-KV-cache training). Built by ``OmniMoTModel.build_memory_state()``. @@ -1000,15 +912,15 @@ def forward( # encode vision tokens original_latent_shapes: List[Tuple[int, int, int]] | None = None if self.config.vision_gen: - original_latent_shapes = self._encode_vision(packed_seq, packed_sequence, target_dtype, fps_vision) + original_latent_shapes = self._encode_vision(packed_seq, packed_sequence, target_dtype) # encode action tokens if self.config.action_gen: - self._encode_action(packed_seq, packed_sequence, target_dtype, fps_action) + self._encode_action(packed_seq, packed_sequence, target_dtype) # encode sound tokens if self.config.sound_gen: - self._encode_sound(packed_seq, packed_sequence, target_dtype, fps_sound) + self._encode_sound(packed_seq, packed_sequence, target_dtype) assert packed_seq.attn_modes is not None assert packed_seq.split_lens is not None diff --git a/cosmos_framework/model/vfm/mot/modeling_utils.py b/cosmos_framework/model/vfm/mot/modeling_utils.py index bcd42d91..2d543a86 100644 --- a/cosmos_framework/model/vfm/mot/modeling_utils.py +++ b/cosmos_framework/model/vfm/mot/modeling_utils.py @@ -2,13 +2,9 @@ # SPDX-License-Identifier: OpenMDW-1.1 import math -from typing import Optional -import numpy as np import torch -from einops import rearrange, repeat from torch import nn -from torch.distributed import ProcessGroup from transformers.activations import ACT2FN from cosmos_framework.data.vfm.sequence_packing import ModalityData @@ -24,307 +20,6 @@ def has_noisy_tokens(modality_data: ModalityData | None) -> bool: ) -# -------------------------------------------------------- -# 2D sine-cosine position embedding (flattened) -# References: -# DiT: https://github.com/facebookresearch/DiT/blob/main/models.py -# -------------------------------------------------------- -def get_2d_sincos_pos_embed( - embed_dim: int, grid_size_h: int, grid_size_w: int, cls_token: bool = False, extra_tokens: int = 0 -) -> np.ndarray: - grid_h = np.arange(grid_size_h, dtype=np.float32) - grid_w = np.arange(grid_size_w, dtype=np.float32) - grid = np.meshgrid(grid_w, grid_h) # here w goes first - grid = np.stack(grid, axis=0) - - grid = grid.reshape([2, 1, grid_size_h, grid_size_w]) - pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) - if cls_token and extra_tokens > 0: - pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0) - return pos_embed - - -def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray: - assert embed_dim % 2 == 0 - - # use half of dimensions to encode grid_h - emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # [H*W,D/2] - emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # [H*W,D/2] - - emb = np.concatenate([emb_h, emb_w], axis=1) # [H*W,D] - return emb - - -def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray: - """ - embed_dim: output dimension for each position - pos: a list of positions to be encoded: size [M] - out: [M,D] - """ - assert embed_dim % 2 == 0 - omega = np.arange(embed_dim // 2, dtype=np.float64) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega # [D/2] - - pos = pos.reshape(-1) # [M] - out = np.einsum("m,d->md", pos, omega) # [M,D/2], outer product - - emb_sin = np.sin(out) # [M,D/2] - emb_cos = np.cos(out) # [M,D/2] - - emb = np.concatenate([emb_sin, emb_cos], axis=1) # [M,D] - return emb - - -class FlattenedSinCosPositionEmbedding(nn.Module): - # This module creates a flattened sin-cos position embedding for a given number of patches per side. - # Indices are created for 2D array and flattened into 1D array. - - def __init__(self, max_latent_h: int, max_latent_w: int, hidden_size: int, interpolate_pos: bool = False): - super().__init__() - self.max_latent_h = max_latent_h - self.max_latent_w = max_latent_w - self.hidden_size = hidden_size - self.interpolate_pos = interpolate_pos - self.pos_embed = nn.Parameter(torch.zeros(max_latent_h * max_latent_w, hidden_size), requires_grad=False) - self._init_weights() - - def _get_flattened_position_ids_extrapolate(self, latent_dim_h: int, latent_dim_w: int) -> torch.Tensor: - coords_h = torch.arange(0, latent_dim_h) # [H] - coords_w = torch.arange(0, latent_dim_w) # [W] - pos_ids = (coords_h[:, None] * self.max_latent_w + coords_w).flatten() # [H*W] - return pos_ids - - def _get_flattened_position_ids_interpolate(self, latent_dim_h: int, latent_dim_w: int) -> torch.Tensor: - boundaries = torch.arange(1 / self.max_latent_w, 1.0, 1 / self.max_latent_w) # [max_latent_w-1] - fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / latent_dim_h) # [H] - fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / latent_dim_w) # [W] - bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True) # [H] - bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True) # [W] - pos_ids = (bucket_coords_h[:, None] * self.max_latent_w + bucket_coords_w).flatten() # [H*W] - return pos_ids - - def _create_flattened_position_ids_packed(self, token_shapes_vision: list[tuple[int, int]]) -> torch.Tensor: - flattened_position_ids = [] - for t, h, w in token_shapes_vision: - if self.interpolate_pos: - flattened_position_ids.append(self._get_flattened_position_ids_interpolate(h, w)) # [H*W] - else: - flattened_position_ids.append(self._get_flattened_position_ids_extrapolate(h, w)) # [H*W] - flattened_position_ids_packed = torch.cat(flattened_position_ids, dim=0) # [N_vision] - return flattened_position_ids_packed - - def _init_weights(self): - # Initialize (and freeze) pos_embed by sin-cos embedding: - pos_embed = get_2d_sincos_pos_embed( - embed_dim=self.hidden_size, grid_size_h=self.max_latent_h, grid_size_w=self.max_latent_w - ) - self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float()) - - def forward(self, token_shapes_vision: list[tuple[int, int]], fps: Optional[torch.Tensor] = None) -> torch.Tensor: - # First create 2D index array - flattened_position_ids_packed = self._create_flattened_position_ids_packed(token_shapes_vision) # [N_vision] - return self.pos_embed[flattened_position_ids_packed] # [N_vision,hidden_size] - - -# -------------------------------------------------------- -# 2D / 3D RoPE Position Embedding -# -------------------------------------------------------- - - -class VideoRopePosition3DEmb(nn.Module): - def __init__( - self, - *, # enforce keyword arguments - head_dim: int, - len_h: int, - len_w: int, - len_t: int, - base_fps: int = 24, - base_temporal_compression_factor: int = 4, - temporal_compression_factor: int = 4, - h_extrapolation_ratio: float = 1.0, - w_extrapolation_ratio: float = 1.0, - t_extrapolation_ratio: float = 1.0, - enable_fps_modulation: bool = False, - **kwargs, # used for compatibility with other positional embeddings; unused in this class - ): - del kwargs - super().__init__() - self.base_tps = base_fps / base_temporal_compression_factor - self.temporal_compression_factor = temporal_compression_factor - self.max_h = len_h - self.max_w = len_w - self.max_t = len_t - self.enable_fps_modulation = enable_fps_modulation - dim = head_dim - dim_h = dim // 6 * 2 - dim_w = dim_h - dim_t = dim - 2 * dim_h - assert dim == dim_h + dim_w + dim_t, f"bad dim: {dim} != {dim_h} + {dim_w} + {dim_t}" - self.register_buffer( - "dim_spatial_range", - torch.arange(0, dim_h, 2)[: (dim_h // 2)].float() / dim_h, - persistent=True, - ) - self.register_buffer( - "dim_temporal_range", - torch.arange(0, dim_t, 2)[: (dim_t // 2)].float() / dim_t, - persistent=True, - ) - self._dim_h = dim_h - self._dim_t = dim_t - - self.h_ntk_factor = h_extrapolation_ratio ** (dim_h / (dim_h - 2)) - self.w_ntk_factor = w_extrapolation_ratio ** (dim_w / (dim_w - 2)) - self.t_ntk_factor = t_extrapolation_ratio ** (dim_t / (dim_t - 2)) - self._init_weights() - - def _init_weights(self) -> None: - dim_h = self._dim_h - dim_t = self._dim_t - - self.dim_spatial_range = ( - torch.arange(0, dim_h, 2)[: (dim_h // 2)].float().to(self.dim_spatial_range.device) / dim_h - ) - self.dim_temporal_range = ( - torch.arange(0, dim_t, 2)[: (dim_t // 2)].float().to(self.dim_spatial_range.device) / dim_t - ) - - def enable_context_parallel(self, process_group: ProcessGroup): - pass - - def disable_context_parallel(self): - pass - - def generate_embeddings( - self, - latent_shape: torch.Size, - input_fps: Optional[torch.Tensor] = None, - h_ntk_factor: Optional[float] = None, - w_ntk_factor: Optional[float] = None, - t_ntk_factor: Optional[float] = None, - start_frame_offset: int = 0, - ): - """ - Generate embeddings for the given input size. - - Args: - latent_shape (torch.Size): Input tensor size (Batch, Time, Height, Width). - input_fps (Optional[torch.Tensor], optional): Frames per second. Defaults to None. - h_ntk_factor (Optional[float], optional): Height NTK factor. If None, uses self.h_ntk_factor. - w_ntk_factor (Optional[float], optional): Width NTK factor. If None, uses self.w_ntk_factor. - t_ntk_factor (Optional[float], optional): Time NTK factor. If None, uses self.t_ntk_factor. - start_frame_offset (int, optional): Offset for frame indices. Use 1 for action embeddings - so that action frame indices start at 1 instead of 0. Defaults to 0. - - Returns: - Not specified in the original code snippet. - """ - if input_fps is not None: - tps = input_fps / self.temporal_compression_factor - else: - tps = None - - h_ntk_factor = h_ntk_factor if h_ntk_factor is not None else self.h_ntk_factor - w_ntk_factor = w_ntk_factor if w_ntk_factor is not None else self.w_ntk_factor - t_ntk_factor = t_ntk_factor if t_ntk_factor is not None else self.t_ntk_factor - assert h_ntk_factor is not None and w_ntk_factor is not None and t_ntk_factor is not None - - h_theta = 10000.0 * h_ntk_factor - w_theta = 10000.0 * w_ntk_factor - t_theta = 10000.0 * t_ntk_factor - - h_spatial_freqs = 1.0 / (h_theta ** self.dim_spatial_range.float()) # [dim_h/2] - w_spatial_freqs = 1.0 / (w_theta ** self.dim_spatial_range.float()) # [dim_w/2] - temporal_freqs = 1.0 / (t_theta ** self.dim_temporal_range.float()) # [dim_t/2] - - B, T, H, W = latent_shape - assert H <= self.max_h and W <= self.max_w, ( - f"Input dimensions (H={H}, W={W}) exceed the maximum dimensions (max_h={self.max_h}, max_w={self.max_w})" - ) - - # Re-allocate buffer if current video needs more indices than what we have for self.seq - # Only rellocate when needed. - max_needed = max(T, H, W) - seq = torch.arange(max_needed, device=self.dim_spatial_range.device, dtype=torch.float) - - half_emb_h = torch.outer(seq[:H], h_spatial_freqs) # [H,dim_h/2] - half_emb_w = torch.outer(seq[:W], w_spatial_freqs) # [W,dim_w/2] - - # Frame indices for the embedding (always 0, 1, 2, ...) - frame_indices = seq[:T] # [T] - - if self.enable_fps_modulation: - uniform_tps = tps is None or tps.shape == (1,) - assert uniform_tps or B == 1 or T == 1, ( - "For video batch, B should be 1 for non-uniform fps. For image batch, T should be 1." - ) - - # apply sequence scaling in temporal dimension - if tps is None: # image case - assert T == 1, "T should be 1 for image batch." - half_emb_t = torch.outer(frame_indices, temporal_freqs) # [T,dim_t/2] - else: - # Calculate scaled time indices - # Apply start_frame_offset to the time calculation (not frame indices) - # This allows one to manipulate the start frame index of embeddings for cross-modality alignment. - scaled_time = (frame_indices + start_frame_offset) / tps[:1] * self.base_tps # [T] - half_emb_t = torch.outer(scaled_time, temporal_freqs) # [T,dim_t/2] - else: - half_emb_t = torch.outer(frame_indices, temporal_freqs) # [T,dim_t/2] - - rope_embed = torch.cat( - [ - repeat(half_emb_t, "t d -> t h w d", h=H, w=W), # [T,H,W,dim_t/2] - repeat(half_emb_h, "h d -> t h w d", t=T, w=W), # [T,H,W,dim_h/2] - repeat(half_emb_w, "w d -> t h w d", t=T, h=H), # [T,H,W,dim_w/2] - ] - * 2, - dim=-1, - ) # [T,H,W,head_dim] - - return rearrange(rope_embed, "t h w d -> (t h w) d").float() # [T*H*W,head_dim] - - def forward( - self, - token_shapes_vision: list[tuple[int, int, int]], - fps: Optional[torch.Tensor] = None, - start_frame_offset: int = 0, - ) -> torch.Tensor: - """ - With CP, the function assume that the input tensor is already split. - It delegates the embedding generation to generate_embeddings function. - - Args: - token_shapes_vision: List of (t, h, w) tuples for each latent. - fps: Frames per second tensor. - start_frame_offset: Offset for frame indices. Use 1 for action embeddings - so that action frame indices start at 1 instead of 0. Defaults to 0. - """ - - embeddings_packed = [] - for i, latent_shape in enumerate(token_shapes_vision): - # latent_shape: (t, h, w) - shape = (1, latent_shape[0], latent_shape[1], latent_shape[2]) - - # Extract FPS for this specific video - video_fps = None - if fps is not None: - assert i < fps.shape[0], f"Index {i} out of bounds for fps tensor of shape {fps.shape}" - video_fps = fps[i : i + 1] - - embeddings = self.generate_embeddings(shape, input_fps=video_fps, start_frame_offset=start_frame_offset) - embeddings_packed.append(embeddings) - - embeddings_packed = torch.cat(embeddings_packed, dim=0) # [N_vision,head_dim] - return embeddings_packed - - @property - def seq_dim(self): - return 0 - - # -------------------------------------------------------- # TimestepEmbedder # Reference: diff --git a/cosmos_framework/model/vfm/mot/parallelize_unified_mot.py b/cosmos_framework/model/vfm/mot/parallelize_unified_mot.py index a4601d18..a011e62f 100644 --- a/cosmos_framework/model/vfm/mot/parallelize_unified_mot.py +++ b/cosmos_framework/model/vfm/mot/parallelize_unified_mot.py @@ -75,6 +75,7 @@ def forward( attention_mask: SplitInfo, natten_metadata: dict | None = None, memory_value: MemoryValue | None = None, + packed_key_states_normalized: SequencePack | None = None, ) -> tuple[SequencePack, KVToStore | None]: if memory_value is not None and not memory_value.supports_context_parallel_attention: raise ValueError("Context-parallel doesn't work when training with a KV-cache.") @@ -88,6 +89,7 @@ def forward( attention_function=self.wrapped_dispatch, natten_metadata=natten_metadata, memory_value=memory_value, + packed_key_states_normalized=packed_key_states_normalized, ) @@ -183,6 +185,7 @@ def forward( attention_mask: SplitInfo, natten_metadata: dict | None = None, memory_value: MemoryValue | None = None, + packed_key_states_normalized: SequencePack | None = None, ) -> tuple[SequencePack, KVToStore | None]: if memory_value is None or getattr(memory_value, "frame_idx", 0) <= 0: return self.wrapped_dispatch( @@ -192,6 +195,7 @@ def forward( attention_mask, natten_metadata=natten_metadata, memory_value=memory_value, + packed_key_states_normalized=packed_key_states_normalized, ) if getattr(memory_value, "for_cuda_graphs", False): raise ValueError("replicated attention_io_layout does not support ARMemoryState(for_cuda_graphs=True)") @@ -201,6 +205,13 @@ def forward( packed_key_states, packed_value_states, ) + local_key_pack_normalized: SequencePack | None = None + if packed_key_states_normalized is not None: + _, local_key_pack_normalized, _ = self._slice_local_heads( + packed_query_states, + packed_key_states_normalized, + packed_value_states, + ) local_output_pack, kv_to_store = self.wrapped_dispatch( local_query_pack, local_key_pack, @@ -208,6 +219,7 @@ def forward( attention_mask, natten_metadata=natten_metadata, memory_value=memory_value, + packed_key_states_normalized=local_key_pack_normalized, ) return local_output_pack, kv_to_store diff --git a/cosmos_framework/model/vfm/mot/parallelize_vfm_network.py b/cosmos_framework/model/vfm/mot/parallelize_vfm_network.py index 746afa04..335234a5 100644 --- a/cosmos_framework/model/vfm/mot/parallelize_vfm_network.py +++ b/cosmos_framework/model/vfm/mot/parallelize_vfm_network.py @@ -83,8 +83,6 @@ def parallelize_vfm_network( if parallel_dims is not None and parallel_dims.dp_enabled: # Collect parameters to ignore during FSDP wrapping ignored_params = set() - if model.latent_pos_embed is not None: - ignored_params.update(model.latent_pos_embed.parameters()) model = fully_shard( module=model, diff --git a/cosmos_framework/model/vfm/mot/und_k_norm_example_test.py b/cosmos_framework/model/vfm/mot/und_k_norm_example_test.py new file mode 100644 index 00000000..05e6cc83 --- /dev/null +++ b/cosmos_framework/model/vfm/mot/und_k_norm_example_test.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +""" +Unit tests for the und_k_norm_for_gen QK-norm fix in PackedAttentionMoT. + +Background +---------- +Nemotron-3-2B has qk_norm_for_text=False (reasoner: Identity) but +qk_norm_for_diffusion=True (generator: RMSNorm). In packed joint attention +this means the generator's full attention computes norm(Q_gen) · K_und_raw^T +where K_und_raw has uncontrolled magnitude, dominating the attention logits. + +The fix (use_und_k_norm_for_gen=True) adds k_norm_und_for_gen (RMSNorm) that +normalises K_und specifically for the gen→und cross-attention path, while the +reasoner's own causal self-attention continues to use raw K_und. + +Tests +----- +1. Init: k_norm_und_for_gen is None when disabled or for Qwen3 (both norms present). +2. Backward compat: passing packed_key_states_normalized=None is a no-op. +3. Gen scale invariance (two_way): gen output unchanged when K_und is scaled but + normed keys are provided. +4. Reasoner uses raw K (two_way): reasoner output changes when K_und is scaled. +5. Gen scale invariance (three_way): same as (3) but for three_way attention. +""" + +import pytest +import torch + +from cosmos_framework.model.vfm.mot.attention import ( + build_packed_sequence, + three_way_attention, + two_way_attention, +) +from cosmos_framework.model.vfm.mot.unified_mot import ( + LayerTypes, + PackedAttentionMoT, +) +from cosmos_framework.model.vfm.reasoner.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( + Nemotron3DenseVLTextConfig, +) +from cosmos_framework.model.vfm.reasoner.nemotron_3_dense_vl.nemotron_3_dense_vl import ( + Nemotron3DenseVLRMSNorm, +) +from cosmos_framework.model.vfm.reasoner.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig +from cosmos_framework.data.vfm.sequence_packing.runtime import ( + get_gen_seq, + get_und_seq, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +NUM_Q_HEADS = 8 +NUM_KV_HEADS = 2 +HEAD_DIM = 64 +SEQS_PER_BATCH = 2 +MAX_SEQ_LEN = 16 +UND_SCALE = 100.0 # scale factor used to test K_und invariance + + +def _make_tiny_nemotron_config() -> Nemotron3DenseVLTextConfig: + return Nemotron3DenseVLTextConfig( + hidden_size=NUM_Q_HEADS * HEAD_DIM, + num_attention_heads=NUM_Q_HEADS, + num_key_value_heads=NUM_KV_HEADS, + num_hidden_layers=1, + attention_bias=False, + ) + + +def _make_tiny_qwen3_config() -> Qwen3VLTextConfig: + return Qwen3VLTextConfig( + hidden_size=NUM_Q_HEADS * HEAD_DIM, + num_attention_heads=NUM_Q_HEADS, + num_key_value_heads=NUM_KV_HEADS, + num_hidden_layers=1, + attention_bias=False, + ) + + +def _make_packs(device: torch.device, impl: str, seed: int = 42): + """ + Build minimal QKV packs for two_way / three_way attention tests. + + Returns + ------- + q_pack, k_pack_raw, v_pack, k_pack_normed, split_info, packed_und_idx + k_pack_raw : K pack where causal (und) tokens have scale UND_SCALE + k_pack_normed: K pack where causal (und) tokens are RMS-normalised + packed_und_idx: LongTensor of und token positions in the flat sequence + """ + torch.manual_seed(seed) + sample_lens = [10, 8] + full_len = sum(sample_lens) + + # Build und/gen splits: first half of each sample → und (causal), rest → gen (full) + split_lens = [] + attn_modes = [] + packed_und_idx = [] + packed_gen_idx = [] + token_shapes = [] + pos = 0 + for slen in sample_lens: + und_len = slen // 2 + gen_len = slen - und_len + split_lens.extend([und_len, gen_len]) + attn_modes.extend(["causal", "full"]) + packed_und_idx.extend(range(pos, pos + und_len)) + packed_gen_idx.extend(range(pos + und_len, pos + slen)) + token_shapes.append((1, gen_len)) # (H, W) for three_way NATTEN metadata + pos += slen + + und_idx_t = torch.tensor(packed_und_idx, dtype=torch.long, device=device) + gen_idx_t = torch.tensor(packed_gen_idx, dtype=torch.long, device=device) + + def make_pack(tensor): + pack, meta, _ = build_packed_sequence( + impl, + packed_sequence=tensor, + attn_modes=attn_modes, + split_lens=split_lens, + sample_lens=sample_lens, + packed_und_token_indexes=und_idx_t, + packed_gen_token_indexes=gen_idx_t, + num_heads=NUM_Q_HEADS, + head_dim=HEAD_DIM, + num_layers=1, + token_shapes=token_shapes, + ) + return pack, meta + + q_raw = torch.randn(full_len, NUM_Q_HEADS, HEAD_DIM, device=device, dtype=torch.bfloat16) + k_raw = torch.randn(full_len, NUM_KV_HEADS, HEAD_DIM, device=device, dtype=torch.bfloat16) + v_raw = torch.randn(full_len, NUM_KV_HEADS, HEAD_DIM, device=device, dtype=torch.bfloat16) + + # Scale up und portion of K to simulate large-magnitude und keys (the bug scenario) + k_scaled = k_raw.clone() + k_scaled[und_idx_t] = k_scaled[und_idx_t] * UND_SCALE + + # Manually RMS-normalise the und portion for the normed pack + k_normed_flat = k_scaled.clone() + k_und = k_normed_flat[und_idx_t].float() # [N_und, kv_heads, head_dim] + rms = k_und.pow(2).mean(-1, keepdim=True).sqrt() + 1e-5 + k_normed_flat[und_idx_t] = (k_und / rms).to(k_raw.dtype) + + q_pack, split_info = make_pack(q_raw) + k_pack_raw, _ = make_pack(k_scaled) + v_pack, _ = make_pack(v_raw) + k_pack_normed, _ = make_pack(k_normed_flat) + + # Also build an unscaled k pack (for backward-compat comparison) + k_pack_unscaled, _ = make_pack(k_raw) + + return q_pack, k_pack_raw, v_pack, k_pack_normed, k_pack_unscaled, split_info, und_idx_t + + +# --------------------------------------------------------------------------- +# 1. Init tests (CPU, L0) +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_k_norm_und_for_gen_disabled_by_default(): + """k_norm_und_for_gen must be None when use_und_k_norm_for_gen=False (default).""" + config = _make_tiny_nemotron_config() + layer_types = LayerTypes("nemotron_dense") + attn = PackedAttentionMoT( + config, + layer_idx=0, + layer_types=layer_types, + qk_norm_for_text=False, + qk_norm_for_diffusion=True, + use_und_k_norm_for_gen=False, + ) + assert attn.k_norm_und_for_gen is None + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_k_norm_und_for_gen_not_created_when_both_pathways_have_norm(): + """k_norm_und_for_gen must be None when qk_norm_for_text=True (Qwen3 config). + + When the reasoner already has QK norm, K_und is already normalised — no fix needed. + """ + config = _make_tiny_qwen3_config() + layer_types = LayerTypes("qwen3_vl_dense") + attn = PackedAttentionMoT( + config, + layer_idx=0, + layer_types=layer_types, + qk_norm_for_text=True, + qk_norm_for_diffusion=True, + use_und_k_norm_for_gen=True, # flag on, but condition not met + ) + assert attn.k_norm_und_for_gen is None + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_k_norm_und_for_gen_created_for_nemotron_when_enabled(): + """k_norm_und_for_gen must be an RMSNorm with weight=1 for Nemotron + flag enabled.""" + config = _make_tiny_nemotron_config() + layer_types = LayerTypes("nemotron_dense") + attn = PackedAttentionMoT( + config, + layer_idx=0, + layer_types=layer_types, + qk_norm_for_text=False, + qk_norm_for_diffusion=True, + use_und_k_norm_for_gen=True, + ) + assert attn.k_norm_und_for_gen is not None + assert isinstance(attn.k_norm_und_for_gen, Nemotron3DenseVLRMSNorm) + assert attn.k_norm_und_for_gen.weight.shape == (attn.head_dim,) + # weight initialised to 1 → pure RMSNorm at step 0 + torch.testing.assert_close(attn.k_norm_und_for_gen.weight, torch.ones(attn.head_dim)) + + +# --------------------------------------------------------------------------- +# 2. Backward compat: packed_key_states_normalized=None is a no-op (GPU, L0) +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_two_way_attention_normed_none_is_noop(): + """two_way_attention output must be identical when packed_key_states_normalized=None.""" + device = torch.device("cuda") + q_pack, k_pack_raw, v_pack, _, _, split_info, _ = _make_packs(device, "two_way") + + out_baseline = two_way_attention(q_pack, k_pack_raw, v_pack) + out_explicit_none = two_way_attention(q_pack, k_pack_raw, v_pack, packed_key_states_normalized=None) + + torch.testing.assert_close( + get_und_seq(out_baseline), + get_und_seq(out_explicit_none), + atol=0, + rtol=0, + ) + torch.testing.assert_close( + get_gen_seq(out_baseline), + get_gen_seq(out_explicit_none), + atol=0, + rtol=0, + ) + + +@pytest.mark.L0 +def test_three_way_attention_normed_none_is_noop(): + """three_way_attention output must be identical when packed_key_states_normalized=None.""" + device = torch.device("cuda") + q_pack, k_pack_raw, v_pack, _, _, split_info, _ = _make_packs(device, "three_way") + + out_baseline = three_way_attention(q_pack, k_pack_raw, v_pack, natten_metadata=None) + out_explicit_none = three_way_attention( + q_pack, k_pack_raw, v_pack, natten_metadata=None, packed_key_states_normalized=None + ) + + torch.testing.assert_close(get_und_seq(out_baseline), get_und_seq(out_explicit_none), atol=0, rtol=0) + torch.testing.assert_close(get_gen_seq(out_baseline), get_gen_seq(out_explicit_none), atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# 3 & 4. Scale invariance and reasoner raw-K tests (GPU, L1) +# --------------------------------------------------------------------------- + + +@pytest.mark.L1 +def test_two_way_gen_output_independent_of_raw_und_key_scale(): + """Gen output must be identical when raw K_und scale changes but normed keys are fixed. + + When packed_key_states_normalized is provided, two_way_attention routes gen full-attention + through get_all_seq(packed_key_states_normalized), not packed_key_states. So regardless + of whether packed_key_states has 1x or 100x K_und, the gen output is the same. + Reasoner output (causal path) uses packed_key_states and therefore differs. + """ + device = torch.device("cuda") + q_pack, k_pack_raw, v_pack, k_pack_normed, k_pack_unscaled, _, _ = _make_packs(device, "two_way") + + # Both calls use the SAME k_pack_normed → gen sees identical keys → identical gen output. + out_raw = two_way_attention(q_pack, k_pack_raw, v_pack, packed_key_states_normalized=k_pack_normed) + out_unscaled = two_way_attention(q_pack, k_pack_unscaled, v_pack, packed_key_states_normalized=k_pack_normed) + + # Gen outputs must be bitwise identical (same normed keys). + torch.testing.assert_close( + get_gen_seq(out_raw), + get_gen_seq(out_unscaled), + atol=0, + rtol=0, + msg="Gen output must be identical when only raw K_und scale differs (normed keys fixed)", + ) + + # Reasoner outputs must differ (different raw K_und: 100x vs 1x). + und_diff = (get_und_seq(out_raw) - get_und_seq(out_unscaled)).abs().max() + assert und_diff > 1e-2, f"Reasoner output should differ with scaled K_und (max diff {und_diff:.4f})" + + +@pytest.mark.L1 +def test_three_way_gen_output_independent_of_raw_und_key_scale(): + """Same property as two_way: gen output unchanged when raw K_und differs but normed keys fixed.""" + device = torch.device("cuda") + q_pack, k_pack_raw, v_pack, k_pack_normed, k_pack_unscaled, _, _ = _make_packs(device, "three_way") + + out_raw = three_way_attention( + q_pack, k_pack_raw, v_pack, natten_metadata=None, packed_key_states_normalized=k_pack_normed + ) + out_unscaled = three_way_attention( + q_pack, k_pack_unscaled, v_pack, natten_metadata=None, packed_key_states_normalized=k_pack_normed + ) + + # Gen outputs must be bitwise identical (same normed keys). + torch.testing.assert_close( + get_gen_seq(out_raw), + get_gen_seq(out_unscaled), + atol=0, + rtol=0, + msg="Gen output must be identical when only raw K_und scale differs (three_way, normed keys fixed)", + ) + + # Reasoner outputs must differ. + und_diff = (get_und_seq(out_raw) - get_und_seq(out_unscaled)).abs().max() + assert und_diff > 1e-2, f"Reasoner output should differ with scaled K_und (max diff {und_diff:.4f})" diff --git a/cosmos_framework/model/vfm/mot/unified_mot.py b/cosmos_framework/model/vfm/mot/unified_mot.py index 5c76a15e..7d5f9757 100644 --- a/cosmos_framework/model/vfm/mot/unified_mot.py +++ b/cosmos_framework/model/vfm/mot/unified_mot.py @@ -5,7 +5,6 @@ import time from collections.abc import Mapping from dataclasses import dataclass -from pathlib import Path from typing import Any import torch @@ -19,13 +18,12 @@ AttentionMaskType, dispatch_attention, ) -from cosmos_framework.model.vfm.utils.memory import KVToStore, MemoryState, MemoryValue # Nemotron 3 Dense VL imports -from cosmos_framework.model.vfm.vlm.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( +from cosmos_framework.model.vfm.reasoner.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( Nemotron3DenseVLTextConfig, ) -from cosmos_framework.model.vfm.vlm.nemotron_3_dense_vl.nemotron_3_dense_vl import ( +from cosmos_framework.model.vfm.reasoner.nemotron_3_dense_vl.nemotron_3_dense_vl import ( MultiModalRotaryEmbedding, Nemotron3DenseVLMLP, Nemotron3DenseVLPreTrainedModel, @@ -34,32 +32,32 @@ ) # Qwen3-VL imports -from cosmos_framework.model.vfm.vlm.qwen3_vl.configuration_qwen3_vl import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.configuration_qwen3_vl import ( Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl.qwen3_vl import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.qwen3_vl import ( Qwen3VLPreTrainedModel, Qwen3VLTextMLP, Qwen3VLTextRMSNorm, Qwen3VLTextRotaryEmbedding, Qwen3VLVisionModel, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl.qwen3_vl import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.qwen3_vl import ( apply_rotary_pos_emb as qwen3_vl_apply_rotary_pos_emb, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( prepare_multimodal_reasoner_inputs, ) # Qwen3-VL-MoE imports -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.configuration_qwen3_vl_moe import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( Qwen3VLMoeConfig, Qwen3VLMoeTextConfig, Qwen3VLMoeVisionConfig, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.qwen3_vl_moe import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.qwen3_vl_moe import ( LBLMetadata, Qwen3VLMoePreTrainedModel, Qwen3VLMoeTextMLP, @@ -68,6 +66,7 @@ Qwen3VLMoeTextSparseMoeBlock, Qwen3VLMoeVisionModel, ) +from cosmos_framework.model.vfm.utils.memory import KVToStore, MemoryState, MemoryValue from cosmos_framework.data.vfm.sequence_packing.runtime import ( SequencePack, from_all_seq, @@ -141,27 +140,6 @@ def is_moe(self) -> bool: # MoT wrapper configs — one per architecture family # ----------------------------------------------------------------------------- -# Package root = parent of the top-level package directory; the shipped -# model-config JSONs live under it (cosmos_framework/model/... in releases; -# cosmos_framework/model/vfm/... in i4 source). -_PACKAGE_ROOT = Path(__file__).resolve().parents[4] - - -def _resolve_packaged_config_path(json_file: str) -> str: - """Resolve a model-config JSON path so it loads regardless of CWD. - - Absolute paths and paths that already exist relative to the CWD are returned - unchanged (preserving existing behavior when launched from the repo root). - A relative path that does not exist against the CWD is resolved against the - installed package root. If that candidate is missing too, the original path - is returned so ``open()`` raises the familiar ``FileNotFoundError``. - """ - path = Path(json_file) - if path.is_absolute() or path.exists(): - return json_file - candidate = _PACKAGE_ROOT / json_file - return str(candidate) if candidate.exists() else json_file - class _MoTConfigBase(object): """Shared MoT wrapper logic for all three architecture families. @@ -212,7 +190,7 @@ class _MoTConfigBase(object): ``text_config`` access picks it up. Post-construction overrides via plain ``setattr`` (the - ``create_vlm_config`` flow in ``configs/base/defaults/vlm.py``) + ``create_vlm_config`` flow in ``configs/base/defaults/reasoner.py``) just update the same plain attributes, so the next property access picks up the latest values. No cache, no ``__setattr__`` interception, no override bucket — the property rebuild is cheap @@ -389,7 +367,7 @@ def from_json_file(cls, json_file: str) -> "_MoTConfigBase": :pyattr:`vision_config` and by HF downstream consumers reading the dict directly. """ - with open(_resolve_packaged_config_path(json_file), encoding="utf-8") as reader: + with open(json_file, encoding="utf-8") as reader: config_dict = json.load(reader) return cls(config_dict=config_dict) @@ -488,6 +466,7 @@ def __init__( layer_types: LayerTypes, qk_norm_for_text: bool, qk_norm_for_diffusion: bool, + use_und_k_norm_for_gen: bool = False, ): super().__init__() self.config = config @@ -524,6 +503,19 @@ def __init__( self.q_norm_moe_gen = nn.Identity() self.k_norm_moe_gen = nn.Identity() + # Cross-attention K norm: normalises und K tokens seen by the generator in the + # gen→und cross-attention path. Only needed when the generation pathway has QK + # norm (qk_norm_for_diffusion=True) but the understanding pathway does not + # (qk_norm_for_text=False), i.e. the Nemotron-3 configuration. Without this, + # the joint softmax computes norm(Q_gen) · K_und_raw^T where K_und_raw has large + # uncontrolled magnitude and dominates attention over the gen self-attention path. + # When both pathways share the same QK norm (or neither has one) k_norm_und_for_gen + # is None and the standard packed K tensor is used for all paths unchanged. + if use_und_k_norm_for_gen and qk_norm_for_diffusion and not qk_norm_for_text: + self.k_norm_und_for_gen: nn.Module | None = layer_types.rms_norm(self.head_dim, eps=eps) + else: + self.k_norm_und_for_gen = None + # Generation pathway linear projections self.q_proj_moe_gen = nn.Linear( self.hidden_size, self.num_attention_heads * self.head_dim, bias=config.attention_bias @@ -644,6 +636,23 @@ def forward( packed_key_states_ = from_und_gen_splits(k_und_, k_gen_, pack) # [N_und+N_gen,num_kv_heads,head_dim] packed_value_states_ = from_und_gen_splits(v_und, v_gen, pack) # [N_und+N_gen,num_kv_heads,head_dim] + # Build a separate K pack where the und tokens are normalised for gen→und + # cross-attention (fixes scale mismatch when qk_norm_for_diffusion=True but + # qk_norm_for_text=False, i.e. the Nemotron-3 config). The raw k_und_ is kept + # in packed_key_states_ for the reasoner's own causal self-attention. + if self.k_norm_und_for_gen is not None: + k_und_normalized = self.k_norm_und_for_gen(k_und) # RMSNorm before RoPE + _, k_und_for_gen_ = self._apply_rotary_pos_emb( + q_und_, # dummy q — already RoPE-applied, output discarded; only k is used + k_und_normalized, + get_und_seq(packed_cos), + get_und_seq(packed_sin), + unsqueeze_dim=1, + ) # k_und_for_gen_: [N_und,num_kv_heads,head_dim] + packed_key_states_normalized_: SequencePack | None = from_und_gen_splits(k_und_for_gen_, k_gen_, pack) + else: + packed_key_states_normalized_ = None + packed_attn_output, kv_to_store = self.dispatch_attention_fn( packed_query_states_, packed_key_states_, @@ -651,6 +660,7 @@ def forward( attention_mask, natten_metadata=natten_metadata, memory_value=memory_value, + packed_key_states_normalized=packed_key_states_normalized_, ) # Produce kv_to_store for MemoryState.write_for_layer() when the @@ -664,10 +674,15 @@ def forward( if memory_value is not None and kv_to_store is None: und_len = pack["_num_causal_tokens"] gen_len = pack["_num_full_tokens"] + # When und K-norm is active, AR frame 1+ gen→und cross-attention uses + # the normalised K, so cache k_und_for_gen_ (RMSNorm+RoPE applied) instead + # of raw k_und_. Without the norm, k_und_for_gen_ is not defined, so + # fall back to k_und_. + k_und_to_store = k_und_for_gen_ if self.k_norm_und_for_gen is not None else k_und_ kv_to_store = ( k_gen_[:gen_len].unsqueeze(0), v_gen[:gen_len].unsqueeze(0), - k_und_[:und_len].unsqueeze(0), + k_und_to_store[:und_len].unsqueeze(0), v_und[:und_len].unsqueeze(0), ) @@ -787,6 +802,7 @@ def _impl_init( qk_norm_for_text: bool, qk_norm_for_diffusion: bool, gen_noisy_gating: bool = False, + use_und_k_norm_for_gen: bool = False, ): """Shared ``__init__`` body for the three MoT text-model variants. @@ -809,6 +825,7 @@ def _impl_init( qk_norm_for_text=qk_norm_for_text, qk_norm_for_diffusion=qk_norm_for_diffusion, gen_noisy_gating=gen_noisy_gating, + use_und_k_norm_for_gen=use_und_k_norm_for_gen, ) ) @@ -986,6 +1003,7 @@ def __init__( qk_norm_for_text: bool, qk_norm_for_diffusion: bool, gen_noisy_gating: bool = False, + use_und_k_norm_for_gen: bool = False, ): super().__init__() self.hidden_size = config.hidden_size @@ -995,6 +1013,7 @@ def __init__( layer_idx=layer_idx, qk_norm_for_text=qk_norm_for_text, qk_norm_for_diffusion=qk_norm_for_diffusion, + use_und_k_norm_for_gen=use_und_k_norm_for_gen, ) if ( @@ -1220,6 +1239,7 @@ def __init__( *, qk_norm_for_text: bool, qk_norm_for_diffusion: bool, + use_und_k_norm_for_gen: bool, ): super().__init__(config) _impl_init( @@ -1228,6 +1248,7 @@ def __init__( layer_types=LayerTypes("qwen3_vl_dense"), qk_norm_for_text=qk_norm_for_text, qk_norm_for_diffusion=qk_norm_for_diffusion, + use_und_k_norm_for_gen=use_und_k_norm_for_gen, ) def init_taylorseer(self, cache_dic=None, current=None): @@ -1254,6 +1275,7 @@ def __init__( qk_norm_for_text: bool, qk_norm_for_diffusion: bool, gen_noisy_gating: bool = False, + use_und_k_norm_for_gen: bool, ): super().__init__(config) _impl_init( @@ -1263,6 +1285,7 @@ def __init__( qk_norm_for_text=qk_norm_for_text, qk_norm_for_diffusion=qk_norm_for_diffusion, gen_noisy_gating=gen_noisy_gating, + use_und_k_norm_for_gen=use_und_k_norm_for_gen, ) def init_taylorseer(self, cache_dic=None, current=None): @@ -1288,6 +1311,7 @@ def __init__( *, qk_norm_for_text: bool, qk_norm_for_diffusion: bool, + use_und_k_norm_for_gen: bool, ): super().__init__(config) _impl_init( @@ -1296,6 +1320,7 @@ def __init__( layer_types=LayerTypes("nemotron_dense"), qk_norm_for_text=qk_norm_for_text, qk_norm_for_diffusion=qk_norm_for_diffusion, + use_und_k_norm_for_gen=use_und_k_norm_for_gen, ) def init_taylorseer(self, cache_dic=None, current=None) -> None: @@ -1596,6 +1621,8 @@ def _impl_generate_reasoner_text( *, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, @@ -1652,10 +1679,9 @@ def _impl_generate_reasoner_text( ``Qwen3VLProcessor`` emits — pass it through unchanged. Moved to the prompt's device internally. ``None`` (default) means text-only prompt; in that case the multimodal prefill - path is skipped entirely. Videos are *not* supported here — - this function has no ``pixel_values_videos`` / ``video_grid_thw`` - parameters; for I2V conditioning, frames must be passed as - images. + path is skipped entirely. For video conditioning, pass ``pixel_values_videos`` + + ``video_grid_thw`` instead (mutually exclusive with the image + pair). image_grid_thw: Optional ``[num_images, 3]`` long tensor giving ``(t, h, w)`` — the temporal / height / width feature-grid size per image as produced by ``Qwen3VLProcessor`` (``t`` is @@ -1747,11 +1773,15 @@ def _impl_generate_reasoner_text( if (pixel_values is None) != (image_grid_thw is None): raise ValueError("pixel_values and image_grid_thw must be provided together.") + if (pixel_values_videos is None) != (video_grid_thw is None): + raise ValueError("pixel_values_videos and video_grid_thw must be provided together.") + if pixel_values is not None and pixel_values_videos is not None: + raise ValueError("Reasoner conditions on one medium at a time: pass image OR video, not both.") _prefill_start = time.time() mrope_position_deltas: torch.Tensor | None = None - if pixel_values is None: + if pixel_values is None and pixel_values_videos is None: hidden = model.reasoner_forward(input_ids, cache=cache) # [B,T_prompt,hidden_size] else: if not hasattr(causal_lm, "visual"): @@ -1767,6 +1797,8 @@ def _impl_generate_reasoner_text( input_ids=input_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, attention_mask=attention_mask, ) hidden = model.reasoner_forward( @@ -1942,7 +1974,8 @@ class Qwen3VLTextForCausalLM(Qwen3VLPreTrainedModel): This variant is used for dense-only MLP models. """ - _tied_weights_keys = ["lm_head.weight"] + # lm_head.weight is tied to model.embed_tokens.weight + _tied_weights_keys: list[str] = ["lm_head.weight"] def __init__(self, config: Qwen3VLMoTConfig): # Materialize a fresh HF ``Qwen3VLTextConfig`` from the wrapper. @@ -1961,6 +1994,7 @@ def __init__(self, config: Qwen3VLMoTConfig): text_config, qk_norm_for_text=config.qk_norm_for_text, qk_norm_for_diffusion=config.qk_norm_for_diffusion, + use_und_k_norm_for_gen=getattr(config, "use_und_k_norm_for_gen", False), ) self.vocab_size = text_config.vocab_size self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) @@ -2033,6 +2067,8 @@ def generate_reasoner_text( *, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, @@ -2053,6 +2089,8 @@ def generate_reasoner_text( the Qwen3-VL visual encoder; omit them for text-only prefill. The two arguments are mutually required: passing exactly one raises ``ValueError`` inside :func:`_impl_generate_reasoner_text`. + Video conditioning is also supported via ``pixel_values_videos`` + + ``video_grid_thw``; the image and video pairs are mutually exclusive. Uses the und-pathway weights (those WITHOUT the ``_moe_gen`` suffix) plus the model-level ``embed_tokens`` / ``norm`` / ``lm_head``, and — @@ -2067,6 +2105,8 @@ def generate_reasoner_text( max_new_tokens=max_new_tokens, pixel_values=pixel_values, image_grid_thw=image_grid_thw, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, attention_mask=attention_mask, eos_token_id=eos_token_id, pad_token_id=pad_token_id, @@ -2087,7 +2127,8 @@ class Qwen3VLMoeTextForCausalLM(Qwen3VLMoePreTrainedModel): This variant is used for MoE MLP models. """ - _tied_weights_keys = ["lm_head.weight"] + # lm_head.weight is tied to model.embed_tokens.weight + _tied_weights_keys: list[str] = ["lm_head.weight"] def __init__(self, config: Qwen3VLMoeMoTConfig): super().__init__(config.full_config) @@ -2098,6 +2139,7 @@ def __init__(self, config: Qwen3VLMoeMoTConfig): qk_norm_for_text=config.qk_norm_for_text, qk_norm_for_diffusion=config.qk_norm_for_diffusion, gen_noisy_gating=config.gen_noisy_gating, + use_und_k_norm_for_gen=getattr(config, "use_und_k_norm_for_gen", False), ) self.vocab_size = text_config.vocab_size self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) @@ -2166,6 +2208,8 @@ def generate_reasoner_text( *, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, @@ -2186,6 +2230,8 @@ def generate_reasoner_text( the Qwen3-VL visual encoder; omit them for text-only prefill. The two arguments are mutually required: passing exactly one raises ``ValueError`` inside :func:`_impl_generate_reasoner_text`. + Video conditioning is also supported via ``pixel_values_videos`` + + ``video_grid_thw``; the image and video pairs are mutually exclusive. Uses the und-pathway weights (those WITHOUT the ``_moe_gen`` suffix) plus the model-level ``embed_tokens`` / ``norm`` / ``lm_head``, and — @@ -2201,6 +2247,8 @@ def generate_reasoner_text( max_new_tokens=max_new_tokens, pixel_values=pixel_values, image_grid_thw=image_grid_thw, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, attention_mask=attention_mask, eos_token_id=eos_token_id, pad_token_id=pad_token_id, @@ -2236,6 +2284,7 @@ def __init__(self, config: Nemotron3DenseVLMoTConfig) -> None: text_config, qk_norm_for_text=config.qk_norm_for_text, qk_norm_for_diffusion=config.qk_norm_for_diffusion, + use_und_k_norm_for_gen=getattr(config, "use_und_k_norm_for_gen", False), ) self.vocab_size = text_config.vocab_size self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) diff --git a/cosmos_framework/model/vfm/mot/unified_mot_test.py b/cosmos_framework/model/vfm/mot/unified_mot_test.py deleted file mode 100644 index 1d246622..00000000 --- a/cosmos_framework/model/vfm/mot/unified_mot_test.py +++ /dev/null @@ -1,68 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -"""Regression tests for package-relative model-config resolution. - -The shipped config defaults (cosmos_framework/configs/base/defaults/vlm.py) -reference the per-architecture JSONs by a repo-root-relative path, e.g. -``"cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json"``. -``_MoTConfigBase.from_json_file`` used to ``open()`` that string verbatim, so it -only resolved when the process CWD was the framework repo root. Launching -``cosmos_framework.scripts.train`` from any other directory (e.g. a cookbook -folder) raised ``FileNotFoundError`` during model construction. These tests pin -the package-root fallback that fixes it. -""" - -from pathlib import Path - -import cosmos_framework -from cosmos_framework.model.vfm.mot.unified_mot import ( - Qwen3VLMoTConfig, - _PACKAGE_ROOT, - _resolve_packaged_config_path, -) - -# A config JSON that ships inside the package, named exactly as the config -# defaults reference it (relative to the package root). -_SHIPPED_REL = "cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" - - -def test_package_root_contains_cosmos_framework(): - # _PACKAGE_ROOT must be the directory that *contains* the cosmos_framework - # package, so that "/cosmos_framework/..." resolves in both editable - # and wheel installs. - assert (_PACKAGE_ROOT / "cosmos_framework").is_dir() - assert Path(cosmos_framework.__file__).resolve().parent == _PACKAGE_ROOT / "cosmos_framework" - - -def test_resolve_absolute_path_passes_through(tmp_path): - abs_path = tmp_path / "cfg.json" - abs_path.write_text("{}") - assert _resolve_packaged_config_path(str(abs_path)) == str(abs_path) - - -def test_resolve_existing_relative_path_passes_through(tmp_path, monkeypatch): - # A relative path that exists against the CWD is returned unchanged — the - # package-root fallback must not shadow a real working-directory file. - monkeypatch.chdir(tmp_path) - (tmp_path / "local.json").write_text("{}") - assert _resolve_packaged_config_path("local.json") == "local.json" - - -def test_resolve_shipped_config_from_foreign_cwd(tmp_path, monkeypatch): - # The actual regression: from a directory that is not the repo root, the - # shipped repo-root-relative path still resolves to the packaged file. - monkeypatch.chdir(tmp_path) - resolved = Path(_resolve_packaged_config_path(_SHIPPED_REL)) - assert resolved.is_absolute() - assert resolved.is_file() - assert resolved == _PACKAGE_ROOT / _SHIPPED_REL - - -def test_from_json_file_loads_shipped_config_from_foreign_cwd(tmp_path, monkeypatch): - # End-to-end: from_json_file (the call made during model construction) loads - # the shipped config from a foreign CWD instead of raising FileNotFoundError. - monkeypatch.chdir(tmp_path) - config = Qwen3VLMoTConfig.from_json_file(_SHIPPED_REL) - assert isinstance(config.config_dict, dict) - assert config.config_dict.get("model_type") # sanity: a real Qwen3-VL config diff --git a/cosmos_framework/model/vfm/omni_mot_model.py b/cosmos_framework/model/vfm/omni_mot_model.py index d136a801..19aac5f1 100644 --- a/cosmos_framework/model/vfm/omni_mot_model.py +++ b/cosmos_framework/model/vfm/omni_mot_model.py @@ -36,6 +36,7 @@ from cosmos_framework.model.vfm.mot.cosmos3_vfm_network import Cosmos3VFMNetwork, Cosmos3VFMNetworkConfig from cosmos_framework.model.vfm.mot.modeling_utils import has_noisy_tokens from cosmos_framework.model.vfm.mot.parallelize_vfm_network import parallelize_vfm_network +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption from cosmos_framework.model.vfm.utils.data_and_condition import ( GenerationDataClean, GenerationDataNoised, @@ -44,8 +45,9 @@ unwrap_and_densify, ) from cosmos_framework.model.vfm.utils.memory import MemoryState -from cosmos_framework.model.vfm.utils.safetensors_loader import load_language_model as load_language_model_safetensors -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.vfm.utils.safetensors_loader import ( + load_language_model as load_language_model_safetensors, +) from cosmos_framework.data.vfm.sequence_packing import ( PackedSequence, SequencePlan, @@ -171,11 +173,11 @@ def set_up_tokenizers(self) -> None: self.tokenizer_sound_gen = None - def build_net(self, dtype: torch.dtype): + def build_net(self, dtype: torch.dtype, *, lora_enabled: bool | None = None) -> torch.nn.Module: # Build model network and parallelize it. + lora_enabled = self.config.lora_enabled if lora_enabled is None else lora_enabled with torch.device("meta"): assert self.vlm_config.model_instance is not None, "Model instance should be specified" - language_model = lazy_instantiate(self.vlm_config.model_instance) # NOTE: We pass "RF timesteps" to the network in the same scale as the scheduler @@ -191,15 +193,11 @@ def build_net(self, dtype: torch.dtype): max_latent_h=self.config.diffusion_expert_config.max_vae_latent_side_after_patchify, max_latent_w=self.config.diffusion_expert_config.max_vae_latent_side_after_patchify, max_latent_t=self.config.state_t, - rope_h_extrapolation_ratio=self.config.diffusion_expert_config.rope_h_extrapolation_ratio, - rope_w_extrapolation_ratio=self.config.diffusion_expert_config.rope_w_extrapolation_ratio, - rope_t_extrapolation_ratio=self.config.diffusion_expert_config.rope_t_extrapolation_ratio, enable_fps_modulation=self.config.diffusion_expert_config.enable_fps_modulation, base_fps=self.config.diffusion_expert_config.base_fps, vision_gen=self.config.vision_gen, action_gen=self.config.action_gen, sound_gen=self.config.sound_gen, - position_embedding_type=self.config.diffusion_expert_config.position_embedding_type, joint_attn_implementation=self.config.joint_attn_implementation, timestep_scale=1.0 / float(num_train_timesteps) * self.config.diffusion_expert_config.timestep_range, action_dim=self.config.max_action_dim, @@ -222,7 +220,7 @@ def build_net(self, dtype: torch.dtype): # injector must see unsharded Linear shapes; injecting post-FSDP causes # lora_B to be created at the per-rank shard size and crashes at # forward time. See `OmniMoTModel.add_lora` for details. - if getattr(self.config, "lora_enabled", False): + if lora_enabled: net = self.add_lora( net, lora_rank=self.config.lora_rank, @@ -248,7 +246,7 @@ def build_net(self, dtype: torch.dtype): # meta), since they are only for checkpoint conversion and smoke # tests. net.init_weights(buffer_device=DEVICE) - if getattr(self.config, "lora_enabled", False): + if lora_enabled: self._init_lora_weights_post_materialization(net) return net @@ -543,7 +541,7 @@ def _pack_input_sequence( ) -> PackedSequence: """Wrap ``pack_input_sequence`` with all config-derived args pre-filled. - Centralises the 10 config-derived positional/embedding args so callers only + Centralises the config-derived positional/embedding args so callers only supply the four per-call arguments (sequence_plans, text tokens, data, timesteps) plus three optional flags. """ @@ -557,7 +555,6 @@ def _pack_input_sequence( latent_patch_size=self.config.diffusion_expert_config.patch_spatial, skip_text_tokens=skip_text_tokens, include_end_of_generation_token=include_end_of_generation_token, - position_embedding_type=self.config.diffusion_expert_config.position_embedding_type, unified_3d_mrope_reset_spatial_ids=self.config.diffusion_expert_config.unified_3d_mrope_reset_spatial_ids, unified_3d_mrope_temporal_modality_margin=self.config.diffusion_expert_config.unified_3d_mrope_temporal_modality_margin, enable_fps_modulation=self.config.diffusion_expert_config.enable_fps_modulation, @@ -907,9 +904,6 @@ def training_step( memory = self.build_memory_state(packed_sequence, memory_info) # pylint: disable=assignment-from-none out_net = self.denoise( data_batch_packed=packed_sequence, - fps_vision=gen_data_clean.fps_vision, - fps_action=gen_data_clean.fps_action, - fps_sound=gen_data_clean.fps_sound, memory=memory, ) @@ -1948,14 +1942,9 @@ def _get_velocity( packed_sequence.to_cuda() # --- Network forward --- - fps_action = gen_data_clean.fps_action if has_action else None - fps_sound = gen_data_clean.fps_sound if has_sound else None out = self.denoise( net=net, data_batch_packed=packed_sequence, - fps_vision=gen_data_clean.fps_vision, - fps_action=fps_action, - fps_sound=fps_sound, ) # --- Apply velocity masks --- @@ -2432,7 +2421,6 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool): # Peers needed CFG so we ran the uncond forward to keep # FSDP allgather aligned; locally we still return cond. return cond_v - v_pred = [u_i + guidance * (c_i - u_i) for c_i, u_i in zip(cond_v, uncond_v)] if normalize_cfg: v_pred = [ @@ -2874,10 +2862,8 @@ def get_data_and_condition(self, data_batch: dict[str, torch.Tensor], iteration: else: x0_tokens_sound = None - # We pass the conditioning FPS along to the denoising function - # It will not be used for RoPE FPS modulation unless enabled in the training config - # Note: conditioning_fps from data is converted to TPS via temporal_compression_factor - # in VideoRopePosition3DEmb. + # FPS metadata is used by the sequence packer for mRoPE temporal IDs when + # FPS modulation is enabled in the training config. fps_raw = data_batch.get("conditioning_fps", None) if isinstance(fps_raw, list): fps_raw = torch.stack(fps_raw).flatten() # list of scalar tensors -> (B,) @@ -3609,9 +3595,6 @@ def denoise( self, net: torch.nn.Module | None = None, data_batch_packed: PackedSequence | None = None, - fps_vision: torch.Tensor | None = None, - fps_action: torch.Tensor | None = None, - fps_sound: torch.Tensor | None = None, memory: MemoryState | None = None, ) -> dict: """ @@ -3619,9 +3602,6 @@ def denoise( Args: data_batch_packed: PackedSequence from `pack_input_sequence(...)`. - fps_vision: Optional FPS tensor used for RoPE FPS modulation (if enabled in config). - fps_action: Optional FPS tensor used for action RoPE FPS modulation (if enabled in config). - fps_sound: Optional FPS tensor for sound RoPE modulation (e.g., sound_latent_fps=25). memory: Optional pre-built MemoryState for autoregressive generation or KV-cache training. @@ -3636,9 +3616,6 @@ def denoise( net = net or self.net out_net = net( packed_seq=data_batch_packed, - fps_vision=fps_vision, - fps_action=fps_action, - fps_sound=fps_sound, memory=memory, ) output_dict = dict() @@ -3855,6 +3832,7 @@ def generate_reasoner_text( max_new_tokens: int, *, images: list[Any] | None = None, + videos: list[Any] | None = None, prompt_builder: Callable[[str], list[dict[str, Any]]] | None = None, do_sample: bool = False, temperature: float | None = 1.0, @@ -3871,8 +3849,10 @@ def generate_reasoner_text( (or wraps the prompt as a single user message when no callback is given), (b) tokenizes it — text-only via :meth:`tokenize_text`, or multimodal via ``self.vlm_processor.apply_chat_template`` when - ``images`` is supplied (which lowers the chat into ``input_ids``, - ``attention_mask``, ``pixel_values``, and ``image_grid_thw``), (c) + ``images`` or ``videos`` is supplied (the image path lowers the chat + into ``input_ids``, ``attention_mask``, ``pixel_values``, and + ``image_grid_thw``; the video path yields ``pixel_values_videos`` and + ``video_grid_thw`` instead), (c) runs the reasoner-only AR decode loop through ``self.net.generate_reasoner_text`` (the lower-level token-driven pass-through that delegates to ``unified_mot._impl_generate_reasoner_text``), @@ -3927,6 +3907,14 @@ def generate_reasoner_text( ``processor.apply_chat_template``, so any input it accepts works (file path ``str``, ``PIL.Image.Image``, ``np.ndarray``, or a CHW / HWC tensor). + videos: Optional per-prompt conditioning videos (mutually + exclusive with ``images``). Each entry must be a + ``{"frames": [...PIL...], "fps": float}`` payload + (pre-decoded by the caller, e.g. via + ``_decode_reasoner_video``). The frames list and fps are + forwarded into the ``{"type": "video", "video": frames, + "fps": fps}`` chat block so the processor produces + ``pixel_values_videos`` / ``video_grid_thw``. prompt_builder: Optional callback that maps a raw prompt string to a chat-style messages list (e.g. :func:`cosmos_framework.model.vfm.upsampler.prompts.build_messages` @@ -3987,32 +3975,42 @@ def generate_reasoner_text( Raises: ValueError: If ``images`` length does not match ``inputs`` - length. - RuntimeError: If ``images`` is supplied but the live VLM - processor does not implement ``apply_chat_template`` - (i.e., the VLM is configured as text-only). + length, or if ``videos`` length does not match ``inputs`` + length. Also raised if both ``images`` and ``videos`` are + supplied simultaneously (only one medium is allowed per + call). + RuntimeError: If ``images`` or ``videos`` is supplied but the + live VLM processor does not implement + ``apply_chat_template`` (i.e., the VLM is configured as + text-only). """ # Decide whether the multimodal flow is in play, and validate the # image-list contract here so the failure happens before any # decoding work — far easier to debug than a downstream # ``apply_chat_template`` error. - use_multimodal = images is not None + if images is not None and videos is not None: + raise ValueError( + "generate_reasoner_text conditions on one medium at a time: pass `images` OR `videos`, not both." + ) + use_image = images is not None + use_video = videos is not None + use_multimodal = use_image or use_video + media = images if use_image else videos if use_multimodal: - assert images is not None # narrowed by `use_multimodal` - if len(images) != len(inputs): + assert media is not None # narrowed by `use_multimodal` + if len(media) != len(inputs): raise ValueError( - f"generate_reasoner_text: `images` length ({len(images)}) " + f"generate_reasoner_text: media length ({len(media)}) " f"must equal `inputs` length ({len(inputs)}) for the " - "image-conditioned flow." + "vision-conditioned flow." ) if not callable(getattr(self.vlm_processor, "apply_chat_template", None)): raise RuntimeError( - "generate_reasoner_text(images=...) requires a multimodal " + "generate_reasoner_text(images=... / videos=...) requires a multimodal " "VLM processor (e.g. Qwen3VLProcessor) but the live processor " f"{type(self.vlm_processor).__name__!r} does not implement " "apply_chat_template — the live VLM is configured as text-only." ) - # Resolve EOS / pad ids internally so callers don't have to know # about VLM-specific id wiring. EOS comes from the cached VLM # special-tokens dict (set in ``set_up_tokenizers``); pad mirrors @@ -4049,22 +4047,22 @@ def generate_reasoner_text( messages = [{"role": "user", "content": prompt}] if use_multimodal: - assert images is not None # narrowed by `use_multimodal` - # Replace the LAST user message's content with a Qwen3-VL - # multimodal block (image + text). Earlier messages - # (system, prior turns) are kept verbatim so any chat - # scaffolding the callback added still governs the - # assistant response. + assert media is not None # narrowed by `use_multimodal` last_user = messages[-1] last_text = last_user["content"] if isinstance(last_user.get("content"), str) else "" + if use_video: + media_item: dict[str, Any] = { + "type": "video", + "video": media[idx]["frames"], + "fps": media[idx]["fps"], + } + else: + media_item = {"type": "image", "image": media[idx]} multimodal_messages = list(messages[:-1]) multimodal_messages.append( { "role": "user", - "content": [ - {"type": "image", "image": images[idx]}, - {"type": "text", "text": last_text}, - ], + "content": [media_item, {"type": "text", "text": last_text}], } ) processor_inputs = self.vlm_processor.apply_chat_template( @@ -4073,31 +4071,48 @@ def generate_reasoner_text( add_generation_prompt=True, return_tensors="pt", ) - # ``Qwen3VLProcessor.apply_chat_template`` strips the - # leading batch dim from ``input_ids`` / ``attention_mask`` - # (see its inline comment); restore it so the inner - # token-level call sees ``[B=1, T_prompt]``. inner_input_ids = processor_inputs["input_ids"].to(device).unsqueeze(0) inner_attention_mask = processor_inputs["attention_mask"].to(device).unsqueeze(0) - inner_pixel_values = processor_inputs["pixel_values"].to(device) # [N_patches,C,H,W] - inner_image_grid_thw = processor_inputs["image_grid_thw"].to(device) # [num_images,3] - out_ids = self.net.generate_reasoner_text( - input_ids=inner_input_ids, - max_new_tokens=max_new_tokens, - pixel_values=inner_pixel_values, - image_grid_thw=inner_image_grid_thw, - attention_mask=inner_attention_mask, - eos_token_id=eos_id, - pad_token_id=pad_id, - do_sample=do_sample, - temperature=temperature if temperature is not None else 1.0, - top_k=top_k, - top_p=top_p, - repetition_penalty=repetition_penalty, - presence_penalty=presence_penalty, - seed=seed, - return_only_new_tokens=True, - ) + if use_video: + inner_pixel_values_videos = processor_inputs["pixel_values_videos"].to(device) + inner_video_grid_thw = processor_inputs["video_grid_thw"].to(device) + out_ids = self.net.generate_reasoner_text( + input_ids=inner_input_ids, + max_new_tokens=max_new_tokens, + pixel_values_videos=inner_pixel_values_videos, + video_grid_thw=inner_video_grid_thw, + attention_mask=inner_attention_mask, + eos_token_id=eos_id, + pad_token_id=pad_id, + do_sample=do_sample, + temperature=temperature if temperature is not None else 1.0, + top_k=top_k, + top_p=top_p, + repetition_penalty=repetition_penalty, + presence_penalty=presence_penalty, + seed=seed, + return_only_new_tokens=True, + ) + else: + inner_pixel_values = processor_inputs["pixel_values"].to(device) # [N_patches,C,H,W] + inner_image_grid_thw = processor_inputs["image_grid_thw"].to(device) # [num_images,3] + out_ids = self.net.generate_reasoner_text( + input_ids=inner_input_ids, + max_new_tokens=max_new_tokens, + pixel_values=inner_pixel_values, + image_grid_thw=inner_image_grid_thw, + attention_mask=inner_attention_mask, + eos_token_id=eos_id, + pad_token_id=pad_id, + do_sample=do_sample, + temperature=temperature if temperature is not None else 1.0, + top_k=top_k, + top_p=top_p, + repetition_penalty=repetition_penalty, + presence_penalty=presence_penalty, + seed=seed, + return_only_new_tokens=True, + ) else: # Text-only path. Pull the system prompt (if any) and # the last user message text out of the messages list, diff --git a/cosmos_framework/model/vfm/parallelize_vlm.py b/cosmos_framework/model/vfm/parallelize_vlm.py index 04dcca3d..c91e48db 100644 --- a/cosmos_framework/model/vfm/parallelize_vlm.py +++ b/cosmos_framework/model/vfm/parallelize_vlm.py @@ -13,9 +13,11 @@ and its meshes — stays in ``vfm/utils/parallelism.py``. """ +import torch.nn as nn from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard from cosmos_framework.utils import log +from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ( PRECISION_TO_TORCH_DTYPE, ParallelismConfig, @@ -24,7 +26,108 @@ from cosmos_framework.utils.vfm.parallelism import ParallelDims -def parallelize( +def _collect_repeated_blocks(inner: nn.Module) -> tuple[list[nn.Module], set[str]]: + """Collect the repeated transformer blocks by their ORIGINAL type name. + + Matches ``inner._no_split_modules`` — the decoder layers (+ vision blocks for + VLMs, e.g. Qwen3-VL ``visual.blocks``). MUST run before any ``fully_shard`` + call: ``fully_shard`` mutates each block's ``__class__`` to a dynamically + created ``FSDP`` type, after which the typename match finds nothing. + In-place ``.compile()`` preserves the type, so :func:`apply_fsdp` can collect + again after :func:`apply_compile` has run. + + Returns the matched blocks and the ``_no_split_modules`` names (the latter is + reported in the no-blocks warning so a misconfiguration is visible). + """ + no_split_names = set(getattr(inner, "_no_split_modules", [])) + blocks = [m for m in inner.modules() if type(m).__name__ in no_split_names] + return blocks, no_split_names + + +def apply_compile(model: HFModel, compile_config: CompileConfig | None) -> None: + """In-place ``torch.compile`` each repeated transformer block, if enabled. + + No-op when ``compile_config`` is ``None`` or ``compile_config.enabled`` is + False. When compile IS enabled but the ``_no_split_modules`` typename scan + matched nothing, logs a warning and returns rather than silently proceeding + uncompiled. Independent of FSDP (peer of :func:`apply_fsdp`), so compile also + applies on the single-GPU / replicate-only path. + + Uses the **in-place** ``nn.Module.compile`` (NOT the ``torch.compile(block)`` + wrapper) on purpose. The wrapper would replace each block with an + ``OptimizedModule`` whose child is ``_orig_mod``, renaming every parameter to + ``...layers.N._orig_mod.*``. That rename would break (a) the suffix-matching + safetensors loader (``HFModel.load_weights``, run AFTER this in + ``_init_vlm`` step g), (b) ``tie_embeddings``, and (c) DCP checkpoint + save/resume (key set must match the prod checkpoint). The in-place variant + compiles ``block._call_impl`` without inserting a wrapper, so ``state_dict`` + keys and module types are unchanged — the typename-based FSDP collection and + all downstream loaders keep working. + + Applied BEFORE ``fully_shard`` (called first in :func:`parallelize`): FSDP2 + ``fully_shard`` SWAPS each block's ``__class__`` to a dynamically-created + ``FSDP`` type, so a typename match against ``_no_split_modules`` + AFTER sharding finds nothing — the blocks must be collected (and compiled) + while their original types are intact. In-place ``.compile()`` sets + ``_compiled_call_impl`` as an instance attribute, which survives the later + ``__class__`` swap; ``nn.Module.__call__`` still routes to it. PyTorch dynamo + skips FSDP2 collective hooks (``torch._dynamo.config.skip_fsdp_hooks`` defaults + True), so the param all-gather / grad reduce-scatter stay in eager AROUND the + compiled forward body (FSDP-outside-compile, the torchtitan ordering) without + the OptimizedModule key-rename. + + Mirrors the per-block strategy of the MoT path + (``parallelize_unified_mot.apply_compile``): the repeated block compiles once + and is reused across all layers. ``fullgraph=False`` because the "cosmos" + attention backend (NATTEN / blackwell-fmha) is an opaque kernel that forces a + graph break; compile still fuses the surrounding pointwise/norm regions (the + ~6.8k tiny elementwise kernels) where the win is. + + Args: + model: HFModel whose ``model._no_split_modules`` blocks (decoder + layers + vision blocks for VLMs) are compiled in place. + compile_config: torch.compile knobs (``compile_dynamic``, + ``max_autotune_pointwise``, ``coordinate_descent_tuning``), + or ``None`` to skip compilation entirely. + """ + if compile_config is None or not compile_config.enabled: + return + + blocks, no_split_names = _collect_repeated_blocks(model.model) + if not blocks: + # Compile was requested but the typename scan matched nothing: either + # _no_split_modules is empty/absent, or its names don't match any module + # type. Without this warning the run would silently proceed uncompiled + # ("0 block(s)" looking like success), so surface the misconfiguration. + log.warning( + "parallelize: compile_config.enabled=True but no _no_split_modules " + f"blocks were found (no_split_names={no_split_names!r}) — torch.compile " + "is a no-op (the model exposes no compilable repeated blocks)" + ) + return + + # max_autotune_pointwise / coordinate_descent_tuning map straight to the + # torch.compile ``options=`` dict (same as the MoT apply_compile). + options: dict[str, bool] = {} + if compile_config.max_autotune_pointwise: + options["max_autotune_pointwise"] = True + if compile_config.coordinate_descent_tuning: + options["coordinate_descent_tuning"] = True + + for block in blocks: + # In-place: sets block._compiled_call_impl; keeps type + state_dict keys. + block.compile( + fullgraph=False, + dynamic=compile_config.compile_dynamic, + options=options or None, + ) + log.info( + f"parallelize: torch.compile applied in-place to {len(blocks)} block(s) " + f"(dynamic={compile_config.compile_dynamic}, options={options or None})" + ) + + +def apply_fsdp( model: HFModel, parallel_dims: ParallelDims, parallelism_config: ParallelismConfig, @@ -40,7 +143,9 @@ def parallelize( - Language models: ``inner.model.layers`` (standard HF LLM structure) - Vision-language models: additionally ``inner.visual.blocks`` (Qwen3-VL) - No-op when FSDP is not needed (single-GPU or replicate-only). + No-op when there is no shard axis (``dp_shard <= 1``): single-GPU, or + replicate-only (``dp_replicate > 1, dp_shard == 1``) which uses DDP outside + this function. Args: model: HFModel instance (``model`` attribute must be on meta or CPU device). @@ -52,9 +157,6 @@ def parallelize( (``"bfloat16"``, ``"float16"``, or ``"float32"``). """ if not parallel_dims.dp_shard_enabled: - # No shard axis: dp_shard <= 1. FSDP2 (fully_shard) has nothing to do. - # For replicate-only (dp_replicate > 1, dp_shard == 1), use DDP outside - # this function. log.info("parallelize: dp_shard <= 1 — skipping FSDP2 wrapping") return @@ -74,13 +176,18 @@ def parallelize( inner = model.model - no_split_names = set(getattr(inner, "_no_split_modules", [])) - wrapped = 0 - for module in reversed(list(inner.modules())): - if type(module).__name__ in no_split_names: - fully_shard(module, **fsdp_kwargs) - wrapped += 1 - log.info(f"Wrapped {wrapped} sub-modules.") + # Collect the repeated blocks by their ORIGINAL type name BEFORE any + # fully_shard call (see _collect_repeated_blocks). apply_compile (if it ran + # first) used in-place compile, which preserves the type names, so this + # re-collection still matches. + blocks, _ = _collect_repeated_blocks(inner) + + # Shard each collected block (reversed = leaf-first), then the root. Iterating + # the collected list (not re-scanning by typename) is robust to the in-place + # compile above and to fully_shard's per-block __class__ swap. + for block in reversed(blocks): + fully_shard(block, **fsdp_kwargs) + log.info(f"Wrapped {len(blocks)} sub-modules.") # Wrap the full inner model to cover remaining parameters # (embed_tokens, final layer norm, lm_head, visual projector stem, etc.) @@ -90,3 +197,32 @@ def parallelize( # "FSDP-2 CPU offload") for how to re-enable it. fully_shard(inner, **fsdp_kwargs) log.info("parallelize: FSDP2 applied to HFModel.model") + + +def parallelize( + model: HFModel, + parallel_dims: ParallelDims, + parallelism_config: ParallelismConfig, + precision: str, + compile_config: CompileConfig | None = None, +) -> None: + """Optimize an HFModel in place: ``torch.compile`` (optional) then FSDP2. + + Mirrors ``parallelize_unified_mot``: :func:`apply_compile` and + :func:`apply_fsdp` are peer passes, each a no-op when its feature is disabled + (compile when ``compile_config`` is ``None``/disabled; FSDP when + ``dp_shard <= 1``). Compile runs FIRST so the in-place block compile is + collected with the blocks' original type names (before ``fully_shard`` swaps + them) and dynamo's ``skip_fsdp_hooks`` keeps the FSDP collectives eager around + the compiled forward body. + + Args: + model: HFModel instance (``model`` attribute on meta or CPU device). + parallel_dims: ParallelDims with meshes already built via + :meth:`ParallelDims.build_meshes`. + parallelism_config: Source of FSDP master dtype (``fsdp_master_dtype``). + precision: FSDP MixedPrecisionPolicy parameter dtype. + compile_config: Optional ``CompileConfig``; ``None``/``enabled=False`` skips compile. + """ + apply_compile(model, compile_config) + apply_fsdp(model, parallel_dims, parallelism_config, precision) diff --git a/cosmos_framework/model/vfm/vlm/__init__.py b/cosmos_framework/model/vfm/reasoner/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/__init__.py rename to cosmos_framework/model/vfm/reasoner/__init__.py diff --git a/cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/__init__.py b/cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/__init__.py rename to cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/__init__.py diff --git a/cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json b/cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json similarity index 100% rename from cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json rename to cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json diff --git a/cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py b/cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py rename to cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py diff --git a/cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/nemotron_3_dense_vl.py b/cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py similarity index 98% rename from cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/nemotron_3_dense_vl.py rename to cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py index aaf42c45..eb91cb93 100644 --- a/cosmos_framework/model/vfm/vlm/nemotron_3_dense_vl/nemotron_3_dense_vl.py +++ b/cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py @@ -17,7 +17,7 @@ from transformers.modeling_rope_utils import dynamic_rope_update from transformers.modeling_utils import PreTrainedModel -from cosmos_framework.model.vfm.vlm.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( +from cosmos_framework.model.vfm.reasoner.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( Nemotron3DenseVLTextConfig, ) diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl/__init__.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl/__init__.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl/__init__.py diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json b/cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json rename to cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json b/cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json rename to cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json b/cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json rename to cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json b/cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json rename to cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl/configs/__init__.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl/configs/__init__.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/__init__.py diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl/configuration_qwen3_vl.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl/configuration_qwen3_vl.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl/configuration_qwen3_vl.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl/configuration_qwen3_vl.py diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl/qwen3_vl.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl/qwen3_vl.py similarity index 99% rename from cosmos_framework/model/vfm/vlm/qwen3_vl/qwen3_vl.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl/qwen3_vl.py index aa92a547..38e784cc 100644 --- a/cosmos_framework/model/vfm/vlm/qwen3_vl/qwen3_vl.py +++ b/cosmos_framework/model/vfm/reasoner/qwen3_vl/qwen3_vl.py @@ -40,16 +40,16 @@ def _default_rope_init(config, device=None, **kwargs): from transformers.utils.deprecation import deprecate_kwarg # Import masking functions from utils for compatibility -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( create_causal_mask, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( get_image_features as _get_image_features, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( get_placeholder_mask as _get_placeholder_mask, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( get_rope_index as _get_rope_index, ) diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl/utils.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl/utils.py similarity index 90% rename from cosmos_framework/model/vfm/vlm/qwen3_vl/utils.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl/utils.py index 0b012fbf..851e8f6e 100644 --- a/cosmos_framework/model/vfm/vlm/qwen3_vl/utils.py +++ b/cosmos_framework/model/vfm/reasoner/qwen3_vl/utils.py @@ -494,8 +494,10 @@ def get_placeholder_mask( def prepare_multimodal_reasoner_inputs( causal_lm: Any, input_ids: torch.Tensor, # [B,T_prompt] - pixel_values: torch.Tensor, # [N_patches,C,H,W] - image_grid_thw: torch.Tensor, # [num_images,3] + pixel_values: torch.Tensor | None = None, # [N_patches,C,H,W] + image_grid_thw: torch.Tensor | None = None, # [num_images,3] + pixel_values_videos: torch.Tensor | None = None, # [N_patches,C,H,W] + video_grid_thw: torch.Tensor | None = None, # [num_videos,3] attention_mask: Optional[torch.Tensor] = None, ) -> tuple[ torch.Tensor, # inputs_embeds [B,T_prompt,hidden_size] @@ -522,11 +524,11 @@ def prepare_multimodal_reasoner_inputs( ``*TextModel.reasoner_forward`` instead of HF's full ``self.language_model(...)`` forward, so HF's ``past_key_values`` / ``cache_position`` lifecycle is replaced by - the AR loop's :class:`ReasonerKVCache` lifecycle. Videos and - dual image+video paths are not supported here; only - ``image_grid_thw`` is consumed — matching the public - ``generate_reasoner_text`` API, which has no - ``pixel_values_videos`` / ``video_grid_thw`` parameters. + the AR loop's :class:`ReasonerKVCache` lifecycle. Either the + image pair (``pixel_values`` + ``image_grid_thw``) or the + video pair (``pixel_values_videos`` + ``video_grid_thw``) is consumed — + not both. The video recipe mirrors the image recipe but routes through + the video placeholder mask and ``video_grid_thw`` rope index. Validation: ``get_placeholder_mask`` raises ``ValueError`` if the number of image placeholder tokens in ``input_ids`` does not match @@ -571,21 +573,41 @@ def prepare_multimodal_reasoner_inputs( mrope_position_deltas: Per-sample rope delta used by the caller to extend positions during decode. """ + if pixel_values is None and pixel_values_videos is None: + raise ValueError( + "prepare_multimodal_reasoner_inputs: exactly one of (pixel_values, image_grid_thw) " + "or (pixel_values_videos, video_grid_thw) must be provided." + ) + is_video = pixel_values_videos is not None inputs_embeds = causal_lm.model.embed_tokens(input_ids).clone() # [B,T_prompt,hidden_size] - pixel_values = pixel_values.to(device=inputs_embeds.device) - image_grid_thw = image_grid_thw.to(device=inputs_embeds.device) - image_embeds, deepstack_visual_embeds = get_image_features(causal_lm, pixel_values, image_grid_thw) - image_embeds = torch.cat(image_embeds, dim=0).to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) - - image_mask, _video_mask = get_placeholder_mask( - causal_lm, - input_ids, - inputs_embeds=inputs_embeds, - image_features=image_embeds, - ) - inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) # [B,T_prompt,hidden_size] - visual_pos_masks = image_mask[..., 0] # [B,T_prompt] + if is_video: + pixel_values_videos = pixel_values_videos.to(device=inputs_embeds.device) + video_grid_thw = video_grid_thw.to(device=inputs_embeds.device) + # get_video_features == get_image_features (same visual tower); reuse the free helper. + video_embeds, deepstack_visual_embeds = get_image_features(causal_lm, pixel_values_videos, video_grid_thw) + video_embeds = torch.cat(video_embeds, dim=0).to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + _image_mask, video_mask = get_placeholder_mask( + causal_lm, + input_ids, + inputs_embeds=inputs_embeds, + video_features=video_embeds, + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) # [B,T_prompt,hidden_size] + visual_pos_masks = video_mask[..., 0] # [B,T_prompt] + else: + pixel_values = pixel_values.to(device=inputs_embeds.device) + image_grid_thw = image_grid_thw.to(device=inputs_embeds.device) + image_embeds, deepstack_visual_embeds = get_image_features(causal_lm, pixel_values, image_grid_thw) + image_embeds = torch.cat(image_embeds, dim=0).to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + image_mask, _video_mask = get_placeholder_mask( + causal_lm, + input_ids, + inputs_embeds=inputs_embeds, + image_features=image_embeds, + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) # [B,T_prompt,hidden_size] + visual_pos_masks = image_mask[..., 0] # [B,T_prompt] deepstack_visual_embeds = [ embed.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) for embed in deepstack_visual_embeds @@ -594,7 +616,8 @@ def prepare_multimodal_reasoner_inputs( position_ids, mrope_position_deltas = get_rope_index( causal_lm, input_ids=input_ids, - image_grid_thw=image_grid_thw, + image_grid_thw=None if is_video else image_grid_thw, + video_grid_thw=video_grid_thw if is_video else None, attention_mask=attention_mask, ) diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl/video_processing_qwen3_vl.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl/video_processing_qwen3_vl.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl/video_processing_qwen3_vl.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl/video_processing_qwen3_vl.py diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/__init__.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl_moe/__init__.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/__init__.py diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json rename to cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json rename to cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configs/__init__.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configs/__init__.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/__init__.py diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl_moe/configuration_qwen3_vl_moe.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe.py similarity index 98% rename from cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe.py index 57944573..9c463204 100644 --- a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe.py +++ b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe.py @@ -7,10 +7,10 @@ import torch.nn as nn from transformers.activations import ACT2FN -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.configuration_qwen3_vl_moe import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( Qwen3VLMoeTextConfig, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_kernels import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_kernels import ( TOKEN_GROUP_ALIGN_SIZE_M, _generate_permute_indices, ) diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_bench.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_bench.py similarity index 94% rename from cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_bench.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_bench.py index 035bf014..48394820 100644 --- a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_bench.py +++ b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_bench.py @@ -5,30 +5,30 @@ Usage: # Default benchmark (forward only, compiled, bf16): - python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench + python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench # Forward + backward: - python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench --backward + python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --backward # Compare grouped_mm vs naive: - python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench --compare + python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --compare # Disable torch.compile: - python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench --no-compile + python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --no-compile # Custom sweep: - python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench --backward \ + python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --backward \ --hidden-size 4096 \ --moe-intermediate-size 1536 # Capture a torch profiler trace (Chrome trace JSON): - python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench --profile + python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --profile # Profile to a custom directory: - python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench --profile --profile-dir ./my_traces + python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --profile --profile-dir ./my_traces # All options: - python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench --help + python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --help """ import itertools @@ -40,10 +40,10 @@ import torch import tyro -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.configuration_qwen3_vl_moe import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( Qwen3VLMoeTextConfig, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe import ( Qwen3VLMoeTextExpertsGroupedMm, Qwen3VLMoeTextExpertsNaive, ) diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_kernels.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_kernels.py similarity index 100% rename from cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_kernels.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_kernels.py diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_test.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_test.py similarity index 93% rename from cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_test.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_test.py index fdb455c2..37430fec 100644 --- a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_test.py +++ b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_test.py @@ -6,10 +6,10 @@ import torch from torch import nn -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.configuration_qwen3_vl_moe import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( Qwen3VLMoeTextConfig, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe import create_text_experts +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe import create_text_experts def run_moe(mod: nn.Module, hidden_states: torch.Tensor, topk_scores: torch.Tensor, expert_indices: torch.Tensor): diff --git a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/qwen3_vl_moe.py b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/qwen3_vl_moe.py similarity index 99% rename from cosmos_framework/model/vfm/vlm/qwen3_vl_moe/qwen3_vl_moe.py rename to cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/qwen3_vl_moe.py index 16144cc8..f25506b4 100644 --- a/cosmos_framework/model/vfm/vlm/qwen3_vl_moe/qwen3_vl_moe.py +++ b/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/qwen3_vl_moe.py @@ -22,21 +22,21 @@ from transformers.utils import TransformersKwargs, is_torchdynamo_compiling from transformers.utils.deprecation import deprecate_kwarg -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( get_image_features as _get_image_features, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( get_placeholder_mask as _get_placeholder_mask, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( get_rope_index as _get_rope_index, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.configuration_qwen3_vl_moe import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( Qwen3VLMoeConfig, Qwen3VLMoeTextConfig, Qwen3VLMoeVisionConfig, ) -from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe import ( +from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe import ( create_text_experts, ) diff --git a/cosmos_framework/model/vfm/tokenizers/uniae/noncausal_4x16x16.py b/cosmos_framework/model/vfm/tokenizers/uniae/noncausal_4x16x16.py index 445f88f4..de4cfbf7 100644 --- a/cosmos_framework/model/vfm/tokenizers/uniae/noncausal_4x16x16.py +++ b/cosmos_framework/model/vfm/tokenizers/uniae/noncausal_4x16x16.py @@ -24,8 +24,6 @@ from cosmos_framework.utils import log from cosmos_framework.utils.distributed import get_rank, sync_model_states from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.model.tokenizer.models.dense_runtime import DenseAutoencoderRuntime -from cosmos_framework.model.tokenizer.models.sparse_autoencoder import AutoencoderKL from cosmos_framework.model.vfm.tokenizers.interface import VideoTokenizerInterface from cosmos_framework.model.vfm.tokenizers.uniae.frame_math import ( get_uniae_latent_num_frames, @@ -34,6 +32,8 @@ normalize_resolution_int_mapping, ) from cosmos_framework.utils.vfm.data_utils import get_vision_data_resolution +from cosmos_framework.model.tokenizer.models.dense_runtime import DenseAutoencoderRuntime +from cosmos_framework.model.tokenizer.models.sparse_autoencoder import AutoencoderKL # S3 architecture config (avoids importing configs/base which pulls in loss deps) _S3_ARCH = dict( diff --git a/cosmos_framework/model/vfm/utils/dcp_loader.py b/cosmos_framework/model/vfm/utils/dcp_loader.py deleted file mode 100644 index 93dd9e28..00000000 --- a/cosmos_framework/model/vfm/utils/dcp_loader.py +++ /dev/null @@ -1,120 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -""" -Checkpoint utilities for VFM models. -Contains utilities for loading model checkpoints with key remapping. -""" - -import time - -import torch -import torch.distributed.checkpoint as dcp -from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE, Metadata - -from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader -from cosmos_framework.utils import log - - -class RenameLoadPlanner(dcp.DefaultLoadPlanner): - """ - RenameLoadPlanner that does two renaming operations during pretrained model load. - 1. Renames model parameters from "model." to "model.language_model." if the core model is a VLM. - 2. Renames model parameters from "_orig_mod." to "." since the pretrained model does not have torch compiled modules. - """ - - def set_up_planner( - self, - state_dict: STATE_DICT_TYPE, - metadata: Metadata | None = None, - is_coordinator: bool = False, - ) -> None: - state_dict = remove_orig_mode_from_state_dict(state_dict) - - # Renames model parameters from "model." to "model.language_model." if the core model is a VLM. - # This is necessary because the checkpoint is saved with "model.language_model." for a VLM. - # If the core model is an LLM, this renaming is not necessary. - has_language_model = any("model.language_model." in key for key in metadata.state_dict_metadata) - if has_language_model: - state_dict = insert_language_model_in_state_dict(state_dict) - - # Perform more error checking to ensure the checkpoint is valid. If the keys are missing, - # then the missing keys should be from the generation pathway. All keys from the - # understanding pathway must be present in the checkpoint. Additionally, for 2B and 4B - # dense Qwen VLMs, the `lm_head.weight` key is not present in the checkpoint. For these - # models, the input embedding and generation layer share the same params due to - # `tie_word_embeddings` being set to True in the configs. For the 0.6B LLM, 8B and 32B dense - # VLMs, and the 30B and 235B MoE VLMs, the `lm_head.weight` key is present in the - # checkpoint. - for key in state_dict: - if key not in metadata.state_dict_metadata and "_moe_gen" not in key and "lm_head.weight" not in key: - raise ValueError(f"Missing key: {key} in checkpoint.") - - # If the keys are unexpected, then the unexpected keys should be from the visual part of the - # VLM. All keys in the language part should be used by the Cosmos3 model. - for key in metadata.state_dict_metadata: - if key not in state_dict and "model.visual." not in key: - raise ValueError(f"Unexpected key: {key} in checkpoint.") - - super().set_up_planner( - state_dict=state_dict, - metadata=metadata, - is_coordinator=is_coordinator, - ) - - -def remove_orig_mode_from_state_dict(state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE: - """Renames model parameters from "_orig_mod." to "." since the pretrained model does not have torch compiled modules.""" - return {key.replace("_orig_mod.", ""): value for key, value in state_dict.items()} - - -def insert_language_model_in_state_dict(state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE: - """Renames model parameters from "model." to "model.language_model." if the core model is a VLM.""" - new_state_dict = {} - for key, value in state_dict.items(): - if key.startswith("model."): - new_state_dict[key.replace("model.", "model.language_model.")] = value - else: - new_state_dict[key] = value - return new_state_dict - - -def load_language_model( - model: torch.nn.Module, checkpoint_path: str, credential_path: str, enable_gcs_patch_in_boto3: bool = False -) -> None: - """ - Universal language model loading function using DCP format. - Handles key remapping for "model.language_model." -> "model." by default. - - Args: - model: The language model to load weights into - checkpoint_path: Path to checkpoint (must end with .safetensors or dcp) - credential_path: Path to S3 credentials - enable_gcs_patch_in_boto3: Whether to enable GCS patch in boto3 for DCP loading from GCS - """ - start_time = time.time() - if not checkpoint_path.strip("/").endswith("dcp"): - raise ValueError(f"Checkpoint path {checkpoint_path} must end with dcp") - - log.info(f"Loading language model weights in DCP format from: {checkpoint_path}") - state_dict = model.state_dict() - - # Choose storage reader - if checkpoint_path.startswith("s3://"): - storage_reader = S3StorageReader( - credential_path=credential_path, - path=checkpoint_path, - enable_gcs_patch_in_boto3=enable_gcs_patch_in_boto3, - ) - else: - storage_reader = dcp.FileSystemReader(checkpoint_path) - - # Load checkpoint - dcp.load( - state_dict=state_dict, - storage_reader=storage_reader, - planner=RenameLoadPlanner(allow_partial_load=False), - ) - - log.info(f"Successfully loaded language model from {checkpoint_path}") - log.info(f"Time taken to load language model: {time.time() - start_time} seconds") diff --git a/cosmos_framework/model/vfm/utils/safetensors_loader.py b/cosmos_framework/model/vfm/utils/safetensors_loader.py index 15b87744..22c2d6fc 100644 --- a/cosmos_framework/model/vfm/utils/safetensors_loader.py +++ b/cosmos_framework/model/vfm/utils/safetensors_loader.py @@ -1016,7 +1016,11 @@ def load_language_model( # `tie_word_embeddings` being set to True in the configs. For the 0.6B LLM, 8B and 32B dense # VLMs, and the 30B and 235B MoE VLMs, the `lm_head.weight` key is present in the # checkpoint. - real_keys_missing = {k for k in keys_missing if "_moe_gen" not in k} + # Keys that are expected to be absent from the HF backbone checkpoint: + # - "_moe_gen": generation-pathway parameters, initialised by init_moe() + # - "k_norm_und_for_gen": new und-K normalisation for gen cross-attention, + # only present when use_und_k_norm_for_gen=True; always init'd from scratch + real_keys_missing = {k for k in keys_missing if "_moe_gen" not in k and "k_norm_und_for_gen" not in k} if real_keys_missing: raise ValueError( f"load_language_model: {len(real_keys_missing)} required model " @@ -1209,9 +1213,8 @@ def load_vfm_model( (``language_model.lm_head.*``) and — for Qwen3-VL-based variants — the visual encoder (ViT) under ``language_model.visual.*``; - the VFM-specific top-level parameters: ``time_embedder.*``, - ``vae2llm.*`` / ``llm2vae.*``, ``latent_pos_embed.*`` (when present), - plus the optional ``action2llm.*`` / ``llm2action.*`` / - ``action_pos_embed.*`` / ``action_modality_embed`` and + ``vae2llm.*`` / ``llm2vae.*``, plus the optional + ``action2llm.*`` / ``llm2action.*`` / ``action_modality_embed`` and ``sound2llm.*`` / ``llm2sound.*`` / ``sound_modality_embed`` heads. Checkpoint keys are interpreted in the **native VFM state-dict @@ -1269,7 +1272,7 @@ def load_vfm_model( partial checkpoint, e.g. when warm-starting from a checkpoint that has only the language tower and you want to keep the freshly-init'd VFM heads (pass ``r"^(time_embedder|vae2llm|llm2vae|"`` - ``r"latent_pos_embed|action[^.]*|llm2action|sound[^.]*|llm2sound)\..*"``). + ``r"action[^.]*|llm2action|sound[^.]*|llm2sound)\..*"``). - Missing model keys whose name contains ``_moe_gen`` are silently tolerated regardless of ``skip_patterns`` — they are populated downstream by diff --git a/cosmos_framework/model/vfm/vlm_model.py b/cosmos_framework/model/vfm/vlm_model.py index e3d311a4..ec62b62f 100644 --- a/cosmos_framework/model/vfm/vlm_model.py +++ b/cosmos_framework/model/vfm/vlm_model.py @@ -28,14 +28,14 @@ from cosmos_framework.utils.lazy_config import instantiate from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.utils import log -from cosmos_framework.model.vfm.algorithm.loss.cross_entropy import cross_entropy_loss +from cosmos_framework.model.vfm.algorithm.loss.cross_entropy import cross_entropy_loss, weighted_cross_entropy_loss from cosmos_framework.configs.base.defaults.parallelism import PRECISION_TO_TORCH_DTYPE -from cosmos_framework.configs.base.vlm.defaults.policy_config import VLMModelConfig +from cosmos_framework.configs.base.reasoner.defaults.policy_config import VLMModelConfig from cosmos_framework.model.vfm.hf_model import HFModel from cosmos_framework.model.vfm.parallelize_vlm import parallelize from cosmos_framework.utils.vfm.parallelism import ParallelDims -from cosmos_framework.utils.vfm.vlm.constant import IGNORE_INDEX -from cosmos_framework.utils.vfm.vlm.create_position_ids import get_position_ids +from cosmos_framework.utils.vfm.reasoner.constant import IGNORE_INDEX +from cosmos_framework.utils.vfm.reasoner.create_position_ids import get_position_ids # Model-type dispatch sets. Using hf_config.model_type (stable HF-defined string) # rather than backbone.model_name avoids the brittleness of substring-matching a local @@ -271,13 +271,24 @@ def __init__(self, config: VLMModelConfig, checkpoint): if self.parallel_dims.cp_enabled: cp_group = self.parallel_dims.cp_mesh.get_group() - self._loss_fn = partial( - cross_entropy_loss, - loss_scaling_factor=1.0, - dp_group=dp_group, - cp_group=cp_group, - ignore_index=IGNORE_INDEX, - ) + if config.policy.use_weighted_ce: + log.info(f"Using weighted CE loss with exponent={config.policy.weighted_ce_exponent}") + self._loss_fn = partial( + weighted_cross_entropy_loss, + exponent=config.policy.weighted_ce_exponent, + loss_scaling_factor=1.0, + dp_group=dp_group, + cp_group=cp_group, + ignore_index=IGNORE_INDEX, + ) + else: + self._loss_fn = partial( + cross_entropy_loss, + loss_scaling_factor=1.0, + dp_group=dp_group, + cp_group=cp_group, + ignore_index=IGNORE_INDEX, + ) def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: """Initialize VLM without the legacy ModelRegistry (Phase 2+). @@ -292,7 +303,7 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: g. Load pretrain weights into sharded CUDA tensors. h. Apply gradient checkpointing if configured. """ - from cosmos_framework.utils.vfm.vlm.pretrained_models_downloader import ( + from cosmos_framework.utils.vfm.reasoner.pretrained_models_downloader import ( maybe_download_hf_model_from_s3, ) @@ -372,13 +383,17 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: "Use dp_shard > 1 for FSDP2. DDP support is planned for Phase 3." ) - # ── d. Apply FSDP2 ── + # ── d. Apply FSDP2 (+ optional torch.compile of the repeated blocks) ── + # config.compile is threaded through so model.config.compile.enabled=True + # actually compiles each block in place (was previously a dead config on + # the VLM path — only the MoT path consumed it). See parallelize_vlm. if torch.distributed.is_initialized(): parallelize( hf_model, parallel_dims, config.parallelism, config.precision, + compile_config=config.compile, ) # ── e. Materialize meta tensors on CUDA ── diff --git a/cosmos_framework/utils/config.py b/cosmos_framework/utils/config.py index 5078fbc2..cf2f152f 100644 --- a/cosmos_framework/utils/config.py +++ b/cosmos_framework/utils/config.py @@ -271,6 +271,17 @@ class CheckpointConfig: # for dcp, whether to use async mode dcp_async_mode_enabled: bool = False + # For dcp load, whether to deduplicate redundant storage reads of replicated state and + # broadcast the data over the (NCCL) device mesh instead. Two replication patterns waste + # reads at large world sizes: (1) fully-replicated entries (e.g. optimizer scalar `step` + # tensors) are saved on global rank 0 yet read by *every* rank — a single-object S3 hotspot; + # (2) sharded DTensors are byte-identical across their mesh's replicate dims (e.g. HSDP's + # `dp_replicate`), so each shard file is read by `dp_replicate` ranks. When enabled, each leaf + # is read by exactly one rank — global rank 0 for fully-replicated non-DTensor leaves, and + # local rank 0 along the DTensor's own replicate mesh dims otherwise — then broadcast to the + # rest, cutting reads from O(world_size) to O(dp_shard). + dcp_load_dedup: bool = False + # Configs for saving the checkpoints to object store. save_to_object_store: ObjectStoreConfig = attrs.field(factory=ObjectStoreConfig) @@ -517,12 +528,7 @@ def validate(self) -> None: assert self.job.name != "" -def load_config( - config_path: str, - opts: list[str], - enable_one_logger: bool = False, -) -> Config: - """Load a config from a ``.yaml`` or ``.py`` path and apply ``opts``.""" +def load_config(config_path: str, opts: list[str], enable_one_logger: bool = False) -> Config: from cosmos_framework.utils.serialization import from_yaml, load_callable t1 = time.monotonic_ns() @@ -554,11 +560,7 @@ def load_config( return config -def _load_py_config( - config_path: str, - opts: list[str], - validate: bool = True, -) -> Config: +def _load_py_config(config_path: str, opts: list[str], validate: bool = True) -> Config: # NOTE: circular dependency from cosmos_framework.utils.config_helper import get_config_module, override diff --git a/cosmos_framework/utils/env_parsers/cred_env_parser.py b/cosmos_framework/utils/env_parsers/cred_env_parser.py index 68e7ede0..04810d9c 100644 --- a/cosmos_framework/utils/env_parsers/cred_env_parser.py +++ b/cosmos_framework/utils/env_parsers/cred_env_parser.py @@ -9,7 +9,7 @@ class CredentialEnvParser(EnvParser): APP_ENV = String(default="") PROD_FT_AWS_CREDS_ACCESS_KEY_ID = String(default="") PROD_FT_AWS_CREDS_SECRET_ACCESS_KEY = String(default="") - PROD_FT_AWS_CREDS_ENDPOINT_URL = String(default="https://invalid_url") + PROD_FT_AWS_CREDS_ENDPOINT_URL = String(default="https://s3.us-west-2.amazonaws.com") PROD_FT_AWS_CREDS_REGION_NAME = String(default="us-west-2") PROD_S3_CHECKPOINT_ACCESS_KEY_ID = String(default="") @@ -33,7 +33,7 @@ class CredentialEnvParser(EnvParser): PROD_TEAM_DIR_REGION_NAME = String(default="") PICASSO_AUTH_MODEL_REGISTRY_API_KEY = String(default="") - PICASSO_API_ENDPOINT_URL = String(default="https://invalid_url") + PICASSO_API_ENDPOINT_URL = String(default="https://invalid") CRED_ENVS = CredentialEnvParser() diff --git a/cosmos_framework/utils/vfm/flash_attn.py b/cosmos_framework/utils/vfm/flash_attn.py index 078dab91..6d0fd4c7 100644 --- a/cosmos_framework/utils/vfm/flash_attn.py +++ b/cosmos_framework/utils/vfm/flash_attn.py @@ -4,7 +4,7 @@ """Flash attention initialization for the vfm/ unified VLM training path. This module replaces `cosmos_rl.policy.kernel.modeling_utils.init_flash_attn_meta` -in VLMModel.__init__, removing the last cosmos_rl import from projects/cosmos3/vfm/. +in VLMModel.__init__, removing the last cosmos_rl import from projects/cosmos3/cosmos3/. Why a stub for the FlashAttnMeta singleton part? HFModel uses HuggingFace's flash_attention_2 implementation diff --git a/cosmos_framework/utils/vfm/fused_adam.py b/cosmos_framework/utils/vfm/fused_adam.py index b7ac17dd..72063b7d 100644 --- a/cosmos_framework/utils/vfm/fused_adam.py +++ b/cosmos_framework/utils/vfm/fused_adam.py @@ -177,7 +177,12 @@ def step(self, closure=None, grads=None, output_params=None, scale=None, grad_no for p, p_master in zip(group["params"], group_master["params"]): if p.grad is None: continue - if p.grad.data.is_sparse: + # Unwrap DTensor grads to their local shard before checking + # sparsity. Touching ``p.grad.data`` directly on a DTensor + # dispatches ``aten.detach`` through ``__torch_dispatch__``, + # which operates on a partially-built DTensor shell and raises + # ``'DTensor' object has no attribute '_local_tensor'``. + if get_local_tensor_if_DTensor(p.grad).is_sparse: raise RuntimeError( "FusedAdam does not support sparse gradients, please consider SparseAdam instead" ) diff --git a/cosmos_framework/utils/vfm/model_loader.py b/cosmos_framework/utils/vfm/model_loader.py index 6e6a0ddc..fd8b9560 100644 --- a/cosmos_framework/utils/vfm/model_loader.py +++ b/cosmos_framework/utils/vfm/model_loader.py @@ -397,6 +397,9 @@ def load_model_from_checkpoint( if hasattr(config.model.config, "load_teacher_weights"): log.info("Setting load_teacher_weights=False for inference to skip teacher checkpoint download.") config.model.config.load_teacher_weights = False + if getattr(config.model.config, "student_load_from", None) is not None: + log.info("Setting student_load_from=None for inference to skip train-time student warm-start download.") + config.model.config.student_load_from = None if ( config.model.config.exclude_reasoner_weights_from_checkpoint diff --git a/cosmos_framework/utils/vfm/monkey_patch.py b/cosmos_framework/utils/vfm/monkey_patch.py index 1320d0b6..5e694e7c 100644 --- a/cosmos_framework/utils/vfm/monkey_patch.py +++ b/cosmos_framework/utils/vfm/monkey_patch.py @@ -18,12 +18,13 @@ def patch_qwen3_vl_forward(model): """Monkey-patch a ``Qwen3VLModel`` **instance's** forward: - **Dummy visual forward for pure-text batches**: Under FSDP, every rank - must call ``self.visual(...)`` each forward step so that collective - all-gather operations stay in sync. When a batch contains only text, - a lightweight dummy image (16x16 zeros) is pushed through the full - ViT -> merger -> deepstack pipeline, then outputs are sliced to ``[0:0]`` - so they carry ``grad_fn`` but contribute no features. + **Single visual forward per batch**: Under FSDP, every rank must call + ``self.visual(...)`` the same number of times each forward step so that + collective all-gather operations stay in sync. Image and video inputs are + encoded together in one visual call. When a batch contains only text, a + lightweight dummy image (16x16 zeros) is pushed through the full ViT -> + merger -> deepstack pipeline, then outputs are sliced to ``[0:0]`` so + they carry ``grad_fn`` but contribute no features. Args: model: The ``Qwen3VLModel`` instance (i.e. ``model.model.model`` when the outer model is ``HFModel``). @@ -68,49 +69,80 @@ def patched_forward( image_mask = None video_mask = None + deepstack_image_embeds = None + deepstack_video_embeds = None + + visual_pixel_values_list: list[torch.Tensor] = [] + visual_grid_thw_list: list[torch.Tensor] = [] + image_embed_len = 0 + video_embed_len = 0 + has_image = pixel_values is not None + has_video = pixel_values_videos is not None + + if has_image: + if image_grid_thw is None: + raise ValueError("image_grid_thw must be provided when pixel_values is provided") + visual_pixel_values_list.append(pixel_values) + visual_grid_thw_list.append(image_grid_thw) + image_embed_len = int((image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).sum().item()) + + if has_video: + if video_grid_thw is None: + raise ValueError("video_grid_thw must be provided when pixel_values_videos is provided") + visual_pixel_values_list.append(pixel_values_videos) + visual_grid_thw_list.append(video_grid_thw) + video_embed_len = int((video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).sum().item()) + + if has_image or has_video: + visual_pixel_values = torch.cat(visual_pixel_values_list, dim=0).type(self.visual.dtype) # [N_patch,D] + visual_grid_thw = torch.cat(visual_grid_thw_list, dim=0) # [N_media,3] + visual_embeds, deepstack_visual_feature_lists = self.visual( + visual_pixel_values, grid_thw=visual_grid_thw + ) # visual_embeds: [N_visual,C] + image_embeds = visual_embeds[:image_embed_len] # [N_image_visual,C] + video_embeds = visual_embeds[image_embed_len : image_embed_len + video_embed_len] # [N_video_visual,C] + deepstack_image_embeds = [ + deepstack_visual_embeds[:image_embed_len] for deepstack_visual_embeds in deepstack_visual_feature_lists + ] # each: [N_image_visual,C] + deepstack_video_embeds = [ + deepstack_visual_embeds[image_embed_len : image_embed_len + video_embed_len] + for deepstack_visual_embeds in deepstack_visual_feature_lists + ] # each: [N_video_visual,C] + + if has_image: + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) # [N_image_visual,C] + image_mask, _ = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) # [B,N_token,C] - if pixel_values is not None: - image_embeds, deepstack_image_embeds = self.get_image_features( - pixel_values, image_grid_thw - ) # list of (T, C) tensors per batch element - image_embeds = torch.cat(image_embeds, dim=0).to( - inputs_embeds.device, inputs_embeds.dtype - ) # concat along token dim across the batch - image_mask, _ = self.get_placeholder_mask( - input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds - ) - inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) - - if pixel_values_videos is not None: - video_embeds, deepstack_video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) - video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) - _, video_mask = self.get_placeholder_mask( - input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds - ) - inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + if has_video: + video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) # [N_video_visual,C] + _, video_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) # [B,N_token,C] # Dummy visual forward for text-only data - if pixel_values is None and pixel_values_videos is None: + else: dummy_h, dummy_w = 16, 16 dummy_pixels = torch.zeros( dummy_h * dummy_w, self.visual.config.temporal_patch_size * self.visual.config.patch_size**2 * 3, device=inputs_embeds.device, dtype=self.visual.dtype, - ) - dummy_thw = torch.tensor([[1, dummy_h, dummy_w]], device=inputs_embeds.device) - image_embeds, deepstack_image_embeds = self.get_image_features(dummy_pixels, dummy_thw) - image_embeds = [e[0:0] for e in image_embeds] - deepstack_image_embeds = [e[0:0] for e in deepstack_image_embeds] + ) # [N_dummy_patch,D] + dummy_thw = torch.tensor([[1, dummy_h, dummy_w]], device=inputs_embeds.device) # [1,3] + image_embeds, deepstack_image_embeds = self.visual(dummy_pixels, grid_thw=dummy_thw) + image_embeds = image_embeds[0:0] # [0,C] + deepstack_image_embeds = [e[0:0] for e in deepstack_image_embeds] # each: [0,C] # no-op to mask scatter empty embeddings into inputs to preserve computation graph - image_embeds = torch.cat(image_embeds, dim=0).to( - inputs_embeds.device, inputs_embeds.dtype - ) # concat along token dim across the batch + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) # [0,C] image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds ) - inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) # [B,N_token,C] visual_pos_masks = None deepstack_visual_embeds = None @@ -200,4 +232,4 @@ def patched_forward( # Replace the forward method model.forward = patched_forward.__get__(model, type(model)) - log.critical(f"Patched {type(model).__name__} instance forward with pure-text dummy forward") + log.critical(f"Patched {type(model).__name__} instance forward with one visual call per forward") diff --git a/cosmos_framework/utils/vfm/vlm/__init__.py b/cosmos_framework/utils/vfm/reasoner/__init__.py similarity index 100% rename from cosmos_framework/utils/vfm/vlm/__init__.py rename to cosmos_framework/utils/vfm/reasoner/__init__.py diff --git a/cosmos_framework/utils/vfm/vlm/constant.py b/cosmos_framework/utils/vfm/reasoner/constant.py similarity index 100% rename from cosmos_framework/utils/vfm/vlm/constant.py rename to cosmos_framework/utils/vfm/reasoner/constant.py diff --git a/cosmos_framework/utils/vfm/vlm/create_position_ids.py b/cosmos_framework/utils/vfm/reasoner/create_position_ids.py similarity index 100% rename from cosmos_framework/utils/vfm/vlm/create_position_ids.py rename to cosmos_framework/utils/vfm/reasoner/create_position_ids.py diff --git a/cosmos_framework/utils/vfm/vlm/flop_calculator.py b/cosmos_framework/utils/vfm/reasoner/flop_calculator.py similarity index 100% rename from cosmos_framework/utils/vfm/vlm/flop_calculator.py rename to cosmos_framework/utils/vfm/reasoner/flop_calculator.py diff --git a/cosmos_framework/utils/vfm/vlm/pretrained_models_downloader.py b/cosmos_framework/utils/vfm/reasoner/pretrained_models_downloader.py similarity index 100% rename from cosmos_framework/utils/vfm/vlm/pretrained_models_downloader.py rename to cosmos_framework/utils/vfm/reasoner/pretrained_models_downloader.py diff --git a/cosmos_framework/utils/vlm/configs_defaults/checkpointer.py b/cosmos_framework/utils/vlm/configs_defaults/checkpointer.py index 29f9f270..b7bfab67 100644 --- a/cosmos_framework/utils/vlm/configs_defaults/checkpointer.py +++ b/cosmos_framework/utils/vlm/configs_defaults/checkpointer.py @@ -19,15 +19,21 @@ s3_object_store = config.ObjectStoreConfig( enabled=True, - credentials="credentials/s3_training.secret", + credentials="credentials/s3_checkpoint.secret", bucket="bucket4", ) +s3_east2_object_store = config.ObjectStoreConfig( + enabled=True, + credentials="credentials/s3_east2_checkpoint.secret", + bucket="bucket", +) + # Permanent store for initial HF model weights on AWS (different bucket from training checkpoints). # Used by train.py to download pretrained weights; NOT used by the checkpointer for auto-resume. aws_load_from_object_store_permanent = config.ObjectStoreConfig( enabled=True, - credentials="credentials/s3_training.secret", + credentials="credentials/s3_checkpoint.secret", bucket="nv-cosmos-vlm", ) @@ -45,7 +51,7 @@ gcp_object_store = config.ObjectStoreConfig( enabled=True, - credentials="credentials/gcp_training.secret", + credentials="credentials/gcp_checkpoint.secret", bucket="bucket1", ) @@ -53,7 +59,7 @@ # Used by train.py to download pretrained weights; NOT used by the checkpointer for auto-resume. gcp_load_from_object_store_permanent = config.ObjectStoreConfig( enabled=True, - credentials="credentials/gcp_training.secret", + credentials="credentials/gcp_checkpoint.secret", bucket="bucket0", ) @@ -71,6 +77,13 @@ broadcast_via_filesystem=True, ) +CHECKPOINT_S3_EAST2 = CheckpointConfig( + save_to_object_store=s3_east2_object_store, + load_from_object_store=s3_east2_object_store, + save_iter=2000, + broadcast_via_filesystem=True, +) + CHECKPOINT_NEB_EU = CheckpointConfig( save_to_object_store=neb_eu_object_store, load_from_object_store=neb_eu_object_store, @@ -90,6 +103,7 @@ def register_checkpoint(): cs = ConfigStore.instance() cs.store(group="checkpoint", package="checkpoint", name="pdx", node=CHECKPOINT_PDX) cs.store(group="checkpoint", package="checkpoint", name="s3", node=CHECKPOINT_S3) + cs.store(group="checkpoint", package="checkpoint", name="s3_east2", node=CHECKPOINT_S3_EAST2) cs.store(group="checkpoint", package="checkpoint", name="neb_eu", node=CHECKPOINT_NEB_EU) cs.store(group="checkpoint", package="checkpoint", name="gcp", node=CHECKPOINT_GCP) diff --git a/examples/integration/net_level.py b/examples/integration/net_level.py index 6890299c..1c75e747 100644 --- a/examples/integration/net_level.py +++ b/examples/integration/net_level.py @@ -112,7 +112,7 @@ from cosmos_framework.data.vfm.sequence_packing import SequencePlan, build_sequence_plans_from_data_batch from cosmos_framework.inference.args import DEFAULT_CHECKPOINT from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption def _load_omni_model(*, config_dir_arg: str | None): diff --git a/examples/integration/trainer_level_inference.py b/examples/integration/trainer_level_inference.py index 7ef01310..6b4b81bf 100644 --- a/examples/integration/trainer_level_inference.py +++ b/examples/integration/trainer_level_inference.py @@ -34,7 +34,7 @@ cosmos_framework.inference.{args,inference} → OmniSampleOverrides + get_sample_data (T2I/T2V only) cosmos_framework.data.vfm.{action,sequence_packing} → SequencePlan helpers (action/sound) - cosmos_framework.model.vfm.vlm.qwen3_vl.utils.tokenize_caption + cosmos_framework.model.vfm.reasoner.qwen3_vl.utils.tokenize_caption model.generate_samples_from_batch(batch, seed) → THE inference call (CFG + sampler) model.decode(latent) → VAE decode @@ -71,7 +71,7 @@ from cosmos_framework.inference.args import DEFAULT_CHECKPOINT, OmniSampleOverrides from cosmos_framework.inference.inference import get_sample_data from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption from cosmos_framework.tools.visualize.video import save_img_or_video diff --git a/examples/integration/trainer_level_training.py b/examples/integration/trainer_level_training.py index 10d50b47..2656bce4 100644 --- a/examples/integration/trainer_level_training.py +++ b/examples/integration/trainer_level_training.py @@ -35,7 +35,7 @@ cosmos_framework.inference.model.Cosmos3OmniModel → model class (random-init in this demo; use `.from_pretrained_dcp(...)` for real weights) cosmos_framework.inference.common.init.init_script → 1-line torch.distributed init - cosmos_framework.model.vfm.vlm.qwen3_vl.utils.tokenize_caption + cosmos_framework.model.vfm.reasoner.qwen3_vl.utils.tokenize_caption → text tokenizer (modelling pkg) model.training_step(batch, iteration) → THE training step (flow-matching loss) model.config.{action_gen,sound_gen,vision_gen,…} → modality flags @@ -126,7 +126,7 @@ from cosmos_framework.data.vfm.sequence_packing import SequencePlan from cosmos_framework.inference.args import DEFAULT_CHECKPOINT from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel -from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption def _load_omni_model(*, config_dir_arg: str | None): From 9b02fa8541f25e8b941e0df10db7b665e3cad7f8 Mon Sep 17 00:00:00 2001 From: Liangkai Zhang Date: Thu, 2 Jul 2026 00:01:25 -0700 Subject: [PATCH 03/26] [Cosmos3 OSS]Add more action datasets (#72) Dataset support added: 1. Fractal (FractalLeRobotDataset class) 2. RoboMind Frank dual arm (through existing RoboMINDFrankaDataset class) 3. RoboMind UR (RoboMINDURDataset class) Other changes: 1. Add corresponding stats files for the newly added datasets 2. Folder structure refactor (minor) --------- Co-authored-by: lfengad --- .../data/vfm/action/agibot_spec.py | 2 +- .../data/vfm/action/datasets/__init__.py | 4 + .../agibotworld_beta_lerobot_dataset.py | 2 +- .../data/vfm/action/datasets/base_dataset.py | 9 +- .../datasets/bridge_orig_lerobot_dataset.py | 2 +- .../action/datasets/droid_lerobot_dataset.py | 2 +- .../datasets/fractal_lerobot_dataset.py | 188 ++++++++++ .../datasets/robomind_franka_dataset.py | 139 ++++++-- .../action/datasets/robomind_ur_dataset.py | 203 +++++++++++ .../datasets/stats/umi_lerobot_stats.json | 4 - .../action/datasets/umi_lerobot_dataset.py | 2 +- .../agibotworld_beta_lerobot_stats.json | 0 .../bridge_orig_lerobot_stats.json | 0 .../droid_lerobot_stats.json | 0 .../fractal_lerobot_stats.json | 4 + .../robomind_franka_dual_stats.json} | 0 .../robomind_franka_stats.json | 4 + .../normalizer_stats/robomind_ur_stats.json | 4 + .../normalizer_stats/umi_lerobot_stats.json | 4 + .../G1_omnipicker_calibrated.urdf | 0 .../action/robot_assets/ur5e_robotiq_2f85.xml | 323 ++++++++++++++++++ 21 files changed, 850 insertions(+), 46 deletions(-) create mode 100644 cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py create mode 100644 cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py delete mode 100644 cosmos_framework/data/vfm/action/datasets/stats/umi_lerobot_stats.json rename cosmos_framework/data/vfm/action/{datasets/stats => normalizer_stats}/agibotworld_beta_lerobot_stats.json (100%) rename cosmos_framework/data/vfm/action/{datasets/stats => normalizer_stats}/bridge_orig_lerobot_stats.json (100%) rename cosmos_framework/data/vfm/action/{datasets/stats => normalizer_stats}/droid_lerobot_stats.json (100%) create mode 100644 cosmos_framework/data/vfm/action/normalizer_stats/fractal_lerobot_stats.json rename cosmos_framework/data/vfm/action/{datasets/stats/robomind_franka_stats.json => normalizer_stats/robomind_franka_dual_stats.json} (100%) create mode 100644 cosmos_framework/data/vfm/action/normalizer_stats/robomind_franka_stats.json create mode 100644 cosmos_framework/data/vfm/action/normalizer_stats/robomind_ur_stats.json create mode 100644 cosmos_framework/data/vfm/action/normalizer_stats/umi_lerobot_stats.json rename cosmos_framework/data/vfm/action/{urdf_visualizer => robot_assets}/G1_omnipicker_calibrated.urdf (100%) create mode 100644 cosmos_framework/data/vfm/action/robot_assets/ur5e_robotiq_2f85.xml diff --git a/cosmos_framework/data/vfm/action/agibot_spec.py b/cosmos_framework/data/vfm/action/agibot_spec.py index 0abfe22a..853afd58 100644 --- a/cosmos_framework/data/vfm/action/agibot_spec.py +++ b/cosmos_framework/data/vfm/action/agibot_spec.py @@ -126,4 +126,4 @@ def get_agibot_world_kind(embodiment_type: str) -> AgibotWorldKind: def get_agibot_world_urdf_path() -> Path: """Return the committed AgiBot G1 omnipicker URDF path.""" - return Path(__file__).resolve().parent / "urdf_visualizer" / AGIBOT_WORLD_URDF_FILENAME + return Path(__file__).resolve().parent / "robot_assets" / AGIBOT_WORLD_URDF_FILENAME diff --git a/cosmos_framework/data/vfm/action/datasets/__init__.py b/cosmos_framework/data/vfm/action/datasets/__init__.py index 64b1b278..b94b1440 100644 --- a/cosmos_framework/data/vfm/action/datasets/__init__.py +++ b/cosmos_framework/data/vfm/action/datasets/__init__.py @@ -12,7 +12,9 @@ from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset from cosmos_framework.data.vfm.action.datasets.bridge_orig_lerobot_dataset import BridgeOrigLeRobotDataset from cosmos_framework.data.vfm.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset +from cosmos_framework.data.vfm.action.datasets.fractal_lerobot_dataset import FractalLeRobotDataset from cosmos_framework.data.vfm.action.datasets.robomind_franka_dataset import RoboMINDFrankaDataset +from cosmos_framework.data.vfm.action.datasets.robomind_ur_dataset import RoboMINDURDataset from cosmos_framework.data.vfm.action.datasets.umi_lerobot_dataset import UMILeRobotDataset __all__ = [ @@ -20,6 +22,8 @@ "AgiBotWorldBetaLeRobotDataset", "BridgeOrigLeRobotDataset", "DROIDLeRobotDataset", + "FractalLeRobotDataset", "RoboMINDFrankaDataset", + "RoboMINDURDataset", "UMILeRobotDataset", ] diff --git a/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py index f95feea8..ce22773c 100644 --- a/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py @@ -39,7 +39,7 @@ _ROBOT_POSITION_KEY = "observation.states.robot.position" _ROBOT_ORIENTATION_KEY = "observation.states.robot.orientation" -_NORMALIZER_PATH = Path(__file__).parent / "stats/agibotworld_beta_lerobot_stats.json" +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/agibotworld_beta_lerobot_stats.json" def _split_task_for_caption(task: str) -> tuple[str, str]: diff --git a/cosmos_framework/data/vfm/action/datasets/base_dataset.py b/cosmos_framework/data/vfm/action/datasets/base_dataset.py index 2c9c4cb2..d7cdcc23 100644 --- a/cosmos_framework/data/vfm/action/datasets/base_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/base_dataset.py @@ -12,6 +12,7 @@ from typing import Any import numpy as np +import pandas as pd import pyarrow.parquet as pq import torch from torch.utils.data import Dataset @@ -69,9 +70,13 @@ def __init__( for path in sorted((self._root / "meta" / "episodes").glob("chunk-*/file-*.parquet")) for row in pq.read_table(path).to_pylist() } + tasks_df = pd.read_parquet(self._root / "meta" / "tasks.parquet") + # LeRobot v2.x stores task text in a "task" column; v3.0 stores it as the + # (unnamed) DataFrame index and keeps only "task_index" as a column. + task_texts = tasks_df["task"] if "task" in tasks_df.columns else tasks_df.index self._tasks = { - int(row["task_index"]): str(row["task"]) - for row in pq.read_table(self._root / "meta" / "tasks.parquet").to_pylist() + int(task_index): str(task) + for task, task_index in zip(task_texts, tasks_df["task_index"]) } self._rows = sorted( ( diff --git a/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py index 5992ce67..25e167b0 100644 --- a/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py @@ -51,7 +51,7 @@ dtype=np.float32, ) -_NORMALIZER_PATH = Path(__file__).parent / "stats/bridge_orig_lerobot_stats.json" +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/bridge_orig_lerobot_stats.json" class BridgeOrigLeRobotDataset(ActionBaseDataset): diff --git a/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py index 3bd18598..f3faa91f 100644 --- a/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py @@ -52,7 +52,7 @@ dtype=np.float32, ) -_NORMALIZER_PATH = Path(__file__).parent / "stats/droid_lerobot_stats.json" +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/droid_lerobot_stats.json" class DROIDLeRobotDataset(ActionBaseDataset): diff --git a/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py new file mode 100644 index 00000000..10628db1 --- /dev/null +++ b/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Fractal LeRobot dataset — Google Robot RT-1. + +Robot: google_robot +87,212 episodes, 3,786,400 frames, 599 tasks, fps=3 +state: [x, y, z, rx, ry, rz, rw, gripper] (8D, quaternion) +action: [x, y, z, roll, pitch, yaw, gripper] (7D, delta) +video: observation.images.image (256×320) +""" + +from __future__ import annotations + +import random +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch +from lerobot.datasets.video_utils import decode_video_frames + +from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.vfm.action.pose_utils import ( + build_abs_pose_from_components, + pose_abs_to_rel, +) + +PoseConvention = Literal["backward_framewise"] +Viewpoint = Literal["ego_view"] + +_IMAGE_FEATURE = "observation.images.image" +_STATE_FEATURE = "observation.state" +_ACTION_FEATURE = "action" + +# These episodes contain base motion, which breaks the fixed-base Google Robot +# action assumption used by training and the viewer. +_SKIPPED_EPISODE_IDS: frozenset[int] = frozenset({29, 189, 382}) + +# Google Robot raw EE frame has x/y axes rotated ~90° around z compared to +# OpenCV convention. Rz(-90°) as a right-multiply corrects this: +# new_x = -old_y (rightward), new_y = old_x (downward), z unchanged (approach). +_GOOGLE_ROBOT_TO_OPENCV = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]], dtype=np.float32) + +# --------------------------------------------------------------------------- +# TCP → flange (gripper body) offset +# --------------------------------------------------------------------------- +# The fractal dataset records EE poses at ``link_gripper_tcp`` — a calibrated +# tool-center-point 164 mm past the gripper body (``link_gripper``), roughly +# at the fingertip. Re-referencing to ``link_gripper`` removes the +# calibration-dependent tilt and places the frame at the last actuated link. +# +# T = oMf[link_gripper_tcp]⁻¹ · oMf[link_gripper], computed via pinocchio FK +# at the neutral config from the SimplerEnv URDF. +# fmt: off +_TCP_TO_FLANGE = np.array([ + [+0.9999897671, -0.0008686425, +0.0044397163, -0.0050618476], + [+0.0008745501, +0.9999987346, -0.0013288658, -0.0016717725], + [-0.0044385564, +0.0013327349, +0.9999892615, -0.1635144743], + [+0.0000000000, +0.0000000000, +0.0000000000, +1.0000000000], +], dtype=np.float32) +# fmt: on + +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/fractal_lerobot_stats.json" + + +class FractalLeRobotDataset(ActionBaseDataset): + """Fractal (Google RT-1) dataset with 10D cartesian actions: + + [pos_delta(3), rot6d_delta(6), gripper(1)] + + Expects a LeRobot v2 dataset with: + * ``observation.images.image``: ego-view RGB video (256×320). + * ``observation.state``: 8D EE pose ``[x, y, z, rx, ry, rz, rw, gripper]`` + in TCP frame with quaternion (x, y, z, w) order. + * ``action``: 7D delta ``[x, y, z, roll, pitch, yaw, gripper]``; only the + gripper column (index 6) is used — SE(3) actions are derived from state. + + Episodes in ``_SKIPPED_EPISODE_IDS`` (base-motion outliers) are dropped. + """ + + def __init__( + self, + root: str, + fps: float = 3.0, + chunk_length: int = 16, + mode: str = "joint", + pose_convention: PoseConvention = "backward_framewise", + tolerance_s: float = 1e-4, + viewpoint: Viewpoint = "ego_view", + action_normalization: str | None = None, + sample_stride: int = 1, + ) -> None: + if viewpoint != "ego_view": + raise NotImplementedError("FractalLeRobotDataset only supports ego_view.") + super().__init__( + root=root, + domain_name="fractal", + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention=pose_convention, + tolerance_s=tolerance_s, + viewpoint=viewpoint, + action_normalization=action_normalization, + sample_stride=sample_stride, + ) + + # Drop rows belonging to known-bad episodes (base motion outliers). + n_before = len(self._rows) + self._rows = [row for row in self._rows if int(row["episode_index"]) not in _SKIPPED_EPISODE_IDS] + n_dropped = n_before - len(self._rows) + if n_dropped: + import logging + logging.getLogger(__name__).info( + "FractalLeRobotDataset: dropped %d / %d rows from episodes %s", + n_dropped, + n_before, + sorted(_SKIPPED_EPISODE_IDS), + ) + + @property + def action_dim(self) -> int: + """Action dimensionality: position(3) + 6D rotation(6) + gripper(1) = 10.""" + return 10 + + def _action_spec(self) -> ActionSpec: + return build_action_spec(Pos(dim=3), Rot("rot6d"), Gripper()) + + @classmethod + def _stats_path(cls) -> Path: + return _NORMALIZER_PATH + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + idx = int(idx) + row_idx = idx * self._sample_stride + observation_rows = self._rows[row_idx : row_idx + self._chunk_length + 1] + action_rows = observation_rows[: self._chunk_length] + + episode = self._episodes[int(observation_rows[0]["episode_index"])] + task = self._tasks[int(observation_rows[0]["task_index"])] + ai_caption = random.choice([part.strip() for part in task.split(" | ") if part.strip()] or [task]) + + video = self._load_video(episode, observation_rows) + raw_action, initial_pose = self._build_raw_action(observation_rows, action_rows) + + return self._build_result( + mode=mode, + video=video, + action=raw_action, + ai_caption=ai_caption, + initial_pose=initial_pose, + ) + + def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + timestamps = [float(row["timestamp"]) for row in observation_rows] + return decode_video_frames( + self._video_path(episode, _IMAGE_FEATURE), + [float(episode.get(f"videos/{_IMAGE_FEATURE}/from_timestamp", 0.0)) + ts for ts in timestamps], + self._tolerance_s, + ) + + def _build_raw_action( + self, + observation_rows: list[dict[str, Any]], + action_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor]: + # State layout: [x, y, z, rx, ry, rz, rw, gripper] (T+1 frames) + # Quaternion order: (rx, ry, rz, rw) matches scipy's (x, y, z, w). + state = np.asarray([row[_STATE_FEATURE] for row in observation_rows], dtype=np.float32) # [T+1, 8] + poses_abs = build_abs_pose_from_components(state[:, 0:3], state[:, 3:7], "quat_xyzw") # [T+1, 4, 4] + + # 1. TCP → flange: shift from link_gripper_tcp to link_gripper + poses_abs = poses_abs @ _TCP_TO_FLANGE + # 2. Kinematics → OpenCV convention (rotation only) + poses_abs[:, :3, :3] = poses_abs[:, :3, :3] @ _GOOGLE_ROBOT_TO_OPENCV + + initial_pose = torch.from_numpy(poses_abs[0].copy()).float() + poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention=self._pose_convention) + + gripper = np.asarray( + [row[_ACTION_FEATURE][6] for row in action_rows], dtype=np.float32 + ).reshape(-1, 1) # [T, 1] + + action = np.concatenate([poses_rel[-self._chunk_length :], gripper[-self._chunk_length :]], axis=-1) + return torch.from_numpy(action).float(), initial_pose diff --git a/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py b/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py index 136cd6c0..53b0a2c6 100644 --- a/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -"""RoboMIND Franka LeRobot dataset.""" +"""RoboMIND Franka LeRobot dataset — single-arm and dual-arm variants.""" from __future__ import annotations @@ -14,6 +14,7 @@ import torch.nn.functional as F from lerobot.datasets.video_utils import decode_video_frames +from cosmos_framework.data.vfm.action.action_normalization import load_action_stats from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset from cosmos_framework.data.vfm.action.pose_utils import ( @@ -22,10 +23,16 @@ ) PoseConvention = Literal["backward_framewise"] -Viewpoint = Literal["concat_view"] +Viewpoint = Literal["concat_view", "third_person_view"] -_IMAGE_FEATURES = { - "front": "observation.images.camera_front", +# Dual-arm uses camera_front as the top view; single-arm uses camera_top. +_IMAGE_FEATURES_DUAL = { + "top": "observation.images.camera_front", + "left": "observation.images.camera_left", + "right": "observation.images.camera_right", +} +_IMAGE_FEATURES_SINGLE = { + "top": "observation.images.camera_top", "left": "observation.images.camera_left", "right": "observation.images.camera_right", } @@ -39,7 +46,8 @@ dtype=np.float32, ) -_NORMALIZER_PATH = Path(__file__).parent / "stats/robomind_franka_stats.json" +_NORMALIZER_PATH_DUAL = Path(__file__).parent.parent / "normalizer_stats/robomind_franka_dual_stats.json" +_NORMALIZER_PATH_SINGLE = Path(__file__).parent.parent / "normalizer_stats/robomind_franka_stats.json" def _dual_arm_action_spec(): @@ -54,15 +62,15 @@ def _dual_arm_action_spec(): class RoboMINDFrankaDataset(ActionBaseDataset): - """RoboMIND Franka dual-arm dataset with 20D cartesian actions:: - - [left_pos_delta(3), left_rot6d_delta(6), left_gripper(1), - right_pos_delta(3), right_rot6d_delta(6), right_gripper(1)] + """RoboMIND Franka dataset — single-arm (10D) or dual-arm (20D) cartesian actions. - Single-arm shards, split/filter logic, image augmentation, fast - initialization, and alternate viewpoints are omitted. + Single-arm (``robomind-franka``): ``[pos_delta(3), rot6d_delta(6), gripper(1)]`` + Dual-arm (``robomind-franka-dual``): + ``[left_pos(3), left_rot6d(6), left_gripper(1), + right_pos(3), right_rot6d(6), right_gripper(1)]`` """ + _SUPPORTED_EMBODIMENTS = ("robomind-franka", "robomind-franka-dual") def __init__( self, @@ -77,11 +85,17 @@ def __init__( action_normalization: str | None = "quantile", sample_stride: int = 1, ) -> None: - if embodiment_type != "robomind-franka-dual": - raise NotImplementedError("This minimal RoboMIND dataset only supports robomind-franka-dual.") - if viewpoint != "concat_view": - raise NotImplementedError("This minimal RoboMIND dataset only supports concat_view.") + if embodiment_type not in self._SUPPORTED_EMBODIMENTS: + raise ValueError( + f"RoboMINDFrankaDataset only supports {self._SUPPORTED_EMBODIMENTS}, " + f"got embodiment_type={embodiment_type!r}." + ) + if viewpoint not in ("concat_view", "third_person_view"): + raise NotImplementedError(f"RoboMINDFrankaDataset does not support viewpoint={viewpoint!r}.") self._embodiment_type = embodiment_type + self._image_features = ( + _IMAGE_FEATURES_SINGLE if embodiment_type == "robomind-franka" else _IMAGE_FEATURES_DUAL + ) super().__init__( root=root, domain_name=embodiment_type, @@ -97,41 +111,82 @@ def __init__( @property def action_dim(self) -> int: - return 20 + return 10 if self._embodiment_type == "robomind-franka" else 20 def _action_spec(self) -> ActionSpec: + if self._embodiment_type == "robomind-franka": + return build_action_spec(Pos(), Rot("rot6d"), Gripper()) return _dual_arm_action_spec() @classmethod def _stats_path(cls) -> Path: - return _NORMALIZER_PATH + return _NORMALIZER_PATH_DUAL + + def _load_norm_stats(self) -> dict[str, torch.Tensor]: + if self._norm_stats is None: + path = _NORMALIZER_PATH_SINGLE if self._embodiment_type == "robomind-franka" else _NORMALIZER_PATH_DUAL + self._norm_stats = { + key: torch.from_numpy(value).float() + for key, value in load_action_stats(str(path)).items() + } + return self._norm_stats def __getitem__(self, idx: int) -> dict[str, Any]: mode = self._choose_mode() idx = int(idx) - first_row = self._rows[idx] - episode = self._episodes[int(first_row["episode_index"])] - row_idx = idx * self._sample_stride observation_rows = self._rows[row_idx : row_idx + self._chunk_length + 1] action_rows = observation_rows[: self._chunk_length] - video = self._load_concat_video(episode, observation_rows) - raw_action, initial_pose_left, initial_pose_right = self._build_raw_action(observation_rows, action_rows) + episode = self._episodes[int(observation_rows[0]["episode_index"])] task = self._tasks[int(observation_rows[0]["task_index"])] ai_caption = random.choice([part.strip() for part in task.split(" | ") if part.strip()] or [task]) + if self._viewpoint == "concat_view": + video = self._load_concat_video(episode, observation_rows) + if self._embodiment_type == "robomind-franka": + view_desc = ( + "The top row shows a third-person perspective looking towards the single-arm Franka robot from above. " + "The bottom-left view looks at the scene from the left side, and the bottom-right view looks at the scene from the right side." + ) + else: + view_desc = ( + "The top row shows a third-person perspective looking towards the dual-arm Franka robot from the front. " + "The bottom-left view looks at the scene from the left side, and the bottom-right view looks at the scene from the right side." + ) + else: + video = self._load_single_video(episode, observation_rows) + view_desc = None + + if self._embodiment_type == "robomind-franka": + raw_action, initial_pose = self._build_raw_action_single(observation_rows, action_rows) + extras: dict[str, Any] = {"initial_pose": initial_pose} + else: + raw_action, initial_pose_left, initial_pose_right = self._build_raw_action_dual(observation_rows, action_rows) + extras = {"initial_pose": initial_pose_left, "initial_pose_right": initial_pose_right} + + if view_desc is not None: + extras["additional_view_description"] = view_desc + return self._build_result( mode=mode, video=video, action=raw_action, ai_caption=ai_caption, - initial_pose=initial_pose_left, - initial_pose_right=initial_pose_right, - additional_view_description=( - "The top row shows a third-person perspective looking towards the dual-arm Franka robot from the front. " - "The bottom-left view looks at the scene from the left side, and the bottom-right view looks at the scene from the right side." - ), + **extras, + ) + + def _load_single_video( + self, + episode: dict[str, Any], + observation_rows: list[dict[str, Any]], + ) -> torch.Tensor: + video_key = self._image_features["top"] + timestamps = [float(row["timestamp"]) for row in observation_rows] + return decode_video_frames( + self._video_path(episode, video_key), + [float(episode.get(f"videos/{video_key}/from_timestamp", 0.0)) + ts for ts in timestamps], + self._tolerance_s, ) def _load_concat_video( @@ -146,18 +201,18 @@ def _load_concat_video( [float(episode.get(f"videos/{video_key}/from_timestamp", 0.0)) + ts for ts in timestamps], self._tolerance_s, ) - for name, video_key in _IMAGE_FEATURES.items() + for name, video_key in self._image_features.items() } - front = frames_by_view["front"] + top = frames_by_view["top"] left = frames_by_view["left"] right = frames_by_view["right"] - _, _, h_front, w_front = front.shape - half_h, half_w = h_front // 2, w_front // 2 + _, _, h_top, w_top = top.shape + half_h, half_w = h_top // 2, w_top // 2 left = F.interpolate(left, size=(half_h, half_w), mode="bilinear", align_corners=False) right = F.interpolate(right, size=(half_h, half_w), mode="bilinear", align_corners=False) bottom = torch.cat([left, right], dim=-1) - return torch.cat([front, bottom], dim=-2) + return torch.cat([top, bottom], dim=-2) def _build_relative_poses( self, @@ -170,7 +225,21 @@ def _build_relative_poses( poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention=self._pose_convention) return poses_rel, initial_pose - def _build_raw_action( + def _build_raw_action_single( + self, + observation_rows: list[dict[str, Any]], + action_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor]: + state = np.asarray([row[_STATE_FEATURE] for row in observation_rows], dtype=np.float32) + gripper = np.asarray([row[_ACTION_FEATURE] for row in action_rows], dtype=np.float32) + poses_rel, initial_pose = self._build_relative_poses(state[:, 0:3], state[:, 3:6]) + action = np.concatenate( + [poses_rel[-self._chunk_length :], 1.0 - gripper[-self._chunk_length :, [7]]], + axis=-1, + ) # [T, 10] + return torch.from_numpy(action).float(), initial_pose + + def _build_raw_action_dual( self, observation_rows: list[dict[str, Any]], action_rows: list[dict[str, Any]], @@ -188,5 +257,5 @@ def _build_raw_action( 1.0 - gripper[-self._chunk_length :, [15]], ], axis=-1, - ) + ) # [T, 20] return torch.from_numpy(action).float(), initial_pose_left, initial_pose_right diff --git a/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py b/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py new file mode 100644 index 00000000..317525f7 --- /dev/null +++ b/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""RoboMIND UR LeRobot dataset for single-arm UR5e embodiment.""" + +from __future__ import annotations + +import random +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch +from lerobot.datasets.video_utils import decode_video_frames + +from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.vfm.action.pose_utils import pose_abs_to_rel + +PoseConvention = Literal["backward_framewise"] +Viewpoint = Literal["third_person_view"] + +_IMAGE_FEATURE = "observation.images.camera_top" +_ACTION_JOINT_FEATURE = "actions.joint_position" + +# UR EE frame → OpenCV convention rotation (3×3, post-multiplied). +# Identity: attachment_site (quat="-1 1 0 0" in ur5e_robotiq_2f85.xml) already +# satisfies OpenCV convention (z = approach). +_ROBOMIND_UR_TO_OPENCV: np.ndarray = np.eye(3, dtype=np.float32) + +_UR5E_ARM_JOINTS = 6 # shoulder_pan … wrist_3 +_UR5E_EE_SITE = "attachment_site" # flange site in ur5e_robotiq_2f85.xml + +_MJCF_PATH = Path(__file__).resolve().parent.parent / "robot_assets" / "ur5e_robotiq_2f85.xml" + +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/robomind_ur_stats.json" + + +class RoboMINDURDataset(ActionBaseDataset): + """RoboMIND UR dataset with 10D cartesian actions: + + [pos_delta(3), rot6d_delta(6), gripper(1)] + + Uses FK on ``actions.joint_position`` (6 arm joints → MuJoCo + ``attachment_site`` SE(3) pose). ``observation.states.end_effector`` is + NOT used — it is recorded incorrectly (constant) in ~89 % of UR episodes; + ``actions.joint_position`` is valid for 100 % of episodes. + + The sample also includes ``joint_configs``: absolute joint angles ``(T, 7)`` + from ``actions.joint_position[1:T+1]`` for FK-based robot mesh animation. + """ + + def __init__( + self, + root: str, + fps: float = 10.0, + chunk_length: int = 16, + mode: str = "joint", + pose_convention: PoseConvention = "backward_framewise", + tolerance_s: float = 1e-4, + viewpoint: Viewpoint = "third_person_view", + action_normalization: str | None = None, + sample_stride: int = 1, + ) -> None: + if viewpoint != "third_person_view": + raise NotImplementedError("RoboMINDURDataset only supports third_person_view.") + super().__init__( + root=root, + domain_name="robomind-ur", + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention=pose_convention, + tolerance_s=tolerance_s, + viewpoint=viewpoint, + action_normalization=action_normalization, + sample_stride=sample_stride, + ) + self._mj_model, self._mj_data, self._ee_site_id = self._init_mujoco() + + @staticmethod + def _init_mujoco(): + """Load UR5e+Robotiq MuJoCo model (kinematics-only) and locate the EE site. + + Strips all geoms and mesh/texture/material assets from the MJCF via + ``MjSpec`` before compile so the model loads without mesh files on disk. + FK only needs the kinematic tree, so ``mj_forward`` + site poses still + produce identical EE poses. + """ + import mujoco + + spec = mujoco.MjSpec.from_file(str(_MJCF_PATH)) + + def _strip_geoms(body): + for g in list(body.geoms): + spec.delete(g) + for child in body.bodies: + _strip_geoms(child) + + _strip_geoms(spec.worldbody) + for m in list(spec.meshes): + spec.delete(m) + for t in list(spec.textures): + spec.delete(t) + for mat in list(spec.materials): + spec.delete(mat) + + mj_model = spec.compile() + mj_data = mujoco.MjData(mj_model) + ee_site_id = mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_SITE, _UR5E_EE_SITE) + if ee_site_id < 0: + raise RuntimeError(f"EE site '{_UR5E_EE_SITE}' not found in {_MJCF_PATH}") + return mj_model, mj_data, ee_site_id + + def _fk_ee_poses(self, arm_q: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Run MuJoCo FK for T+1 arm configs → EE site positions and rotations. + + Args: + arm_q: ``(T+1, 6)`` arm joint angles in radians. + + Returns: + ``(positions (T+1, 3), rotations (T+1, 3, 3))`` in MuJoCo world frame. + """ + import mujoco + + T1 = len(arm_q) + positions = np.empty((T1, 3), dtype=np.float32) + rotations = np.empty((T1, 3, 3), dtype=np.float32) + for t in range(T1): + self._mj_data.qpos[:_UR5E_ARM_JOINTS] = arm_q[t] + mujoco.mj_forward(self._mj_model, self._mj_data) + positions[t] = self._mj_data.site_xpos[self._ee_site_id] + rotations[t] = self._mj_data.site_xmat[self._ee_site_id].reshape(3, 3) + return positions, rotations + + @property + def action_dim(self) -> int: + return 10 + + def _action_spec(self) -> ActionSpec: + return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + + @classmethod + def _stats_path(cls) -> Path: + return _NORMALIZER_PATH + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + idx = int(idx) + row_idx = idx * self._sample_stride + # T+1 rows: current frame + T future frames (needed for FK EE poses and joint_configs) + observation_rows = self._rows[row_idx : row_idx + self._chunk_length + 1] + + episode = self._episodes[int(observation_rows[0]["episode_index"])] + task = self._tasks[int(observation_rows[0]["task_index"])] + ai_caption = random.choice([part.strip() for part in task.split(" | ") if part.strip()] or [task]) + + video = self._load_video(episode, observation_rows) + action, initial_pose, joint_configs = self._build_raw_action(observation_rows) + + return self._build_result( + mode=mode, + video=video, + action=action, + ai_caption=ai_caption, + initial_pose=initial_pose, + joint_configs=joint_configs, + ) + + def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + timestamps = [float(row["timestamp"]) for row in observation_rows] + return decode_video_frames( + self._video_path(episode, _IMAGE_FEATURE), + [float(episode.get(f"videos/{_IMAGE_FEATURE}/from_timestamp", 0.0)) + ts for ts in timestamps], + self._tolerance_s, + ) + + def _build_raw_action( + self, + observation_rows: list[dict[str, Any]], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # [T+1, 7]: 6 arm joints + 1 gripper command + q = np.asarray([row[_ACTION_JOINT_FEATURE] for row in observation_rows], dtype=np.float32) + T = len(q) - 1 + + # FK EE trajectory: T+1 absolute poses from arm joints via MuJoCo + fk_pos, fk_rot = self._fk_ee_poses(q[:, :_UR5E_ARM_JOINTS]) + poses_abs = np.tile(np.eye(4, dtype=np.float32), (T + 1, 1, 1)) + poses_abs[:, :3, 3] = fk_pos + poses_abs[:, :3, :3] = fk_rot @ _ROBOMIND_UR_TO_OPENCV + + initial_pose = torch.from_numpy(poses_abs[0].copy()).float() + poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention=self._pose_convention) + + # Raw UR gripper: 0=open, 1=closed → invert so action convention is 0=closed, 1=open. + # joint_configs keeps the raw value; FK mesh uses raw * 255 → Robotiq ctrl. + gripper = torch.from_numpy(1.0 - q[:T, 6:7]).float() + action = torch.cat([torch.from_numpy(poses_rel).float(), gripper], dim=-1) # [T, 10] + + # Mesh animation: frames 1..T of joint position (post-action states) + joint_configs = torch.from_numpy(q[1 : 1 + T].copy()) # [T, 7] + + return action, initial_pose, joint_configs diff --git a/cosmos_framework/data/vfm/action/datasets/stats/umi_lerobot_stats.json b/cosmos_framework/data/vfm/action/datasets/stats/umi_lerobot_stats.json deleted file mode 100644 index 44b9c8ce..00000000 --- a/cosmos_framework/data/vfm/action/datasets/stats/umi_lerobot_stats.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "q01": [-0.035246, -0.037122, -0.035762, 0.984050, -0.108706, -0.065188, -0.110908, 0.982889, -0.085106, 0.000000, -0.027468, -0.036971, -0.029396, 0.993522, -0.076207, -0.061227, -0.074231, 0.992173, -0.076929, 0.000000], - "q99": [ 0.038095, 0.033082, 0.032447, 1.000000, 0.110573, 0.068087, 0.108972, 1.000000, 0.089037, 0.096749, 0.033588, 0.032840, 0.027391, 1.000000, 0.073697, 0.057541, 0.075834, 1.000000, 0.086257, 0.085000] -} diff --git a/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py index cec95876..59395961 100644 --- a/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py @@ -36,7 +36,7 @@ FORWARD_EEF_IN_CAMERA_FRAME_XYZ_WXYZ: tuple[float, ...] = (0.0, 0.086, 0.056, 1.0, 0.0, 0.0, 0.0) """EEF offset for touch_in_the_wild / FastUMI rigs (camera mounted slightly forward).""" -_NORMALIZER_PATH = Path(__file__).parent / "stats/umi_lerobot_stats.json" +_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/umi_lerobot_stats.json" # Action layout: single-arm is the first 10D of the 20D bimanual stats file # (right_eef_poses(9) + right_eef_commands(1)). diff --git a/cosmos_framework/data/vfm/action/datasets/stats/agibotworld_beta_lerobot_stats.json b/cosmos_framework/data/vfm/action/normalizer_stats/agibotworld_beta_lerobot_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/datasets/stats/agibotworld_beta_lerobot_stats.json rename to cosmos_framework/data/vfm/action/normalizer_stats/agibotworld_beta_lerobot_stats.json diff --git a/cosmos_framework/data/vfm/action/datasets/stats/bridge_orig_lerobot_stats.json b/cosmos_framework/data/vfm/action/normalizer_stats/bridge_orig_lerobot_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/datasets/stats/bridge_orig_lerobot_stats.json rename to cosmos_framework/data/vfm/action/normalizer_stats/bridge_orig_lerobot_stats.json diff --git a/cosmos_framework/data/vfm/action/datasets/stats/droid_lerobot_stats.json b/cosmos_framework/data/vfm/action/normalizer_stats/droid_lerobot_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/datasets/stats/droid_lerobot_stats.json rename to cosmos_framework/data/vfm/action/normalizer_stats/droid_lerobot_stats.json diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/fractal_lerobot_stats.json b/cosmos_framework/data/vfm/action/normalizer_stats/fractal_lerobot_stats.json new file mode 100644 index 00000000..d54b1720 --- /dev/null +++ b/cosmos_framework/data/vfm/action/normalizer_stats/fractal_lerobot_stats.json @@ -0,0 +1,4 @@ +{ + "q01": [-0.039816, -0.049270, -0.056266, 0.983667, -0.134543, -0.107048, -0.126518, 0.977277, -0.091363, 0.000000], + "q99": [0.043860, 0.050352, 0.072505, 1.000000, 0.127404, 0.107273, 0.134140, 1.000000, 0.179731, 1.000000] +} diff --git a/cosmos_framework/data/vfm/action/datasets/stats/robomind_franka_stats.json b/cosmos_framework/data/vfm/action/normalizer_stats/robomind_franka_dual_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/datasets/stats/robomind_franka_stats.json rename to cosmos_framework/data/vfm/action/normalizer_stats/robomind_franka_dual_stats.json diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/robomind_franka_stats.json b/cosmos_framework/data/vfm/action/normalizer_stats/robomind_franka_stats.json new file mode 100644 index 00000000..c2d8432e --- /dev/null +++ b/cosmos_framework/data/vfm/action/normalizer_stats/robomind_franka_stats.json @@ -0,0 +1,4 @@ +{ + "q01": [-0.065029, -0.030683, -0.075321, 0.981664, -0.137429, -0.069593, -0.140220, 0.976885, -0.140399, 0.000000], + "q99": [0.068546, 0.036309, 0.051772, 1.000000, 0.140290, 0.079942, 0.137529, 1.000000, 0.113651, 1.000000] +} diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/robomind_ur_stats.json b/cosmos_framework/data/vfm/action/normalizer_stats/robomind_ur_stats.json new file mode 100644 index 00000000..60964383 --- /dev/null +++ b/cosmos_framework/data/vfm/action/normalizer_stats/robomind_ur_stats.json @@ -0,0 +1,4 @@ +{ + "q01": [-0.050087, -0.032055, -0.045222, 0.986554, -0.114343, -0.074432, -0.128455, 0.982298, -0.132509, 0.000000], + "q99": [0.040385, 0.036577, 0.032376, 1.000000, 0.129744, 0.095558, 0.115080, 1.000000, 0.121987, 1.000000] +} diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/umi_lerobot_stats.json b/cosmos_framework/data/vfm/action/normalizer_stats/umi_lerobot_stats.json new file mode 100644 index 00000000..5df9734a --- /dev/null +++ b/cosmos_framework/data/vfm/action/normalizer_stats/umi_lerobot_stats.json @@ -0,0 +1,4 @@ +{ + "q01": [-0.035246, -0.037122, -0.035762, 0.984050, -0.108706, -0.065188, -0.110908, 0.982889, -0.085106, 0.000000, -0.027468, -0.036971, -0.029396, 0.993522, -0.076207, -0.061227, -0.074231, 0.992173, -0.076929, 0.000000], + "q99": [ 0.038095, 0.033082, 0.032447, 1.000000, 0.110573, 0.068087, 0.108972, 1.000000, 0.089037, 0.096749, 0.033588, 0.032840, 0.027391, 1.000000, 0.073697, 0.057541, 0.075834, 1.000000, 0.086257, 0.085000] +} diff --git a/cosmos_framework/data/vfm/action/urdf_visualizer/G1_omnipicker_calibrated.urdf b/cosmos_framework/data/vfm/action/robot_assets/G1_omnipicker_calibrated.urdf similarity index 100% rename from cosmos_framework/data/vfm/action/urdf_visualizer/G1_omnipicker_calibrated.urdf rename to cosmos_framework/data/vfm/action/robot_assets/G1_omnipicker_calibrated.urdf diff --git a/cosmos_framework/data/vfm/action/robot_assets/ur5e_robotiq_2f85.xml b/cosmos_framework/data/vfm/action/robot_assets/ur5e_robotiq_2f85.xml new file mode 100644 index 00000000..75025458 --- /dev/null +++ b/cosmos_framework/data/vfm/action/robot_assets/ur5e_robotiq_2f85.xml @@ -0,0 +1,323 @@ + + + + + + + From 33ea6a7fdbc4bfa11864768a72741819bd1662f0 Mon Sep 17 00:00:00 2001 From: LiangHao Date: Thu, 2 Jul 2026 15:59:28 +0800 Subject: [PATCH 04/26] Add Cosmos3-Nano LIBERO-10 action-policy SFT recipe, config, eval harness, and doc (#61) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What Adds the **Cosmos3-Nano LIBERO-10 action-policy SFT** surface, mirroring the existing DROID counterpart (`action_policy_droid_nano` + toml + launcher + doc). ### Feature (net-new) - **Experiment configs** `action_policy_libero_nano` (libero_10-only) and `action_policy_libero_all_nano` (equal 4-suite mix) — gen + action heads from the public Cosmos3-Nano base. - **Dataset** `LIBEROLeRobotDataset` + `get_action_libero_sft_dataset` — frame_wise_relative rot6d, `quantile_rot`, concat_view (third-person + wrist), 20 fps. - `base_dataset` `tasks.parquet` fallback for community LIBERO layouts. - Resample-on-decode-failure guard so one undecodable packed-mp4 frame can't crash a multi-node run (matches i4 behavior). - **Closed-loop eval harness** with vectorized sim, batched `/predict_batch`, single-rank `no_dist` checkpoint load. - **Structured-prompt serving** in the policy server (`--format-prompt-as-json`), so eval matches the training prompt format; the recipe defaults to it. ### Recipe + doc — two presets (to match the Cosmos3 LIBERO-10 result) Both lr 5e-5, warmup 500, cycle 16000, global batch 2048 (HSDP 2x8): - **(A) libero_10-only** — `action_policy_libero_repro.toml` + `launch_sft_action_policy_libero.sh` (max_iter 2000). - **(B) libero-all (4-suite equal mix)** — `action_policy_libero_all_repro.toml` + `launch_sft_action_policy_libero_all.sh` (max_iter 5000; `LIBERO_ROOT` = LIBERO_LeRobot_v3 parent dir). - `docs/action_policy_libero_sft.md` documents both. ## Notes - Scoped to LIBERO only; broader action-dataloader/model changes are intentionally not included here. - Based on `main`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- cosmos_framework/configs/base/config.py | 5 +- .../action_policy_libero_all_nano.py | 233 +++ .../action_policy_libero_nano.py | 233 +++ .../data/vfm/action/datasets/__init__.py | 2 + .../vfm/action/datasets/action_sft_dataset.py | 74 +- .../action/datasets/libero_lerobot_dataset.py | 328 ++++ .../data/vfm/action/libero_pose_utils.py | 69 + ...bero_native_frame_wise_relative_rot6d.json | 37 + .../scripts/action_policy_server_libero.py | 272 +++- cosmos_framework/simulation/__init__.py | 2 + .../simulation/libero/__init__.py | 3 + .../simulation/libero/closed_loop_eval.py | 1343 +++++++++++++++++ cosmos_framework/utils/vfm/model_loader.py | 7 + docs/action_policy_libero_sft.md | 139 ++ examples/launch_sft_action_policy_libero.sh | 46 + .../launch_sft_action_policy_libero_all.sh | 49 + .../action_policy_libero_all_repro.toml | 49 + .../action_policy_libero_repro.toml | 46 + 18 files changed, 2865 insertions(+), 72 deletions(-) create mode 100644 cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py create mode 100644 cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py create mode 100644 cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py create mode 100644 cosmos_framework/data/vfm/action/libero_pose_utils.py create mode 100644 cosmos_framework/data/vfm/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json create mode 100644 cosmos_framework/simulation/__init__.py create mode 100644 cosmos_framework/simulation/libero/__init__.py create mode 100644 cosmos_framework/simulation/libero/closed_loop_eval.py create mode 100644 docs/action_policy_libero_sft.md create mode 100755 examples/launch_sft_action_policy_libero.sh create mode 100755 examples/launch_sft_action_policy_libero_all.sh create mode 100644 examples/toml/sft_config/action_policy_libero_all_repro.toml create mode 100644 examples/toml/sft_config/action_policy_libero_repro.toml diff --git a/cosmos_framework/configs/base/config.py b/cosmos_framework/configs/base/config.py index 99fbd808..ec65e32e 100644 --- a/cosmos_framework/configs/base/config.py +++ b/cosmos_framework/configs/base/config.py @@ -94,7 +94,10 @@ def make_config() -> Config: # Register shipped experiments explicitly. (vision_sft_nano also defines # vision_sft_nano_mapstyle_dataloader — the CosmosDataLoader variant — in the same module.) + import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_droid_nano # noqa: F401 + import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_libero_all_nano # noqa: F401 + import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_libero_nano # noqa: F401 import cosmos_framework.configs.base.experiment.sft.vision_sft_nano # noqa: F401 import cosmos_framework.configs.base.experiment.sft.vision_sft_super # noqa: F401 - import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_droid_nano # noqa: F401 + return c diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py new file mode 100644 index 00000000..5df504ce --- /dev/null +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""``action_policy_libero_all_nano`` — Cosmos3-Nano LIBERO-all (4-suite) action-policy SFT recipe. + +Feeds ``LIBEROLeRobotDataset`` (frame-wise-relative rot6d, ``quantile_rot``, +concat_view third-person + wrist) and trains the generation + action heads from +the public ``nvidia/Cosmos3-Nano`` base. Trains on all 4 LIBERO suites (equal mix); ``LIBERO_ROOT`` is the +LIBERO_LeRobot_v3 parent dir. See docs/action_policy_libero_sft.md. +""" + +import copy + +from hydra.core.config_store import ConfigStore + +from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG +from cosmos_framework.data.vfm.action.datasets.action_sft_dataset import get_action_libero_sft_dataset +from cosmos_framework.data.vfm.joint_dataloader import ( + PackingDataLoader, + RankPartitionedDataLoader, +) +from cosmos_framework.utils.lazy_config import LazyCall as L +from cosmos_framework.utils.lazy_config import LazyDict + +cs = ConfigStore.instance() + + +def _action_policy_libero_nano_model_config() -> dict: + """LIBERO model config: capped packed tokens, selective activation + checkpointing, fresh diffusion-expert init, 10x vision flow-matching loss. + Keep ``encode_exact_durations=[17, 61, 73]`` to match the Cosmos3-Nano base.""" + cfg = copy.deepcopy(NANO_MODEL_CONFIG) # action_gen=True, max_action_dim=64 + # Cap the packed sequence. Uncapped (-1) + a large max_samples_per_batch packs + # one very long sequence and OOMs even on H200; 74000 keeps the GA-validated bound. + cfg["max_num_tokens_after_packing"] = 74000 + cfg["activation_checkpointing"]["mode"] = "selective" + cfg["diffusion_expert_config"]["load_weights_from_pretrained"] = False + cfg["rectified_flow_training_config"]["loss_scale"] = 10.0 + cfg["rectified_flow_training_config"]["image_loss_scale"] = None + cfg["tokenizer"]["encode_exact_durations"] = [17, 61, 73] # match Cosmos3 base + reference SFT (do NOT reduce) + return cfg + + +action_policy_libero_all_nano = LazyDict( + dict( + defaults=[ + {"override /model": "mot_fsdp"}, + {"override /data_train": None}, + {"override /data_val": None}, + # FusedAdam with fp32 master_weights + eps 1e-8 (bf16 params + eps 1e-6 + # diverged on the action loss). + {"override /optimizer": "fusedadamw"}, + {"override /scheduler": "lambdalinear"}, # linear LR decay + {"override /checkpoint": "s3"}, + { + "override /callbacks": [ + "basic", + "optimization", + "job_monitor", + ] + }, + {"override /ema": "power"}, + {"override /tokenizer": "wan2pt2_tokenizer"}, + {"override /sound_tokenizer": None}, + {"override /vlm_config": None}, + {"override /ckpt_type": "dcp"}, + "_self_", + ], + job=dict( + project="cosmos3", + group="action_sft", + name="action_policy_libero_all_nano", + wandb_mode="disabled", + ), + model=dict( + config=_action_policy_libero_nano_model_config(), + ), + optimizer=dict( + betas=[0.9, 0.99], + eps=1.0e-08, + fused=True, # popped by build_optimizer for FusedAdam (fused by construction) + # Train the generation + action heads. + keys_to_select=[ + "moe_gen", + "time_embedder", + "vae2llm", + "llm2vae", + "action2llm", + "llm2action", + "action_modality_embed", + ], + lr=5.0e-05, + lr_multipliers={ + "action2llm": 5.0, + "llm2action": 5.0, + "action_modality_embed": 5.0, + }, + optimizer_type="FusedAdam", + weight_decay=0.05, + ), + scheduler=dict( + lr_scheduler_type="LambdaLinear", + cycle_lengths=[100], # smoke: 100 iters (real run sets via TOML, GA=10000) + f_max=[1.0], + f_min=[0.0], + f_start=[1.0e-06], + verbosity_interval=0, + warm_up_steps=[0], # smoke (real run sets via TOML, GA=2000) + ), + trainer=dict( + distributed_parallelism="fsdp", + grad_accum_iter=1, # real run sets via TOML (GA=2) + logging_iter=1, + max_iter=100, # smoke + max_val_iter=None, + run_validation=False, + run_validation_on_start=False, + save_zero_checkpoint=False, + seed=42, + timeout_period=999999999, + validation_iter=100, + compile_config=dict(recompile_limit=8, use_duck_shape=False), + cudnn=dict(benchmark=True, deterministic=False), + ddp=dict(broadcast_buffers=True, find_unused_parameters=False, static_graph=True), + grad_scaler_args=dict(enabled=False), + callbacks=dict( + dataloader_speed=dict(every_n=100, save_s3=False, step_size=1), + device_monitor=dict( + every_n=200, log_memory_detail=True, save_s3=False, step_size=1, upload_every_n_mul=5 + ), + grad_clip=dict(clip_norm=1.0, force_finite=True), + heart_beat=dict(every_n=200, save_s3=False, step_size=1, update_interval_in_minute=20), + iter_speed=dict(every_n=1, hit_thres=50, save_s3=False, save_s3_every_log_n=500), + low_precision=dict(update_iter=1), + manual_gc=dict(every_n=5, gc_level=1, warm_up=1), + param_count=dict(save_s3=False), + skip_nan_step=dict(max_consecutive_nan=100), + training_stats=dict(log_freq=100), + ), + ), + checkpoint=dict( + broadcast_via_filesystem=False, + dcp_async_mode_enabled=False, + enable_gcs_patch_in_boto3=True, + keys_not_to_resume=[], + # Skip net_ema (EMA warm-starts from net, see dcp.py) and the action + # heads, so they init fresh from the base (the public Cosmos3-Nano base + # has no LIBERO-trained action heads). + keys_to_skip_loading=[ + "net_ema.", + "action2llm", + "llm2action", + "action_modality_embed", + "action_pos_embed", + ], + load_ema_to_reg=False, + load_path="???", # Cosmos3-Nano DCP dir; supply via TOML/env + load_training_state=False, + only_load_scheduler_state=False, + save_iter=100, + strict_resume=False, # base init: tolerate key set differences + verbose=True, + hf_export=dict( + enabled=False, + export_every_n=1, + hf_repo_id=None, + upload_to_object_store=dict(bucket="", credentials="", enabled=False), + ), + jit=dict(device="cuda", dtype="bfloat16", enabled=False, input_shape=None, strict=True), + load_from_object_store=dict(bucket="", credentials="", enabled=False), + save_to_object_store=dict(bucket="", credentials="", enabled=False), + ), + dataloader_train=L(PackingDataLoader)( + audio_sample_rate=48000, + dataset_name="action_libero_all", + max_samples_per_batch=128, # peak-mem bound (256 OOMs on H200); global = 128 x DP8 x grad_accum2 = 2048 + max_sequence_length=None, # None disables token packing (TOML can't express null) + patch_spatial=2, + sound_latent_fps=0, + tokenizer_spatial_compression_factor=16, + tokenizer_temporal_compression_factor=4, + dataloader=L(RankPartitionedDataLoader)( + batch_size=1, + in_order=False, + num_workers=4, + persistent_workers=True, + pin_memory=True, + prefetch_factor=4, + sampler=None, + # Shuffling is handled by the dataset (iterable_shuffle=True below): + # ActionIterableShuffleDataset streams rank x worker-sharded, episode-order- + # shuffled, sequential-within-episode. + # libero-all: equal 1:1:1:1 mix over the 4 LIBERO suites. LIBERO_ROOT is the + # LIBERO_LeRobot_v3 PARENT dir; each suite reads ${LIBERO_ROOT}/. Use the + # 20 FPS nvidia/LIBERO_LeRobot_v3 (matches the bundled stats + 20 Hz eval). + datasets={ + _suite: dict( + ratio=1, + dataset=L(get_action_libero_sft_dataset)( + root="${oc.env:LIBERO_ROOT}/" + _suite, + fps=20, + chunk_length=16, + image_size=256, # concat_view -> 256x512 + mode="policy", + camera_mode="concat_view", + action_space="frame_wise_relative", + rotation_space="6d", + pose_coordinate_frame="native", + action_normalization="quantile_rot", + val_ratio=0.01, + iterable_shuffle=True, + episode_shuffle_seed=42, + resolution=None, + max_action_dim="${model.config.max_action_dim}", + cfg_dropout_rate=0.1, + format_prompt_as_json=True, + tokenizer_config="${model.config.vlm_config.tokenizer}", + ), + ) + for _suite in ("libero_spatial", "libero_object", "libero_goal", "libero_10") + }, + ), + ), + dataloader_val=None, + upload_reproducible_setup=False, + ), + flags={"allow_objects": True}, +) + + +for _item in [action_policy_libero_all_nano]: + _name = [k for k, v in globals().items() if v is _item][0] + cs.store(group="experiment", package="_global_", name=_name, node=_item) diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py new file mode 100644 index 00000000..2be73ab9 --- /dev/null +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""``action_policy_libero_nano`` — Cosmos3-Nano LIBERO-10 action-policy SFT recipe. + +Feeds ``LIBEROLeRobotDataset`` (frame-wise-relative rot6d, ``quantile_rot``, +concat_view third-person + wrist) and trains the generation + action heads from +the public ``nvidia/Cosmos3-Nano`` base. Train on ``libero_10`` alone +(``LIBERO_ROOT``). See docs/action_policy_libero_sft.md. +""" + +import copy + +from hydra.core.config_store import ConfigStore + +from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG +from cosmos_framework.data.vfm.action.datasets.action_sft_dataset import get_action_libero_sft_dataset +from cosmos_framework.data.vfm.joint_dataloader import ( + PackingDataLoader, + RankPartitionedDataLoader, +) +from cosmos_framework.utils.lazy_config import LazyCall as L +from cosmos_framework.utils.lazy_config import LazyDict + +cs = ConfigStore.instance() + + +def _action_policy_libero_nano_model_config() -> dict: + """LIBERO model config: capped packed tokens, selective activation + checkpointing, fresh diffusion-expert init, 10x vision flow-matching loss. + Keep ``encode_exact_durations=[17, 61, 73]`` to match the Cosmos3-Nano base.""" + cfg = copy.deepcopy(NANO_MODEL_CONFIG) # action_gen=True, max_action_dim=64 + # Cap the packed sequence. Uncapped (-1) + a large max_samples_per_batch packs + # one very long sequence and OOMs even on H200; 74000 keeps the GA-validated bound. + cfg["max_num_tokens_after_packing"] = 74000 + cfg["activation_checkpointing"]["mode"] = "selective" + cfg["diffusion_expert_config"]["load_weights_from_pretrained"] = False + cfg["rectified_flow_training_config"]["loss_scale"] = 10.0 + cfg["rectified_flow_training_config"]["image_loss_scale"] = None + cfg["tokenizer"]["encode_exact_durations"] = [17, 61, 73] # match Cosmos3 base + reference SFT (do NOT reduce) + return cfg + + +action_policy_libero_nano = LazyDict( + dict( + defaults=[ + {"override /model": "mot_fsdp"}, + {"override /data_train": None}, + {"override /data_val": None}, + # FusedAdam with fp32 master_weights + eps 1e-8 (bf16 params + eps 1e-6 + # diverged on the action loss). + {"override /optimizer": "fusedadamw"}, + {"override /scheduler": "lambdalinear"}, # linear LR decay + {"override /checkpoint": "s3"}, + { + "override /callbacks": [ + "basic", + "optimization", + "job_monitor", + ] + }, + {"override /ema": "power"}, + {"override /tokenizer": "wan2pt2_tokenizer"}, + {"override /sound_tokenizer": None}, + {"override /vlm_config": None}, + {"override /ckpt_type": "dcp"}, + "_self_", + ], + job=dict( + project="cosmos3", + group="action_sft", + name="action_policy_libero_nano", + wandb_mode="disabled", + ), + model=dict( + config=_action_policy_libero_nano_model_config(), + ), + optimizer=dict( + betas=[0.9, 0.99], + eps=1.0e-08, + fused=True, # popped by build_optimizer for FusedAdam (fused by construction) + # Train the generation + action heads. + keys_to_select=[ + "moe_gen", + "time_embedder", + "vae2llm", + "llm2vae", + "action2llm", + "llm2action", + "action_modality_embed", + ], + lr=5.0e-05, + lr_multipliers={ + "action2llm": 5.0, + "llm2action": 5.0, + "action_modality_embed": 5.0, + }, + optimizer_type="FusedAdam", + weight_decay=0.05, + ), + scheduler=dict( + lr_scheduler_type="LambdaLinear", + cycle_lengths=[100], # smoke: 100 iters (real run sets via TOML, GA=10000) + f_max=[1.0], + f_min=[0.0], + f_start=[1.0e-06], + verbosity_interval=0, + warm_up_steps=[0], # smoke (real run sets via TOML, GA=2000) + ), + trainer=dict( + distributed_parallelism="fsdp", + grad_accum_iter=1, # real run sets via TOML (GA=2) + logging_iter=1, + max_iter=100, # smoke + max_val_iter=None, + run_validation=False, + run_validation_on_start=False, + save_zero_checkpoint=False, + seed=42, + timeout_period=999999999, + validation_iter=100, + compile_config=dict(recompile_limit=8, use_duck_shape=False), + cudnn=dict(benchmark=True, deterministic=False), + ddp=dict(broadcast_buffers=True, find_unused_parameters=False, static_graph=True), + grad_scaler_args=dict(enabled=False), + callbacks=dict( + dataloader_speed=dict(every_n=100, save_s3=False, step_size=1), + device_monitor=dict( + every_n=200, log_memory_detail=True, save_s3=False, step_size=1, upload_every_n_mul=5 + ), + grad_clip=dict(clip_norm=1.0, force_finite=True), + heart_beat=dict(every_n=200, save_s3=False, step_size=1, update_interval_in_minute=20), + iter_speed=dict(every_n=1, hit_thres=50, save_s3=False, save_s3_every_log_n=500), + low_precision=dict(update_iter=1), + manual_gc=dict(every_n=5, gc_level=1, warm_up=1), + param_count=dict(save_s3=False), + skip_nan_step=dict(max_consecutive_nan=100), + training_stats=dict(log_freq=100), + ), + ), + checkpoint=dict( + broadcast_via_filesystem=False, + dcp_async_mode_enabled=False, + enable_gcs_patch_in_boto3=True, + keys_not_to_resume=[], + # Skip net_ema (EMA warm-starts from net, see dcp.py) and the action + # heads, so they init fresh from the base (the public Cosmos3-Nano base + # has no LIBERO-trained action heads). + keys_to_skip_loading=[ + "net_ema.", + "action2llm", + "llm2action", + "action_modality_embed", + "action_pos_embed", + ], + load_ema_to_reg=False, + load_path="???", # Cosmos3-Nano DCP dir; supply via TOML/env + load_training_state=False, + only_load_scheduler_state=False, + save_iter=100, + strict_resume=False, # base init: tolerate key set differences + verbose=True, + hf_export=dict( + enabled=False, + export_every_n=1, + hf_repo_id=None, + upload_to_object_store=dict(bucket="", credentials="", enabled=False), + ), + jit=dict(device="cuda", dtype="bfloat16", enabled=False, input_shape=None, strict=True), + load_from_object_store=dict(bucket="", credentials="", enabled=False), + save_to_object_store=dict(bucket="", credentials="", enabled=False), + ), + dataloader_train=L(PackingDataLoader)( + audio_sample_rate=48000, + dataset_name="action_libero", + max_samples_per_batch=128, # peak-mem bound (256 OOMs on H200); global = 128 x DP8 x grad_accum2 = 2048 + max_sequence_length=None, # None disables token packing (TOML can't express null) + patch_spatial=2, + sound_latent_fps=0, + tokenizer_spatial_compression_factor=16, + tokenizer_temporal_compression_factor=4, + dataloader=L(RankPartitionedDataLoader)( + batch_size=1, + in_order=False, + num_workers=4, + persistent_workers=True, + pin_memory=True, + prefetch_factor=4, + sampler=None, + # Shuffling is handled by the dataset (iterable_shuffle=True below): + # ActionIterableShuffleDataset streams rank x worker-sharded, episode-order- + # shuffled, sequential-within-episode. + datasets=dict( + libero=dict( + ratio=1, + dataset=L(get_action_libero_sft_dataset)( + # Local LeRobot dir for the libero_10 suite ONLY. Use the + # 20 FPS nvidia/LIBERO_LeRobot_v3 (matches the bundled stats + 20 Hz eval): + # hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset \ + # --include 'libero_10/**' --local-dir # LIBERO_ROOT=/libero_10 + root="${oc.env:LIBERO_ROOT}", + fps=20, # metadata only (FPS-agnostic loader reads native fps from info.json) + chunk_length=16, + image_size=256, # concat_view -> 256x512 + mode="policy", + camera_mode="concat_view", + action_space="frame_wise_relative", + rotation_space="6d", + pose_coordinate_frame="native", + action_normalization="quantile_rot", + val_ratio=0.01, + iterable_shuffle=True, + episode_shuffle_seed=42, + resolution=None, + max_action_dim="${model.config.max_action_dim}", + cfg_dropout_rate=0.1, + format_prompt_as_json=True, # structured JSON prompts (set False for plain-text) + tokenizer_config="${model.config.vlm_config.tokenizer}", + ), + ), + ), + ), + ), + dataloader_val=None, + upload_reproducible_setup=False, + ), + flags={"allow_objects": True}, +) + + +for _item in [action_policy_libero_nano]: + _name = [k for k, v in globals().items() if v is _item][0] + cs.store(group="experiment", package="_global_", name=_name, node=_item) diff --git a/cosmos_framework/data/vfm/action/datasets/__init__.py b/cosmos_framework/data/vfm/action/datasets/__init__.py index b94b1440..b1357d1a 100644 --- a/cosmos_framework/data/vfm/action/datasets/__init__.py +++ b/cosmos_framework/data/vfm/action/datasets/__init__.py @@ -13,6 +13,7 @@ from cosmos_framework.data.vfm.action.datasets.bridge_orig_lerobot_dataset import BridgeOrigLeRobotDataset from cosmos_framework.data.vfm.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset from cosmos_framework.data.vfm.action.datasets.fractal_lerobot_dataset import FractalLeRobotDataset +from cosmos_framework.data.vfm.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset from cosmos_framework.data.vfm.action.datasets.robomind_franka_dataset import RoboMINDFrankaDataset from cosmos_framework.data.vfm.action.datasets.robomind_ur_dataset import RoboMINDURDataset from cosmos_framework.data.vfm.action.datasets.umi_lerobot_dataset import UMILeRobotDataset @@ -23,6 +24,7 @@ "BridgeOrigLeRobotDataset", "DROIDLeRobotDataset", "FractalLeRobotDataset", + "LIBEROLeRobotDataset", "RoboMINDFrankaDataset", "RoboMINDURDataset", "UMILeRobotDataset", diff --git a/cosmos_framework/data/vfm/action/datasets/action_sft_dataset.py b/cosmos_framework/data/vfm/action/datasets/action_sft_dataset.py index 1790de55..a36e4357 100644 --- a/cosmos_framework/data/vfm/action/datasets/action_sft_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/action_sft_dataset.py @@ -12,6 +12,7 @@ to ``RankPartitionedDataLoader`` (mirroring how the vision recipe uses ``get_sft_dataset``). """ + from __future__ import annotations from typing import Any @@ -19,6 +20,7 @@ from torch.utils.data import Dataset, IterableDataset, get_worker_info from cosmos_framework.data.vfm.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset +from cosmos_framework.data.vfm.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset from cosmos_framework.data.vfm.action.transforms import ActionTransformPipeline @@ -42,7 +44,6 @@ def get_shuffle_blocks(self): return self._dataset.get_shuffle_blocks() - class ActionIterableShuffleDataset(IterableDataset): """Streaming view of a map-style ``ActionSFTDataset``. @@ -139,3 +140,74 @@ def get_action_droid_sft_dataset( if iterable_shuffle: return ActionIterableShuffleDataset(sft, seed=episode_shuffle_seed) return sft + + +def get_action_libero_sft_dataset( + *, + root: str, + fps: float = 20.0, + chunk_length: int = 16, + image_size: int = 256, + mode: str = "policy", + camera_mode: str = "concat_view", + action_space: str = "frame_wise_relative", + rotation_space: str = "6d", + pose_coordinate_frame: str = "native", + action_normalization: str | None = "quantile_rot", + action_stats_path: str | None = None, + split: str = "train", + val_ratio: float = 0.01, + seed: int = 0, + resolution: str | int | None = None, + max_action_dim: int = 64, + tokenizer_config: dict | None = None, + cfg_dropout_rate: float = 0.1, + append_viewpoint_info: bool = True, + append_duration_fps_timestamps: bool = True, + append_resolution_info: bool = True, + append_idle_frames: bool = True, + format_prompt_as_json: bool = False, + iterable_shuffle: bool = False, + episode_shuffle_seed: int = 42, +) -> Dataset: + """Build the LIBERO action-policy SFT dataset (GA reproduction defaults). + + Feeds ``LIBEROLeRobotDataset`` (frame-wise-relative rot6d actions, + ``quantile_rot``-normalized, concat_view third-person + wrist at 256x256 each + → 256x512) through ``ActionTransformPipeline``. ``root`` is a LOCAL LeRobot dir + (read parquet + video directly); pre-sync the HF dataset once, e.g. + ``hf download lerobot/libero_10 --repo-type dataset --local-dir ``. Point + ``root`` at libero_10 alone. The + dataset is FPS-agnostic (decodes at real frame timestamps); ``fps`` is metadata + for ``conditioning_fps`` / prompt duration. + """ + dataset = LIBEROLeRobotDataset( + root=root, + image_size=image_size, + chunk_length=chunk_length, + fps=fps, + mode=mode, + split=split, + val_ratio=val_ratio, + seed=seed, + camera_mode=camera_mode, + action_space=action_space, + rotation_space=rotation_space, + pose_coordinate_frame=pose_coordinate_frame, + action_normalization=action_normalization, + action_stats_path=action_stats_path, + ) + transform = ActionTransformPipeline( + tokenizer_config=tokenizer_config, + cfg_dropout_rate=cfg_dropout_rate, + max_action_dim=max_action_dim, + append_viewpoint_info=append_viewpoint_info, + append_duration_fps_timestamps=append_duration_fps_timestamps, + append_resolution_info=append_resolution_info, + append_idle_frames=append_idle_frames, + format_prompt_as_json=format_prompt_as_json, + ) + sft = ActionSFTDataset(dataset, transform, resolution) + if iterable_shuffle: + return ActionIterableShuffleDataset(sft, seed=episode_shuffle_seed) + return sft diff --git a/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py new file mode 100644 index 00000000..fb8d5e8c --- /dev/null +++ b/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""LIBERO LeRobot dataset (frame-wise-relative action policy). + +Mirrors ``DROIDLeRobotDataset``: reads the LeRobot parquet directly, windows by +frame index, and decodes video at each frame's REAL timestamp. That makes it +FPS-agnostic — it works with the 10 FPS community ``lerobot/libero_*`` datasets +and a 20 FPS conversion alike, without LeRobot's ``delta_timestamps`` grid (which +rejects any window whose synthetic timestamps don't land on real frames). + +Action layout (``frame_wise_relative``): the stored 7D ``action`` is already a +per-frame delta ``[dpos(3), drot_axisangle(3), gripper(1)]``; only the rotation is +re-encoded to the requested ``rotation_space`` -> ``[dpos(3), rot6d(6), gripper(1)]`` +(10D for ``6d``). + +NOTE on FPS / stats fidelity: the bundled ``quantile_rot`` stats were computed on +a 20 FPS conversion. Per-frame deltas at 10 FPS span 2x the wall-clock motion, so +use a 20 FPS LIBERO dataset (or recompute stats for the dataset's FPS). +Loading/training is correct at any FPS regardless. +""" + +from __future__ import annotations + +import json +import random +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pyarrow.parquet as pq +import torch +import torch.nn.functional as F +from lerobot.datasets.video_utils import decode_video_frames + +from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.vfm.action.libero_pose_utils import libero_action_dim, libero_rotation_format +from cosmos_framework.data.vfm.action.pose_utils import convert_rotation +from cosmos_framework.utils import log + +CameraMode = Literal["image", "wrist_image", "concat_view"] +RotationSpace = Literal["3d", "6d", "9d"] + +_ACTION_FEATURE = "action" +_IMAGE_FEATURE = "observation.images.image" +_WRIST_FEATURE = "observation.images.wrist_image" +_STAT_KEYS = ("mean", "std", "min", "max", "q01", "q99") +_NORMALIZERS_DIR = Path(__file__).parent.parent / "normalizer_stats" + +_VIEWPOINT_BY_CAMERA = { + "image": "third_person_view", + "wrist_image": "wrist_view", + "concat_view": "concat_view", +} + + +class LIBEROLeRobotDataset(ActionBaseDataset): + """LIBERO action-policy dataset with frame-wise-relative rot6d actions. + + 10D ``[pos_delta(3), rot6d_delta(6), gripper(1)]`` (for ``rotation_space='6d'``), + ``concat_view`` third-person + wrist video, and ``quantile_rot`` normalization + against the bundled stats. Reads parquet + decodes video at real timestamps, + so the requested ``fps`` is metadata only (it sets ``conditioning_fps`` and the + prompt duration); frame windows always use the data's actual frames. + """ + + def __init__( + self, + root: str, + fps: float = 20.0, + chunk_length: int = 16, + mode: str = "policy", + tolerance_s: float = 1e-4, + camera_mode: CameraMode = "concat_view", + image_size: int = 256, + action_space: str = "frame_wise_relative", + rotation_space: RotationSpace = "6d", + pose_coordinate_frame: str = "native", + embodiment_type: str = "libero", + action_normalization: str | None = "quantile_rot", + action_stats_path: str | None = None, + split: str = "train", + val_ratio: float = 0.01, + seed: int = 0, + sample_stride: int = 1, + ) -> None: + if action_space != "frame_wise_relative": + raise NotImplementedError( + f"This LIBERO dataset only supports action_space='frame_wise_relative', got {action_space!r}." + ) + if camera_mode not in _VIEWPOINT_BY_CAMERA: + raise ValueError(f"Unsupported camera_mode={camera_mode!r}. Use image/wrist_image/concat_view.") + split = split.lower().strip() + if split not in {"train", "val", "valid", "validation", "eval", "test", "full"}: + raise ValueError(f"Unsupported split={split!r}. Use train/val/full.") + if chunk_length % 4 != 0: + raise ValueError(f"chunk_length must be divisible by 4, got {chunk_length}.") + + super().__init__( + root=root, + domain_name=embodiment_type, + fps=fps, + chunk_length=chunk_length, + mode=mode, + pose_convention="backward_framewise", # unused for frame_wise deltas; satisfies the base assert + tolerance_s=tolerance_s, + viewpoint=_VIEWPOINT_BY_CAMERA[camera_mode], + # frame_wise_relative ⇔ backward_framewise idle semantics. quantile_rot is a + # LIBERO convention -> normalize with the "quantile" formula on raw-rotation + # stats (see _load_norm_stats); pass the method the base will call. + action_normalization=None if action_normalization is None else "quantile", + sample_stride=sample_stride, + ) + # FPS-agnostic loader: trust the dataset's NATIVE fps for conditioning_fps / + # prompt duration so the metadata is truthful (10 for the public + # lerobot/libero_*, 20 for a 20 FPS conversion). Frame sampling uses each + # frame's real timestamp regardless, so the requested ``fps`` is ignored here. + info_fps = self._info.get("fps") + if info_fps: + if int(info_fps) != int(fps): + log.info(f"Using dataset native fps={info_fps} for conditioning (requested {fps}).") + self._fps = float(info_fps) + self._dt = 1.0 / self._fps + self._camera_mode = camera_mode + self._image_size = int(image_size) + self._rotation_space = rotation_space.lower().strip() + self._pose_coordinate_frame = pose_coordinate_frame + self._embodiment_type = embodiment_type + self._requested_normalization = action_normalization + # quantile_rot normalizes against the raw (un-orthonormalized) rotation stats + # under "global_raw"; everything else uses "global". + self._stats_key = "global_raw" if action_normalization == "quantile_rot" else "global" + self._stats_file = self._resolve_stats_file(action_stats_path) + + if self._camera_mode == "image": + self._video_keys = [_IMAGE_FEATURE] + elif self._camera_mode == "wrist_image": + self._video_keys = [_WRIST_FEATURE] + else: + self._video_keys = [_IMAGE_FEATURE, _WRIST_FEATURE] + + # Compact, lazy frame index (mirrors DROIDLeRobotDataset): read only the + # columns the sample builder needs into contiguous arrays, ordered by global + # frame index, so DataLoader worker forks share them copy-on-write. + index_parts, episode_parts, task_parts, ts_parts, action_parts = [], [], [], [], [] + for path in sorted((self._root / "data").glob("chunk-*/file-*.parquet")): + table = pq.read_table(path, columns=["index", "episode_index", "task_index", "timestamp", _ACTION_FEATURE]) + index_parts.append(table["index"].to_numpy()) + episode_parts.append(table["episode_index"].to_numpy()) + task_parts.append(table["task_index"].to_numpy()) + ts_parts.append(table["timestamp"].to_numpy()) + action_parts.append(np.asarray(table[_ACTION_FEATURE].to_pylist(), dtype=np.float32)) + if not index_parts: + raise FileNotFoundError(f"No data parquet found under {self._root / 'data'}.") + order = np.argsort(np.concatenate(index_parts).astype(np.int64), kind="stable") + self._row_episode = np.concatenate(episode_parts).astype(np.int64)[order] + self._row_task = np.concatenate(task_parts).astype(np.int64)[order] + self._row_timestamp = np.concatenate(ts_parts).astype(np.float64)[order] + self._row_action = np.concatenate(action_parts, axis=0).astype(np.float32)[order] + + assert np.all(np.diff(self._row_episode) >= 0), "episode_index not contiguous after sorting by frame index" + ep_vals, ep_starts, ep_counts = np.unique(self._row_episode, return_index=True, return_counts=True) + + # Deterministic per-episode train/val split (seeded; same on every rank). + keep = self._split_episode_ids(ep_vals.tolist(), split, val_ratio, seed) + kept = np.array([int(v) in keep for v in ep_vals], dtype=bool) + self._ep_vals = ep_vals.astype(np.int64)[kept] + self._ep_starts = ep_starts.astype(np.int64)[kept] + kept_counts = ep_counts.astype(np.int64)[kept] + # Within-episode windows only: total - n_kept_episodes * chunk_length valid samples. + self._valid_cum = np.cumsum(np.maximum(0, kept_counts - self._chunk_length)).astype(np.int64) + + log.info( + f"Loaded LIBERO dataset root={self._root} split={split!r} camera_mode={camera_mode!r} " + f"fps={self._fps} kept_episodes={len(self._ep_vals)}/{len(ep_vals)} " + f"valid_indices={int(self._valid_cum[-1]) if self._valid_cum.size else 0}" + ) + + # ---- spec / dims ------------------------------------------------------- + + @property + def action_dim(self) -> int: + return libero_action_dim(self._rotation_space) + + def _action_spec(self) -> ActionSpec: + return build_action_spec(Pos(), Rot(libero_rotation_format(self._rotation_space)), Gripper()) + + @classmethod + def _stats_path(cls) -> Path: + # Base classmethod fallback; the instance uses self._stats_file (which also + # honors action_stats_path + the rotation/coordinate-frame-specific filename). + return _NORMALIZERS_DIR / "libero_native_frame_wise_relative_rot6d.json" + + # ---- normalization (nested global/global_raw + quantile_rot) ------------ + + def _bundled_stats_filename(self) -> str: + rotation_suffix = {"3d": "3d", "6d": "rot6d", "9d": "rot9d"}.get(self._rotation_space) + if rotation_suffix is None: + raise ValueError(f"Unsupported rotation_space={self._rotation_space!r}.") + action_space = "frame_wise_relative" + return f"{self._embodiment_type}_{self._pose_coordinate_frame}_{action_space}_{rotation_suffix}.json" + + def _resolve_stats_file(self, action_stats_path: str | None) -> Path: + if action_stats_path: + p = Path(action_stats_path) + if not p.is_absolute(): + p = _NORMALIZERS_DIR / p.name + if not p.exists(): + raise FileNotFoundError(f"action_stats_path not found: {action_stats_path!r}") + return p + p = _NORMALIZERS_DIR / self._bundled_stats_filename() + if not p.exists(): + raise FileNotFoundError( + f"Bundled LIBERO stats not found at {p}. Pass action_stats_path or recompute stats." + ) + return p + + def _load_norm_stats(self) -> dict[str, torch.Tensor]: + if self._norm_stats is None: + raw = json.loads(self._stats_file.read_text())[self._stats_key] + self._norm_stats = {k: torch.tensor(v, dtype=torch.float32) for k, v in raw.items() if k in _STAT_KEYS} + return self._norm_stats + + # ---- index helpers ----------------------------------------------------- + + @staticmethod + def _split_episode_ids(ep_ids: list[int], split: str, val_ratio: float, seed: int) -> set[int]: + if split == "full": + return set(int(v) for v in ep_ids) + if not (0.0 < val_ratio < 1.0): + raise ValueError(f"val_ratio must be in (0, 1), got {val_ratio}.") + n_val = max(1, int(round(len(ep_ids) * val_ratio))) + rng = random.Random(seed) # identical selection on every rank + val = set(int(v) for v in rng.sample(list(ep_ids), n_val)) + if split == "train": + return set(int(v) for v in ep_ids) - val + return val # val/valid/validation/eval/test + + def __len__(self) -> int: + return int(self._valid_cum[-1]) if self._valid_cum.size else 0 + + def get_shuffle_blocks(self) -> list[tuple[int, int]]: + """Per-episode ``(start, length)`` flat-index blocks for + ``ActionIterableShuffleDataset`` (shuffle block ORDER + shard across + ranks, sequential within a block).""" + blocks: list[tuple[int, int]] = [] + prev = 0 + for c in np.asarray(self._valid_cum).tolist(): + c = int(c) + if c > prev: + blocks.append((prev, c - prev)) + prev = c + return blocks + + # ---- sample build ------------------------------------------------------ + + def __getitem__(self, idx: int) -> dict[str, Any]: + # Resample a different valid window if a frame fails to decode (bounded retries). + n = len(self) + last_err: Exception | None = None + for _attempt in range(8): + try: + return self._build_item(idx) + except Exception as e: # noqa: BLE001 — skip past undecodable frames + last_err = e + log.warning(f"LIBERO: sample idx={idx} failed to load ({type(e).__name__}: {e}); resampling") + if n > 0: + idx = random.randint(0, n - 1) + raise RuntimeError(f"LIBERO: failed to load a sample after 8 resamples; last error: {last_err}") + + def _build_item(self, idx: int) -> dict[str, Any]: + mode = self._choose_mode() + idx = int(idx) + ep = int(np.searchsorted(self._valid_cum, idx, side="right")) + prev = int(self._valid_cum[ep - 1]) if ep > 0 else 0 + start = int(self._ep_starts[ep]) + (idx - prev) + episode_index = int(self._ep_vals[ep]) + episode = self._episodes[episode_index] + + stop = start + self._chunk_length + 1 + timestamps = [float(self._row_timestamp[j]) for j in range(start, stop)] + video = self._load_video(episode, timestamps) + + # frame_wise_relative: chunk per-frame deltas are the stored actions directly. + raw = self._row_action[start : start + self._chunk_length] # [chunk, 7] + action = self._build_frame_wise_action(raw) + + task = self._tasks[int(self._row_task[start])] + ai_caption = random.choice([p.strip() for p in task.split(" | ") if p.strip()] or [task]) + + extras: dict[str, Any] = {} + if self._camera_mode == "concat_view": + extras["additional_view_description"] = ( + "The left half shows the third-person view; the right half shows the wrist-mounted camera." + ) + return self._build_result(mode=mode, video=video, action=action, ai_caption=ai_caption, **extras) + + def _build_frame_wise_action(self, raw: np.ndarray) -> torch.Tensor: + raw_t = torch.from_numpy(np.ascontiguousarray(raw)).float() # [chunk, 7] + translation = raw_t[:, 0:3] + rotation_matrix = convert_rotation(raw_t[:, 3:6], input_format="axisangle", output_format="matrix") + rotation = convert_rotation( + rotation_matrix, input_format="matrix", output_format=libero_rotation_format(self._rotation_space) + ) + gripper = raw_t[:, 6:7] + return torch.cat([translation, rotation, gripper], dim=-1) # [chunk, action_dim] + + def _load_video(self, episode: dict[str, Any], timestamps: list[float]) -> torch.Tensor: + frames_by_view = {} + for key in self._video_keys: + from_ts = float(episode.get(f"videos/{key}/from_timestamp", 0.0)) + frames = decode_video_frames( + self._video_path(episode, key), + [from_ts + ts for ts in timestamps], + self._tolerance_s, + ) # [T, C, H, W] in [0, 1] + frames = self._resize(frames) + frames_by_view[key] = frames + if self._camera_mode == "concat_view": + # third-person (left) + wrist (right), horizontally concatenated -> [T, C, H, 2W] + return torch.cat([frames_by_view[_IMAGE_FEATURE], frames_by_view[_WRIST_FEATURE]], dim=-1) + return frames_by_view[self._video_keys[0]] + + def _resize(self, frames: torch.Tensor) -> torch.Tensor: + if frames.shape[-1] == self._image_size and frames.shape[-2] == self._image_size: + return frames + return F.interpolate(frames, size=(self._image_size, self._image_size), mode="bilinear", align_corners=False) diff --git a/cosmos_framework/data/vfm/action/libero_pose_utils.py b/cosmos_framework/data/vfm/action/libero_pose_utils.py new file mode 100644 index 00000000..3a4fd8e2 --- /dev/null +++ b/cosmos_framework/data/vfm/action/libero_pose_utils.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Small LIBERO pose helpers shared by training and closed-loop eval.""" + +from __future__ import annotations + +import numpy as np +import torch + +from cosmos_framework.data.vfm.action.pose_utils import ( + RotationConvention, + build_abs_pose_from_components, +) + +# Local-frame post-rotation pattern: +# R_opencv = R_native @ *_TO_OPENCV. +LIBERO_TO_OPENCV: np.ndarray = np.array( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + dtype=np.float32, +) + +LIBERO_ROTATION_FORMATS: dict[str, RotationConvention] = { + "3d": "axisangle", + "6d": "rot6d", + "9d": "rot9d", +} +LIBERO_ACTION_DIMS: dict[str, int] = {"3d": 7, "6d": 10, "9d": 13} + + +def libero_rotation_format(rotation_space: str) -> RotationConvention: + """Return the shared ``pose_utils`` rotation format for a LIBERO setting.""" + rotation_format = LIBERO_ROTATION_FORMATS.get(rotation_space) + if rotation_format is None: + raise ValueError(f"Unsupported rotation_space={rotation_space!r}. Use 3d/6d/9d.") + return rotation_format + + +def libero_action_dim(rotation_space: str) -> int: + """Return ``[xyz, rotation, gripper]`` action width for LIBERO.""" + action_dim = LIBERO_ACTION_DIMS.get(rotation_space) + if action_dim is None: + raise ValueError(f"Unsupported rotation_space={rotation_space!r}. Use 3d/6d/9d.") + return action_dim + + +def libero_rotation_space_from_action_dim(action_dim: int) -> str: + """Infer LIBERO rotation space from unpadded action width.""" + for rotation_space, dim in LIBERO_ACTION_DIMS.items(): + if dim == action_dim: + return rotation_space + raise ValueError(f"Unable to infer rotation_space from action_dim={action_dim}.") + + +def build_libero_abs_pose(state_raw: torch.Tensor | np.ndarray, *, to_opencv: bool) -> np.ndarray: + """Build absolute LIBERO EE poses from state rows. + + ``state_raw`` is ``[x,y,z,axisangle(3),gripper(2)]``. When requested, the + local EE frame is post-rotated into the shared OpenCV-style action frame. + """ + if isinstance(state_raw, torch.Tensor): + state_np = state_raw.detach().cpu().numpy().astype(np.float32, copy=False) + else: + state_np = np.asarray(state_raw, dtype=np.float32) + + poses_abs = build_abs_pose_from_components(state_np[:, :3], state_np[:, 3:6], "axisangle") + if to_opencv: + poses_abs[:, :3, :3] = poses_abs[:, :3, :3] @ LIBERO_TO_OPENCV + return poses_abs diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json b/cosmos_framework/data/vfm/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json new file mode 100644 index 00000000..a705e7cf --- /dev/null +++ b/cosmos_framework/data/vfm/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json @@ -0,0 +1,37 @@ +{ + "metadata": { + "embodiment_type": "libero", + "pose_convention": "frame_wise_relative", + "pose_coordinate_frame": "native", + "rotation_format": "6d", + "action_dim": 10, + "skip_rotation_dims": [3, 4, 5, 6, 7, 8], + "chunk_length": 16, + "sample_stride": null, + "dataset_name": "libero", + "dataset_class": "LIBEROLeRobotDataset", + "dataset_root": ["outputs/libero_datasets/libero_10", "outputs/libero_datasets/libero_object", "outputs/libero_datasets/libero_spatial", "outputs/libero_datasets/libero_goal"], + "_comment": "Dataset paths are placeholders; the statistics values are independent of local dataset location.", + "split": "train", + "num_samples_stats": 10000, + "reservoir_size": 50000, + "max_samples": 10000, + "sampling_seed": 42 + }, + "global": { + "mean": [ 0.050704, 0.097407, -0.094833, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.476725], + "std": [ 0.333621, 0.387175, 0.457140, 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, 0.499460], + "min": [-0.937500, -0.937500, -0.937500, -1.000000, -1.000000, -1.000000, -1.000000, -1.000000, -1.000000, 0.000000], + "max": [ 0.937500, 0.937500, 0.937500, 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, 1.000000], + "q01": [-0.723214, -0.808929, -0.937500, -1.000000, -1.000000, -1.000000, -1.000000, -1.000000, -1.000000, 0.000000], + "q99": [ 0.937500, 0.870536, 0.937500, 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, 1.000000, 1.000000] + }, + "global_raw": { + "mean": [ 0.050704, 0.097407, -0.094833, 0.994873, -0.004579, -0.004288, 0.004389, 0.996104, 0.001109, 0.476725], + "std": [ 0.333621, 0.387175, 0.457140, 0.010807, 0.077802, 0.063386, 0.078571, 0.009994, 0.038504, 0.499460], + "min": [-0.937500, -0.937500, -0.937500, 0.902028, -0.356085, -0.367416, -0.370434, 0.921907, -0.255000, 0.000000], + "max": [ 0.937500, 0.937500, 0.937500, 1.000000, 0.368853, 0.341214, 0.356395, 1.000000, 0.348251, 1.000000], + "q01": [-0.723214, -0.808929, -0.937500, 0.934955, -0.223431, -0.189878, -0.334735, 0.938516, -0.107736, 0.000000], + "q99": [ 0.937500, 0.870536, 0.937500, 1.000000, 0.331000, 0.163153, 0.226216, 1.000000, 0.127158, 1.000000] + } +} diff --git a/cosmos_framework/scripts/action_policy_server_libero.py b/cosmos_framework/scripts/action_policy_server_libero.py index 7382b978..a7be2348 100644 --- a/cosmos_framework/scripts/action_policy_server_libero.py +++ b/cosmos_framework/scripts/action_policy_server_libero.py @@ -63,7 +63,12 @@ # Action-specific helpers live in the in-tree project tree. Imports stay as # `projects.cosmos3.vfm.*` and are auto-rewritten to `cosmos3._src.vfm.*` by the # cosmos-framework release script. +from cosmos_framework.data.vfm.action.action_processing import ( + ActionProcessingRecord, + make_batched_action_processing_fields, +) from cosmos_framework.data.vfm.action.domain_utils import get_domain_id +from cosmos_framework.data.vfm.action.json_formatter import ActionPromptJsonFormatter from cosmos_framework.data.vfm.action.transforms import ( build_sequence_plan_from_mode, find_closest_target_size, @@ -92,6 +97,11 @@ _DURATION_FPS_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS." _RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution." +# Viewpoint tag for the concat_view (third-person + wrist) eval the LIBERO client runs; +# matches LIBEROLeRobotDataset's _VIEWPOINT_BY_CAMERA["concat_view"]. Used only when the +# experiment trains with JSON-structured prompts (format_prompt_as_json=True). +_LIBERO_JSON_VIEWPOINT = "concat_view" + # --------------------------------------------------------------------------- # Pre/post processing helpers (copied verbatim from the previous server, with @@ -431,7 +441,8 @@ class ActionServerArgs(pydantic.BaseModel): # ``OmniSetupOverrides`` programmatically in ``build_setup_overrides``. checkpoint: tyro.conf.OmitArgPrefixes[CheckpointOverrides] = CheckpointOverrides.model_construct() - """Checkpoint and config loading configuration.""" + """Checkpoint and config loading configuration. ``use_ema_weights`` lives here and + defaults True at inference (suppressed from CLI) -> evals load net_ema by default.""" output_dir: Path | None = None """Output directory for ``OmniInference`` (saved config.yaml, benchmarks). @@ -469,6 +480,12 @@ class ActionServerArgs(pydantic.BaseModel): """Action normalization to invert. ``auto`` reads ``action_normalization`` from the experiment config (default ``minmax`` if unspecified).""" + # ----- prompt format ------------------------------------------------------ + format_prompt_as_json: bool | None = None + """Serve prompts as structured JSON (matching training ``format_prompt_as_json``). + ``None`` reads the flag from the experiment config; set explicitly to override when + the eval experiment differs from the checkpoint's training prompt format.""" + # ----- debug dumps -------------------------------------------------------- dump_dir: Path | None = None """If set, dump observations, predicted actions, and rollout videos under @@ -642,9 +659,23 @@ def __init__(self, args: ActionServerArgs) -> None: self.append_resolution_info = _extract_bool_from_config( self.experiment_config, "append_resolution_info", default=True ) + # When the experiment trains with format_prompt_as_json=True, the caption is a + # structured JSON dict (ActionPromptJsonFormatter) and the legacy string appenders + # are skipped. Mirror that at serve time so the prompt format matches training. The + # CLI flag overrides the config when the eval experiment differs from the checkpoint. + if args.format_prompt_as_json is not None: + self.format_prompt_as_json = bool(args.format_prompt_as_json) + else: + self.format_prompt_as_json = _extract_bool_from_config( + self.experiment_config, "format_prompt_as_json", default=False + ) + self._prompt_json_formatter = ( + ActionPromptJsonFormatter(caption_key="ai_caption") if self.format_prompt_as_json else None + ) log.info( f"[action-server] prompt augmentation: " - f"append_duration_fps={self.append_duration_fps}, append_resolution_info={self.append_resolution_info}" + f"append_duration_fps={self.append_duration_fps}, append_resolution_info={self.append_resolution_info}, " + f"format_prompt_as_json={self.format_prompt_as_json}" ) # Action denormalization stats. @@ -804,6 +835,147 @@ def get_info(self) -> dict[str, Any]: # Predict # ------------------------------------------------------------------ + def _input_video_key(self) -> str: + input_video_key = getattr(self.model, "input_video_key", None) + if input_video_key is None: + input_video_key = getattr(self.model, "config", None).input_video_key # type: ignore[union-attr] + return input_video_key + + def _build_json_prompt(self, prompt: str, *, video: torch.Tensor, image_size: torch.Tensor) -> str: + """Reproduce the training-time JSON prompt for format_prompt_as_json=True runs. + + Runs the same ``ActionPromptJsonFormatter`` the training pipeline uses (after + spatial resize/pad), then ``json.dumps`` the dict exactly as + ``TextTokenizerTransform`` does before tokenization. ``idle_frames=0`` matches the + modal active-manipulation chunk (the policy should keep moving); ``viewpoint`` and + the zero ``action`` (total-frame count) mirror the LIBERO concat_view dataset.""" + data_dict: dict[str, Any] = { + "ai_caption": prompt, + "viewpoint": _LIBERO_JSON_VIEWPOINT, + "video": video, # post-pad [C,T,H,W]; formatter reads T for duration + "image_size": image_size, # post-pad [H,W]; formatter reads resolution + "conditioning_fps": torch.tensor(self.cfg.fps, dtype=torch.long), + "mode": "policy", + # Zero action chunk: only its frame count (chunk length) is read, for " out of ". + "action": torch.zeros((self.cfg.action_chunk_size, self.cfg.max_action_dim), dtype=torch.float32), + "idle_frames": torch.tensor(0, dtype=torch.long), + } + formatted = self._prompt_json_formatter(data_dict)["ai_caption"] + return json.dumps(formatted) if isinstance(formatted, dict) else str(formatted) + + def _prep_policy_item(self, req: dict[str, Any]) -> dict[str, Any]: + """Validate one request and build the per-sample model inputs (video pad, + prompt augmentation, sequence_plan). Shared by predict_policy (batch=1) and + predict_policy_batch (batch=N) so the two paths stay byte-identical per item.""" + image_b64 = req.get("image") + if not isinstance(image_b64, str): + raise ValueError("'image' must be a base64 string") + prompt = req.get("prompt") + if not isinstance(prompt, str): + raise ValueError("'prompt' must be a string") + domain_name = req.get("domain_name") + if not isinstance(domain_name, str): + raise ValueError("'domain_name' must be a string") + image_size = req.get("image_size") + if not isinstance(image_size, int) or image_size <= 0: + raise ValueError("'image_size' must be a positive integer") + + img_chw_uint8 = _decode_base64_png_to_rgb_uint8(image_b64) + img_h, img_w = img_chw_uint8.shape[-2:] + # Multi-view (non-square) images: scale proportionally, matching height to image_size. + if img_h != image_size: + scale = image_size / img_h + new_w = int(round(img_w * scale)) + hwc = img_chw_uint8.permute(1, 2, 0).cpu().numpy() + resized = Image.fromarray(hwc).resize((new_w, image_size), resample=Image.Resampling.BILINEAR) + arr = np.asarray(resized, dtype=np.uint8).copy() + img_chw_uint8 = torch.from_numpy(arr).permute(2, 0, 1).contiguous() + + t_frames = self.cfg.action_chunk_size + 1 + _, final_h, final_w = img_chw_uint8.shape + video_c_t_h_w_uint8 = img_chw_uint8.unsqueeze(1).repeat(1, t_frames, 1, 1) # [3,T,H,W] + resolution = get_vision_data_resolution((final_h, final_w)) + target_w, target_h = find_closest_target_size(final_h, final_w, resolution) + pad_dict: dict[str, Any] = {"video": video_c_t_h_w_uint8} + reflection_pad_to_target(pad_dict, ["video"], True, target_w, target_h) + sequence_plan = build_sequence_plan_from_mode( + mode="policy", + video_length=self.cfg.action_chunk_size + 1, + action_length=self.cfg.action_chunk_size, + has_text=True, + ) + if self._prompt_json_formatter is not None: + augmented_prompt = self._build_json_prompt( + prompt, video=pad_dict["video"], image_size=pad_dict["image_size"] + ) + else: + augmented_prompt = _augment_prompt_with_metadata( + prompt, + t_frames=t_frames, + fps=self.cfg.fps, + height=final_h, + width=final_w, + append_duration_fps=self.append_duration_fps, + append_resolution_info=self.append_resolution_info, + ) + return { + "img_chw_uint8": img_chw_uint8, + "video_padded": pad_dict["video"], + "padded_image_size": pad_dict["image_size"], + "augmented_prompt": augmented_prompt, + "sequence_plan": sequence_plan, + "domain_name": domain_name, + "image_size": image_size, + } + + def predict_policy_batch(self, reqs: list[dict[str, Any]]) -> dict[str, Any]: + """Batched policy inference: N requests -> ONE diffusion forward (batch_size=N) + -> N denormalized action chunks. Skips vision decode (the vectorized eval client + only needs actions), so it is ~N x faster than N serial /predict calls.""" + t0 = time.monotonic() + if not isinstance(reqs, list) or not reqs: + raise ValueError("'items' must be a non-empty list of policy requests") + preps = [self._prep_policy_item(r) for r in reqs] + n = len(preps) + action_t_d = torch.zeros((self.cfg.action_chunk_size, self.cfg.max_action_dim), dtype=torch.float32) + input_video_key = self._input_video_key() + batch: dict[str, Any] = { + input_video_key: [[p["video_padded"]] for p in preps], + **make_batched_action_processing_fields( + ActionProcessingRecord(raw_action_dim=self.raw_action_dim, action_normalizer=None), + batch_size=n, + ), + "action": [[action_t_d] for _ in preps], + "mode": ["policy"] * n, + "ai_caption": [p["augmented_prompt"] for p in preps], + "prompt": [p["augmented_prompt"] for p in preps], + "conditioning_fps": [torch.tensor(self.cfg.fps, dtype=torch.long) for _ in preps], + "image_size": torch.stack([p["padded_image_size"] for p in preps]).to(device="cuda"), + "domain_id": [torch.tensor(get_domain_id(p["domain_name"]), dtype=torch.long) for p in preps], + "sequence_plan": [p["sequence_plan"] for p in preps], + } + t_inf0 = time.monotonic() + with self._lock: + with torch.inference_mode(): + samples = self.model.generate_samples_from_batch( + batch, + guidance=self.cfg.guidance, + seed=[self.cfg.seed] * n, + num_steps=self.cfg.num_steps, + has_negative_prompt=False, + ) + t_inf1 = time.monotonic() + actions: list[list[list[float]]] = [] + for i in range(n): + pred = samples["action"][i].float().squeeze(0) # [T,D] + pred = self._denormalize_action(pred) + actions.append(pred.detach().cpu().numpy().tolist()) + log.info( + f"[action-server] predict_batch n={n} steps={self.cfg.num_steps} " + f"ms_total={(time.monotonic() - t0) * 1000.0:.1f} ms_infer={(t_inf1 - t_inf0) * 1000.0:.1f}" + ) + return {"actions": actions} + def predict_policy(self, req: dict[str, Any]) -> dict[str, Any]: """ Run policy inference: given an observation image and prompt, predict actions. @@ -835,52 +1007,17 @@ def predict_policy(self, req: dict[str, Any]) -> dict[str, Any]: self._req_id += 1 request_id = int(self._req_id) - # Validate request - image_b64 = req.get("image") - if not isinstance(image_b64, str): - raise ValueError("'image' must be a base64 string") - - prompt = req.get("prompt") - if not isinstance(prompt, str): - raise ValueError("'prompt' must be a string") - - domain_name = req.get("domain_name") - if not isinstance(domain_name, str): - raise ValueError("'domain_name' must be a string") - - image_size = req.get("image_size") - if not isinstance(image_size, int) or image_size <= 0: - raise ValueError("'image_size' must be a positive integer") - - # Decode image + # Per-item preprocessing (validation, decode/resize/pad, prompt, sequence_plan). t_decode0 = time.monotonic() - img_chw_uint8 = _decode_base64_png_to_rgb_uint8(image_b64) - img_h, img_w = img_chw_uint8.shape[-2:] - - # Handle resizing: for multi-view (non-square) images, scale proportionally - # to maintain aspect ratio while matching height to image_size - if img_h != image_size: - # Calculate new width to maintain aspect ratio - scale = image_size / img_h - new_w = int(round(img_w * scale)) - hwc = img_chw_uint8.permute(1, 2, 0).cpu().numpy() # [H,W,3] - resized = Image.fromarray(hwc).resize((new_w, image_size), resample=Image.Resampling.BILINEAR) - arr = np.asarray(resized, dtype=np.uint8).copy() - img_chw_uint8 = torch.from_numpy(arr).permute(2, 0, 1).contiguous() # [3,H,W] # [3,H,W] + prep = self._prep_policy_item(req) t_decode1 = time.monotonic() - - # Construct batch in IterativeJointDataLoader format (list-of-lists for multi-item keys) - t_frames = self.cfg.action_chunk_size + 1 - _, final_h, final_w = img_chw_uint8.shape - video_c_t_h_w_uint8 = img_chw_uint8.unsqueeze(1).repeat(1, t_frames, 1, 1) # [3,T,H,W] - - # Apply reflection padding to match closest predefined resolution - resolution = get_vision_data_resolution((final_h, final_w)) - target_w, target_h = find_closest_target_size(final_h, final_w, resolution) - pad_dict: dict[str, Any] = {"video": video_c_t_h_w_uint8} - reflection_pad_to_target(pad_dict, ["video"], True, target_w, target_h) - video_padded = pad_dict["video"] # (C, T, target_h, target_w) - padded_image_size = pad_dict["image_size"] # (4,) + img_chw_uint8 = prep["img_chw_uint8"] + video_padded = prep["video_padded"] + padded_image_size = prep["padded_image_size"] + augmented_prompt = prep["augmented_prompt"] + sequence_plan = prep["sequence_plan"] + domain_name = prep["domain_name"] + image_size = prep["image_size"] # Action: zeros tensor as noise starting point for policy mode action_t_d = torch.zeros( @@ -888,30 +1025,17 @@ def predict_policy(self, req: dict[str, Any]) -> dict[str, Any]: dtype=torch.float32, ) # [T,action_dim] - input_video_key = getattr(self.model, "input_video_key", None) - if input_video_key is None: - input_video_key = getattr(self.model, "config", None).input_video_key # type: ignore[union-attr] - - sequence_plan = build_sequence_plan_from_mode( - mode="policy", - video_length=self.cfg.action_chunk_size + 1, - action_length=self.cfg.action_chunk_size, - has_text=True, - ) - - augmented_prompt = _augment_prompt_with_metadata( - prompt, - t_frames=t_frames, - fps=self.cfg.fps, - height=final_h, - width=final_w, - append_duration_fps=self.append_duration_fps, - append_resolution_info=self.append_resolution_info, - ) + input_video_key = self._input_video_key() batch: dict[str, Any] = { input_video_key: [[video_padded]], - "raw_action_dim": [torch.tensor(self.raw_action_dim, dtype=torch.long)], + # Provide BOTH raw_action_dim and the action_processing_record the model + # needs to externalize (invert) the generated action; building the batch + # by hand previously omitted the record -> "cannot be externalized". + **make_batched_action_processing_fields( + ActionProcessingRecord(raw_action_dim=self.raw_action_dim, action_normalizer=None), + batch_size=1, + ), "action": [[action_t_d]], "mode": ["policy"], "ai_caption": [augmented_prompt], @@ -1103,7 +1227,7 @@ def do_GET(self) -> None: # noqa: N802 self._send_json(404, {"error": "Not found"}) def do_POST(self) -> None: # noqa: N802 - if self.path not in ("/", "/predict"): + if self.path not in ("/", "/predict", "/predict_batch"): self._send_json(404, {"error": "Not found"}) return @@ -1147,13 +1271,21 @@ def do_POST(self) -> None: # noqa: N802 f"path={self.path} bytes={length}" ) + is_batch = self.path == "/predict_batch" try: - out = service.predict_policy(req) + if is_batch: + out = service.predict_policy_batch(req.get("items", [])) + else: + out = service.predict_policy(req) except Exception as e: err = str(e) traceback.print_exc() - payload = {"action": [], "error": err, "request_id": req.get("request_id")} + payload = ( + {"actions": [], "error": err} + if is_batch + else {"action": [], "error": err, "request_id": req.get("request_id")} + ) log.error(f"[action-server] request_id={req.get('request_id')} ERROR: {err}") # Dump failed request for offline debugging if enabled. diff --git a/cosmos_framework/simulation/__init__.py b/cosmos_framework/simulation/__init__.py new file mode 100644 index 00000000..28a81be6 --- /dev/null +++ b/cosmos_framework/simulation/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 diff --git a/cosmos_framework/simulation/libero/__init__.py b/cosmos_framework/simulation/libero/__init__.py new file mode 100644 index 00000000..503ec1b1 --- /dev/null +++ b/cosmos_framework/simulation/libero/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + diff --git a/cosmos_framework/simulation/libero/closed_loop_eval.py b/cosmos_framework/simulation/libero/closed_loop_eval.py new file mode 100644 index 00000000..660be360 --- /dev/null +++ b/cosmos_framework/simulation/libero/closed_loop_eval.py @@ -0,0 +1,1343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +""" +Closed-loop evaluation for LIBERO using the Action HTTP inference server. + +# Single-view example (agentview camera): +PYTHONPATH=. python cosmos_framework/simulation/libero/closed_loop_eval.py \ + --server_url http://localhost:8000 \ + --task_suite libero_10 \ + --num_trials_per_task 10 \ + --action_horizon 16 \ + --camera agentview \ + --save_gifs --gif_fps 20 \ + --action_space frame_wise_relative \ + --rotation_space 6d \ + --action_dim 10 \ + --output_dir results/libero_closed_loop_10_single_view + +# Multi-view example (agentview + wrist cameras): +PYTHONPATH=. python cosmos_framework/simulation/libero/closed_loop_eval.py \ + --server_url http://localhost:8000 \ + --task_suite libero_goal \ + --num_trials_per_task 2 \ + --action_horizon 16 \ + --camera agentview,wrist \ + --save_gifs --gif_fps 20 \ + --action_space frame_wise_relative \ + --rotation_space 6d \ + --action_dim 10 \ + --output_dir results/libero_closed_loop_goal_multiview +""" + +from __future__ import annotations + +import argparse +import base64 +import io +import json +import os +import random +import sys +import time +from dataclasses import dataclass +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import numpy as np +import requests +from PIL import Image +from scipy.spatial.transform import Rotation as R + +from cosmos_framework.data.vfm.action.libero_pose_utils import ( + libero_rotation_format, + libero_rotation_space_from_action_dim, +) +from cosmos_framework.data.vfm.action.pose_utils import convert_rotation +from cosmos_framework.data.vfm.action.viewpoint_utils import DEFAULT_VIEWPOINT_TEMPLATES + +benchmark: Any +get_libero_path: Any +OffScreenRenderEnv: Any + + +TASK_MAX_STEPS: dict[str, int] = { + "libero_spatial": 220, + "libero_object": 280, + "libero_goal": 300, + "libero_10": 520, + "libero_90": 400, +} + + +_CAMERA_PROMPT_NAMES: dict[str, str] = { + "agentview": "third-person view", + "wrist": "wrist-mounted camera", +} + + +def _append_prompt_sentence(prompt: str, sentence: str) -> str: + """Append one metadata sentence using the same separator convention as training augmentors.""" + if sentence in prompt: + return prompt + prompt = prompt.rstrip() + if not prompt: + return sentence.rstrip() + separator = " " if prompt.rstrip().endswith(".") else ". " + return prompt + separator + sentence.rstrip() + + +def _concat_view_layout_description(cameras: list[str]) -> str: + """Describe the horizontal camera layout sent by ``ActionEnvironmentClient``.""" + camera_names = [_CAMERA_PROMPT_NAMES[camera] for camera in cameras] + if len(camera_names) == 2: + return f"The left half shows the {camera_names[0]}; the right half shows the {camera_names[1]}." + layout = ", ".join(camera_names) + return f"The views are concatenated horizontally from left to right as: {layout}." + + +def _augment_task_prompt_with_viewpoint(task_description: str, cameras: list[str]) -> str: + """Concat-view caption augmentation for closed-loop LIBERO eval.""" + if len(cameras) <= 1: + return task_description + prompt = _append_prompt_sentence(task_description, DEFAULT_VIEWPOINT_TEMPLATES["concat_view"]) + return _append_prompt_sentence(prompt, _concat_view_layout_description(cameras)) + + +def _rotation_repr_to_mat(rotation: np.ndarray, rotation_space: str) -> np.ndarray: + """Convert a single LIBERO rotation block to a 3x3 rotation matrix.""" + matrix = convert_rotation( + rotation, + libero_rotation_format(rotation_space), + "matrix", + normalize_matrix=rotation_space != "3d", + ) + if not isinstance(matrix, np.ndarray): + raise TypeError(f"Expected NumPy rotation matrix, got {type(matrix)!r}") + return matrix + + +@dataclass +class EpisodeResult: + success: bool + steps: int + error: str | None + actions: list[list[float]] + + +class ActionEnvironmentClient: + """Client for interacting with the Action model server.""" + + server_url: str + domain_name: str + prompt: str + image_size: int + timeout: float + + def __init__( + self, + server_url: str, + domain_name: str, + prompt: str, + image_size: int, + timeout: float, + ) -> None: + self.server_url = server_url.rstrip("/") + self.domain_name = domain_name + self.prompt = prompt + self.image_size = image_size + self.timeout = timeout + + def check_health(self) -> bool: + """Check if the model server is healthy.""" + try: + resp = requests.get(f"{self.server_url}/", timeout=5.0) + return resp.status_code == 200 + except requests.RequestException: + return False + + def get_info(self) -> dict[str, str]: + """Get model server info.""" + resp = requests.get(f"{self.server_url}/info", timeout=5.0) + resp.raise_for_status() + return resp.json() + + def notify_next_episode(self) -> None: + """Notify server to advance to next episode (used with dataset action server).""" + try: + requests.post( + f"{self.server_url}/next_episode", + json={"prompt": self.prompt}, + timeout=5.0, + ) + except requests.RequestException: + pass + + def encode_image(self, image: np.ndarray) -> str: + """Encode a numpy image (H, W, 3) uint8 to base64 PNG, resizing to image_size.""" + if image.dtype != np.uint8: + if image.max() <= 1.0: + image = (image * 255.0).round().astype(np.uint8) + else: + image = image.astype(np.uint8) + pil_img = Image.fromarray(image) + if pil_img.size != (self.image_size, self.image_size): + pil_img = pil_img.resize( + (self.image_size, self.image_size), + resample=Image.Resampling.BILINEAR, + ) + buf = io.BytesIO() + pil_img.save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode("ascii") + + def encode_image_raw(self, image: np.ndarray) -> str: + """Encode a numpy image (H, W, 3) uint8 to base64 PNG without resizing.""" + if image.dtype != np.uint8: + if image.max() <= 1.0: + image = (image * 255.0).round().astype(np.uint8) + else: + image = image.astype(np.uint8) + pil_img = Image.fromarray(image) + buf = io.BytesIO() + pil_img.save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode("ascii") + + def resize_image(self, image: np.ndarray) -> np.ndarray: + """Resize image to model input size.""" + if image.dtype != np.uint8: + if image.max() <= 1.0: + image = (image * 255.0).round().astype(np.uint8) + else: + image = image.astype(np.uint8) + pil_img = Image.fromarray(image) + if pil_img.size != (self.image_size, self.image_size): + pil_img = pil_img.resize( + (self.image_size, self.image_size), + resample=Image.Resampling.BILINEAR, + ) + return np.array(pil_img) + + def concatenate_images(self, images: list[np.ndarray]) -> np.ndarray: + """Resize each image and concatenate horizontally (side-by-side). + + Args: + images: List of images with shape (H, W, 3). + + Returns: + Concatenated image with shape (image_size, image_size*num_views, 3). + """ + resized = [self.resize_image(img) for img in images] + return np.concatenate(resized, axis=1) + + def predict(self, observation: np.ndarray | list[np.ndarray]) -> dict[str, Any]: + """Send observation(s) to model server and get predicted actions. + + Args: + observation: Single image as np.ndarray or list of images for multi-view. + For multi-view, images are resized and concatenated horizontally before sending. + """ + if isinstance(observation, list): + # Multi-view: resize each, concatenate horizontally, and send as single image + concatenated = self.concatenate_images(observation) + encoded = self.encode_image_raw(concatenated) + else: + # Single view: send single image + encoded = self.encode_image(observation) + + payload = { + "image": encoded, + "prompt": self.prompt, + "domain_name": self.domain_name, + "image_size": self.image_size, + } + + resp = requests.post( + f"{self.server_url}/predict", + json=payload, + headers={"Content-Type": "application/json"}, + timeout=self.timeout, + ) + resp.raise_for_status() + + result = resp.json() + if "error" in result and result["error"]: + raise RuntimeError(f"Model server error: {result['error']}") + return result + + def predict_batch(self, observations: list[list[np.ndarray]]) -> list[list[list[float]]]: + """Batched inference: a list of per-env multi-view observations -> ONE + POST /predict_batch -> a list of action chunks (one per env). Used by the + vectorized eval so N parallel envs share a single diffusion forward.""" + items = [] + for obs_imgs in observations: + concat = self.concatenate_images(obs_imgs) if len(obs_imgs) > 1 else self.resize_image(obs_imgs[0]) + items.append( + { + "image": self.encode_image_raw(concat), + "prompt": self.prompt, + "domain_name": self.domain_name, + "image_size": self.image_size, + } + ) + resp = requests.post( + f"{self.server_url}/predict_batch", + json={"items": items}, + headers={"Content-Type": "application/json"}, + timeout=max(self.timeout, 300.0), + ) + resp.raise_for_status() + result = resp.json() + if "error" in result and result["error"]: + raise RuntimeError(f"Model server error: {result['error']}") + return result["actions"] + + +def _find_accessible_dri_nodes() -> list[Path]: + dri_path = Path("/dev/dri") + if not dri_path.exists(): + return [] + nodes = list(dri_path.glob("renderD*")) + list(dri_path.glob("card*")) + return [node for node in nodes if os.access(node, os.R_OK | os.W_OK)] + + +def _resolve_mujoco_backend(requested_backend: str) -> tuple[str, str]: + requested_backend = requested_backend.lower() + if requested_backend != "auto": + return requested_backend, "requested" + + env_backend = os.environ.get("MUJOCO_GL") + if env_backend: + return env_backend.lower(), "env" + + if _find_accessible_dri_nodes(): + return "egl", "auto-gpu" + return "osmesa", "auto-cpu" + + +def _configure_mujoco_env(requested_backend: str) -> str: + backend, source = _resolve_mujoco_backend(requested_backend) + if backend not in {"egl", "osmesa", "glfw"}: + raise ValueError(f"Unsupported MuJoCo GL backend: {backend!r}. Use auto, egl, osmesa, or glfw.") + + os.environ["MUJOCO_GL"] = backend + if backend == "egl": + os.environ["PYOPENGL_PLATFORM"] = "egl" + elif backend == "osmesa": + os.environ["PYOPENGL_PLATFORM"] = "osmesa" + return f"{backend} ({source})" + + +def _import_libero() -> None: + global benchmark, get_libero_path, OffScreenRenderEnv + try: + from libero.libero import benchmark as libero_benchmark + from libero.libero import get_libero_path as libero_get_libero_path + from libero.libero.envs import OffScreenRenderEnv as libero_offscreen_render_env + except ImportError as exc: # pragma: no cover - environment-specific dependency + raise RuntimeError( + "Failed to import LIBERO. Make sure the LIBERO environment is activated. " + f"python={sys.executable!r}, import_error={exc!r}" + ) from exc + + benchmark = libero_benchmark + get_libero_path = libero_get_libero_path + OffScreenRenderEnv = libero_offscreen_render_env + + +def _wait_for_server(client: ActionEnvironmentClient, timeout_s: float) -> None: + start = time.perf_counter() + while time.perf_counter() - start < timeout_s: + if client.check_health(): + return + time.sleep(1.0) + raise RuntimeError(f"Timed out waiting for server at {client.server_url}") + + +def _get_libero_env( + task: Any, + *, + resolution: int, + seed: int, + render_gpu_device_id: int, +) -> tuple[Any, str]: + task_description = str(task.language) + task_bddl_file = os.path.join(get_libero_path("bddl_files"), task.problem_folder, task.bddl_file) + env_args = { + "bddl_file_name": task_bddl_file, + "camera_heights": resolution, + "camera_widths": resolution, + "render_gpu_device_id": render_gpu_device_id, + } + env = OffScreenRenderEnv(**env_args) + env.seed(seed) + return env, task_description + + +def _get_libero_dummy_action() -> list[float]: + return [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -1.0] + + +def _get_libero_image( + obs: dict[str, Any], + camera: str, + *, + flip_images: bool, + rotate_180: bool, +) -> np.ndarray: + if camera == "agentview": + image = obs["agentview_image"] + elif camera == "wrist": + image = obs["robot0_eye_in_hand_image"] + else: + raise ValueError(f"Unsupported camera={camera!r}. Use 'agentview' or 'wrist'.") + + if rotate_180: + image = image[::-1, ::-1] + if flip_images: + image = np.flipud(image) + return image + + +def _get_libero_images( + obs: dict[str, Any], + cameras: list[str], + *, + flip_images: bool, + rotate_180: bool, +) -> list[np.ndarray]: + """Get images from multiple cameras.""" + return [_get_libero_image(obs, camera, flip_images=flip_images, rotate_180=rotate_180) for camera in cameras] + + +def _ensure_uint8_image(image: np.ndarray) -> np.ndarray: + if image.dtype != np.uint8: + if image.max() <= 1.0: + image = (image * 255.0).round().astype(np.uint8) + else: + image = image.astype(np.uint8) + return image + + +def _save_gif(frames: list[Image.Image], output_path: Path, fps: int) -> None: + if not frames: + return + duration_ms = int(1000 / fps) if fps > 0 else 100 + output_path.parent.mkdir(parents=True, exist_ok=True) + first, *rest = frames + first.save( + output_path, + save_all=True, + append_images=rest, + duration=duration_ms, + loop=0, + ) + + +def _decode_b64_frames(b64_frames: list[str]) -> list[Image.Image]: + """Decode a list of base64-encoded PNG strings into PIL Images.""" + images: list[Image.Image] = [] + for b64 in b64_frames: + raw = base64.b64decode(b64) + images.append(Image.open(io.BytesIO(raw)).convert("RGB")) + return images + + +def _save_comparison_gif( + comparison_windows: list[tuple[list[Image.Image], list[Image.Image]]], + output_path: Path, + fps: int, + target_height: int = 256, + separator_width: int = 4, +) -> None: + """Create and save a side-by-side comparison GIF (Action prediction | env rollout). + + Each window is a (action_frames, env_frames) pair from one prediction call. + Frames are paired index-by-index; the conditioning frame (index 0) of + subsequent windows is skipped to avoid duplicating the boundary frame. + """ + from PIL import ImageDraw + + combined_frames: list[Image.Image] = [] + banner_h = 16 + + for window_idx, (action_frames, env_frames) in enumerate(comparison_windows): + n = min(len(action_frames), len(env_frames)) + start = 1 if window_idx > 0 else 0 + for i in range(start, n): + action_img = action_frames[i] + env_img = env_frames[i] + + action_w = int(action_img.width * target_height / action_img.height) + env_w = int(env_img.width * target_height / env_img.height) + action_resized = action_img.resize((action_w, target_height), Image.Resampling.BILINEAR) + env_resized = env_img.resize((env_w, target_height), Image.Resampling.BILINEAR) + + total_w = action_w + separator_width + env_w + total_h = target_height + banner_h + combined = Image.new("RGB", (total_w, total_h), color=0) + + draw = ImageDraw.Draw(combined) + draw.rectangle([(0, 0), (action_w, banner_h)], fill=(30, 30, 60)) + draw.rectangle([(action_w + separator_width, 0), (total_w, banner_h)], fill=(30, 60, 30)) + draw.text((4, 1), "Action Prediction", fill=(100, 180, 255)) + draw.text((action_w + separator_width + 4, 1), "Environment", fill=(100, 255, 100)) + + combined.paste(action_resized, (0, banner_h)) + combined.paste(env_resized, (action_w + separator_width, banner_h)) + combined_frames.append(combined) + + if combined_frames: + _save_gif(combined_frames, output_path, fps) + + +def _select_action_chunk(actions: list[list[float]], action_horizon: int) -> list[list[float]]: + if action_horizon <= 0 or action_horizon >= len(actions): + return actions + return actions[:action_horizon] + + +def _format_action(action: list[float], action_dim: int) -> list[float]: + if len(action) < action_dim: + raise ValueError(f"Action dimension {len(action)} smaller than expected {action_dim}") + return action[:action_dim] + + +def _remap_gripper(action: list[float], mode: str) -> list[float]: + """Map the model's gripper command to the LIBERO env's [-1, 1] (negative = open). + + The right mapping depends on the gripper convention of the dataset the policy + was trained on (the server denormalizes back to that raw convention): + + * ``zero_one`` (NVIDIA LIBERO_LeRobot_v3): raw gripper in [0, 1]; the env wants + [-1, 1] with negative=open. The i4/cosmos-rl reference BINARIZES this to hard + {-1, +1} via ``-sign(2g - 1)`` (not the continuous ``1 - 2g`` from issue #50). + For a confident policy the two agree (g~0/1), but an undertrained policy emits + g~0.5 where continuous ``1-2g``~0 never actuates the gripper -> grasps fail. + Binarizing matches the reference and is robust to weak checkpoints. + * ``pm_one`` (community ``lerobot/libero_*``): raw gripper already in {-1, +1} + (robosuite convention) -> pass through (clamped). + * ``pm_one_flip``: {-1, +1} but with inverted open/close sign. + """ + action = list(action) # avoid mutating the caller's list + g = action[-1] + if mode == "zero_one": + action[-1] = max(-1.0, min(1.0, g * 2.0 - 1.0)) * -1.0 # [0,1] -> [-1,1], negative=open (issue #50) + elif mode == "pm_one": + action[-1] = max(-1.0, min(1.0, g)) + elif mode == "pm_one_flip": + action[-1] = max(-1.0, min(1.0, -g)) + else: + raise ValueError(f"Unknown gripper_mode={mode!r}. Use zero_one/pm_one/pm_one_flip.") + return action + + +def _infer_rotation_space(action_dim: int, rotation_space: str) -> str: + if rotation_space != "auto": + return rotation_space + return libero_rotation_space_from_action_dim(action_dim) + + +def _obs_to_pose(obs: dict[str, Any]) -> tuple[np.ndarray, np.ndarray]: + position = np.asarray(obs["robot0_eef_pos"], dtype=np.float32) + quat = np.asarray(obs["robot0_eef_quat"], dtype=np.float32) + rotation = R.from_quat(quat).as_matrix() + return position, rotation + + +def _anchored_action_to_delta( + anchored_action: np.ndarray, + base_pose: tuple[np.ndarray, np.ndarray], + current_pose: tuple[np.ndarray, np.ndarray], + rotation_space: str, +) -> np.ndarray: + anchored_translation = anchored_action[:3] + rotation_dim = anchored_action.shape[0] - 4 + anchored_rotation = anchored_action[3 : 3 + rotation_dim] + gripper = anchored_action[3 + rotation_dim : 4 + rotation_dim] + + base_pos, base_rot = base_pose + current_pos, current_rot = current_pose + + if rotation_space == "3d": + anchored_rot = R.from_rotvec(anchored_rotation).as_matrix() + elif rotation_space == "6d": + anchored_rot = _rotation_repr_to_mat(anchored_rotation, rotation_space) + elif rotation_space == "9d": + anchored_rot = anchored_rotation.reshape(3, 3) + else: + raise ValueError(f"Unsupported rotation_space={rotation_space!r}. Use 3d/6d/9d.") + target_rot = base_rot @ anchored_rot + target_pos = base_pos + base_rot @ anchored_translation + delta_pos = target_pos - current_pos + delta_rot = target_rot @ current_rot.T + delta_rotvec = R.from_matrix(delta_rot).as_rotvec() + + return np.concatenate([delta_pos, delta_rotvec, gripper], axis=0) + + +def _framewise_action_to_delta( + framewise_action: np.ndarray, + rotation_space: str, +) -> np.ndarray: + """Convert a frame-wise policy action to LIBERO's 7D simulator command. + + Frame-wise actions are already per-step deltas in the LIBERO controller's + convention (see ``LiberoDataset`` with ``action_space='frame_wise_relative'``), + so the only conversion required is decoding the chosen rotation + representation back to a rotation vector. No anchor/current pose is needed. + """ + if rotation_space == "3d": + return framewise_action + + translation = framewise_action[:3] + rotation_dim = framewise_action.shape[0] - 4 + rotation_repr = framewise_action[3 : 3 + rotation_dim] + gripper = framewise_action[3 + rotation_dim : 4 + rotation_dim] + rotation_delta = _rotation_repr_to_mat(rotation_repr, rotation_space) + + delta_pos = translation + delta_rotvec = R.from_matrix(rotation_delta).as_rotvec() + return np.concatenate([delta_pos, delta_rotvec, gripper], axis=0) + + +def _run_episode( + env: Any, + client: ActionEnvironmentClient, + *, + cameras: list[str], + flip_images: bool, + rotate_180: bool, + action_horizon: int, + action_dim: int, + action_space: str, + rotation_space: str, + gripper_mode: str, + max_steps: int, + warmup_steps: int, + initial_state: np.ndarray | None, + gif_path: Path | None, + gif_fps: int, + comparison_path: Path | None = None, +) -> EpisodeResult: + env.reset() + if initial_state is not None: + obs = env.set_init_state(initial_state) + else: + obs = env.get_observation() + + action_queue: list[list[float]] = [] + base_pose: tuple[np.ndarray, np.ndarray] | None = None + step = 0 + success = False + gif_frames: list[Image.Image] = [] + action_log: list[list[float]] = [] + is_multi_view = len(cameras) > 1 + resolved_rotation_space = _infer_rotation_space(action_dim, rotation_space) + + comparison_windows: list[tuple[list[Image.Image], list[Image.Image]]] = [] + + def record_frame(current_obs: dict[str, Any]) -> None: + if gif_path is None: + return + image = _get_libero_image( + current_obs, + cameras[0], + flip_images=flip_images, + rotate_180=rotate_180, + ) + image = _ensure_uint8_image(image) + gif_frames.append(Image.fromarray(image).convert("RGB")) + + def capture_comparison_frame(current_obs: dict[str, Any]) -> Image.Image: + """Capture an env frame matching Action's input view (multi-view concatenated if applicable).""" + if is_multi_view: + imgs = _get_libero_images(current_obs, cameras, flip_images=flip_images, rotate_180=rotate_180) + concat = client.concatenate_images(imgs) + return Image.fromarray(_ensure_uint8_image(concat)).convert("RGB") + img = _get_libero_image(current_obs, cameras[0], flip_images=flip_images, rotate_180=rotate_180) + return Image.fromarray(_ensure_uint8_image(img)).convert("RGB") + + record_frame(obs) + + while step < max_steps: + if step < warmup_steps: + dummy = _get_libero_dummy_action() + obs, _, _, _ = env.step(dummy) + action_log.append(dummy) + step += 1 + record_frame(obs) + continue + + if not action_queue: + if is_multi_view: + observation_imgs = _get_libero_images( + obs, + cameras, + flip_images=flip_images, + rotate_180=rotate_180, + ) + result = client.predict(observation_imgs) + else: + observation_img = _get_libero_image( + obs, + cameras[0], + flip_images=flip_images, + rotate_180=rotate_180, + ) + result = client.predict(observation_img) + actions = result.get("action", []) + if not actions: + return EpisodeResult(False, step, "Empty action chunk from server", action_log) + action_queue = _select_action_chunk(actions, action_horizon) + + if comparison_path is not None: + action_video_b64 = result.get("video", []) + if action_video_b64: + action_frames = _decode_b64_frames(action_video_b64) + env_comparison_frames = [capture_comparison_frame(obs)] + comparison_windows.append((action_frames, env_comparison_frames)) + + if action_space == "relative": + base_pose = _obs_to_pose(obs) + + raw_action = _format_action(action_queue.pop(0), action_dim) + if action_space == "relative": + if base_pose is None: + raise RuntimeError("Missing base pose for relative action conversion") + current_pose = _obs_to_pose(obs) + action = _anchored_action_to_delta( + np.asarray(raw_action, dtype=np.float32), + base_pose, + current_pose, + resolved_rotation_space, + ) + action_list = action.tolist() + else: + action = _framewise_action_to_delta( + np.asarray(raw_action, dtype=np.float32), + resolved_rotation_space, + ) + action_list = action.tolist() + + # Map the model's gripper command to the env's [-1, 1] per the dataset convention. + action_list = _remap_gripper(action_list, gripper_mode) + + action_log.append(action_list) + obs, _, done, info = env.step(action_list) + step += 1 + record_frame(obs) + + if comparison_path is not None and comparison_windows: + comparison_windows[-1][1].append(capture_comparison_frame(obs)) + + if isinstance(info, dict) and info.get("success"): + success = True + break + if done: + success = True if not isinstance(info, dict) else bool(info.get("success", True)) + break + + if gif_path is not None: + _save_gif(gif_frames, gif_path, gif_fps) + if comparison_path is not None and comparison_windows: + _save_comparison_gif(comparison_windows, comparison_path, gif_fps) + return EpisodeResult(success, step, None, action_log) + + +def _load_initial_states( + task_suite: Any, + task_id: int, + *, + task_description: str, + initial_states_path: str, + episode_idx: int, +) -> np.ndarray | None: + default_initial_states = task_suite.get_task_init_states(task_id) + + if initial_states_path == "DEFAULT": + return np.array(default_initial_states[episode_idx]) + + with open(initial_states_path, "r", encoding="utf-8") as f: + all_initial_states = json.load(f) + + task_key = task_description.replace(" ", "_") + episode_key = f"demo_{episode_idx}" + if not all_initial_states[task_key][episode_key]["success"]: + return None + return np.array(all_initial_states[task_key][episode_key]["initial_state"]) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="LIBERO closed-loop evaluation via Action HTTP server") + parser.add_argument( + "--server_url", type=str, required=True, help="Base URL for Action server (e.g., http://host:8000)" + ) + parser.add_argument("--task_suite", type=str, default="libero_spatial", choices=sorted(TASK_MAX_STEPS.keys())) + parser.add_argument("--num_trials_per_task", type=int, default=10) + parser.add_argument("--task_ids", type=str, default="", help="Comma-separated task IDs to evaluate (default: all)") + parser.add_argument("--image_size", type=int, default=256, help="Model input image size") + parser.add_argument("--env_image_size", type=int, default=256, help="Environment render resolution") + parser.add_argument("--action_horizon", type=int, default=0, help="Actions to execute per request (0=full chunk)") + parser.add_argument("--action_dim", type=int, default=10, help="Action dimension for LIBERO") + parser.add_argument( + "--action_space", + type=str, + default="frame_wise_relative", + choices=["relative", "frame_wise_relative"], + help="Action space expected from the model (relative=anchored, frame_wise_relative=framewise deltas).", + ) + parser.add_argument( + "--rotation_space", + type=str, + default="auto", + choices=["auto", "3d", "6d", "9d"], + help="Rotation representation for anchored actions (auto infers from action_dim).", + ) + parser.add_argument( + "--gripper_mode", + type=str, + default="zero_one", + choices=["zero_one", "pm_one", "pm_one_flip"], + help="Gripper convention of the training data: 'zero_one' = [0,1] (NVIDIA " + "LIBERO_LeRobot_v3, mapped 1-2g); 'pm_one' = {-1,+1} (community lerobot/libero_*, " + "pass-through); 'pm_one_flip' = {-1,+1} with inverted sign.", + ) + parser.add_argument("--domain_name", type=str, default="libero") + parser.add_argument( + "--camera", + type=str, + default="agentview", + help="Camera(s) to use. Single camera: 'agentview' or 'wrist'. Multiple cameras: comma-separated, e.g., 'agentview,wrist'.", + ) + parser.add_argument("--flip_images", action="store_true", help="Flip images vertically before encoding") + parser.add_argument( + "--rotate_180", + action=argparse.BooleanOptionalAction, + default=True, + help="Rotate images by 180 degrees before encoding (default: True; pass --no-rotate-180 to disable)", + ) + parser.add_argument("--warmup_steps", type=int, default=10, help="Stabilization steps with dummy actions") + parser.add_argument("--max_steps", type=int, default=0, help="Override max steps per episode (0=default)") + parser.add_argument("--timeout", type=float, default=30.0, help="HTTP request timeout in seconds") + parser.add_argument("--wait_timeout", type=float, default=60.0, help="Seconds to wait for server health") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--save_gifs", action="store_true", help="Save per-episode GIFs of rendered frames") + parser.add_argument( + "--save_comparison", + action="store_true", + help="Save side-by-side comparison GIFs (Action prediction vs environment rollout)", + ) + parser.add_argument("--gif_fps", type=int, default=20, help="Frames per second for saved GIFs") + parser.add_argument( + "--mujoco_gl", + type=str, + default="auto", + choices=["auto", "egl", "osmesa", "glfw"], + help="MuJoCo GL backend (auto picks egl if /dev/dri is accessible, else osmesa).", + ) + parser.add_argument( + "--render_gpu_device_id", + type=int, + default=-1, + help="GPU device index for EGL rendering (-1 uses default device).", + ) + parser.add_argument( + "--initial_states_path", + type=str, + default="DEFAULT", + help='Path to initial states JSON. Use "DEFAULT" for benchmark defaults.', + ) + parser.add_argument( + "--num_envs", + type=int, + default=1, + help="Number of parallel LIBERO envs (SubprocVectorEnv). >1 runs trials in waves " + "with ONE batched /predict_batch per control step (~num_envs x faster). 1 = serial.", + ) + parser.add_argument("--output_dir", type=str, default="", help="Directory to save evaluation summary JSON") + return parser.parse_args() + + +class _LiberoEnvFactory: + """Picklable env factory for SubprocVectorEnv under the spawn start method. + + spawn pickles each env_fn and re-imports this module in the child, so the + factory must be a top-level class (lambdas/closures are not picklable). The + child sets the GL backend and imports OffScreenRenderEnv locally so its EGL + context is created fresh in the worker process.""" + + def __init__( + self, + *, + bddl_file_name: str, + camera_heights: int, + camera_widths: int, + render_gpu_device_id: int, + mujoco_gl: str, + ) -> None: + self.bddl_file_name = bddl_file_name + self.camera_heights = camera_heights + self.camera_widths = camera_widths + self.render_gpu_device_id = render_gpu_device_id + self.mujoco_gl = mujoco_gl + + def __call__(self) -> Any: + # Resolve to a concrete GPU; -1 (auto) makes EGL device selection race/fail + # across spawned workers (EGLError / "'EGLGLContext' object has no attribute + # '_context'"). Set the GL backend + pin the EGL device BEFORE importing + # OffScreenRenderEnv (which dlopen's the GL stack at import). + dev = self.render_gpu_device_id if self.render_gpu_device_id >= 0 else 0 + os.environ["MUJOCO_GL"] = self.mujoco_gl + if self.mujoco_gl == "egl": + os.environ["PYOPENGL_PLATFORM"] = "egl" + os.environ["MUJOCO_EGL_DEVICE_ID"] = str(dev) + os.environ["EGL_DEVICE_ID"] = str(dev) + elif self.mujoco_gl == "osmesa": + os.environ["PYOPENGL_PLATFORM"] = "osmesa" + from libero.libero.envs import OffScreenRenderEnv as _OffScreenRenderEnv + + return _OffScreenRenderEnv( + bddl_file_name=self.bddl_file_name, + camera_heights=self.camera_heights, + camera_widths=self.camera_widths, + render_gpu_device_id=dev, + ) + + +def _run_task_vectorized( + task: Any, + task_description: str, + *, + num_trials: int, + num_envs: int, + env_image_size: int, + seed: int, + render_gpu_device_id: int, + client: ActionEnvironmentClient, + cameras: list[str], + flip_images: bool, + rotate_180: bool, + action_horizon: int, + action_dim: int, + rotation_space: str, + gripper_mode: str, + max_steps: int, + warmup_steps: int, + init_states: list[np.ndarray | None], +) -> list[dict[str, Any]]: + """Run all `num_trials` of one task across `num_envs` parallel LIBERO envs + (SubprocVectorEnv), in waves. Each control step gathers obs from the ACTIVE + (not-done) envs, issues ONE batched /predict_batch, and steps all active envs; + done envs are masked out. Returns per-trial result dicts in trial order with the + same shape as the serial path's episode_results.""" + import multiprocessing as _mp + + from libero.libero.envs.venv import SubprocVectorEnv + + # LIBERO's SubprocVectorEnv defaults to the fork start method; forked children + # inherit the parent's already-dlopen'd EGL/GL state, which corrupts per-child + # render-context creation (EGLError / 'EGLGLContext' has no attribute '_context'). + # Force spawn so each env worker starts clean — exactly like the (working) serial + # single-process path. spawn pickles env_fns, so the factory below is picklable. + try: + _mp.set_start_method("spawn", force=True) + except RuntimeError: # pragma: no cover - already set + pass + + resolved_rotation_space = _infer_rotation_space(action_dim, rotation_space) + bddl = os.path.join(get_libero_path("bddl_files"), task.problem_folder, task.bddl_file) + + results: list[dict[str, Any]] = [None] * num_trials # type: ignore[list-item] + for t in range(num_trials): + if init_states[t] is None: + results[t] = { + "episode": t, + "success": False, + "steps": 0, + "error": "Skipped due to failed expert demo", + "elapsed_s": 0.0, + } + runnable = [t for t in range(num_trials) if init_states[t] is not None] + if not runnable: + return results + + n = min(num_envs, len(runnable)) + + mujoco_gl = os.environ.get("MUJOCO_GL", "egl") + env_fn = _LiberoEnvFactory( + bddl_file_name=bddl, + camera_heights=env_image_size, + camera_widths=env_image_size, + render_gpu_device_id=render_gpu_device_id, + mujoco_gl=mujoco_gl, + ) + venv = SubprocVectorEnv([env_fn for _ in range(n)]) + try: + venv.seed(seed) + for w0 in range(0, len(runnable), n): + wave = runnable[w0 : w0 + n] # trial indices for this wave + slots = list(range(len(wave))) # env slots in use + t_wave0 = time.perf_counter() + venv.reset(id=slots) + states = np.stack([np.asarray(init_states[t], dtype=np.float64) for t in wave]) + obs_arr = venv.set_init_state(states, id=slots) + obs_by_slot = {s: obs_arr[i] for i, s in enumerate(slots)} + done = {s: False for s in slots} + succ = {s: False for s in slots} + err: dict[int, str | None] = {s: None for s in slots} + nsteps = {s: max_steps for s in slots} + step = 0 + + for _ in range(warmup_steps): + act = np.stack([_get_libero_dummy_action() for _ in slots]) + obs_arr, _, _, _ = venv.step(act, id=slots) + for i, s in enumerate(slots): + obs_by_slot[s] = obs_arr[i] + step += 1 + + while step < max_steps: + active = [s for s in slots if not done[s]] + if not active: + break + obs_batch = [ + _get_libero_images(obs_by_slot[s], cameras, flip_images=flip_images, rotate_180=rotate_180) + for s in active + ] + try: + chunks = client.predict_batch(obs_batch) + except Exception as e: # noqa: BLE001 + for s in active: + done[s] = True + err[s] = f"server error: {e}" + nsteps[s] = step + break + if not chunks or len(chunks) != len(active): + for s in active: + done[s] = True + err[s] = "bad batch response from server" + nsteps[s] = step + break + chunk_by_slot = {s: chunks[k] for k, s in enumerate(active)} + horizon = action_horizon if action_horizon > 0 else len(chunks[0]) + for h in range(horizon): + cur = [s for s in slots if not done[s]] + if not cur or step >= max_steps: + break + env_actions = [] + for s in cur: + raw = _format_action(chunk_by_slot[s][h], action_dim) + a = _framewise_action_to_delta(np.asarray(raw, dtype=np.float32), resolved_rotation_space) + env_actions.append(_remap_gripper(a.tolist(), gripper_mode)) + obs_arr, _, d, info = venv.step(np.stack(env_actions), id=cur) + step += 1 + for i, s in enumerate(cur): + obs_by_slot[s] = obs_arr[i] + di = bool(d[i]) + ii = info[i] if isinstance(info, (list, np.ndarray)) else info + is_succ = bool(ii.get("success")) if isinstance(ii, dict) else False + if is_succ: + done[s], succ[s], nsteps[s] = True, True, step + elif di: + # mirror serial: done w/o explicit success defaults to success + done[s] = True + succ[s] = ii.get("success", True) if isinstance(ii, dict) else True + nsteps[s] = step + per_ep_elapsed = round((time.perf_counter() - t_wave0) / max(1, len(wave)), 3) + for s, t in zip(slots, wave): + results[t] = { + "episode": t, + "success": bool(succ[s]), + "steps": int(nsteps[s]), + "error": err[s], + "elapsed_s": per_ep_elapsed, + } + finally: + try: + venv.close() + except Exception: # noqa: BLE001 + pass + return results + + +def main() -> None: + args = _parse_args() + random.seed(args.seed) + np.random.seed(args.seed) + + if args.save_gifs and not args.output_dir: + raise ValueError("--save_gifs requires --output_dir to be set") + if args.save_comparison and not args.output_dir: + raise ValueError("--save_comparison requires --output_dir to be set") + + # Parse cameras from comma-separated string + cameras = [c.strip() for c in args.camera.split(",") if c.strip()] + if not cameras: + raise ValueError("At least one camera must be specified") + for cam in cameras: + if cam not in ("agentview", "wrist"): + raise ValueError(f"Unsupported camera={cam!r}. Use 'agentview' or 'wrist'.") + + mujoco_backend = _configure_mujoco_env(args.mujoco_gl) + _import_libero() + + client = ActionEnvironmentClient( + server_url=args.server_url, + domain_name=args.domain_name, + prompt="", + image_size=args.image_size, + timeout=args.timeout, + ) + print(f"MuJoCo GL backend: {mujoco_backend}", flush=True) + print("Waiting for model server...", flush=True) + _wait_for_server(client, args.wait_timeout) + print(f"Connected to model server: {client.get_info()}", flush=True) + + benchmark_dict = benchmark.get_benchmark_dict() + task_suite = benchmark_dict[args.task_suite]() + num_tasks = int(task_suite.n_tasks) + + if args.task_ids: + selected_task_ids = [int(t) for t in args.task_ids.split(",") if t.strip()] + else: + selected_task_ids = list(range(num_tasks)) + + max_steps = args.max_steps if args.max_steps > 0 else TASK_MAX_STEPS[args.task_suite] + + total_episodes = 0 + total_successes = 0 + task_results: list[dict[str, Any]] = [] + + output_dir = Path(args.output_dir) if args.output_dir else None + gif_root = output_dir / "gifs" if output_dir and args.save_gifs else None + comparison_root = output_dir / "comparisons" if output_dir and args.save_comparison else None + + for task_id in selected_task_ids: + task = task_suite.get_task(task_id) + + # ---- Vectorized path: N parallel envs + one batched /predict_batch per step ---- + if args.num_envs > 1: + task_description = str(task.language) + client.prompt = _augment_task_prompt_with_viewpoint(task_description, cameras) + init_states = [ + _load_initial_states( + task_suite, + task_id, + task_description=task_description, + initial_states_path=args.initial_states_path, + episode_idx=e, + ) + for e in range(args.num_trials_per_task) + ] + episode_results = _run_task_vectorized( + task, + task_description, + num_trials=args.num_trials_per_task, + num_envs=args.num_envs, + env_image_size=args.env_image_size, + seed=args.seed, + render_gpu_device_id=args.render_gpu_device_id, + client=client, + cameras=cameras, + flip_images=args.flip_images, + rotate_180=args.rotate_180, + action_horizon=args.action_horizon, + action_dim=args.action_dim, + rotation_space=args.rotation_space, + gripper_mode=args.gripper_mode, + max_steps=max_steps, + warmup_steps=args.warmup_steps, + init_states=init_states, + ) + task_episodes = 0 + task_successes = 0 + for er in episode_results: + task_episodes += 1 + total_episodes += 1 + if er["success"]: + task_successes += 1 + total_successes += 1 + print( + f"Task {task_id} | Episode {er['episode'] + 1}/{args.num_trials_per_task} | " + f"success={er['success']} steps={er['steps']} elapsed_s={er['elapsed_s']:.1f} | " + f"task SR {task_successes}/{task_episodes} ({100.0 * task_successes / max(1, task_episodes):.1f}%) | " + f"overall SR {total_successes}/{total_episodes} " + f"({100.0 * total_successes / max(1, total_episodes):.1f}%)", + flush=True, + ) + task_success_rate = float(task_successes) / float(task_episodes) if task_episodes > 0 else 0.0 + task_results.append( + { + "task_id": task_id, + "task_description": task_description, + "episodes": task_episodes, + "successes": task_successes, + "success_rate": task_success_rate, + "episode_results": episode_results, + } + ) + print( + f"Task {task_id} summary: {task_successes}/{task_episodes} ({task_success_rate * 100:.1f}%)", + flush=True, + ) + continue + + env, task_description = _get_libero_env( + task, + resolution=args.env_image_size, + seed=args.seed, + render_gpu_device_id=args.render_gpu_device_id, + ) + + task_episodes = 0 + task_successes = 0 + episode_results: list[dict[str, Any]] = [] + + for episode_idx in range(args.num_trials_per_task): + episode_t0 = time.perf_counter() + client.prompt = _augment_task_prompt_with_viewpoint(task_description, cameras) + initial_state = _load_initial_states( + task_suite, + task_id, + task_description=task_description, + initial_states_path=args.initial_states_path, + episode_idx=episode_idx, + ) + if initial_state is None: + episode_elapsed_s = time.perf_counter() - episode_t0 + episode_results.append( + { + "episode": episode_idx, + "success": False, + "steps": 0, + "error": "Skipped due to failed expert demo", + "elapsed_s": round(episode_elapsed_s, 3), + } + ) + print( + f"Task {task_id} | Episode {episode_idx + 1}/{args.num_trials_per_task} | " + "success=False steps=0 " + f"elapsed_s={episode_elapsed_s:.1f} " + "error='Skipped due to failed expert demo'", + flush=True, + ) + continue + + gif_path = ( + gif_root / f"task_{task_id:03d}" / f"episode_{episode_idx:03d}.gif" if gif_root is not None else None + ) + comparison_path = ( + comparison_root / f"task_{task_id:03d}" / f"episode_{episode_idx:03d}.gif" + if comparison_root is not None + else None + ) + try: + result = _run_episode( + env, + client, + cameras=cameras, + flip_images=args.flip_images, + rotate_180=args.rotate_180, + action_horizon=args.action_horizon, + action_dim=args.action_dim, + action_space=args.action_space, + rotation_space=args.rotation_space, + gripper_mode=args.gripper_mode, + max_steps=max_steps, + warmup_steps=args.warmup_steps, + initial_state=initial_state, + gif_path=gif_path, + gif_fps=args.gif_fps, + comparison_path=comparison_path, + ) + except Exception as exc: + result = EpisodeResult(False, 0, str(exc), []) + episode_elapsed_s = time.perf_counter() - episode_t0 + + task_episodes += 1 + total_episodes += 1 + if result.success: + task_successes += 1 + total_successes += 1 + + episode_results.append( + { + "episode": episode_idx, + "success": result.success, + "steps": result.steps, + "error": result.error, + "elapsed_s": round(episode_elapsed_s, 3), + } + ) + + # Save per-episode action log as JSON + if output_dir is not None and result.actions: + action_log_dir = output_dir / "actions" / f"task_{task_id:03d}" + action_log_dir.mkdir(parents=True, exist_ok=True) + action_log_path = action_log_dir / f"episode_{episode_idx:03d}.json" + action_log_path.write_text( + json.dumps(result.actions, indent=2), + encoding="utf-8", + ) + + client.notify_next_episode() + + print( + f"Task {task_id} | Episode {episode_idx + 1}/{args.num_trials_per_task} | " + f"success={result.success} steps={result.steps} elapsed_s={episode_elapsed_s:.1f} | " + f"task SR {task_successes}/{task_episodes} ({100.0 * task_successes / max(1, task_episodes):.1f}%) | " + f"overall SR {total_successes}/{total_episodes} ({100.0 * total_successes / max(1, total_episodes):.1f}%)", + flush=True, + ) + + task_success_rate = float(task_successes) / float(task_episodes) if task_episodes > 0 else 0.0 + task_results.append( + { + "task_id": task_id, + "task_description": task_description, + "episodes": task_episodes, + "successes": task_successes, + "success_rate": task_success_rate, + "episode_results": episode_results, + } + ) + print( + f"Task {task_id} summary: {task_successes}/{task_episodes} ({task_success_rate * 100:.1f}%)", + flush=True, + ) + # Close the env (and its EGL/MuJoCo render context) before the next task. + # Leaving it open leaks one EGL context per task and hangs after ~8 tasks. + try: + env.close() + except Exception: + pass + + overall_success_rate = float(total_successes) / float(total_episodes) if total_episodes > 0 else 0.0 + summary = { + "task_suite": args.task_suite, + "total_episodes": total_episodes, + "total_successes": total_successes, + "overall_success_rate": overall_success_rate, + "num_trials_per_task": args.num_trials_per_task, + "selected_task_ids": selected_task_ids, + "action_space": args.action_space, + "rotation_space": _infer_rotation_space(args.action_dim, args.rotation_space), + "action_dim": args.action_dim, + "task_results": task_results, + } + + print( + f"Overall success rate: {total_successes}/{total_episodes} ({overall_success_rate * 100:.1f}%)", + flush=True, + ) + + if output_dir is not None: + output_dir.mkdir(parents=True, exist_ok=True) + summary_path = output_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") + print(f"Saved summary to {summary_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/cosmos_framework/utils/vfm/model_loader.py b/cosmos_framework/utils/vfm/model_loader.py index fd8b9560..92e67ab0 100644 --- a/cosmos_framework/utils/vfm/model_loader.py +++ b/cosmos_framework/utils/vfm/model_loader.py @@ -252,10 +252,17 @@ def _load_model( keys_to_skip_loading=keys_to_skip_loading or [], ) + # Single-rank load (e.g. the action-policy inference server): force no_dist so + # ``dcp.load`` skips the collective ``gather_object`` over the load plan, which + # pickles the plan and can fail on training/EMA DCPs. Multi-rank loads keep the + # default distributed path. + no_dist = not (dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1) + dcp.load( state_dict=state_dict, storage_reader=storage_reader, planner=load_planner, + no_dist=no_dist, ) log.info(f"Successfully loaded model from {checkpoint_path}") diff --git a/docs/action_policy_libero_sft.md b/docs/action_policy_libero_sft.md new file mode 100644 index 00000000..b1de5df7 --- /dev/null +++ b/docs/action_policy_libero_sft.md @@ -0,0 +1,139 @@ +# Cosmos3-Nano LIBERO-10 action-policy SFT + +Full SFT of the public `nvidia/Cosmos3-Nano` base into a LIBERO-10 action +policy: vision + language in, action chunks out. + +To match the LIBERO-10 SR reported in Cosmos3, we provide **two presets** (both +lr 5e-5, warmup 500, cycle 16000, gbs 2048): + +- **(A) libero_10-only** — trains on `libero_10` alone; peaks by ~iter 1500 + (max_iter 2000). Fast. + `action_policy_libero_nano` + `action_policy_libero_repro.toml` + + `launch_sft_action_policy_libero.sh`. +- **(B) libero-all** — equal mix of all 4 LIBERO suites; needs longer training + (max_iter 5000). + `action_policy_libero_all_nano` + `action_policy_libero_all_repro.toml` + + `launch_sft_action_policy_libero_all.sh`. + +| Piece | Path | +| ---------------- | ----------------------------------------------------------------------------------------------- | +| Dataset | `cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py` (`LIBEROLeRobotDataset`) | +| SFT wrapper | `get_action_libero_sft_dataset` in `.../datasets/action_sft_dataset.py` | +| Norm stats | `.../normalizer_stats/libero_native_frame_wise_relative_rot6d.json` | +| Experiment | `cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py` | +| Run TOML | `examples/toml/sft_config/action_policy_libero_repro.toml` | +| Launch | `examples/launch_sft_action_policy_libero.sh` | +| Inference server | `cosmos_framework/scripts/action_policy_server_libero.py` | +| Closed-loop eval | `cosmos_framework/simulation/libero/closed_loop_eval.py` | + +## 1. Data + +`LIBEROLeRobotDataset` reads a local LeRobot dir. Use the 20 FPS +[`nvidia/LIBERO_LeRobot_v3`](https://huggingface.co/datasets/nvidia/LIBERO_LeRobot_v3), +which the bundled `quantile_rot` stats and the 20 Hz eval assume. + +**Preset A (libero_10-only)** — `LIBERO_ROOT` points at the `libero_10` suite dir: + +```bash +hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset \ + --include 'libero_10/**' --local-dir /LIBERO_LeRobot_v3 +export LIBERO_ROOT=/LIBERO_LeRobot_v3/libero_10 +``` + +**Preset B (libero-all)** — download all 4 suites; `LIBERO_ROOT` is the **parent** dir: + +```bash +hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset --local-dir /LIBERO_LeRobot_v3 +export LIBERO_ROOT=/LIBERO_LeRobot_v3 # parent of libero_spatial/object/goal/10 +``` + +Actions are `frame_wise_relative` rot6d (10D = pos 3 + rot6d 6 + gripper 1), +`concat_view` (third-person + wrist, each 256×256 → 256×512), `quantile_rot` +normalized. The pipeline snaps the 256×512 concat to a 192×320 model canvas; the +eval server reproduces the same snap (§4). + +## 2. Train + +Common env, then pick a preset launcher: + +```bash +export LD_LIBRARY_PATH='' # NGC container: avoid torch._C import error +export BASE_CHECKPOINT_PATH= +export WAN_VAE_PATH= +export IMAGINAIRE_OUTPUT_ROOT=/path/to/output_root + +# Preset A — libero_10-only (LIBERO_ROOT = the libero_10 suite dir): +export LIBERO_ROOT=/LIBERO_LeRobot_v3/libero_10 +bash examples/launch_sft_action_policy_libero.sh # HSDP 2x8; set NNODES/NODE_RANK/MASTER_ADDR per node + +# Preset B — libero-all 4-suite (LIBERO_ROOT = the LIBERO_LeRobot_v3 parent dir): +export LIBERO_ROOT=/LIBERO_LeRobot_v3 +bash examples/launch_sft_action_policy_libero_all.sh # HSDP 2x8; needs ~4500 iters to converge +``` + +Both recipes set lr 5e-5, warmup 500, cycle 16000, `save_iter=500`, HSDP 2x8 (global +batch 2048 = `max_samples_per_batch` 128 × 16 ranks × grad_accum 1). They differ only in +`max_iter`: **2000** for libero_10-only (peaks ~iter 1500), **5000** for libero-all +(the 4-suite mix takes longer to converge on libero_10, ~iter 4500). + +## 3. Closed-loop eval + +Start the policy server on a **trained** checkpoint (the base DCP has no action +heads), then run the LIBERO simulator client against it. Same for both presets — +`action_policy_libero_nano` supplies the model config for either run's +checkpoint; just point `--checkpoint-path` at the one you trained. + +```bash +python -m cosmos_framework.scripts.action_policy_server_libero \ + --experiment action_policy_libero_nano \ + --experiment-overrides "model.config.tokenizer.vae_path=$WAN_VAE_PATH" \ + --checkpoint-path /checkpoints/iter_000001500 \ + --action-normalization quantile_rot \ + --action-stats-path cosmos_framework/data/vfm/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json \ + --raw-action-dim 10 --fps 20 --port 8000 +``` + +The LIBERO sim needs a separate venv (robosuite/mujoco pins conflict with the +training env): + +```bash +# Optional — only on a headless container without working GPU EGL: +# export NVIDIA_DRIVER_CAPABILITIES=all +# apt-get install -y libegl1 libglvnd0 libgl1 libglib2.0-0 ffmpeg +# mkdir -p /usr/share/glvnd/egl_vendor.d +# echo '{"file_format_version":"1.0.0","ICD":{"library_path":"libEGL_nvidia.so.0"}}' \ +# > /usr/share/glvnd/egl_vendor.d/10_nvidia.json + +uv venv --python 3.10 .libenv && VV=.libenv/bin/python +git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git && \ + uv pip install -p $VV -e LIBERO -r LIBERO/requirements.txt +uv pip install -p $VV "robosuite==1.4.1" "mujoco==2.3.7" "torch<2.6" loguru requests scipy pillow numpy +mkdir -p ~/.libero && touch ~/.libero/config.yaml +RS=$($VV -c "import robosuite,os;print(os.path.dirname(robosuite.__file__))"); $VV "$RS/scripts/setup_macros.py" +$VV -c "from libero.libero import set_libero_default_path; set_libero_default_path()" + +MUJOCO_GL=egl PYTHONPATH=$PWD:$PWD/LIBERO $VV \ + cosmos_framework/simulation/libero/closed_loop_eval.py \ + --server_url http://localhost:8000 \ + --task_suite libero_10 --num_trials_per_task 50 --num_envs 8 \ + --camera agentview,wrist --image_size 256 \ + --action_space frame_wise_relative --rotation_space 6d --action_dim 10 \ + --output_dir results/libero_closed_loop_10 +``` + +## 4. Heads-up + +- **Lower-memory GPUs** — reduce the per-rank batch: + `--opts dataloader_train.max_samples_per_batch=64` (scale `replicate` to keep + global batch 2048). + +Eval parity — the client/server already handle these; verify if accuracy is low: + +- **Concat layout** — run with `--camera agentview,wrist --image_size 256` so the + 256×512 concat matches training (the server snaps it to 192×320 identically). +- **Gripper** — model emits `[0, 1]`; the env wants `[-1, 1]` (negative = open). + The client applies `1 − 2·g`; flip the sign if the gripper never opens. +- **Image orientation** — sim frames are rotated 180° vs training; the client + rotates them back. +- **Normalization** — start the server with `--action-normalization quantile_rot` + and the bundled rot6d stats, or actions come out at the wrong scale. diff --git a/examples/launch_sft_action_policy_libero.sh b/examples/launch_sft_action_policy_libero.sh new file mode 100755 index 00000000..29a9a16a --- /dev/null +++ b/examples/launch_sft_action_policy_libero.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# Structured-TOML launch for action_policy_libero_nano — Cosmos3-Nano LIBERO +# action-policy SFT (HSDP, full SFT). Drives cosmos_framework.scripts.train +# against examples/toml/sft_config/action_policy_libero_repro.toml. +# +# Point LIBERO_ROOT at the libero_10 suite ONLY. Use the 20 FPS +# nvidia/LIBERO_LeRobot_v3. The default recipe is HSDP 2x8 (global batch 2048); +# set NNODES/NODE_RANK/MASTER_ADDR per node. +# See docs/action_policy_libero_sft.md. +# +# Required env vars: +# LIBERO_ROOT local LIBERO-10 LeRobot dataset dir, e.g. /libero_10 (no default) +# Optional env vars (defaults below; override to relocate data/checkpoints): +# BASE_CHECKPOINT_PATH default: examples/checkpoints/Cosmos3-Nano +# WAN_VAE_PATH default: examples/checkpoints/wan22_vae/Wan2.2_VAE.pth +# HF_TOKEN if any tokenizer download requires gated HF access +# OUTPUT_ROOT default: outputs/train +# +# Pre-sync the 20 FPS suite once: +# hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset --include 'libero_10/**' --local-dir +# export LIBERO_ROOT=/libero_10 +# +# Usage (HSDP 2x8; set NNODES/NODE_RANK/MASTER_ADDR per node): +# LIBERO_ROOT=/libero_10 bash examples/launch_sft_action_policy_libero.sh + +TOML_FILE="examples/toml/sft_config/action_policy_libero_repro.toml" +: "${BASE_CHECKPOINT_PATH:=examples/checkpoints/Cosmos3-Nano}" + +# LIBEROLeRobotDataset reads ${oc.env:LIBERO_ROOT} directly (a LOCAL LeRobot dir); +# export it so torchrun (launched in this shell) inherits it. +export LIBERO_ROOT="${LIBERO_ROOT:-}" + +EXTRA_DATASET_CHECK='[[ -f "$LIBERO_ROOT/meta/info.json" ]] || { echo "ERROR: LIBERO_ROOT must be a local LeRobot dir containing meta/info.json (got: '\''$LIBERO_ROOT'\''). Pre-sync: hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset --include '\''libero_10/**'\'' --local-dir (then LIBERO_ROOT=/libero_10). See docs/action_policy_libero_sft.md" >&2; exit 1; }' + +# Extra Hydra overrides from the environment: a space-separated string word-split into +# the TAIL_OVERRIDES array. An exported string survives `bash ` (a child +# process), unlike a TAIL_OVERRIDES array set in your shell. Use it for smoke runs, +# e.g. EXTRA_TAIL_OVERRIDES="trainer.max_iter=5 job.wandb_mode=offline". +TAIL_OVERRIDES=( + ${EXTRA_TAIL_OVERRIDES:-} +) + +source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh" diff --git a/examples/launch_sft_action_policy_libero_all.sh b/examples/launch_sft_action_policy_libero_all.sh new file mode 100755 index 00000000..67eeb6f4 --- /dev/null +++ b/examples/launch_sft_action_policy_libero_all.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# Structured-TOML launch for action_policy_libero_all_nano — Cosmos3-Nano +# LIBERO-all (4-suite) action-policy SFT (HSDP, full SFT). Drives +# cosmos_framework.scripts.train against +# examples/toml/sft_config/action_policy_libero_all_repro.toml. +# +# Trains on all 4 LIBERO suites (equal mix). Point LIBERO_ROOT at the +# LIBERO_LeRobot_v3 PARENT dir (containing libero_spatial/object/goal/10), NOT a +# single suite. Use the 20 FPS nvidia/LIBERO_LeRobot_v3. Default recipe is +# HSDP 2x8 (global batch 2048); set NNODES/NODE_RANK/MASTER_ADDR per node. +# The 4-suite mix is coverage-limited — it needs ~4500 iters to reach ~95% on +# libero_10 (max_iter defaults to 5000). See docs/action_policy_libero_sft.md. +# +# Required env vars: +# LIBERO_ROOT local LIBERO_LeRobot_v3 PARENT dir (no default) +# Optional env vars (defaults below; override to relocate data/checkpoints): +# BASE_CHECKPOINT_PATH default: examples/checkpoints/Cosmos3-Nano +# WAN_VAE_PATH default: examples/checkpoints/wan22_vae/Wan2.2_VAE.pth +# HF_TOKEN if any tokenizer download requires gated HF access +# OUTPUT_ROOT default: outputs/train +# +# Pre-sync all 4 suites once: +# hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset --local-dir +# export LIBERO_ROOT= +# +# Usage (HSDP 2x8; set NNODES/NODE_RANK/MASTER_ADDR per node): +# LIBERO_ROOT= bash examples/launch_sft_action_policy_libero_all.sh + +TOML_FILE="examples/toml/sft_config/action_policy_libero_all_repro.toml" +: "${BASE_CHECKPOINT_PATH:=examples/checkpoints/Cosmos3-Nano}" + +# The libero-all experiment reads ${oc.env:LIBERO_ROOT}/ for each of the 4 +# suites; export the PARENT dir so torchrun (launched in this shell) inherits it. +export LIBERO_ROOT="${LIBERO_ROOT:-}" + +EXTRA_DATASET_CHECK='for _s in libero_spatial libero_object libero_goal libero_10; do [[ -f "$LIBERO_ROOT/$_s/meta/info.json" ]] || { echo "ERROR: LIBERO_ROOT must be the LIBERO_LeRobot_v3 parent dir containing all 4 suites (missing $_s; got: '\''$LIBERO_ROOT'\''). Pre-sync: hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset --local-dir (then LIBERO_ROOT=). See docs/action_policy_libero_sft.md" >&2; exit 1; }; done' + +# Extra Hydra overrides from the environment: a space-separated string word-split into +# the TAIL_OVERRIDES array. An exported string survives `bash ` (a child +# process), unlike a TAIL_OVERRIDES array set in your shell. Use it for smoke runs, +# e.g. EXTRA_TAIL_OVERRIDES="trainer.max_iter=5 job.wandb_mode=offline". +TAIL_OVERRIDES=( + ${EXTRA_TAIL_OVERRIDES:-} +) + +source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh" diff --git a/examples/toml/sft_config/action_policy_libero_all_repro.toml b/examples/toml/sft_config/action_policy_libero_all_repro.toml new file mode 100644 index 00000000..bb4d7c6a --- /dev/null +++ b/examples/toml/sft_config/action_policy_libero_all_repro.toml @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# LIBERO-all (4-suite) action-policy SFT run config for the +# `action_policy_libero_all_nano` experiment. Trains on all 4 LIBERO suites +# (equal mix); LIBERO_ROOT is the LIBERO_LeRobot_v3 PARENT dir. The 4-suite mix +# is coverage-limited — it needs ~4500 iters to reach ~95% on libero_10 (vs the +# libero_10-only recipe, which peaks by ~1500) — so max_iter is longer here. +# Env: LIBERO_ROOT, BASE_CHECKPOINT_PATH, WAN_VAE_PATH, IMAGINAIRE_OUTPUT_ROOT. +# See docs/action_policy_libero_sft.md. + +[job] +task = "vfm" +experiment = "action_policy_libero_all_nano" +project = "cosmos3_action_libero" +group = "action_sft" +name = "action_policy_libero_all_repro" +wandb_mode = "online" + +[model] +precision = "bfloat16" +max_num_tokens_after_packing = 74000 + +[model.parallelism] +data_parallel_shard_degree = 8 +data_parallel_replicate_degree = 2 # HSDP 2x8 = 16 ranks (2 nodes); minimum for gbs 2048 at grad_accum 1 + +[model.activation_checkpointing] +mode = "selective" +save_ops_regex = ["fmha"] + +[model.tokenizer] +vae_path = "${oc.env:WAN_VAE_PATH}" + +[optimizer] +lr = 5.0e-05 + +[scheduler] +cycle_lengths = [16000] +warm_up_steps = [500] + +[trainer] +max_iter = 5000 +logging_iter = 50 +grad_accum_iter = 1 # global batch = max_samples 128 x (shard 8 x replicate 2) x 1 = 2048 + +[checkpoint] +load_path = "${oc.env:BASE_CHECKPOINT_PATH}" +save_iter = 500 diff --git a/examples/toml/sft_config/action_policy_libero_repro.toml b/examples/toml/sft_config/action_policy_libero_repro.toml new file mode 100644 index 00000000..a0c49c7a --- /dev/null +++ b/examples/toml/sft_config/action_policy_libero_repro.toml @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# LIBERO-10 action-policy SFT run config for the `action_policy_libero_nano` +# experiment. Train on libero_10 alone (HSDP 2x8, global batch 2048). +# Env: LIBERO_ROOT, BASE_CHECKPOINT_PATH, WAN_VAE_PATH, IMAGINAIRE_OUTPUT_ROOT. +# See docs/action_policy_libero_sft.md. + +[job] +task = "vfm" +experiment = "action_policy_libero_nano" +project = "cosmos3_action_libero" +group = "action_sft" +name = "action_policy_libero_repro" +wandb_mode = "online" + +[model] +precision = "bfloat16" +max_num_tokens_after_packing = 74000 + +[model.parallelism] +data_parallel_shard_degree = 8 +data_parallel_replicate_degree = 2 # HSDP 2x8 = 16 ranks (2 nodes); minimum for gbs 2048 at grad_accum 1 + +[model.activation_checkpointing] +mode = "selective" +save_ops_regex = ["fmha"] + +[model.tokenizer] +vae_path = "${oc.env:WAN_VAE_PATH}" + +[optimizer] +lr = 5.0e-05 + +[scheduler] +cycle_lengths = [16000] +warm_up_steps = [500] + +[trainer] +max_iter = 2000 +logging_iter = 50 +grad_accum_iter = 1 # global batch = max_samples 128 x (shard 8 x replicate 2) x 1 = 2048 + +[checkpoint] +load_path = "${oc.env:BASE_CHECKPOINT_PATH}" +save_iter = 500 From 90290a26c1e5af6437c0939e63c598c6368690f1 Mon Sep 17 00:00:00 2001 From: Xiangyu Lu <169013972+xlu451@users.noreply.github.com> Date: Thu, 2 Jul 2026 01:41:27 -0700 Subject: [PATCH 05/26] Remove eval CLI references (#5) Co-authored-by: Xiangyu Lu Co-authored-by: lfengad --- .agents/skills/cosmos3-codebase-nav/SKILL.md | 1 - AGENTS.md | 2 +- docs/code_structure.md | 4 ++-- docs/setup.md | 1 - 4 files changed, 3 insertions(+), 5 deletions(-) diff --git a/.agents/skills/cosmos3-codebase-nav/SKILL.md b/.agents/skills/cosmos3-codebase-nav/SKILL.md index 2969c414..eaaf7c9d 100644 --- a/.agents/skills/cosmos3-codebase-nav/SKILL.md +++ b/.agents/skills/cosmos3-codebase-nav/SKILL.md @@ -95,7 +95,6 @@ Users can also supply a custom defaults file per-request via the `defaults_file` | ---------------------------------- | -------------------------------------------------------------------------------------------- | | Batch inference | `python -m cosmos_framework.scripts.inference` | | Training | `python -m cosmos_framework.scripts.train --sft-toml=examples/toml/sft_config/.toml` | -| Action evaluation | `python -m cosmos_framework.scripts.eval` | | Online serving (Ray) | `python -m cosmos_framework.inference.ray.serve` | | Submit to Ray server | `python -m cosmos_framework.inference.ray.submit` | | Gradio UI | `python -m cosmos_framework.inference.ray.gradio` | diff --git a/AGENTS.md b/AGENTS.md index f395478b..7af614bd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -7,7 +7,7 @@ Read this file first — it is the canonical map for navigating the Cosmos repos - **Training infrastructure** — top-level subpackages under `cosmos_framework/` (data, model, trainer, callbacks, checkpoint, …). - **Inference infrastructure** — `cosmos_framework/inference/` (Diffusers / Transformers / vLLM-friendly inference core, online serving via Ray + Gradio). - **Backend packages** — `packages/{diffusers,transformers,vllm}-cosmos3/` provide library-style shims that load Cosmos3 checkpoints into the respective ecosystems. -- **Entry-point scripts** — `cosmos_framework/scripts/` (`train.py`, `inference.py`, `eval.py`, `export_model.py`, …) invoked as `python -m cosmos_framework.scripts.`. Primary training entry point: `cosmos_framework.scripts.train` driven by a structured, pydantic-validated TOML interface (`--sft-toml=`); the schema lives at [`cosmos_framework/configs/toml_config/sft_config.py`](./cosmos_framework/configs/toml_config/sft_config.py) and the canonical recipe pattern is documented in [`examples/README.md`](./examples/README.md). +- **Entry-point scripts** — `cosmos_framework/scripts/` (`train.py`, `inference.py`, `export_model.py`, …) invoked as `python -m cosmos_framework.scripts.`. Primary training entry point: `cosmos_framework.scripts.train` driven by a structured, pydantic-validated TOML interface (`--sft-toml=`); the schema lives at [`cosmos_framework/configs/toml_config/sft_config.py`](./cosmos_framework/configs/toml_config/sft_config.py) and the canonical recipe pattern is documented in [`examples/README.md`](./examples/README.md). > All paths below are relative to the repository root (the directory containing `pyproject.toml`, the `cosmos_framework/` Python package, and `packages/`). diff --git a/docs/code_structure.md b/docs/code_structure.md index a4453933..3c7ecd39 100644 --- a/docs/code_structure.md +++ b/docs/code_structure.md @@ -37,7 +37,7 @@ Cosmos/ ├── cosmos_framework/ # Main package (training infra + inference subpackage) │ ├── inference/ # Inference subpackage (model, args, defaults, Ray serving, common helpers, SFT experiment configs) │ └── ... # Training-infra subpackages: data, model, trainer, callbacks, checkpoint, … -│ └── scripts/ # CLI entry-point scripts: train.py, _train.py, inference.py, eval.py, export_model.py, … +│ └── scripts/ # CLI entry-point scripts: train.py, _train.py, inference.py, export_model.py, … ├── packages/ # Backend shim packages: diffusers-cosmos3, transformers-cosmos3, vllm-cosmos3 ├── docs/ # User documentation (you are here) ├── docker/ # Dockerfiles for reproducible environments @@ -196,7 +196,7 @@ Add new worker types as sibling subpackages — each owns its own startup, messa - `tests/` — pytest tests, mirroring the `cosmos_framework/` package layout. - `examples/` — runnable end-to-end examples; see `examples/README.md`. - `docker/` — Dockerfiles and image build helpers; see `docker/README.md`. -- `cosmos_framework/scripts/` — CLI entry-point scripts (`train.py`, `inference.py`, `eval.py`, `export_model.py`, …); invoke as `python -m cosmos_framework.scripts.`. Primary training entry point: `cosmos_framework.scripts.train` driven by a structured, pydantic-validated TOML interface (`--sft-toml=`); recipe TOMLs live under [`examples/toml/sft_config/`](../examples/toml/sft_config/) and the schema is defined in [`cosmos_framework/configs/toml_config/sft_config.py`](../cosmos_framework/configs/toml_config/sft_config.py) — see [`examples/README.md`](../examples/README.md) and [`docs/training.md`](./training.md). +- `cosmos_framework/scripts/` — CLI entry-point scripts (`train.py`, `inference.py`, `export_model.py`, …); invoke as `python -m cosmos_framework.scripts.`. Primary training entry point: `cosmos_framework.scripts.train` driven by a structured, pydantic-validated TOML interface (`--sft-toml=`); recipe TOMLs live under [`examples/toml/sft_config/`](../examples/toml/sft_config/) and the schema is defined in [`cosmos_framework/configs/toml_config/sft_config.py`](../cosmos_framework/configs/toml_config/sft_config.py) — see [`examples/README.md`](../examples/README.md) and [`docs/training.md`](./training.md). - `packages/` — library-style backend shims: `packages/{diffusers,transformers,vllm}-cosmos3/`, each installable independently. ## Where to Add New Code diff --git a/docs/setup.md b/docs/setup.md index 9e89ff43..299476be 100644 --- a/docs/setup.md +++ b/docs/setup.md @@ -216,7 +216,6 @@ If there is no attention backend wheel for your torch/cuda versions, you can bui Optional package extras: - `train`: Training infrastructure (FSDP, parallelism, checkpointing, datasets) -- `eval`: Evaluation harnesses for trained checkpoints #### CUDA Variants From ea63366a71e9c5a1948e235facf32e2faf309a2b Mon Sep 17 00:00:00 2001 From: lfengad Date: Thu, 2 Jul 2026 19:28:29 +0800 Subject: [PATCH 06/26] Add multi-control transfer inference smoke test (#77) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What Adds `test_nano_inference_multi_control_transfer` to `tests/nano_inference_smoke_test.py`, giving the Cosmos3-Nano generator inference smoke test coverage for the **multi-control transfer** feature (two control hints blended by `multi_control_two_way_attention`). It mirrors the existing single-control transfer run (same `latency` preset, 4 ranks → cfgp=2/cp=2), but the generated spec sets **two** hints (`edge` + `blur`) with per-hint `weight` and **no** `control_path`, so both controls are computed on the fly from a single `vision_path` (the pinned public `robot_pouring.mp4` clip) and aggregated by the weighted N-pass multi-control attention. ## What it verifies - **Multi-control path executes**: 2 hints → `control_weights` set → the network routes to `multi_control_two_way_attention` (both controls computed on the fly). - **Process stays within expected invariants**: the framework's multi-control runtime asserts (weights length == #controls, packing consistency, batch_size==1) must hold or the run exits non-zero → the test fails. Plus the test's own asserts: exactly one output; both `edge` and `blur` active with a `weight` and no `control_path`; `vision_path` set; `control_guidance` and `guidance` > 1.0. - **Output is valid**: non-degenerate `vision.mp4` via the existing `_assert_video_has_content` (≥16 frames, finite pixels, pixel std > 3). Smoke-level (output validity + path executed under asserts), not numeric goldens — consistent with the rest of the module. ## Test Runs on the same 8-GPU gate as `test_nano_inference_omni`: ``` pytest -s tests/nano_inference_smoke_test.py --num-gpus=8 --levels=2 -o addopts= ``` Verified end to end on a 4-rank `latency` run (exit 0; log shows both `Computing edge control input on the fly` and `Computing blur control input on the fly`; all assertions pass). Collection confirmed both tests register (`MAX_GPUS=8`). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 (1M context) --- tests/nano_inference_smoke_test.py | 115 +++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/tests/nano_inference_smoke_test.py b/tests/nano_inference_smoke_test.py index 14defa15..098d8510 100644 --- a/tests/nano_inference_smoke_test.py +++ b/tests/nano_inference_smoke_test.py @@ -101,6 +101,54 @@ "edge": {"control_path": _TRANSFER_CONTROL_URL, "preset_edge_threshold": "medium"}, } +# Multi-control transfer (video2video, edge + blur) input, written to a temp file +# at run time. Mirrors the cookbook +# ``cookbooks/cosmos3/generator/transfer/specs/multi_control.json`` — two control +# hints (edge + blur) computed on the fly from a single source video (``vision_path``) +# and blended by ``multi_control_two_way_attention`` (N independent maskless SDPA +# passes, one per control, summed by the per-hint ``weight``) — but downscaled +# (480p / 10 steps / single 29-frame chunk) for a fast smoke run. The source clip is +# the exact one the cookbook uses (a robot arm pouring into a glass), pinned to a +# public raw URL; the prompt is a compact caption of it. Unlike ``_TRANSFER_SPEC`` +# (a single pre-computed ``control_path``), both controls here are derived on the +# fly, so this exercises the transfer control augmentor in addition to the weighted +# multi-control aggregation. ``guidance`` + ``control_guidance`` > 1.0 also keep the +# text-CFG and control-CFG branches active. +_MULTI_CONTROL_VISION_URL = ( + "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/" + "2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/vision/robot_pouring.mp4" +) +_MULTI_CONTROL_SPEC = { + "name": "transfer_multi_control", + "model_mode": "video2video", + "resolution": "480", + "aspect_ratio": "16,9", + "num_frames": 29, + "fps": 30, + "shift": 10.0, + "num_steps": 10, + "seed": 2026, + "num_video_frames_per_chunk": 29, + "max_frames": 29, + "num_conditional_frames": 1, + "num_first_chunk_conditional_frames": 0, + "share_vision_temporal_positions": True, + "guidance": 3.0, + "control_guidance": 1.5, + "vision_path": _MULTI_CONTROL_VISION_URL, + "prompt": ( + "A white robotic arm with black joints and cables carefully pours a clear liquid from a " + "small light-green pitcher into a glass on a white tabletop, in a clean, brightly lit " + "modern indoor setting." + ), + "negative_prompt": "blurry, distorted, deformed, low quality, flickering, artifacts", + # Two hints, no control_path -> both derived on the fly from vision_path; the + # per-hint weights drive the weighted multi-control attention aggregation. + "edge": {"weight": 0.5, "preset_edge_threshold": "medium"}, + "blur": {"weight": 0.5, "preset_blur_strength": "medium"}, + "emphasize_control_in_prompt": False, +} + # Audio sanity thresholds for the muxed sound track. _RMS_SILENCE_FLOOR = 1e-4 # below this the track is effectively silence _PEAK_SANITY_CEIL = 1.5 # decoded float audio should sit within ~[-1, 1] @@ -355,3 +403,70 @@ def test_nano_inference_omni(tmp_path: Path) -> None: transfer_video = so.parent / "vision.mp4" assert transfer_video.is_file(), f"transfer run produced no vision.mp4 ({so})" _assert_video_has_content(transfer_video) + + @pytest.mark.level(2) + @pytest.mark.gpus(8) + def test_nano_inference_multi_control_transfer(tmp_path: Path) -> None: + """Multi-control transfer: edge + blur derived on the fly from ONE source + video, blended by ``multi_control_two_way_attention``. + + Mirrors ``test_nano_inference_omni``'s single-control transfer run (same + ``latency`` preset, 4 ranks -> cfgp=2, cp=2 -- the cookbook Cosmos3-Super + transfer layout), but the generated spec sets TWO control hints (edge + + blur) each with a per-hint ``weight`` and no ``control_path``, so both + controls are computed on the fly from ``vision_path`` and aggregated by the + weighted multi-control attention path (``multi_control_two_way_attention``: + N maskless SDPA passes summed by weight). A non-degenerate clip confirms + that path ran end to end -- a broken multi-control route would raise + mid-sampling, and a numerically broken one would collapse the output + (caught by ``_assert_video_has_content``). The on-the-fly derivation also + exercises the transfer control augmentor (opencv), unlike the single-control + run above which loads a pre-computed control_path.""" + spec_file = tmp_path / "transfer_multi_control.json" + spec_file.write_text(json.dumps(_MULTI_CONTROL_SPEC)) + out_dir = tmp_path / "out_multi_control" + cmd = [ + "torchrun", + "--nproc_per_node=4", + f"--master_port={_free_port()}", + "-m", + "cosmos_framework.scripts.inference", + "--parallelism-preset=latency", + "-i", + str(spec_file), + "-o", + str(out_dir), + "--checkpoint-path", + "Cosmos3-Nano", + "--seed=0", + ] + _run(cmd, tmp_path / "inference_multi_control.log") + + results = sorted(out_dir.rglob("sample_outputs.json")) + assert len(results) == 1, ( + f"expected 1 multi-control sample_outputs.json, found {[str(p) for p in results]}" + ) + so = results[0] + args = json.loads(so.read_text()).get("args", {}) + # Multi-control-specific: BOTH edge and blur hints are active (2 controls -> + # the weighted multi_control_two_way_attention path), each carries a weight, + # and neither has a control_path (both derived on the fly from vision_path). + edge = args.get("edge") or {} + blur = args.get("blur") or {} + assert edge and blur, f"expected both edge and blur hints active ({so}); edge={edge} blur={blur}" + assert edge.get("weight") is not None and blur.get("weight") is not None, ( + f"expected a per-hint weight on both controls ({so}); edge={edge} blur={blur}" + ) + assert not edge.get("control_path") and not blur.get("control_path"), ( + f"expected on-the-fly controls (no control_path) ({so}); edge={edge} blur={blur}" + ) + assert args.get("vision_path"), f"multi-control run missing vision_path ({so})" + assert args.get("control_guidance", 1.0) > 1.0, ( + f"expected control-CFG (control_guidance > 1.0), got {args.get('control_guidance')} ({so})" + ) + assert (args.get("guidance") or 1.0) > 1.0, ( + f"expected text-CFG (guidance > 1.0), got {args.get('guidance')} ({so})" + ) + video = so.parent / "vision.mp4" + assert video.is_file(), f"multi-control run produced no vision.mp4 ({so})" + _assert_video_has_content(video) From fbb5c9bf4b1298a09cabbe8d60389ef06ab60821 Mon Sep 17 00:00:00 2001 From: "Hongyu, Chiu" Date: Thu, 2 Jul 2026 22:15:39 +0800 Subject: [PATCH 07/26] Lazily import `lerobot` in datasets. (#76) `lerobot` is a heavy package that pins a lot of python packages. We can lazily import it to support a lightweight env setup. Signed-off-by: Hong-Yu Chiu Co-authored-by: lfengad --- .../action/datasets/agibotworld_beta_lerobot_dataset.py | 4 +++- .../vfm/action/datasets/bridge_orig_lerobot_dataset.py | 4 +++- .../data/vfm/action/datasets/droid_lerobot_dataset.py | 4 +++- .../data/vfm/action/datasets/fractal_lerobot_dataset.py | 4 +++- .../data/vfm/action/datasets/libero_lerobot_dataset.py | 4 +++- .../data/vfm/action/datasets/robomind_franka_dataset.py | 7 ++++++- .../data/vfm/action/datasets/robomind_ur_dataset.py | 4 +++- .../data/vfm/action/datasets/umi_lerobot_dataset.py | 4 +++- 8 files changed, 27 insertions(+), 8 deletions(-) diff --git a/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py index ce22773c..81a9a364 100644 --- a/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py @@ -11,7 +11,6 @@ import numpy as np import torch import torch.nn.functional as F -from lerobot.datasets.video_utils import decode_video_frames from cosmos_framework.data.vfm.action.agibot_fk import ( AGIBOT_WORLD_GRIPPER_TO_OPENCV_BY_WRIST, @@ -232,6 +231,9 @@ def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, return self._compose_multi_view(top, left, right) def _load_video_key(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]], key: str) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + timestamps = [float(row["timestamp"]) for row in observation_rows] return decode_video_frames( self._video_path(episode, key), diff --git a/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py index 25e167b0..e6ea9d2b 100644 --- a/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py @@ -11,7 +11,6 @@ import numpy as np import torch -from lerobot.datasets.video_utils import decode_video_frames from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset @@ -126,6 +125,9 @@ def __getitem__(self, idx: int) -> dict[str, Any]: ) def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + timestamps = [float(row["timestamp"]) for row in observation_rows] return decode_video_frames( self._video_path(episode, _IMAGE_FEATURE), diff --git a/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py index f3faa91f..5d31a999 100644 --- a/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py @@ -15,7 +15,6 @@ import torch import torch.nn.functional as F import torchvision.transforms as T -from lerobot.datasets.video_utils import decode_video_frames from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Joint, Pos, Rot, build_action_spec from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset @@ -290,6 +289,9 @@ def _load_concat_video( episode: dict[str, Any], observation_rows: list[dict[str, Any]], ) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + timestamps = [float(row["timestamp"]) for row in observation_rows] frames_by_view = { name: decode_video_frames( diff --git a/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py index 10628db1..b9dd23d7 100644 --- a/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py @@ -18,7 +18,6 @@ import numpy as np import torch -from lerobot.datasets.video_utils import decode_video_frames from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset @@ -155,6 +154,9 @@ def __getitem__(self, idx: int) -> dict[str, Any]: ) def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + timestamps = [float(row["timestamp"]) for row in observation_rows] return decode_video_frames( self._video_path(episode, _IMAGE_FEATURE), diff --git a/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py index fb8d5e8c..fe31f9d5 100644 --- a/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py @@ -31,7 +31,6 @@ import pyarrow.parquet as pq import torch import torch.nn.functional as F -from lerobot.datasets.video_utils import decode_video_frames from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset @@ -307,6 +306,9 @@ def _build_frame_wise_action(self, raw: np.ndarray) -> torch.Tensor: return torch.cat([translation, rotation, gripper], dim=-1) # [chunk, action_dim] def _load_video(self, episode: dict[str, Any], timestamps: list[float]) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + frames_by_view = {} for key in self._video_keys: from_ts = float(episode.get(f"videos/{key}/from_timestamp", 0.0)) diff --git a/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py b/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py index 53b0a2c6..7bda3053 100644 --- a/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py @@ -12,7 +12,6 @@ import numpy as np import torch import torch.nn.functional as F -from lerobot.datasets.video_utils import decode_video_frames from cosmos_framework.data.vfm.action.action_normalization import load_action_stats from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec @@ -181,6 +180,9 @@ def _load_single_video( episode: dict[str, Any], observation_rows: list[dict[str, Any]], ) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + video_key = self._image_features["top"] timestamps = [float(row["timestamp"]) for row in observation_rows] return decode_video_frames( @@ -194,6 +196,9 @@ def _load_concat_video( episode: dict[str, Any], observation_rows: list[dict[str, Any]], ) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + timestamps = [float(row["timestamp"]) for row in observation_rows] frames_by_view = { name: decode_video_frames( diff --git a/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py b/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py index 317525f7..ab7bb433 100644 --- a/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py @@ -11,7 +11,6 @@ import numpy as np import torch -from lerobot.datasets.video_utils import decode_video_frames from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset @@ -168,6 +167,9 @@ def __getitem__(self, idx: int) -> dict[str, Any]: ) def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + timestamps = [float(row["timestamp"]) for row in observation_rows] return decode_video_frames( self._video_path(episode, _IMAGE_FEATURE), diff --git a/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py index 59395961..61a31497 100644 --- a/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py +++ b/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py @@ -11,7 +11,6 @@ import numpy as np import torch -from lerobot.datasets.video_utils import decode_video_frames from cosmos_framework.data.vfm.action.action_normalization import load_action_stats from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec @@ -140,6 +139,9 @@ def __getitem__(self, idx: int) -> dict[str, Any]: ) def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor: + # lerobot is a heavy, optional ("train" extra) dependency; import lazily. + from lerobot.datasets.video_utils import decode_video_frames + timestamps = [float(row["timestamp"]) for row in observation_rows] return decode_video_frames( self._video_path(episode, self._image_key), From df864d8c4c8b13a8024dfaf66a97ff8044870294 Mon Sep 17 00:00:00 2001 From: yy-code-nv Date: Fri, 3 Jul 2026 00:06:09 +0800 Subject: [PATCH 08/26] =?UTF-8?q?Release:=20vfm=20=E2=86=92=20generator,?= =?UTF-8?q?=20vlm=20=E2=86=92=20reasoner=20dest=20rename=20(moves=20+=20im?= =?UTF-8?q?port=20=E2=80=A6=20(#78)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit …paths only) Follow-up to the initial vlm→reasoner sync (PR #70) — this PR renames the top-level CF dest folders and rewrites every import that references them. Contents are strictly file relocations + path-string rewrites; no code-behavior edits. Dest-folder rename: * cosmos_framework/data/vfm/ → cosmos_framework/data/generator/ * cosmos_framework/model/vfm/ → cosmos_framework/model/generator/ * cosmos_framework/utils/vfm/ → cosmos_framework/utils/generator/ * cosmos_framework/data/vlm/ → cosmos_framework/data/reasoner/ * cosmos_framework/utils/vlm/ → cosmos_framework/utils/reasoner/ * cosmos_framework/configs/base/vlm/{defaults,experiment}/ → cosmos_framework/configs/base/reasoner/{defaults,experiment}/ Import-path rewrites — every reference under ``cosmos_framework.{model,data,utils}.{vfm,vlm}.*`` and ``cosmos_framework.configs.base.vlm.*`` is retargeted to the new module names, in: * Release-managed files produced by cosmos-framework-release. * CF-owned production code (inference/, callbacks/, configs/, model/, utils/, data/generator/*/, data/reasoner/*/, tools/, scripts/). * Runtime rewriters: cosmos_framework/inference/common/config.py::CONFIG_REPLACEMENTS_INVERSE and cosmos_framework/inference/common/public_model_config.py (``_canonicalize_module_path``, ``_replace_vfm_module_prefix``, ``_replace_vfm_file_prefix``, ``_module_exists`` probes, {package}-templated dest strings) — old checkpoint JSONs' canonical ``cosmos3._src.vfm.*`` / ``projects.cosmos3.vfm.*`` target strings now map to the new module locations at load time. Verified: * Fresh AST-walker sweep across ``cosmos_framework/`` → **0 dangling ``cosmos_framework.*`` imports**. * Every ``M`` (modified-content) file's diff is verified path-only (each +/- line contains a ``vfm|vlm|generator|reasoner`` substring or is whitespace). * Every ``R`` is a git-detected rename. * 4-GPU GB200 regression: ``test_launch_regression[vision_sft_nano]`` passes and matches gb200 goldens exactly. * Unit suite (venv Python 3.13): 327 passed / 1 failed (the failure is the ``convert_model_to_dcp`` script test, which needs ``HF_TOKEN`` for the private ``nvidia/Cosmos3-Experimental`` snapshot — not fixable in-repo). --- .agents/skills/cosmos3-codebase-nav/SKILL.md | 4 +- .claude/skills/cosmos3-codebase-nav/SKILL.md | 4 +- .github/workflows/gpu-tests.yml | 4 +- NOTICE | 14 ++-- .../callbacks/compile_tokenizer.py | 4 +- .../callbacks/every_n_draw_sample.py | 2 +- cosmos_framework/callbacks/expert_heatmap.py | 2 +- cosmos_framework/callbacks/hf_export.py | 2 +- .../callbacks/moe_specialization_callback.py | 2 +- .../callbacks/moe_stability_callback.py | 2 +- cosmos_framework/callbacks/norm_monitor.py | 2 +- .../callbacks/sequence_packing_padding.py | 2 +- cosmos_framework/callbacks/tokens_per_sec.py | 2 +- cosmos_framework/callbacks/training_stats.py | 2 +- cosmos_framework/checkpoint/dcp.py | 2 +- .../base/defaults/activation_checkpointing.py | 4 +- .../configs/base/defaults/model.py | 2 +- .../base/defaults/open_source_dataloader.py | 6 +- .../configs/base/defaults/optimizer.py | 2 +- .../configs/base/defaults/reasoner.py | 64 ++++++++--------- .../configs/base/defaults/tokenizer.py | 14 ++-- .../action_policy_droid_nano.py | 4 +- .../action_policy_libero_all_nano.py | 4 +- .../action_policy_libero_nano.py | 4 +- .../sft/models/nano_model_config.py | 4 +- .../sft/models/super_model_config.py | 4 +- .../base/experiment/sft/vision_sft_nano.py | 8 +-- .../base/experiment/sft/vision_sft_super.py | 6 +- .../configs/base/reasoner/defaults/model.py | 2 +- .../reasoner/experiment/dataflow_roles.py | 4 +- .../base/reasoner/experiment/llava_ov_vlm.py | 8 +-- .../experiment/videophy2_dataflow_roles.py | 4 +- .../reasoner/experiment/videophy2_sft_nano.py | 10 +-- .../configs/toml_config/sft_config.py | 4 +- .../data/{vfm => generator}/__init__.py | 0 .../{vfm => generator}/action/__init__.py | 0 .../action/action_normalization.py | 0 .../action/action_normalization_test.py | 2 +- .../action/action_processing.py | 0 .../{vfm => generator}/action/action_spec.py | 2 +- .../{vfm => generator}/action/agibot_fk.py | 4 +- .../{vfm => generator}/action/agibot_spec.py | 0 .../generator/action/datasets/__init__.py | 31 ++++++++ .../action/datasets/action_sft_dataset.py | 6 +- .../agibotworld_beta_lerobot_dataset.py | 8 +-- .../action/datasets/base_dataset.py | 8 +-- .../datasets/bridge_orig_lerobot_dataset.py | 6 +- .../action/datasets/droid_lerobot_dataset.py | 6 +- .../datasets/fractal_lerobot_dataset.py | 6 +- .../action/datasets/libero_lerobot_dataset.py | 8 +-- .../datasets/robomind_franka_dataset.py | 8 +-- .../action/datasets/robomind_ur_dataset.py | 6 +- .../action/datasets/umi_lerobot_dataset.py | 8 +-- .../{vfm => generator}/action/domain_utils.py | 0 .../action/json_formatter.py | 4 +- .../action/libero_pose_utils.py | 2 +- .../agibotworld_beta_lerobot_stats.json | 0 .../bridge_orig_lerobot_stats.json | 0 .../normalizer_stats/droid_lerobot_stats.json | 0 .../fractal_lerobot_stats.json | 0 ...bero_native_frame_wise_relative_rot6d.json | 0 .../robomind_franka_dual_stats.json | 0 .../robomind_franka_stats.json | 0 .../normalizer_stats/robomind_ur_stats.json | 0 .../normalizer_stats/umi_lerobot_stats.json | 0 .../{vfm => generator}/action/pose_utils.py | 2 +- .../action/pose_utils_test.py | 2 +- .../G1_omnipicker_calibrated.urdf | 0 .../action/robot_assets/ur5e_robotiq_2f85.xml | 0 .../{vfm => generator}/action/transforms.py | 20 +++--- .../action/transforms_test.py | 8 +-- .../action/viewpoint_utils.py | 0 .../{vfm => generator}/augmentor_provider.py | 30 ++++---- .../{vfm => generator}/augmentors/__init__.py | 0 .../augmentors/append_fps_frames_for_image.py | 0 .../augmentors/audio_caption.py | 0 .../augmentors/audio_parsing.py | 0 .../augmentors/caption_filter.py | 0 .../{vfm => generator}/augmentors/cropping.py | 0 .../duration_fps_text_timestamps.py | 0 .../augmentors/idle_frames_text_info.py | 2 +- .../augmentors/image_editing_transform.py | 2 +- .../image_editing_transform_test.py | 2 +- .../augmentors/image_resolution_filter.py | 2 +- .../augmentors/interleaved_image_transform.py | 0 .../augmentors/interleaved_video_parsing.py | 2 +- .../augmentors/merge_datadict.py | 0 .../augmentors/pkl_to_media.py | 2 +- .../augmentors/reasoner/__init__.py | 0 .../augmentors/reasoner/bytes_to_media.py | 6 +- .../augmentors/reasoner/filter_output_key.py | 0 .../augmentors/reasoner/filter_seq_length.py | 2 +- .../reasoner/floating_number_format.py | 0 .../reasoner/format_describe_anything.py | 2 +- .../reasoner/nvlm_data_to_conversation.py | 0 .../augmentors/reasoner/prompt_format.py | 0 .../reasoner/shuffle_text_media_order.py | 0 .../augmentors/reasoner/timestamp.py | 0 .../timestamp_with_subject_tracking.py | 2 +- .../timestamp_without_augment_message.py | 2 +- .../reasoner/timestamp_without_end_time.py | 2 +- .../augmentors/reasoner/tokenize_data.py | 6 +- .../reasoner/user_prompt_caption_general.json | 0 .../augmentors/reasoner/user_prompt_ocr.json | 0 .../augmentors/resolution_text_info.py | 0 .../augmentors/sequence_plan.py | 4 +- .../augmentors/sound_sequence_plan.py | 4 +- .../augmentors/text_tokenizer.py | 0 .../augmentors/text_transforms_for_image.py | 0 .../augmentors/text_transforms_for_video.py | 0 .../transfer_control_input/__init__.py | 2 +- .../augmentors/transfer_control_input/blur.py | 2 +- .../transfer_control_input/control_input.py | 4 +- .../transfer_control_input/fast_blur.py | 0 .../augmentors/transfer_control_input/seg.py | 0 .../augmentors/transfer_control_transform.py | 6 +- .../augmentors/video_parsing.py | 4 +- .../{vfm => generator}/dataflow/__init__.py | 12 ++-- .../data/{vfm => generator}/dataflow/base.py | 0 .../{vfm => generator}/dataflow/batchers.py | 2 +- .../{vfm => generator}/dataflow/collators.py | 2 +- .../dataflow/distributors.py | 2 +- .../dataflow/golden_vfm_test.py | 4 +- .../{vfm => generator}/dataflow/loader.py | 8 +-- .../{vfm => generator}/dataflow/processors.py | 2 +- .../dataflow/resume_test.py | 2 +- .../{vfm => generator}/joint_dataloader.py | 2 +- .../local_datasets/__init__.py | 0 .../local_datasets/helper.py | 0 .../local_datasets/sft_dataset.py | 10 +-- .../sft_dataset_caption_test.py | 2 +- .../packing_iterable_dataset.py | 2 +- .../{vfm => generator}/processors/__init__.py | 12 ++-- .../{vfm => generator}/processors/base.py | 4 +- .../processors/nemotron3densevl_processor.py | 2 +- .../processors/nemotronvl_processor.py | 4 +- .../processors/qwen3vl_processor.py | 4 +- .../reasoner/video_decoder_qwen.py | 2 +- .../{vfm => generator}/recipe_database_api.py | 0 .../sequence_packing/__init__.py | 4 +- .../sequence_packing/modalities.py | 4 +- .../sequence_packing/mrope.py | 0 .../sequence_packing/natten.py | 0 .../sequence_packing/packers.py | 8 +-- .../sequence_packing/runtime.py | 0 .../sequence_packing/temporal_causal.py | 4 +- .../sequence_packing/types.py | 2 +- .../{vfm => generator}/sound_data_utils.py | 6 +- .../data/{vfm => generator}/utils.py | 0 .../data/{vfm => generator}/watchdog.py | 0 .../data/{vlm => reasoner}/__init__.py | 0 .../data_sources_videophy2/__init__.py | 0 .../data_weight/__init__.py | 0 .../data_weight/default.py | 2 +- .../data_sources_videophy2/videophy2.py | 2 +- .../{vlm => reasoner}/local_dataset_utils.py | 0 .../{vlm => reasoner}/local_sft_dataset.py | 0 .../{vlm => reasoner}/processors/__init__.py | 8 +-- .../processors/nemotron3densevl_processor.py | 2 +- .../processors/nemotronvl_processor.py | 4 +- .../processors/qwen3vl_processor.py | 4 +- .../data/vfm/action/datasets/__init__.py | 31 -------- cosmos_framework/inference/action.py | 10 +-- cosmos_framework/inference/args.py | 2 +- cosmos_framework/inference/args_test.py | 2 +- cosmos_framework/inference/common/args.py | 2 +- .../inference/common/checkpoints.py | 4 +- cosmos_framework/inference/common/config.py | 48 ++++++------- cosmos_framework/inference/common/init.py | 2 +- .../inference/common/public_model_config.py | 70 +++++++++---------- .../common/public_model_config_test.py | 10 +-- cosmos_framework/inference/dataset_samples.py | 2 +- cosmos_framework/inference/inference.py | 8 +-- cosmos_framework/inference/model.py | 6 +- cosmos_framework/inference/sound.py | 4 +- cosmos_framework/inference/sound_test.py | 2 +- cosmos_framework/inference/transfer.py | 10 +-- cosmos_framework/inference/vision.py | 4 +- .../model/{vfm => generator}/__init__.py | 0 .../{vfm => generator}/algorithm/__init__.py | 0 .../algorithm/loss/__init__.py | 8 +-- .../algorithm/loss/cross_entropy.py | 2 +- .../algorithm/loss/flow_matching.py | 2 +- .../algorithm/loss/load_balancing.py | 2 +- .../algorithm/loss/time_weight.py | 0 .../{vfm => generator}/diffusion/__init__.py | 0 .../diffusion/rectified_flow.py | 2 +- .../diffusion/samplers/__init__.py | 0 .../diffusion/samplers/edm.py | 0 .../diffusion/samplers/fixed_step.py | 2 +- .../diffusion/samplers/fixed_step_test.py | 2 +- .../diffusion/samplers/fm_solvers_unipc.py | 0 .../diffusion/samplers/unipc.py | 4 +- .../diffusion/samplers/utils.py | 0 .../diffusion/samplers/utils_test.py | 2 +- .../model/{vfm => generator}/hf_model.py | 12 ++-- .../model/{vfm => generator}/mot/__init__.py | 0 .../model/{vfm => generator}/mot/attention.py | 6 +- .../{vfm => generator}/mot/attention_test.py | 6 +- .../{vfm => generator}/mot/cfgp_ar_test.py | 6 +- .../mot/context_parallel_test.py | 18 ++--- .../mot/context_parallel_utils.py | 8 +-- .../mot/cosmos3_vfm_network.py | 14 ++-- .../mot/domain_aware_linear.py | 0 .../mot/dot_product_attention.py | 0 .../{vfm => generator}/mot/modeling_utils.py | 2 +- .../mot/parallelize_unified_mot.py | 10 +-- .../mot/parallelize_vfm_network.py | 4 +- .../mot/qwen3_vl_unified_mot.py | 4 +- .../mot/und_k_norm_example_test.py | 12 ++-- .../{vfm => generator}/mot/unified_mot.py | 22 +++--- .../{vfm => generator}/omni_mot_model.py | 70 +++++++++---------- .../{vfm => generator}/parallelize_vlm.py | 8 +-- .../{vfm => generator}/reasoner/__init__.py | 0 .../reasoner/nemotron_3_dense_vl/__init__.py | 0 .../configs/Nemotron-2B-Dense-VL.json | 0 .../configuration_nemotron_3_dense_vl.py | 0 .../nemotron_3_dense_vl.py | 2 +- .../reasoner/qwen3_vl/__init__.py | 0 .../configs/Qwen3-VL-2B-Instruct.json | 0 .../configs/Qwen3-VL-32B-Instruct.json | 0 .../configs/Qwen3-VL-4B-Instruct.json | 0 .../configs/Qwen3-VL-8B-Instruct.json | 0 .../reasoner/qwen3_vl/configs/__init__.py | 0 .../qwen3_vl/configuration_qwen3_vl.py | 0 .../reasoner/qwen3_vl/qwen3_vl.py | 8 +-- .../reasoner/qwen3_vl/utils.py | 0 .../qwen3_vl/video_processing_qwen3_vl.py | 0 .../reasoner/qwen3_vl_moe/__init__.py | 0 .../configs/Qwen3-VL-235B-A22B-Instruct.json | 0 .../configs/Qwen3-VL-30B-A3B-Instruct.json | 0 .../reasoner/qwen3_vl_moe/configs/__init__.py | 0 .../configuration_qwen3_vl_moe.py | 0 .../reasoner/qwen3_vl_moe/moe.py | 4 +- .../reasoner/qwen3_vl_moe/moe_bench.py | 20 +++--- .../reasoner/qwen3_vl_moe/moe_kernels.py | 0 .../reasoner/qwen3_vl_moe/moe_test.py | 4 +- .../reasoner/qwen3_vl_moe/qwen3_vl_moe.py | 10 +-- .../{vfm => generator}/tokenizers/__init__.py | 0 .../tokenizers/audio/__init__.py | 0 .../tokenizers/audio/avae.py | 6 +- .../tokenizers/audio/avae_utils/__init__.py | 4 +- .../audio/avae_utils/activations.py | 0 .../avae_utils/alias_free_torch/__init__.py | 0 .../audio/avae_utils/alias_free_torch/act.py | 0 .../avae_utils/alias_free_torch/filter.py | 0 .../avae_utils/alias_free_torch/resample.py | 0 .../audio/avae_utils/bottlenecks.py | 0 .../tokenizers/audio/avae_utils/env.py | 0 .../tokenizers/audio/avae_utils/models.py | 0 .../tokenizers/audio/avae_utils/modules.py | 0 .../audio/avae_utils/modules_encodec.py | 0 .../tokenizers/dc_ae/__init__.py | 2 +- .../tokenizers/dc_ae/convert_pt_to_distcp.py | 2 +- .../tokenizers/dc_ae/dc_ae_4x32x32.py | 6 +- .../tokenizers/dc_ae/dc_ae_v.py | 2 +- .../tokenizers/dc_ae/dc_ae_v_ops.py | 2 +- .../dc_ae/dc_ae_v_triton_rms_norm.py | 0 .../tokenizers/flux_vae_8x8.py | 2 +- .../tokenizers/interface.py | 0 .../tokenizers/stable_diffusion_vae_8x8.py | 2 +- .../tokenizers/tokenization_qwen2.py | 0 .../tokenizers/uniae/__init__.py | 0 .../tokenizers/uniae/frame_math.py | 2 +- .../tokenizers/uniae/frame_math_test.py | 2 +- .../tokenizers/uniae/noncausal_4x16x16.py | 8 +-- .../tokenizers/wan2pt1_vae_4x8x8.py | 2 +- .../tokenizers/wan2pt2_vae_4x16x16.py | 6 +- .../{vfm => generator}/upsampler/__init__.py | 0 .../{vfm => generator}/upsampler/prompts.py | 4 +- .../{vfm => generator}/utils/__init__.py | 0 .../utils/data_and_condition.py | 0 .../model/{vfm => generator}/utils/memory.py | 0 .../utils/safetensors_loader.py | 6 +- .../utils/safetensors_loader_test.py | 8 +-- .../{vfm => generator}/utils/taylorseer.py | 0 .../model/{vfm => generator}/vlm_model.py | 16 ++--- .../scripts/_convert_model_to_diffusers.py | 2 +- cosmos_framework/scripts/_train.py | 2 +- .../scripts/action_policy_server_libero.py | 10 +-- .../scripts/action_policy_server_robolab.py | 8 +-- cosmos_framework/scripts/export_model.py | 2 +- cosmos_framework/scripts/train.py | 2 +- cosmos_framework/scripts/upsample_prompts.py | 2 +- .../simulation/libero/closed_loop_eval.py | 6 +- .../utils/{vfm => generator}/__init__.py | 0 .../utils/{vfm => generator}/data_utils.py | 0 .../{vfm => generator}/dtensor_helper.py | 0 .../utils/{vfm => generator}/flash_attn.py | 0 .../utils/{vfm => generator}/fused_adam.py | 0 .../{vfm => generator}/hf_attention_cosmos.py | 0 .../utils/{vfm => generator}/lora.py | 0 .../utils/{vfm => generator}/model_loader.py | 11 +-- .../{vfm => generator}/model_weights_stats.py | 0 .../utils/{vfm => generator}/monkey_patch.py | 0 .../utils/{vfm => generator}/optimizer.py | 4 +- .../utils/{vfm => generator}/parallelism.py | 4 +- .../utils/{vfm => generator}/rand_state.py | 0 .../{vfm => generator}/reasoner/__init__.py | 0 .../{vfm => generator}/reasoner/constant.py | 0 .../reasoner/create_position_ids.py | 0 .../reasoner/flop_calculator.py | 0 .../reasoner/pretrained_models_downloader.py | 2 +- .../{vfm => generator}/video_preprocess.py | 0 .../utils/{vlm => reasoner}/__init__.py | 0 .../compute_flops_qwen3vl.py | 0 .../configs_defaults/__init__.py | 0 .../configs_defaults/checkpointer.py | 2 +- .../utils/{vlm => reasoner}/constant.py | 0 .../{vlm => reasoner}/create_position_ids.py | 0 .../{vlm => reasoner}/dcp_checkpointer.py | 6 +- .../utils/{vlm => reasoner}/distributed.py | 0 .../{vlm => reasoner}/flop_calculator.py | 2 +- .../utils/{vlm => reasoner}/fused_adam.py | 2 +- .../utils/{vlm => reasoner}/model_wrapper.py | 0 .../utils/{vlm => reasoner}/optimizer.py | 2 +- .../utils/{vlm => reasoner}/planner.py | 0 .../pretrained_models_downloader.py | 4 +- .../{vlm => reasoner}/video_preprocess.py | 0 docs/action_policy_libero_sft.md | 22 +++--- docs/custom_dataset.md | 18 ++--- docs/faq.md | 2 +- docs/sft_config.md | 16 ++--- examples/integration/README.md | 14 ++-- examples/integration/net_level.py | 24 +++---- .../integration/trainer_level_inference.py | 16 ++--- .../integration/trainer_level_training.py | 24 +++---- examples/launch_sft_llava_ov.sh | 2 +- examples/launch_sft_videophy2_nano.sh | 2 +- examples/toml/sft_config/llava_ov.toml | 4 +- .../llava_ov_mapstyle_dataloader.toml | 2 +- .../toml/sft_config/videophy2_sft_nano.toml | 2 +- tests/_reasoner_logits_probe.py | 2 +- tests/_stage_h100_inputs.sh | 4 +- 334 files changed, 684 insertions(+), 691 deletions(-) rename cosmos_framework/data/{vfm => generator}/__init__.py (100%) rename cosmos_framework/data/{vfm => generator}/action/__init__.py (100%) rename cosmos_framework/data/{vfm => generator}/action/action_normalization.py (100%) rename cosmos_framework/data/{vfm => generator}/action/action_normalization_test.py (98%) rename cosmos_framework/data/{vfm => generator}/action/action_processing.py (100%) rename cosmos_framework/data/{vfm => generator}/action/action_spec.py (99%) rename cosmos_framework/data/{vfm => generator}/action/agibot_fk.py (99%) rename cosmos_framework/data/{vfm => generator}/action/agibot_spec.py (100%) create mode 100644 cosmos_framework/data/generator/action/datasets/__init__.py rename cosmos_framework/data/{vfm => generator}/action/datasets/action_sft_dataset.py (96%) rename cosmos_framework/data/{vfm => generator}/action/datasets/agibotworld_beta_lerobot_dataset.py (97%) rename cosmos_framework/data/{vfm => generator}/action/datasets/base_dataset.py (95%) rename cosmos_framework/data/{vfm => generator}/action/datasets/bridge_orig_lerobot_dataset.py (95%) rename cosmos_framework/data/{vfm => generator}/action/datasets/droid_lerobot_dataset.py (98%) rename cosmos_framework/data/{vfm => generator}/action/datasets/fractal_lerobot_dataset.py (96%) rename cosmos_framework/data/{vfm => generator}/action/datasets/libero_lerobot_dataset.py (97%) rename cosmos_framework/data/{vfm => generator}/action/datasets/robomind_franka_dataset.py (96%) rename cosmos_framework/data/{vfm => generator}/action/datasets/robomind_ur_dataset.py (96%) rename cosmos_framework/data/{vfm => generator}/action/datasets/umi_lerobot_dataset.py (95%) rename cosmos_framework/data/{vfm => generator}/action/domain_utils.py (100%) rename cosmos_framework/data/{vfm => generator}/action/json_formatter.py (98%) rename cosmos_framework/data/{vfm => generator}/action/libero_pose_utils.py (97%) rename cosmos_framework/data/{vfm => generator}/action/normalizer_stats/agibotworld_beta_lerobot_stats.json (100%) rename cosmos_framework/data/{vfm => generator}/action/normalizer_stats/bridge_orig_lerobot_stats.json (100%) rename cosmos_framework/data/{vfm => generator}/action/normalizer_stats/droid_lerobot_stats.json (100%) rename cosmos_framework/data/{vfm => generator}/action/normalizer_stats/fractal_lerobot_stats.json (100%) rename cosmos_framework/data/{vfm => generator}/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json (100%) rename cosmos_framework/data/{vfm => generator}/action/normalizer_stats/robomind_franka_dual_stats.json (100%) rename cosmos_framework/data/{vfm => generator}/action/normalizer_stats/robomind_franka_stats.json (100%) rename cosmos_framework/data/{vfm => generator}/action/normalizer_stats/robomind_ur_stats.json (100%) rename cosmos_framework/data/{vfm => generator}/action/normalizer_stats/umi_lerobot_stats.json (100%) rename cosmos_framework/data/{vfm => generator}/action/pose_utils.py (99%) rename cosmos_framework/data/{vfm => generator}/action/pose_utils_test.py (99%) rename cosmos_framework/data/{vfm => generator}/action/robot_assets/G1_omnipicker_calibrated.urdf (100%) rename cosmos_framework/data/{vfm => generator}/action/robot_assets/ur5e_robotiq_2f85.xml (100%) rename cosmos_framework/data/{vfm => generator}/action/transforms.py (97%) rename cosmos_framework/data/{vfm => generator}/action/transforms_test.py (96%) rename cosmos_framework/data/{vfm => generator}/action/viewpoint_utils.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentor_provider.py (97%) rename cosmos_framework/data/{vfm => generator}/augmentors/__init__.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/append_fps_frames_for_image.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/audio_caption.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/audio_parsing.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/caption_filter.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/cropping.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/duration_fps_text_timestamps.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/idle_frames_text_info.py (98%) rename cosmos_framework/data/{vfm => generator}/augmentors/image_editing_transform.py (99%) rename cosmos_framework/data/{vfm => generator}/augmentors/image_editing_transform_test.py (97%) rename cosmos_framework/data/{vfm => generator}/augmentors/image_resolution_filter.py (96%) rename cosmos_framework/data/{vfm => generator}/augmentors/interleaved_image_transform.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/interleaved_video_parsing.py (99%) rename cosmos_framework/data/{vfm => generator}/augmentors/merge_datadict.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/pkl_to_media.py (99%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/__init__.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/bytes_to_media.py (98%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/filter_output_key.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/filter_seq_length.py (97%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/floating_number_format.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/format_describe_anything.py (99%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/nvlm_data_to_conversation.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/prompt_format.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/shuffle_text_media_order.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/timestamp.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/timestamp_with_subject_tracking.py (99%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/timestamp_without_augment_message.py (99%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/timestamp_without_end_time.py (99%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/tokenize_data.py (98%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/user_prompt_caption_general.json (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/reasoner/user_prompt_ocr.json (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/resolution_text_info.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/sequence_plan.py (98%) rename cosmos_framework/data/{vfm => generator}/augmentors/sound_sequence_plan.py (94%) rename cosmos_framework/data/{vfm => generator}/augmentors/text_tokenizer.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/text_transforms_for_image.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/text_transforms_for_video.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/transfer_control_input/__init__.py (73%) rename cosmos_framework/data/{vfm => generator}/augmentors/transfer_control_input/blur.py (98%) rename cosmos_framework/data/{vfm => generator}/augmentors/transfer_control_input/control_input.py (99%) rename cosmos_framework/data/{vfm => generator}/augmentors/transfer_control_input/fast_blur.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/transfer_control_input/seg.py (100%) rename cosmos_framework/data/{vfm => generator}/augmentors/transfer_control_transform.py (98%) rename cosmos_framework/data/{vfm => generator}/augmentors/video_parsing.py (99%) rename cosmos_framework/data/{vfm => generator}/dataflow/__init__.py (56%) rename cosmos_framework/data/{vfm => generator}/dataflow/base.py (100%) rename cosmos_framework/data/{vfm => generator}/dataflow/batchers.py (99%) rename cosmos_framework/data/{vfm => generator}/dataflow/collators.py (99%) rename cosmos_framework/data/{vfm => generator}/dataflow/distributors.py (98%) rename cosmos_framework/data/{vfm => generator}/dataflow/golden_vfm_test.py (98%) rename cosmos_framework/data/{vfm => generator}/dataflow/loader.py (97%) rename cosmos_framework/data/{vfm => generator}/dataflow/processors.py (85%) rename cosmos_framework/data/{vfm => generator}/dataflow/resume_test.py (96%) rename cosmos_framework/data/{vfm => generator}/joint_dataloader.py (99%) rename cosmos_framework/data/{vfm => generator}/local_datasets/__init__.py (100%) rename cosmos_framework/data/{vfm => generator}/local_datasets/helper.py (100%) rename cosmos_framework/data/{vfm => generator}/local_datasets/sft_dataset.py (98%) rename cosmos_framework/data/{vfm => generator}/local_datasets/sft_dataset_caption_test.py (96%) rename cosmos_framework/data/{vfm => generator}/packing_iterable_dataset.py (99%) rename cosmos_framework/data/{vfm => generator}/processors/__init__.py (93%) rename cosmos_framework/data/{vfm => generator}/processors/base.py (97%) rename cosmos_framework/data/{vfm => generator}/processors/nemotron3densevl_processor.py (99%) rename cosmos_framework/data/{vfm => generator}/processors/nemotronvl_processor.py (99%) rename cosmos_framework/data/{vfm => generator}/processors/qwen3vl_processor.py (98%) rename cosmos_framework/data/{vfm => generator}/reasoner/video_decoder_qwen.py (99%) rename cosmos_framework/data/{vfm => generator}/recipe_database_api.py (100%) rename cosmos_framework/data/{vfm => generator}/sequence_packing/__init__.py (73%) rename cosmos_framework/data/{vfm => generator}/sequence_packing/modalities.py (99%) rename cosmos_framework/data/{vfm => generator}/sequence_packing/mrope.py (100%) rename cosmos_framework/data/{vfm => generator}/sequence_packing/natten.py (100%) rename cosmos_framework/data/{vfm => generator}/sequence_packing/packers.py (98%) rename cosmos_framework/data/{vfm => generator}/sequence_packing/runtime.py (100%) rename cosmos_framework/data/{vfm => generator}/sequence_packing/temporal_causal.py (98%) rename cosmos_framework/data/{vfm => generator}/sequence_packing/types.py (99%) rename cosmos_framework/data/{vfm => generator}/sound_data_utils.py (95%) rename cosmos_framework/data/{vfm => generator}/utils.py (100%) rename cosmos_framework/data/{vfm => generator}/watchdog.py (100%) rename cosmos_framework/data/{vlm => reasoner}/__init__.py (100%) rename cosmos_framework/data/{vlm => reasoner}/data_sources_videophy2/__init__.py (100%) rename cosmos_framework/data/{vlm => reasoner}/data_sources_videophy2/data_weight/__init__.py (100%) rename cosmos_framework/data/{vlm => reasoner}/data_sources_videophy2/data_weight/default.py (84%) rename cosmos_framework/data/{vlm => reasoner}/data_sources_videophy2/videophy2.py (95%) rename cosmos_framework/data/{vlm => reasoner}/local_dataset_utils.py (100%) rename cosmos_framework/data/{vlm => reasoner}/local_sft_dataset.py (100%) rename cosmos_framework/data/{vlm => reasoner}/processors/__init__.py (73%) rename cosmos_framework/data/{vlm => reasoner}/processors/nemotron3densevl_processor.py (99%) rename cosmos_framework/data/{vlm => reasoner}/processors/nemotronvl_processor.py (99%) rename cosmos_framework/data/{vlm => reasoner}/processors/qwen3vl_processor.py (98%) delete mode 100644 cosmos_framework/data/vfm/action/datasets/__init__.py rename cosmos_framework/model/{vfm => generator}/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/algorithm/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/algorithm/loss/__init__.py (51%) rename cosmos_framework/model/{vfm => generator}/algorithm/loss/cross_entropy.py (98%) rename cosmos_framework/model/{vfm => generator}/algorithm/loss/flow_matching.py (98%) rename cosmos_framework/model/{vfm => generator}/algorithm/loss/load_balancing.py (97%) rename cosmos_framework/model/{vfm => generator}/algorithm/loss/time_weight.py (100%) rename cosmos_framework/model/{vfm => generator}/diffusion/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/diffusion/rectified_flow.py (99%) rename cosmos_framework/model/{vfm => generator}/diffusion/samplers/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/diffusion/samplers/edm.py (100%) rename cosmos_framework/model/{vfm => generator}/diffusion/samplers/fixed_step.py (98%) rename cosmos_framework/model/{vfm => generator}/diffusion/samplers/fixed_step_test.py (98%) rename cosmos_framework/model/{vfm => generator}/diffusion/samplers/fm_solvers_unipc.py (100%) rename cosmos_framework/model/{vfm => generator}/diffusion/samplers/unipc.py (95%) rename cosmos_framework/model/{vfm => generator}/diffusion/samplers/utils.py (100%) rename cosmos_framework/model/{vfm => generator}/diffusion/samplers/utils_test.py (97%) rename cosmos_framework/model/{vfm => generator}/hf_model.py (96%) rename cosmos_framework/model/{vfm => generator}/mot/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/mot/attention.py (99%) rename cosmos_framework/model/{vfm => generator}/mot/attention_test.py (98%) rename cosmos_framework/model/{vfm => generator}/mot/cfgp_ar_test.py (97%) rename cosmos_framework/model/{vfm => generator}/mot/context_parallel_test.py (97%) rename cosmos_framework/model/{vfm => generator}/mot/context_parallel_utils.py (98%) rename cosmos_framework/model/{vfm => generator}/mot/cosmos3_vfm_network.py (98%) rename cosmos_framework/model/{vfm => generator}/mot/domain_aware_linear.py (100%) rename cosmos_framework/model/{vfm => generator}/mot/dot_product_attention.py (100%) rename cosmos_framework/model/{vfm => generator}/mot/modeling_utils.py (97%) rename cosmos_framework/model/{vfm => generator}/mot/parallelize_unified_mot.py (97%) rename cosmos_framework/model/{vfm => generator}/mot/parallelize_vfm_network.py (97%) rename cosmos_framework/model/{vfm => generator}/mot/qwen3_vl_unified_mot.py (77%) rename cosmos_framework/model/{vfm => generator}/mot/und_k_norm_example_test.py (96%) rename cosmos_framework/model/{vfm => generator}/mot/unified_mot.py (99%) rename cosmos_framework/model/{vfm => generator}/omni_mot_model.py (98%) rename cosmos_framework/model/{vfm => generator}/parallelize_vlm.py (97%) rename cosmos_framework/model/{vfm => generator}/reasoner/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/nemotron_3_dense_vl/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py (98%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl/configs/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl/configuration_qwen3_vl.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl/qwen3_vl.py (99%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl/utils.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl/video_processing_qwen3_vl.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl_moe/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl_moe/configs/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl_moe/moe.py (98%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl_moe/moe_bench.py (93%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl_moe/moe_kernels.py (100%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl_moe/moe_test.py (93%) rename cosmos_framework/model/{vfm => generator}/reasoner/qwen3_vl_moe/qwen3_vl_moe.py (99%) rename cosmos_framework/model/{vfm => generator}/tokenizers/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae.py (98%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/__init__.py (64%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/activations.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/alias_free_torch/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/alias_free_torch/act.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/alias_free_torch/filter.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/alias_free_torch/resample.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/bottlenecks.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/env.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/models.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/modules.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/audio/avae_utils/modules_encodec.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/dc_ae/__init__.py (83%) rename cosmos_framework/model/{vfm => generator}/tokenizers/dc_ae/convert_pt_to_distcp.py (92%) rename cosmos_framework/model/{vfm => generator}/tokenizers/dc_ae/dc_ae_4x32x32.py (96%) rename cosmos_framework/model/{vfm => generator}/tokenizers/dc_ae/dc_ae_v.py (99%) rename cosmos_framework/model/{vfm => generator}/tokenizers/dc_ae/dc_ae_v_ops.py (99%) rename cosmos_framework/model/{vfm => generator}/tokenizers/dc_ae/dc_ae_v_triton_rms_norm.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/flux_vae_8x8.py (99%) rename cosmos_framework/model/{vfm => generator}/tokenizers/interface.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/stable_diffusion_vae_8x8.py (98%) rename cosmos_framework/model/{vfm => generator}/tokenizers/tokenization_qwen2.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/uniae/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/tokenizers/uniae/frame_math.py (99%) rename cosmos_framework/model/{vfm => generator}/tokenizers/uniae/frame_math_test.py (93%) rename cosmos_framework/model/{vfm => generator}/tokenizers/uniae/noncausal_4x16x16.py (98%) rename cosmos_framework/model/{vfm => generator}/tokenizers/wan2pt1_vae_4x8x8.py (99%) rename cosmos_framework/model/{vfm => generator}/tokenizers/wan2pt2_vae_4x16x16.py (99%) rename cosmos_framework/model/{vfm => generator}/upsampler/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/upsampler/prompts.py (99%) rename cosmos_framework/model/{vfm => generator}/utils/__init__.py (100%) rename cosmos_framework/model/{vfm => generator}/utils/data_and_condition.py (100%) rename cosmos_framework/model/{vfm => generator}/utils/memory.py (100%) rename cosmos_framework/model/{vfm => generator}/utils/safetensors_loader.py (99%) rename cosmos_framework/model/{vfm => generator}/utils/safetensors_loader_test.py (98%) rename cosmos_framework/model/{vfm => generator}/utils/taylorseer.py (100%) rename cosmos_framework/model/{vfm => generator}/vlm_model.py (97%) rename cosmos_framework/utils/{vfm => generator}/__init__.py (100%) rename cosmos_framework/utils/{vfm => generator}/data_utils.py (100%) rename cosmos_framework/utils/{vfm => generator}/dtensor_helper.py (100%) rename cosmos_framework/utils/{vfm => generator}/flash_attn.py (100%) rename cosmos_framework/utils/{vfm => generator}/fused_adam.py (100%) rename cosmos_framework/utils/{vfm => generator}/hf_attention_cosmos.py (100%) rename cosmos_framework/utils/{vfm => generator}/lora.py (100%) rename cosmos_framework/utils/{vfm => generator}/model_loader.py (97%) rename cosmos_framework/utils/{vfm => generator}/model_weights_stats.py (100%) rename cosmos_framework/utils/{vfm => generator}/monkey_patch.py (100%) rename cosmos_framework/utils/{vfm => generator}/optimizer.py (99%) rename cosmos_framework/utils/{vfm => generator}/parallelism.py (98%) rename cosmos_framework/utils/{vfm => generator}/rand_state.py (100%) rename cosmos_framework/utils/{vfm => generator}/reasoner/__init__.py (100%) rename cosmos_framework/utils/{vfm => generator}/reasoner/constant.py (100%) rename cosmos_framework/utils/{vfm => generator}/reasoner/create_position_ids.py (100%) rename cosmos_framework/utils/{vfm => generator}/reasoner/flop_calculator.py (100%) rename cosmos_framework/utils/{vfm => generator}/reasoner/pretrained_models_downloader.py (99%) rename cosmos_framework/utils/{vfm => generator}/video_preprocess.py (100%) rename cosmos_framework/utils/{vlm => reasoner}/__init__.py (100%) rename cosmos_framework/utils/{vlm => reasoner}/compute_flops_qwen3vl.py (100%) rename cosmos_framework/utils/{vlm => reasoner}/configs_defaults/__init__.py (100%) rename cosmos_framework/utils/{vlm => reasoner}/configs_defaults/checkpointer.py (97%) rename cosmos_framework/utils/{vlm => reasoner}/constant.py (100%) rename cosmos_framework/utils/{vlm => reasoner}/create_position_ids.py (100%) rename cosmos_framework/utils/{vlm => reasoner}/dcp_checkpointer.py (99%) rename cosmos_framework/utils/{vlm => reasoner}/distributed.py (100%) rename cosmos_framework/utils/{vlm => reasoner}/flop_calculator.py (98%) rename cosmos_framework/utils/{vlm => reasoner}/fused_adam.py (99%) rename cosmos_framework/utils/{vlm => reasoner}/model_wrapper.py (100%) rename cosmos_framework/utils/{vlm => reasoner}/optimizer.py (99%) rename cosmos_framework/utils/{vlm => reasoner}/planner.py (100%) rename cosmos_framework/utils/{vlm => reasoner}/pretrained_models_downloader.py (97%) rename cosmos_framework/utils/{vlm => reasoner}/video_preprocess.py (100%) diff --git a/.agents/skills/cosmos3-codebase-nav/SKILL.md b/.agents/skills/cosmos3-codebase-nav/SKILL.md index eaaf7c9d..cb2bab03 100644 --- a/.agents/skills/cosmos3-codebase-nav/SKILL.md +++ b/.agents/skills/cosmos3-codebase-nav/SKILL.md @@ -22,7 +22,7 @@ All paths below are relative to this file's location (`.agents/skills/cosmos3-co - `cosmos_framework/` — main training package (data, model, trainer, callbacks, checkpoint, utils, …). - `cosmos_framework/configs/base/experiment/` — vfm (generator) experiment SKUs referenced by `[train.train_policy].experiment` in the recipe TOMLs. -- `cosmos_framework/configs/base/vlm/experiment/` — vlm (reasoner) experiment SKUs. +- `cosmos_framework/configs/base/reasoner/experiment/` — vlm (reasoner) experiment SKUs. - `cosmos_framework/inference/` — inference subpackage (args, model, inference engine, defaults, Ray serving, common helpers). - `cosmos_framework/scripts/` — top-level entry-point scripts (train, inference, eval, export_model, convert_model_to_dcp, upsample_prompts, caption_from_video, captions_to_sft_jsonl, action_policy_server, …). Invoked as `python -m cosmos_framework.scripts.`. - `examples/toml/sft_config/.toml` + `examples/launch_sft_.sh` — paired SFT recipes (training entry-point input). The shell sources `examples/_sft_launcher_common.sh`, which forwards into `cosmos_framework.scripts.train --sft-toml=...`. @@ -45,7 +45,7 @@ All paths below are relative to this file's location (`.agents/skills/cosmos3-co | SFT recipe TOMLs (paired with `examples/launch_sft_*.sh`) | `../../../examples/toml/sft_config/.toml` | | SFT pydantic schema (validates the recipe TOML) | `../../../cosmos_framework/configs/toml_config/sft_config.py` | | Training experiment SKUs (vfm) | `../../../cosmos_framework/configs/base/experiment/` | -| Training experiment SKUs (vlm / reasoner) | `../../../cosmos_framework/configs/base/vlm/experiment/` | +| Training experiment SKUs (vlm / reasoner) | `../../../cosmos_framework/configs/base/reasoner/experiment/` | | Example inputs | `../../../inputs/omni/t2i.json`, `../../../inputs/omni/t2v.json`, `../../../inputs/omni/i2v.json`, … | Available modality modes for defaults: `text2image`, `text2video`, `image2video`, `image2image`, `video2video`, `forward_dynamics`, `inverse_dynamics`, `policy`. diff --git a/.claude/skills/cosmos3-codebase-nav/SKILL.md b/.claude/skills/cosmos3-codebase-nav/SKILL.md index 964c760c..f3f070df 100644 --- a/.claude/skills/cosmos3-codebase-nav/SKILL.md +++ b/.claude/skills/cosmos3-codebase-nav/SKILL.md @@ -22,7 +22,7 @@ All paths below are relative to this file's location (`.claude/skills/cosmos3-co - `cosmos_framework/` — main training package (data, model, trainer, callbacks, checkpoint, utils, …). - `cosmos_framework/configs/base/experiment/` — vfm (generator) experiment SKUs referenced by `[train.train_policy].experiment` in the recipe TOMLs. -- `cosmos_framework/configs/base/vlm/experiment/` — vlm (reasoner) experiment SKUs. +- `cosmos_framework/configs/base/reasoner/experiment/` — vlm (reasoner) experiment SKUs. - `cosmos_framework/inference/` — inference subpackage (args, model, inference engine, defaults, Ray serving, common helpers). - `cosmos_framework/scripts/` — top-level entry-point scripts (train, inference, eval, export_model, convert_model_to_dcp, upsample_prompts, caption_from_video, captions_to_sft_jsonl, action_policy_server, …). Invoked as `python -m cosmos_framework.scripts.`. - `examples/toml/sft_config/.toml` + `examples/launch_sft_.sh` — paired SFT recipes (training entry-point input). The shell sources `examples/_sft_launcher_common.sh`, which forwards into `cosmos_framework.scripts.train --sft-toml=...`. @@ -45,7 +45,7 @@ All paths below are relative to this file's location (`.claude/skills/cosmos3-co | SFT recipe TOMLs (paired with `examples/launch_sft_*.sh`) | `../../../examples/toml/sft_config/.toml` | | SFT pydantic schema (validates the recipe TOML) | `../../../cosmos_framework/configs/toml_config/sft_config.py` | | Training experiment SKUs (vfm) | `../../../cosmos_framework/configs/base/experiment/` | -| Training experiment SKUs (vlm / reasoner) | `../../../cosmos_framework/configs/base/vlm/experiment/` | +| Training experiment SKUs (vlm / reasoner) | `../../../cosmos_framework/configs/base/reasoner/experiment/` | | Example inputs | `../../../inputs/omni/t2i.json`, `../../../inputs/omni/t2v.json`, `../../../inputs/omni/i2v.json`, … | Available modality modes for defaults: `text2image`, `text2video`, `image2video`, `image2image`, `video2video`, `forward_dynamics`, `inverse_dynamics`, `policy`. diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index 72ff4f46..5142f40a 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -256,13 +256,13 @@ jobs: run: | export LD_LIBRARY_PATH= uv run --all-extras --group=cu128-train torchrun --nproc_per_node=2 -m pytest -v \ - cosmos_framework/model/vfm/mot/cfgp_ar_test.py -o addopts= + cosmos_framework/model/generator/mot/cfgp_ar_test.py -o addopts= - name: Distributed unit tests - context_parallel (torchrun, 4 ranks) run: | export LD_LIBRARY_PATH= uv run --all-extras --group=cu128-train torchrun --nproc_per_node=4 -m pytest -v \ - cosmos_framework/model/vfm/mot/context_parallel_test.py -o addopts= + cosmos_framework/model/generator/mot/context_parallel_test.py -o addopts= # Clear everything the suite writes into the working tree (all gitignored # scratch): pytest tmp dirs (DCP checkpoint, logs), the script-test diff --git a/NOTICE b/NOTICE index 9cbe5c66..1e1f5ad2 100644 --- a/NOTICE +++ b/NOTICE @@ -15,10 +15,10 @@ Notices: Copyright 2025 The Qwen Team and The HuggingFace Inc. team. Copyright 2025 HuggingFace Inc. team. Files: - cosmos_framework/model/vfm/vlm/qwen3_vl/configuration_qwen3_vl.py - cosmos_framework/model/vfm/vlm/qwen3_vl/qwen3_vl.py - cosmos_framework/model/vfm/vlm/qwen3_vl/utils.py - cosmos_framework/model/vfm/vlm/qwen3_vl/video_processing_qwen3_vl.py + cosmos_framework/model/generator/vlm/qwen3_vl/configuration_qwen3_vl.py + cosmos_framework/model/generator/vlm/qwen3_vl/qwen3_vl.py + cosmos_framework/model/generator/vlm/qwen3_vl/utils.py + cosmos_framework/model/generator/vlm/qwen3_vl/video_processing_qwen3_vl.py ByteDance-Seed Bagel / Qwen2 tokenizer Source: https://github.com/ByteDance-Seed/Bagel @@ -26,7 +26,7 @@ License: Apache-2.0 Notice: Copyright 2024 The Qwen Team and The HuggingFace Inc. team. File: - cosmos_framework/model/vfm/tokenizers/tokenization_qwen2.py + cosmos_framework/model/generator/tokenizers/tokenization_qwen2.py HuggingFace Diffusers UniPC scheduler Source: https://github.com/huggingface/diffusers @@ -35,7 +35,7 @@ Notices: Copyright 2024 TSAIL Team and The HuggingFace Team. All rights reserved. Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. File: - cosmos_framework/model/vfm/diffusion/samplers/fm_solvers_unipc.py + cosmos_framework/model/generator/diffusion/samplers/fm_solvers_unipc.py Detectron2 lazy config helpers Source: https://github.com/facebookresearch/detectron2 @@ -51,4 +51,4 @@ License: GPL-3.0 Notice: Upstream source files do not contain an explicit copyright notice. Files: - cosmos_framework/model/vfm/utils/taylorseer.py + cosmos_framework/model/generator/utils/taylorseer.py diff --git a/cosmos_framework/callbacks/compile_tokenizer.py b/cosmos_framework/callbacks/compile_tokenizer.py index 3ee15d24..aa0de519 100644 --- a/cosmos_framework/callbacks/compile_tokenizer.py +++ b/cosmos_framework/callbacks/compile_tokenizer.py @@ -4,7 +4,7 @@ """Training callback that defers AOT compilation of the VAE tokenizer. The actual compilation logic lives in -:meth:`~cosmos_framework.model.vfm.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface.compile_encode`. +:meth:`~cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface.compile_encode`. This module provides a :class:`CompileTokenizer` callback that invokes it at the right point during training (after ``compile_after_iterations`` steps, to avoid NCCL timeouts during CUDA/cuDNN warm-up). @@ -27,7 +27,7 @@ from cosmos_framework.utils import log from cosmos_framework.utils.callback import Callback -from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel +from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel class CompileTokenizer(Callback): diff --git a/cosmos_framework/callbacks/every_n_draw_sample.py b/cosmos_framework/callbacks/every_n_draw_sample.py index 9aa96fa5..7d181d6f 100644 --- a/cosmos_framework/callbacks/every_n_draw_sample.py +++ b/cosmos_framework/callbacks/every_n_draw_sample.py @@ -21,7 +21,7 @@ from cosmos_framework.utils import distributed, log, misc from cosmos_framework.utils.easy_io import easy_io from cosmos_framework.tools.visualize.video import save_img_or_video -from cosmos_framework.utils.vfm.data_utils import slice_data_batch +from cosmos_framework.utils.generator.data_utils import slice_data_batch def resize_image(image: torch.Tensor, size: int = 1024) -> torch.Tensor: diff --git a/cosmos_framework/callbacks/expert_heatmap.py b/cosmos_framework/callbacks/expert_heatmap.py index f3262fc7..888d8949 100644 --- a/cosmos_framework/callbacks/expert_heatmap.py +++ b/cosmos_framework/callbacks/expert_heatmap.py @@ -10,7 +10,7 @@ from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.trainer import ImaginaireTrainer from cosmos_framework.utils import distributed -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock def compute_expert_heatmap(vfm: torch.nn.Module) -> dict[str, torch.Tensor]: diff --git a/cosmos_framework/callbacks/hf_export.py b/cosmos_framework/callbacks/hf_export.py index 6d235688..97de1fd8 100644 --- a/cosmos_framework/callbacks/hf_export.py +++ b/cosmos_framework/callbacks/hf_export.py @@ -133,7 +133,7 @@ def on_save_checkpoint(self, model: Any, state_dict: dict[str, Any]) -> None: return # Deferred import to avoid circular dependency at module load time. - from cosmos_framework.model.vfm.vlm_model import VLMModel + from cosmos_framework.model.generator.vlm_model import VLMModel if not isinstance(model, VLMModel): # The legacy vlm/train.py path passes model_parts: list[nn.Module] (raw HF diff --git a/cosmos_framework/callbacks/moe_specialization_callback.py b/cosmos_framework/callbacks/moe_specialization_callback.py index 4d5d4233..c5fa622f 100644 --- a/cosmos_framework/callbacks/moe_specialization_callback.py +++ b/cosmos_framework/callbacks/moe_specialization_callback.py @@ -39,7 +39,7 @@ from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.trainer import ImaginaireTrainer from cosmos_framework.utils import distributed -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock def _get_device_mesh(vfm: torch.nn.Module): diff --git a/cosmos_framework/callbacks/moe_stability_callback.py b/cosmos_framework/callbacks/moe_stability_callback.py index 8e241071..8071fc3a 100644 --- a/cosmos_framework/callbacks/moe_stability_callback.py +++ b/cosmos_framework/callbacks/moe_stability_callback.py @@ -107,7 +107,7 @@ from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.trainer import ImaginaireTrainer from cosmos_framework.utils import distributed -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock def _effective_experts( diff --git a/cosmos_framework/callbacks/norm_monitor.py b/cosmos_framework/callbacks/norm_monitor.py index 460804cf..36e6ea00 100644 --- a/cosmos_framework/callbacks/norm_monitor.py +++ b/cosmos_framework/callbacks/norm_monitor.py @@ -14,7 +14,7 @@ from cosmos_framework.utils import distributed, log, misc from cosmos_framework.utils.callback import Callback from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.data.vfm.sequence_packing.runtime import get_gen_seq +from cosmos_framework.data.generator.sequence_packing.runtime import get_gen_seq try: from apex.contrib.layer_norm import FastLayerNorm diff --git a/cosmos_framework/callbacks/sequence_packing_padding.py b/cosmos_framework/callbacks/sequence_packing_padding.py index 91f3b7ed..15c4088b 100644 --- a/cosmos_framework/callbacks/sequence_packing_padding.py +++ b/cosmos_framework/callbacks/sequence_packing_padding.py @@ -7,7 +7,7 @@ from cosmos_framework.callbacks.every_n import EveryN from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.trainer import ImaginaireTrainer -from cosmos_framework.data.vfm.sequence_packing.runtime import get_padding_stats +from cosmos_framework.data.generator.sequence_packing.runtime import get_padding_stats class SequencePackingPadding(EveryN): diff --git a/cosmos_framework/callbacks/tokens_per_sec.py b/cosmos_framework/callbacks/tokens_per_sec.py index 1634db95..d4bc2319 100644 --- a/cosmos_framework/callbacks/tokens_per_sec.py +++ b/cosmos_framework/callbacks/tokens_per_sec.py @@ -11,7 +11,7 @@ ``MFUCallback`` is written for the OmniMoT/VFM network -- it reads ``model.net.language_model.config`` and per-modality token counts (``output_batch["und_token_length"]`` etc.). ``VLMModel`` -(``cosmos_framework/model/vfm/vlm_model.py``) is a plain HF wrapper: it exposes +(``cosmos_framework/model/generator/vlm_model.py``) is a plain HF wrapper: it exposes ``self.model`` (no ``.net``) and its ``training_step`` returns only ``{"loss", "loss_avg", "labels"}``. ``MFUCallback`` would therefore silently no-op (token length is ``None``) or fail on ``model.net``. This callback instead diff --git a/cosmos_framework/callbacks/training_stats.py b/cosmos_framework/callbacks/training_stats.py index 9eb02129..1e7f3a7e 100644 --- a/cosmos_framework/callbacks/training_stats.py +++ b/cosmos_framework/callbacks/training_stats.py @@ -11,7 +11,7 @@ from cosmos_framework.utils import distributed from cosmos_framework.utils.callback import Callback from cosmos_framework.callbacks.wandb_log import _LossRecord -from cosmos_framework.data.vfm.action.domain_utils import EMBODIMENT_TO_DOMAIN_ID +from cosmos_framework.data.generator.action.domain_utils import EMBODIMENT_TO_DOMAIN_ID # Build inverse mapping: domain_id -> embodiment_type. First occurrence wins when multiple embodiment names share the # same domain id. diff --git a/cosmos_framework/checkpoint/dcp.py b/cosmos_framework/checkpoint/dcp.py index d3cad0ea..43145c88 100644 --- a/cosmos_framework/checkpoint/dcp.py +++ b/cosmos_framework/checkpoint/dcp.py @@ -75,7 +75,7 @@ from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.utils import callback, distributed, log, misc from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.utils.vfm.rand_state import get_rand_state_dict, set_rand_state_dict +from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict class ModelWrapper(Stateful): diff --git a/cosmos_framework/configs/base/defaults/activation_checkpointing.py b/cosmos_framework/configs/base/defaults/activation_checkpointing.py index 543c0296..765dcc87 100644 --- a/cosmos_framework/configs/base/defaults/activation_checkpointing.py +++ b/cosmos_framework/configs/base/defaults/activation_checkpointing.py @@ -35,9 +35,9 @@ class ActivationCheckpointingConfig: Read sites: - MoT path consumes every field — see - cosmos_framework/model/vfm/mot/parallelize_unified_mot.py. + cosmos_framework/model/generator/mot/parallelize_unified_mot.py. - VLM path consumes only ``mode`` (and only ``"full"`` enables - checkpointing) — see cosmos_framework/model/vfm/vlm_model.py. + checkpointing) — see cosmos_framework/model/generator/vlm_model.py. """ # AC mode: diff --git a/cosmos_framework/configs/base/defaults/model.py b/cosmos_framework/configs/base/defaults/model.py index d725e0a1..43fecea9 100644 --- a/cosmos_framework/configs/base/defaults/model.py +++ b/cosmos_framework/configs/base/defaults/model.py @@ -6,7 +6,7 @@ from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.configs.base.defaults.model_config import OmniMoTModelConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig -from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel +from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel MOT_DDP_CONFIG = dict( trainer=dict( diff --git a/cosmos_framework/configs/base/defaults/open_source_dataloader.py b/cosmos_framework/configs/base/defaults/open_source_dataloader.py index 4dd42b33..50417c6d 100644 --- a/cosmos_framework/configs/base/defaults/open_source_dataloader.py +++ b/cosmos_framework/configs/base/defaults/open_source_dataloader.py @@ -21,17 +21,17 @@ Original YAML reference target paths use the ``cosmos3._src.vfm.*`` namespace (the OSS-release form of ``projects.cosmos3.vfm.*``); inside this released -tree the same modules live under ``cosmos_framework.data.vfm.*``. +tree the same modules live under ``cosmos_framework.data.generator.*``. """ from hydra.core.config_store import ConfigStore from cosmos_framework.configs.base.defaults.reasoner import create_qwen2_tokenizer_with_download -from cosmos_framework.data.vfm.joint_dataloader import ( +from cosmos_framework.data.generator.joint_dataloader import ( PackingDataLoader, RankPartitionedDataLoader, ) -from cosmos_framework.data.vfm.local_datasets.sft_dataset import get_sft_dataset +from cosmos_framework.data.generator.local_datasets.sft_dataset import get_sft_dataset from cosmos_framework.utils.lazy_config import LazyCall as L # --------------------------------------------------------------------------- diff --git a/cosmos_framework/configs/base/defaults/optimizer.py b/cosmos_framework/configs/base/defaults/optimizer.py index 4e254e6a..dc56906c 100644 --- a/cosmos_framework/configs/base/defaults/optimizer.py +++ b/cosmos_framework/configs/base/defaults/optimizer.py @@ -8,7 +8,7 @@ from cosmos_framework.utils.lazy_config import PLACEHOLDER from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.utils.config_helper import ConfigStore -from cosmos_framework.utils.vfm.optimizer import build_lr_scheduler, build_optimizer +from cosmos_framework.utils.generator.optimizer import build_lr_scheduler, build_optimizer OPTIMIZER_KWARGS: dict[str, Any] = dict( # Learning rate for the optimizer. diff --git a/cosmos_framework/configs/base/defaults/reasoner.py b/cosmos_framework/configs/base/defaults/reasoner.py index 7293b6e6..bee81b63 100644 --- a/cosmos_framework/configs/base/defaults/reasoner.py +++ b/cosmos_framework/configs/base/defaults/reasoner.py @@ -16,7 +16,7 @@ from cosmos_framework.utils import log from cosmos_framework.utils.config_helper import ConfigStore from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.model.vfm.mot.unified_mot import ( +from cosmos_framework.model.generator.mot.unified_mot import ( Nemotron3DenseVLMoTConfig, Nemotron3DenseVLTextForCausalLM, Qwen3MoTConfig, @@ -25,8 +25,8 @@ Qwen3VLMoTConfig, Qwen3VLTextForCausalLM, ) -from cosmos_framework.data.vfm.processors import LLMTokenizerProcessor, build_processor_lazy -from cosmos_framework.model.vfm.tokenizers.tokenization_qwen2 import Qwen2Tokenizer +from cosmos_framework.data.generator.processors import LLMTokenizerProcessor, build_processor_lazy +from cosmos_framework.model.generator.tokenizers.tokenization_qwen2 import Qwen2Tokenizer def create_vlm_config(base_config: LazyDict, **overrides): @@ -177,7 +177,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3MoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/llm/qwen3/configs/Qwen3-0.6B.json" + json_file="cosmos_framework/model/generator/llm/qwen3/configs/Qwen3-0.6B.json" ), qk_norm_for_text=True, ), @@ -197,7 +197,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3MoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/llm/qwen3/configs/Qwen3-0.6B.json" + json_file="cosmos_framework/model/generator/llm/qwen3/configs/Qwen3-0.6B.json" ), qk_norm_for_text=True, ), @@ -217,7 +217,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -243,7 +243,7 @@ class VLMConfig: model_instance=L(Qwen3VLMoeTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoeMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json" ), layer_module="Qwen3VLMoeTextMoTDecoderLayer", qk_norm_for_text=True, @@ -266,7 +266,7 @@ class VLMConfig: model_instance=L(Qwen3VLMoeTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoeMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json" ), layer_module="Qwen3VLMoeTextMoTDecoderLayer", qk_norm_for_text=True, @@ -291,7 +291,7 @@ class VLMConfig: model_instance=L(Qwen3VLMoeTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoeMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json" ), layer_module="Qwen3VLMoeTextMoTDecoderLayer", qk_norm_for_text=True, @@ -317,7 +317,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" ), qk_norm_for_text=True, ), @@ -337,7 +337,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" ), qk_norm_for_text=True, ), @@ -358,7 +358,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" ), qk_norm_for_text=True, ), @@ -374,7 +374,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -396,7 +396,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -418,7 +418,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json" ), qk_norm_for_text=True, ), @@ -441,7 +441,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json" ), qk_norm_for_text=True, ), @@ -461,7 +461,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json" ), qk_norm_for_text=True, ), @@ -484,7 +484,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -504,7 +504,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -525,7 +525,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -546,7 +546,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -567,7 +567,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -588,7 +588,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json" ), qk_norm_for_text=True, ), @@ -610,7 +610,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" ), qk_norm_for_text=True, ), @@ -630,7 +630,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" ), qk_norm_for_text=True, ), @@ -651,7 +651,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" ), qk_norm_for_text=True, ), @@ -672,7 +672,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" ), qk_norm_for_text=True, ), @@ -693,7 +693,7 @@ class VLMConfig: model_instance=L(Qwen3VLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" + json_file="cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json" ), qk_norm_for_text=True, ), @@ -720,7 +720,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -744,7 +744,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -768,7 +768,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, ), @@ -792,7 +792,7 @@ class VLMConfig: model_instance=L(Nemotron3DenseVLTextForCausalLM)( config=L(create_vlm_config)( base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( - json_file="cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + json_file="cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" ), qk_norm_for_text=False, use_und_k_norm_for_gen=True, diff --git a/cosmos_framework/configs/base/defaults/tokenizer.py b/cosmos_framework/configs/base/defaults/tokenizer.py index 530907d4..04c08873 100644 --- a/cosmos_framework/configs/base/defaults/tokenizer.py +++ b/cosmos_framework/configs/base/defaults/tokenizer.py @@ -5,13 +5,13 @@ from cosmos_framework.utils.lazy_config import PLACEHOLDER, LazyDict from cosmos_framework.utils.lazy_config import LazyCall as L -from cosmos_framework.model.vfm.tokenizers.audio.avae import AVAEInterface -from cosmos_framework.model.vfm.tokenizers.dc_ae.dc_ae_4x32x32 import DCAE4x32x32Interface -from cosmos_framework.model.vfm.tokenizers.flux_vae_8x8 import FluxVAEInterface -from cosmos_framework.model.vfm.tokenizers.stable_diffusion_vae_8x8 import StableDiffusionVAEInterface -from cosmos_framework.model.vfm.tokenizers.uniae.noncausal_4x16x16 import UniAEVAEInterface -from cosmos_framework.model.vfm.tokenizers.wan2pt1_vae_4x8x8 import Wan2pt1VAEInterface -from cosmos_framework.model.vfm.tokenizers.wan2pt2_vae_4x16x16 import Wan2pt2VAEInterface +from cosmos_framework.model.generator.tokenizers.audio.avae import AVAEInterface +from cosmos_framework.model.generator.tokenizers.dc_ae.dc_ae_4x32x32 import DCAE4x32x32Interface +from cosmos_framework.model.generator.tokenizers.flux_vae_8x8 import FluxVAEInterface +from cosmos_framework.model.generator.tokenizers.stable_diffusion_vae_8x8 import StableDiffusionVAEInterface +from cosmos_framework.model.generator.tokenizers.uniae.noncausal_4x16x16 import UniAEVAEInterface +from cosmos_framework.model.generator.tokenizers.wan2pt1_vae_4x8x8 import Wan2pt1VAEInterface +from cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16 import Wan2pt2VAEInterface PRETRAINED_TOKENIZER_SD_VAE_REPO = "stabilityai/sd-vae-ft-ema" PRETRAINED_TOKENIZER_WAN2PT1_VAE_PTH = "pretrained/tokenizers/video/wan2pt1/Wan2.1_VAE.pth" diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py index 92cb2818..926afea9 100644 --- a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py @@ -25,11 +25,11 @@ from cosmos_framework.utils.lazy_config import LazyDict from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG -from cosmos_framework.data.vfm.joint_dataloader import ( +from cosmos_framework.data.generator.joint_dataloader import ( PackingDataLoader, RankPartitionedDataLoader, ) -from cosmos_framework.data.vfm.action.datasets.action_sft_dataset import get_action_droid_sft_dataset +from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_droid_sft_dataset cs = ConfigStore.instance() diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py index 5df504ce..8cec6a0b 100644 --- a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py @@ -14,8 +14,8 @@ from hydra.core.config_store import ConfigStore from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG -from cosmos_framework.data.vfm.action.datasets.action_sft_dataset import get_action_libero_sft_dataset -from cosmos_framework.data.vfm.joint_dataloader import ( +from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_libero_sft_dataset +from cosmos_framework.data.generator.joint_dataloader import ( PackingDataLoader, RankPartitionedDataLoader, ) diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py index 2be73ab9..594c3460 100644 --- a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py @@ -14,8 +14,8 @@ from hydra.core.config_store import ConfigStore from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG -from cosmos_framework.data.vfm.action.datasets.action_sft_dataset import get_action_libero_sft_dataset -from cosmos_framework.data.vfm.joint_dataloader import ( +from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_libero_sft_dataset +from cosmos_framework.data.generator.joint_dataloader import ( PackingDataLoader, RankPartitionedDataLoader, ) diff --git a/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py b/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py index cd7b6a3f..dd473be5 100644 --- a/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py +++ b/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py @@ -12,7 +12,7 @@ create_qwen2_tokenizer_with_download, create_vlm_config, ) -from cosmos_framework.model.vfm.mot.unified_mot import Qwen3VLMoTConfig, Qwen3VLTextForCausalLM +from cosmos_framework.model.generator.mot.unified_mot import Qwen3VLMoTConfig, Qwen3VLTextForCausalLM from cosmos_framework.utils.lazy_config import LazyCall as L NANO_MODEL_CONFIG = dict( @@ -124,7 +124,7 @@ config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( json_file=( - "cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/" + "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/" "Qwen3-VL-8B-Instruct.json" ), ), diff --git a/cosmos_framework/configs/base/experiment/sft/models/super_model_config.py b/cosmos_framework/configs/base/experiment/sft/models/super_model_config.py index a10cee1a..a4f65298 100644 --- a/cosmos_framework/configs/base/experiment/sft/models/super_model_config.py +++ b/cosmos_framework/configs/base/experiment/sft/models/super_model_config.py @@ -30,7 +30,7 @@ create_qwen2_tokenizer_with_download, create_vlm_config, ) -from cosmos_framework.model.vfm.mot.unified_mot import Qwen3VLMoTConfig, Qwen3VLTextForCausalLM +from cosmos_framework.model.generator.mot.unified_mot import Qwen3VLMoTConfig, Qwen3VLTextForCausalLM SUPER_MODEL_CONFIG = dict( @@ -146,7 +146,7 @@ config=L(create_vlm_config)( base_config=L(Qwen3VLMoTConfig.from_json_file)( json_file=( - "cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/" + "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/" "Qwen3-VL-32B-Instruct.json" ), ), diff --git a/cosmos_framework/configs/base/experiment/sft/vision_sft_nano.py b/cosmos_framework/configs/base/experiment/sft/vision_sft_nano.py index 46957d19..6597ffb2 100644 --- a/cosmos_framework/configs/base/experiment/sft/vision_sft_nano.py +++ b/cosmos_framework/configs/base/experiment/sft/vision_sft_nano.py @@ -35,18 +35,18 @@ from hydra.core.config_store import ConfigStore from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG -from cosmos_framework.data.vfm.joint_dataloader import ( +from cosmos_framework.data.generator.joint_dataloader import ( PackingDataLoader, RankPartitionedDataLoader, ) -from cosmos_framework.data.vfm.dataflow import ( +from cosmos_framework.data.generator.dataflow import ( CosmosDataLoader, IdentityProcessor, RankPartitionedDistributor, SequentialPackingBatcher, VFMListCollator, ) -from cosmos_framework.data.vfm.local_datasets.sft_dataset import get_sft_dataset +from cosmos_framework.data.generator.local_datasets.sft_dataset import get_sft_dataset from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.utils.lazy_config import LazyDict @@ -61,7 +61,7 @@ {"override /data_val": None}, {"override /optimizer": "adamw"}, # YAML used `scheduler: warmup_cosine_lr` but that group is only - # registered in cosmos_framework/configs/base/vlm/defaults/optimizer.py + # registered in cosmos_framework/configs/base/reasoner/defaults/optimizer.py # (reachable from the vlm config tree). The base vfm config path # only knows `lambdacosine`, which also sets # lr_scheduler_type="LambdaCosine" — behaviorally identical. diff --git a/cosmos_framework/configs/base/experiment/sft/vision_sft_super.py b/cosmos_framework/configs/base/experiment/sft/vision_sft_super.py index a49bb3d9..7c18dfbf 100644 --- a/cosmos_framework/configs/base/experiment/sft/vision_sft_super.py +++ b/cosmos_framework/configs/base/experiment/sft/vision_sft_super.py @@ -51,11 +51,11 @@ from hydra.core.config_store import ConfigStore from cosmos_framework.configs.base.experiment.sft.models.super_model_config import SUPER_MODEL_CONFIG -from cosmos_framework.data.vfm.joint_dataloader import ( +from cosmos_framework.data.generator.joint_dataloader import ( PackingDataLoader, RankPartitionedDataLoader, ) -from cosmos_framework.data.vfm.local_datasets.sft_dataset import get_sft_dataset +from cosmos_framework.data.generator.local_datasets.sft_dataset import get_sft_dataset from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.utils.lazy_config import LazyDict @@ -70,7 +70,7 @@ {"override /data_val": None}, {"override /optimizer": "adamw"}, # YAML used `scheduler: warmup_cosine_lr` but that group is only - # registered in cosmos_framework/configs/base/vlm/defaults/optimizer.py + # registered in cosmos_framework/configs/base/reasoner/defaults/optimizer.py # (reachable from the vlm config tree). The base vfm config path # only knows `lambdacosine`, which also sets # lr_scheduler_type="LambdaCosine" — behaviorally identical. diff --git a/cosmos_framework/configs/base/reasoner/defaults/model.py b/cosmos_framework/configs/base/reasoner/defaults/model.py index 0c991f08..8d7b5415 100644 --- a/cosmos_framework/configs/base/reasoner/defaults/model.py +++ b/cosmos_framework/configs/base/reasoner/defaults/model.py @@ -6,7 +6,7 @@ from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig from cosmos_framework.configs.base.reasoner.defaults.policy_config import VLMModelConfig -from cosmos_framework.model.vfm.vlm_model import VLMModel +from cosmos_framework.model.generator.vlm_model import VLMModel VLM_FSDP_CONFIG = dict( trainer=dict( diff --git a/cosmos_framework/configs/base/reasoner/experiment/dataflow_roles.py b/cosmos_framework/configs/base/reasoner/experiment/dataflow_roles.py index 38281827..cffceb80 100644 --- a/cosmos_framework/configs/base/reasoner/experiment/dataflow_roles.py +++ b/cosmos_framework/configs/base/reasoner/experiment/dataflow_roles.py @@ -11,8 +11,8 @@ import torch from torch.utils.data._utils.collate import default_collate -from cosmos_framework.data.vfm.dataflow.base import BatchCollator, RawItemProcessor -from cosmos_framework.utils.vlm.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD +from cosmos_framework.data.generator.dataflow.base import BatchCollator, RawItemProcessor +from cosmos_framework.utils.reasoner.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD class VLMProcessor(RawItemProcessor): diff --git a/cosmos_framework/configs/base/reasoner/experiment/llava_ov_vlm.py b/cosmos_framework/configs/base/reasoner/experiment/llava_ov_vlm.py index 1510ea1b..eef59003 100644 --- a/cosmos_framework/configs/base/reasoner/experiment/llava_ov_vlm.py +++ b/cosmos_framework/configs/base/reasoner/experiment/llava_ov_vlm.py @@ -35,7 +35,7 @@ Usage (smoke test):: torchrun --nproc_per_node=4 --master_port=12344 -m cosmos_framework.scripts.train \\ - --config=cosmos_framework/configs/base/vlm/config.py -- \\ + --config=cosmos_framework/configs/base/reasoner/config.py -- \\ experiment=pre_exp012_llava_ov \\ "model.config.policy.backbone.model_name=/path/to/Siglip2-Qwen3-1.7B-BF16-Alignment" \\ trainer.max_iter=10 trainer.logging_iter=1 \\ @@ -53,14 +53,14 @@ from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.utils.lazy_config import LazyDict -from cosmos_framework.data.vfm.dataflow import ( +from cosmos_framework.data.generator.dataflow import ( CosmosDataLoader, IterableDistributor, MapDistributor, PoolPackingBatcher, ) -from cosmos_framework.data.vfm.processors import build_processor -from cosmos_framework.utils.vlm.constant import IGNORE_INDEX +from cosmos_framework.data.generator.processors import build_processor +from cosmos_framework.utils.reasoner.constant import IGNORE_INDEX from cosmos_framework.configs.base.reasoner.experiment.dataflow_roles import VLMProcessor, VLMCollator from cosmos_framework.callbacks.cosmos_dataloader_state import CosmosDataLoaderStateCallback diff --git a/cosmos_framework/configs/base/reasoner/experiment/videophy2_dataflow_roles.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_dataflow_roles.py index 813ed301..0d0a548d 100644 --- a/cosmos_framework/configs/base/reasoner/experiment/videophy2_dataflow_roles.py +++ b/cosmos_framework/configs/base/reasoner/experiment/videophy2_dataflow_roles.py @@ -8,8 +8,8 @@ import io from typing import Any -from cosmos_framework.data.vfm.dataflow.base import RawItemProcessor -from cosmos_framework.utils.vlm.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD +from cosmos_framework.data.generator.dataflow.base import RawItemProcessor +from cosmos_framework.utils.reasoner.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD _MAX_VIDEO_FRAMES = 32 _TARGET_VIDEO_FPS = 2.0 diff --git a/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py index baa438af..1405ac7f 100644 --- a/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py +++ b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py @@ -19,12 +19,12 @@ from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.utils.lazy_config import LazyDict -from cosmos_framework.data.vfm.dataflow import CosmosDataLoader, IterableDistributor, PoolPackingBatcher -from cosmos_framework.data.vfm.processors import build_processor -from cosmos_framework.data.vlm.local_sft_dataset import LocalSFTDataset -from cosmos_framework.data.vlm.data_sources_videophy2.videophy2 import DATAINFO +from cosmos_framework.data.generator.dataflow import CosmosDataLoader, IterableDistributor, PoolPackingBatcher +from cosmos_framework.data.generator.processors import build_processor +from cosmos_framework.data.reasoner.local_sft_dataset import LocalSFTDataset +from cosmos_framework.data.reasoner.data_sources_videophy2.videophy2 import DATAINFO from cosmos_framework.utils import log -from cosmos_framework.utils.vlm.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD +from cosmos_framework.utils.reasoner.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD from cosmos_framework.configs.base.reasoner.experiment.dataflow_roles import VLMCollator from cosmos_framework.configs.base.reasoner.experiment.videophy2_dataflow_roles import VideoPhy2Processor diff --git a/cosmos_framework/configs/toml_config/sft_config.py b/cosmos_framework/configs/toml_config/sft_config.py index 5d17f2b2..9ce57728 100644 --- a/cosmos_framework/configs/toml_config/sft_config.py +++ b/cosmos_framework/configs/toml_config/sft_config.py @@ -39,7 +39,7 @@ class JobConfig(BaseModel): description=( "META — chooses which make_config() to call: " "'vfm' → cosmos_framework/configs/base/config.py (video foundation model), " - "'vlm' → cosmos_framework/configs/base/vlm/config.py (vision-language model). " + "'vlm' → cosmos_framework/configs/base/reasoner/config.py (vision-language model). " "Also picks the path-remap rules in toml_config_helper.PATH_REMAPS." ), ) @@ -692,7 +692,7 @@ def load_experiment_from_toml( The base config module is picked from ``[job].task`` in the TOML: - ``task = "vfm"`` → ``cosmos_framework/configs/base/config.py`` - - ``task = "vlm"`` → ``cosmos_framework/configs/base/vlm/config.py`` + - ``task = "vlm"`` → ``cosmos_framework/configs/base/reasoner/config.py`` ``extra_overrides`` is appended after the TOML-derived Hydra overrides, so command-line entries take precedence over TOML values. Each entry must be diff --git a/cosmos_framework/data/vfm/__init__.py b/cosmos_framework/data/generator/__init__.py similarity index 100% rename from cosmos_framework/data/vfm/__init__.py rename to cosmos_framework/data/generator/__init__.py diff --git a/cosmos_framework/data/vfm/action/__init__.py b/cosmos_framework/data/generator/action/__init__.py similarity index 100% rename from cosmos_framework/data/vfm/action/__init__.py rename to cosmos_framework/data/generator/action/__init__.py diff --git a/cosmos_framework/data/vfm/action/action_normalization.py b/cosmos_framework/data/generator/action/action_normalization.py similarity index 100% rename from cosmos_framework/data/vfm/action/action_normalization.py rename to cosmos_framework/data/generator/action/action_normalization.py diff --git a/cosmos_framework/data/vfm/action/action_normalization_test.py b/cosmos_framework/data/generator/action/action_normalization_test.py similarity index 98% rename from cosmos_framework/data/vfm/action/action_normalization_test.py rename to cosmos_framework/data/generator/action/action_normalization_test.py index bbfff16d..378665e9 100644 --- a/cosmos_framework/data/vfm/action/action_normalization_test.py +++ b/cosmos_framework/data/generator/action/action_normalization_test.py @@ -7,7 +7,7 @@ import pytest import torch -from cosmos_framework.data.vfm.action.action_normalization import ( +from cosmos_framework.data.generator.action.action_normalization import ( denormalize_action, load_action_stats, normalize_action, diff --git a/cosmos_framework/data/vfm/action/action_processing.py b/cosmos_framework/data/generator/action/action_processing.py similarity index 100% rename from cosmos_framework/data/vfm/action/action_processing.py rename to cosmos_framework/data/generator/action/action_processing.py diff --git a/cosmos_framework/data/vfm/action/action_spec.py b/cosmos_framework/data/generator/action/action_spec.py similarity index 99% rename from cosmos_framework/data/vfm/action/action_spec.py rename to cosmos_framework/data/generator/action/action_spec.py index c9c0f384..07e15e7b 100644 --- a/cosmos_framework/data/vfm/action/action_spec.py +++ b/cosmos_framework/data/generator/action/action_spec.py @@ -34,7 +34,7 @@ from enum import Enum from typing import ClassVar -from cosmos_framework.data.vfm.action.pose_utils import ( +from cosmos_framework.data.generator.action.pose_utils import ( RotationConvention, _identity_rotation_vector, ) diff --git a/cosmos_framework/data/vfm/action/agibot_fk.py b/cosmos_framework/data/generator/action/agibot_fk.py similarity index 99% rename from cosmos_framework/data/vfm/action/agibot_fk.py rename to cosmos_framework/data/generator/action/agibot_fk.py index 10382844..56651565 100644 --- a/cosmos_framework/data/vfm/action/agibot_fk.py +++ b/cosmos_framework/data/generator/action/agibot_fk.py @@ -10,7 +10,7 @@ import numpy as np -from cosmos_framework.data.vfm.action.agibot_spec import ( +from cosmos_framework.data.generator.action.agibot_spec import ( AGIBOT_WORLD_ARM_JOINT_NAMES_LEFT, AGIBOT_WORLD_ARM_JOINT_NAMES_RIGHT, AGIBOT_WORLD_ARM_STATE_SLICE, @@ -42,7 +42,7 @@ get_agibot_world_kind_spec, get_agibot_world_urdf_path, ) -from cosmos_framework.data.vfm.action.pose_utils import convert_rotation +from cosmos_framework.data.generator.action.pose_utils import convert_rotation _GRIPPER_VALUE_EPS = 1e-4 _QUATERNION_NORM_EPS = 1e-8 diff --git a/cosmos_framework/data/vfm/action/agibot_spec.py b/cosmos_framework/data/generator/action/agibot_spec.py similarity index 100% rename from cosmos_framework/data/vfm/action/agibot_spec.py rename to cosmos_framework/data/generator/action/agibot_spec.py diff --git a/cosmos_framework/data/generator/action/datasets/__init__.py b/cosmos_framework/data/generator/action/datasets/__init__.py new file mode 100644 index 00000000..62437797 --- /dev/null +++ b/cosmos_framework/data/generator/action/datasets/__init__.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Action dataset wrappers for Cosmos Action. + +All concrete datasets inherit from :class:`ActionBaseDataset` and expose a +``load_action_stats()`` classmethod for retrieving pre-computed normalization +statistics without instantiating the dataset. +""" + +from cosmos_framework.data.generator.action.datasets.agibotworld_beta_lerobot_dataset import AgiBotWorldBetaLeRobotDataset +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.datasets.bridge_orig_lerobot_dataset import BridgeOrigLeRobotDataset +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset +from cosmos_framework.data.generator.action.datasets.fractal_lerobot_dataset import FractalLeRobotDataset +from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset +from cosmos_framework.data.generator.action.datasets.robomind_franka_dataset import RoboMINDFrankaDataset +from cosmos_framework.data.generator.action.datasets.robomind_ur_dataset import RoboMINDURDataset +from cosmos_framework.data.generator.action.datasets.umi_lerobot_dataset import UMILeRobotDataset + +__all__ = [ + "ActionBaseDataset", + "AgiBotWorldBetaLeRobotDataset", + "BridgeOrigLeRobotDataset", + "DROIDLeRobotDataset", + "FractalLeRobotDataset", + "LIBEROLeRobotDataset", + "RoboMINDFrankaDataset", + "RoboMINDURDataset", + "UMILeRobotDataset", +] diff --git a/cosmos_framework/data/vfm/action/datasets/action_sft_dataset.py b/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py similarity index 96% rename from cosmos_framework/data/vfm/action/datasets/action_sft_dataset.py rename to cosmos_framework/data/generator/action/datasets/action_sft_dataset.py index a36e4357..8a424bc9 100644 --- a/cosmos_framework/data/vfm/action/datasets/action_sft_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py @@ -19,9 +19,9 @@ from torch.utils.data import Dataset, IterableDataset, get_worker_info -from cosmos_framework.data.vfm.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset -from cosmos_framework.data.vfm.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset -from cosmos_framework.data.vfm.action.transforms import ActionTransformPipeline +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset +from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset +from cosmos_framework.data.generator.action.transforms import ActionTransformPipeline class ActionSFTDataset(Dataset): diff --git a/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/agibotworld_beta_lerobot_dataset.py similarity index 97% rename from cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py rename to cosmos_framework/data/generator/action/datasets/agibotworld_beta_lerobot_dataset.py index 81a9a364..ec4a0e3c 100644 --- a/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/agibotworld_beta_lerobot_dataset.py @@ -12,16 +12,16 @@ import torch import torch.nn.functional as F -from cosmos_framework.data.vfm.action.agibot_fk import ( +from cosmos_framework.data.generator.action.agibot_fk import ( AGIBOT_WORLD_GRIPPER_TO_OPENCV_BY_WRIST, apply_agibot_gripper_to_opencv, apply_robot_base_motion_to_poses, compute_fk_transforms_batch, convert_gripper_state_to_open_fraction, ) -from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec -from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset -from cosmos_framework.data.vfm.action.pose_utils import pose_abs_to_rel +from cosmos_framework.data.generator.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.pose_utils import pose_abs_to_rel PoseConvention = Literal["backward_framewise"] Viewpoint = Literal["concat_view", "ego_view"] diff --git a/cosmos_framework/data/vfm/action/datasets/base_dataset.py b/cosmos_framework/data/generator/action/datasets/base_dataset.py similarity index 95% rename from cosmos_framework/data/vfm/action/datasets/base_dataset.py rename to cosmos_framework/data/generator/action/datasets/base_dataset.py index d7cdcc23..e3d2ec81 100644 --- a/cosmos_framework/data/vfm/action/datasets/base_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/base_dataset.py @@ -17,10 +17,10 @@ import torch from torch.utils.data import Dataset -from cosmos_framework.data.vfm.action.action_normalization import load_action_stats, normalize_action -from cosmos_framework.data.vfm.action.action_spec import ActionSpec -from cosmos_framework.data.vfm.action.domain_utils import get_domain_id -from cosmos_framework.data.vfm.action.pose_utils import compute_idle_frames +from cosmos_framework.data.generator.action.action_normalization import load_action_stats, normalize_action +from cosmos_framework.data.generator.action.action_spec import ActionSpec +from cosmos_framework.data.generator.action.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.pose_utils import compute_idle_frames _MODE_CHOICES = ("forward_dynamics", "inverse_dynamics", "policy") diff --git a/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/bridge_orig_lerobot_dataset.py similarity index 95% rename from cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py rename to cosmos_framework/data/generator/action/datasets/bridge_orig_lerobot_dataset.py index e6ea9d2b..50e604b3 100644 --- a/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/bridge_orig_lerobot_dataset.py @@ -12,9 +12,9 @@ import numpy as np import torch -from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec -from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset -from cosmos_framework.data.vfm.action.pose_utils import ( +from cosmos_framework.data.generator.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.pose_utils import ( build_abs_pose_from_components, pose_abs_to_rel, ) diff --git a/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py similarity index 98% rename from cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py rename to cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py index 5d31a999..7e432ea3 100644 --- a/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py @@ -16,9 +16,9 @@ import torch.nn.functional as F import torchvision.transforms as T -from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Joint, Pos, Rot, build_action_spec -from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset -from cosmos_framework.data.vfm.action.pose_utils import ( +from cosmos_framework.data.generator.action.action_spec import ActionSpec, Gripper, Joint, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.pose_utils import ( build_abs_pose_from_components, pose_abs_to_rel, ) diff --git a/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/fractal_lerobot_dataset.py similarity index 96% rename from cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py rename to cosmos_framework/data/generator/action/datasets/fractal_lerobot_dataset.py index b9dd23d7..99bd5461 100644 --- a/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/fractal_lerobot_dataset.py @@ -19,9 +19,9 @@ import numpy as np import torch -from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec -from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset -from cosmos_framework.data.vfm.action.pose_utils import ( +from cosmos_framework.data.generator.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.pose_utils import ( build_abs_pose_from_components, pose_abs_to_rel, ) diff --git a/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/libero_lerobot_dataset.py similarity index 97% rename from cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py rename to cosmos_framework/data/generator/action/datasets/libero_lerobot_dataset.py index fe31f9d5..a9f0a313 100644 --- a/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/libero_lerobot_dataset.py @@ -32,10 +32,10 @@ import torch import torch.nn.functional as F -from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec -from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset -from cosmos_framework.data.vfm.action.libero_pose_utils import libero_action_dim, libero_rotation_format -from cosmos_framework.data.vfm.action.pose_utils import convert_rotation +from cosmos_framework.data.generator.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.libero_pose_utils import libero_action_dim, libero_rotation_format +from cosmos_framework.data.generator.action.pose_utils import convert_rotation from cosmos_framework.utils import log CameraMode = Literal["image", "wrist_image", "concat_view"] diff --git a/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py b/cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py similarity index 96% rename from cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py rename to cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py index 7bda3053..615d5fce 100644 --- a/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py @@ -13,10 +13,10 @@ import torch import torch.nn.functional as F -from cosmos_framework.data.vfm.action.action_normalization import load_action_stats -from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec -from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset -from cosmos_framework.data.vfm.action.pose_utils import ( +from cosmos_framework.data.generator.action.action_normalization import load_action_stats +from cosmos_framework.data.generator.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.pose_utils import ( build_abs_pose_from_components, pose_abs_to_rel, ) diff --git a/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py b/cosmos_framework/data/generator/action/datasets/robomind_ur_dataset.py similarity index 96% rename from cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py rename to cosmos_framework/data/generator/action/datasets/robomind_ur_dataset.py index ab7bb433..5e25e12c 100644 --- a/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/robomind_ur_dataset.py @@ -12,9 +12,9 @@ import numpy as np import torch -from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec -from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset -from cosmos_framework.data.vfm.action.pose_utils import pose_abs_to_rel +from cosmos_framework.data.generator.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.pose_utils import pose_abs_to_rel PoseConvention = Literal["backward_framewise"] Viewpoint = Literal["third_person_view"] diff --git a/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/umi_lerobot_dataset.py similarity index 95% rename from cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py rename to cosmos_framework/data/generator/action/datasets/umi_lerobot_dataset.py index 61a31497..be88b214 100644 --- a/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/umi_lerobot_dataset.py @@ -12,10 +12,10 @@ import numpy as np import torch -from cosmos_framework.data.vfm.action.action_normalization import load_action_stats -from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec -from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset -from cosmos_framework.data.vfm.action.pose_utils import ( +from cosmos_framework.data.generator.action.action_normalization import load_action_stats +from cosmos_framework.data.generator.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec +from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +from cosmos_framework.data.generator.action.pose_utils import ( build_abs_pose_from_components, pose_abs_to_rel, ) diff --git a/cosmos_framework/data/vfm/action/domain_utils.py b/cosmos_framework/data/generator/action/domain_utils.py similarity index 100% rename from cosmos_framework/data/vfm/action/domain_utils.py rename to cosmos_framework/data/generator/action/domain_utils.py diff --git a/cosmos_framework/data/vfm/action/json_formatter.py b/cosmos_framework/data/generator/action/json_formatter.py similarity index 98% rename from cosmos_framework/data/vfm/action/json_formatter.py rename to cosmos_framework/data/generator/action/json_formatter.py index 201a76e5..d51c8da6 100644 --- a/cosmos_framework/data/vfm/action/json_formatter.py +++ b/cosmos_framework/data/generator/action/json_formatter.py @@ -8,8 +8,8 @@ import torch from cosmos_framework.utils import log -from cosmos_framework.data.vfm.action.viewpoint_utils import DEFAULT_VIEWPOINT_TEMPLATES -from cosmos_framework.data.vfm.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.data.generator.action.viewpoint_utils import DEFAULT_VIEWPOINT_TEMPLATES +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO def _should_append_idle_frame_info(mode: object) -> bool: diff --git a/cosmos_framework/data/vfm/action/libero_pose_utils.py b/cosmos_framework/data/generator/action/libero_pose_utils.py similarity index 97% rename from cosmos_framework/data/vfm/action/libero_pose_utils.py rename to cosmos_framework/data/generator/action/libero_pose_utils.py index 3a4fd8e2..a6552c7b 100644 --- a/cosmos_framework/data/vfm/action/libero_pose_utils.py +++ b/cosmos_framework/data/generator/action/libero_pose_utils.py @@ -8,7 +8,7 @@ import numpy as np import torch -from cosmos_framework.data.vfm.action.pose_utils import ( +from cosmos_framework.data.generator.action.pose_utils import ( RotationConvention, build_abs_pose_from_components, ) diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/agibotworld_beta_lerobot_stats.json b/cosmos_framework/data/generator/action/normalizer_stats/agibotworld_beta_lerobot_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/normalizer_stats/agibotworld_beta_lerobot_stats.json rename to cosmos_framework/data/generator/action/normalizer_stats/agibotworld_beta_lerobot_stats.json diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/bridge_orig_lerobot_stats.json b/cosmos_framework/data/generator/action/normalizer_stats/bridge_orig_lerobot_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/normalizer_stats/bridge_orig_lerobot_stats.json rename to cosmos_framework/data/generator/action/normalizer_stats/bridge_orig_lerobot_stats.json diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/droid_lerobot_stats.json b/cosmos_framework/data/generator/action/normalizer_stats/droid_lerobot_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/normalizer_stats/droid_lerobot_stats.json rename to cosmos_framework/data/generator/action/normalizer_stats/droid_lerobot_stats.json diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/fractal_lerobot_stats.json b/cosmos_framework/data/generator/action/normalizer_stats/fractal_lerobot_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/normalizer_stats/fractal_lerobot_stats.json rename to cosmos_framework/data/generator/action/normalizer_stats/fractal_lerobot_stats.json diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json b/cosmos_framework/data/generator/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json similarity index 100% rename from cosmos_framework/data/vfm/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json rename to cosmos_framework/data/generator/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/robomind_franka_dual_stats.json b/cosmos_framework/data/generator/action/normalizer_stats/robomind_franka_dual_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/normalizer_stats/robomind_franka_dual_stats.json rename to cosmos_framework/data/generator/action/normalizer_stats/robomind_franka_dual_stats.json diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/robomind_franka_stats.json b/cosmos_framework/data/generator/action/normalizer_stats/robomind_franka_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/normalizer_stats/robomind_franka_stats.json rename to cosmos_framework/data/generator/action/normalizer_stats/robomind_franka_stats.json diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/robomind_ur_stats.json b/cosmos_framework/data/generator/action/normalizer_stats/robomind_ur_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/normalizer_stats/robomind_ur_stats.json rename to cosmos_framework/data/generator/action/normalizer_stats/robomind_ur_stats.json diff --git a/cosmos_framework/data/vfm/action/normalizer_stats/umi_lerobot_stats.json b/cosmos_framework/data/generator/action/normalizer_stats/umi_lerobot_stats.json similarity index 100% rename from cosmos_framework/data/vfm/action/normalizer_stats/umi_lerobot_stats.json rename to cosmos_framework/data/generator/action/normalizer_stats/umi_lerobot_stats.json diff --git a/cosmos_framework/data/vfm/action/pose_utils.py b/cosmos_framework/data/generator/action/pose_utils.py similarity index 99% rename from cosmos_framework/data/vfm/action/pose_utils.py rename to cosmos_framework/data/generator/action/pose_utils.py index d6d26a4e..4780135e 100644 --- a/cosmos_framework/data/vfm/action/pose_utils.py +++ b/cosmos_framework/data/generator/action/pose_utils.py @@ -685,7 +685,7 @@ def compute_idle_frames( # Import locally to avoid a circular import at module load time # (action_spec.py imports RotationConvention from this file). - from cosmos_framework.data.vfm.action.action_spec import DimType + from cosmos_framework.data.generator.action.action_spec import DimType pos_idx = [i for i, t in enumerate(spec.types) if t == DimType.POS] rot_idx = [i for i, t in enumerate(spec.types) if t == DimType.ROT] diff --git a/cosmos_framework/data/vfm/action/pose_utils_test.py b/cosmos_framework/data/generator/action/pose_utils_test.py similarity index 99% rename from cosmos_framework/data/vfm/action/pose_utils_test.py rename to cosmos_framework/data/generator/action/pose_utils_test.py index 3cdc20bb..08fe9c0c 100644 --- a/cosmos_framework/data/vfm/action/pose_utils_test.py +++ b/cosmos_framework/data/generator/action/pose_utils_test.py @@ -6,7 +6,7 @@ import torch from scipy.spatial.transform import Rotation as R -from cosmos_framework.data.vfm.action.pose_utils import ( +from cosmos_framework.data.generator.action.pose_utils import ( _normalize_rotation_matrices, _to_numpy_float32, build_abs_pose_from_components, diff --git a/cosmos_framework/data/vfm/action/robot_assets/G1_omnipicker_calibrated.urdf b/cosmos_framework/data/generator/action/robot_assets/G1_omnipicker_calibrated.urdf similarity index 100% rename from cosmos_framework/data/vfm/action/robot_assets/G1_omnipicker_calibrated.urdf rename to cosmos_framework/data/generator/action/robot_assets/G1_omnipicker_calibrated.urdf diff --git a/cosmos_framework/data/vfm/action/robot_assets/ur5e_robotiq_2f85.xml b/cosmos_framework/data/generator/action/robot_assets/ur5e_robotiq_2f85.xml similarity index 100% rename from cosmos_framework/data/vfm/action/robot_assets/ur5e_robotiq_2f85.xml rename to cosmos_framework/data/generator/action/robot_assets/ur5e_robotiq_2f85.xml diff --git a/cosmos_framework/data/vfm/action/transforms.py b/cosmos_framework/data/generator/action/transforms.py similarity index 97% rename from cosmos_framework/data/vfm/action/transforms.py rename to cosmos_framework/data/generator/action/transforms.py index 96f4f0ba..6417a6e9 100644 --- a/cosmos_framework/data/vfm/action/transforms.py +++ b/cosmos_framework/data/generator/action/transforms.py @@ -20,19 +20,19 @@ import torchvision.transforms.functional as transforms_F from cosmos_framework.utils import log -from cosmos_framework.data.vfm.action.action_processing import ( +from cosmos_framework.data.generator.action.action_processing import ( ActionNormalizer, ActionProcessor, ) -from cosmos_framework.data.vfm.action.json_formatter import ActionPromptJsonFormatter -from cosmos_framework.data.vfm.action.viewpoint_utils import ViewpointTextInfo -from cosmos_framework.data.vfm.augmentors.duration_fps_text_timestamps import DurationFPSTextTimeStamps -from cosmos_framework.data.vfm.augmentors.idle_frames_text_info import IdleFramesTextInfo -from cosmos_framework.data.vfm.augmentors.resolution_text_info import ResolutionTextInfo -from cosmos_framework.data.vfm.augmentors.text_tokenizer import TextTokenizerTransform -from cosmos_framework.data.vfm.utils import VIDEO_RES_SIZE_INFO -from cosmos_framework.data.vfm.sequence_packing import SequencePlan -from cosmos_framework.utils.vfm.data_utils import get_vision_data_resolution +from cosmos_framework.data.generator.action.json_formatter import ActionPromptJsonFormatter +from cosmos_framework.data.generator.action.viewpoint_utils import ViewpointTextInfo +from cosmos_framework.data.generator.augmentors.duration_fps_text_timestamps import DurationFPSTextTimeStamps +from cosmos_framework.data.generator.augmentors.idle_frames_text_info import IdleFramesTextInfo +from cosmos_framework.data.generator.augmentors.resolution_text_info import ResolutionTextInfo +from cosmos_framework.data.generator.augmentors.text_tokenizer import TextTokenizerTransform +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.data.generator.sequence_packing import SequencePlan +from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution def _should_append_idle_frame_info(mode: object) -> bool: diff --git a/cosmos_framework/data/vfm/action/transforms_test.py b/cosmos_framework/data/generator/action/transforms_test.py similarity index 96% rename from cosmos_framework/data/vfm/action/transforms_test.py rename to cosmos_framework/data/generator/action/transforms_test.py index 759ec21d..93b62f05 100644 --- a/cosmos_framework/data/vfm/action/transforms_test.py +++ b/cosmos_framework/data/generator/action/transforms_test.py @@ -8,14 +8,14 @@ import pytest import torch -from cosmos_framework.data.vfm.action.json_formatter import ActionPromptJsonFormatter -from cosmos_framework.data.vfm.action.transforms import ( +from cosmos_framework.data.generator.action.json_formatter import ActionPromptJsonFormatter +from cosmos_framework.data.generator.action.transforms import ( ActionTransformPipeline, reflection_pad_to_target, remove_reflection_padding, ) -from cosmos_framework.data.vfm.augmentors.duration_fps_text_timestamps import DurationFPSTextTimeStamps -from cosmos_framework.data.vfm.augmentors.resolution_text_info import ResolutionTextInfo +from cosmos_framework.data.generator.augmentors.duration_fps_text_timestamps import DurationFPSTextTimeStamps +from cosmos_framework.data.generator.augmentors.resolution_text_info import ResolutionTextInfo @pytest.mark.L0 diff --git a/cosmos_framework/data/vfm/action/viewpoint_utils.py b/cosmos_framework/data/generator/action/viewpoint_utils.py similarity index 100% rename from cosmos_framework/data/vfm/action/viewpoint_utils.py rename to cosmos_framework/data/generator/action/viewpoint_utils.py diff --git a/cosmos_framework/data/vfm/augmentor_provider.py b/cosmos_framework/data/generator/augmentor_provider.py similarity index 97% rename from cosmos_framework/data/vfm/augmentor_provider.py rename to cosmos_framework/data/generator/augmentor_provider.py index fa859e42..8840d160 100644 --- a/cosmos_framework/data/vfm/augmentor_provider.py +++ b/cosmos_framework/data/generator/augmentor_provider.py @@ -7,24 +7,24 @@ import cosmos_framework.data.imaginaire.webdataset.augmentors.image.normalize as normalize import cosmos_framework.data.imaginaire.webdataset.augmentors.image.padding as padding import cosmos_framework.data.imaginaire.webdataset.augmentors.image.resize as resize -import cosmos_framework.data.vfm.augmentors.append_fps_frames_for_image as append_fps_frames_for_image -import cosmos_framework.data.vfm.augmentors.audio_caption as audio_caption -import cosmos_framework.data.vfm.augmentors.caption_filter as caption_filter -import cosmos_framework.data.vfm.augmentors.cropping as cosmos_cropping -import cosmos_framework.data.vfm.augmentors.duration_fps_text_timestamps as duration_fps_text_timestamps -import cosmos_framework.data.vfm.augmentors.image_resolution_filter as image_resolution_filter -import cosmos_framework.data.vfm.augmentors.merge_datadict as merge_datadict -import cosmos_framework.data.vfm.augmentors.resolution_text_info as resolution_text_info -import cosmos_framework.data.vfm.augmentors.sound_sequence_plan as sound_sequence_plan -import cosmos_framework.data.vfm.augmentors.text_tokenizer as text_tokenizer -import cosmos_framework.data.vfm.augmentors.text_transforms_for_image as text_transforms_for_image -import cosmos_framework.data.vfm.augmentors.text_transforms_for_video as text_transforms_for_video -import cosmos_framework.data.vfm.augmentors.video_parsing as video_parsing +import cosmos_framework.data.generator.augmentors.append_fps_frames_for_image as append_fps_frames_for_image +import cosmos_framework.data.generator.augmentors.audio_caption as audio_caption +import cosmos_framework.data.generator.augmentors.caption_filter as caption_filter +import cosmos_framework.data.generator.augmentors.cropping as cosmos_cropping +import cosmos_framework.data.generator.augmentors.duration_fps_text_timestamps as duration_fps_text_timestamps +import cosmos_framework.data.generator.augmentors.image_resolution_filter as image_resolution_filter +import cosmos_framework.data.generator.augmentors.merge_datadict as merge_datadict +import cosmos_framework.data.generator.augmentors.resolution_text_info as resolution_text_info +import cosmos_framework.data.generator.augmentors.sound_sequence_plan as sound_sequence_plan +import cosmos_framework.data.generator.augmentors.text_tokenizer as text_tokenizer +import cosmos_framework.data.generator.augmentors.text_transforms_for_image as text_transforms_for_image +import cosmos_framework.data.generator.augmentors.text_transforms_for_video as text_transforms_for_video +import cosmos_framework.data.generator.augmentors.video_parsing as video_parsing from cosmos_framework.utils.lazy_config import LazyCall as L from cosmos_framework.utils.lazy_config import LazyDict from cosmos_framework.utils import log -from cosmos_framework.data.vfm.augmentors import sequence_plan -from cosmos_framework.data.vfm.utils import IMAGE_RES_SIZE_INFO, VIDEO_RES_SIZE_INFO +from cosmos_framework.data.generator.augmentors import sequence_plan +from cosmos_framework.data.generator.utils import IMAGE_RES_SIZE_INFO, VIDEO_RES_SIZE_INFO # UniAE requires spatial dimensions divisible by (spatial_compression * patch_spatial) = 16 * 2 = 32. UNIAE_SPATIAL_MULTIPLE = 32 diff --git a/cosmos_framework/data/vfm/augmentors/__init__.py b/cosmos_framework/data/generator/augmentors/__init__.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/__init__.py rename to cosmos_framework/data/generator/augmentors/__init__.py diff --git a/cosmos_framework/data/vfm/augmentors/append_fps_frames_for_image.py b/cosmos_framework/data/generator/augmentors/append_fps_frames_for_image.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/append_fps_frames_for_image.py rename to cosmos_framework/data/generator/augmentors/append_fps_frames_for_image.py diff --git a/cosmos_framework/data/vfm/augmentors/audio_caption.py b/cosmos_framework/data/generator/augmentors/audio_caption.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/audio_caption.py rename to cosmos_framework/data/generator/augmentors/audio_caption.py diff --git a/cosmos_framework/data/vfm/augmentors/audio_parsing.py b/cosmos_framework/data/generator/augmentors/audio_parsing.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/audio_parsing.py rename to cosmos_framework/data/generator/augmentors/audio_parsing.py diff --git a/cosmos_framework/data/vfm/augmentors/caption_filter.py b/cosmos_framework/data/generator/augmentors/caption_filter.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/caption_filter.py rename to cosmos_framework/data/generator/augmentors/caption_filter.py diff --git a/cosmos_framework/data/vfm/augmentors/cropping.py b/cosmos_framework/data/generator/augmentors/cropping.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/cropping.py rename to cosmos_framework/data/generator/augmentors/cropping.py diff --git a/cosmos_framework/data/vfm/augmentors/duration_fps_text_timestamps.py b/cosmos_framework/data/generator/augmentors/duration_fps_text_timestamps.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/duration_fps_text_timestamps.py rename to cosmos_framework/data/generator/augmentors/duration_fps_text_timestamps.py diff --git a/cosmos_framework/data/vfm/augmentors/idle_frames_text_info.py b/cosmos_framework/data/generator/augmentors/idle_frames_text_info.py similarity index 98% rename from cosmos_framework/data/vfm/augmentors/idle_frames_text_info.py rename to cosmos_framework/data/generator/augmentors/idle_frames_text_info.py index 34f6116f..59b6f13a 100644 --- a/cosmos_framework/data/vfm/augmentors/idle_frames_text_info.py +++ b/cosmos_framework/data/generator/augmentors/idle_frames_text_info.py @@ -8,7 +8,7 @@ frames (i.e. the relative-pose delta is close to identity and the gripper command does not change). The upstream dataset is responsible for populating ``data_dict[idle_frames_key]`` via -:func:`cosmos_framework.data.vfm.action.pose_utils.compute_idle_frames`. +:func:`cosmos_framework.data.generator.action.pose_utils.compute_idle_frames`. Per-field dropout (default 5%) is applied here, matching Pi0.7's approach of independently dropping each metadata component. This is complementary to the diff --git a/cosmos_framework/data/vfm/augmentors/image_editing_transform.py b/cosmos_framework/data/generator/augmentors/image_editing_transform.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/image_editing_transform.py rename to cosmos_framework/data/generator/augmentors/image_editing_transform.py index 4af344b7..1571fd44 100644 --- a/cosmos_framework/data/vfm/augmentors/image_editing_transform.py +++ b/cosmos_framework/data/generator/augmentors/image_editing_transform.py @@ -457,7 +457,7 @@ def __call__(self, data_dict: dict) -> dict | None: # by GenerationDataClean.num_vision_items_per_sample (set in get_data_and_condition). # In pack_input_sequence, all items except the last are fully conditioned; # the last item uses condition_frame_indexes_vision ([] = fully generated). - from cosmos_framework.data.vfm.sequence_packing import SequencePlan + from cosmos_framework.data.generator.sequence_packing import SequencePlan data_dict["sequence_plan"] = SequencePlan( has_text=True, diff --git a/cosmos_framework/data/vfm/augmentors/image_editing_transform_test.py b/cosmos_framework/data/generator/augmentors/image_editing_transform_test.py similarity index 97% rename from cosmos_framework/data/vfm/augmentors/image_editing_transform_test.py rename to cosmos_framework/data/generator/augmentors/image_editing_transform_test.py index 849ad006..6d621abd 100644 --- a/cosmos_framework/data/vfm/augmentors/image_editing_transform_test.py +++ b/cosmos_framework/data/generator/augmentors/image_editing_transform_test.py @@ -6,7 +6,7 @@ import pytest from PIL import Image -from cosmos_framework.data.vfm.augmentors.image_editing_transform import ExtractImageEditingConversation +from cosmos_framework.data.generator.augmentors.image_editing_transform import ExtractImageEditingConversation _STRUCTURED_KEY = "edit_schema_all_inputs_qwen3-vl-235b-a22b-instruct" diff --git a/cosmos_framework/data/vfm/augmentors/image_resolution_filter.py b/cosmos_framework/data/generator/augmentors/image_resolution_filter.py similarity index 96% rename from cosmos_framework/data/vfm/augmentors/image_resolution_filter.py rename to cosmos_framework/data/generator/augmentors/image_resolution_filter.py index 50fb7fe7..b53040b0 100644 --- a/cosmos_framework/data/vfm/augmentors/image_resolution_filter.py +++ b/cosmos_framework/data/generator/augmentors/image_resolution_filter.py @@ -4,7 +4,7 @@ from typing import Optional from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor -from cosmos_framework.data.vfm.utils import IMAGE_RES_SIZE_INFO +from cosmos_framework.data.generator.utils import IMAGE_RES_SIZE_INFO # Map dataset_resolution_type to resolution tier key in IMAGE_RES_SIZE_INFO _DATASET_RESOLUTION_TIER: dict[str, str] = {"gt480p": "480", "gt720p": "720", "gt1080p": "1080"} diff --git a/cosmos_framework/data/vfm/augmentors/interleaved_image_transform.py b/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/interleaved_image_transform.py rename to cosmos_framework/data/generator/augmentors/interleaved_image_transform.py diff --git a/cosmos_framework/data/vfm/augmentors/interleaved_video_parsing.py b/cosmos_framework/data/generator/augmentors/interleaved_video_parsing.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/interleaved_video_parsing.py rename to cosmos_framework/data/generator/augmentors/interleaved_video_parsing.py index 41e4e4c6..c505c865 100644 --- a/cosmos_framework/data/vfm/augmentors/interleaved_video_parsing.py +++ b/cosmos_framework/data/generator/augmentors/interleaved_video_parsing.py @@ -14,7 +14,7 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.image.misc import obtain_augmentation_size from cosmos_framework.utils import log -from cosmos_framework.data.vfm.augmentors.video_parsing import VideoParsingWithFullFrames +from cosmos_framework.data.generator.augmentors.video_parsing import VideoParsingWithFullFrames # Local copies of the torchcodec decoder helpers so this module does not depend on # private symbols of ``video_parsing.py``. Behavior matches the originals. diff --git a/cosmos_framework/data/vfm/augmentors/merge_datadict.py b/cosmos_framework/data/generator/augmentors/merge_datadict.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/merge_datadict.py rename to cosmos_framework/data/generator/augmentors/merge_datadict.py diff --git a/cosmos_framework/data/vfm/augmentors/pkl_to_media.py b/cosmos_framework/data/generator/augmentors/pkl_to_media.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/pkl_to_media.py rename to cosmos_framework/data/generator/augmentors/pkl_to_media.py index aa9eb214..c66ea953 100644 --- a/cosmos_framework/data/vfm/augmentors/pkl_to_media.py +++ b/cosmos_framework/data/generator/augmentors/pkl_to_media.py @@ -17,7 +17,7 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.utils.vfm.video_preprocess import tensor_to_pil_images +from cosmos_framework.utils.generator.video_preprocess import tensor_to_pil_images Image.MAX_IMAGE_PIXELS = 933120000 _VIDEO_EXTENSIONS = "mp4 avi webm mov".split() diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/__init__.py b/cosmos_framework/data/generator/augmentors/reasoner/__init__.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/reasoner/__init__.py rename to cosmos_framework/data/generator/augmentors/reasoner/__init__.py diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/bytes_to_media.py b/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py similarity index 98% rename from cosmos_framework/data/vfm/augmentors/reasoner/bytes_to_media.py rename to cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py index db670b30..03a61f5e 100644 --- a/cosmos_framework/data/vfm/augmentors/reasoner/bytes_to_media.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py @@ -16,9 +16,9 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.reasoner.video_decoder_qwen import _video_decoder_qwen_func -from cosmos_framework.data.vfm.processors.qwen3vl_processor import Qwen3VLProcessor -from cosmos_framework.utils.vfm.video_preprocess import tensor_to_pil_images +from cosmos_framework.data.generator.reasoner.video_decoder_qwen import _video_decoder_qwen_func +from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor +from cosmos_framework.utils.generator.video_preprocess import tensor_to_pil_images class BytesToMedia(Augmentor): diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/filter_output_key.py b/cosmos_framework/data/generator/augmentors/reasoner/filter_output_key.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/reasoner/filter_output_key.py rename to cosmos_framework/data/generator/augmentors/reasoner/filter_output_key.py diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/filter_seq_length.py b/cosmos_framework/data/generator/augmentors/reasoner/filter_seq_length.py similarity index 97% rename from cosmos_framework/data/vfm/augmentors/reasoner/filter_seq_length.py rename to cosmos_framework/data/generator/augmentors/reasoner/filter_seq_length.py index 99e8a465..ed95ca36 100644 --- a/cosmos_framework/data/vfm/augmentors/reasoner/filter_seq_length.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/filter_seq_length.py @@ -7,7 +7,7 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.processors.qwen3vl_processor import Qwen3VLProcessor +from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor class FilterSeqLength(Augmentor): diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/floating_number_format.py b/cosmos_framework/data/generator/augmentors/reasoner/floating_number_format.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/reasoner/floating_number_format.py rename to cosmos_framework/data/generator/augmentors/reasoner/floating_number_format.py diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/format_describe_anything.py b/cosmos_framework/data/generator/augmentors/reasoner/format_describe_anything.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/reasoner/format_describe_anything.py rename to cosmos_framework/data/generator/augmentors/reasoner/format_describe_anything.py index 45028366..a955f3e9 100644 --- a/cosmos_framework/data/vfm/augmentors/reasoner/format_describe_anything.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/format_describe_anything.py @@ -13,7 +13,7 @@ from typing import Dict, List, Literal from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor -from cosmos_framework.data.vfm.augmentors.reasoner.timestamp import markdown_to_list +from cosmos_framework.data.generator.augmentors.reasoner.timestamp import markdown_to_list # reorder dict entries diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/nvlm_data_to_conversation.py b/cosmos_framework/data/generator/augmentors/reasoner/nvlm_data_to_conversation.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/reasoner/nvlm_data_to_conversation.py rename to cosmos_framework/data/generator/augmentors/reasoner/nvlm_data_to_conversation.py diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/prompt_format.py b/cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/reasoner/prompt_format.py rename to cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/shuffle_text_media_order.py b/cosmos_framework/data/generator/augmentors/reasoner/shuffle_text_media_order.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/reasoner/shuffle_text_media_order.py rename to cosmos_framework/data/generator/augmentors/reasoner/shuffle_text_media_order.py diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/timestamp.py b/cosmos_framework/data/generator/augmentors/reasoner/timestamp.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/reasoner/timestamp.py rename to cosmos_framework/data/generator/augmentors/reasoner/timestamp.py diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_with_subject_tracking.py b/cosmos_framework/data/generator/augmentors/reasoner/timestamp_with_subject_tracking.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/reasoner/timestamp_with_subject_tracking.py rename to cosmos_framework/data/generator/augmentors/reasoner/timestamp_with_subject_tracking.py index 34851186..88ebe8ce 100644 --- a/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_with_subject_tracking.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/timestamp_with_subject_tracking.py @@ -16,7 +16,7 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.augmentors.reasoner.timestamp import ( +from cosmos_framework.data.generator.augmentors.reasoner.timestamp import ( json_to_markdown, markdown_to_list, overlay_text, diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_augment_message.py b/cosmos_framework/data/generator/augmentors/reasoner/timestamp_without_augment_message.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_augment_message.py rename to cosmos_framework/data/generator/augmentors/reasoner/timestamp_without_augment_message.py index 8ca4984b..fc2fa3a0 100644 --- a/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_augment_message.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/timestamp_without_augment_message.py @@ -13,7 +13,7 @@ from typing import Dict, List, Literal, Tuple from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor -from cosmos_framework.data.vfm.augmentors.reasoner.timestamp import overlay_text +from cosmos_framework.data.generator.augmentors.reasoner.timestamp import overlay_text def list_to_markdown(conversation_data: List[Dict]) -> str: diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_end_time.py b/cosmos_framework/data/generator/augmentors/reasoner/timestamp_without_end_time.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_end_time.py rename to cosmos_framework/data/generator/augmentors/reasoner/timestamp_without_end_time.py index 0c9fb05d..46fff7f9 100644 --- a/cosmos_framework/data/vfm/augmentors/reasoner/timestamp_without_end_time.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/timestamp_without_end_time.py @@ -16,7 +16,7 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.augmentors.reasoner.timestamp import ( +from cosmos_framework.data.generator.augmentors.reasoner.timestamp import ( json_to_markdown, markdown_to_list, overlay_text, diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/tokenize_data.py b/cosmos_framework/data/generator/augmentors/reasoner/tokenize_data.py similarity index 98% rename from cosmos_framework/data/vfm/augmentors/reasoner/tokenize_data.py rename to cosmos_framework/data/generator/augmentors/reasoner/tokenize_data.py index d3446726..bc53e13f 100644 --- a/cosmos_framework/data/vfm/augmentors/reasoner/tokenize_data.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/tokenize_data.py @@ -12,9 +12,9 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.reasoner.video_decoder_qwen import token_to_pixels -from cosmos_framework.data.vfm.processors.qwen3vl_processor import Qwen3VLProcessor as Processor -from cosmos_framework.utils.vfm.reasoner.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD +from cosmos_framework.data.generator.reasoner.video_decoder_qwen import token_to_pixels +from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor as Processor +from cosmos_framework.utils.generator.reasoner.constant import IGNORE_INDEX, PROCESSOR_KEYS_TO_ADD def maybe_subsample_frames(model_name_or_path, list_of_pil_image, max_video_token_length, processor): diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/user_prompt_caption_general.json b/cosmos_framework/data/generator/augmentors/reasoner/user_prompt_caption_general.json similarity index 100% rename from cosmos_framework/data/vfm/augmentors/reasoner/user_prompt_caption_general.json rename to cosmos_framework/data/generator/augmentors/reasoner/user_prompt_caption_general.json diff --git a/cosmos_framework/data/vfm/augmentors/reasoner/user_prompt_ocr.json b/cosmos_framework/data/generator/augmentors/reasoner/user_prompt_ocr.json similarity index 100% rename from cosmos_framework/data/vfm/augmentors/reasoner/user_prompt_ocr.json rename to cosmos_framework/data/generator/augmentors/reasoner/user_prompt_ocr.json diff --git a/cosmos_framework/data/vfm/augmentors/resolution_text_info.py b/cosmos_framework/data/generator/augmentors/resolution_text_info.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/resolution_text_info.py rename to cosmos_framework/data/generator/augmentors/resolution_text_info.py diff --git a/cosmos_framework/data/vfm/augmentors/sequence_plan.py b/cosmos_framework/data/generator/augmentors/sequence_plan.py similarity index 98% rename from cosmos_framework/data/vfm/augmentors/sequence_plan.py rename to cosmos_framework/data/generator/augmentors/sequence_plan.py index 5a2202fd..e6e38532 100644 --- a/cosmos_framework/data/vfm/augmentors/sequence_plan.py +++ b/cosmos_framework/data/generator/augmentors/sequence_plan.py @@ -17,8 +17,8 @@ import torch from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor -from cosmos_framework.data.vfm.sequence_packing import SequencePlan -from cosmos_framework.model.vfm.tokenizers.uniae.frame_math import ( +from cosmos_framework.data.generator.sequence_packing import SequencePlan +from cosmos_framework.model.generator.tokenizers.uniae.frame_math import ( get_uniae_chunk_frames, get_uniae_latent_num_frames, normalize_uniae_chunk_frames, diff --git a/cosmos_framework/data/vfm/augmentors/sound_sequence_plan.py b/cosmos_framework/data/generator/augmentors/sound_sequence_plan.py similarity index 94% rename from cosmos_framework/data/vfm/augmentors/sound_sequence_plan.py rename to cosmos_framework/data/generator/augmentors/sound_sequence_plan.py index c3590536..3d7cafbc 100644 --- a/cosmos_framework/data/vfm/augmentors/sound_sequence_plan.py +++ b/cosmos_framework/data/generator/augmentors/sound_sequence_plan.py @@ -14,7 +14,7 @@ from typing import Optional from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor -from cosmos_framework.data.vfm.sound_data_utils import VALID_SOUND_MODES, build_sequence_plan_for_sound +from cosmos_framework.data.generator.sound_data_utils import VALID_SOUND_MODES, build_sequence_plan_for_sound class SoundSequencePlanBuilder(Augmentor): @@ -84,7 +84,7 @@ def __call__(self, data_dict: dict) -> dict | None: sound_latent_length=0, ) else: - from cosmos_framework.data.vfm.sequence_packing import SequencePlan + from cosmos_framework.data.generator.sequence_packing import SequencePlan data_dict["sequence_plan"] = SequencePlan( has_text=True, diff --git a/cosmos_framework/data/vfm/augmentors/text_tokenizer.py b/cosmos_framework/data/generator/augmentors/text_tokenizer.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/text_tokenizer.py rename to cosmos_framework/data/generator/augmentors/text_tokenizer.py diff --git a/cosmos_framework/data/vfm/augmentors/text_transforms_for_image.py b/cosmos_framework/data/generator/augmentors/text_transforms_for_image.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/text_transforms_for_image.py rename to cosmos_framework/data/generator/augmentors/text_transforms_for_image.py diff --git a/cosmos_framework/data/vfm/augmentors/text_transforms_for_video.py b/cosmos_framework/data/generator/augmentors/text_transforms_for_video.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/text_transforms_for_video.py rename to cosmos_framework/data/generator/augmentors/text_transforms_for_video.py diff --git a/cosmos_framework/data/vfm/augmentors/transfer_control_input/__init__.py b/cosmos_framework/data/generator/augmentors/transfer_control_input/__init__.py similarity index 73% rename from cosmos_framework/data/vfm/augmentors/transfer_control_input/__init__.py rename to cosmos_framework/data/generator/augmentors/transfer_control_input/__init__.py index 99d3ee8f..1cb86f5c 100644 --- a/cosmos_framework/data/vfm/augmentors/transfer_control_input/__init__.py +++ b/cosmos_framework/data/generator/augmentors/transfer_control_input/__init__.py @@ -3,6 +3,6 @@ """Transfer control input augmentors (edge, blur, depth, seg) for cosmos3 VFM; copied from transfer2 to avoid cosmos dependency.""" -from cosmos_framework.data.vfm.augmentors.transfer_control_input.control_input import AddControlInputComb +from cosmos_framework.data.generator.augmentors.transfer_control_input.control_input import AddControlInputComb __all__ = ["AddControlInputComb"] diff --git a/cosmos_framework/data/vfm/augmentors/transfer_control_input/blur.py b/cosmos_framework/data/generator/augmentors/transfer_control_input/blur.py similarity index 98% rename from cosmos_framework/data/vfm/augmentors/transfer_control_input/blur.py rename to cosmos_framework/data/generator/augmentors/transfer_control_input/blur.py index daacd306..1d3181da 100644 --- a/cosmos_framework/data/vfm/augmentors/transfer_control_input/blur.py +++ b/cosmos_framework/data/generator/augmentors/transfer_control_input/blur.py @@ -9,7 +9,7 @@ import numpy as np import torch -from cosmos_framework.data.vfm.augmentors.transfer_control_input.fast_blur import BilateralGaussian +from cosmos_framework.data.generator.augmentors.transfer_control_input.fast_blur import BilateralGaussian @attrs.define diff --git a/cosmos_framework/data/vfm/augmentors/transfer_control_input/control_input.py b/cosmos_framework/data/generator/augmentors/transfer_control_input/control_input.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/transfer_control_input/control_input.py rename to cosmos_framework/data/generator/augmentors/transfer_control_input/control_input.py index 3b46a825..f47ba3b1 100644 --- a/cosmos_framework/data/vfm/augmentors/transfer_control_input/control_input.py +++ b/cosmos_framework/data/generator/augmentors/transfer_control_input/control_input.py @@ -12,8 +12,8 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.augmentors.transfer_control_input.blur import Blur, BlurConfig -from cosmos_framework.data.vfm.augmentors.transfer_control_input.seg import ( +from cosmos_framework.data.generator.augmentors.transfer_control_input.blur import Blur, BlurConfig +from cosmos_framework.data.generator.augmentors.transfer_control_input.seg import ( decode_partial_rle_width1, segmentation_color_mask, ) diff --git a/cosmos_framework/data/vfm/augmentors/transfer_control_input/fast_blur.py b/cosmos_framework/data/generator/augmentors/transfer_control_input/fast_blur.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/transfer_control_input/fast_blur.py rename to cosmos_framework/data/generator/augmentors/transfer_control_input/fast_blur.py diff --git a/cosmos_framework/data/vfm/augmentors/transfer_control_input/seg.py b/cosmos_framework/data/generator/augmentors/transfer_control_input/seg.py similarity index 100% rename from cosmos_framework/data/vfm/augmentors/transfer_control_input/seg.py rename to cosmos_framework/data/generator/augmentors/transfer_control_input/seg.py diff --git a/cosmos_framework/data/vfm/augmentors/transfer_control_transform.py b/cosmos_framework/data/generator/augmentors/transfer_control_transform.py similarity index 98% rename from cosmos_framework/data/vfm/augmentors/transfer_control_transform.py rename to cosmos_framework/data/generator/augmentors/transfer_control_transform.py index c94b8af1..47f7249d 100644 --- a/cosmos_framework/data/vfm/augmentors/transfer_control_transform.py +++ b/cosmos_framework/data/generator/augmentors/transfer_control_transform.py @@ -28,9 +28,9 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -from cosmos_framework.data.vfm.augmentors.transfer_control_input import AddControlInputComb -from cosmos_framework.data.vfm.utils import VIDEO_RES_SIZE_INFO -from cosmos_framework.data.vfm.sequence_packing import SequencePlan +from cosmos_framework.data.generator.augmentors.transfer_control_input import AddControlInputComb +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.data.generator.sequence_packing import SequencePlan class SampleResolution(Augmentor): diff --git a/cosmos_framework/data/vfm/augmentors/video_parsing.py b/cosmos_framework/data/generator/augmentors/video_parsing.py similarity index 99% rename from cosmos_framework/data/vfm/augmentors/video_parsing.py rename to cosmos_framework/data/generator/augmentors/video_parsing.py index 1304f6c3..a443948e 100644 --- a/cosmos_framework/data/vfm/augmentors/video_parsing.py +++ b/cosmos_framework/data/generator/augmentors/video_parsing.py @@ -15,8 +15,8 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.data.imaginaire.webdataset.augmentors.image.misc import obtain_augmentation_size from cosmos_framework.utils import log -from cosmos_framework.data.vfm.utils import VIDEO_RES_SIZE_INFO -from cosmos_framework.model.vfm.tokenizers.uniae.frame_math import ( +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.model.generator.tokenizers.uniae.frame_math import ( align_uniae_num_video_frames, get_uniae_chunk_frames, normalize_uniae_chunk_frames, diff --git a/cosmos_framework/data/vfm/dataflow/__init__.py b/cosmos_framework/data/generator/dataflow/__init__.py similarity index 56% rename from cosmos_framework/data/vfm/dataflow/__init__.py rename to cosmos_framework/data/generator/dataflow/__init__.py index 50b300b3..08070950 100644 --- a/cosmos_framework/data/vfm/dataflow/__init__.py +++ b/cosmos_framework/data/generator/dataflow/__init__.py @@ -6,17 +6,17 @@ from __future__ import annotations -from cosmos_framework.data.vfm.dataflow.base import ( +from cosmos_framework.data.generator.dataflow.base import ( BatchCollator, DataDistributor, RawItemProcessor, SampleBatcher, ) -from cosmos_framework.data.vfm.dataflow.batchers import PoolPackingBatcher, SequentialPackingBatcher, SimpleBatcher -from cosmos_framework.data.vfm.dataflow.collators import DefaultBatchCollator, VFMListCollator -from cosmos_framework.data.vfm.dataflow.distributors import IterableDistributor, MapDistributor, MixtureDistributor, RankPartitionedDistributor -from cosmos_framework.data.vfm.dataflow.loader import CosmosDataLoader, JointCosmosDataLoader -from cosmos_framework.data.vfm.dataflow.processors import IdentityProcessor +from cosmos_framework.data.generator.dataflow.batchers import PoolPackingBatcher, SequentialPackingBatcher, SimpleBatcher +from cosmos_framework.data.generator.dataflow.collators import DefaultBatchCollator, VFMListCollator +from cosmos_framework.data.generator.dataflow.distributors import IterableDistributor, MapDistributor, MixtureDistributor, RankPartitionedDistributor +from cosmos_framework.data.generator.dataflow.loader import CosmosDataLoader, JointCosmosDataLoader +from cosmos_framework.data.generator.dataflow.processors import IdentityProcessor __all__ = [ "BatchCollator", diff --git a/cosmos_framework/data/vfm/dataflow/base.py b/cosmos_framework/data/generator/dataflow/base.py similarity index 100% rename from cosmos_framework/data/vfm/dataflow/base.py rename to cosmos_framework/data/generator/dataflow/base.py diff --git a/cosmos_framework/data/vfm/dataflow/batchers.py b/cosmos_framework/data/generator/dataflow/batchers.py similarity index 99% rename from cosmos_framework/data/vfm/dataflow/batchers.py rename to cosmos_framework/data/generator/dataflow/batchers.py index b6c63dd8..699b21a0 100644 --- a/cosmos_framework/data/vfm/dataflow/batchers.py +++ b/cosmos_framework/data/generator/dataflow/batchers.py @@ -9,7 +9,7 @@ from enum import Enum from typing import Callable, Iterator, Optional -from cosmos_framework.data.vfm.dataflow.base import SampleBatcher +from cosmos_framework.data.generator.dataflow.base import SampleBatcher class SimpleBatcher(SampleBatcher): diff --git a/cosmos_framework/data/vfm/dataflow/collators.py b/cosmos_framework/data/generator/dataflow/collators.py similarity index 99% rename from cosmos_framework/data/vfm/dataflow/collators.py rename to cosmos_framework/data/generator/dataflow/collators.py index 81b3782f..535b428b 100644 --- a/cosmos_framework/data/vfm/dataflow/collators.py +++ b/cosmos_framework/data/generator/dataflow/collators.py @@ -9,7 +9,7 @@ import torch.utils.data from torch.utils.data.dataloader import default_collate -from cosmos_framework.data.vfm.dataflow.base import BatchCollator +from cosmos_framework.data.generator.dataflow.base import BatchCollator class DefaultBatchCollator(BatchCollator): diff --git a/cosmos_framework/data/vfm/dataflow/distributors.py b/cosmos_framework/data/generator/dataflow/distributors.py similarity index 98% rename from cosmos_framework/data/vfm/dataflow/distributors.py rename to cosmos_framework/data/generator/dataflow/distributors.py index e9fd6ad4..6c1a5467 100644 --- a/cosmos_framework/data/vfm/dataflow/distributors.py +++ b/cosmos_framework/data/generator/dataflow/distributors.py @@ -12,7 +12,7 @@ from typing import Any, Iterator -from cosmos_framework.data.vfm.dataflow.base import DataDistributor +from cosmos_framework.data.generator.dataflow.base import DataDistributor class IterableDistributor(DataDistributor): diff --git a/cosmos_framework/data/vfm/dataflow/golden_vfm_test.py b/cosmos_framework/data/generator/dataflow/golden_vfm_test.py similarity index 98% rename from cosmos_framework/data/vfm/dataflow/golden_vfm_test.py rename to cosmos_framework/data/generator/dataflow/golden_vfm_test.py index 131c812e..906c3fc7 100644 --- a/cosmos_framework/data/vfm/dataflow/golden_vfm_test.py +++ b/cosmos_framework/data/generator/dataflow/golden_vfm_test.py @@ -214,11 +214,11 @@ def test_vfm_golden_batches_match(monkeypatch): Asserts exact list[list[Tensor]] nesting for _MULTI_ITEM_KEYS (video, text_token_ids) and flat list[Tensor] for image_size — no nesting normalization. """ - from cosmos_framework.data.vfm.joint_dataloader import ( + from cosmos_framework.data.generator.joint_dataloader import ( PackingDataLoader, RankPartitionedDataLoader, ) - from cosmos_framework.data.vfm.dataflow import ( + from cosmos_framework.data.generator.dataflow import ( CosmosDataLoader, RankPartitionedDistributor, SequentialPackingBatcher, diff --git a/cosmos_framework/data/vfm/dataflow/loader.py b/cosmos_framework/data/generator/dataflow/loader.py similarity index 97% rename from cosmos_framework/data/vfm/dataflow/loader.py rename to cosmos_framework/data/generator/dataflow/loader.py index 3a2a6a47..edf164bc 100644 --- a/cosmos_framework/data/vfm/dataflow/loader.py +++ b/cosmos_framework/data/generator/dataflow/loader.py @@ -13,14 +13,14 @@ import numpy as np from cosmos_framework.utils import log -from cosmos_framework.data.vfm.dataflow.base import ( +from cosmos_framework.data.generator.dataflow.base import ( BatchCollator, DataDistributor, RawItemProcessor, SampleBatcher, ) -from cosmos_framework.data.vfm.dataflow.batchers import SimpleBatcher -from cosmos_framework.data.vfm.dataflow.collators import DefaultBatchCollator +from cosmos_framework.data.generator.dataflow.batchers import SimpleBatcher +from cosmos_framework.data.generator.dataflow.collators import DefaultBatchCollator class _DataflowIterableDataset(torch.utils.data.IterableDataset): @@ -152,7 +152,7 @@ def __init__( dp_world_size=dp_world_size, ) - from cosmos_framework.data.vfm.dataflow.distributors import MapDistributor + from cosmos_framework.data.generator.dataflow.distributors import MapDistributor if isinstance(distributor, MapDistributor) and num_workers > 0 and not persistent_workers: log.info( diff --git a/cosmos_framework/data/vfm/dataflow/processors.py b/cosmos_framework/data/generator/dataflow/processors.py similarity index 85% rename from cosmos_framework/data/vfm/dataflow/processors.py rename to cosmos_framework/data/generator/dataflow/processors.py index fc2b6af0..bf1f4bb6 100644 --- a/cosmos_framework/data/vfm/dataflow/processors.py +++ b/cosmos_framework/data/generator/dataflow/processors.py @@ -7,7 +7,7 @@ from typing import Any -from cosmos_framework.data.vfm.dataflow.base import RawItemProcessor +from cosmos_framework.data.generator.dataflow.base import RawItemProcessor class IdentityProcessor(RawItemProcessor): diff --git a/cosmos_framework/data/vfm/dataflow/resume_test.py b/cosmos_framework/data/generator/dataflow/resume_test.py similarity index 96% rename from cosmos_framework/data/vfm/dataflow/resume_test.py rename to cosmos_framework/data/generator/dataflow/resume_test.py index 2fde6b95..2f3a67a8 100644 --- a/cosmos_framework/data/vfm/dataflow/resume_test.py +++ b/cosmos_framework/data/generator/dataflow/resume_test.py @@ -9,7 +9,7 @@ import torch from cosmos_framework.callbacks.cosmos_dataloader_state import CosmosDataLoaderStateCallback -from cosmos_framework.data.vfm.dataflow import ( +from cosmos_framework.data.generator.dataflow import ( CosmosDataLoader, IdentityProcessor, MapDistributor, diff --git a/cosmos_framework/data/vfm/joint_dataloader.py b/cosmos_framework/data/generator/joint_dataloader.py similarity index 99% rename from cosmos_framework/data/vfm/joint_dataloader.py rename to cosmos_framework/data/generator/joint_dataloader.py index fa27e528..64f430eb 100644 --- a/cosmos_framework/data/vfm/joint_dataloader.py +++ b/cosmos_framework/data/generator/joint_dataloader.py @@ -14,7 +14,7 @@ from cosmos_framework.utils.lazy_config import instantiate from cosmos_framework.utils import log -from cosmos_framework.model.vfm.tokenizers.uniae.frame_math import ( +from cosmos_framework.model.generator.tokenizers.uniae.frame_math import ( get_uniae_chunk_frames, get_uniae_latent_num_frames, normalize_uniae_chunk_frames, diff --git a/cosmos_framework/data/vfm/local_datasets/__init__.py b/cosmos_framework/data/generator/local_datasets/__init__.py similarity index 100% rename from cosmos_framework/data/vfm/local_datasets/__init__.py rename to cosmos_framework/data/generator/local_datasets/__init__.py diff --git a/cosmos_framework/data/vfm/local_datasets/helper.py b/cosmos_framework/data/generator/local_datasets/helper.py similarity index 100% rename from cosmos_framework/data/vfm/local_datasets/helper.py rename to cosmos_framework/data/generator/local_datasets/helper.py diff --git a/cosmos_framework/data/vfm/local_datasets/sft_dataset.py b/cosmos_framework/data/generator/local_datasets/sft_dataset.py similarity index 98% rename from cosmos_framework/data/vfm/local_datasets/sft_dataset.py rename to cosmos_framework/data/generator/local_datasets/sft_dataset.py index 0632c2a1..1a9fcdbc 100644 --- a/cosmos_framework/data/vfm/local_datasets/sft_dataset.py +++ b/cosmos_framework/data/generator/local_datasets/sft_dataset.py @@ -16,7 +16,7 @@ import numpy as np import torch -from cosmos_framework.data.vfm.local_datasets.helper import ( +from cosmos_framework.data.generator.local_datasets.helper import ( client_config, download_from_s3, ffmpeg_decode_video, @@ -24,11 +24,11 @@ get_video_metadata, parse_s3_url, ) -from cosmos_framework.data.vfm.sequence_packing import SequencePlan -from cosmos_framework.data.vfm.sequence_packing.modalities import add_special_tokens -from cosmos_framework.data.vfm.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.data.generator.sequence_packing import SequencePlan +from cosmos_framework.data.generator.sequence_packing.modalities import add_special_tokens +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO from cosmos_framework.inference.structured_caption import CAPTION_JSON_KEY, caption_json_to_prompt -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import tokenize_caption from cosmos_framework.utils import log from cosmos_framework.utils.flags import INTERNAL from cosmos_framework.utils.lazy_config import instantiate as lazy_instantiate diff --git a/cosmos_framework/data/vfm/local_datasets/sft_dataset_caption_test.py b/cosmos_framework/data/generator/local_datasets/sft_dataset_caption_test.py similarity index 96% rename from cosmos_framework/data/vfm/local_datasets/sft_dataset_caption_test.py rename to cosmos_framework/data/generator/local_datasets/sft_dataset_caption_test.py index 6909221c..04560883 100644 --- a/cosmos_framework/data/vfm/local_datasets/sft_dataset_caption_test.py +++ b/cosmos_framework/data/generator/local_datasets/sft_dataset_caption_test.py @@ -4,7 +4,7 @@ import json -from cosmos_framework.data.vfm.local_datasets.sft_dataset import _select_caption +from cosmos_framework.data.generator.local_datasets.sft_dataset import _select_caption from cosmos_framework.inference.structured_caption import CAPTION_JSON_KEY diff --git a/cosmos_framework/data/vfm/packing_iterable_dataset.py b/cosmos_framework/data/generator/packing_iterable_dataset.py similarity index 99% rename from cosmos_framework/data/vfm/packing_iterable_dataset.py rename to cosmos_framework/data/generator/packing_iterable_dataset.py index 8d8d022b..ed7c0d91 100644 --- a/cosmos_framework/data/vfm/packing_iterable_dataset.py +++ b/cosmos_framework/data/generator/packing_iterable_dataset.py @@ -4,7 +4,7 @@ """ Abstract base class for pool-based token-budget bin-packing over multiple datasets. -Extracted from ``cosmos_framework.data.vfm.reasoner.joint_dataset_dynamic_batch_webloader`` +Extracted from ``cosmos_framework.data.generator.reasoner.joint_dataset_dynamic_batch_webloader`` so that both the VLM and VFM internal dataloaders can share a single packing implementation. Usage diff --git a/cosmos_framework/data/vfm/processors/__init__.py b/cosmos_framework/data/generator/processors/__init__.py similarity index 93% rename from cosmos_framework/data/vfm/processors/__init__.py rename to cosmos_framework/data/generator/processors/__init__.py index 73ea3404..46a3a473 100644 --- a/cosmos_framework/data/vfm/processors/__init__.py +++ b/cosmos_framework/data/generator/processors/__init__.py @@ -9,12 +9,12 @@ from transformers import PreTrainedTokenizerFast -from cosmos_framework.data.vfm.processors.base import BaseVLMProcessor -from cosmos_framework.data.vfm.processors.nemotron3densevl_processor import Nemotron3DenseVLProcessor -from cosmos_framework.data.vfm.processors.nemotronvl_processor import NemotronVLProcessor -from cosmos_framework.data.vfm.processors.qwen3vl_processor import Qwen3VLProcessor -from cosmos_framework.model.vfm.tokenizers.tokenization_qwen2 import Qwen2Tokenizer -from cosmos_framework.utils.vfm.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 +from cosmos_framework.data.generator.processors.base import BaseVLMProcessor +from cosmos_framework.data.generator.processors.nemotron3densevl_processor import Nemotron3DenseVLProcessor +from cosmos_framework.data.generator.processors.nemotronvl_processor import NemotronVLProcessor +from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor +from cosmos_framework.model.generator.tokenizers.tokenization_qwen2 import Qwen2Tokenizer +from cosmos_framework.utils.generator.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 _VARIANT_TO_CREDENTIALS = { "s3": ("credentials/s3_training.secret", "bucket4"), diff --git a/cosmos_framework/data/vfm/processors/base.py b/cosmos_framework/data/generator/processors/base.py similarity index 97% rename from cosmos_framework/data/vfm/processors/base.py rename to cosmos_framework/data/generator/processors/base.py index 7b0aa3e0..f944139e 100644 --- a/cosmos_framework/data/vfm/processors/base.py +++ b/cosmos_framework/data/generator/processors/base.py @@ -22,8 +22,8 @@ from transformers.models.auto.processing_auto import AutoProcessor from cosmos_framework.utils import log -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption -from cosmos_framework.utils.vfm.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import tokenize_caption +from cosmos_framework.utils.generator.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 def convert_string_content_to_list_content(messages: List[Dict]) -> List[Dict]: diff --git a/cosmos_framework/data/vfm/processors/nemotron3densevl_processor.py b/cosmos_framework/data/generator/processors/nemotron3densevl_processor.py similarity index 99% rename from cosmos_framework/data/vfm/processors/nemotron3densevl_processor.py rename to cosmos_framework/data/generator/processors/nemotron3densevl_processor.py index 8024c98a..5c923b1e 100644 --- a/cosmos_framework/data/vfm/processors/nemotron3densevl_processor.py +++ b/cosmos_framework/data/generator/processors/nemotron3densevl_processor.py @@ -8,7 +8,7 @@ from PIL import Image from qwen_vl_utils.vision_process import smart_resize -from cosmos_framework.data.vfm.processors.base import ( +from cosmos_framework.data.generator.processors.base import ( BaseVLMProcessor, convert_string_content_to_list_content, maybe_parse_video_content, diff --git a/cosmos_framework/data/vfm/processors/nemotronvl_processor.py b/cosmos_framework/data/generator/processors/nemotronvl_processor.py similarity index 99% rename from cosmos_framework/data/vfm/processors/nemotronvl_processor.py rename to cosmos_framework/data/generator/processors/nemotronvl_processor.py index 077c80e8..5e4cea5b 100644 --- a/cosmos_framework/data/vfm/processors/nemotronvl_processor.py +++ b/cosmos_framework/data/generator/processors/nemotronvl_processor.py @@ -10,7 +10,7 @@ from transformers.video_utils import VideoMetadata from cosmos_framework.utils import log -from cosmos_framework.data.vfm.processors.base import BaseVLMProcessor, convert_string_content_to_list_content +from cosmos_framework.data.generator.processors.base import BaseVLMProcessor, convert_string_content_to_list_content nemotron_chat_template = """ {%- set ns = namespace(enable_thinking=false, has_sys_prompt=false, non_tool_system_content='', has_video=false, explicit_think_requested=false) -%} @@ -387,7 +387,7 @@ def add_assistant_tokens_mask(self, tokens): if __name__ == "__main__": """ - PYTHONPATH=. python3 cosmos_framework/data/vfm/processors/nemotronvl_processor.py + PYTHONPATH=. python3 cosmos_framework/data/generator/processors/nemotronvl_processor.py """ processor = NemotronVLProcessor("nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16") from io import BytesIO diff --git a/cosmos_framework/data/vfm/processors/qwen3vl_processor.py b/cosmos_framework/data/generator/processors/qwen3vl_processor.py similarity index 98% rename from cosmos_framework/data/vfm/processors/qwen3vl_processor.py rename to cosmos_framework/data/generator/processors/qwen3vl_processor.py index 2a8eb59e..863bd18b 100644 --- a/cosmos_framework/data/vfm/processors/qwen3vl_processor.py +++ b/cosmos_framework/data/generator/processors/qwen3vl_processor.py @@ -6,7 +6,7 @@ import numpy as np import torch -from cosmos_framework.data.vfm.processors.base import ( +from cosmos_framework.data.generator.processors.base import ( BaseVLMProcessor, convert_string_content_to_list_content, maybe_get_max_pixels_from_images_kwargs, @@ -158,7 +158,7 @@ def add_assistant_tokens_mask(self, tokens): if __name__ == "__main__": """ - PYTHONPATH=. python3 cosmos_framework/data/vfm/processors/qwen3vl_processor.py + PYTHONPATH=. python3 cosmos_framework/data/generator/processors/qwen3vl_processor.py For image, expected output: input_ids: type: shape: torch.Size([2772]) diff --git a/cosmos_framework/data/vfm/reasoner/video_decoder_qwen.py b/cosmos_framework/data/generator/reasoner/video_decoder_qwen.py similarity index 99% rename from cosmos_framework/data/vfm/reasoner/video_decoder_qwen.py rename to cosmos_framework/data/generator/reasoner/video_decoder_qwen.py index 12c9bcc0..06e3d819 100644 --- a/cosmos_framework/data/vfm/reasoner/video_decoder_qwen.py +++ b/cosmos_framework/data/generator/reasoner/video_decoder_qwen.py @@ -22,7 +22,7 @@ from torchvision.transforms import InterpolationMode from cosmos_framework.utils import log -from cosmos_framework.data.vfm.processors.qwen3vl_processor import Qwen3VLProcessor +from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor Image.MAX_IMAGE_PIXELS = 933120000 _VIDEO_EXTENSIONS = "mp4 avi webm mov".split() diff --git a/cosmos_framework/data/vfm/recipe_database_api.py b/cosmos_framework/data/generator/recipe_database_api.py similarity index 100% rename from cosmos_framework/data/vfm/recipe_database_api.py rename to cosmos_framework/data/generator/recipe_database_api.py diff --git a/cosmos_framework/data/vfm/sequence_packing/__init__.py b/cosmos_framework/data/generator/sequence_packing/__init__.py similarity index 73% rename from cosmos_framework/data/vfm/sequence_packing/__init__.py rename to cosmos_framework/data/generator/sequence_packing/__init__.py index 3e91c62a..33349353 100644 --- a/cosmos_framework/data/vfm/sequence_packing/__init__.py +++ b/cosmos_framework/data/generator/sequence_packing/__init__.py @@ -3,8 +3,8 @@ """High-level entry points for VFM sequence packing.""" -from cosmos_framework.data.vfm.sequence_packing.packers import pack_input_sequence -from cosmos_framework.data.vfm.sequence_packing.types import ( +from cosmos_framework.data.generator.sequence_packing.packers import pack_input_sequence +from cosmos_framework.data.generator.sequence_packing.types import ( ModalityData, PackedSequence, SequencePlan, diff --git a/cosmos_framework/data/vfm/sequence_packing/modalities.py b/cosmos_framework/data/generator/sequence_packing/modalities.py similarity index 99% rename from cosmos_framework/data/vfm/sequence_packing/modalities.py rename to cosmos_framework/data/generator/sequence_packing/modalities.py index 247789a9..b74cf634 100644 --- a/cosmos_framework/data/vfm/sequence_packing/modalities.py +++ b/cosmos_framework/data/generator/sequence_packing/modalities.py @@ -7,11 +7,11 @@ import torch -from cosmos_framework.data.vfm.sequence_packing.mrope import ( +from cosmos_framework.data.generator.sequence_packing.mrope import ( get_3d_mrope_ids_text_tokens, get_3d_mrope_ids_vae_tokens, ) -from cosmos_framework.data.vfm.sequence_packing.types import ModalityData, PackedSequence +from cosmos_framework.data.generator.sequence_packing.types import ModalityData, PackedSequence def prepare_attention_mask_per_sample(split_lens, attn_modes, device="cpu"): diff --git a/cosmos_framework/data/vfm/sequence_packing/mrope.py b/cosmos_framework/data/generator/sequence_packing/mrope.py similarity index 100% rename from cosmos_framework/data/vfm/sequence_packing/mrope.py rename to cosmos_framework/data/generator/sequence_packing/mrope.py diff --git a/cosmos_framework/data/vfm/sequence_packing/natten.py b/cosmos_framework/data/generator/sequence_packing/natten.py similarity index 100% rename from cosmos_framework/data/vfm/sequence_packing/natten.py rename to cosmos_framework/data/generator/sequence_packing/natten.py diff --git a/cosmos_framework/data/vfm/sequence_packing/packers.py b/cosmos_framework/data/generator/sequence_packing/packers.py similarity index 98% rename from cosmos_framework/data/vfm/sequence_packing/packers.py rename to cosmos_framework/data/generator/sequence_packing/packers.py index 643c944d..b6a52234 100644 --- a/cosmos_framework/data/vfm/sequence_packing/packers.py +++ b/cosmos_framework/data/generator/sequence_packing/packers.py @@ -9,17 +9,17 @@ import torch -from cosmos_framework.data.vfm.sequence_packing.modalities import ( +from cosmos_framework.data.generator.sequence_packing.modalities import ( pack_action_tokens, pack_sound_tokens, pack_text_tokens, pack_vision_tokens, ) -from cosmos_framework.data.vfm.sequence_packing.temporal_causal import pack_supertokens_temporal_causal -from cosmos_framework.data.vfm.sequence_packing.types import PackedSequence, SequencePlan +from cosmos_framework.data.generator.sequence_packing.temporal_causal import pack_supertokens_temporal_causal +from cosmos_framework.data.generator.sequence_packing.types import PackedSequence, SequencePlan if TYPE_CHECKING: - from cosmos_framework.model.vfm.utils.data_and_condition import GenerationDataClean + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean def pack_input_sequence( diff --git a/cosmos_framework/data/vfm/sequence_packing/runtime.py b/cosmos_framework/data/generator/sequence_packing/runtime.py similarity index 100% rename from cosmos_framework/data/vfm/sequence_packing/runtime.py rename to cosmos_framework/data/generator/sequence_packing/runtime.py diff --git a/cosmos_framework/data/vfm/sequence_packing/temporal_causal.py b/cosmos_framework/data/generator/sequence_packing/temporal_causal.py similarity index 98% rename from cosmos_framework/data/vfm/sequence_packing/temporal_causal.py rename to cosmos_framework/data/generator/sequence_packing/temporal_causal.py index 2610c26a..bcbe26d2 100644 --- a/cosmos_framework/data/vfm/sequence_packing/temporal_causal.py +++ b/cosmos_framework/data/generator/sequence_packing/temporal_causal.py @@ -7,8 +7,8 @@ import torch -from cosmos_framework.data.vfm.sequence_packing.mrope import get_3d_mrope_ids_vae_tokens -from cosmos_framework.data.vfm.sequence_packing.types import ModalityData, PackedSequence +from cosmos_framework.data.generator.sequence_packing.mrope import get_3d_mrope_ids_vae_tokens +from cosmos_framework.data.generator.sequence_packing.types import ModalityData, PackedSequence def pack_supertokens_temporal_causal( diff --git a/cosmos_framework/data/vfm/sequence_packing/types.py b/cosmos_framework/data/generator/sequence_packing/types.py similarity index 99% rename from cosmos_framework/data/vfm/sequence_packing/types.py rename to cosmos_framework/data/generator/sequence_packing/types.py index 38018bd0..4904c163 100644 --- a/cosmos_framework/data/vfm/sequence_packing/types.py +++ b/cosmos_framework/data/generator/sequence_packing/types.py @@ -11,7 +11,7 @@ import torch if TYPE_CHECKING: - from cosmos_framework.model.vfm.utils.data_and_condition import GenerationDataClean + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean @dataclass diff --git a/cosmos_framework/data/vfm/sound_data_utils.py b/cosmos_framework/data/generator/sound_data_utils.py similarity index 95% rename from cosmos_framework/data/vfm/sound_data_utils.py rename to cosmos_framework/data/generator/sound_data_utils.py index 2d739b0c..9eb4f22a 100644 --- a/cosmos_framework/data/vfm/sound_data_utils.py +++ b/cosmos_framework/data/generator/sound_data_utils.py @@ -4,7 +4,7 @@ """Sound data utilities for building sequence plans and handling audio-video generation modes. This module provides utilities for building SequencePlan objects based on sound generation modes, -similar to how action modes are handled in cosmos_framework/data/vfm/action/data_utils.py. +similar to how action modes are handled in cosmos_framework/data/generator/action/data_utils.py. Supported modes: - t2vs: Text → Video + Sound (joint generation) @@ -13,7 +13,7 @@ - ti2sv: Text + Image → Sound + Video (first frame conditioned, rest + sound generated) """ -from cosmos_framework.data.vfm.sequence_packing import SequencePlan +from cosmos_framework.data.generator.sequence_packing import SequencePlan # Valid generation modes for sound VALID_SOUND_MODES = {"t2vs", "tv2s", "ts2v", "ti2sv"} @@ -29,7 +29,7 @@ def build_sequence_plan_for_sound( This function determines the appropriate condition frame indexes for vision and sound based on the specified mode. It mirrors how `build_sequence_plan_from_mode` works - for action in cosmos_framework/data/vfm/action/data_utils.py. + for action in cosmos_framework/data/generator/action/data_utils.py. Args: mode: Generation mode. One of: diff --git a/cosmos_framework/data/vfm/utils.py b/cosmos_framework/data/generator/utils.py similarity index 100% rename from cosmos_framework/data/vfm/utils.py rename to cosmos_framework/data/generator/utils.py diff --git a/cosmos_framework/data/vfm/watchdog.py b/cosmos_framework/data/generator/watchdog.py similarity index 100% rename from cosmos_framework/data/vfm/watchdog.py rename to cosmos_framework/data/generator/watchdog.py diff --git a/cosmos_framework/data/vlm/__init__.py b/cosmos_framework/data/reasoner/__init__.py similarity index 100% rename from cosmos_framework/data/vlm/__init__.py rename to cosmos_framework/data/reasoner/__init__.py diff --git a/cosmos_framework/data/vlm/data_sources_videophy2/__init__.py b/cosmos_framework/data/reasoner/data_sources_videophy2/__init__.py similarity index 100% rename from cosmos_framework/data/vlm/data_sources_videophy2/__init__.py rename to cosmos_framework/data/reasoner/data_sources_videophy2/__init__.py diff --git a/cosmos_framework/data/vlm/data_sources_videophy2/data_weight/__init__.py b/cosmos_framework/data/reasoner/data_sources_videophy2/data_weight/__init__.py similarity index 100% rename from cosmos_framework/data/vlm/data_sources_videophy2/data_weight/__init__.py rename to cosmos_framework/data/reasoner/data_sources_videophy2/data_weight/__init__.py diff --git a/cosmos_framework/data/vlm/data_sources_videophy2/data_weight/default.py b/cosmos_framework/data/reasoner/data_sources_videophy2/data_weight/default.py similarity index 84% rename from cosmos_framework/data/vlm/data_sources_videophy2/data_weight/default.py rename to cosmos_framework/data/reasoner/data_sources_videophy2/data_weight/default.py index 016f4dc6..666a23e7 100644 --- a/cosmos_framework/data/vlm/data_sources_videophy2/data_weight/default.py +++ b/cosmos_framework/data/reasoner/data_sources_videophy2/data_weight/default.py @@ -3,7 +3,7 @@ """Default weight mix for the VideoPhy-2 dataset (single source).""" -from cosmos_framework.data.vlm.data_sources_videophy2.videophy2 import DATAINFO, url_to_category +from cosmos_framework.data.reasoner.data_sources_videophy2.videophy2 import DATAINFO, url_to_category # Single-source SFT mix — used for both data_train and data_val. The dataloader # registrar picks the right split via the LocalDataSource.manifest_path mapping. diff --git a/cosmos_framework/data/vlm/data_sources_videophy2/videophy2.py b/cosmos_framework/data/reasoner/data_sources_videophy2/videophy2.py similarity index 95% rename from cosmos_framework/data/vlm/data_sources_videophy2/videophy2.py rename to cosmos_framework/data/reasoner/data_sources_videophy2/videophy2.py index f132799e..2d07f7a6 100644 --- a/cosmos_framework/data/vlm/data_sources_videophy2/videophy2.py +++ b/cosmos_framework/data/reasoner/data_sources_videophy2/videophy2.py @@ -16,7 +16,7 @@ import os -from cosmos_framework.data.vlm.local_dataset_utils import LocalDataSource +from cosmos_framework.data.reasoner.local_dataset_utils import LocalDataSource _DEFAULT_ROOT = "examples/data/videophysics" VIDEOPHYSICS_ROOT = os.environ.get("VIDEOPHYSICS_ROOT", _DEFAULT_ROOT) diff --git a/cosmos_framework/data/vlm/local_dataset_utils.py b/cosmos_framework/data/reasoner/local_dataset_utils.py similarity index 100% rename from cosmos_framework/data/vlm/local_dataset_utils.py rename to cosmos_framework/data/reasoner/local_dataset_utils.py diff --git a/cosmos_framework/data/vlm/local_sft_dataset.py b/cosmos_framework/data/reasoner/local_sft_dataset.py similarity index 100% rename from cosmos_framework/data/vlm/local_sft_dataset.py rename to cosmos_framework/data/reasoner/local_sft_dataset.py diff --git a/cosmos_framework/data/vlm/processors/__init__.py b/cosmos_framework/data/reasoner/processors/__init__.py similarity index 73% rename from cosmos_framework/data/vlm/processors/__init__.py rename to cosmos_framework/data/reasoner/processors/__init__.py index 6d6d143d..2fa9d6b0 100644 --- a/cosmos_framework/data/vlm/processors/__init__.py +++ b/cosmos_framework/data/reasoner/processors/__init__.py @@ -3,10 +3,10 @@ from typing import Optional -from cosmos_framework.data.vlm.processors.nemotron3densevl_processor import Nemotron3DenseVLProcessor -from cosmos_framework.data.vlm.processors.nemotronvl_processor import NemotronVLProcessor -from cosmos_framework.data.vlm.processors.qwen3vl_processor import Qwen3VLProcessor -from cosmos_framework.utils.vlm.pretrained_models_downloader import resolve_hf_model_store +from cosmos_framework.data.reasoner.processors.nemotron3densevl_processor import Nemotron3DenseVLProcessor +from cosmos_framework.data.reasoner.processors.nemotronvl_processor import NemotronVLProcessor +from cosmos_framework.data.reasoner.processors.qwen3vl_processor import Qwen3VLProcessor +from cosmos_framework.utils.reasoner.pretrained_models_downloader import resolve_hf_model_store def build_processor( diff --git a/cosmos_framework/data/vlm/processors/nemotron3densevl_processor.py b/cosmos_framework/data/reasoner/processors/nemotron3densevl_processor.py similarity index 99% rename from cosmos_framework/data/vlm/processors/nemotron3densevl_processor.py rename to cosmos_framework/data/reasoner/processors/nemotron3densevl_processor.py index fd8406f8..37795233 100644 --- a/cosmos_framework/data/vlm/processors/nemotron3densevl_processor.py +++ b/cosmos_framework/data/reasoner/processors/nemotron3densevl_processor.py @@ -11,7 +11,7 @@ from transformers.models.auto.processing_auto import AutoProcessor from cosmos_framework.utils import log -from cosmos_framework.utils.vlm.pretrained_models_downloader import maybe_download_hf_model_from_s3 +from cosmos_framework.utils.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 def convert_string_content_to_list_content(messages: List[Dict]) -> List[Dict]: diff --git a/cosmos_framework/data/vlm/processors/nemotronvl_processor.py b/cosmos_framework/data/reasoner/processors/nemotronvl_processor.py similarity index 99% rename from cosmos_framework/data/vlm/processors/nemotronvl_processor.py rename to cosmos_framework/data/reasoner/processors/nemotronvl_processor.py index 1fde099a..6a4d8ce2 100644 --- a/cosmos_framework/data/vlm/processors/nemotronvl_processor.py +++ b/cosmos_framework/data/reasoner/processors/nemotronvl_processor.py @@ -12,7 +12,7 @@ from transformers.video_utils import VideoMetadata from cosmos_framework.utils import log -from cosmos_framework.utils.vlm.pretrained_models_downloader import maybe_download_hf_model_from_s3 +from cosmos_framework.utils.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 nemotron_chat_template = """ {%- set ns = namespace(enable_thinking=false, has_sys_prompt=false, non_tool_system_content='', has_video=false, explicit_think_requested=false) -%} @@ -427,7 +427,7 @@ def decode(self, *args, **kwargs): if __name__ == "__main__": """ - PYTHONPATH=. python3 cosmos_framework/data/vlm/processors/nemotronvl_processor.py + PYTHONPATH=. python3 cosmos_framework/data/reasoner/processors/nemotronvl_processor.py inputs: dict_keys(['input_ids', 'attention_mask', 'pixel_values', 'image_sizes', 'text']) input_ids: type: shape: torch.Size([6699]) diff --git a/cosmos_framework/data/vlm/processors/qwen3vl_processor.py b/cosmos_framework/data/reasoner/processors/qwen3vl_processor.py similarity index 98% rename from cosmos_framework/data/vlm/processors/qwen3vl_processor.py rename to cosmos_framework/data/reasoner/processors/qwen3vl_processor.py index b624d6a1..7a476651 100644 --- a/cosmos_framework/data/vlm/processors/qwen3vl_processor.py +++ b/cosmos_framework/data/reasoner/processors/qwen3vl_processor.py @@ -9,7 +9,7 @@ from transformers.models.auto.processing_auto import AutoProcessor from cosmos_framework.utils import log -from cosmos_framework.utils.vlm.pretrained_models_downloader import maybe_download_hf_model_from_s3 +from cosmos_framework.utils.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 def convert_string_content_to_list_content(messages: List[Dict]) -> List[Dict]: @@ -226,7 +226,7 @@ def decode(self, *args, **kwargs): if __name__ == "__main__": """ - PYTHONPATH=. python3 cosmos_framework/data/vlm/processors/qwen3vl_processor.py + PYTHONPATH=. python3 cosmos_framework/data/reasoner/processors/qwen3vl_processor.py inputs: dict_keys(['input_ids', 'attention_mask', 'pixel_values', 'image_sizes', 'text']) input_ids: type: shape: torch.Size([6699]) diff --git a/cosmos_framework/data/vfm/action/datasets/__init__.py b/cosmos_framework/data/vfm/action/datasets/__init__.py deleted file mode 100644 index b1357d1a..00000000 --- a/cosmos_framework/data/vfm/action/datasets/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -"""Action dataset wrappers for Cosmos Action. - -All concrete datasets inherit from :class:`ActionBaseDataset` and expose a -``load_action_stats()`` classmethod for retrieving pre-computed normalization -statistics without instantiating the dataset. -""" - -from cosmos_framework.data.vfm.action.datasets.agibotworld_beta_lerobot_dataset import AgiBotWorldBetaLeRobotDataset -from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset -from cosmos_framework.data.vfm.action.datasets.bridge_orig_lerobot_dataset import BridgeOrigLeRobotDataset -from cosmos_framework.data.vfm.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset -from cosmos_framework.data.vfm.action.datasets.fractal_lerobot_dataset import FractalLeRobotDataset -from cosmos_framework.data.vfm.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset -from cosmos_framework.data.vfm.action.datasets.robomind_franka_dataset import RoboMINDFrankaDataset -from cosmos_framework.data.vfm.action.datasets.robomind_ur_dataset import RoboMINDURDataset -from cosmos_framework.data.vfm.action.datasets.umi_lerobot_dataset import UMILeRobotDataset - -__all__ = [ - "ActionBaseDataset", - "AgiBotWorldBetaLeRobotDataset", - "BridgeOrigLeRobotDataset", - "DROIDLeRobotDataset", - "FractalLeRobotDataset", - "LIBEROLeRobotDataset", - "RoboMINDFrankaDataset", - "RoboMINDURDataset", - "UMILeRobotDataset", -] diff --git a/cosmos_framework/inference/action.py b/cosmos_framework/inference/action.py index 7f27c351..08dd106b 100644 --- a/cosmos_framework/inference/action.py +++ b/cosmos_framework/inference/action.py @@ -9,21 +9,21 @@ import torch -from cosmos_framework.data.vfm.action.action_processing import ( +from cosmos_framework.data.generator.action.action_processing import ( ActionProcessingRecord, make_batched_action_processing_fields, pad_action_to_max_dim, ) -from cosmos_framework.data.vfm.action.domain_utils import EMBODIMENT_TO_RAW_ACTION_DIM, get_domain_id -from cosmos_framework.data.vfm.action.json_formatter import ActionPromptJsonFormatter -from cosmos_framework.data.vfm.action.transforms import ( +from cosmos_framework.data.generator.action.domain_utils import EMBODIMENT_TO_RAW_ACTION_DIM, get_domain_id +from cosmos_framework.data.generator.action.json_formatter import ActionPromptJsonFormatter +from cosmos_framework.data.generator.action.transforms import ( build_sequence_plan_from_mode, find_closest_target_size, reflection_pad_to_target, ) from cosmos_framework.inference.args import ModelMode from cosmos_framework.inference.vision import read_media_frames -from cosmos_framework.utils.vfm.data_utils import get_vision_data_resolution +from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution def _load_actions( diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 7095fa7e..914d0807 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -419,7 +419,7 @@ def vision_size(self) -> tuple[int, int]: autodetect via cosmos_framework.inference.vision.read_and_resize_media before reaching this property and never observe the fallback. """ - from cosmos_framework.data.vfm.utils import IMAGE_RES_SIZE_INFO, VIDEO_RES_SIZE_INFO + from cosmos_framework.data.generator.utils import IMAGE_RES_SIZE_INFO, VIDEO_RES_SIZE_INFO assert self.resolution aspect_ratio: AspectRatio = self.aspect_ratio or "16,9" diff --git a/cosmos_framework/inference/args_test.py b/cosmos_framework/inference/args_test.py index d373bf2a..eb5002b1 100644 --- a/cosmos_framework/inference/args_test.py +++ b/cosmos_framework/inference/args_test.py @@ -22,7 +22,7 @@ from cosmos_framework.inference.common.config import structure_config if TYPE_CHECKING: - from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel _H100_MEMORY_BYTES = 80 * 1024**3 # Reserved for future use (paired with the reserved memory-based `_get_dp_shard_size` diff --git a/cosmos_framework/inference/common/args.py b/cosmos_framework/inference/common/args.py index 8ac4ec7c..4fa87c97 100644 --- a/cosmos_framework/inference/common/args.py +++ b/cosmos_framework/inference/common/args.py @@ -548,7 +548,7 @@ def _build_checkpoint(self, checkpoints: dict[str, CheckpointConfig]): self.checkpoint_path = self.checkpoint_path.rstrip("/") # Strip '/model' suffix, since it isn't included in checkpoint_db. # Automatically added during checkpoint load by - # 'cosmos_framework.utils.vfm.model_loader.load_model_from_checkpoint'. + # 'cosmos_framework.utils.generator.model_loader.load_model_from_checkpoint'. if not self.checkpoint_path.endswith("/model"): self.checkpoint_path = self.checkpoint_path + "/model" else: diff --git a/cosmos_framework/inference/common/checkpoints.py b/cosmos_framework/inference/common/checkpoints.py index 0a10aae3..62913aa9 100644 --- a/cosmos_framework/inference/common/checkpoints.py +++ b/cosmos_framework/inference/common/checkpoints.py @@ -71,7 +71,7 @@ def _materialize_avae_ckpt(local_dir: str) -> None: The new HF layout ships ``sound_tokenizer/{config.json, diffusion_pytorch_model.safetensors}`` in the diffusers OobleckDecoder layout (``decoder.block.*`` keys, Snake1d ``alpha``/``beta`` shaped ``[1, C, 1]``). The - native loader in ``cosmos_framework/model/vfm/tokenizers/audio/avae.py`` builds an + native loader in ``cosmos_framework/model/generator/tokenizers/audio/avae.py`` builds an ``nn.Sequential`` decoder keyed ``decoder.layers.*`` with Snake params shaped ``[C]`` and loads via ``load_state_dict(strict=False)`` — so without remapping the keys, none match and every decoder weight is silently left at init (noise). @@ -131,7 +131,7 @@ def register_checkpoints(): for s3_prefix in [ # 'cosmos_framework.configs.base.defaults.vlm.download_tokenizer_files' "cosmos3/pretrained/huggingface", - # 'cosmos_framework.utils.vfm.vlm.pretrained_models_downloader.maybe_download_hf_model_from_s3' + # 'cosmos_framework.utils.generator.vlm.pretrained_models_downloader.maybe_download_hf_model_from_s3' "cosmos_reason2/hf_models", ]: register_checkpoint( diff --git a/cosmos_framework/inference/common/config.py b/cosmos_framework/inference/common/config.py index 121b1252..73d65c9b 100644 --- a/cosmos_framework/inference/common/config.py +++ b/cosmos_framework/inference/common/config.py @@ -434,49 +434,49 @@ def _structure_torch_tensor(data: Any, _cls: Any) -> torch.Tensor | None: (r"(? str: "cosmos_framework.model.attention", "cosmos_framework.utils.checkpoint_db", "imaginaire.trainer", - "cosmos_framework.utils.vfm.model_loader", + "cosmos_framework.utils.generator.model_loader", "*.callbacks.*", ] _LOGGER_EXCLUDE = [ diff --git a/cosmos_framework/inference/common/public_model_config.py b/cosmos_framework/inference/common/public_model_config.py index 7a7362f6..66bb9e31 100644 --- a/cosmos_framework/inference/common/public_model_config.py +++ b/cosmos_framework/inference/common/public_model_config.py @@ -180,16 +180,16 @@ def _canonicalize_module_path(path: str) -> str: # so the reasoner subtree isn't shadowed by the general vfm rules below. ("cosmos_framework.configs.base.defaults.reasoner.", "projects.cosmos3.vfm.configs.base.defaults.vlm."), ("cosmos_framework.configs.base.reasoner.", "projects.cosmos3.vfm.configs.base.vlm."), - ("cosmos_framework.model.vfm.reasoner.", "projects.cosmos3.vfm.models.vlm."), - ("cosmos_framework.data.vfm.reasoner.", "projects.cosmos3.vfm.datasets.vlm."), - ("cosmos_framework.data.vfm.augmentors.reasoner.", "projects.cosmos3.vfm.datasets.augmentors.vlm."), - ("cosmos_framework.utils.vfm.reasoner.", "projects.cosmos3.vfm.utils.vlm."), + ("cosmos_framework.model.generator.reasoner.", "projects.cosmos3.vfm.models.vlm."), + ("cosmos_framework.data.generator.reasoner.", "projects.cosmos3.vfm.datasets.vlm."), + ("cosmos_framework.data.generator.augmentors.reasoner.", "projects.cosmos3.vfm.datasets.augmentors.vlm."), + ("cosmos_framework.utils.generator.reasoner.", "projects.cosmos3.vfm.utils.vlm."), ("cosmos3._src.vfm.", "projects.cosmos3.vfm."), ("cosmos_framework.configs.base.", "projects.cosmos3.vfm.configs.base."), - ("cosmos_framework.model.vfm.tokenizers.", "projects.cosmos3.vfm.tokenizers."), - ("cosmos_framework.model.vfm.diffusion.", "projects.cosmos3.vfm.diffusion."), - ("cosmos_framework.model.vfm.", "projects.cosmos3.vfm.models."), - ("cosmos_framework.data.vfm.processors.", "projects.cosmos3.vfm.processors."), + ("cosmos_framework.model.generator.tokenizers.", "projects.cosmos3.vfm.tokenizers."), + ("cosmos_framework.model.generator.diffusion.", "projects.cosmos3.vfm.diffusion."), + ("cosmos_framework.model.generator.", "projects.cosmos3.vfm.models."), + ("cosmos_framework.data.generator.processors.", "projects.cosmos3.vfm.processors."), ("cosmos.model.vfm.tokenizers.", "projects.cosmos3.vfm.tokenizers."), ("cosmos.model.vfm.diffusion.", "projects.cosmos3.vfm.diffusion."), ("cosmos.model.vfm.", "projects.cosmos3.vfm.models."), @@ -202,7 +202,7 @@ def _canonicalize_module_path(path: str) -> str: def _runtime_module_path(canonical_path: str) -> str: - if _module_exists("cosmos_framework.model.vfm"): + if _module_exists("cosmos_framework.model.generator"): return _replace_vfm_module_prefix(canonical_path, package="cosmos_framework") if _module_exists("cosmos.model.vfm"): return _replace_vfm_module_prefix(canonical_path, package="cosmos") @@ -217,18 +217,18 @@ def _replace_vfm_module_prefix(canonical_path: str, *, package: str) -> str: # so the specific vlm→reasoner subtree isn't shadowed by them. ("projects.cosmos3.vfm.configs.base.defaults.vlm.", f"{package}.configs.base.defaults.reasoner."), ("projects.cosmos3.vfm.configs.base.vlm.", f"{package}.configs.base.reasoner."), - ("projects.cosmos3.vfm.models.vlm.", f"{package}.model.vfm.reasoner."), - ("projects.cosmos3.vfm.datasets.vlm.", f"{package}.data.vfm.reasoner."), - ("projects.cosmos3.vfm.datasets.augmentors.vlm.", f"{package}.data.vfm.augmentors.reasoner."), - ("projects.cosmos3.vfm.utils.vlm.", f"{package}.utils.vfm.reasoner."), + ("projects.cosmos3.vfm.models.vlm.", f"{package}.model.generator.reasoner."), + ("projects.cosmos3.vfm.datasets.vlm.", f"{package}.data.generator.reasoner."), + ("projects.cosmos3.vfm.datasets.augmentors.vlm.", f"{package}.data.generator.augmentors.reasoner."), + ("projects.cosmos3.vfm.utils.vlm.", f"{package}.utils.generator.reasoner."), ("projects.cosmos3.vfm.configs.base.", f"{package}.configs.base."), - ("projects.cosmos3.vfm.models.", f"{package}.model.vfm."), - ("projects.cosmos3.vfm.tokenizers.", f"{package}.model.vfm.tokenizers."), - ("projects.cosmos3.vfm.diffusion.", f"{package}.model.vfm.diffusion."), - ("projects.cosmos3.vfm.processors.", f"{package}.data.vfm.processors."), - ("projects.cosmos3.vfm.datasets.", f"{package}.data.vfm."), - ("projects.cosmos3.vfm.scripts.action.", f"{package}.data.vfm.action_scripts."), - ("projects.cosmos3.vfm.utils.", f"{package}.utils.vfm."), + ("projects.cosmos3.vfm.models.", f"{package}.model.generator."), + ("projects.cosmos3.vfm.tokenizers.", f"{package}.model.generator.tokenizers."), + ("projects.cosmos3.vfm.diffusion.", f"{package}.model.generator.diffusion."), + ("projects.cosmos3.vfm.processors.", f"{package}.data.generator.processors."), + ("projects.cosmos3.vfm.datasets.", f"{package}.data.generator."), + ("projects.cosmos3.vfm.scripts.action.", f"{package}.data.generator.action_scripts."), + ("projects.cosmos3.vfm.utils.", f"{package}.utils.generator."), ) for old, new in replacements: if canonical_path.startswith(old): @@ -250,13 +250,13 @@ def _to_public_string(value: str) -> str: def _public_string_from_runtime_file_prefix(value: str, *, package: str) -> str | None: replacements = ( (f"{package}/configs/base/", "configs/base/"), - (f"{package}/model/vfm/tokenizers/", "tokenizers/"), - (f"{package}/model/vfm/diffusion/", "diffusion/"), - (f"{package}/model/vfm/", "models/"), - (f"{package}/data/vfm/processors/", "processors/"), - (f"{package}/data/vfm/action_scripts/", "scripts/action/"), - (f"{package}/data/vfm/", "datasets/"), - (f"{package}/utils/vfm/", "utils/"), + (f"{package}/model/generator/tokenizers/", "tokenizers/"), + (f"{package}/model/generator/diffusion/", "diffusion/"), + (f"{package}/model/generator/", "models/"), + (f"{package}/data/generator/processors/", "processors/"), + (f"{package}/data/generator/action_scripts/", "scripts/action/"), + (f"{package}/data/generator/", "datasets/"), + (f"{package}/utils/generator/", "utils/"), ) for old, new in replacements: if value.startswith(old): @@ -268,7 +268,7 @@ def _from_public_string(value: str) -> str: if not value.startswith(_PUBLIC_VFM_URI_PREFIX): return value suffix = value[len(_PUBLIC_VFM_URI_PREFIX) :] - if _module_exists("cosmos_framework.model.vfm"): + if _module_exists("cosmos_framework.model.generator"): return _replace_vfm_file_prefix(suffix, package="cosmos_framework") if _module_exists("cosmos.model.vfm"): return _replace_vfm_file_prefix(suffix, package="cosmos") @@ -280,13 +280,13 @@ def _from_public_string(value: str) -> str: def _replace_vfm_file_prefix(suffix: str, *, package: str) -> str: replacements = ( ("configs/base/", f"{package}/configs/base/"), - ("models/", f"{package}/model/vfm/"), - ("tokenizers/", f"{package}/model/vfm/tokenizers/"), - ("diffusion/", f"{package}/model/vfm/diffusion/"), - ("processors/", f"{package}/data/vfm/processors/"), - ("datasets/", f"{package}/data/vfm/"), - ("scripts/action/", f"{package}/data/vfm/action_scripts/"), - ("utils/", f"{package}/utils/vfm/"), + ("models/", f"{package}/model/generator/"), + ("tokenizers/", f"{package}/model/generator/tokenizers/"), + ("diffusion/", f"{package}/model/generator/diffusion/"), + ("processors/", f"{package}/data/generator/processors/"), + ("datasets/", f"{package}/data/generator/"), + ("scripts/action/", f"{package}/data/generator/action_scripts/"), + ("utils/", f"{package}/utils/generator/"), ) for old, new in replacements: if suffix.startswith(old): diff --git a/cosmos_framework/inference/common/public_model_config_test.py b/cosmos_framework/inference/common/public_model_config_test.py index feb63ec4..def99d3f 100644 --- a/cosmos_framework/inference/common/public_model_config_test.py +++ b/cosmos_framework/inference/common/public_model_config_test.py @@ -21,7 +21,7 @@ def _has_key(obj, key: str) -> bool: def test_public_model_config_round_trip_removes_internal_metadata(): model_config = { "_recursive_": False, - "_target_": "cosmos_framework.model.vfm.omni_mot_model.OmniMoTModel", + "_target_": "cosmos_framework.model.generator.omni_mot_model.OmniMoTModel", "config": { "_type": "cosmos_framework.configs.base.defaults.model_config.OmniMoTModelConfig", "activation_checkpointing": { @@ -29,18 +29,18 @@ def test_public_model_config_round_trip_removes_internal_metadata(): "mode": "full", }, "tokenizer": { - "_target_": "cosmos_framework.model.vfm.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface", + "_target_": "cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface", "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth", }, "vlm_config": { "_type": "cosmos_framework.configs.base.defaults.reasoner.VLMConfig", "model_instance": { - "_target_": "cosmos_framework.model.vfm.mot.unified_mot.Qwen3VLTextForCausalLM", + "_target_": "cosmos_framework.model.generator.mot.unified_mot.Qwen3VLTextForCausalLM", "config": { "_target_": "cosmos_framework.configs.base.defaults.reasoner.create_vlm_config", "base_config": { - "_target_": "cosmos_framework.model.vfm.mot.unified_mot.Qwen3VLMoTConfig.from_json_file", - "json_file": "cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json", + "_target_": "cosmos_framework.model.generator.mot.unified_mot.Qwen3VLMoTConfig.from_json_file", + "json_file": "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json", }, }, }, diff --git a/cosmos_framework/inference/dataset_samples.py b/cosmos_framework/inference/dataset_samples.py index 5910b524..90f5aeee 100644 --- a/cosmos_framework/inference/dataset_samples.py +++ b/cosmos_framework/inference/dataset_samples.py @@ -15,7 +15,7 @@ from cosmos_framework.inference.args import OmniSampleOverrides from cosmos_framework.scripts.dataset_utils import set_dataset_mode -from cosmos_framework.utils.vfm.data_utils import get_vision_data_resolution +from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution Sample = tuple[OmniSampleOverrides, dict[str, Any]] diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index a35968e5..53ee1a57 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -50,9 +50,9 @@ from cosmos_framework.tools.visualize.video import save_img_or_video from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig -from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import _SYSTEM_PROMPT_IMAGE_EDITING -from cosmos_framework.model.vfm.upsampler.prompts import is_upsampled_prompt +from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import _SYSTEM_PROMPT_IMAGE_EDITING +from cosmos_framework.model.generator.upsampler.prompts import is_upsampled_prompt if TYPE_CHECKING: from cosmos_framework.configs.base.defaults.model_config import OmniMoTModelConfig @@ -1033,7 +1033,7 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: compile_config = cls._get_compile_config(setup_args) if setup_args.checkpoint_type == CheckpointType.DCP and setup_args.config_file_type == ConfigFileType.MODULE: from cosmos_framework.inference.common.config import save_config - from cosmos_framework.utils.vfm.model_loader import load_model_from_checkpoint + from cosmos_framework.utils.generator.model_loader import load_model_from_checkpoint if not setup_args.experiment: raise ValueError("'experiment' is required") diff --git a/cosmos_framework/inference/model.py b/cosmos_framework/inference/model.py index 8904e6ea..e9c38b48 100644 --- a/cosmos_framework/inference/model.py +++ b/cosmos_framework/inference/model.py @@ -46,17 +46,17 @@ from cosmos_framework.utils.flags import SMOKE if TYPE_CHECKING: - from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel # Resolve to the release-tree root so relative-path checkpoint config entries -# (e.g. `cosmos_framework/model/vfm/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json`) load +# (e.g. `cosmos_framework/model/generator/vlm/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json`) load # correctly under contextlib.chdir(_ROOT_DIR) in __init__. In the original # cosmos3 release the cosmos3 package lives at the tree root, so parents[1] is # the release root. In the cosmos_training release the package lives at # cosmos_framework/inference/, so the release root is one level higher (parents[2]). try: - import cosmos_framework.model.vfm # noqa: F401 + import cosmos_framework.model.generator # noqa: F401 _ROOT_DIR = Path(__file__).parents[2].absolute() except ImportError: diff --git a/cosmos_framework/inference/sound.py b/cosmos_framework/inference/sound.py index 59d008e8..cebf39a4 100644 --- a/cosmos_framework/inference/sound.py +++ b/cosmos_framework/inference/sound.py @@ -9,8 +9,8 @@ import torch -from cosmos_framework.data.vfm.sequence_packing import SequencePlan -from cosmos_framework.data.vfm.sound_data_utils import build_sequence_plan_for_sound +from cosmos_framework.data.generator.sequence_packing import SequencePlan +from cosmos_framework.data.generator.sound_data_utils import build_sequence_plan_for_sound @dataclass diff --git a/cosmos_framework/inference/sound_test.py b/cosmos_framework/inference/sound_test.py index 06664294..5add45d5 100644 --- a/cosmos_framework/inference/sound_test.py +++ b/cosmos_framework/inference/sound_test.py @@ -38,7 +38,7 @@ def test_load_conditioning_audio_trims(tmp_path: Path): import types -from cosmos_framework.data.vfm.sequence_packing import SequencePlan +from cosmos_framework.data.generator.sequence_packing import SequencePlan from cosmos_framework.inference.sound import inject_sound_into_batch diff --git a/cosmos_framework/inference/transfer.py b/cosmos_framework/inference/transfer.py index 17aead8e..5d38c9a0 100644 --- a/cosmos_framework/inference/transfer.py +++ b/cosmos_framework/inference/transfer.py @@ -26,10 +26,10 @@ uint8_to_normalized_float, ) from cosmos_framework.utils import log -from cosmos_framework.data.vfm.sequence_packing import SequencePlan -from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel -from cosmos_framework.model.vfm.utils.data_and_condition import GenerationDataClean -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import _SYSTEM_PROMPT_TRANSFER +from cosmos_framework.data.generator.sequence_packing import SequencePlan +from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel +from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import _SYSTEM_PROMPT_TRANSFER @dataclass @@ -62,7 +62,7 @@ def apply_transfer_control_augmentor( preset_blur_strength: PresetBlurStrength, ) -> torch.Tensor: """Compute edge/blur transfer controls on the fly from uint8 input frames.""" - from cosmos_framework.data.vfm.augmentors.transfer_control_input.control_input import ( + from cosmos_framework.data.generator.augmentors.transfer_control_input.control_input import ( AddControlInputBlur, AddControlInputEdge, ) diff --git a/cosmos_framework/inference/vision.py b/cosmos_framework/inference/vision.py index ce3b62bc..bdb2d734 100644 --- a/cosmos_framework/inference/vision.py +++ b/cosmos_framework/inference/vision.py @@ -11,8 +11,8 @@ import torchvision.transforms.functional as TF from PIL import Image -from cosmos_framework.data.vfm.sequence_packing import SequencePlan -from cosmos_framework.data.vfm.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.data.generator.sequence_packing import SequencePlan +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO def resize_pil_image(image: Image.Image, max_size: int, padding_constant: int) -> Image.Image: diff --git a/cosmos_framework/model/vfm/__init__.py b/cosmos_framework/model/generator/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/__init__.py rename to cosmos_framework/model/generator/__init__.py diff --git a/cosmos_framework/model/vfm/algorithm/__init__.py b/cosmos_framework/model/generator/algorithm/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/algorithm/__init__.py rename to cosmos_framework/model/generator/algorithm/__init__.py diff --git a/cosmos_framework/model/vfm/algorithm/loss/__init__.py b/cosmos_framework/model/generator/algorithm/loss/__init__.py similarity index 51% rename from cosmos_framework/model/vfm/algorithm/loss/__init__.py rename to cosmos_framework/model/generator/algorithm/loss/__init__.py index 4a67ce95..b640ffa1 100644 --- a/cosmos_framework/model/vfm/algorithm/loss/__init__.py +++ b/cosmos_framework/model/generator/algorithm/loss/__init__.py @@ -5,18 +5,18 @@ __all__: list[str] = [] -from cosmos_framework.model.vfm.algorithm.loss.cross_entropy import cross_entropy_loss, weighted_cross_entropy_loss +from cosmos_framework.model.generator.algorithm.loss.cross_entropy import cross_entropy_loss, weighted_cross_entropy_loss __all__ += ["cross_entropy_loss", "weighted_cross_entropy_loss"] -from cosmos_framework.model.vfm.algorithm.loss.load_balancing import compute_load_balancing_loss +from cosmos_framework.model.generator.algorithm.loss.load_balancing import compute_load_balancing_loss __all__ += ["compute_load_balancing_loss"] -from cosmos_framework.model.vfm.algorithm.loss.time_weight import TrainTimeWeight +from cosmos_framework.model.generator.algorithm.loss.time_weight import TrainTimeWeight __all__ += ["TrainTimeWeight"] -from cosmos_framework.model.vfm.algorithm.loss.flow_matching import compute_flow_matching_loss +from cosmos_framework.model.generator.algorithm.loss.flow_matching import compute_flow_matching_loss __all__ += ["compute_flow_matching_loss"] diff --git a/cosmos_framework/model/vfm/algorithm/loss/cross_entropy.py b/cosmos_framework/model/generator/algorithm/loss/cross_entropy.py similarity index 98% rename from cosmos_framework/model/vfm/algorithm/loss/cross_entropy.py rename to cosmos_framework/model/generator/algorithm/loss/cross_entropy.py index 9e3c9b5b..283b1366 100644 --- a/cosmos_framework/model/vfm/algorithm/loss/cross_entropy.py +++ b/cosmos_framework/model/generator/algorithm/loss/cross_entropy.py @@ -32,7 +32,7 @@ import torch.distributed as dist import torch.nn.functional as F -from cosmos_framework.utils.vfm.reasoner.constant import IGNORE_INDEX +from cosmos_framework.utils.generator.reasoner.constant import IGNORE_INDEX def cross_entropy_loss( diff --git a/cosmos_framework/model/vfm/algorithm/loss/flow_matching.py b/cosmos_framework/model/generator/algorithm/loss/flow_matching.py similarity index 98% rename from cosmos_framework/model/vfm/algorithm/loss/flow_matching.py rename to cosmos_framework/model/generator/algorithm/loss/flow_matching.py index ecece254..d0fce869 100644 --- a/cosmos_framework/model/vfm/algorithm/loss/flow_matching.py +++ b/cosmos_framework/model/generator/algorithm/loss/flow_matching.py @@ -12,7 +12,7 @@ import torch -from cosmos_framework.model.vfm.diffusion.rectified_flow import RectifiedFlow +from cosmos_framework.model.generator.diffusion.rectified_flow import RectifiedFlow def compute_flow_matching_loss( diff --git a/cosmos_framework/model/vfm/algorithm/loss/load_balancing.py b/cosmos_framework/model/generator/algorithm/loss/load_balancing.py similarity index 97% rename from cosmos_framework/model/vfm/algorithm/loss/load_balancing.py rename to cosmos_framework/model/generator/algorithm/loss/load_balancing.py index 206d1a39..8d073fb5 100644 --- a/cosmos_framework/model/vfm/algorithm/loss/load_balancing.py +++ b/cosmos_framework/model/generator/algorithm/loss/load_balancing.py @@ -5,7 +5,7 @@ from torch.distributed.tensor import DTensor, Partial from torch.distributed.tensor.device_mesh import DeviceMesh -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.qwen3_vl_moe import LBLMetadata +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import LBLMetadata def compute_load_balancing_loss( diff --git a/cosmos_framework/model/vfm/algorithm/loss/time_weight.py b/cosmos_framework/model/generator/algorithm/loss/time_weight.py similarity index 100% rename from cosmos_framework/model/vfm/algorithm/loss/time_weight.py rename to cosmos_framework/model/generator/algorithm/loss/time_weight.py diff --git a/cosmos_framework/model/vfm/diffusion/__init__.py b/cosmos_framework/model/generator/diffusion/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/diffusion/__init__.py rename to cosmos_framework/model/generator/diffusion/__init__.py diff --git a/cosmos_framework/model/vfm/diffusion/rectified_flow.py b/cosmos_framework/model/generator/diffusion/rectified_flow.py similarity index 99% rename from cosmos_framework/model/vfm/diffusion/rectified_flow.py rename to cosmos_framework/model/generator/diffusion/rectified_flow.py index bea0e04c..c5ce7237 100644 --- a/cosmos_framework/model/vfm/diffusion/rectified_flow.py +++ b/cosmos_framework/model/generator/diffusion/rectified_flow.py @@ -7,7 +7,7 @@ import torch.distributed from diffusers import FlowMatchEulerDiscreteScheduler -from cosmos_framework.model.vfm.algorithm.loss.time_weight import TrainTimeWeight +from cosmos_framework.model.generator.algorithm.loss.time_weight import TrainTimeWeight class TrainTimeSampler: diff --git a/cosmos_framework/model/vfm/diffusion/samplers/__init__.py b/cosmos_framework/model/generator/diffusion/samplers/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/diffusion/samplers/__init__.py rename to cosmos_framework/model/generator/diffusion/samplers/__init__.py diff --git a/cosmos_framework/model/vfm/diffusion/samplers/edm.py b/cosmos_framework/model/generator/diffusion/samplers/edm.py similarity index 100% rename from cosmos_framework/model/vfm/diffusion/samplers/edm.py rename to cosmos_framework/model/generator/diffusion/samplers/edm.py diff --git a/cosmos_framework/model/vfm/diffusion/samplers/fixed_step.py b/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py similarity index 98% rename from cosmos_framework/model/vfm/diffusion/samplers/fixed_step.py rename to cosmos_framework/model/generator/diffusion/samplers/fixed_step.py index f33b9177..706b4d4e 100644 --- a/cosmos_framework/model/vfm/diffusion/samplers/fixed_step.py +++ b/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py @@ -20,7 +20,7 @@ import torch -from cosmos_framework.model.vfm.diffusion.samplers.utils import run_multiseed +from cosmos_framework.model.generator.diffusion.samplers.utils import run_multiseed class FixedStepSampler: diff --git a/cosmos_framework/model/vfm/diffusion/samplers/fixed_step_test.py b/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py similarity index 98% rename from cosmos_framework/model/vfm/diffusion/samplers/fixed_step_test.py rename to cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py index d7f43439..dd214018 100644 --- a/cosmos_framework/model/vfm/diffusion/samplers/fixed_step_test.py +++ b/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py @@ -4,7 +4,7 @@ import pytest import torch -from cosmos_framework.model.vfm.diffusion.samplers.fixed_step import FixedStepSampler +from cosmos_framework.model.generator.diffusion.samplers.fixed_step import FixedStepSampler @pytest.mark.L0 diff --git a/cosmos_framework/model/vfm/diffusion/samplers/fm_solvers_unipc.py b/cosmos_framework/model/generator/diffusion/samplers/fm_solvers_unipc.py similarity index 100% rename from cosmos_framework/model/vfm/diffusion/samplers/fm_solvers_unipc.py rename to cosmos_framework/model/generator/diffusion/samplers/fm_solvers_unipc.py diff --git a/cosmos_framework/model/vfm/diffusion/samplers/unipc.py b/cosmos_framework/model/generator/diffusion/samplers/unipc.py similarity index 95% rename from cosmos_framework/model/vfm/diffusion/samplers/unipc.py rename to cosmos_framework/model/generator/diffusion/samplers/unipc.py index db5f3d49..f192c4a2 100644 --- a/cosmos_framework/model/vfm/diffusion/samplers/unipc.py +++ b/cosmos_framework/model/generator/diffusion/samplers/unipc.py @@ -8,8 +8,8 @@ from cosmos_framework.utils.config import make_freezable from cosmos_framework.utils.progress_bar import progress_bar -from cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc import FlowUniPCMultistepScheduler -from cosmos_framework.model.vfm.diffusion.samplers.utils import run_multiseed +from cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc import FlowUniPCMultistepScheduler +from cosmos_framework.model.generator.diffusion.samplers.utils import run_multiseed @make_freezable diff --git a/cosmos_framework/model/vfm/diffusion/samplers/utils.py b/cosmos_framework/model/generator/diffusion/samplers/utils.py similarity index 100% rename from cosmos_framework/model/vfm/diffusion/samplers/utils.py rename to cosmos_framework/model/generator/diffusion/samplers/utils.py diff --git a/cosmos_framework/model/vfm/diffusion/samplers/utils_test.py b/cosmos_framework/model/generator/diffusion/samplers/utils_test.py similarity index 97% rename from cosmos_framework/model/vfm/diffusion/samplers/utils_test.py rename to cosmos_framework/model/generator/diffusion/samplers/utils_test.py index 37c09433..d75c33a9 100644 --- a/cosmos_framework/model/vfm/diffusion/samplers/utils_test.py +++ b/cosmos_framework/model/generator/diffusion/samplers/utils_test.py @@ -3,7 +3,7 @@ import pytest -from cosmos_framework.model.vfm.diffusion.samplers.utils import run_multiseed +from cosmos_framework.model.generator.diffusion.samplers.utils import run_multiseed # --------------------------------------------------------------------------- # Single-call (no list kwargs) passthrough diff --git a/cosmos_framework/model/vfm/hf_model.py b/cosmos_framework/model/generator/hf_model.py similarity index 96% rename from cosmos_framework/model/vfm/hf_model.py rename to cosmos_framework/model/generator/hf_model.py index 7c3fa839..6482479b 100644 --- a/cosmos_framework/model/vfm/hf_model.py +++ b/cosmos_framework/model/generator/hf_model.py @@ -27,8 +27,8 @@ from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoModelForImageTextToText from cosmos_framework.utils import log -from cosmos_framework.model.vfm.utils.safetensors_loader import load_language_model, load_vlm_model -from cosmos_framework.utils.vfm.parallelism import ParallelDims +from cosmos_framework.model.generator.utils.safetensors_loader import load_language_model, load_vlm_model +from cosmos_framework.utils.generator.parallelism import ParallelDims def _tensor_names_to_skip_for(model_type: str) -> list[str]: @@ -103,7 +103,7 @@ def __init__( if attn_implementation == "cosmos": from transformers import AttentionInterface - from cosmos_framework.utils.vfm.hf_attention_cosmos import hf_attention_cosmos + from cosmos_framework.utils.generator.hf_attention_cosmos import hf_attention_cosmos AttentionInterface.register("cosmos", hf_attention_cosmos) @@ -163,7 +163,7 @@ def __init__( # image and slices the output to [0:0] so it contributes no features. # Must happen BEFORE parallelize() so FSDP captures the patched forward. if hf_config.model_type == "qwen3_vl" and hasattr(self.model, "model"): - from cosmos_framework.utils.vfm.monkey_patch import patch_qwen3_vl_forward + from cosmos_framework.utils.generator.monkey_patch import patch_qwen3_vl_forward patch_qwen3_vl_forward(self.model.model) log.info("HFModel: applied patch_qwen3_vl_forward for text-only batch support") @@ -260,8 +260,8 @@ def load_weights( ``org/model`` repo IDs fall back to Hugging Face. credential_path: S3 credential file, or None for local/HF. parallel_dims: ``ParallelDims`` instance (from - ``cosmos_framework.utils.vfm.parallelism``). The loader uses - it via :func:`~cosmos_framework.model.vfm.utils.safetensors_loader._get_dp_shard_mesh` + ``cosmos_framework.utils.generator.parallelism``). The loader uses + it via :func:`~cosmos_framework.model.generator.utils.safetensors_loader._get_dp_shard_mesh` to obtain the 1-D ``dp_shard`` sub-mesh (or None when ``dp_shard <= 1``) for striping checkpoint reads across FSDP shard ranks. When non-None, the caller MUST have diff --git a/cosmos_framework/model/vfm/mot/__init__.py b/cosmos_framework/model/generator/mot/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/mot/__init__.py rename to cosmos_framework/model/generator/mot/__init__.py diff --git a/cosmos_framework/model/vfm/mot/attention.py b/cosmos_framework/model/generator/mot/attention.py similarity index 99% rename from cosmos_framework/model/vfm/mot/attention.py rename to cosmos_framework/model/generator/mot/attention.py index babcdd94..8d7bbfee 100644 --- a/cosmos_framework/model/vfm/mot/attention.py +++ b/cosmos_framework/model/generator/mot/attention.py @@ -9,7 +9,7 @@ multi_dimensional_attention_varlen, ) from cosmos_framework.model.attention.masks import CausalType -from cosmos_framework.model.vfm.utils.memory import KVToStore, MemoryValue +from cosmos_framework.model.generator.utils.memory import KVToStore, MemoryValue class SplitInfo: @@ -71,11 +71,11 @@ def __init__( _dotproduct_attention_cache = {} -from cosmos_framework.data.vfm.sequence_packing.natten import ( +from cosmos_framework.data.generator.sequence_packing.natten import ( generate_natten_metadata, generate_temporal_causal_natten_metadata, ) -from cosmos_framework.data.vfm.sequence_packing.runtime import ( +from cosmos_framework.data.generator.sequence_packing.runtime import ( SequencePack, from_mode_splits, get_all_seq, diff --git a/cosmos_framework/model/vfm/mot/attention_test.py b/cosmos_framework/model/generator/mot/attention_test.py similarity index 98% rename from cosmos_framework/model/vfm/mot/attention_test.py rename to cosmos_framework/model/generator/mot/attention_test.py index 53112eed..77c32198 100644 --- a/cosmos_framework/model/vfm/mot/attention_test.py +++ b/cosmos_framework/model/generator/mot/attention_test.py @@ -7,12 +7,12 @@ import pytest import torch -import cosmos_framework.model.vfm.mot.attention as attention +import cosmos_framework.model.generator.mot.attention as attention from cosmos_framework.model.attention.natten import NATTEN_SUPPORTED -from cosmos_framework.model.vfm.mot.attention import ( +from cosmos_framework.model.generator.mot.attention import ( build_packed_sequence, ) -from cosmos_framework.data.vfm.sequence_packing.runtime import ( +from cosmos_framework.data.generator.sequence_packing.runtime import ( get_all_seq, get_gen_seq, get_und_seq, diff --git a/cosmos_framework/model/vfm/mot/cfgp_ar_test.py b/cosmos_framework/model/generator/mot/cfgp_ar_test.py similarity index 97% rename from cosmos_framework/model/vfm/mot/cfgp_ar_test.py rename to cosmos_framework/model/generator/mot/cfgp_ar_test.py index b0fd2a78..2c6bcef6 100644 --- a/cosmos_framework/model/vfm/mot/cfgp_ar_test.py +++ b/cosmos_framework/model/generator/mot/cfgp_ar_test.py @@ -4,7 +4,7 @@ """Multi-rank tests for CFGP (CFG Parallelism) in AR inference. Run with: - torchrun --nproc_per_node=2 -m pytest cosmos_framework/model/vfm/mot/cfgp_ar_test.py -v -m L0 + torchrun --nproc_per_node=2 -m pytest cosmos_framework/model/generator/mot/cfgp_ar_test.py -v -m L0 """ import os @@ -13,7 +13,7 @@ import torch import torch.distributed as dist -from cosmos_framework.utils.vfm.parallelism import ParallelDims +from cosmos_framework.utils.generator.parallelism import ParallelDims def setup_distributed_environment(): @@ -197,7 +197,7 @@ def test_cfgp_seed_broadcast(): cfgp_rank = parallel_dims.cfgp_rank cfgp_size = parallel_dims.cfgp_size - from cosmos_framework.model.vfm.omni_mot_model import _broadcast_seed + from cosmos_framework.model.generator.omni_mot_model import _broadcast_seed # Ranks start with different local seeds local_seed = 42 if cfgp_rank == 0 else 99 diff --git a/cosmos_framework/model/vfm/mot/context_parallel_test.py b/cosmos_framework/model/generator/mot/context_parallel_test.py similarity index 97% rename from cosmos_framework/model/vfm/mot/context_parallel_test.py rename to cosmos_framework/model/generator/mot/context_parallel_test.py index 3ba64339..3114b244 100644 --- a/cosmos_framework/model/vfm/mot/context_parallel_test.py +++ b/cosmos_framework/model/generator/mot/context_parallel_test.py @@ -10,24 +10,24 @@ import torch import torch.distributed as dist -from cosmos_framework.data.vfm.joint_dataloader import IterativeJointDataLoader -from cosmos_framework.model.vfm.mot.attention import ( +from cosmos_framework.data.generator.joint_dataloader import IterativeJointDataLoader +from cosmos_framework.model.generator.mot.attention import ( SplitInfo, dispatch_attention, ) -from cosmos_framework.model.vfm.mot.context_parallel_utils import ( +from cosmos_framework.model.generator.mot.context_parallel_utils import ( context_parallel_attention, get_context_parallel_sharded_sequence, ) -from cosmos_framework.model.vfm.mot.parallelize_unified_mot import ARReplicatedIODispatch -from cosmos_framework.model.vfm.mot.unified_mot import _apply_head_sharded_o_proj -from cosmos_framework.model.vfm.utils.data_and_condition import GenerationDataClean -from cosmos_framework.data.vfm.sequence_packing import ( +from cosmos_framework.model.generator.mot.parallelize_unified_mot import ARReplicatedIODispatch +from cosmos_framework.model.generator.mot.unified_mot import _apply_head_sharded_o_proj +from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean +from cosmos_framework.data.generator.sequence_packing import ( PackedSequence, build_sequence_plans_from_data_batch, pack_input_sequence, ) -from cosmos_framework.data.vfm.sequence_packing.runtime import ( +from cosmos_framework.data.generator.sequence_packing.runtime import ( SequencePack, from_all_seq, from_mode_splits, @@ -38,7 +38,7 @@ set_gen_seq, set_und_seq, ) -from cosmos_framework.utils.vfm.parallelism import ParallelDims +from cosmos_framework.utils.generator.parallelism import ParallelDims def _broadcast_test_object(data: Any, parallel_dims: ParallelDims, iteration: int) -> Any: diff --git a/cosmos_framework/model/vfm/mot/context_parallel_utils.py b/cosmos_framework/model/generator/mot/context_parallel_utils.py similarity index 98% rename from cosmos_framework/model/vfm/mot/context_parallel_utils.py rename to cosmos_framework/model/generator/mot/context_parallel_utils.py index 0195b835..fdc17fb2 100644 --- a/cosmos_framework/model/vfm/mot/context_parallel_utils.py +++ b/cosmos_framework/model/generator/mot/context_parallel_utils.py @@ -38,9 +38,9 @@ from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor, Replicate, Shard -from cosmos_framework.model.vfm.mot.attention import SplitInfo -from cosmos_framework.model.vfm.utils.memory import KVToStore, MemoryValue -from cosmos_framework.data.vfm.sequence_packing.runtime import ( +from cosmos_framework.model.generator.mot.attention import SplitInfo +from cosmos_framework.model.generator.utils.memory import KVToStore, MemoryValue +from cosmos_framework.data.generator.sequence_packing.runtime import ( SequencePack, from_mode_splits, get_all_seq, @@ -51,7 +51,7 @@ get_und_position_ids, get_und_seq, ) -from cosmos_framework.utils.vfm.parallelism import ParallelDims +from cosmos_framework.utils.generator.parallelism import ParallelDims def _pad_to_N(N, x: torch.Tensor) -> torch.Tensor: diff --git a/cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py b/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py similarity index 98% rename from cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py rename to cosmos_framework/model/generator/mot/cosmos3_vfm_network.py index 0a326996..39245765 100644 --- a/cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py +++ b/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py @@ -9,16 +9,16 @@ from transformers.configuration_utils import PretrainedConfig from transformers.modeling_utils import PreTrainedModel -from cosmos_framework.model.vfm.mot.attention import SplitInfo, build_packed_sequence -from cosmos_framework.model.vfm.mot.context_parallel_utils import ( +from cosmos_framework.model.generator.mot.attention import SplitInfo, build_packed_sequence +from cosmos_framework.model.generator.mot.context_parallel_utils import ( get_context_parallel_last_hidden_state, get_context_parallel_sharded_sequence, ) -from cosmos_framework.model.vfm.mot.domain_aware_linear import DomainAwareLinear -from cosmos_framework.model.vfm.mot.modeling_utils import TimestepEmbedder -from cosmos_framework.model.vfm.utils.memory import MemoryState -from cosmos_framework.data.vfm.sequence_packing import ModalityData, PackedSequence -from cosmos_framework.data.vfm.sequence_packing.natten import verify_natten_parameter_list +from cosmos_framework.model.generator.mot.domain_aware_linear import DomainAwareLinear +from cosmos_framework.model.generator.mot.modeling_utils import TimestepEmbedder +from cosmos_framework.model.generator.utils.memory import MemoryState +from cosmos_framework.data.generator.sequence_packing import ModalityData, PackedSequence +from cosmos_framework.data.generator.sequence_packing.natten import verify_natten_parameter_list class Cosmos3VFMNetworkConfig(PretrainedConfig): diff --git a/cosmos_framework/model/vfm/mot/domain_aware_linear.py b/cosmos_framework/model/generator/mot/domain_aware_linear.py similarity index 100% rename from cosmos_framework/model/vfm/mot/domain_aware_linear.py rename to cosmos_framework/model/generator/mot/domain_aware_linear.py diff --git a/cosmos_framework/model/vfm/mot/dot_product_attention.py b/cosmos_framework/model/generator/mot/dot_product_attention.py similarity index 100% rename from cosmos_framework/model/vfm/mot/dot_product_attention.py rename to cosmos_framework/model/generator/mot/dot_product_attention.py diff --git a/cosmos_framework/model/vfm/mot/modeling_utils.py b/cosmos_framework/model/generator/mot/modeling_utils.py similarity index 97% rename from cosmos_framework/model/vfm/mot/modeling_utils.py rename to cosmos_framework/model/generator/mot/modeling_utils.py index 2d543a86..037cd863 100644 --- a/cosmos_framework/model/vfm/mot/modeling_utils.py +++ b/cosmos_framework/model/generator/mot/modeling_utils.py @@ -7,7 +7,7 @@ from torch import nn from transformers.activations import ACT2FN -from cosmos_framework.data.vfm.sequence_packing import ModalityData +from cosmos_framework.data.generator.sequence_packing import ModalityData def has_noisy_tokens(modality_data: ModalityData | None) -> bool: diff --git a/cosmos_framework/model/vfm/mot/parallelize_unified_mot.py b/cosmos_framework/model/generator/mot/parallelize_unified_mot.py similarity index 97% rename from cosmos_framework/model/vfm/mot/parallelize_unified_mot.py rename to cosmos_framework/model/generator/mot/parallelize_unified_mot.py index a011e62f..d8f97a8c 100644 --- a/cosmos_framework/model/vfm/mot/parallelize_unified_mot.py +++ b/cosmos_framework/model/generator/mot/parallelize_unified_mot.py @@ -28,16 +28,16 @@ from cosmos_framework.utils import log from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig from cosmos_framework.configs.base.defaults.compile import CompileConfig -from cosmos_framework.model.vfm.mot.attention import SplitInfo, dispatch_attention -from cosmos_framework.model.vfm.mot.context_parallel_utils import context_parallel_attention -from cosmos_framework.model.vfm.utils.memory import KVToStore, MemoryValue -from cosmos_framework.data.vfm.sequence_packing.runtime import ( +from cosmos_framework.model.generator.mot.attention import SplitInfo, dispatch_attention +from cosmos_framework.model.generator.mot.context_parallel_utils import context_parallel_attention +from cosmos_framework.model.generator.utils.memory import KVToStore, MemoryValue +from cosmos_framework.data.generator.sequence_packing.runtime import ( SequencePack, from_und_gen_splits, get_gen_seq, get_und_seq, ) -from cosmos_framework.utils.vfm.parallelism import ParallelDims +from cosmos_framework.utils.generator.parallelism import ParallelDims class ContextParallelDispatch(nn.Module): diff --git a/cosmos_framework/model/vfm/mot/parallelize_vfm_network.py b/cosmos_framework/model/generator/mot/parallelize_vfm_network.py similarity index 97% rename from cosmos_framework/model/vfm/mot/parallelize_vfm_network.py rename to cosmos_framework/model/generator/mot/parallelize_vfm_network.py index 335234a5..e2279cfa 100644 --- a/cosmos_framework/model/vfm/mot/parallelize_vfm_network.py +++ b/cosmos_framework/model/generator/mot/parallelize_vfm_network.py @@ -6,8 +6,8 @@ from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig from cosmos_framework.configs.base.defaults.compile import CompileConfig -from cosmos_framework.model.vfm.mot.parallelize_unified_mot import parallelize_unified_mot -from cosmos_framework.utils.vfm.parallelism import ParallelDims +from cosmos_framework.model.generator.mot.parallelize_unified_mot import parallelize_unified_mot +from cosmos_framework.utils.generator.parallelism import ParallelDims def apply_compile(model: torch.nn.Module, config: CompileConfig): diff --git a/cosmos_framework/model/vfm/mot/qwen3_vl_unified_mot.py b/cosmos_framework/model/generator/mot/qwen3_vl_unified_mot.py similarity index 77% rename from cosmos_framework/model/vfm/mot/qwen3_vl_unified_mot.py rename to cosmos_framework/model/generator/mot/qwen3_vl_unified_mot.py index e6554e2f..1c25c453 100644 --- a/cosmos_framework/model/vfm/mot/qwen3_vl_unified_mot.py +++ b/cosmos_framework/model/generator/mot/qwen3_vl_unified_mot.py @@ -4,8 +4,8 @@ # Backward-compatibility shim: this module was renamed to unified_mot.py. # Existing serialized configs / checkpoints may reference the old module path, # so we re-export everything from the new location. -from cosmos_framework.model.vfm.mot.unified_mot import * # noqa: F401, F403 -from cosmos_framework.model.vfm.mot.unified_mot import ( # noqa: F401 # explicit re-exports for type checkers +from cosmos_framework.model.generator.mot.unified_mot import * # noqa: F401, F403 +from cosmos_framework.model.generator.mot.unified_mot import ( # noqa: F401 # explicit re-exports for type checkers LayerTypes, MoTDecoderLayer, Nemotron3DenseVLTextConfig, diff --git a/cosmos_framework/model/vfm/mot/und_k_norm_example_test.py b/cosmos_framework/model/generator/mot/und_k_norm_example_test.py similarity index 96% rename from cosmos_framework/model/vfm/mot/und_k_norm_example_test.py rename to cosmos_framework/model/generator/mot/und_k_norm_example_test.py index 05e6cc83..ca8c3ba0 100644 --- a/cosmos_framework/model/vfm/mot/und_k_norm_example_test.py +++ b/cosmos_framework/model/generator/mot/und_k_norm_example_test.py @@ -27,23 +27,23 @@ import pytest import torch -from cosmos_framework.model.vfm.mot.attention import ( +from cosmos_framework.model.generator.mot.attention import ( build_packed_sequence, three_way_attention, two_way_attention, ) -from cosmos_framework.model.vfm.mot.unified_mot import ( +from cosmos_framework.model.generator.mot.unified_mot import ( LayerTypes, PackedAttentionMoT, ) -from cosmos_framework.model.vfm.reasoner.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( +from cosmos_framework.model.generator.reasoner.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( Nemotron3DenseVLTextConfig, ) -from cosmos_framework.model.vfm.reasoner.nemotron_3_dense_vl.nemotron_3_dense_vl import ( +from cosmos_framework.model.generator.reasoner.nemotron_3_dense_vl.nemotron_3_dense_vl import ( Nemotron3DenseVLRMSNorm, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig -from cosmos_framework.data.vfm.sequence_packing.runtime import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig +from cosmos_framework.data.generator.sequence_packing.runtime import ( get_gen_seq, get_und_seq, ) diff --git a/cosmos_framework/model/vfm/mot/unified_mot.py b/cosmos_framework/model/generator/mot/unified_mot.py similarity index 99% rename from cosmos_framework/model/vfm/mot/unified_mot.py rename to cosmos_framework/model/generator/mot/unified_mot.py index 7d5f9757..2bdebfdc 100644 --- a/cosmos_framework/model/vfm/mot/unified_mot.py +++ b/cosmos_framework/model/generator/mot/unified_mot.py @@ -14,16 +14,16 @@ from cosmos_framework.model.attention import attention as imaginaire_attention from cosmos_framework.model.attention.masks import CausalType from cosmos_framework.utils import log -from cosmos_framework.model.vfm.mot.attention import ( +from cosmos_framework.model.generator.mot.attention import ( AttentionMaskType, dispatch_attention, ) # Nemotron 3 Dense VL imports -from cosmos_framework.model.vfm.reasoner.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( +from cosmos_framework.model.generator.reasoner.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( Nemotron3DenseVLTextConfig, ) -from cosmos_framework.model.vfm.reasoner.nemotron_3_dense_vl.nemotron_3_dense_vl import ( +from cosmos_framework.model.generator.reasoner.nemotron_3_dense_vl.nemotron_3_dense_vl import ( MultiModalRotaryEmbedding, Nemotron3DenseVLMLP, Nemotron3DenseVLPreTrainedModel, @@ -32,32 +32,32 @@ ) # Qwen3-VL imports -from cosmos_framework.model.vfm.reasoner.qwen3_vl.configuration_qwen3_vl import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.configuration_qwen3_vl import ( Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl.qwen3_vl import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.qwen3_vl import ( Qwen3VLPreTrainedModel, Qwen3VLTextMLP, Qwen3VLTextRMSNorm, Qwen3VLTextRotaryEmbedding, Qwen3VLVisionModel, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl.qwen3_vl import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.qwen3_vl import ( apply_rotary_pos_emb as qwen3_vl_apply_rotary_pos_emb, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import ( prepare_multimodal_reasoner_inputs, ) # Qwen3-VL-MoE imports -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( Qwen3VLMoeConfig, Qwen3VLMoeTextConfig, Qwen3VLMoeVisionConfig, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.qwen3_vl_moe import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import ( LBLMetadata, Qwen3VLMoePreTrainedModel, Qwen3VLMoeTextMLP, @@ -66,8 +66,8 @@ Qwen3VLMoeTextSparseMoeBlock, Qwen3VLMoeVisionModel, ) -from cosmos_framework.model.vfm.utils.memory import KVToStore, MemoryState, MemoryValue -from cosmos_framework.data.vfm.sequence_packing.runtime import ( +from cosmos_framework.model.generator.utils.memory import KVToStore, MemoryState, MemoryValue +from cosmos_framework.data.generator.sequence_packing.runtime import ( SequencePack, from_all_seq, from_und_gen_splits, diff --git a/cosmos_framework/model/vfm/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py similarity index 98% rename from cosmos_framework/model/vfm/omni_mot_model.py rename to cosmos_framework/model/generator/omni_mot_model.py index 19aac5f1..164ca3de 100644 --- a/cosmos_framework/model/vfm/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -23,44 +23,44 @@ from cosmos_framework.utils import log, misc from cosmos_framework.utils.count_params import count_params from cosmos_framework.utils.timer import Timer -from cosmos_framework.model.vfm.algorithm.loss.flow_matching import compute_flow_matching_loss -from cosmos_framework.model.vfm.algorithm.loss.load_balancing import compute_load_balancing_loss +from cosmos_framework.model.generator.algorithm.loss.flow_matching import compute_flow_matching_loss +from cosmos_framework.model.generator.algorithm.loss.load_balancing import compute_load_balancing_loss from cosmos_framework.configs.base.defaults.model_config import OmniMoTModelConfig -from cosmos_framework.data.vfm.action.action_processing import ActionProcessor, get_action_processing_records -from cosmos_framework.data.vfm.utils import IMAGE_RES_SIZE_INFO, VIDEO_RES_SIZE_INFO -from cosmos_framework.model.vfm.diffusion.rectified_flow import RectifiedFlow -from cosmos_framework.model.vfm.diffusion.samplers.edm import EDMSampler -from cosmos_framework.model.vfm.diffusion.samplers.fixed_step import FixedStepSampler -from cosmos_framework.model.vfm.diffusion.samplers.unipc import UniPCSampler, UniPCSamplerConfig -from cosmos_framework.model.vfm.mot.context_parallel_utils import context_parallel_broadcast_tensor_list -from cosmos_framework.model.vfm.mot.cosmos3_vfm_network import Cosmos3VFMNetwork, Cosmos3VFMNetworkConfig -from cosmos_framework.model.vfm.mot.modeling_utils import has_noisy_tokens -from cosmos_framework.model.vfm.mot.parallelize_vfm_network import parallelize_vfm_network -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption -from cosmos_framework.model.vfm.utils.data_and_condition import ( +from cosmos_framework.data.generator.action.action_processing import ActionProcessor, get_action_processing_records +from cosmos_framework.data.generator.utils import IMAGE_RES_SIZE_INFO, VIDEO_RES_SIZE_INFO +from cosmos_framework.model.generator.diffusion.rectified_flow import RectifiedFlow +from cosmos_framework.model.generator.diffusion.samplers.edm import EDMSampler +from cosmos_framework.model.generator.diffusion.samplers.fixed_step import FixedStepSampler +from cosmos_framework.model.generator.diffusion.samplers.unipc import UniPCSampler, UniPCSamplerConfig +from cosmos_framework.model.generator.mot.context_parallel_utils import context_parallel_broadcast_tensor_list +from cosmos_framework.model.generator.mot.cosmos3_vfm_network import Cosmos3VFMNetwork, Cosmos3VFMNetworkConfig +from cosmos_framework.model.generator.mot.modeling_utils import has_noisy_tokens +from cosmos_framework.model.generator.mot.parallelize_vfm_network import parallelize_vfm_network +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.generator.utils.data_and_condition import ( GenerationDataClean, GenerationDataNoised, _expand_per_sample_to_per_vision_item, build_dense_sound_schedule, unwrap_and_densify, ) -from cosmos_framework.model.vfm.utils.memory import MemoryState -from cosmos_framework.model.vfm.utils.safetensors_loader import ( +from cosmos_framework.model.generator.utils.memory import MemoryState +from cosmos_framework.model.generator.utils.safetensors_loader import ( load_language_model as load_language_model_safetensors, ) -from cosmos_framework.data.vfm.sequence_packing import ( +from cosmos_framework.data.generator.sequence_packing import ( PackedSequence, SequencePlan, build_sequence_plans_from_data_batch, pack_input_sequence, ) -from cosmos_framework.data.vfm.sequence_packing.modalities import add_special_tokens -from cosmos_framework.model.vfm.tokenizers.interface import VideoTokenizerInterface -from cosmos_framework.model.vfm.upsampler.prompts import build_messages, clean_response -from cosmos_framework.utils.vfm.data_utils import get_vision_data_resolution -from cosmos_framework.utils.vfm.dtensor_helper import DTensorFastEmaModelUpdater -from cosmos_framework.utils.vfm.model_weights_stats import WeightTrainingStat -from cosmos_framework.utils.vfm.parallelism import ParallelDims +from cosmos_framework.data.generator.sequence_packing.modalities import add_special_tokens +from cosmos_framework.model.generator.tokenizers.interface import VideoTokenizerInterface +from cosmos_framework.model.generator.upsampler.prompts import build_messages, clean_response +from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution +from cosmos_framework.utils.generator.dtensor_helper import DTensorFastEmaModelUpdater +from cosmos_framework.utils.generator.model_weights_stats import WeightTrainingStat +from cosmos_framework.utils.generator.parallelism import ParallelDims class OmniMoTModel(ImaginaireModel): @@ -2529,7 +2529,7 @@ def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: # run. Avoids two EDM-specific footguns: # (1) ``EDMSampler._forward_impl`` always runs an extra # ``sample_clean`` denoiser forward (see - # ``cosmos_framework/model/vfm/diffusion/samplers/edm.py``). + # ``cosmos_framework/model/generator/diffusion/samplers/edm.py``). # A nested sampler call would add one too many # forwards on fast ranks, since the slow rank's # single call also pays the ``sample_clean`` cost. @@ -3314,7 +3314,7 @@ def _extract_upsample_video_specs( tensor (``shape[-1]`` for width, ``shape[-2]`` for height), and the ``aspect_ratio`` string is reverse-looked-up against the canonical ``{IMAGE,VIDEO}_RES_SIZE_INFO`` tables in - :mod:`cosmos_framework.data.vfm.utils` — image table for + :mod:`cosmos_framework.data.generator.utils` — image table for ``"t2i"``, video table otherwise. Note these tables are ``{res: {ar: (W, H)}}`` (the first entry is *width*); the existing logging-only lookup in @@ -3329,7 +3329,7 @@ def _extract_upsample_video_specs( where ``num_frames`` is the temporal dimension (``shape[-3]``) of the same vision tensor. For ``"t2i"`` both fields are returned as ``None`` so - :func:`cosmos_framework.model.vfm.upsampler.prompts.build_user_text`'s + :func:`cosmos_framework.model.generator.upsampler.prompts.build_user_text`'s ``t2i``-must-have-no-video-args contract is satisfied. Args: @@ -3385,7 +3385,7 @@ def _extract_upsample_video_specs( f"Cannot infer aspect_ratio for upsample task={task!r}: vision tensor has " f"(W={w}, H={h}) which is not a canonical entry in {res_info_name}. " "Use a supported (resolution, aspect_ratio) combination from " - "cosmos_framework.data.vfm.utils, or pass explicit specs to " + "cosmos_framework.data.generator.utils, or pass explicit specs to " "`upsample_captions(...)`." ) @@ -3404,7 +3404,7 @@ def _extract_upsample_video_specs( raise ValueError(f"upsample task={task!r}: conditioning_fps must be positive; got {fps_int}.") num_frames = int(sample.shape[-3]) # Integer-floor seconds matches the canonical V4.2 ``M:SS`` rendering - # in :func:`cosmos_framework.model.vfm.upsampler.prompts._format_duration`, + # in :func:`cosmos_framework.model.generator.upsampler.prompts._format_duration`, # which expects an int and rejects fractional seconds. duration_secs = max(1, num_frames // fps_int) return aspect_ratio, w, h, fps_int, duration_secs @@ -3917,7 +3917,7 @@ def generate_reasoner_text( ``pixel_values_videos`` / ``video_grid_thw``. prompt_builder: Optional callback that maps a raw prompt string to a chat-style messages list (e.g. - :func:`cosmos_framework.model.vfm.upsampler.prompts.build_messages` + :func:`cosmos_framework.model.generator.upsampler.prompts.build_messages` for V4.2 caption upsampling). When ``None``, prompts are wrapped as ``[{"role": "user", "content": prompt}]`` with no system message. @@ -4189,7 +4189,7 @@ def upsample_captions( prompt-driven branch. The only thing this method adds on top of the generic per-prompt loop is the V4.2 chat-template injection: each caption is wrapped via - :func:`cosmos_framework.model.vfm.upsampler.prompts.build_messages` + :func:`cosmos_framework.model.generator.upsampler.prompts.build_messages` (which returns ``[system, user]`` with the user content embedding the caption inside the canonical V4.2 template — instructions, task constraints, and output JSON schema for the requested task). @@ -4212,7 +4212,7 @@ def upsample_captions( position ids) before kicking off the AR decode loop. Each raw reasoner output is post-processed by - :func:`cosmos_framework.model.vfm.upsampler.prompts.clean_response` + :func:`cosmos_framework.model.generator.upsampler.prompts.clean_response` before being returned. The cleaner strips ```` / ```` / ```` / etc. reasoning blocks and any prose preamble that appears before the @@ -4264,7 +4264,7 @@ def upsample_captions( fps: Target frames-per-second for the generated clip. Required for the video tasks (``"t2v"``, ``"i2v"``) and must be ``None`` for ``"t2i"`` — the underlying - :func:`cosmos_framework.model.vfm.upsampler.prompts.build_user_text` + :func:`cosmos_framework.model.generator.upsampler.prompts.build_user_text` raises ``ValueError`` if a video task is missing ``fps`` or ``duration_secs``. duration_secs: Clip duration in whole seconds (rendered as @@ -4490,7 +4490,7 @@ def add_lora( lora_target_modules: str, ) -> torch.nn.Module: """Pre-FSDP LoRA injection — see :func:`inject_lora_pre_fsdp` for details.""" - from cosmos_framework.utils.vfm.lora import inject_lora_pre_fsdp + from cosmos_framework.utils.generator.lora import inject_lora_pre_fsdp self.lora_alpha = lora_alpha return inject_lora_pre_fsdp( @@ -4502,7 +4502,7 @@ def add_lora( def _init_lora_weights_post_materialization(self, network: torch.nn.Module) -> None: """Post-materialization LoRA init — see :func:`init_lora_weights_post_materialization`.""" - from cosmos_framework.utils.vfm.lora import init_lora_weights_post_materialization + from cosmos_framework.utils.generator.lora import init_lora_weights_post_materialization init_lora_weights_post_materialization(network) diff --git a/cosmos_framework/model/vfm/parallelize_vlm.py b/cosmos_framework/model/generator/parallelize_vlm.py similarity index 97% rename from cosmos_framework/model/vfm/parallelize_vlm.py rename to cosmos_framework/model/generator/parallelize_vlm.py index c91e48db..57a80c59 100644 --- a/cosmos_framework/model/vfm/parallelize_vlm.py +++ b/cosmos_framework/model/generator/parallelize_vlm.py @@ -4,12 +4,12 @@ """FSDP2 wrapping for Cosmos3 VLM ``HFModel`` instances. Hosts the single VLM-specific ``parallelize`` entry point used by -``vlm_model.VLMModel._init_vlm``. Lives under ``cosmos_framework/model/vfm/`` +``vlm_model.VLMModel._init_vlm``. Lives under ``cosmos_framework/model/generator/`` so the FSDP wrapping concern sits next to the model class it operates on (mirroring the layout of ``models/mot/parallelize_unified_mot.py`` for the MoT path). -Pure parallelism plumbing — :class:`~cosmos_framework.utils.vfm.parallelism.ParallelDims` +Pure parallelism plumbing — :class:`~cosmos_framework.utils.generator.parallelism.ParallelDims` and its meshes — stays in ``vfm/utils/parallelism.py``. """ @@ -22,8 +22,8 @@ PRECISION_TO_TORCH_DTYPE, ParallelismConfig, ) -from cosmos_framework.model.vfm.hf_model import HFModel -from cosmos_framework.utils.vfm.parallelism import ParallelDims +from cosmos_framework.model.generator.hf_model import HFModel +from cosmos_framework.utils.generator.parallelism import ParallelDims def _collect_repeated_blocks(inner: nn.Module) -> tuple[list[nn.Module], set[str]]: diff --git a/cosmos_framework/model/vfm/reasoner/__init__.py b/cosmos_framework/model/generator/reasoner/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/__init__.py rename to cosmos_framework/model/generator/reasoner/__init__.py diff --git a/cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/__init__.py b/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/__init__.py rename to cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/__init__.py diff --git a/cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json b/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json similarity index 100% rename from cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json rename to cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json diff --git a/cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py b/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py rename to cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py diff --git a/cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py b/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py similarity index 98% rename from cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py rename to cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py index eb91cb93..6fc56f7e 100644 --- a/cosmos_framework/model/vfm/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py +++ b/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py @@ -17,7 +17,7 @@ from transformers.modeling_rope_utils import dynamic_rope_update from transformers.modeling_utils import PreTrainedModel -from cosmos_framework.model.vfm.reasoner.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( +from cosmos_framework.model.generator.reasoner.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import ( Nemotron3DenseVLTextConfig, ) diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl/__init__.py b/cosmos_framework/model/generator/reasoner/qwen3_vl/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl/__init__.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl/__init__.py diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json b/cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json rename to cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json b/cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json rename to cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json b/cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json rename to cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json b/cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json rename to cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/__init__.py b/cosmos_framework/model/generator/reasoner/qwen3_vl/configs/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl/configs/__init__.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl/configs/__init__.py diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl/configuration_qwen3_vl.py b/cosmos_framework/model/generator/reasoner/qwen3_vl/configuration_qwen3_vl.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl/configuration_qwen3_vl.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl/configuration_qwen3_vl.py diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl/qwen3_vl.py b/cosmos_framework/model/generator/reasoner/qwen3_vl/qwen3_vl.py similarity index 99% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl/qwen3_vl.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl/qwen3_vl.py index 38e784cc..d442752a 100644 --- a/cosmos_framework/model/vfm/reasoner/qwen3_vl/qwen3_vl.py +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl/qwen3_vl.py @@ -40,16 +40,16 @@ def _default_rope_init(config, device=None, **kwargs): from transformers.utils.deprecation import deprecate_kwarg # Import masking functions from utils for compatibility -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import ( create_causal_mask, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import ( get_image_features as _get_image_features, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import ( get_placeholder_mask as _get_placeholder_mask, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import ( get_rope_index as _get_rope_index, ) diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl/utils.py b/cosmos_framework/model/generator/reasoner/qwen3_vl/utils.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl/utils.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl/utils.py diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl/video_processing_qwen3_vl.py b/cosmos_framework/model/generator/reasoner/qwen3_vl/video_processing_qwen3_vl.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl/video_processing_qwen3_vl.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl/video_processing_qwen3_vl.py diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/__init__.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/__init__.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl_moe/__init__.py diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json rename to cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json rename to cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/__init__.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configs/__init__.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/__init__.py diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe.py similarity index 98% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe.py index 9c463204..e0faf542 100644 --- a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe.py +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe.py @@ -7,10 +7,10 @@ import torch.nn as nn from transformers.activations import ACT2FN -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( Qwen3VLMoeTextConfig, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_kernels import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe_kernels import ( TOKEN_GROUP_ALIGN_SIZE_M, _generate_permute_indices, ) diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_bench.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_bench.py similarity index 93% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_bench.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_bench.py index 48394820..d7630fad 100644 --- a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_bench.py +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_bench.py @@ -5,30 +5,30 @@ Usage: # Default benchmark (forward only, compiled, bf16): - python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench + python -m cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe_bench # Forward + backward: - python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --backward + python -m cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe_bench --backward # Compare grouped_mm vs naive: - python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --compare + python -m cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe_bench --compare # Disable torch.compile: - python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --no-compile + python -m cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe_bench --no-compile # Custom sweep: - python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --backward \ + python -m cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe_bench --backward \ --hidden-size 4096 \ --moe-intermediate-size 1536 # Capture a torch profiler trace (Chrome trace JSON): - python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --profile + python -m cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe_bench --profile # Profile to a custom directory: - python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --profile --profile-dir ./my_traces + python -m cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe_bench --profile --profile-dir ./my_traces # All options: - python -m cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe_bench --help + python -m cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe_bench --help """ import itertools @@ -40,10 +40,10 @@ import torch import tyro -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( Qwen3VLMoeTextConfig, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import ( Qwen3VLMoeTextExpertsGroupedMm, Qwen3VLMoeTextExpertsNaive, ) diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_kernels.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_kernels.py similarity index 100% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_kernels.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_kernels.py diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_test.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py similarity index 93% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_test.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py index 37430fec..bc570db3 100644 --- a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/moe_test.py +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py @@ -6,10 +6,10 @@ import torch from torch import nn -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( Qwen3VLMoeTextConfig, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe import create_text_experts +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import create_text_experts def run_moe(mod: nn.Module, hidden_states: torch.Tensor, topk_scores: torch.Tensor, expert_indices: torch.Tensor): diff --git a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/qwen3_vl_moe.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/qwen3_vl_moe.py similarity index 99% rename from cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/qwen3_vl_moe.py rename to cosmos_framework/model/generator/reasoner/qwen3_vl_moe/qwen3_vl_moe.py index f25506b4..0add9ba7 100644 --- a/cosmos_framework/model/vfm/reasoner/qwen3_vl_moe/qwen3_vl_moe.py +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/qwen3_vl_moe.py @@ -22,21 +22,21 @@ from transformers.utils import TransformersKwargs, is_torchdynamo_compiling from transformers.utils.deprecation import deprecate_kwarg -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import ( get_image_features as _get_image_features, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import ( get_placeholder_mask as _get_placeholder_mask, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import ( get_rope_index as _get_rope_index, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( Qwen3VLMoeConfig, Qwen3VLMoeTextConfig, Qwen3VLMoeVisionConfig, ) -from cosmos_framework.model.vfm.reasoner.qwen3_vl_moe.moe import ( +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import ( create_text_experts, ) diff --git a/cosmos_framework/model/vfm/tokenizers/__init__.py b/cosmos_framework/model/generator/tokenizers/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/__init__.py rename to cosmos_framework/model/generator/tokenizers/__init__.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/__init__.py b/cosmos_framework/model/generator/tokenizers/audio/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/__init__.py rename to cosmos_framework/model/generator/tokenizers/audio/__init__.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae.py b/cosmos_framework/model/generator/tokenizers/audio/avae.py similarity index 98% rename from cosmos_framework/model/vfm/tokenizers/audio/avae.py rename to cosmos_framework/model/generator/tokenizers/audio/avae.py index d28e9108..293b2f34 100644 --- a/cosmos_framework/model/vfm/tokenizers/audio/avae.py +++ b/cosmos_framework/model/generator/tokenizers/audio/avae.py @@ -14,9 +14,9 @@ from cosmos_framework.utils.flags import DEVICE, INTERNAL from cosmos_framework.utils import log from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.model.vfm.tokenizers.audio.avae_utils.env import AttrDict -from cosmos_framework.model.vfm.tokenizers.audio.avae_utils.models import load_generator -from cosmos_framework.model.vfm.tokenizers.interface import AudioTokenizerInterface +from cosmos_framework.model.generator.tokenizers.audio.avae_utils.env import AttrDict +from cosmos_framework.model.generator.tokenizers.audio.avae_utils.models import load_generator +from cosmos_framework.model.generator.tokenizers.interface import AudioTokenizerInterface def _load_avae_model( diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/__init__.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/__init__.py similarity index 64% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/__init__.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/__init__.py index 9f1cdafc..2bae7776 100644 --- a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/__init__.py +++ b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/__init__.py @@ -8,8 +8,8 @@ oobleck decoder, and VAE bottleneck configuration. """ -from cosmos_framework.model.vfm.tokenizers.audio.avae_utils.env import AttrDict -from cosmos_framework.model.vfm.tokenizers.audio.avae_utils.models import LatentAutoEncoderV2, load_generator +from cosmos_framework.model.generator.tokenizers.audio.avae_utils.env import AttrDict +from cosmos_framework.model.generator.tokenizers.audio.avae_utils.models import LatentAutoEncoderV2, load_generator __all__ = [ "LatentAutoEncoderV2", diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/activations.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/activations.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/activations.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/activations.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/alias_free_torch/__init__.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/alias_free_torch/__init__.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/__init__.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/alias_free_torch/act.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/act.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/alias_free_torch/act.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/act.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/alias_free_torch/filter.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/filter.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/alias_free_torch/filter.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/filter.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/alias_free_torch/resample.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/resample.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/alias_free_torch/resample.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/resample.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/bottlenecks.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/bottlenecks.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/bottlenecks.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/bottlenecks.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/env.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/env.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/env.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/env.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/models.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/models.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/models.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/models.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/modules.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/modules.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/modules.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/modules.py diff --git a/cosmos_framework/model/vfm/tokenizers/audio/avae_utils/modules_encodec.py b/cosmos_framework/model/generator/tokenizers/audio/avae_utils/modules_encodec.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/audio/avae_utils/modules_encodec.py rename to cosmos_framework/model/generator/tokenizers/audio/avae_utils/modules_encodec.py diff --git a/cosmos_framework/model/vfm/tokenizers/dc_ae/__init__.py b/cosmos_framework/model/generator/tokenizers/dc_ae/__init__.py similarity index 83% rename from cosmos_framework/model/vfm/tokenizers/dc_ae/__init__.py rename to cosmos_framework/model/generator/tokenizers/dc_ae/__init__.py index 3f85ab5a..66eecb5b 100644 --- a/cosmos_framework/model/vfm/tokenizers/dc_ae/__init__.py +++ b/cosmos_framework/model/generator/tokenizers/dc_ae/__init__.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -from cosmos_framework.model.vfm.tokenizers.dc_ae.dc_ae_v import ( +from cosmos_framework.model.generator.tokenizers.dc_ae.dc_ae_v import ( DCAEV, CompilableDCAEVEncoder, DCAEVConfig, diff --git a/cosmos_framework/model/vfm/tokenizers/dc_ae/convert_pt_to_distcp.py b/cosmos_framework/model/generator/tokenizers/dc_ae/convert_pt_to_distcp.py similarity index 92% rename from cosmos_framework/model/vfm/tokenizers/dc_ae/convert_pt_to_distcp.py rename to cosmos_framework/model/generator/tokenizers/dc_ae/convert_pt_to_distcp.py index 67ddc380..a01df405 100644 --- a/cosmos_framework/model/vfm/tokenizers/dc_ae/convert_pt_to_distcp.py +++ b/cosmos_framework/model/generator/tokenizers/dc_ae/convert_pt_to_distcp.py @@ -28,7 +28,7 @@ def main(): """ -python cosmos_framework/model/vfm/tokenizers/dc_ae/convert_pt_to_distcp.py \ +python cosmos_framework/model/generator/tokenizers/dc_ae/convert_pt_to_distcp.py \ input_path=checkpoints/cosmos_8b_wan22_vae/iter_000020000/model_ema_fp32.pt \ output_dir=checkpoints/cosmos_8b_wan22_vae/iter_000020000/model_dcp_from_torch_save """ diff --git a/cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_4x32x32.py b/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_4x32x32.py similarity index 96% rename from cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_4x32x32.py rename to cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_4x32x32.py index 0a1f76ac..9b93c259 100644 --- a/cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_4x32x32.py +++ b/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_4x32x32.py @@ -8,13 +8,13 @@ from cosmos_framework.utils import log from cosmos_framework.utils.distributed import get_rank, sync_model_states from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.data.vfm.utils import VIDEO_RES_SIZE_INFO -from cosmos_framework.model.vfm.tokenizers.dc_ae.dc_ae_v import ( +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.model.generator.tokenizers.dc_ae.dc_ae_v import ( DCAEV, DCAEVConfig, dc_ae_v_f32t4_encoder_causal_decoder_chunk_causal_4, ) -from cosmos_framework.model.vfm.tokenizers.interface import VideoTokenizerInterface +from cosmos_framework.model.generator.tokenizers.interface import VideoTokenizerInterface DEFAULT_MODEL_NAME = "dcae4x32x32_c64_t120_256p_fps_all_encoder_causal_decoder_chunk_causal_4_nogan_cosmos_pad_7_v0.2" diff --git a/cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_v.py b/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v.py similarity index 99% rename from cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_v.py rename to cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v.py index 976f7021..708039dd 100644 --- a/cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_v.py +++ b/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v.py @@ -10,7 +10,7 @@ from omegaconf import MISSING from tqdm import tqdm -from cosmos_framework.model.vfm.tokenizers.dc_ae.dc_ae_v_ops import ( +from cosmos_framework.model.generator.tokenizers.dc_ae.dc_ae_v_ops import ( ChannelDuplicatingPixelShuffleUpSampleLayer3d, CompilableOpSequential3d, CompilableRMSNorm2d, diff --git a/cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_v_ops.py b/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v_ops.py similarity index 99% rename from cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_v_ops.py rename to cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v_ops.py index 6f71d995..2ecb8dac 100644 --- a/cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_v_ops.py +++ b/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v_ops.py @@ -61,7 +61,7 @@ def zero_out(self): nn.init.constant_(self.bias, 0) def forward(self, x: torch.Tensor) -> torch.Tensor: - from cosmos_framework.model.vfm.tokenizers.dc_ae.dc_ae_v_triton_rms_norm import TritonRMSNorm2dFunc + from cosmos_framework.model.generator.tokenizers.dc_ae.dc_ae_v_triton_rms_norm import TritonRMSNorm2dFunc if not torch.compiler.is_compiling(): input_numel = x.numel() diff --git a/cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_v_triton_rms_norm.py b/cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v_triton_rms_norm.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/dc_ae/dc_ae_v_triton_rms_norm.py rename to cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v_triton_rms_norm.py diff --git a/cosmos_framework/model/vfm/tokenizers/flux_vae_8x8.py b/cosmos_framework/model/generator/tokenizers/flux_vae_8x8.py similarity index 99% rename from cosmos_framework/model/vfm/tokenizers/flux_vae_8x8.py rename to cosmos_framework/model/generator/tokenizers/flux_vae_8x8.py index 85188b45..78d0b5d2 100644 --- a/cosmos_framework/model/vfm/tokenizers/flux_vae_8x8.py +++ b/cosmos_framework/model/generator/tokenizers/flux_vae_8x8.py @@ -10,7 +10,7 @@ from torch import Tensor, nn from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.model.vfm.tokenizers.interface import VideoTokenizerInterface +from cosmos_framework.model.generator.tokenizers.interface import VideoTokenizerInterface @dataclass diff --git a/cosmos_framework/model/vfm/tokenizers/interface.py b/cosmos_framework/model/generator/tokenizers/interface.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/interface.py rename to cosmos_framework/model/generator/tokenizers/interface.py diff --git a/cosmos_framework/model/vfm/tokenizers/stable_diffusion_vae_8x8.py b/cosmos_framework/model/generator/tokenizers/stable_diffusion_vae_8x8.py similarity index 98% rename from cosmos_framework/model/vfm/tokenizers/stable_diffusion_vae_8x8.py rename to cosmos_framework/model/generator/tokenizers/stable_diffusion_vae_8x8.py index 2284230e..8b3d98be 100644 --- a/cosmos_framework/model/vfm/tokenizers/stable_diffusion_vae_8x8.py +++ b/cosmos_framework/model/generator/tokenizers/stable_diffusion_vae_8x8.py @@ -7,7 +7,7 @@ import torch -from cosmos_framework.model.vfm.tokenizers.interface import VideoTokenizerInterface +from cosmos_framework.model.generator.tokenizers.interface import VideoTokenizerInterface _DTYPE_BY_NAME: dict[str, torch.dtype] = { "float32": torch.float32, diff --git a/cosmos_framework/model/vfm/tokenizers/tokenization_qwen2.py b/cosmos_framework/model/generator/tokenizers/tokenization_qwen2.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/tokenization_qwen2.py rename to cosmos_framework/model/generator/tokenizers/tokenization_qwen2.py diff --git a/cosmos_framework/model/vfm/tokenizers/uniae/__init__.py b/cosmos_framework/model/generator/tokenizers/uniae/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/tokenizers/uniae/__init__.py rename to cosmos_framework/model/generator/tokenizers/uniae/__init__.py diff --git a/cosmos_framework/model/vfm/tokenizers/uniae/frame_math.py b/cosmos_framework/model/generator/tokenizers/uniae/frame_math.py similarity index 99% rename from cosmos_framework/model/vfm/tokenizers/uniae/frame_math.py rename to cosmos_framework/model/generator/tokenizers/uniae/frame_math.py index 9ab1eb98..f456cba2 100644 --- a/cosmos_framework/model/vfm/tokenizers/uniae/frame_math.py +++ b/cosmos_framework/model/generator/tokenizers/uniae/frame_math.py @@ -5,7 +5,7 @@ from collections.abc import Iterable, Mapping -from cosmos_framework.utils.vfm.data_utils import get_vision_data_resolution +from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution DEFAULT_RESOLUTION_KEYS = ("256", "480") diff --git a/cosmos_framework/model/vfm/tokenizers/uniae/frame_math_test.py b/cosmos_framework/model/generator/tokenizers/uniae/frame_math_test.py similarity index 93% rename from cosmos_framework/model/vfm/tokenizers/uniae/frame_math_test.py rename to cosmos_framework/model/generator/tokenizers/uniae/frame_math_test.py index b233ac1d..363fc4ea 100644 --- a/cosmos_framework/model/vfm/tokenizers/uniae/frame_math_test.py +++ b/cosmos_framework/model/generator/tokenizers/uniae/frame_math_test.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -from cosmos_framework.model.vfm.tokenizers.uniae.frame_math import ( +from cosmos_framework.model.generator.tokenizers.uniae.frame_math import ( align_uniae_num_video_frames, ceil_uniae_num_video_frames, ) diff --git a/cosmos_framework/model/vfm/tokenizers/uniae/noncausal_4x16x16.py b/cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py similarity index 98% rename from cosmos_framework/model/vfm/tokenizers/uniae/noncausal_4x16x16.py rename to cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py index de4cfbf7..a9e0dac7 100644 --- a/cosmos_framework/model/vfm/tokenizers/uniae/noncausal_4x16x16.py +++ b/cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py @@ -7,7 +7,7 @@ to provide a VideoTokenizerInterface compatible with diffusion model training. Usage: - from cosmos_framework.model.vfm.tokenizers.uniae.noncausal_4x16x16 import UniAEVAE + from cosmos_framework.model.generator.tokenizers.uniae.noncausal_4x16x16 import UniAEVAE vae = UniAEVAE( vae_pth="s3://bucket0/pretrained/tokenizers/video/cosmos/...", @@ -24,14 +24,14 @@ from cosmos_framework.utils import log from cosmos_framework.utils.distributed import get_rank, sync_model_states from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.model.vfm.tokenizers.interface import VideoTokenizerInterface -from cosmos_framework.model.vfm.tokenizers.uniae.frame_math import ( +from cosmos_framework.model.generator.tokenizers.interface import VideoTokenizerInterface +from cosmos_framework.model.generator.tokenizers.uniae.frame_math import ( get_uniae_latent_num_frames, get_uniae_latent_temporal_positions, get_uniae_pixel_num_frames, normalize_resolution_int_mapping, ) -from cosmos_framework.utils.vfm.data_utils import get_vision_data_resolution +from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution from cosmos_framework.model.tokenizer.models.dense_runtime import DenseAutoencoderRuntime from cosmos_framework.model.tokenizer.models.sparse_autoencoder import AutoencoderKL diff --git a/cosmos_framework/model/vfm/tokenizers/wan2pt1_vae_4x8x8.py b/cosmos_framework/model/generator/tokenizers/wan2pt1_vae_4x8x8.py similarity index 99% rename from cosmos_framework/model/vfm/tokenizers/wan2pt1_vae_4x8x8.py rename to cosmos_framework/model/generator/tokenizers/wan2pt1_vae_4x8x8.py index f1368a08..7f655e24 100644 --- a/cosmos_framework/model/vfm/tokenizers/wan2pt1_vae_4x8x8.py +++ b/cosmos_framework/model/generator/tokenizers/wan2pt1_vae_4x8x8.py @@ -11,7 +11,7 @@ from cosmos_framework.utils import log from cosmos_framework.utils.distributed import get_rank, sync_model_states from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.model.vfm.tokenizers.interface import VideoTokenizerInterface +from cosmos_framework.model.generator.tokenizers.interface import VideoTokenizerInterface # For sequential decoding, CACHE_T is the number of frames to cache. CACHE_T = 2 diff --git a/cosmos_framework/model/vfm/tokenizers/wan2pt2_vae_4x16x16.py b/cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py similarity index 99% rename from cosmos_framework/model/vfm/tokenizers/wan2pt2_vae_4x16x16.py rename to cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py index a7968307..fea6320d 100644 --- a/cosmos_framework/model/vfm/tokenizers/wan2pt2_vae_4x16x16.py +++ b/cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py @@ -15,9 +15,9 @@ from cosmos_framework.utils import log from cosmos_framework.utils.distributed import get_rank, sync_model_states from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.data.vfm.utils import VIDEO_RES_SIZE_INFO -from cosmos_framework.model.vfm.tokenizers.interface import VideoTokenizerInterface -from cosmos_framework.utils.vfm.data_utils import get_vision_data_resolution +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.model.generator.tokenizers.interface import VideoTokenizerInterface +from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution # For sequential decoding, CACHE_T is the number of frames to cache. CACHE_T = 2 diff --git a/cosmos_framework/model/vfm/upsampler/__init__.py b/cosmos_framework/model/generator/upsampler/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/upsampler/__init__.py rename to cosmos_framework/model/generator/upsampler/__init__.py diff --git a/cosmos_framework/model/vfm/upsampler/prompts.py b/cosmos_framework/model/generator/upsampler/prompts.py similarity index 99% rename from cosmos_framework/model/vfm/upsampler/prompts.py rename to cosmos_framework/model/generator/upsampler/prompts.py index 527fe27d..4d5b6a76 100644 --- a/cosmos_framework/model/vfm/upsampler/prompts.py +++ b/cosmos_framework/model/generator/upsampler/prompts.py @@ -9,7 +9,7 @@ Usage:: - from cosmos_framework.model.vfm.upsampler.prompts import ( + from cosmos_framework.model.generator.upsampler.prompts import ( build_user_text, build_messages, clean_response, is_upsampled_prompt, ) @@ -1218,7 +1218,7 @@ def is_upsampled_prompt(prompt: str) -> bool: # ----- Self-test ---------------------------------------------------------- # -# Run as ``python -m cosmos_framework.model.vfm.upsampler.prompts`` to verify that +# Run as ``python -m cosmos_framework.model.generator.upsampler.prompts`` to verify that # the templates round-trip and the ``{description}`` placeholder is honoured. diff --git a/cosmos_framework/model/vfm/utils/__init__.py b/cosmos_framework/model/generator/utils/__init__.py similarity index 100% rename from cosmos_framework/model/vfm/utils/__init__.py rename to cosmos_framework/model/generator/utils/__init__.py diff --git a/cosmos_framework/model/vfm/utils/data_and_condition.py b/cosmos_framework/model/generator/utils/data_and_condition.py similarity index 100% rename from cosmos_framework/model/vfm/utils/data_and_condition.py rename to cosmos_framework/model/generator/utils/data_and_condition.py diff --git a/cosmos_framework/model/vfm/utils/memory.py b/cosmos_framework/model/generator/utils/memory.py similarity index 100% rename from cosmos_framework/model/vfm/utils/memory.py rename to cosmos_framework/model/generator/utils/memory.py diff --git a/cosmos_framework/model/vfm/utils/safetensors_loader.py b/cosmos_framework/model/generator/utils/safetensors_loader.py similarity index 99% rename from cosmos_framework/model/vfm/utils/safetensors_loader.py rename to cosmos_framework/model/generator/utils/safetensors_loader.py index 22c2d6fc..7fc12180 100644 --- a/cosmos_framework/model/vfm/utils/safetensors_loader.py +++ b/cosmos_framework/model/generator/utils/safetensors_loader.py @@ -53,7 +53,7 @@ from cosmos_framework.utils.flags import INTERNAL from cosmos_framework.utils import log from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.utils.vfm.parallelism import ParallelDims +from cosmos_framework.utils.generator.parallelism import ParallelDims # Prefixes stripped when matching checkpoint keys to model state-dict keys. # Order matters: longest first. For each model key, the longest matching @@ -988,7 +988,7 @@ def load_language_model( if tie_embeddings: # The `*ForCausalLM` classes in - # `cosmos_framework/model/vfm/mot/unified_mot.py` override + # `cosmos_framework/model/generator/mot/unified_mot.py` override # `get_input_embeddings` (canonical HF idiom) to return the inner # `model.embed_tokens`, so this call returns a real `nn.Embedding` # rather than raising `NotImplementedError`. @@ -1204,7 +1204,7 @@ def load_vfm_model( r"""Load a complete Cosmos3 VFM checkpoint (safetensors) into a Cosmos3VFMNetwork. Loads the *entire* state of a - :class:`~cosmos_framework.model.vfm.mot.cosmos3_vfm_network.Cosmos3VFMNetwork` + :class:`~cosmos_framework.model.generator.mot.cosmos3_vfm_network.Cosmos3VFMNetwork` in one shot: - the language tower (``language_model.*``), which carries the diff --git a/cosmos_framework/model/vfm/utils/safetensors_loader_test.py b/cosmos_framework/model/generator/utils/safetensors_loader_test.py similarity index 98% rename from cosmos_framework/model/vfm/utils/safetensors_loader_test.py rename to cosmos_framework/model/generator/utils/safetensors_loader_test.py index 657bcccb..aaa26445 100644 --- a/cosmos_framework/model/vfm/utils/safetensors_loader_test.py +++ b/cosmos_framework/model/generator/utils/safetensors_loader_test.py @@ -4,7 +4,7 @@ """ Unit tests for safetensors_loader helpers and load_vlm_model. -pytest cosmos_framework/model/vfm/utils/safetensors_loader_test.py -v +pytest cosmos_framework/model/generator/utils/safetensors_loader_test.py -v """ from pathlib import Path @@ -15,7 +15,7 @@ import torch from safetensors.torch import save_file -from cosmos_framework.model.vfm.utils.safetensors_loader import ( +from cosmos_framework.model.generator.utils.safetensors_loader import ( MultiRankCheckpointLoader, _build_model_key_by_tail, _get_dp_shard_mesh, @@ -68,7 +68,7 @@ def _make_safetensors(tmp_path: Path, tensors: dict[str, torch.Tensor]) -> Path: # The single-rank CPU fallback is reached by passing ``parallel_dims=None`` # (the documented escape hatch — see ``load_vlm_model`` docstring). All # end-to-end tests below use that path; multi-rank behavior is covered in -# the GPU-marked tests under ``cosmos_framework/model/vfm/mot/``. +# the GPU-marked tests under ``cosmos_framework/model/generator/mot/``. # # Do NOT introduce a "fake" ``ParallelDims`` MagicMock fixture for this # fallback: ``MagicMock.__getitem__`` returns another MagicMock rather than @@ -682,7 +682,7 @@ def test_load_vlm_model_qwen2_5_bare_visual_keys_overlay(tmp_path): """ # Import the production helper pattern instead of retyping — guarantees drift # between test and source is impossible. - from cosmos_framework.model.vfm.vlm_model import _get_overlay_config + from cosmos_framework.model.generator.vlm_model import _get_overlay_config overlay_skip_patterns, _ = _get_overlay_config("qwen2_5_vl") diff --git a/cosmos_framework/model/vfm/utils/taylorseer.py b/cosmos_framework/model/generator/utils/taylorseer.py similarity index 100% rename from cosmos_framework/model/vfm/utils/taylorseer.py rename to cosmos_framework/model/generator/utils/taylorseer.py diff --git a/cosmos_framework/model/vfm/vlm_model.py b/cosmos_framework/model/generator/vlm_model.py similarity index 97% rename from cosmos_framework/model/vfm/vlm_model.py rename to cosmos_framework/model/generator/vlm_model.py index ec62b62f..e76970a8 100644 --- a/cosmos_framework/model/vfm/vlm_model.py +++ b/cosmos_framework/model/generator/vlm_model.py @@ -28,14 +28,14 @@ from cosmos_framework.utils.lazy_config import instantiate from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.utils import log -from cosmos_framework.model.vfm.algorithm.loss.cross_entropy import cross_entropy_loss, weighted_cross_entropy_loss +from cosmos_framework.model.generator.algorithm.loss.cross_entropy import cross_entropy_loss, weighted_cross_entropy_loss from cosmos_framework.configs.base.defaults.parallelism import PRECISION_TO_TORCH_DTYPE from cosmos_framework.configs.base.reasoner.defaults.policy_config import VLMModelConfig -from cosmos_framework.model.vfm.hf_model import HFModel -from cosmos_framework.model.vfm.parallelize_vlm import parallelize -from cosmos_framework.utils.vfm.parallelism import ParallelDims -from cosmos_framework.utils.vfm.reasoner.constant import IGNORE_INDEX -from cosmos_framework.utils.vfm.reasoner.create_position_ids import get_position_ids +from cosmos_framework.model.generator.hf_model import HFModel +from cosmos_framework.model.generator.parallelize_vlm import parallelize +from cosmos_framework.utils.generator.parallelism import ParallelDims +from cosmos_framework.utils.generator.reasoner.constant import IGNORE_INDEX +from cosmos_framework.utils.generator.reasoner.create_position_ids import get_position_ids # Model-type dispatch sets. Using hf_config.model_type (stable HF-defined string) # rather than backbone.model_name avoids the brittleness of substring-matching a local @@ -248,7 +248,7 @@ class VLMModel(ImaginaireModel): def __init__(self, config: VLMModelConfig, checkpoint): super().__init__() - from cosmos_framework.utils.vfm.flash_attn import init_flash_attn_meta + from cosmos_framework.utils.generator.flash_attn import init_flash_attn_meta self.config = config # Expose model.precision so LowPrecisionCallback can read it (mirrors OmniMoTModel). @@ -303,7 +303,7 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None: g. Load pretrain weights into sharded CUDA tensors. h. Apply gradient checkpointing if configured. """ - from cosmos_framework.utils.vfm.reasoner.pretrained_models_downloader import ( + from cosmos_framework.utils.generator.reasoner.pretrained_models_downloader import ( maybe_download_hf_model_from_s3, ) diff --git a/cosmos_framework/scripts/_convert_model_to_diffusers.py b/cosmos_framework/scripts/_convert_model_to_diffusers.py index 98503c77..e1c14016 100644 --- a/cosmos_framework/scripts/_convert_model_to_diffusers.py +++ b/cosmos_framework/scripts/_convert_model_to_diffusers.py @@ -17,7 +17,7 @@ from cosmos_framework.inference.model import Cosmos3OmniModel from cosmos_framework.utils import log -from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel +from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel DEFAULT_SOUND_TOKENIZER_CONFIG = { "model_type": "autoencoder_v2", diff --git a/cosmos_framework/scripts/_train.py b/cosmos_framework/scripts/_train.py index b4996f81..6532f40a 100644 --- a/cosmos_framework/scripts/_train.py +++ b/cosmos_framework/scripts/_train.py @@ -38,7 +38,7 @@ from torch.utils.data import DataLoader from cosmos_framework.utils.config import Config - from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel + from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel def _validate_config_file(v: Path) -> Path: diff --git a/cosmos_framework/scripts/action_policy_server_libero.py b/cosmos_framework/scripts/action_policy_server_libero.py index a7be2348..5b75f87f 100644 --- a/cosmos_framework/scripts/action_policy_server_libero.py +++ b/cosmos_framework/scripts/action_policy_server_libero.py @@ -63,13 +63,13 @@ # Action-specific helpers live in the in-tree project tree. Imports stay as # `projects.cosmos3.vfm.*` and are auto-rewritten to `cosmos3._src.vfm.*` by the # cosmos-framework release script. -from cosmos_framework.data.vfm.action.action_processing import ( +from cosmos_framework.data.generator.action.action_processing import ( ActionProcessingRecord, make_batched_action_processing_fields, ) -from cosmos_framework.data.vfm.action.domain_utils import get_domain_id -from cosmos_framework.data.vfm.action.json_formatter import ActionPromptJsonFormatter -from cosmos_framework.data.vfm.action.transforms import ( +from cosmos_framework.data.generator.action.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.json_formatter import ActionPromptJsonFormatter +from cosmos_framework.data.generator.action.transforms import ( build_sequence_plan_from_mode, find_closest_target_size, reflection_pad_to_target, @@ -88,7 +88,7 @@ ) from cosmos_framework.utils import log from cosmos_framework.utils.lazy_config import instantiate -from cosmos_framework.utils.vfm.data_utils import get_vision_data_resolution +from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution _DEFAULT_ACTION_CHUNK_SIZE = 16 ActionNormalization = Literal["auto", "meanstd", "minmax", "quantile", "quantile_rot"] diff --git a/cosmos_framework/scripts/action_policy_server_robolab.py b/cosmos_framework/scripts/action_policy_server_robolab.py index a5cc459b..f76989ab 100644 --- a/cosmos_framework/scripts/action_policy_server_robolab.py +++ b/cosmos_framework/scripts/action_policy_server_robolab.py @@ -37,15 +37,15 @@ import torch.nn.functional as F import tyro -from cosmos_framework.data.vfm.action.domain_utils import get_domain_id -from cosmos_framework.data.vfm.action.pose_utils import ( +from cosmos_framework.data.generator.action.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.pose_utils import ( build_abs_pose_from_components, convert_rotation, pose_abs_to_rel, pose_rel_to_abs, ) -from cosmos_framework.data.vfm.action.transforms import ActionTransformPipeline -from cosmos_framework.data.vfm.joint_dataloader import IterativeJointDataLoader +from cosmos_framework.data.generator.action.transforms import ActionTransformPipeline +from cosmos_framework.data.generator.joint_dataloader import IterativeJointDataLoader from cosmos_framework.inference.args import OmniSetupArgs, OmniSetupOverrides from cosmos_framework.inference.common.args import ConfigFileType, ConfigOverrides, tyro_cli from cosmos_framework.inference.common.config import deserialize_config, deserialize_config_dict, load_config diff --git a/cosmos_framework/scripts/export_model.py b/cosmos_framework/scripts/export_model.py index 02191ae7..d1479784 100644 --- a/cosmos_framework/scripts/export_model.py +++ b/cosmos_framework/scripts/export_model.py @@ -37,7 +37,7 @@ from cosmos_framework.inference.common.init import is_rank0 from cosmos_framework.inference.common.public_model_config import build_public_model_config from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel -from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel +from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel from cosmos_framework.utils import log from cosmos_framework.utils.checkpoint_db import CheckpointConfig, sanitize_uri from cosmos_framework.utils.lazy_config.registry import convert_target_to_string diff --git a/cosmos_framework/scripts/train.py b/cosmos_framework/scripts/train.py index 7d19e1cb..07ef26a5 100644 --- a/cosmos_framework/scripts/train.py +++ b/cosmos_framework/scripts/train.py @@ -15,7 +15,7 @@ raises on unknown keys), then ``cosmos_framework.configs.toml_config.sft_config.load_experiment_from_toml`` picks the base ``config.py`` from ``[job].task`` (``vfm`` → ``cosmos_framework/configs/base/config.py``, -``vlm`` → ``cosmos_framework/configs/base/vlm/config.py``), resolves ``[job].experiment`` +``vlm`` → ``cosmos_framework/configs/base/reasoner/config.py``), resolves ``[job].experiment`` against the Hydra ``ConfigStore``, and overlays every other TOML key as a Hydra override. Trailing ``key.path=value`` positionals are applied last (so they win over TOML). diff --git a/cosmos_framework/scripts/upsample_prompts.py b/cosmos_framework/scripts/upsample_prompts.py index 74435b9a..360ec031 100644 --- a/cosmos_framework/scripts/upsample_prompts.py +++ b/cosmos_framework/scripts/upsample_prompts.py @@ -16,7 +16,7 @@ from tqdm import tqdm from cosmos_framework.inference.args import ModelMode, OmniSampleArgs, OmniSampleOverrides, OmniSetupOverrides -from cosmos_framework.model.vfm.upsampler.prompts import build_messages, clean_response +from cosmos_framework.model.generator.upsampler.prompts import build_messages, clean_response from cosmos_framework.utils import log if TYPE_CHECKING: diff --git a/cosmos_framework/simulation/libero/closed_loop_eval.py b/cosmos_framework/simulation/libero/closed_loop_eval.py index 660be360..885fd683 100644 --- a/cosmos_framework/simulation/libero/closed_loop_eval.py +++ b/cosmos_framework/simulation/libero/closed_loop_eval.py @@ -51,12 +51,12 @@ from PIL import Image from scipy.spatial.transform import Rotation as R -from cosmos_framework.data.vfm.action.libero_pose_utils import ( +from cosmos_framework.data.generator.action.libero_pose_utils import ( libero_rotation_format, libero_rotation_space_from_action_dim, ) -from cosmos_framework.data.vfm.action.pose_utils import convert_rotation -from cosmos_framework.data.vfm.action.viewpoint_utils import DEFAULT_VIEWPOINT_TEMPLATES +from cosmos_framework.data.generator.action.pose_utils import convert_rotation +from cosmos_framework.data.generator.action.viewpoint_utils import DEFAULT_VIEWPOINT_TEMPLATES benchmark: Any get_libero_path: Any diff --git a/cosmos_framework/utils/vfm/__init__.py b/cosmos_framework/utils/generator/__init__.py similarity index 100% rename from cosmos_framework/utils/vfm/__init__.py rename to cosmos_framework/utils/generator/__init__.py diff --git a/cosmos_framework/utils/vfm/data_utils.py b/cosmos_framework/utils/generator/data_utils.py similarity index 100% rename from cosmos_framework/utils/vfm/data_utils.py rename to cosmos_framework/utils/generator/data_utils.py diff --git a/cosmos_framework/utils/vfm/dtensor_helper.py b/cosmos_framework/utils/generator/dtensor_helper.py similarity index 100% rename from cosmos_framework/utils/vfm/dtensor_helper.py rename to cosmos_framework/utils/generator/dtensor_helper.py diff --git a/cosmos_framework/utils/vfm/flash_attn.py b/cosmos_framework/utils/generator/flash_attn.py similarity index 100% rename from cosmos_framework/utils/vfm/flash_attn.py rename to cosmos_framework/utils/generator/flash_attn.py diff --git a/cosmos_framework/utils/vfm/fused_adam.py b/cosmos_framework/utils/generator/fused_adam.py similarity index 100% rename from cosmos_framework/utils/vfm/fused_adam.py rename to cosmos_framework/utils/generator/fused_adam.py diff --git a/cosmos_framework/utils/vfm/hf_attention_cosmos.py b/cosmos_framework/utils/generator/hf_attention_cosmos.py similarity index 100% rename from cosmos_framework/utils/vfm/hf_attention_cosmos.py rename to cosmos_framework/utils/generator/hf_attention_cosmos.py diff --git a/cosmos_framework/utils/vfm/lora.py b/cosmos_framework/utils/generator/lora.py similarity index 100% rename from cosmos_framework/utils/vfm/lora.py rename to cosmos_framework/utils/generator/lora.py diff --git a/cosmos_framework/utils/vfm/model_loader.py b/cosmos_framework/utils/generator/model_loader.py similarity index 97% rename from cosmos_framework/utils/vfm/model_loader.py rename to cosmos_framework/utils/generator/model_loader.py index 92e67ab0..a53089de 100644 --- a/cosmos_framework/utils/vfm/model_loader.py +++ b/cosmos_framework/utils/generator/model_loader.py @@ -41,7 +41,7 @@ def write_lock(self) -> FileLock: from cosmos_framework.utils.config_helper import get_config_module, override from cosmos_framework.utils.easy_io import easy_io from cosmos_framework.checkpoint.dcp import CustomLoadPlanner, CustomSavePlanner, ModelWrapper -from cosmos_framework.model.vfm.utils.safetensors_loader import load_vfm_model +from cosmos_framework.model.generator.utils.safetensors_loader import load_vfm_model ################################################### # below are the load_model function for inference # @@ -252,17 +252,10 @@ def _load_model( keys_to_skip_loading=keys_to_skip_loading or [], ) - # Single-rank load (e.g. the action-policy inference server): force no_dist so - # ``dcp.load`` skips the collective ``gather_object`` over the load plan, which - # pickles the plan and can fail on training/EMA DCPs. Multi-rank loads keep the - # default distributed path. - no_dist = not (dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1) - dcp.load( state_dict=state_dict, storage_reader=storage_reader, planner=load_planner, - no_dist=no_dist, ) log.info(f"Successfully loaded model from {checkpoint_path}") @@ -330,7 +323,7 @@ def load_model_from_checkpoint( * **safetensors**: a directory containing one or more ``*.safetensors`` shards in the native Cosmos3 VFM state-dict layout. Loaded via - :func:`cosmos_framework.model.vfm.utils.safetensors_loader.load_vfm_model`. + :func:`cosmos_framework.model.generator.utils.safetensors_loader.load_vfm_model`. No ``/model`` suffix is appended. credential_path: Path to credentials file (if required for remote storage). Optional. enable_gcs_patch_in_boto3: Whether to enable the boto3 patch for GCS S3-compatibility. diff --git a/cosmos_framework/utils/vfm/model_weights_stats.py b/cosmos_framework/utils/generator/model_weights_stats.py similarity index 100% rename from cosmos_framework/utils/vfm/model_weights_stats.py rename to cosmos_framework/utils/generator/model_weights_stats.py diff --git a/cosmos_framework/utils/vfm/monkey_patch.py b/cosmos_framework/utils/generator/monkey_patch.py similarity index 100% rename from cosmos_framework/utils/vfm/monkey_patch.py rename to cosmos_framework/utils/generator/monkey_patch.py diff --git a/cosmos_framework/utils/vfm/optimizer.py b/cosmos_framework/utils/generator/optimizer.py similarity index 99% rename from cosmos_framework/utils/vfm/optimizer.py rename to cosmos_framework/utils/generator/optimizer.py index 1930cbc8..2c1c798b 100644 --- a/cosmos_framework/utils/vfm/optimizer.py +++ b/cosmos_framework/utils/generator/optimizer.py @@ -44,7 +44,7 @@ def _optimizer_cls( - ``"adam"`` / ``"adamw"``: forwarded to ``torch.optim.Adam`` / ``torch.optim.AdamW``. ``fused`` (if present in ``optimizer_kwargs``) flows through and selects the fused CUDA kernel. - - ``"fusedadam"``: NVIDIA's :class:`cosmos_framework.utils.vfm.fused_adam.FusedAdam`. + - ``"fusedadam"``: NVIDIA's :class:`cosmos_framework.utils.generator.fused_adam.FusedAdam`. It is fused by construction and rejects a ``fused`` kwarg, so any ``fused`` entry is popped before instantiation. We also force ``capturable=True`` and ``master_weights=True`` because those are the @@ -57,7 +57,7 @@ def _optimizer_cls( elif optimizer_type.lower() == "adamw": optimizer = torch.optim.AdamW(params, **optimizer_kwargs) elif optimizer_type.lower() == "fusedadam": - from cosmos_framework.utils.vfm.fused_adam import FusedAdam + from cosmos_framework.utils.generator.fused_adam import FusedAdam # FusedAdam is fused by construction and does not accept a ``fused`` kwarg. optimizer_kwargs.pop("fused", None) diff --git a/cosmos_framework/utils/vfm/parallelism.py b/cosmos_framework/utils/generator/parallelism.py similarity index 98% rename from cosmos_framework/utils/vfm/parallelism.py rename to cosmos_framework/utils/generator/parallelism.py index 58e88f10..70aeee3b 100644 --- a/cosmos_framework/utils/vfm/parallelism.py +++ b/cosmos_framework/utils/generator/parallelism.py @@ -28,8 +28,8 @@ - VFM inference — ``dp_shard`` + cfgp/cp overlays; replicate forced to 1. FSDP wrapping for VLM ``HFModel`` instances lives in -``cosmos_framework.model.vfm.parallelize_vlm``; MoT wrapping lives in -``cosmos_framework.model.vfm.mot.parallelize_unified_mot``. Both consume +``cosmos_framework.model.generator.parallelize_vlm``; MoT wrapping lives in +``cosmos_framework.model.generator.mot.parallelize_unified_mot``. Both consume ``ParallelDims`` from this module. """ diff --git a/cosmos_framework/utils/vfm/rand_state.py b/cosmos_framework/utils/generator/rand_state.py similarity index 100% rename from cosmos_framework/utils/vfm/rand_state.py rename to cosmos_framework/utils/generator/rand_state.py diff --git a/cosmos_framework/utils/vfm/reasoner/__init__.py b/cosmos_framework/utils/generator/reasoner/__init__.py similarity index 100% rename from cosmos_framework/utils/vfm/reasoner/__init__.py rename to cosmos_framework/utils/generator/reasoner/__init__.py diff --git a/cosmos_framework/utils/vfm/reasoner/constant.py b/cosmos_framework/utils/generator/reasoner/constant.py similarity index 100% rename from cosmos_framework/utils/vfm/reasoner/constant.py rename to cosmos_framework/utils/generator/reasoner/constant.py diff --git a/cosmos_framework/utils/vfm/reasoner/create_position_ids.py b/cosmos_framework/utils/generator/reasoner/create_position_ids.py similarity index 100% rename from cosmos_framework/utils/vfm/reasoner/create_position_ids.py rename to cosmos_framework/utils/generator/reasoner/create_position_ids.py diff --git a/cosmos_framework/utils/vfm/reasoner/flop_calculator.py b/cosmos_framework/utils/generator/reasoner/flop_calculator.py similarity index 100% rename from cosmos_framework/utils/vfm/reasoner/flop_calculator.py rename to cosmos_framework/utils/generator/reasoner/flop_calculator.py diff --git a/cosmos_framework/utils/vfm/reasoner/pretrained_models_downloader.py b/cosmos_framework/utils/generator/reasoner/pretrained_models_downloader.py similarity index 99% rename from cosmos_framework/utils/vfm/reasoner/pretrained_models_downloader.py rename to cosmos_framework/utils/generator/reasoner/pretrained_models_downloader.py index 54c18cb3..1bdd964e 100644 --- a/cosmos_framework/utils/vfm/reasoner/pretrained_models_downloader.py +++ b/cosmos_framework/utils/generator/reasoner/pretrained_models_downloader.py @@ -255,7 +255,7 @@ def maybe_download_hf_model_from_s3( if __name__ == "__main__": """ Usage: - PYTHONPATH=. python3 cosmos_framework/utils/vlm/pretrained_models_downloader.py + PYTHONPATH=. python3 cosmos_framework/utils/reasoner/pretrained_models_downloader.py """ cache_dir = maybe_download_model( # noqa: F821 "eagle_er_qwen3_1p7b_siglip_400m", "credentials/s3_training.secret", "bucket4" diff --git a/cosmos_framework/utils/vfm/video_preprocess.py b/cosmos_framework/utils/generator/video_preprocess.py similarity index 100% rename from cosmos_framework/utils/vfm/video_preprocess.py rename to cosmos_framework/utils/generator/video_preprocess.py diff --git a/cosmos_framework/utils/vlm/__init__.py b/cosmos_framework/utils/reasoner/__init__.py similarity index 100% rename from cosmos_framework/utils/vlm/__init__.py rename to cosmos_framework/utils/reasoner/__init__.py diff --git a/cosmos_framework/utils/vlm/compute_flops_qwen3vl.py b/cosmos_framework/utils/reasoner/compute_flops_qwen3vl.py similarity index 100% rename from cosmos_framework/utils/vlm/compute_flops_qwen3vl.py rename to cosmos_framework/utils/reasoner/compute_flops_qwen3vl.py diff --git a/cosmos_framework/utils/vlm/configs_defaults/__init__.py b/cosmos_framework/utils/reasoner/configs_defaults/__init__.py similarity index 100% rename from cosmos_framework/utils/vlm/configs_defaults/__init__.py rename to cosmos_framework/utils/reasoner/configs_defaults/__init__.py diff --git a/cosmos_framework/utils/vlm/configs_defaults/checkpointer.py b/cosmos_framework/utils/reasoner/configs_defaults/checkpointer.py similarity index 97% rename from cosmos_framework/utils/vlm/configs_defaults/checkpointer.py rename to cosmos_framework/utils/reasoner/configs_defaults/checkpointer.py index b7bfab67..67003279 100644 --- a/cosmos_framework/utils/vlm/configs_defaults/checkpointer.py +++ b/cosmos_framework/utils/reasoner/configs_defaults/checkpointer.py @@ -9,7 +9,7 @@ from cosmos_framework.checkpoint.dummy import Checkpointer as DummyCheckpointer from cosmos_framework.utils.config import CheckpointConfig from cosmos_framework.utils.lazy_config import LazyCall as L -from cosmos_framework.utils.vlm.dcp_checkpointer import DistributedCheckpointer +from cosmos_framework.utils.reasoner.dcp_checkpointer import DistributedCheckpointer pdx_object_store = config.ObjectStoreConfig( enabled=True, diff --git a/cosmos_framework/utils/vlm/constant.py b/cosmos_framework/utils/reasoner/constant.py similarity index 100% rename from cosmos_framework/utils/vlm/constant.py rename to cosmos_framework/utils/reasoner/constant.py diff --git a/cosmos_framework/utils/vlm/create_position_ids.py b/cosmos_framework/utils/reasoner/create_position_ids.py similarity index 100% rename from cosmos_framework/utils/vlm/create_position_ids.py rename to cosmos_framework/utils/reasoner/create_position_ids.py diff --git a/cosmos_framework/utils/vlm/dcp_checkpointer.py b/cosmos_framework/utils/reasoner/dcp_checkpointer.py similarity index 99% rename from cosmos_framework/utils/vlm/dcp_checkpointer.py rename to cosmos_framework/utils/reasoner/dcp_checkpointer.py index 6984cae0..442b8a4c 100644 --- a/cosmos_framework/utils/vlm/dcp_checkpointer.py +++ b/cosmos_framework/utils/reasoner/dcp_checkpointer.py @@ -58,9 +58,9 @@ from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.utils import callback, distributed, log, misc from cosmos_framework.utils.easy_io import easy_io -from cosmos_framework.utils.vlm.model_wrapper import ModelWrapper -from cosmos_framework.utils.vlm.optimizer import OptimizersContainer -from cosmos_framework.utils.vlm.planner import RenameLoadPlanner +from cosmos_framework.utils.reasoner.model_wrapper import ModelWrapper +from cosmos_framework.utils.reasoner.optimizer import OptimizersContainer +from cosmos_framework.utils.reasoner.planner import RenameLoadPlanner # (qsh 2025-01-01) the design is from https://github.com/pytorch/torchtitan/blob/1060feacc1b51cb6b339a04e53a5243b8466552b/torchtitan/checkpoint.py # we recreate wrapper when needed instead of creating one from the beginning. diff --git a/cosmos_framework/utils/vlm/distributed.py b/cosmos_framework/utils/reasoner/distributed.py similarity index 100% rename from cosmos_framework/utils/vlm/distributed.py rename to cosmos_framework/utils/reasoner/distributed.py diff --git a/cosmos_framework/utils/vlm/flop_calculator.py b/cosmos_framework/utils/reasoner/flop_calculator.py similarity index 98% rename from cosmos_framework/utils/vlm/flop_calculator.py rename to cosmos_framework/utils/reasoner/flop_calculator.py index 03c4cbe4..193c30a9 100644 --- a/cosmos_framework/utils/vlm/flop_calculator.py +++ b/cosmos_framework/utils/reasoner/flop_calculator.py @@ -14,7 +14,7 @@ import torch -from cosmos_framework.utils.vlm.compute_flops_qwen3vl import compute_qwen3vl_flops_from_config +from cosmos_framework.utils.reasoner.compute_flops_qwen3vl import compute_qwen3vl_flops_from_config class FlopCalculator: diff --git a/cosmos_framework/utils/vlm/fused_adam.py b/cosmos_framework/utils/reasoner/fused_adam.py similarity index 99% rename from cosmos_framework/utils/vlm/fused_adam.py rename to cosmos_framework/utils/reasoner/fused_adam.py index e930da0b..1c6159ea 100644 --- a/cosmos_framework/utils/vlm/fused_adam.py +++ b/cosmos_framework/utils/reasoner/fused_adam.py @@ -6,7 +6,7 @@ from torch.distributed._functional_collectives import AsyncCollectiveTensor from torch.distributed._tensor.api import DTensor -from cosmos_framework.utils.vlm import distributed +from cosmos_framework.utils.reasoner import distributed def get_local_tensor_if_DTensor(tensor: torch.Tensor | DTensor) -> torch.Tensor: diff --git a/cosmos_framework/utils/vlm/model_wrapper.py b/cosmos_framework/utils/reasoner/model_wrapper.py similarity index 100% rename from cosmos_framework/utils/vlm/model_wrapper.py rename to cosmos_framework/utils/reasoner/model_wrapper.py diff --git a/cosmos_framework/utils/vlm/optimizer.py b/cosmos_framework/utils/reasoner/optimizer.py similarity index 99% rename from cosmos_framework/utils/vlm/optimizer.py rename to cosmos_framework/utils/reasoner/optimizer.py index fb74a7ff..36f13fe2 100644 --- a/cosmos_framework/utils/vlm/optimizer.py +++ b/cosmos_framework/utils/reasoner/optimizer.py @@ -18,7 +18,7 @@ from cosmos_framework.utils.config import make_freezable from cosmos_framework.utils.lazy_config import LazyDict from cosmos_framework.utils import log -from cosmos_framework.utils.vlm.fused_adam import FusedAdam +from cosmos_framework.utils.reasoner.fused_adam import FusedAdam @make_freezable diff --git a/cosmos_framework/utils/vlm/planner.py b/cosmos_framework/utils/reasoner/planner.py similarity index 100% rename from cosmos_framework/utils/vlm/planner.py rename to cosmos_framework/utils/reasoner/planner.py diff --git a/cosmos_framework/utils/vlm/pretrained_models_downloader.py b/cosmos_framework/utils/reasoner/pretrained_models_downloader.py similarity index 97% rename from cosmos_framework/utils/vlm/pretrained_models_downloader.py rename to cosmos_framework/utils/reasoner/pretrained_models_downloader.py index 5fff313a..28c2eea6 100644 --- a/cosmos_framework/utils/vlm/pretrained_models_downloader.py +++ b/cosmos_framework/utils/reasoner/pretrained_models_downloader.py @@ -20,7 +20,7 @@ def resolve_hf_model_store(credentials: str, bucket: str) -> tuple[str, str]: AWS training checkpoints → aws_load_from_object_store_permanent (nv-cosmos-vlm) Falls back to the provided credentials/bucket if neither matches. """ - from cosmos_framework.utils.vlm.configs_defaults.checkpointer import ( + from cosmos_framework.utils.reasoner.configs_defaults.checkpointer import ( aws_load_from_object_store_permanent, gcp_load_from_object_store_permanent, ) @@ -201,7 +201,7 @@ def maybe_download_hf_model_from_s3( if __name__ == "__main__": """ Usage: - PYTHONPATH=. python3 cosmos_framework/utils/vlm/pretrained_models_downloader.py + PYTHONPATH=. python3 cosmos_framework/utils/reasoner/pretrained_models_downloader.py """ cache_dir = maybe_download_model( # noqa: F821 "eagle_er_qwen3_1p7b_siglip_400m", "credentials/s3_training.secret", "bucket4" diff --git a/cosmos_framework/utils/vlm/video_preprocess.py b/cosmos_framework/utils/reasoner/video_preprocess.py similarity index 100% rename from cosmos_framework/utils/vlm/video_preprocess.py rename to cosmos_framework/utils/reasoner/video_preprocess.py diff --git a/docs/action_policy_libero_sft.md b/docs/action_policy_libero_sft.md index b1de5df7..f7752e96 100644 --- a/docs/action_policy_libero_sft.md +++ b/docs/action_policy_libero_sft.md @@ -15,16 +15,16 @@ lr 5e-5, warmup 500, cycle 16000, gbs 2048): `action_policy_libero_all_nano` + `action_policy_libero_all_repro.toml` + `launch_sft_action_policy_libero_all.sh`. -| Piece | Path | -| ---------------- | ----------------------------------------------------------------------------------------------- | -| Dataset | `cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py` (`LIBEROLeRobotDataset`) | -| SFT wrapper | `get_action_libero_sft_dataset` in `.../datasets/action_sft_dataset.py` | -| Norm stats | `.../normalizer_stats/libero_native_frame_wise_relative_rot6d.json` | -| Experiment | `cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py` | -| Run TOML | `examples/toml/sft_config/action_policy_libero_repro.toml` | -| Launch | `examples/launch_sft_action_policy_libero.sh` | -| Inference server | `cosmos_framework/scripts/action_policy_server_libero.py` | -| Closed-loop eval | `cosmos_framework/simulation/libero/closed_loop_eval.py` | +| Piece | Path | +| ---------------- | ---------------------------------------------------------------------------------------------------- | +| Dataset | `cosmos_framework/data/generator/action/datasets/libero_lerobot_dataset.py` (`LIBEROLeRobotDataset`) | +| SFT wrapper | `get_action_libero_sft_dataset` in `.../datasets/action_sft_dataset.py` | +| Norm stats | `.../normalizer_stats/libero_native_frame_wise_relative_rot6d.json` | +| Experiment | `cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py` | +| Run TOML | `examples/toml/sft_config/action_policy_libero_repro.toml` | +| Launch | `examples/launch_sft_action_policy_libero.sh` | +| Inference server | `cosmos_framework/scripts/action_policy_server_libero.py` | +| Closed-loop eval | `cosmos_framework/simulation/libero/closed_loop_eval.py` | ## 1. Data @@ -89,7 +89,7 @@ python -m cosmos_framework.scripts.action_policy_server_libero \ --experiment-overrides "model.config.tokenizer.vae_path=$WAN_VAE_PATH" \ --checkpoint-path /checkpoints/iter_000001500 \ --action-normalization quantile_rot \ - --action-stats-path cosmos_framework/data/vfm/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json \ + --action-stats-path cosmos_framework/data/generator/action/normalizer_stats/libero_native_frame_wise_relative_rot6d.json \ --raw-action-dim 10 --fps 20 --port 8000 ``` diff --git a/docs/custom_dataset.md b/docs/custom_dataset.md index be7e1aa2..3f41318b 100644 --- a/docs/custom_dataset.md +++ b/docs/custom_dataset.md @@ -24,7 +24,7 @@ DataDistributor → RawItemProcessor → SampleBatcher → BatchColl together in a batch (fixed size, token-budget packing, …). - **`BatchCollator`** turns a chosen group of samples into one batch dict. -Everything lives in `cosmos_framework.data.vfm.dataflow`. The loader is a +Everything lives in `cosmos_framework.data.generator.dataflow`. The loader is a `torch.utils.data.DataLoader` subclass, so it drops into existing training loops. --- @@ -49,7 +49,7 @@ Everything lives in `cosmos_framework.data.vfm.dataflow`. The loader is a "I have a map-style dataset and just want normal, shuffled, resumable batches": ```python -from cosmos_framework.data.vfm.dataflow import ( +from cosmos_framework.data.generator.dataflow import ( CosmosDataLoader, MapDistributor, IdentityProcessor, ) @@ -72,7 +72,7 @@ fancier. (Passing both `batch_size=` and `batcher=` is an error.) ## 2. The four roles -Each role is a tiny ABC in `cosmos_framework.data.vfm.dataflow.base`. Implement +Each role is a tiny ABC in `cosmos_framework.data.generator.dataflow.base`. Implement the one method (plus, for distributors, optional resume hooks). ```python @@ -196,7 +196,7 @@ processor/batcher/collator (homogeneous join). ### Interleave heterogeneous pipelines (different processors/collators) ```python -from cosmos_framework.data.vfm.dataflow import JointCosmosDataLoader +from cosmos_framework.data.generator.dataflow import JointCosmosDataLoader joint = JointCosmosDataLoader( dataloaders={ @@ -327,7 +327,7 @@ A local image-caption folder, fully custom processor, normal batching: ```python import torch -from cosmos_framework.data.vfm.dataflow import ( +from cosmos_framework.data.generator.dataflow import ( CosmosDataLoader, MapDistributor, RawItemProcessor, DefaultBatchCollator, SimpleBatcher, ) @@ -368,7 +368,7 @@ collator that pads/stacks accordingly — nothing else changes. ### Reasoner (VLM) — HuggingFace image-text dataset, streaming -**File**: `cosmos_framework/configs/base/vlm/experiment/llava_ov_vlm.py` +**File**: `cosmos_framework/configs/base/reasoner/experiment/llava_ov_vlm.py` (`pre_exp012_llava_ov`) ``` @@ -385,7 +385,7 @@ wraps it in a `MapDistributor`, so checkpoint/resume works (see §5). ### Reasoner (VLM) — local video dialog dataset -**File**: `cosmos_framework/configs/base/vlm/experiment/videophy2_sft_nano.py` +**File**: `cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py` (`videophy2_sft_nano`) ``` @@ -447,8 +447,8 @@ collator: VFMListCollator # media kept as p ## Reference: where things live -- ABCs + built-ins: `cosmos_framework/data/vfm/dataflow/` (`base.py`, +- ABCs + built-ins: `cosmos_framework/data/generator/dataflow/` (`base.py`, `distributors.py`, `batchers.py`, `collators.py`, `processors.py`, `loader.py`). -- Public symbols are re-exported from `cosmos_framework.data.vfm.dataflow`. +- Public symbols are re-exported from `cosmos_framework.data.generator.dataflow`. - Live recipes using the loader: `pre_exp012_llava_ov`, `pre_exp012_llava_ov_mapstyle_dataloader`, `videophy2_sft_nano`, and `vision_sft_nano_mapstyle_dataloader`. diff --git a/docs/faq.md b/docs/faq.md index 4ddf47ed..445bab50 100644 --- a/docs/faq.md +++ b/docs/faq.md @@ -262,7 +262,7 @@ Knobs are in the recipe TOML under `[model]`, `[model.parallelism]`, and `[datal - `mode = "full"` — checkpoint every transformer block (largest memory savings, trades extra recompute for memory). - `mode = "selective"` — per-op SAC, MoT only (smaller savings, smaller overhead). Falls back to no checkpointing on the VLM path. -3. **Raise `[model.parallelism].data_parallel_shard_degree`** to shard weights/optimizer state across more ranks via FSDP. Runtime invariant (from `cosmos_framework/utils/vfm/parallelism.py:50-52`): `data_parallel_replicate_degree × data_parallel_shard_degree == WORLD_SIZE` always holds — `context_parallel_shard_degree` and `cfg_parallel_shard_degree` are *overlay* axes that share dp rank slots, not separate mesh dims. Use `-1` to let `data_parallel_shard_degree` auto-fill from `torchrun` world size. +3. **Raise `[model.parallelism].data_parallel_shard_degree`** to shard weights/optimizer state across more ranks via FSDP. Runtime invariant (from `cosmos_framework/utils/generator/parallelism.py:50-52`): `data_parallel_replicate_degree × data_parallel_shard_degree == WORLD_SIZE` always holds — `context_parallel_shard_degree` and `cfg_parallel_shard_degree` are *overlay* axes that share dp rank slots, not separate mesh dims. Use `-1` to let `data_parallel_shard_degree` auto-fill from `torchrun` world size. 4. **Raise `[model.parallelism].context_parallel_shard_degree`** to split the sequence dimension across ranks. Helpful when activations (not weights) drive the OOM — long videos, high resolution. diff --git a/docs/sft_config.md b/docs/sft_config.md index 8df6e915..db0bb83b 100644 --- a/docs/sft_config.md +++ b/docs/sft_config.md @@ -70,14 +70,14 @@ The full pipeline (dataloader class, dataset wiring, model_instance LazyCall, et Run identity + meta-fields that pick the Hydra config tree to load. -| field | default | description | -| ------------ | ------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -| `task` | `"vfm"` | **META** — chooses which `make_config()` to call: `"vfm"` → `cosmos_framework/configs/base/config.py`, `"vlm"` → `cosmos_framework/configs/base/vlm/config.py`. Also picks the path-remap rules in `toml_config_helper.PATH_REMAPS`. | -| `experiment` | `""` | **META** — names the Hydra experiment LazyDict registered in `ConfigStore` under `experiment/`. Resolved at load time via `experiment=` (e.g. `vision_sft_nano`). | -| `project` | `""` | W&B project (team-level bucket). Flows to `config.job.project`. | -| `group` | `""` | W&B sub-label for clustering related runs (e.g. `"sft"`). Flows to `config.job.group`. | -| `name` | `""` | W&B run name; forms part of the output dir `$IMAGINAIRE_OUTPUT_ROOT////`. Leave empty (or use `${now:%Y-%m-%d}_${now:%H-%M-%S}`) for auto-timestamped subdir. | -| `wandb_mode` | `"disabled"` | `"online"` (real-time, needs `WANDB_API_KEY`), `"offline"` (log locally, sync later via `wandb sync`), or `"disabled"`. | +| field | default | description | +| ------------ | ------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `task` | `"vfm"` | **META** — chooses which `make_config()` to call: `"vfm"` → `cosmos_framework/configs/base/config.py`, `"vlm"` → `cosmos_framework/configs/base/reasoner/config.py`. Also picks the path-remap rules in `toml_config_helper.PATH_REMAPS`. | +| `experiment` | `""` | **META** — names the Hydra experiment LazyDict registered in `ConfigStore` under `experiment/`. Resolved at load time via `experiment=` (e.g. `vision_sft_nano`). | +| `project` | `""` | W&B project (team-level bucket). Flows to `config.job.project`. | +| `group` | `""` | W&B sub-label for clustering related runs (e.g. `"sft"`). Flows to `config.job.group`. | +| `name` | `""` | W&B run name; forms part of the output dir `$IMAGINAIRE_OUTPUT_ROOT////`. Leave empty (or use `${now:%Y-%m-%d}_${now:%H-%M-%S}`) for auto-timestamped subdir. | +| `wandb_mode` | `"disabled"` | `"online"` (real-time, needs `WANDB_API_KEY`), `"offline"` (log locally, sync later via `wandb sync`), or `"disabled"`. | ## `[model]` diff --git a/examples/integration/README.md b/examples/integration/README.md index 2dae6ded..fd8ec12e 100644 --- a/examples/integration/README.md +++ b/examples/integration/README.md @@ -189,7 +189,7 @@ loss.backward() # ← Your A pure level-A extraction (zero `import cosmos_framework`) is **not feasible without re-vendoring** — `Cosmos3VFMNetwork.forward` takes a `PackedSequence`, which -~2400 lines of `cosmos_framework/data/vfm/sequence_packing.py` build. These demos show +~2400 lines of `cosmos_framework/data/generator/sequence_packing.py` build. These demos show the realistic options: | Cosmos surface you keep | Trainer-level | Net-level | @@ -290,12 +290,12 @@ at a directory containing that model's `config.json`. | Topic | File | | ------------------------------- | ------------------------------------------------------------------------------ | -| OmniMoTModel definition | `cosmos_framework/model/vfm/omni_mot_model.py` | -| Cosmos3VFMNetwork (`model.net`) | `cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py` | -| PackedSequence + packer | `cosmos_framework/data/vfm/sequence_packing.py` | -| Rectified-flow loss | `cosmos_framework/model/vfm/algorithm/loss/flow_matching.py` | -| UniPC / EDM samplers | `cosmos_framework/model/vfm/diffusion/samplers/` | +| OmniMoTModel definition | `cosmos_framework/model/generator/omni_mot_model.py` | +| Cosmos3VFMNetwork (`model.net`) | `cosmos_framework/model/generator/mot/cosmos3_vfm_network.py` | +| PackedSequence + packer | `cosmos_framework/data/generator/sequence_packing.py` | +| Rectified-flow loss | `cosmos_framework/model/generator/algorithm/loss/flow_matching.py` | +| UniPC / EDM samplers | `cosmos_framework/model/generator/diffusion/samplers/` | | Checkpoint loader | `cosmos_framework/inference/model.py` (`Cosmos3OmniModel.from_pretrained_dcp`) | | Default sample args | `cosmos_framework/inference/defaults//sample_args.json` | -| FSDP / parallelism wrapping | `cosmos_framework/utils/vfm/parallelism.py` (`ParallelDims`) | +| FSDP / parallelism wrapping | `cosmos_framework/utils/generator/parallelism.py` (`ParallelDims`) | | Production trainer (skipped) | `cosmos_framework/scripts/train.py`, `examples/toml/*.toml` | diff --git a/examples/integration/net_level.py b/examples/integration/net_level.py index 1c75e747..6f315690 100644 --- a/examples/integration/net_level.py +++ b/examples/integration/net_level.py @@ -19,12 +19,12 @@ Real training samples σ from a logit-normal (image) / waver (video) distribution per `OmniMoTModel._get_train_noise_level_vision`. 3. LOSS — plain MSE on velocity. - Real training uses `cosmos_framework.model.vfm.algorithm.loss.flow_matching + Real training uses `cosmos_framework.model.generator.algorithm.loss.flow_matching .compute_flow_matching_loss`, which adds per-sample weighting, condition-mask zeroing, and `loss_scale=10` (with separate image/video scaling). 4. SAMPLER — plain Euler, no CFG, ~8 steps. - Real inference uses UniPC (`cosmos_framework.model.vfm.diffusion.samplers.unipc`) + Real inference uses UniPC (`cosmos_framework.model.generator.diffusion.samplers.unipc`) with `guidance=1.5` and 35 steps. If you train or sample with the demo's simplifications you will diverge from @@ -84,7 +84,7 @@ (just packs + calls net) If you wanted ZERO cosmos_framework imports at runtime, you would re-vendor -`cosmos_framework/data/vfm/sequence_packing.py` and the VAE into your own framework. +`cosmos_framework/data/generator/sequence_packing.py` and the VAE into your own framework. ================================================================================ RUN @@ -107,12 +107,12 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig -from cosmos_framework.data.vfm.action.domain_utils import get_domain_id -from cosmos_framework.data.vfm.action.transforms import build_sequence_plan_from_mode -from cosmos_framework.data.vfm.sequence_packing import SequencePlan, build_sequence_plans_from_data_batch +from cosmos_framework.data.generator.action.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.transforms import build_sequence_plan_from_mode +from cosmos_framework.data.generator.sequence_packing import SequencePlan, build_sequence_plans_from_data_batch from cosmos_framework.inference.args import DEFAULT_CHECKPOINT from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import tokenize_caption def _load_omni_model(*, config_dir_arg: str | None): @@ -147,8 +147,8 @@ def _load_omni_model(*, config_dir_arg: str | None): config_text = (config_dir / "config.json").read_text() for _old, _new in [ ("cosmos3._src.vfm.configs.base.", "cosmos_framework.configs.base."), - ("cosmos3._src.vfm.models.", "cosmos_framework.model.vfm."), - ("cosmos3._src.vfm.tokenizers.", "cosmos_framework.model.vfm.tokenizers."), + ("cosmos3._src.vfm.models.", "cosmos_framework.model.generator."), + ("cosmos3._src.vfm.tokenizers.", "cosmos_framework.model.generator.tokenizers."), ("cosmos3._src.imaginaire.", "cosmos_framework."), ]: config_text = config_text.replace(_old, _new) @@ -271,7 +271,7 @@ def train_one_step(model, net, batch, *, iteration: int) -> torch.Tensor: gen_data_clean = model.get_data_and_condition(batch, iteration=iteration) # Pick a mid-range noise level for the demo (real training samples sigma - # from a per-modality distribution; see cosmos_framework.model.vfm.omni_mot_model + # from a per-modality distribution; see cosmos_framework.model.generator.omni_mot_model # `_get_train_noise_level_vision`). B = gen_data_clean.batch_size # tensor_kwargs_fp32 = {"dtype": float32, "device": ...} — keeps demo @@ -294,7 +294,7 @@ def train_one_step(model, net, batch, *, iteration: int) -> torch.Tensor: v_pred = out["preds_vision"] # list of [C, T, H, W] # ── 3. Custom flow-matching loss (MSE on velocity). This is what - # `cosmos_framework.model.vfm.algorithm.loss.flow_matching.compute_flow_matching_loss` + # `cosmos_framework.model.generator.algorithm.loss.flow_matching.compute_flow_matching_loss` # computes, minus the per-sample weighting & condition masking. v_target = gen_data_noised.vt_target_vision # list of [C, T, H, W] loss = sum(F.mse_loss(p.float(), t.float()) for p, t in zip(v_pred, v_target)) @@ -313,7 +313,7 @@ def train_one_step(model, net, batch, *, iteration: int) -> torch.Tensor: def sample(model, net, batch, *, num_steps: int = 12) -> dict: """N-step Euler integration of dx/dt = v(x,t) — no cosmos_framework sampler involved. - Production cosmos_framework uses UniPC/EDM (cosmos_framework.model.vfm.diffusion.samplers.*). + Production cosmos_framework uses UniPC/EDM (cosmos_framework.model.generator.diffusion.samplers.*). Plain Euler keeps the loop on one screen and surfaces where `net` is called. Returns a dict: diff --git a/examples/integration/trainer_level_inference.py b/examples/integration/trainer_level_inference.py index 6b4b81bf..ede1a399 100644 --- a/examples/integration/trainer_level_inference.py +++ b/examples/integration/trainer_level_inference.py @@ -33,8 +33,8 @@ cosmos_framework.inference.common.init.init_script → 1-line torch.distributed init cosmos_framework.inference.{args,inference} → OmniSampleOverrides + get_sample_data (T2I/T2V only) - cosmos_framework.data.vfm.{action,sequence_packing} → SequencePlan helpers (action/sound) - cosmos_framework.model.vfm.reasoner.qwen3_vl.utils.tokenize_caption + cosmos_framework.data.generator.{action,sequence_packing} → SequencePlan helpers (action/sound) + cosmos_framework.model.generator.reasoner.qwen3_vl.utils.tokenize_caption model.generate_samples_from_batch(batch, seed) → THE inference call (CFG + sampler) model.decode(latent) → VAE decode @@ -65,13 +65,13 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig -from cosmos_framework.data.vfm.action.domain_utils import get_domain_id -from cosmos_framework.data.vfm.action.transforms import build_sequence_plan_from_mode -from cosmos_framework.data.vfm.sequence_packing import SequencePlan +from cosmos_framework.data.generator.action.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.transforms import build_sequence_plan_from_mode +from cosmos_framework.data.generator.sequence_packing import SequencePlan from cosmos_framework.inference.args import DEFAULT_CHECKPOINT, OmniSampleOverrides from cosmos_framework.inference.inference import get_sample_data from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import tokenize_caption from cosmos_framework.tools.visualize.video import save_img_or_video @@ -108,8 +108,8 @@ def _load_omni_model(*, config_dir_arg: str | None): config_text = (config_dir / "config.json").read_text() for _old, _new in [ ("cosmos3._src.vfm.configs.base.", "cosmos_framework.configs.base."), - ("cosmos3._src.vfm.models.", "cosmos_framework.model.vfm."), - ("cosmos3._src.vfm.tokenizers.", "cosmos_framework.model.vfm.tokenizers."), + ("cosmos3._src.vfm.models.", "cosmos_framework.model.generator."), + ("cosmos3._src.vfm.tokenizers.", "cosmos_framework.model.generator.tokenizers."), ("cosmos3._src.imaginaire.", "cosmos_framework."), ]: config_text = config_text.replace(_old, _new) diff --git a/examples/integration/trainer_level_training.py b/examples/integration/trainer_level_training.py index 2656bce4..c304bab1 100644 --- a/examples/integration/trainer_level_training.py +++ b/examples/integration/trainer_level_training.py @@ -35,15 +35,15 @@ cosmos_framework.inference.model.Cosmos3OmniModel → model class (random-init in this demo; use `.from_pretrained_dcp(...)` for real weights) cosmos_framework.inference.common.init.init_script → 1-line torch.distributed init - cosmos_framework.model.vfm.reasoner.qwen3_vl.utils.tokenize_caption + cosmos_framework.model.generator.reasoner.qwen3_vl.utils.tokenize_caption → text tokenizer (modelling pkg) model.training_step(batch, iteration) → THE training step (flow-matching loss) model.config.{action_gen,sound_gen,vision_gen,…} → modality flags What we DO NOT use: cosmos_framework.scripts.train, cosmos_framework.trainer.* → CLI + Trainer class - cosmos_framework.data.vfm.joint_dataloader.* → iterative joint dataloader - cosmos_framework.data.vfm.augmentor_provider.* → text/video augmentor pipeline + cosmos_framework.data.generator.joint_dataloader.* → iterative joint dataloader + cosmos_framework.data.generator.augmentor_provider.* → text/video augmentor pipeline cosmos_framework.inference.inference.OmniInference → inference pipeline ================================================================================ @@ -97,7 +97,7 @@ ≥ 8 GPUs and/or LoRA — see `cosmos_framework.scripts.train` and `examples/toml/*.toml`. To make full-fine-tuning fit on real hardware, you would either: - - shard with FSDP (`cosmos_framework.utils.vfm.parallelism.ParallelDims` + FSDP wrap), + - shard with FSDP (`cosmos_framework.utils.generator.parallelism.ParallelDims` + FSDP wrap), - inject LoRA (`model.add_lora(...)`), or - swap the optimizer for one with lower state (Adafactor, 8-bit AdamW). @@ -121,12 +121,12 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig -from cosmos_framework.data.vfm.action.domain_utils import get_domain_id -from cosmos_framework.data.vfm.action.transforms import build_sequence_plan_from_mode -from cosmos_framework.data.vfm.sequence_packing import SequencePlan +from cosmos_framework.data.generator.action.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.transforms import build_sequence_plan_from_mode +from cosmos_framework.data.generator.sequence_packing import SequencePlan from cosmos_framework.inference.args import DEFAULT_CHECKPOINT from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel -from cosmos_framework.model.vfm.reasoner.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import tokenize_caption def _load_omni_model(*, config_dir_arg: str | None): @@ -162,8 +162,8 @@ def _load_omni_model(*, config_dir_arg: str | None): config_text = (config_dir / "config.json").read_text() for _old, _new in [ ("cosmos3._src.vfm.configs.base.", "cosmos_framework.configs.base."), - ("cosmos3._src.vfm.models.", "cosmos_framework.model.vfm."), - ("cosmos3._src.vfm.tokenizers.", "cosmos_framework.model.vfm.tokenizers."), + ("cosmos3._src.vfm.models.", "cosmos_framework.model.generator."), + ("cosmos3._src.vfm.tokenizers.", "cosmos_framework.model.generator.tokenizers."), ("cosmos3._src.imaginaire.", "cosmos_framework."), ]: config_text = config_text.replace(_old, _new) @@ -187,7 +187,7 @@ def _tokenize(model, caption: str, device) -> torch.Tensor: is_video=False, use_system_prompt=model.vlm_config.use_system_prompt, ) - # Shape [1, N_tok]. The collate format in cosmos_framework.data.vfm.joint_dataloader + # Shape [1, N_tok]. The collate format in cosmos_framework.data.generator.joint_dataloader # keeps text_token_ids as a list of [1, N] tensors (one per sample) because # token counts vary across the batch. return torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0) @@ -277,7 +277,7 @@ def make_action_fdm_batch(model, *, caption: str, num_video_frames: int = 5, See `cosmos_framework/inference/action.py: build_action_batch` for the canonical impl. `domain_name` selects the cross-embodiment routing; see - `cosmos_framework/data/vfm/action/domain_utils.py` for the full list of supported + `cosmos_framework/data/generator/action/domain_utils.py` for the full list of supported embodiments. """ # First frame is the conditioning anchor; remaining frames are predicted. diff --git a/examples/launch_sft_llava_ov.sh b/examples/launch_sft_llava_ov.sh index cc56d42b..afb3d99d 100755 --- a/examples/launch_sft_llava_ov.sh +++ b/examples/launch_sft_llava_ov.sh @@ -7,7 +7,7 @@ # cosmos_framework.scripts.train against # examples/toml/sft_config/llava_ov.toml. # -# [job].task = "vlm" — picks cosmos_framework/configs/base/vlm/config.py as the base config. +# [job].task = "vlm" — picks cosmos_framework/configs/base/reasoner/config.py as the base config. # # The dataset streams from the HuggingFace Hub, so DATASET_PATH / # WAN_VAE_PATH / BASE_CHECKPOINT_PATH are NOT required. diff --git a/examples/launch_sft_videophy2_nano.sh b/examples/launch_sft_videophy2_nano.sh index b0818fbc..56046450 100755 --- a/examples/launch_sft_videophy2_nano.sh +++ b/examples/launch_sft_videophy2_nano.sh @@ -6,7 +6,7 @@ # via CosmosDataLoader). Drives cosmos_framework.scripts.train against # examples/toml/sft_config/videophy2_sft_nano.toml. # -# [job].task = "vlm" — picks cosmos_framework/configs/base/vlm/config.py as the base config. +# [job].task = "vlm" — picks cosmos_framework/configs/base/reasoner/config.py as the base config. # # Required env: # VIDEOPHYSICS_ROOT dir containing videophy2_train/ and videophy2_val/ diff --git a/examples/toml/sft_config/llava_ov.toml b/examples/toml/sft_config/llava_ov.toml index f07a07b1..39c81202 100644 --- a/examples/toml/sft_config/llava_ov.toml +++ b/examples/toml/sft_config/llava_ov.toml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: OpenMDW-1.1 # pre_exp012_llava_ov — VLM training on lmms-lab/LLaVA-OneVision-Data -# via CosmosDataLoader. Base config = cosmos_framework/configs/base/vlm/config.py +# via CosmosDataLoader. Base config = cosmos_framework/configs/base/reasoner/config.py # (selected by [job].task="vlm"). # # One knob that the SFTExperimentConfig dataclass does NOT model — supply @@ -93,7 +93,7 @@ enabled = false [trainer.callbacks.grad_clip] clip_norm = 1.0 -force_finite = false # matches VLM default in cosmos_framework/configs/base/vlm/defaults/callbacks.py:55 +force_finite = false # matches VLM default in cosmos_framework/configs/base/reasoner/defaults/callbacks.py:55 [checkpoint] keys_to_skip_loading = [] diff --git a/examples/toml/sft_config/llava_ov_mapstyle_dataloader.toml b/examples/toml/sft_config/llava_ov_mapstyle_dataloader.toml index 5b1a8d9c..8a67acea 100644 --- a/examples/toml/sft_config/llava_ov_mapstyle_dataloader.toml +++ b/examples/toml/sft_config/llava_ov_mapstyle_dataloader.toml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: OpenMDW-1.1 # pre_exp012_llava_ov_mapstyle_dataloader — map-style resumable VLM recipe -# Base config = cosmos_framework/configs/base/vlm/config.py (task="vlm"). +# Base config = cosmos_framework/configs/base/reasoner/config.py (task="vlm"). # # Loads LLaVA-OneVision-Data ai2d(gpt4v) as a real map-style Dataset # (load_dataset(streaming=False)), filters + caps to 4000 rows, and shards it diff --git a/examples/toml/sft_config/videophy2_sft_nano.toml b/examples/toml/sft_config/videophy2_sft_nano.toml index a183db33..431c406d 100644 --- a/examples/toml/sft_config/videophy2_sft_nano.toml +++ b/examples/toml/sft_config/videophy2_sft_nano.toml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: OpenMDW-1.1 # videophy2_sft_nano — VLM dialog SFT on VideoPhy-2 via CosmosDataLoader. -# Base config = cosmos_framework/configs/base/vlm/config.py (selected by [job].task="vlm"). +# Base config = cosmos_framework/configs/base/reasoner/config.py (selected by [job].task="vlm"). # # Dataset prep: # python -m cosmos_framework.scripts.vlm.prepare_videophy2_from_hf \ diff --git a/tests/_reasoner_logits_probe.py b/tests/_reasoner_logits_probe.py index 8830cb68..efa72ba2 100644 --- a/tests/_reasoner_logits_probe.py +++ b/tests/_reasoner_logits_probe.py @@ -41,7 +41,7 @@ def _install_determinism() -> None: def _install_logits_probe(dump_path: str) -> None: import torch - import cosmos_framework.model.vfm.mot.unified_mot as unified_mot + import cosmos_framework.model.generator.mot.unified_mot as unified_mot original = unified_mot._sample_next_token state = {"saved": False} diff --git a/tests/_stage_h100_inputs.sh b/tests/_stage_h100_inputs.sh index d59c34ed..db9d9b09 100755 --- a/tests/_stage_h100_inputs.sh +++ b/tests/_stage_h100_inputs.sh @@ -33,7 +33,7 @@ mkdir -p "$HF_HOME" echo ">>> $(date '+%H:%M:%S') HF_HOME=$HF_HOME STAGE_DIR=$STAGE_DIR REPO_ROOT=$REPO_ROOT" # ---------------------------------------------------------------------------- -# 0. Python env: uv sync + pinned transformers. cosmos_framework/utils/vfm/monkey_patch.py +# 0. Python env: uv sync + pinned transformers. cosmos_framework/utils/generator/monkey_patch.py # hard-rejects every transformers version except 4.57.1 (pyproject's # `>=4.57.1,<5.0` is looser than what actually works at runtime). # ---------------------------------------------------------------------------- @@ -71,7 +71,7 @@ echo "WAN_VAE_PATH=$WAN_VAE_PATH" # ---------------------------------------------------------------------------- # 3. VLM backbone for launch_vlm_llava_ov (Qwen3-VL-8B-Instruct). Cosmos's # tokenizer dispatcher checks for the substring `Qwen/Qwen3-VL` in the path -# (cosmos_framework/data/vfm/processors/__init__.py); HF's cache uses +# (cosmos_framework/data/generator/processors/__init__.py); HF's cache uses # `models--Qwen--Qwen3-VL-8B-Instruct/snapshots/...` which doesn't match. We # add a `$STAGE_DIR/Qwen/Qwen3-VL-8B-Instruct` symlink so the dispatched # substring is present, and point `MODEL_PATH` at the symlink. From 6efeea330369bd882586d070cb61172f383801ce Mon Sep 17 00:00:00 2001 From: lfengad Date: Fri, 3 Jul 2026 00:39:45 +0800 Subject: [PATCH 09/26] Convert Cosmos3-Nano vision tower in reasoner VLM export + fix public-alias remap (#79) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Fixes two issues in the Reasoner / VideoPhy-2 SFT "Step 2" checkpoint prep (`convert_model_to_vlm_safetensors`), so the exported VLM is fully sourced from Cosmos3-Nano, and adds a regression test. ### 1. Fix `vlm`→`reasoner` remap for public-alias file paths (`public_model_config.py`) The i4 `vlm`→`reasoner` rename was applied to the module-path remap helpers but **not** to their file-path siblings. The current released `nvidia/Cosmos3-Nano` snapshot stores paths as public URIs (`cosmos3://vfm/models/vlm/qwen3_vl/configs/...`), which route through the file-path helpers and resolved to a non-existent `cosmos_framework/model/vfm/vlm/qwen3_vl/...`, crashing checkpoint load with `FileNotFoundError`. This affects loading the released checkpoint generally (not just the converter). Added the missing `vlm`↔`reasoner` rules to both `_replace_vfm_file_prefix` (public→runtime) and `_public_string_from_runtime_file_prefix` (runtime→public), mirroring the module-path helpers; the public↔runtime round-trip is stable. ### 2. Convert the Cosmos3-Nano vision tower (`convert_model_to_vlm_safetensors.py`) The converter previously extracted only the language model and left the visual tower as the **stock Qwen3-VL** weights — because the generation model instantiates no `visual` submodule, so the 351 `vision_encoder/` tensors never appear in its state dict. Now the visual tower is loaded directly from the checkpoint's `vision_encoder/` shards (`_load_vision_state`) and overlaid as `model.visual.*`. Checkpoints without a vision tower (Text2Image / Image2Video) keep Qwen3-VL's tower unchanged. The merged export is now 100% Cosmos3-Nano-sourced (399 LM + 351 visual). ### 3. Regression test (`tests/launch_regression_test.py`) `test_convert_reasoner_converts_all_qwen_tensors` asserts: 1. merged tensor set == stock Qwen3-VL set (all included, none extra); 2. every `model.visual.*` tensor matches the Cosmos3-Nano `vision_encoder/` source bit-for-bit; 3. a non-trivial subset differs from stock Qwen3-VL (so it catches the vision-drop regression — not a no-op); 4. the language tower was overlaid too. The ~16 GB output is regenerated per run and removed via a finalizer; inputs stay in the shared HF cache. ## Test plan - ✅ New test passes on a 4-GPU node: `1 passed`, `351 visual tensors all sourced from Cosmos3-Nano; 309 differ from stock Qwen3-VL`, output cleaned up (basetemp 13K after finalizer). - ✅ End-to-end: VideoPhy-2 Reasoner SFT loads `750 keys` from the merged checkpoint and trains cleanly across 4 GPUs (loss ~0.66→~0.15 over 20 iters). - ✅ Existing `public_model_config_test.py` still passes; public↔runtime round-trip verified stable. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- .../inference/common/public_model_config.py | 18 ++ .../convert_model_to_vlm_safetensors.py | 69 ++++++-- .../scripts/{vlm => reasoner}/__init__.py | 0 .../{vlm => reasoner}/eval_videophy2.py | 4 +- .../prepare_videophy2_from_hf.py | 2 +- docs/training.md | 2 +- examples/launch_sft_videophy2_nano.sh | 2 +- .../toml/sft_config/videophy2_sft_nano.toml | 2 +- tests/launch_regression_test.py | 162 ++++++++++++++++++ 9 files changed, 245 insertions(+), 16 deletions(-) rename cosmos_framework/scripts/{vlm => reasoner}/__init__.py (100%) rename cosmos_framework/scripts/{vlm => reasoner}/eval_videophy2.py (98%) rename cosmos_framework/scripts/{vlm => reasoner}/prepare_videophy2_from_hf.py (99%) diff --git a/cosmos_framework/inference/common/public_model_config.py b/cosmos_framework/inference/common/public_model_config.py index 66bb9e31..4bd2c679 100644 --- a/cosmos_framework/inference/common/public_model_config.py +++ b/cosmos_framework/inference/common/public_model_config.py @@ -249,6 +249,15 @@ def _to_public_string(value: str) -> str: def _public_string_from_runtime_file_prefix(value: str, *, package: str) -> str | None: replacements = ( + # reasoner → vlm (inverse of the upstream rename). MUST precede the + # general vfm rules so the reasoner subtree isn't shadowed by them. + # Mirrors the module-path rules in ``_canonicalize_module_path``. + (f"{package}/configs/base/defaults/reasoner/", "configs/base/defaults/vlm/"), + (f"{package}/configs/base/reasoner/", "configs/base/vlm/"), + (f"{package}/model/generator/reasoner/", "models/vlm/"), + (f"{package}/data/generator/augmentors/reasoner/", "datasets/augmentors/vlm/"), + (f"{package}/data/generator/reasoner/", "datasets/vlm/"), + (f"{package}/utils/generator/reasoner/", "utils/vlm/"), (f"{package}/configs/base/", "configs/base/"), (f"{package}/model/generator/tokenizers/", "tokenizers/"), (f"{package}/model/generator/diffusion/", "diffusion/"), @@ -279,6 +288,15 @@ def _from_public_string(value: str) -> str: def _replace_vfm_file_prefix(suffix: str, *, package: str) -> str: replacements = ( + # vlm → reasoner (upstream rename). MUST precede the general vfm rules + # so the specific vlm→reasoner subtree isn't shadowed by them. Mirrors + # the module-path rules in ``_replace_vfm_module_prefix``. + ("configs/base/defaults/vlm/", f"{package}/configs/base/defaults/reasoner/"), + ("configs/base/vlm/", f"{package}/configs/base/reasoner/"), + ("models/vlm/", f"{package}/model/generator/reasoner/"), + ("datasets/augmentors/vlm/", f"{package}/data/generator/augmentors/reasoner/"), + ("datasets/vlm/", f"{package}/data/generator/reasoner/"), + ("utils/vlm/", f"{package}/utils/generator/reasoner/"), ("configs/base/", f"{package}/configs/base/"), ("models/", f"{package}/model/generator/"), ("tokenizers/", f"{package}/model/generator/tokenizers/"), diff --git a/cosmos_framework/scripts/convert_model_to_vlm_safetensors.py b/cosmos_framework/scripts/convert_model_to_vlm_safetensors.py index 1795cf93..528f630f 100644 --- a/cosmos_framework/scripts/convert_model_to_vlm_safetensors.py +++ b/cosmos_framework/scripts/convert_model_to_vlm_safetensors.py @@ -4,9 +4,13 @@ """Convert a Cosmos3 OmniMoT checkpoint into a Qwen3-VL HF safetensors directory. The result is a complete HF directory (config.json + safetensors + tokenizer) -shaped like a Qwen3VLForConditionalGeneration release, but with the language-model -tensors sourced from the Cosmos3 OmniMoT and the visual tower kept from the -Qwen3-VL-Instruct release (Cosmos3 has no visual counterpart). +shaped like a Qwen3VLForConditionalGeneration release, with the language-model +tensors sourced from the Cosmos3 OmniMoT. When the Cosmos3 checkpoint ships its +own (Qwen3-VL-shaped) visual tower — as the base ``Cosmos3-Nano`` release does — +those visual tensors are overlaid too, so the exported VLM uses Cosmos3's trained +vision encoder. Task-specialized checkpoints (e.g. Text2Image, Image2Video) omit +vision weights; for those the visual tower is kept from the Qwen3-VL-Instruct +release instead. Pass the resulting path as ``[model.backbone].safetensors_path`` in a VLM SFT TOML to bootstrap training from Cosmos3 weights while keeping the public HF @@ -26,24 +30,35 @@ } ) +import collections +from pathlib import Path from typing import Annotated import pydantic import torch import tyro +from safetensors import safe_open from torch.distributed.checkpoint.state_dict import get_model_state_dict from transformers import AutoProcessor, AutoTokenizer, Qwen3VLForConditionalGeneration from cosmos_framework.inference.args import OmniSetupOverrides from cosmos_framework.inference.common.args import CheckpointOverrides, ResolvedPath -from cosmos_framework.inference.model import Cosmos3OmniModel +from cosmos_framework.inference.model import ( + Cosmos3OmniModel, + _diffusers_weight_map, + _is_diffusers_checkpoint, +) # Cosmos3 OmniMoT exposes its inner LM under ``net.language_model.*``; Qwen3VL # expects ``lm_head.*`` (top-level) and ``model.language_model.*`` (text-decoder # sub-tree of the VLM). The OmniMoT MoE-generation pathway (``*_moe_gen``) has -# no Qwen3VL counterpart and is dropped. +# no Qwen3VL counterpart and is dropped. The visual tower is NOT part of the +# generation model's state dict (Cosmos3-Nano instantiates no ``visual`` +# submodule); it is read straight from the checkpoint's ``vision_encoder/`` +# shards by ``_load_vision_state`` instead. _OMNIMOT_LM_PREFIX = "net.language_model." +_VISION_ENCODER_PREFIX = "vision_encoder/" def _remap_to_qwen3vl(key: str) -> str | None: @@ -60,6 +75,31 @@ def _remap_to_qwen3vl(key: str) -> str | None: return None +def _load_vision_state(checkpoint_path: Path) -> dict[str, torch.Tensor]: + """Read the checkpoint's Qwen3-VL-shaped visual tower as ``model.visual.*``. + + The visual tower ships as a standalone ``vision_encoder/`` tree in the + diffusers checkpoint and is absent from the generation model's state dict, + so it must be loaded directly from the safetensors shards. Its tensor names + already match Qwen3VL's visual tower verbatim, so they only need the + ``model.visual.`` prefix. Returns ``{}`` for non-diffusers checkpoints or + checkpoints that ship no vision tower (e.g. Text2Image / Image2Video). + """ + if not _is_diffusers_checkpoint(checkpoint_path): + return {} + weight_map = _diffusers_weight_map(checkpoint_path) + keys_by_file: dict[str, list[str]] = collections.defaultdict(list) + for diff_key, rel_path in weight_map.items(): + if rel_path.startswith(_VISION_ENCODER_PREFIX): + keys_by_file[rel_path].append(diff_key) + vision_state: dict[str, torch.Tensor] = {} + for rel_path, diff_keys in keys_by_file.items(): + with safe_open(str(checkpoint_path / rel_path), framework="pt") as f: + for diff_key in diff_keys: + vision_state["model.visual." + diff_key] = f.get_tensor(diff_key).to(torch.bfloat16) + return vision_state + + class Args(pydantic.BaseModel): checkpoint: CheckpointOverrides """Cosmos3 OmniMoT checkpoint (e.g. Cosmos3-Nano).""" @@ -72,7 +112,7 @@ class Args(pydantic.BaseModel): def convert_model_to_vlm_safetensors(args: Args) -> None: print(f"Loading Cosmos3 checkpoint via CheckpointOverrides...") cosmos3_config = args.checkpoint.build_checkpoint(checkpoints=OmniSetupOverrides.CHECKPOINTS) - cosmos3_path = cosmos3_config.download_checkpoint() + cosmos3_path = Path(cosmos3_config.download_checkpoint()) cosmos3_model = Cosmos3OmniModel.from_pretrained_dcp(cosmos3_path) cosmos3_state = get_model_state_dict(cosmos3_model.model) @@ -81,9 +121,18 @@ def convert_model_to_vlm_safetensors(args: Args) -> None: new_k = _remap_to_qwen3vl(k) if new_k is not None and isinstance(v, torch.Tensor): lm_state[new_k] = v - print(f" extracted {len(lm_state)} LM tensors from {len(cosmos3_state)} OmniMoT tensors") del cosmos3_state, cosmos3_model + # The visual tower is a standalone tree in the checkpoint, not part of the + # generation model's state dict, so load it directly. Empty when the + # checkpoint ships no vision tower (keeps Qwen3-VL's tower in that case). + vision_state = _load_vision_state(cosmos3_path) + lm_state.update(vision_state) + print( + f" extracted {len(lm_state)} tensors " + f"({len(lm_state) - len(vision_state)} LM + {len(vision_state)} visual)" + ) + print(f"Loading {args.vlm_model_name} (visual tower + LM defaults, bf16, CPU)...") model = Qwen3VLForConditionalGeneration.from_pretrained( args.vlm_model_name, dtype=torch.bfloat16 @@ -91,12 +140,12 @@ def convert_model_to_vlm_safetensors(args: Args) -> None: incompatible = model.load_state_dict(lm_state, strict=False) n_overlaid = len(lm_state) - len(incompatible.unexpected_keys) - print(f" overlaid {n_overlaid}/{len(lm_state)} LM tensors " + print(f" overlaid {n_overlaid}/{len(lm_state)} Cosmos3 tensors " f"(unexpected={len(incompatible.unexpected_keys)}, " - f"missing-in-LM-state={len(incompatible.missing_keys)} — these are visual/etc kept from HF)") + f"missing-in-cosmos3-state={len(incompatible.missing_keys)} — kept from HF)") if incompatible.unexpected_keys: raise RuntimeError( - f"Cosmos3 LM tensors not present in Qwen3VL: " + f"Cosmos3 tensors not present in Qwen3VL: " f"{incompatible.unexpected_keys[:5]}{'...' if len(incompatible.unexpected_keys) > 5 else ''}" ) diff --git a/cosmos_framework/scripts/vlm/__init__.py b/cosmos_framework/scripts/reasoner/__init__.py similarity index 100% rename from cosmos_framework/scripts/vlm/__init__.py rename to cosmos_framework/scripts/reasoner/__init__.py diff --git a/cosmos_framework/scripts/vlm/eval_videophy2.py b/cosmos_framework/scripts/reasoner/eval_videophy2.py similarity index 98% rename from cosmos_framework/scripts/vlm/eval_videophy2.py rename to cosmos_framework/scripts/reasoner/eval_videophy2.py index dec6e650..35bfef82 100644 --- a/cosmos_framework/scripts/vlm/eval_videophy2.py +++ b/cosmos_framework/scripts/reasoner/eval_videophy2.py @@ -20,12 +20,12 @@ Single-GPU example:: - python -m cosmos_framework.scripts.vlm.eval_videophy2 \\ + python -m cosmos_framework.scripts.reasoner.eval_videophy2 \\ --hf_ckpt $HF_CKPT --val_root $VAL_ROOT --results_dir $OUT 8-GPU data-parallel example:: - torchrun --nproc_per_node=8 -m cosmos_framework.scripts.vlm.eval_videophy2 \\ + torchrun --nproc_per_node=8 -m cosmos_framework.scripts.reasoner.eval_videophy2 \\ --hf_ckpt $HF_CKPT --val_root $VAL_ROOT --results_dir $OUT --batch_size 2 The inference path here is intentionally lightweight — it is expected to be diff --git a/cosmos_framework/scripts/vlm/prepare_videophy2_from_hf.py b/cosmos_framework/scripts/reasoner/prepare_videophy2_from_hf.py similarity index 99% rename from cosmos_framework/scripts/vlm/prepare_videophy2_from_hf.py rename to cosmos_framework/scripts/reasoner/prepare_videophy2_from_hf.py index a5c79346..7a40e65b 100644 --- a/cosmos_framework/scripts/vlm/prepare_videophy2_from_hf.py +++ b/cosmos_framework/scripts/reasoner/prepare_videophy2_from_hf.py @@ -24,7 +24,7 @@ Example:: - python -m cosmos_framework.scripts.vlm.prepare_videophy2_from_hf \\ + python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \\ --out_root examples/data/videophysics \\ --split both """ diff --git a/docs/training.md b/docs/training.md index fe07f572..5b11e5ff 100644 --- a/docs/training.md +++ b/docs/training.md @@ -109,7 +109,7 @@ Launch shell: `examples/launch_sft_videophy2_nano.sh` ```shell # Step 1 (data): materialize the public HF dataset into the canonical local layout # (videophy2_{train,val}/{meta.json, media/, text/}). -python -m cosmos_framework.scripts.vlm.prepare_videophy2_from_hf \ +python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ --out_root examples/data/videophysics --split both ``` diff --git a/examples/launch_sft_videophy2_nano.sh b/examples/launch_sft_videophy2_nano.sh index 56046450..a9302fb3 100755 --- a/examples/launch_sft_videophy2_nano.sh +++ b/examples/launch_sft_videophy2_nano.sh @@ -11,7 +11,7 @@ # Required env: # VIDEOPHYSICS_ROOT dir containing videophy2_train/ and videophy2_val/ # (each with meta.json + media/ + text/). Populate via -# `python -m cosmos_framework.scripts.vlm.prepare_videophy2_from_hf`. +# `python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf`. # # Optional env: # HF_TOKEN for gated Qwen3-VL-8B-Instruct downloads. diff --git a/examples/toml/sft_config/videophy2_sft_nano.toml b/examples/toml/sft_config/videophy2_sft_nano.toml index 431c406d..35b8e785 100644 --- a/examples/toml/sft_config/videophy2_sft_nano.toml +++ b/examples/toml/sft_config/videophy2_sft_nano.toml @@ -5,7 +5,7 @@ # Base config = cosmos_framework/configs/base/reasoner/config.py (selected by [job].task="vlm"). # # Dataset prep: -# python -m cosmos_framework.scripts.vlm.prepare_videophy2_from_hf \ +# python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ # --out_root $VIDEOPHYSICS_ROOT --split train # and again with --split val # # Required env at launch: VIDEOPHYSICS_ROOT (read by the experiment Python). diff --git a/tests/launch_regression_test.py b/tests/launch_regression_test.py index 04322e5d..24aec9d6 100644 --- a/tests/launch_regression_test.py +++ b/tests/launch_regression_test.py @@ -471,6 +471,41 @@ def _make_dcp() -> Path: shutil.rmtree(stage, ignore_errors=True) +@pytest.fixture(scope="module") +def qwen_vl_model_path() -> str: + """Local Qwen3-VL-8B-Instruct snapshot for the convert-reasoner test. + + Honors a pre-set ``MODEL_PATH`` (so ``source env.sh`` / the ``h100_inputs`` + staging is reused); otherwise downloads the pinned revision from the Hub. + Independent of ``h100_inputs`` so this test does not depend on the Nano→DCP + conversion that ``h100_inputs`` also performs for the training specs. + """ + if shutil.which("uvx") is None: + pytest.skip("uvx not on PATH -- required to stage Qwen3-VL") + return os.environ.get("MODEL_PATH") or _hf_download( + ["Qwen/Qwen3-VL-8B-Instruct", "--revision", _QWEN_VL_REVISION] + ) + + +def _weight_map(directory: Path) -> dict[str, str]: + """``key -> shard filename`` for a (sharded or single-file) safetensors dir.""" + index = directory / "model.safetensors.index.json" + if index.exists(): + import json + + return json.loads(index.read_text())["weight_map"] + return {} # single-file: callers fall back to model.safetensors + + +def _load_st_tensor(directory: Path, key: str, weight_map: dict[str, str]): + """Lazily read a single tensor by key (keeps peak memory to one tensor).""" + from safetensors import safe_open + + rel = weight_map.get(key, "model.safetensors") + with safe_open(str(directory / rel), framework="pt") as f: + return f.get_tensor(key) + + # --- tests ------------------------------------------------------------------- @@ -558,6 +593,133 @@ def test_launch_regression(spec_key: str, tmp_path: Path, h100_inputs: dict[str, """Re-run ``spec``'s torchrun command and check loss / grad-norm against goldens.""" _assert_spec_matches_goldens(spec_key, tmp_path, h100_inputs) + @pytest.mark.level(2) + @pytest.mark.gpus(4) + def test_convert_reasoner_converts_all_qwen_tensors( + tmp_path: Path, qwen_vl_model_path: str, request: pytest.FixtureRequest + ) -> None: + """Regression guard for ``convert_model_to_vlm_safetensors`` (Reasoner / + VideoPhy-2 SFT "Step 2": merge Cosmos3-Nano onto the Qwen3-VL shell). + + The converter always *saves* a full ``Qwen3VLForConditionalGeneration``, + so its tensor set matches Qwen3-VL by construction — a bare key-coverage + check is trivially true. The invariant that actually regressed once (the + visual tower was silently kept as the stock Qwen3-VL weights) is that + **every** Qwen3-VL tensor is sourced from Cosmos3-Nano. This asserts: + + 1. the merged tensor set == the stock Qwen3-VL tensor set (all + included, none extra); + 2. the whole visual tower (``model.visual.*``) matches the Cosmos3-Nano + ``vision_encoder/`` source bit-for-bit (converted from Cosmos3, not + left as the stock Qwen3-VL default); + 3. that source genuinely differs from stock Qwen3-VL for a non-trivial + number of visual tensors — so check (2) has teeth against the + vision-drop regression. (Cosmos3's tower is derived from Qwen3-VL, so + a subset of tensors, e.g. some biases, legitimately coincide, which + is why this is a "some differ", not "all differ", check); and + 4. the language tower was overlaid too (layer-0 projection *weights* + differ from stock Qwen3-VL). + + CPU-only (the converter forces ``COSMOS_DEVICE=cpu``); it carries the + ``gpus(4)`` marker only so it is collected by the same 4-GPU regression + invocation as the launch specs. + """ + import torch + + qwen_dir = Path(qwen_vl_model_path) + out_dir = tmp_path / "Cosmos3-Nano-VLM" + # The merged output is ~16 GB; pytest only rotates ``tmp_path`` (keeps the + # last few runs), so delete it unconditionally after the test to avoid + # filling the temp filesystem. Inputs (the Qwen3-VL / Cosmos3-Nano HF + # caches) are shared and intentionally left in place. + request.addfinalizer(lambda: shutil.rmtree(out_dir, ignore_errors=True)) + + # Run the real conversion, reusing the staged Qwen3-VL snapshot as the + # shell so the test does not download an 8B model a second time. + env = os.environ.copy() + env["PYTHONPATH"] = f".:{env.get('PYTHONPATH', '')}" + result = subprocess.run( + [ + sys.executable, "-m", "cosmos_framework.scripts.convert_model_to_vlm_safetensors", + "--checkpoint-path", "Cosmos3-Nano", + "-o", str(out_dir), + "--vlm-model-name", str(qwen_dir), + ], + cwd=str(REPO_ROOT), + env=env, + ) + assert result.returncode == 0, ( + f"convert_model_to_vlm_safetensors failed with exit code {result.returncode}" + ) + + merged_wm = _weight_map(out_dir) + qwen_wm = _weight_map(qwen_dir) + assert merged_wm and qwen_wm, "both checkpoints must be sharded safetensors with an index" + + # (1) Coverage: identical tensor sets — every Qwen3-VL tensor included, none extra. + missing = sorted(set(qwen_wm) - set(merged_wm)) + extra = sorted(set(merged_wm) - set(qwen_wm)) + assert not missing and not extra, ( + f"merged tensor set != stock Qwen3-VL: missing={missing[:10]} extra={extra[:10]}" + ) + + # (2)+(3) Visual tower: every tensor is converted from the Cosmos3-Nano + # ``vision_encoder/`` source bit-for-bit, and a non-trivial subset differs + # from stock Qwen3-VL (so (2) actually distinguishes a converted tower + # from a stock one — the vision-drop regression). + from cosmos_framework.inference.args import OmniSetupOverrides + from cosmos_framework.inference.common.args import CheckpointOverrides + from safetensors import safe_open + + nano_dir = Path( + CheckpointOverrides(checkpoint_path="Cosmos3-Nano") + .build_checkpoint(checkpoints=OmniSetupOverrides.CHECKPOINTS) + .download_checkpoint() + ) + vision_src = nano_dir / "vision_encoder" / "model.safetensors" + assert vision_src.exists(), f"Cosmos3-Nano vision tower not found at {vision_src}" + + visual_keys = sorted(k for k in merged_wm if k.startswith("model.visual.")) + assert visual_keys, "no model.visual.* tensors in merged checkpoint" + n_differ_from_stock = 0 + with safe_open(str(vision_src), framework="pt") as src: + src_keys = set(src.keys()) + for k in visual_keys: + sub = k[len("model.visual."):] + assert sub in src_keys, f"{k} has no Cosmos3-Nano vision_encoder counterpart" + got = _load_st_tensor(out_dir, k, merged_wm).float() + want = src.get_tensor(sub).float() + assert torch.equal(got, want), f"visual tensor {k} not sourced from Cosmos3-Nano" + stock = _load_st_tensor(qwen_dir, k, qwen_wm).float() + if got.shape != stock.shape or not torch.equal(got, stock): + n_differ_from_stock += 1 + # Cosmos3's tower is derived from Qwen3-VL, so some tensors coincide; but a + # meaningful fraction must differ, else keeping the stock tower (the bug) + # would be indistinguishable from converting. + assert n_differ_from_stock > 0, ( + f"all {len(visual_keys)} visual tensors equal stock Qwen3-VL — " + "conversion is indistinguishable from keeping the stock tower" + ) + print( + f"\nconvert-reasoner: {len(visual_keys)} visual tensors all sourced from " + f"Cosmos3-Nano; {n_differ_from_stock} differ from stock Qwen3-VL." + ) + + # (4) Language tower overlaid too: layer-0 projection *weights* (a distinct + # model, so never bit-identical) differ from stock Qwen3-VL. Restricted to + # ``*proj.weight`` to avoid biases/norms that can coincide by init. + lm_sample = [ + k for k in sorted(merged_wm) + if ".layers.0." in k and not k.startswith("model.visual.") and k.endswith("proj.weight") + ][:6] + assert lm_sample, "no layer-0 language-model projection weights found in merged checkpoint" + for k in lm_sample: + got = _load_st_tensor(out_dir, k, merged_wm).float() + stock = _load_st_tensor(qwen_dir, k, qwen_wm).float() + assert got.shape != stock.shape or not torch.equal(got, stock), ( + f"LM tensor {k} still equals stock Qwen3-VL (not converted from Cosmos3-Nano)" + ) + if MAX_GPUS == 8: From 7127c4df58997ba5dd00aa79da23073c399cacda Mon Sep 17 00:00:00 2001 From: Liangkai Zhang Date: Thu, 2 Jul 2026 21:24:12 -0700 Subject: [PATCH 10/26] [Cosmos3 OSS] Fix RoboMind Dataset (#81) Fix RoboMind Dataset - RoboMINDFrankaDataset supports loading from both `robomind-franka` and `robomind-franka-dual`, but it only load `norm_stats` for `robomind-franka-dual` - With this fix it will load `norm_stats` for `robomind-franka-dual` or `robomind-franka depends on the embodiment_type. --- .../datasets/robomind_franka_dataset.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py b/cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py index 615d5fce..683a975a 100644 --- a/cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/robomind_franka_dataset.py @@ -119,15 +119,27 @@ def _action_spec(self) -> ActionSpec: @classmethod def _stats_path(cls) -> Path: + # Class-level default (no instance to branch on) — matches the + # constructor's default embodiment_type ("robomind-franka-dual"). return _NORMALIZER_PATH_DUAL + def load_action_stats(self) -> dict[str, torch.Tensor]: + """Instance-aware override: respects ``self._embodiment_type``. + + The inherited classmethod always resolves via ``_stats_path()``, which + has no instance to branch on and is hardcoded to the dual-arm file. A + single-arm instance calling ``load_action_stats()`` would otherwise + silently get 20D dual-arm stats instead of its own 10D stats. + """ + path = _NORMALIZER_PATH_DUAL if self._embodiment_type == "robomind-franka-dual" else _NORMALIZER_PATH_SINGLE + return { + key: torch.from_numpy(value).float() + for key, value in load_action_stats(str(path)).items() + } + def _load_norm_stats(self) -> dict[str, torch.Tensor]: if self._norm_stats is None: - path = _NORMALIZER_PATH_SINGLE if self._embodiment_type == "robomind-franka" else _NORMALIZER_PATH_DUAL - self._norm_stats = { - key: torch.from_numpy(value).float() - for key, value in load_action_stats(str(path)).items() - } + self._norm_stats = self.load_action_stats() return self._norm_stats def __getitem__(self, idx: int) -> dict[str, Any]: From 02159672c0c51ff1f2bfc5a7bcc0baaa9047a5ea Mon Sep 17 00:00:00 2001 From: Trung Pham Date: Fri, 3 Jul 2026 00:54:25 -0400 Subject: [PATCH 11/26] Fix multi-control transfer: honor control weights under torch.compile (#82) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Multi-control transfer inference silently ignored per-control weights. ## Changes - **`omni_mot_model.py`**: propagate `control_weights` into `gen_data_for_packing` - **`attention.py`**: make `multi_control_two_way_attention` `torch.compile`-safe. Its per-segment lengths come from data-dependent unpadding (unbacked symints), which tripped several Dynamo guards under `fullgraph=True`: - `torch._check(k.shape[0] == v.shape[0])` — frontend K/V-length guard - `torch._check(n_q > 0)` / `torch._check(n_kv > 0)` — NATTEN `max_seqlen == 0` / `< 1` varlen guards - `torch._check(n_full == noisy_e)` — makes per-segment `[cs:ce]` slices concrete-length, fixing the in-place write shape guard - pass `cumulative_seqlen_{Q,KV}` + `max_seqlen_{Q,KV}` instead of `seqlens_{Q,KV}` to avoid the disallowed `generate_varlen_parameters` device-host sync inside the compiled region ## Test plan - [x] PAI-Bench-C multi_control eval runs end-to-end under compiled attention (previously crashed with data-dependent Dynamo guard errors). - [x] Diagnostics confirm `control_weights` reach the network and the weighted-sum path executes. - [x] A/B check: extreme weights (e.g. edge=1 vs seg=1) now produce different outputs (previously byte-identical). - [x] Single-control transfer paths unaffected (they use the static-shape `two_way_attention`). Co-authored-by: Maosheng Liao --- .../model/generator/mot/attention.py | 39 +++++++++++++++++-- .../model/generator/omni_mot_model.py | 5 +++ 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/cosmos_framework/model/generator/mot/attention.py b/cosmos_framework/model/generator/mot/attention.py index 8d7bbfee..f0d87660 100644 --- a/cosmos_framework/model/generator/mot/attention.py +++ b/cosmos_framework/model/generator/mot/attention.py @@ -336,6 +336,14 @@ def multi_control_two_way_attention( n_text = int(causal_k_offsets[-1]) n_full = int(full_q_offsets[-1]) + # `n_full` comes from int(full_q_offsets[-1]) → an unbacked symint under + # torch.compile. The control ranges + noisy range partition the full/gen + # segment with noisy last, so `noisy_e` (a concrete int from SplitInfo) is + # exactly the number of valid gen tokens == n_full. Binding them lets Dynamo + # treat the per-segment `full_*_v[cs:ce]` slices below as concrete-length, so + # the in-place writes `full_out_v[cs:ce] = _sdpa(...)` don't raise + # data-dependent `Eq(slice_len, out_len)` guards. + torch._check(n_full == noisy_e) # Unpad to avoid padded rows entering the softmax denominator. causal_k_v = causal_k[:n_text] # [N_text, Hkv, D] @@ -350,15 +358,38 @@ def multi_control_two_way_attention( def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: """Maskless attention using cosmos_framework.model.attention() → [N_q, Hq*D].""" + # K and V are built by concatenating the SAME [text | ctrl_i | noisy] + # slices, so their sequence lengths are always equal. Under + # torch.compile (fullgraph=True) those lengths are unbacked symints + # (from data-dependent unpadding), and the attention frontend's + # `if key_shape[1] != value_shape[1]` guard (attention/checks.py) cannot + # be resolved symbolically. Assert the invariant so Dynamo can discharge + # the guard statically instead of raising a data-dependent error. + torch._check(k.shape[0] == v.shape[0]) n_q, n_kv = q.shape[0], k.shape[0] - seqlens_q = torch.tensor([n_q], dtype=torch.int32, device=q.device) - seqlens_kv = torch.tensor([n_kv], dtype=torch.int32, device=k.device) + # These lengths come from data-dependent unpadding, so they are unbacked + # symints under torch.compile. The selected attention backend (NATTEN) + # validates varlen inputs with `max_seqlen == 0` / `max_seqlen < 1` + # guards; without a positivity fact Dynamo cannot discharge `Eq(n, 0)`. + # Every control/noisy segment always has at least one token, so assert it. + torch._check(n_q > 0) + torch._check(n_kv > 0) + # Pass cumulative_seqlen_{Q,KV} + max_seqlen_{Q,KV} directly instead of + # seqlens_{Q,KV}. The frontend derives cumulative offsets from seqlens via + # `generate_varlen_parameters`, which calls `.max().item()` (a device-host + # sync) and is explicitly disallowed inside a torch.compile region. Each + # pass here is a single (batch=1) packed sequence, so the cumulative + # offsets are simply [0, n]. Building them ourselves keeps the whole path + # inside the compiled graph. + zero = torch.zeros(1, dtype=torch.int32, device=q.device) + cu_seqlens_q = torch.cat([zero, torch.tensor([n_q], dtype=torch.int32, device=q.device)]) + cu_seqlens_kv = torch.cat([zero, torch.tensor([n_kv], dtype=torch.int32, device=q.device)]) res = attention( q.unsqueeze(0), # [1, N_q, Hq, D] k.unsqueeze(0), # [1, N_kv, Hkv, D] v.unsqueeze(0), # [1, N_kv, Hkv, D] - seqlens_Q=seqlens_q, - seqlens_KV=seqlens_kv, + cumulative_seqlen_Q=cu_seqlens_q, + cumulative_seqlen_KV=cu_seqlens_kv, max_seqlen_Q=n_q, max_seqlen_KV=n_kv, ) # [1, N_q, Hq, D] diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index 164ca3de..771427b2 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -1915,6 +1915,11 @@ def _get_velocity( x0_tokens_sound=noise_x_sound if has_sound else None, fps_sound=gen_data_clean.fps_sound if has_sound else None, num_vision_items_per_sample=num_items, + # Multi-control transfer: carry per-control weights so the packer can + # populate vision_item_split_lens / control_weights on the packed + # sequence. Without this, multi_control_two_way_attention never runs + # and all controls are blended equally (weights ignored). + control_weights=gen_data_clean.control_weights, ) packed_sequence = self._pack_input_sequence( From 576bb790f636bea05a9ccd483dc047d304a6049f Mon Sep 17 00:00:00 2001 From: yy-code-nv Date: Fri, 3 Jul 2026 14:53:26 +0800 Subject: [PATCH 12/26] =?UTF-8?q?Release:=20sync=20from=20internal=20?= =?UTF-8?q?=EF=BC=88clean=20run)=20(#83)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Standard cosmos-framework-release pipeline run against current main (generator/reasoner dest layout). Brings the latest i4 changes into CF. Highlights: * New leaf helper shipped: cosmos_framework/utils/generator/torchcodec_video.py (mapped via a new [[overrides]] entry — it lives at the top-level projects/cosmos3/utils/, outside the bulk cosmos3/cosmos3/utils mapping, and is imported by the pkl_to_media / reasoner.bytes_to_media augmentors). * New augmentors: data/generator/augmentors/multi_reference_transform.py, torchcodec_callers_test.py. * 13 existing shipped files refreshed to match i4 head. _source_commit (i4): 546faba4 _dest_commit (cf base): 6efeea33 Co-authored-by: lfengad --- .../callbacks/learning_rate_logger.py | 46 +- cosmos_framework/callbacks/norm_monitor.py | 6 + .../configs/base/reasoner/defaults/config.py | 1 + .../base/reasoner/defaults/optimizer.py | 4 +- .../augmentors/image_editing_transform.py | 40 +- .../augmentors/interleaved_image_transform.py | 75 +++ .../augmentors/multi_reference_transform.py | 447 ++++++++++++++++++ .../data/generator/augmentors/pkl_to_media.py | 19 +- .../augmentors/reasoner/bytes_to_media.py | 26 +- .../augmentors/torchcodec_callers_test.py | 208 ++++++++ .../generator/augmentors/video_parsing.py | 2 +- .../diffusion/samplers/fixed_step.py | 54 ++- .../diffusion/samplers/fixed_step_test.py | 39 ++ .../model/generator/omni_mot_model.py | 60 ++- .../utils/generator/data_utils.py | 5 + .../utils/generator/torchcodec_video.py | 192 ++++++++ 16 files changed, 1154 insertions(+), 70 deletions(-) create mode 100644 cosmos_framework/data/generator/augmentors/multi_reference_transform.py create mode 100644 cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py create mode 100644 cosmos_framework/utils/generator/torchcodec_video.py diff --git a/cosmos_framework/callbacks/learning_rate_logger.py b/cosmos_framework/callbacks/learning_rate_logger.py index 223d827f..657e3181 100644 --- a/cosmos_framework/callbacks/learning_rate_logger.py +++ b/cosmos_framework/callbacks/learning_rate_logger.py @@ -4,21 +4,21 @@ import torch import wandb +from cosmos_framework.utils import distributed from cosmos_framework.utils.callback import Callback class LearningRateLogger(Callback): - """Logs per-model-part learning rate every ``every_n × logging_iter`` steps. + """Log reasoner default and LR-multiplier groups. - Designed for VLM training where the optimizer is an - ``OptimizersContainer`` exposing ``.optimizers`` (list of single-element - optimizer lists) paired with ``.model_part_names``. Silently no-ops when - those attributes are absent so it can be registered alongside plain - ``torch.optim.Optimizer`` setups without harm. + Parameters that do not match a configured ``optimizer.lr_multipliers`` key + are logged as ``optim/lr_default``. Named multiplier groups, such as + ``model.visual``, are logged as ``optim/lr_model.visual`` every + ``every_n × logging_iter`` steps. """ - def __init__(self, every_n: int = 10): - self.every_n = every_n + def __init__(self, every_n: int = 10) -> None: + self.every_n: int = every_n def on_before_optimizer_step( self, @@ -32,16 +32,30 @@ def on_before_optimizer_step( gate = self.config.trainer.logging_iter * self.every_n if not (iteration == 1 or (gate > 0 and iteration % gate == 0)): return - if not wandb.run: + if not distributed.is_rank0() or not wandb.run: return - if not (hasattr(optimizer, "optimizers") and hasattr(optimizer, "model_part_names")): + if not (hasattr(optimizer, "optimizers") and hasattr(optimizer, "model")): return - unique_lr: dict[str, float] = {} - for optim_per_model, name in zip(optimizer.optimizers, optimizer.model_part_names): - if not optim_per_model: - continue - for pg in optim_per_model[0].param_groups: - unique_lr[f"optim/lr_{name}"] = pg["lr"] + + lr_multiplier_keys = list(self.config.optimizer.get("lr_multipliers", {})) + optimizer_net = getattr(optimizer.model, "net", None) + if optimizer_net is None: + return + + lr_key_by_param_id: dict[int, str] = {} + for param_name, param in optimizer_net.named_parameters(): + lr_key_by_param_id[id(param)] = next( + (lr_key for lr_key in lr_multiplier_keys if lr_key in param_name), + "default", + ) + + unique_lr: dict[str, float | torch.Tensor] = {} + for inner_optimizer in optimizer.optimizers: + for param_group in inner_optimizer.param_groups: + for param in param_group["params"]: + lr_key = lr_key_by_param_id.get(id(param)) + if lr_key is not None: + unique_lr[f"optim/lr_{lr_key}"] = param_group["lr"] if not unique_lr: return wandb.log(unique_lr, step=iteration) diff --git a/cosmos_framework/callbacks/norm_monitor.py b/cosmos_framework/callbacks/norm_monitor.py index 36e6ea00..39a2ec27 100644 --- a/cosmos_framework/callbacks/norm_monitor.py +++ b/cosmos_framework/callbacks/norm_monitor.py @@ -229,6 +229,12 @@ def _compute_l2_stats(self, tensor: torch.Tensor, detach: bool = True) -> dict[s if isinstance(data, DTensor): data = data.to_local() + if data.numel() == 0: + return { + "sq_sum": torch.zeros((), device=data.device, dtype=torch.float32), # [] + "max": torch.zeros((), device=data.device, dtype=data.dtype), # [] + } + return { "sq_sum": (data.float() ** 2).sum(), "max": data.abs().max(), diff --git a/cosmos_framework/configs/base/reasoner/defaults/config.py b/cosmos_framework/configs/base/reasoner/defaults/config.py index 6da4659a..f6c3e21c 100644 --- a/cosmos_framework/configs/base/reasoner/defaults/config.py +++ b/cosmos_framework/configs/base/reasoner/defaults/config.py @@ -52,6 +52,7 @@ class DataSetting: num_data_workers: int = 8 data_prefetch_factor: int = 1 val_split_ratio: float = 0.0 + recipe_name: str | None = None @attrs.define(slots=False) diff --git a/cosmos_framework/configs/base/reasoner/defaults/optimizer.py b/cosmos_framework/configs/base/reasoner/defaults/optimizer.py index dbd5829c..e73066f6 100644 --- a/cosmos_framework/configs/base/reasoner/defaults/optimizer.py +++ b/cosmos_framework/configs/base/reasoner/defaults/optimizer.py @@ -20,7 +20,7 @@ betas=(0.9, 0.95), fused=True, keys_to_select=[], - lr_multipliers={"vision_encoder": 0.1}, + lr_multipliers={"model.visual": 0.1}, ) # ``f_start`` / ``f_min`` are ratios of the optimizer's base ``lr``: @@ -32,7 +32,7 @@ cycle_lengths=["${trainer.max_iter}"], f_start=[0.01], f_max=[1.0], - f_min=[0.5], + f_min=[0.1], ) diff --git a/cosmos_framework/data/generator/augmentors/image_editing_transform.py b/cosmos_framework/data/generator/augmentors/image_editing_transform.py index 1571fd44..d0a513c1 100644 --- a/cosmos_framework/data/generator/augmentors/image_editing_transform.py +++ b/cosmos_framework/data/generator/augmentors/image_editing_transform.py @@ -375,13 +375,19 @@ class ImageEditingToTrainingFormat(Augmentor): (e.g. ``OmniInterleavedMediaResize``). This augmentor only normalises the PIL images to tensors and assembles the remaining metadata fields. + Supports both single-source image editing and multi-reference generation: + ``source_image`` may be a single ``PIL.Image`` (one reference) or a + ``list[PIL.Image]`` (N references). The output ``images`` always places the + target last, so the resulting layout is ``[ref_1, ..., ref_N, target]`` with + ``num_frames = N + 1``. + Input (from data_dict): - - source_image: PIL.Image (already resized by upstream augmentor) + - source_image: PIL.Image | list[PIL.Image] (already resized by upstream augmentor) - target_image: PIL.Image (already resized by upstream augmentor) - editing_instruction: str Output (added to data_dict): - - images: list[torch.Tensor] — ``[source (C,H_s,W_s), target (C,H_t,W_t)]`` + - images: list[torch.Tensor] — ``[ref_1, ..., ref_N, target]`` (each C,H,W) - ai_caption: str - selected_caption_type: str - fps: float @@ -415,21 +421,28 @@ def __call__(self, data_dict: dict) -> dict | None: Returns: Updated data_dict with training-compatible fields, or None on error. """ - source_image: Image.Image = data_dict.get("source_image") + source_image = data_dict.get("source_image") target_image: Image.Image = data_dict.get("target_image") editing_instruction: str = data_dict.get("editing_instruction", "") if source_image is None or target_image is None: return None + # Support both single-source editing (one PIL image) and multi-reference + # generation (a list of PIL images). Normalise to a list of references. + source_images = [source_image] if isinstance(source_image, Image.Image) else list(source_image) + if not source_images: + return None + try: - # Normalize PIL images to tensors (upstream augmentor already handled resizing) - source_tensor = self._normalize_image(source_image) # [C,H_s,W_s] - target_tensor = self._normalize_image(target_image) # [C,H_t,W_t] + # Normalize PIL images to tensors (upstream augmentor already handled resizing). + # Each image keeps its own spatial size; the model encodes them separately. + # The target is placed last: [ref_1, ..., ref_N, target]. + images = [self._normalize_image(src) for src in source_images] # each [C,H_s,W_s] + images.append(self._normalize_image(target_image)) # [C,H_t,W_t] # Store as list of tensors for the batch collation. - # Each image keeps its own spatial size; the model encodes them separately. - data_dict["images"] = [source_tensor, target_tensor] + data_dict["images"] = images # Set text fields data_dict["ai_caption"] = editing_instruction @@ -437,16 +450,13 @@ def __call__(self, data_dict: dict) -> dict | None: # Set metadata data_dict["fps"] = 30.0 # Same as standard image training - data_dict["num_frames"] = 2 # Source + target = 2 frames + data_dict["num_frames"] = len(images) # N references + target data_dict["image_size"] = [ torch.tensor( - [source_image.height, source_image.width, source_image.height, source_image.width], - dtype=torch.float, - ), # [4] - torch.tensor( - [target_image.height, target_image.width, target_image.height, target_image.width], + [img.height, img.width, img.height, img.width], dtype=torch.float, - ), # [4] + ) # [4] + for img in (*source_images, target_image) ] # Set the dataset name if not already present if "dataset_name" not in data_dict: diff --git a/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py b/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py index 8faa9721..24f61a07 100644 --- a/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py +++ b/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py @@ -15,6 +15,8 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.data.imaginaire.webdataset.augmentors.image.misc import obtain_image_size +Image.MAX_IMAGE_PIXELS = 933120000 + class ResizeToPaddingDivisor(Augmentor): """Resize images so that both width and height are multiples of padding_divisor.""" @@ -264,3 +266,76 @@ def _resize_image( resized_image = resized_image.resize((final_width, final_height), Image.Resampling.LANCZOS) return resized_image + + +class InterleavedMediaResizeByMaxPixels(Augmentor): + """Resize interleaved media by constraining total pixel area (max_pixels), preserving aspect ratio. + + Unlike :class:`InterleavedMediaResize` (which caps the longest side), this augmentor scales + each image so its total pixel area does not exceed ``max_pixels`` while keeping the aspect + ratio, then aligns both dimensions down to multiples of ``padding_divisor`` (default 16). + Images are only scaled down, never up. + + Mirrors the FLUX2 max-pixel resize behavior. Produces ``diffusion_media_list`` keyed the same + way as the input ``media_list`` dict, compatible with the downstream + ``ExtractMultiReferenceConversation`` / ``ExtractImageEditingConversation`` augmentors. + + Args: + input_keys: List with the single key holding the media dict. Defaults to ``["media_list"]``. + max_pixels: Maximum total pixel area per image after resize. + padding_divisor: Dimension alignment divisor (both width and height become multiples of it). + args: Additional arguments passed to the parent class. + """ + + def __init__( + self, + input_keys: Optional[List] = None, + max_pixels: int = 1048576, + padding_divisor: int = 16, + args: Optional[dict] = None, + ) -> None: + input_keys = input_keys or ["media_list"] + super().__init__(input_keys, None, args) + self.max_pixels = max_pixels + self.padding_divisor = padding_divisor + + def _compute_target_size(self, width: int, height: int) -> tuple[int, int]: + """Scale to fit within max_pixels, then align down to padding_divisor.""" + total_pixels = width * height + if total_pixels > self.max_pixels: + scale = math.sqrt(self.max_pixels / total_pixels) + width = int(width * scale) + height = int(height * scale) + + width = max(self.padding_divisor, (width // self.padding_divisor) * self.padding_divisor) + height = max(self.padding_divisor, (height // self.padding_divisor) * self.padding_divisor) + return width, height + + def _resize_image(self, img: Image.Image) -> Image.Image: + w, h = img.size + target_w, target_h = self._compute_target_size(w, h) + if (target_w, target_h) != (w, h): + img = img.resize((target_w, target_h), Image.Resampling.LANCZOS) + return img + + def _process_media(self, media) -> Image.Image | list: + if isinstance(media, list): + return [self._resize_image(frame) for frame in media] + return self._resize_image(media) + + def __call__(self, data_dict: Dict) -> Optional[Dict]: + assert len(self.input_keys) == 1, ( + "This transform only supports one input key. Try to organize all the media contents under one key." + ) + media_key = self.input_keys[0] + media_list = data_dict.get(media_key) + if media_list is None: + print(f"Input key {media_key} not found in data_dict: {data_dict.get('__key__', 'unknown')}") + return None + + diffusion_media_content = {} + for key, media in media_list.items(): + diffusion_media_content[key] = self._process_media(media) + + data_dict["diffusion_media_list"] = diffusion_media_content + return data_dict diff --git a/cosmos_framework/data/generator/augmentors/multi_reference_transform.py b/cosmos_framework/data/generator/augmentors/multi_reference_transform.py new file mode 100644 index 00000000..02d38e3f --- /dev/null +++ b/cosmos_framework/data/generator/augmentors/multi_reference_transform.py @@ -0,0 +1,447 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Augmentors for multi-reference image generation. + +Multi-reference data layout (per sample): + annotations: dict with keys {instruction, raw_instruction, in_instruction, + input_images, output_image, editing_type, dataset_name, split} + images: dict mapping {input_01, input_02, ..., input_NN, output} -> raw JPEG bytes + +The ``instruction`` field may be either a single string (legacy format) or a +dict of prompt variants ``{original, short, medium, detailed}`` (new format). +For the dict form, one variant is sampled uniformly at random per sample. + +These augmentors transform that on-disk format into the same in-memory layout +expected by ``ImageEditingToTrainingFormat`` (i.e. ``source_image`` is a list of +PIL images, ``target_image`` is a single PIL image, ``editing_instruction`` is a +string), so that the downstream image-editing augmentors can be reused unchanged. +""" + +from __future__ import annotations + +import io +import random +import re +from typing import Optional + +from PIL import Image, UnidentifiedImageError + +from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor +from cosmos_framework.utils import log + +Image.MAX_IMAGE_PIXELS = 933120000 + +_INPUT_KEY_RE = re.compile(r"^input_(\d+)$") +_OUTPUT_KEY = "output" + +# The ``instruction`` annotation field may be either a single string (legacy +# format) or a dict of prompt variants at different lengths (new format). When +# it is a dict, one variant is sampled uniformly at random per sample. These are +# the variant keys we sample over, in their canonical order. +_PROMPT_VARIANT_KEYS = ("original", "short", "medium", "detailed") + + +def _sorted_input_keys(keys: list[str]) -> list[str]: + """Return ``input_NN`` keys sorted by their numeric index.""" + indexed: list[tuple[int, str]] = [] + for key in keys: + match = _INPUT_KEY_RE.match(key) + if match is not None: + indexed.append((int(match.group(1)), key)) + indexed.sort(key=lambda x: x[0]) + return [key for _, key in indexed] + + +class MultiReferencePKLToMedia(Augmentor): + """Decode the multi-reference image bundle into PIL images. + + Reads ``data_dict[input_key]`` (a ``dict[str, bytes]`` whose keys are + ``input_01``..``input_NN`` plus ``output``) and writes the decoded PIL + images into ``data_dict[output_key]`` as ``dict[str, PIL.Image]`` with the + same keys. + + Returns ``None`` (skip sample) if any image fails to decode or if the + ``output`` image is missing — both are required for training. + """ + + def __init__( + self, + input_keys: list | None = None, + input_key: str = "images", + output_key: str = "media_list", + args: Optional[dict] = None, + ) -> None: + super().__init__(input_keys or [], args=args) + self.input_key = input_key + self.output_key = output_key + + def _bytes_to_pil(self, image_bytes: bytes, identifier: str) -> Image.Image | None: + try: + with io.BytesIO(image_bytes) as stream: + img = Image.open(stream) + img.load() + return img.convert("RGB") + except UnidentifiedImageError: + log.warning( + f"Skipping item '{identifier}': cannot identify image bytes.", + rank0_only=False, + ) + except Exception as e: + log.warning( + f"Skipping item '{identifier}': error decoding image bytes: {e}", + rank0_only=False, + ) + return None + + def __call__(self, data_dict: dict) -> dict | None: + if self.input_key not in data_dict: + log.warning( + f"Input key '{self.input_key}' not found in data_dict (keys={list(data_dict.keys())})", + rank0_only=False, + ) + return None + + bundle = data_dict[self.input_key] + if not isinstance(bundle, dict): + log.warning( + f"Expected dict at data_dict['{self.input_key}'], got {type(bundle)}", + rank0_only=False, + ) + return None + + media: dict[str, Image.Image] = {} + for name, item in bundle.items(): + if not isinstance(item, (bytes, bytearray)): + log.warning( + f"Skipping item '{name}': expected bytes, got {type(item)}", + rank0_only=False, + ) + return None + decoded = self._bytes_to_pil(bytes(item), identifier=f"{self.input_key}['{name}']") + if decoded is None: + return None + media[name] = decoded + + if _OUTPUT_KEY not in media: + log.warning( + f"Multi-reference sample missing '{_OUTPUT_KEY}' image: {data_dict.get('__key__', 'unknown')}", + rank0_only=False, + ) + return None + if not _sorted_input_keys(list(media.keys())): + log.warning( + f"Multi-reference sample has no 'input_NN' images: {data_dict.get('__key__', 'unknown')}", + rank0_only=False, + ) + return None + + data_dict[self.output_key] = media + if self.input_key != self.output_key: + del data_dict[self.input_key] + + return data_dict + + +class ExtractMultiReferenceConversation(Augmentor): + """Extract source/target images and instruction from multi-reference annotation. + + Expected ``data_dict`` state on entry: + annotations: dict with at least ``instruction`` and (implicitly via the + image bundle) ``input_NN``/``output`` images. ``in_instruction`` is + present but intentionally ignored — per pipeline design we keep all + input images in the order indicated by their ``input_NN`` key. + diffusion_media_list: dict[str, PIL.Image] keyed by the same + ``input_NN``/``output`` names, populated by + ``InterleavedMediaResize`` upstream. + + The ``instruction`` field supports two formats: + - str: a single instruction (legacy format), used directly. + - dict: prompt variants keyed by ``original``/``short``/``medium``/ + ``detailed`` (new format); one non-empty variant is sampled uniformly + at random per sample. + + Output additions to ``data_dict``: + source_image: list[PIL.Image] in input-index order, truncated to + ``max_reference_images`` if necessary. + target_image: PIL.Image (the ``output`` image). + editing_instruction: str (the sampled ``instruction``, with + ``, , ...`` markers preserved). + dataset_name: ``f"{annotations.dataset_name}/{annotations.split}"`` + (used for logging only). + + Returns ``None`` when required fields are missing. + """ + + def __init__( + self, + input_keys: list | None = None, + max_reference_images: int = 20, + annotation_key: str = "annotations", + media_key: str = "diffusion_media_list", + instruction_key: str = "instruction", + prompt_variant_keys: tuple[str, ...] | list[str] = _PROMPT_VARIANT_KEYS, + args: Optional[dict] = None, + ) -> None: + super().__init__(input_keys or [], args=args) + if max_reference_images <= 0: + raise ValueError(f"max_reference_images must be positive, got {max_reference_images}") + self.max_reference_images = max_reference_images + self.annotation_key = annotation_key + self.media_key = media_key + self.instruction_key = instruction_key + self.prompt_variant_keys = tuple(prompt_variant_keys) + + def _resolve_instruction(self, annotation: dict) -> str | None: + """Resolve the instruction string, sampling a variant when given a dict. + + Returns a non-empty, stripped instruction string, or ``None`` when no + usable instruction is present. + """ + raw = annotation.get(self.instruction_key) + + if isinstance(raw, str): + return raw.strip() or None + + if isinstance(raw, dict): + # Prefer the canonical variant keys (in order), then fall back to any + # other string values so we never silently drop a usable prompt. + candidates = [ + raw[key].strip() + for key in self.prompt_variant_keys + if isinstance(raw.get(key), str) and raw[key].strip() + ] + if not candidates: + candidates = [value.strip() for value in raw.values() if isinstance(value, str) and value.strip()] + if candidates: + return random.choice(candidates) + + return None + + def __call__(self, data_dict: dict) -> dict | None: + for required_key in (self.annotation_key, self.media_key): + if required_key not in data_dict: + log.warning( + f"'{required_key}' not found in data_dict: {data_dict.get('__key__', 'unknown')}", + rank0_only=False, + ) + return None + + annotation = data_dict[self.annotation_key] + if not isinstance(annotation, dict): + log.warning( + f"Expected dict for '{self.annotation_key}', got {type(annotation)}: " + f"{data_dict.get('__key__', 'unknown')}", + rank0_only=False, + ) + return None + + instruction = self._resolve_instruction(annotation) + if not instruction: + log.warning( + f"Missing/empty '{self.instruction_key}' in annotation: {data_dict.get('__key__', 'unknown')}", + rank0_only=False, + ) + return None + + media_dict = data_dict[self.media_key] + if not isinstance(media_dict, dict): + log.warning( + f"Expected dict for '{self.media_key}', got {type(media_dict)}: {data_dict.get('__key__', 'unknown')}", + rank0_only=False, + ) + return None + + target_image = media_dict.get(_OUTPUT_KEY) + if isinstance(target_image, list): + target_image = target_image[0] if target_image else None + if target_image is None: + log.warning( + f"Missing '{_OUTPUT_KEY}' image after resize: {data_dict.get('__key__', 'unknown')}", + rank0_only=False, + ) + return None + + ordered_input_keys = _sorted_input_keys(list(media_dict.keys())) + if not ordered_input_keys: + log.warning( + f"No 'input_NN' images after resize: {data_dict.get('__key__', 'unknown')}", + rank0_only=False, + ) + return None + + if len(ordered_input_keys) > self.max_reference_images: + ordered_input_keys = ordered_input_keys[: self.max_reference_images] + + source_images: list[Image.Image] = [] + for key in ordered_input_keys: + ref = media_dict[key] + if isinstance(ref, list): + ref = ref[0] if ref else None + if ref is None: + log.warning( + f"Reference image '{key}' is None after resize: {data_dict.get('__key__', 'unknown')}", + rank0_only=False, + ) + return None + source_images.append(ref) + + data_dict["source_image"] = source_images + data_dict["target_image"] = target_image + data_dict["editing_instruction"] = instruction + + ds_name = annotation.get("dataset_name") + split = annotation.get("split") + if isinstance(ds_name, str) and ds_name: + data_dict["dataset_name"] = f"{ds_name}/{split}" if isinstance(split, str) and split else ds_name + + return data_dict + + +class RandomResizeReferenceImages(Augmentor): + """Randomly rescale each reference image (aspect ratio preserved); leave the target untouched. + + Operates on the ``media_list`` dict produced by :class:`MultiReferencePKLToMedia`, + i.e. before :class:`InterleavedMediaResize`. With probability ``resize_prob`` (one + Bernoulli draw per sample), every key matching ``input_NN`` gets its own independently + sampled ratio in ``[min_resize_ratio, max_resize_ratio]`` and is resized via LANCZOS. + The ``output`` key (target image) is skipped so the target resolution is unchanged. + + Padding-divisor alignment and the side-length cap are handled by the downstream + :class:`InterleavedMediaResize` stage, so this augmentor does not need to worry + about VAE divisibility. + """ + + def __init__( + self, + input_keys: list | None = None, + min_resize_ratio: float = 0.5, + max_resize_ratio: float = 1.5, + resize_prob: float = 0.0, + media_key: str = "media_list", + output_key_pattern: str = _OUTPUT_KEY, + args: Optional[dict] = None, + ) -> None: + if not 0.0 <= resize_prob <= 1.0: + raise ValueError(f"resize_prob must be in [0, 1], got {resize_prob}") + if min_resize_ratio <= 0: + raise ValueError(f"min_resize_ratio must be > 0, got {min_resize_ratio}") + if max_resize_ratio < min_resize_ratio: + raise ValueError(f"max_resize_ratio ({max_resize_ratio}) must be >= min_resize_ratio ({min_resize_ratio})") + super().__init__(input_keys or [media_key], args=args) + self.min_resize_ratio = float(min_resize_ratio) + self.max_resize_ratio = float(max_resize_ratio) + self.resize_prob = float(resize_prob) + self.media_key = media_key + self.output_key_pattern = output_key_pattern + + def _resize_one(self, img: Image.Image) -> Image.Image: + ratio = random.uniform(self.min_resize_ratio, self.max_resize_ratio) + w, h = img.size + new_w = max(1, round(w * ratio)) + new_h = max(1, round(h * ratio)) + if (new_w, new_h) == (w, h): + return img + return img.resize((new_w, new_h), Image.LANCZOS) + + def __call__(self, data_dict: dict) -> dict | None: + if self.resize_prob <= 0.0 or random.random() >= self.resize_prob: + return data_dict + + media_list = data_dict.get(self.media_key) + if not isinstance(media_list, dict): + return data_dict + + for key, media in media_list.items(): + if key == self.output_key_pattern or _INPUT_KEY_RE.match(key) is None: + continue + if isinstance(media, list): + media_list[key] = [self._resize_one(frame) for frame in media] + elif isinstance(media, Image.Image): + media_list[key] = self._resize_one(media) + + return data_dict + + +_MARKER_RE: re.Pattern[str] = re.compile(r"") + + +class ReorderReferenceImages(Augmentor): + """Shuffle the reference list and rewrite ```` markers in the instruction. + + Operates on the ``source_image`` list and ``editing_instruction`` string produced by + :class:`ExtractMultiReferenceConversation`. With probability ``shuffle_prob`` (one + Bernoulli draw per sample), shuffles ``source_image`` and consistently renumbers the + ```` markers inside ``editing_instruction`` so that the same physical image is + still referenced after the shuffle. + + Shuffling is **only** applied when ``editing_instruction`` contains at least one + ```` marker. Some samples refer to references by natural-language phrases + (``"Image 1"``, ``"first image"``, etc.) that we cannot reliably renumber, so + shuffling those would silently desync the instruction from the image order. Such + samples are returned unchanged. + + The marker rewrite is done in a single ``re.sub`` pass to avoid cascading rewrites + (e.g. ```` -> ```` -> ````). Markers whose index falls outside + ``[1, len(source_image)]`` are left untouched. + """ + + def __init__( + self, + input_keys: list | None = None, + shuffle_prob: float = 0.0, + source_key: str = "source_image", + instruction_key: str = "editing_instruction", + args: Optional[dict] = None, + ) -> None: + if not 0.0 <= shuffle_prob <= 1.0: + raise ValueError(f"shuffle_prob must be in [0, 1], got {shuffle_prob}") + super().__init__(input_keys or [], args=args) + self.shuffle_prob = float(shuffle_prob) + self.source_key = source_key + self.instruction_key = instruction_key + + def __call__(self, data_dict: dict) -> dict | None: + if self.shuffle_prob <= 0.0 or random.random() >= self.shuffle_prob: + return data_dict + + source_image = data_dict.get(self.source_key) + instruction = data_dict.get(self.instruction_key) + if not isinstance(source_image, list) or len(source_image) <= 1: + return data_dict + if not isinstance(instruction, str) or not instruction: + return data_dict + + # Only shuffle when the instruction actually uses ```` markers. + # Otherwise the references are by natural-language phrases (e.g. "Image 1", + # "the first image") that we can't reliably renumber, so reordering would + # silently desync the instruction from the image order. + if _MARKER_RE.search(instruction) is None: + return data_dict + + n = len(source_image) + # Try a few times to get a non-identity permutation; otherwise accept identity. + perm = list(range(n)) + for _ in range(5): + candidate = random.sample(range(n), n) + if candidate != list(range(n)): + perm = candidate + break + + if perm == list(range(n)): + return data_dict + + inv = [0] * n + for new_idx, old_idx in enumerate(perm): + inv[old_idx] = new_idx + + def _remap(match: re.Match) -> str: + old_idx = int(match.group(1)) - 1 + if 0 <= old_idx < n: + return f"" + return match.group(0) + + data_dict[self.source_key] = [source_image[perm[i]] for i in range(n)] + data_dict[self.instruction_key] = _MARKER_RE.sub(_remap, instruction) + + return data_dict diff --git a/cosmos_framework/data/generator/augmentors/pkl_to_media.py b/cosmos_framework/data/generator/augmentors/pkl_to_media.py index c66ea953..4b0fda79 100644 --- a/cosmos_framework/data/generator/augmentors/pkl_to_media.py +++ b/cosmos_framework/data/generator/augmentors/pkl_to_media.py @@ -18,6 +18,7 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log from cosmos_framework.utils.generator.video_preprocess import tensor_to_pil_images +from cosmos_framework.utils.generator.torchcodec_video import decode_frames_tchw_uint8, probe_video Image.MAX_IMAGE_PIXELS = 933120000 _VIDEO_EXTENSIONS = "mp4 avi webm mov".split() @@ -64,7 +65,7 @@ def _video_decoder_qwen_func( target_fps (float, optional): Target FPS. Defaults to 2.0. min_video_token_length (int, optional): Minimum token length. Defaults to 16. max_video_token_length (int, optional): Maximum token length. Defaults to 8192. - num_threads (int, optional): Number of threads for decord. Defaults to 0. + num_threads (int, optional): Number of video decoding threads. Defaults to 0. random_augmentation (bool, optional): Whether to randomize the FPS and max_video_token_length. Defaults to False. fps_random_range (list[float], optional): Random FPS range. Defaults to [10.0, 24.0]. max_video_token_length_random_range (list[float], optional): Random max_video_token_length range. Defaults to [0.75, 1.25]. @@ -80,17 +81,14 @@ def _video_decoder_qwen_func( Returns: dict | None: Dictionary with video frames tensor and target FPS """ - import decord - # Check video extension extension = re.sub(r".*[.]", "", key) if extension.lower() not in _VIDEO_EXTENSIONS: return None # Read video - video_buffer = io.BytesIO(data) - video_reader = decord.VideoReader(video_buffer, num_threads=num_threads) - total_frames, video_fps = len(video_reader), video_reader.get_avg_fps() + metadata = probe_video(data, num_threads=num_threads) + total_frames, video_fps = metadata.num_frames, metadata.average_fps if start_frame is not None and end_frame is not None: total_frames = end_frame - start_frame @@ -138,8 +136,7 @@ def _video_decoder_qwen_func( idx = torch.linspace(start_frame, end_frame - 1, nframes).round().long().tolist() # [nframes] else: idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() # [nframes] - video_frames = video_reader.get_batch(idx).asnumpy() - video_frames = torch.tensor(video_frames).permute(0, 3, 1, 2) # [T,C,H,W] + video_frames, _ = decode_frames_tchw_uint8(data, idx, num_threads=num_threads) # [T,C,H,W] sample_fps = nframes / max(total_frames, 1e-6) * video_fps # recompute max_pixels based on number of sampled frames @@ -156,13 +153,9 @@ def _video_decoder_qwen_func( [resized_height, resized_width], interpolation=InterpolationMode.BICUBIC, antialias=True, - ).float() + ).float() # [T,C,H,W] video_frames = video_frames.permute(1, 0, 2, 3) # [C,T,H,W] - # Clean up - video_reader.seek(0) # set video reader point back to 0 to clean up cache - del video_reader # delete the reader to avoid memory leak - return dict(videos=video_frames, fps=sample_fps) diff --git a/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py b/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py index 03a61f5e..dddbadc2 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py @@ -19,6 +19,7 @@ from cosmos_framework.data.generator.reasoner.video_decoder_qwen import _video_decoder_qwen_func from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor from cosmos_framework.utils.generator.video_preprocess import tensor_to_pil_images +from cosmos_framework.utils.generator.torchcodec_video import probe_video class BytesToMedia(Augmentor): @@ -104,21 +105,16 @@ def _probe_video_duration_seconds( ) -> float | None: """Probe the effective video duration in seconds used for proportional budget allocation.""" try: - import decord - - with io.BytesIO(video_bytes) as stream: - video_reader = decord.VideoReader(stream, num_threads=self.video_decoder_params["num_threads"]) - frame_count = ( - max(end_frame - start_frame, 0) - if start_frame is not None and end_frame is not None - else len(video_reader) - ) - video_fps = video_reader.get_avg_fps() - video_reader.seek(0) - del video_reader - if frame_count <= 0 or video_fps <= 0: - return None - return frame_count / video_fps + metadata = probe_video(video_bytes, num_threads=self.video_decoder_params["num_threads"]) + frame_count = ( + max(end_frame - start_frame, 0) + if start_frame is not None and end_frame is not None + else metadata.num_frames + ) + video_fps = metadata.average_fps + if frame_count <= 0 or video_fps <= 0: + return None + return frame_count / video_fps except Exception as e: log.warning(f"Could not probe video duration for '{identifier}': {e}") return None diff --git a/cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py b/cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py new file mode 100644 index 00000000..6ddebd71 --- /dev/null +++ b/cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Tests for Cosmos3 dataset call sites that use TorchCodec video helpers.""" + +from __future__ import annotations + +from importlib import import_module +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +pytestmark = [pytest.mark.L1, pytest.mark.CPU] + + +def _import_or_skip(module_name: str) -> Any: + try: + return import_module(module_name) + except ModuleNotFoundError as exc: + pytest.skip(f"Optional dependency unavailable while importing {module_name}: {exc.name}") + + +def test_pkl_qwen_decoder_uses_probe_and_decodes_selected_window(monkeypatch: pytest.MonkeyPatch) -> None: + module = _import_or_skip("cosmos_framework.data.generator.augmentors.pkl_to_media") + video_bytes = b"video-bytes" + calls: dict[str, object] = {} + + def fake_probe_video(source: object, num_threads: int = 0, **_: object) -> SimpleNamespace: + calls["probe"] = (source, num_threads) + return SimpleNamespace(num_frames=10, average_fps=12.0, height=8, width=10) + + def fake_decode_frames_tchw_uint8( + source: object, + indices: list[int], + num_threads: int = 0, + **_: object, + ) -> tuple[torch.Tensor, SimpleNamespace]: + calls["decode"] = (source, indices, num_threads) + frames = torch.arange(3 * 3 * 8 * 10, dtype=torch.uint8).reshape(3, 3, 8, 10) # [T,C,H,W] + return frames, fake_probe_video(source, num_threads=num_threads) + + monkeypatch.setattr(module, "probe_video", fake_probe_video) + monkeypatch.setattr(module, "decode_frames_tchw_uint8", fake_decode_frames_tchw_uint8) + monkeypatch.setattr(module, "smart_nframes", lambda *_args, **_kwargs: 3) + monkeypatch.setattr(module, "smart_resize", lambda height, width, **_kwargs: (height, width)) + + result = module._video_decoder_qwen_func( + key="clip.mp4", + data=video_bytes, + min_fps_thres=1, + max_fps_thres=60, + target_fps=3.0, + min_video_token_length=1, + max_video_token_length=1024, + num_threads=5, + start_frame=2, + end_frame=8, + ) + + assert calls["probe"] == (video_bytes, 5) + assert calls["decode"] == (video_bytes, [2, 4, 7], 5) + assert result is not None + assert tuple(result["videos"].shape) == (3, 3, 8, 10) + assert result["videos"].dtype == torch.float32 + assert result["fps"] == pytest.approx(6.0) + + +@pytest.mark.parametrize( + "module_name", + [ + "cosmos_framework.data.generator.augmentors.reasoner.bytes_to_media", + "cosmos_framework.utils.datasets.augmentors.bytes_to_media", + ], +) +def test_bytes_to_media_uses_probe_durations_for_token_budget( + monkeypatch: pytest.MonkeyPatch, + module_name: str, +) -> None: + module = _import_or_skip(module_name) + seen_sources: list[tuple[object, int]] = [] + + def fake_probe_video(source: object, num_threads: int = 0, **_: object) -> SimpleNamespace: + seen_sources.append((source, num_threads)) + frame_counts = {b"first": 80, b"second": 40} + return SimpleNamespace(num_frames=frame_counts[source], average_fps=20.0, height=16, width=16) + + monkeypatch.setattr(module, "probe_video", fake_probe_video) + augmentor = module.BytesToMedia( + min_video_token_length=8, + max_video_token_length=80, + num_threads=3, + is_input_pickle_byptes=False, + ) + + durations = augmentor._get_video_durations( + { + "video_first.mp4": b"first", + "control_input_depth": b"second", + "image.jpg": b"image", + }, + {}, + ) + + assert durations == {"video_first.mp4": 4.0, "control_input_depth": 2.0} + assert seen_sources == [(b"first", 3), (b"second", 3)] + weighted_params = augmentor._get_decoder_params( + video_count=2, + video_duration=durations["video_first.mp4"], + total_video_duration=sum(durations.values()), + ) + assert weighted_params["max_video_token_length"] == 53 + + +@pytest.mark.parametrize( + "module_name", + [ + "cosmos_framework.data.generator.augmentors.reasoner.bytes_to_media", + "cosmos_framework.utils.datasets.augmentors.bytes_to_media", + ], +) +def test_bytes_to_media_probe_duration_respects_start_end_frame( + monkeypatch: pytest.MonkeyPatch, + module_name: str, +) -> None: + module = _import_or_skip(module_name) + monkeypatch.setattr( + module, + "probe_video", + lambda *_args, **_kwargs: SimpleNamespace(num_frames=100, average_fps=25.0, height=16, width=16), + ) + augmentor = module.BytesToMedia(num_threads=2, is_input_pickle_byptes=False) + + duration = augmentor._probe_video_duration_seconds( + b"video", + identifier="media['video.mp4']", + start_frame=7, + end_frame=32, + ) + empty_duration = augmentor._probe_video_duration_seconds( + b"video", + identifier="media['video.mp4']", + start_frame=7, + end_frame=7, + ) + + assert duration == pytest.approx(1.0) + assert empty_duration is None + + +def test_multiview_extract_frames_resizes_torchcodec_frames(monkeypatch: pytest.MonkeyPatch) -> None: + module = _import_or_skip("cosmos_framework.data.generator.multiview.multiview_dataset") + calls: list[tuple[object, list[int]]] = [] + + def fake_decode_frames_tchw_uint8( + source: object, + indices: list[int], + **_: object, + ) -> tuple[torch.Tensor, SimpleNamespace]: + calls.append((source, indices)) + frames = torch.arange(len(indices) * 3 * 4 * 6, dtype=torch.uint8).reshape(len(indices), 3, 4, 6) # [T,C,H,W] + return frames, SimpleNamespace(average_fps=29.97) + + monkeypatch.setattr(module, "decode_frames_tchw_uint8", fake_decode_frames_tchw_uint8) + + frames, fps, original_hw = module.ExtractFramesAndCaptions._extract_frames( # [T,C,H,W] + b"video", + [1, 3, 5], + (8, 12), + ) + + assert calls == [(b"video", [1, 3, 5])] + assert tuple(frames.shape) == (3, 3, 8, 12) + assert frames.dtype == torch.uint8 + assert fps == pytest.approx(29.97) + assert original_hw == (4, 6) + + +def test_sekai_frame_count_uses_probe_video(monkeypatch: pytest.MonkeyPatch) -> None: + module = _import_or_skip("projects.cosmos3.sil.omnidreams.datasets.sekai") + seen: list[bytes] = [] + + def fake_probe_video(source: bytes) -> SimpleNamespace: + seen.append(source) + return SimpleNamespace(num_frames=37, average_fps=30.0, height=16, width=16) + + monkeypatch.setattr(module, "probe_video", fake_probe_video) + + assert module._num_available_video_frames(b"sekai-video") == 37 + assert seen == [b"sekai-video"] + + +def test_sekai_fit_window_clamps_to_available_span() -> None: + module = _import_or_skip("projects.cosmos3.sil.omnidreams.datasets.sekai") + + assert module._fit_window_to_available_video( + frame_start=20, + num_video_frames=8, + stride=2, + num_available_frames=30, + ) == (20, 5) + assert module._fit_window_to_available_video( + frame_start=100, + num_video_frames=8, + stride=3, + num_available_frames=10, + ) == (0, 4) diff --git a/cosmos_framework/data/generator/augmentors/video_parsing.py b/cosmos_framework/data/generator/augmentors/video_parsing.py index a443948e..fb5808a6 100644 --- a/cosmos_framework/data/generator/augmentors/video_parsing.py +++ b/cosmos_framework/data/generator/augmentors/video_parsing.py @@ -35,7 +35,7 @@ class VideoParsing(Augmentor): This augmentor is used to parse the video bytes and get the video frames. the return dict is back-compatible with old datasets, which video decoding happens in the decoder stage. - Now uses torchcodec instead of decord for video decoding, with optional audio extraction. + Uses TorchCodec for video decoding, with optional audio extraction. """ def __init__(self, input_keys: list, output_keys: Optional[list] = None, args: Optional[dict] = None) -> None: diff --git a/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py b/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py index 706b4d4e..714268e2 100644 --- a/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py +++ b/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py @@ -53,6 +53,8 @@ def __call__( num_steps: int | None = None, shift: float | None = None, seed: int | list[int] | None = None, + condition_reference: torch.Tensor | list[torch.Tensor] | None = None, + condition_mask: torch.Tensor | list[torch.Tensor] | None = None, ) -> torch.Tensor | list[torch.Tensor]: """Run the fixed-step sampling loop. @@ -77,15 +79,37 @@ def __call__( ``len(t_list) - 1`` when provided). shift: When set, derive the sigma schedule dynamically using the flow-matching shift formula instead of ``self.t_list``. + condition_reference: Optional clean reference tensor(s) to preserve + where ``condition_mask`` is 1. + condition_mask: Optional mask tensor(s), same shape as ``noise``, + where 1 marks clean conditioning values and 0 marks generated + values. Returns: Denoised sample(s). A single ``torch.Tensor`` when ``noise`` is a tensor, or a ``list[torch.Tensor]`` when ``noise`` is a list. """ + assert (condition_reference is None) == (condition_mask is None), ( + "condition_reference and condition_mask must be both set or both None" + ) if isinstance(noise, list): device = noise[0].device + if seed is None: + seed = [None] * len(noise) + assert isinstance(seed, list), "seed must be a list when noise is a list" + if condition_reference is None: + condition_reference = [None] * len(noise) + condition_mask = [None] * len(noise) + else: + assert isinstance(condition_reference, list), "condition_reference must be a list when noise is a list" + assert isinstance(condition_mask, list), "condition_mask must be a list when noise is a list" else: device = noise.device + assert not isinstance(seed, list), "seed must not be a list when noise is a tensor" + assert not isinstance(condition_reference, list), ( + "condition_reference must not be a list when noise is a tensor" + ) + assert not isinstance(condition_mask, list), "condition_mask must not be a list when noise is a tensor" if shift is not None: assert num_steps is not None, "num_steps is required when shift is provided" @@ -105,27 +129,43 @@ def __call__( timestep = torch.tensor(sigma_cur * self.num_train_timesteps, device=device) v_pred = velocity_fn(latent, timestep.reshape(1, 1)) - def _sde_step(seed: int | None, latent: torch.Tensor, v_pred: torch.Tensor) -> torch.Tensor: - x0_pred = latent - sigma_cur * v_pred + def _sde_step( + seed: int | None, + latent: torch.Tensor, + v_pred: torch.Tensor, + condition_reference: torch.Tensor | None, + condition_mask: torch.Tensor | None, + ) -> torch.Tensor: + x0_pred = latent - sigma_cur * v_pred # [...,D] if sigma_next > 0: if self.sample_type == "ode": # Euler ODE step - latent = latent + (sigma_next - sigma_cur) * v_pred + latent_next = latent + (sigma_next - sigma_cur) * v_pred # [...,D] else: if seed is not None: torch.manual_seed(seed + step_idx) - eps_fresh = torch.randn_like(x0_pred) - latent = (1.0 - sigma_next) * x0_pred + sigma_next * eps_fresh + eps_fresh = torch.randn_like(x0_pred) # [...,D] + latent_next = (1.0 - sigma_next) * x0_pred + sigma_next * eps_fresh # [...,D] else: - latent = x0_pred - return latent + latent_next = x0_pred # [...,D] + + if condition_reference is not None: + assert condition_mask is not None, "condition_mask is required when condition_reference is set" + condition_reference = condition_reference.to( + dtype=latent_next.dtype, device=latent_next.device + ) # [...,D] + condition_mask = condition_mask.to(dtype=latent_next.dtype, device=latent_next.device) # [...,D] + latent_next = condition_mask * condition_reference + (1.0 - condition_mask) * latent_next # [...,D] + return latent_next latent = run_multiseed( _sde_step, seed=seed, latent=latent, v_pred=v_pred, + condition_reference=condition_reference, + condition_mask=condition_mask, ) return latent diff --git a/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py b/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py index dd214018..a6aff60e 100644 --- a/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py +++ b/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py @@ -90,6 +90,45 @@ def velocity_fn(x, _t): assert torch.allclose(result_a, result_b) +@pytest.mark.L0 +def test_sde_preserves_conditioned_entries(): + sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.0], sample_type="sde") + noise = [torch.tensor([5.0, 1.0])] # [N] + condition_reference = [torch.tensor([5.0, 0.0])] # [N] + condition_mask = [torch.tensor([1.0, 0.0])] # [N] + + def velocity_fn(x, _t): + return [torch.zeros_like(x_i) for x_i in x] + + result = sampler( + velocity_fn, + noise, + seed=[123], + condition_reference=condition_reference, + condition_mask=condition_mask, + ) + assert torch.allclose(result[0][0], condition_reference[0][0]) + + +@pytest.mark.L0 +def test_ode_preserves_conditioned_entries(): + sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.0], sample_type="ode") + noise = torch.tensor([5.0, 1.0]) # [N] + condition_reference = torch.tensor([5.0, 0.0]) # [N] + condition_mask = torch.tensor([1.0, 0.0]) # [N] + + def velocity_fn(x, _t): + return torch.ones_like(x) # [N] + + result = sampler( + velocity_fn, + noise, + condition_reference=condition_reference, + condition_mask=condition_mask, + ) + assert torch.allclose(result[0], condition_reference[0]) + + @pytest.mark.L0 def test_output_shape_preserved(): sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.25, 0.0]) diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index 771427b2..d3a544c4 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -1625,6 +1625,8 @@ def _prepare_inference_data( list[list[int]], list[list[int]], list[torch.Tensor], + list[torch.Tensor], + list[torch.Tensor], ]: """ Prepare all data needed for inference sampling. @@ -1651,6 +1653,10 @@ def _prepare_inference_data( - uncond_text_tokens: Unconditional text tokens (for CFG) - initial_noise: List of noise tensors (one per sample), each containing flattened vision (and optionally action) noise concatenated + - condition_reference: List of clean reference tensors flattened + in the same order as initial_noise + - condition_mask: List of masks flattened in the same order as + initial_noise, where 1 keeps condition_reference values fixed """ # 1. Build sequence plans (same as training) sequence_plans = build_sequence_plans_from_data_batch( @@ -1786,28 +1792,67 @@ def _prepare_inference_data( # noise_action_list and noise_sound_list are dense (only modality-having samples), # so we use separate indexes. initial_noise: list[torch.Tensor] = [] + condition_reference: list[torch.Tensor] = [] + condition_mask: list[torch.Tensor] = [] idx_vision = 0 idx_action = 0 idx_sound = 0 for i in range(n_sample): parts = [] + condition_reference_parts = [] + condition_mask_parts = [] # Flatten and concatenate all vision items for this sample num_vis = num_items_per_sample[i] if num_items_per_sample is not None else 1 for _ in range(num_vis): parts.append(noise_vision_list[idx_vision].reshape(-1)) + x0_vision = gen_data_clean.x0_tokens_vision[idx_vision] # [C,T,H,W] + mask_vision = packed_sequence.vision.condition_mask[idx_vision].to( # [T,1,1] + dtype=x0_vision.dtype, device=x0_vision.device + ) + condition_reference_parts.append(x0_vision.reshape(-1)) # [N_vision] + condition_mask_parts.append((mask_vision * torch.ones_like(x0_vision)).reshape(-1)) # [N_vision] idx_vision += 1 if noise_action_list is not None and sequence_plans[i].has_action: + assert packed_sequence.action is not None + assert packed_sequence.action.condition_mask is not None + assert gen_data_clean.x0_tokens_action is not None parts.append(noise_action_list[idx_action].reshape(-1)) + x0_action = gen_data_clean.x0_tokens_action[idx_action].to( # [T,D] + dtype=noise_action_list[idx_action].dtype, + device=noise_action_list[idx_action].device, + ) + if gen_data_clean.raw_action_dim is not None and gen_data_clean.raw_action_dim[idx_action] is not None: + x0_action = x0_action.clone() # [T,D] + x0_action[:, gen_data_clean.raw_action_dim[idx_action] :] = 0 # [T,D] + mask_action = packed_sequence.action.condition_mask[idx_action].to( # [T,1] + dtype=x0_action.dtype, device=x0_action.device + ) + condition_reference_parts.append(x0_action.reshape(-1)) # [N_action] + condition_mask_parts.append((mask_action * torch.ones_like(x0_action)).reshape(-1)) # [N_action] idx_action += 1 if noise_sound_list is not None and sequence_plans[i].has_sound: + assert packed_sequence.sound is not None + assert packed_sequence.sound.condition_mask is not None + assert gen_data_clean.x0_tokens_sound is not None parts.append(noise_sound_list[idx_sound].reshape(-1)) + x0_sound = gen_data_clean.x0_tokens_sound[idx_sound].to( # [C_sound,T_sound] + dtype=noise_sound_list[idx_sound].dtype, + device=noise_sound_list[idx_sound].device, + ) + mask_sound = packed_sequence.sound.condition_mask[idx_sound].T.to( # [1,T_sound] + dtype=x0_sound.dtype, device=x0_sound.device + ) + condition_reference_parts.append(x0_sound.reshape(-1)) # [N_sound] + condition_mask_parts.append((mask_sound * torch.ones_like(x0_sound)).reshape(-1)) # [N_sound] idx_sound += 1 initial_noise.append(torch.cat(parts, dim=0)) # [N_tokens_flat] + condition_reference.append(torch.cat(condition_reference_parts, dim=0)) # [N_tokens_flat] + condition_mask.append(torch.cat(condition_mask_parts, dim=0)) # [N_tokens_flat] return ( sequence_plans, @@ -1815,6 +1860,8 @@ def _prepare_inference_data( cond_text_tokens, uncond_text_tokens, initial_noise, + condition_reference, + condition_mask, ) def _get_velocity( @@ -2304,6 +2351,8 @@ def generate_samples_from_batch( cond_tokens, uncond_tokens, initial_noise, + condition_reference, + condition_mask, ) = self._prepare_inference_data(data_batch, seed, has_negative_prompt) if n_sample is not None: @@ -2477,6 +2526,13 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool): else: log.info(f"Using sampler: EDM (sigma_max={sigma_max}, num_steps={num_steps})") + fixed_step_sampler_kwargs = {} + if isinstance(sampler, FixedStepSampler): + fixed_step_sampler_kwargs = { + "condition_reference": condition_reference, + "condition_mask": condition_mask, + } + if isinstance(sampler, FixedStepSampler) or scheduler_type == "unipc": latents = sampler( velocity_fn, @@ -2484,6 +2540,7 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool): num_steps=num_steps, shift=shift, seed=seed, + **fixed_step_sampler_kwargs, ) if _extra_num_steps > 0: # Dummy sampler call to issue (_extra_num_steps × per-step) @@ -2501,6 +2558,7 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool): num_steps=_extra_num_steps, shift=shift, seed=seed, + **fixed_step_sampler_kwargs, ) else: # EDM Sampler @@ -2534,7 +2592,7 @@ def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: # run. Avoids two EDM-specific footguns: # (1) ``EDMSampler._forward_impl`` always runs an extra # ``sample_clean`` denoiser forward (see - # ``cosmos_framework/model/generator/diffusion/samplers/edm.py``). + # ``cosmos_framework/model/vfm/diffusion/samplers/edm.py``). # A nested sampler call would add one too many # forwards on fast ranks, since the slow rank's # single call also pays the ``sample_clean`` cost. diff --git a/cosmos_framework/utils/generator/data_utils.py b/cosmos_framework/utils/generator/data_utils.py index 4ffb1e24..91d83193 100644 --- a/cosmos_framework/utils/generator/data_utils.py +++ b/cosmos_framework/utils/generator/data_utils.py @@ -45,6 +45,11 @@ def get_vision_data_resolution(spatial_shape: tuple[int, int]) -> str: return "480" elif min_dim <= 960: return "720" + elif min_dim <= 2048: + # Free-form inputs above the 720 tier (e.g. multi-reference generation + # producing shapes like (992, 1024)) that are not a canonical 768 shape: + # route to the closest defined higher tier "768". + return "768" else: raise ValueError(f"Unsupported resolution: {spatial_shape}") diff --git a/cosmos_framework/utils/generator/torchcodec_video.py b/cosmos_framework/utils/generator/torchcodec_video.py new file mode 100644 index 00000000..6a27ac0d --- /dev/null +++ b/cosmos_framework/utils/generator/torchcodec_video.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""TorchCodec helpers for Cosmos3 video decoding.""" + +from __future__ import annotations + +import io +from dataclasses import dataclass +from pathlib import Path +from typing import Any, BinaryIO + +import numpy as np +import torch + +VideoSource = str | Path | bytes | io.BytesIO | BinaryIO + + +@dataclass(frozen=True) +class VideoMetadata: + num_frames: int + average_fps: float + height: int | None = None + width: int | None = None + + +def _normalize_source(source: VideoSource) -> VideoSource: + if isinstance(source, Path): + return source.as_posix() + if isinstance(source, io.BytesIO): + source.seek(0) + return source + + +def _get_video_decoder_cls() -> Any: + try: + from torchcodec.decoders import VideoDecoder + except ImportError as e: + raise ImportError("TorchCodec is required for Cosmos3 video decoding. Install the torchcodec package.") from e + return VideoDecoder + + +def _build_decoder( + source: VideoSource, + *, + num_threads: int = 1, + seek_mode: str = "exact", + device: str = "cpu", +) -> Any: + normalized_source = _normalize_source(source) + # Preserve FFmpeg/TorchCodec's 0 sentinel so callers can request automatic thread selection. + num_ffmpeg_threads = 0 if num_threads == 0 else max(num_threads, 1) + kwargs: dict[str, Any] = {"seek_mode": seek_mode, "num_ffmpeg_threads": num_ffmpeg_threads} + if device != "cpu": + kwargs["device"] = device + video_decoder_cls = _get_video_decoder_cls() + return video_decoder_cls(normalized_source, **kwargs) + + +def _read_basic_metadata(decoder: Any) -> tuple[int, float]: + metadata = decoder.metadata + num_frames = metadata.num_frames + average_fps = metadata.average_fps + if num_frames is None or average_fps is None: + raise ValueError(f"TorchCodec missing metadata (num_frames={num_frames}, average_fps={average_fps})") + return int(num_frames), float(average_fps) + + +def _metadata_from_frame( + decoder: Any, + first_frame_tchw: torch.Tensor | None = None, + *, + include_dimensions: bool = True, +) -> VideoMetadata: + num_frames, average_fps = _read_basic_metadata(decoder) + if not include_dimensions: + return VideoMetadata(num_frames=num_frames, average_fps=average_fps) + if first_frame_tchw is None: + first_frame_tchw = decoder.get_frames_at([0]).data.cpu() # [1,C,H,W] + _, _, height, width = first_frame_tchw.shape # [T,C,H,W] + return VideoMetadata(num_frames=num_frames, average_fps=average_fps, height=int(height), width=int(width)) + + +def probe_video( + source: VideoSource, + *, + num_threads: int = 1, + seek_mode: str = "exact", + device: str = "cpu", + include_dimensions: bool = False, +) -> VideoMetadata: + """Read video metadata, optionally decoding frame 0 to get frame dimensions.""" + decoder = _build_decoder(source, num_threads=num_threads, seek_mode=seek_mode, device=device) + return _metadata_from_frame(decoder, include_dimensions=include_dimensions) + + +class TorchCodecVideoReader: + """Reusable indexed video reader backed by one TorchCodec decoder.""" + + metadata: VideoMetadata + + def __init__( + self, + source: VideoSource, + *, + num_threads: int = 1, + seek_mode: str = "exact", + device: str = "cpu", + include_dimensions: bool = False, + ) -> None: + self._decoder = _build_decoder(source, num_threads=num_threads, seek_mode=seek_mode, device=device) + self.metadata = _metadata_from_frame(self._decoder, include_dimensions=include_dimensions) + + def __len__(self) -> int: + return self.metadata.num_frames + + def __getitem__(self, index: int) -> np.ndarray: + return self.get_frame_nhwc_uint8(index) # [H,W,C] + + def get_avg_fps(self) -> float: + return self.metadata.average_fps + + def get_frames_tchw_uint8(self, indices: list[int]) -> torch.Tensor: + frames_tchw = self._decoder.get_frames_at(indices).data.cpu() # [T,C,H,W] + return frames_tchw # [T,C,H,W] + + def get_frames_nhwc_uint8(self, indices: list[int]) -> np.ndarray: + frames_tchw = self.get_frames_tchw_uint8(indices) # [T,C,H,W] + frames_nhwc = frames_tchw.permute(0, 2, 3, 1).contiguous().numpy() # [T,H,W,C] + return frames_nhwc # [T,H,W,C] + + def get_frame_nhwc_uint8(self, index: int) -> np.ndarray: + frames_nhwc = self.get_frames_nhwc_uint8([index]) # [1,H,W,C] + return frames_nhwc[0] # [H,W,C] + + +def decode_frames_tchw_uint8( + source: VideoSource, + indices: list[int], + *, + num_threads: int = 1, + seek_mode: str = "exact", + device: str = "cpu", +) -> tuple[torch.Tensor, VideoMetadata]: + decoder = _build_decoder(source, num_threads=num_threads, seek_mode=seek_mode, device=device) + frames_tchw = decoder.get_frames_at(indices).data.cpu() # [T,C,H,W] + metadata = _metadata_from_frame(decoder, frames_tchw[:1]) # frames_tchw[:1]: [1,C,H,W] + return frames_tchw, metadata + + +def decode_frames_cthw_uint8( + source: VideoSource, + indices: list[int], + *, + num_threads: int = 1, + seek_mode: str = "exact", + device: str = "cpu", +) -> tuple[torch.Tensor, VideoMetadata]: + frames_tchw, metadata = decode_frames_tchw_uint8( + source, indices, num_threads=num_threads, seek_mode=seek_mode, device=device + ) + frames_cthw = frames_tchw.permute(1, 0, 2, 3).contiguous() # [C,T,H,W] + return frames_cthw, metadata + + +def decode_frames_nhwc_uint8( + source: VideoSource, + indices: list[int], + *, + num_threads: int = 1, + seek_mode: str = "exact", + device: str = "cpu", +) -> tuple[np.ndarray, VideoMetadata]: + frames_tchw, metadata = decode_frames_tchw_uint8( + source, indices, num_threads=num_threads, seek_mode=seek_mode, device=device + ) + frames_nhwc = frames_tchw.permute(0, 2, 3, 1).contiguous().numpy() # [T,H,W,C] + return frames_nhwc, metadata + + +def decode_frame_nhwc_uint8( + source: VideoSource, + index: int, + *, + num_threads: int = 1, + seek_mode: str = "exact", + device: str = "cpu", +) -> tuple[np.ndarray, VideoMetadata]: + frames_nhwc, metadata = decode_frames_nhwc_uint8( + source, [index], num_threads=num_threads, seek_mode=seek_mode, device=device + ) + return frames_nhwc[0], metadata # frames_nhwc[0]: [H,W,C] From dff6966446f2325f69e9a65c396908019b721df1 Mon Sep 17 00:00:00 2001 From: pengcuo Date: Fri, 3 Jul 2026 16:13:58 +0800 Subject: [PATCH 13/26] deps: use natten gb300 (sm_103) wheel for x86_64 cu130 (#84) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem On **x86_64**, the `cu130` dependency group pinned the stock natten wheel `natten==0.21.6.dev6+cu130.torch210`, which does **not** contain `sm_103a` kernels. On B300 / GB300 (compute capability 10.3) this fails at runtime with `no kernel image is available for execution on the device`. ## Change - Point x86_64 cu130 natten at `==0.21.6.dev6+cu130.torch210.gb300`, mirroring the existing aarch64 entry. That build appends `10.3 -> sm_103a` to `NATTEN_CUDA_ARCH`. - Regenerated `uv.lock` (now resolves the x86_64 gb300 wheel, sha256 `04ace5d5…`). The wheel is published in the cosmos index via nvidia-cosmos/cosmos-dependencies#60. ## Verification - `uv lock --check` passes (lock consistent with pyproject). - End-to-end: `Cosmos3-Nano` t2i inference on a B300 (sm_103) completes successfully — 960×960 image generated, no "no kernel image" error. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: lfengad --- pyproject.toml | 4 ++-- uv.lock | 42 +++++++++++++++--------------------------- 2 files changed, 17 insertions(+), 29 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f3fffd9c..2130b4f9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -203,9 +203,9 @@ cu128-train = [ cu130 = [ "flash-attn-3-nv==1.0.3+cu130.torch210; platform_machine == 'x86_64'", "flash-attn==2.7.4.post1+cu130.torch210; platform_machine == 'x86_64'", - # aarch64 (GB300/sm_103a): use the custom build that includes sm_103a kernels. + # GB300/B300 (sm_103a): use the custom build that includes sm_103a kernels (both arches). "natten==0.21.6.dev6+cu130.torch210.gb300; platform_machine == 'aarch64'", - "natten==0.21.6.dev6+cu130.torch210; platform_machine == 'x86_64'", + "natten==0.21.6.dev6+cu130.torch210.gb300; platform_machine == 'x86_64'", "torch==2.10.0+cu130", "torchcodec==0.10.0+cu130", # https://github.com/meta-pytorch/torchcodec/releases "torchvision==0.25.0+cu130", # https://github.com/pytorch/vision/releases diff --git a/uv.lock b/uv.lock index aaf0e7cb..81fa99bd 100644 --- a/uv.lock +++ b/uv.lock @@ -1922,8 +1922,7 @@ cu128-train = [ cu130 = [ { name = "flash-attn", version = "2.7.4.post1+cu130.torch210", source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, marker = "platform_machine == 'x86_64'" }, { name = "flash-attn-3-nv", version = "1.0.3+cu130.torch210", source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, marker = "platform_machine == 'x86_64'" }, - { name = "natten", version = "0.21.6.dev6+cu130.torch210", source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, marker = "(platform_machine == 'x86_64' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-vllm')" }, - { name = "natten", version = "0.21.6.dev6+cu130.torch210.gb300", source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, marker = "(platform_machine == 'aarch64' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-vllm')" }, + { name = "natten", version = "0.21.6.dev6+cu130.torch210.gb300", source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, { name = "nvidia-cublas" }, { name = "nvidia-cuda-cupti" }, { name = "nvidia-cuda-nvrtc" }, @@ -1947,8 +1946,7 @@ cu130 = [ cu130-train = [ { name = "flash-attn", version = "2.7.4.post1+cu130.torch210", source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, marker = "platform_machine == 'x86_64'" }, { name = "flash-attn-3-nv", version = "1.0.3+cu130.torch210", source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, marker = "platform_machine == 'x86_64'" }, - { name = "natten", version = "0.21.6.dev6+cu130.torch210", source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, marker = "(platform_machine == 'x86_64' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-vllm')" }, - { name = "natten", version = "0.21.6.dev6+cu130.torch210.gb300", source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, marker = "(platform_machine == 'aarch64' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-vllm')" }, + { name = "natten", version = "0.21.6.dev6+cu130.torch210.gb300", source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, { name = "nvidia-cublas" }, { name = "nvidia-cuda-cupti" }, { name = "nvidia-cuda-nvrtc" }, @@ -2178,7 +2176,7 @@ cu130 = [ { name = "flash-attn", marker = "platform_machine == 'x86_64'", specifier = "==2.7.4.post1+cu130.torch210", index = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, { name = "flash-attn-3-nv", marker = "platform_machine == 'x86_64'", specifier = "==1.0.3+cu130.torch210", index = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, { name = "natten", marker = "platform_machine == 'aarch64'", specifier = "==0.21.6.dev6+cu130.torch210.gb300", index = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, - { name = "natten", marker = "platform_machine == 'x86_64'", specifier = "==0.21.6.dev6+cu130.torch210", index = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, + { name = "natten", marker = "platform_machine == 'x86_64'", specifier = "==0.21.6.dev6+cu130.torch210.gb300", index = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, { name = "nvidia-cublas", specifier = "==13.1.0.3" }, { name = "nvidia-cuda-cupti", specifier = "==13.0.85" }, { name = "nvidia-cuda-nvrtc", specifier = "==13.0.88" }, @@ -2203,7 +2201,7 @@ cu130-train = [ { name = "flash-attn", marker = "platform_machine == 'x86_64'", specifier = "==2.7.4.post1+cu130.torch210", index = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, { name = "flash-attn-3-nv", marker = "platform_machine == 'x86_64'", specifier = "==1.0.3+cu130.torch210", index = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, { name = "natten", marker = "platform_machine == 'aarch64'", specifier = "==0.21.6.dev6+cu130.torch210.gb300", index = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, - { name = "natten", marker = "platform_machine == 'x86_64'", specifier = "==0.21.6.dev6+cu130.torch210", index = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, + { name = "natten", marker = "platform_machine == 'x86_64'", specifier = "==0.21.6.dev6+cu130.torch210.gb300", index = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" }, { name = "nvidia-cublas", specifier = "==13.1.0.3" }, { name = "nvidia-cuda-cupti", specifier = "==13.0.85" }, { name = "nvidia-cuda-nvrtc", specifier = "==13.0.88" }, @@ -7118,45 +7116,35 @@ wheels = [ { url = "https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/v1.5.0/natten-0.21.6.dev6%2Bcu128.torch210-cp313-cp313-linux_x86_64.whl", hash = "sha256:164b4cd16d36f747ddd19c6f0cb1cd0a275e758e01f2402d4e957be5d7818a0d" }, ] -[[package]] -name = "natten" -version = "0.21.6.dev6+cu130.torch210" -source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" } -resolution-markers = [ - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'linux'", - "python_full_version == '3.13.*' and platform_machine == 'x86_64' and sys_platform == 'linux'", - "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform != 'linux'", - "python_full_version == '3.13.*' and platform_machine == 'x86_64' and sys_platform != 'linux'", - "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'linux'", - "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform != 'linux'", - "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform != 'linux'", - "python_full_version < '3.11' and platform_machine == 'x86_64' and sys_platform == 'linux'", - "python_full_version < '3.11' and platform_machine == 'x86_64' and sys_platform != 'linux'", -] -wheels = [ - { url = "https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/v1.5.0/natten-0.21.6.dev6%2Bcu130.torch210-cp313-cp313-linux_aarch64.whl", hash = "sha256:d157abf86d3f01aa683ef6637abce4f1fdf2b7aa86b5af52736e5d53e9c42da5" }, - { url = "https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/v1.5.0/natten-0.21.6.dev6%2Bcu130.torch210-cp313-cp313-linux_x86_64.whl", hash = "sha256:13d81645a33e58500c33f4c8840c992cab1bf43b2f476b96e7a70c3864906511" }, -] - [[package]] name = "natten" version = "0.21.6.dev6+cu130.torch210.gb300" source = { registry = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.5.0" } resolution-markers = [ + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform == 'linux'", "python_full_version >= '3.14' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.13.*' and platform_machine == 'x86_64' and sys_platform == 'linux'", "python_full_version == '3.13.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.14' and platform_machine == 'x86_64' and sys_platform != 'linux'", "python_full_version >= '3.14' and platform_machine == 'aarch64' and sys_platform != 'linux'", + "python_full_version == '3.13.*' and platform_machine == 'x86_64' and sys_platform != 'linux'", "python_full_version == '3.13.*' and platform_machine == 'aarch64' and sys_platform != 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform == 'linux'", "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'x86_64' and sys_platform != 'linux'", "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform != 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform == 'linux'", "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'x86_64' and sys_platform != 'linux'", "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform != 'linux'", + "python_full_version < '3.11' and platform_machine == 'x86_64' and sys_platform == 'linux'", "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'x86_64' and sys_platform != 'linux'", "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform != 'linux'", ] wheels = [ { url = "https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/v1.5.0/natten-0.21.6.dev6%2Bcu130.torch210.gb300-cp313-cp313-linux_aarch64.whl", hash = "sha256:cd63a31be0ea73c3f5d8bded442508e02a74fc9827f56f92b89314498a2fde9b" }, + { url = "https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/v1.5.0/natten-0.21.6.dev6%2Bcu130.torch210.gb300-cp313-cp313-linux_x86_64.whl", hash = "sha256:04ace5d591b2775149d7a9adc8f3e6099b18fedad364c48e70b4a998ee598b9c" }, ] [[package]] From 463ac142b47c637098a7e86d0631252f82b65418 Mon Sep 17 00:00:00 2001 From: Xiangyu Lu <169013972+xlu451@users.noreply.github.com> Date: Mon, 6 Jul 2026 15:16:36 -0700 Subject: [PATCH 14/26] Allow config aliases as public targets (#4) Co-authored-by: Xiangyu Lu --- cosmos_framework/inference/common/public_model_config.py | 4 ++++ .../inference/common/public_model_config_test.py | 5 +++++ 2 files changed, 9 insertions(+) diff --git a/cosmos_framework/inference/common/public_model_config.py b/cosmos_framework/inference/common/public_model_config.py index 4bd2c679..584a65c9 100644 --- a/cosmos_framework/inference/common/public_model_config.py +++ b/cosmos_framework/inference/common/public_model_config.py @@ -41,6 +41,10 @@ "projects.cosmos3.vfm.configs.base.defaults.vlm.VLMConfig": "vlm_config", } +# Config objects can be serialized as either `_type` or `_target_` depending on +# whether they came from structured config metadata or LazyCall construction. +_TARGET_ALIASES.update(_TYPE_ALIASES) + _ALIAS_TARGETS = {alias: target for target, alias in _TARGET_ALIASES.items()} _ALIAS_TYPES = {alias: target for target, alias in _TYPE_ALIASES.items()} diff --git a/cosmos_framework/inference/common/public_model_config_test.py b/cosmos_framework/inference/common/public_model_config_test.py index def99d3f..40e7e852 100644 --- a/cosmos_framework/inference/common/public_model_config_test.py +++ b/cosmos_framework/inference/common/public_model_config_test.py @@ -28,6 +28,10 @@ def test_public_model_config_round_trip_removes_internal_metadata(): "_type": "cosmos_framework.configs.base.defaults.activation_checkpointing.ActivationCheckpointingConfig", "mode": "full", }, + "compile": { + "_target_": "cosmos_framework.configs.base.defaults.compile.CompileConfig", + "enabled": False, + }, "tokenizer": { "_target_": "cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface", "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth", @@ -56,6 +60,7 @@ def test_public_model_config_round_trip_removes_internal_metadata(): assert not _has_key(public_model_config, "config_name") assert public_model_config["_target"] == "omni_mot_model" assert public_model_config["config"]["_type"] == "omni_mot_model_config" + assert public_model_config["config"]["compile"]["_target"] == "compile_config" assert "projects.cosmos3" not in text assert "projects/cosmos3" not in text assert "cosmos3._src" not in text From 81394820122519e4cdd290c74e491d171f21a1bf Mon Sep 17 00:00:00 2001 From: yy-code-nv Date: Tue, 7 Jul 2026 16:54:39 +0800 Subject: [PATCH 15/26] Release 2026-07-07 (i4 c2c7152e) (#91) Automated release from i4. _source_commit: `c2c7152e59e4090a61822dc22a1691e7afb44702-dirty` _dest_commit (base): `463ac142b47c637098a7e86d0631252f82b65418` --- .file_mapping.json | 468 ++++++++++++++++++ .../callbacks/every_n_draw_sample.py | 306 +++++++++++- cosmos_framework/callbacks/grad_clip.py | 14 +- .../configs/base/defaults/model_config.py | 4 +- .../configs/base/reasoner/config.py | 1 + .../augmentors/text_transforms_for_image.py | 1 - cosmos_framework/data/generator/utils.py | 39 +- .../processors/nemotronvl_processor.py | 6 +- .../diffusion/samplers/fixed_step.py | 55 +- .../diffusion/samplers/fixed_step_test.py | 163 +----- .../generator/mot/context_parallel_test.py | 1 + .../model/generator/omni_mot_model.py | 10 +- .../evaluation/reconstruction_metrics.py | 1 - 13 files changed, 814 insertions(+), 255 deletions(-) create mode 100644 .file_mapping.json diff --git a/.file_mapping.json b/.file_mapping.json new file mode 100644 index 00000000..503fa929 --- /dev/null +++ b/.file_mapping.json @@ -0,0 +1,468 @@ +{ + "_source_commit": "c2c7152e59e4090a61822dc22a1691e7afb44702-dirty", + "_dest_commit": "463ac142b47c637098a7e86d0631252f82b65418", + "_generated_at": "2026-07-07T05:54:17Z", + "files": { + "imaginaire/__init__.py": "cosmos_framework/__init__.py", + "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", + "imaginaire/attention/backends.py": "cosmos_framework/model/attention/backends.py", + "imaginaire/attention/benchmarks/benchmark_fmha_head_sharding.py": "cosmos_framework/model/attention/benchmarks/benchmark_fmha_head_sharding.py", + "imaginaire/attention/checks.py": "cosmos_framework/model/attention/checks.py", + "imaginaire/attention/cudnn/__init__.py": "cosmos_framework/model/attention/cudnn/__init__.py", + "imaginaire/attention/cudnn/checks.py": "cosmos_framework/model/attention/cudnn/checks.py", + "imaginaire/attention/cudnn/cudnn_forward.py": "cosmos_framework/model/attention/cudnn/cudnn_forward.py", + "imaginaire/attention/cudnn/functions.py": "cosmos_framework/model/attention/cudnn/functions.py", + "imaginaire/attention/cudnn/meta.py": "cosmos_framework/model/attention/cudnn/meta.py", + "imaginaire/attention/cudnn/stubs.py": "cosmos_framework/model/attention/cudnn/stubs.py", + "imaginaire/attention/flash2/__init__.py": "cosmos_framework/model/attention/flash2/__init__.py", + "imaginaire/attention/flash2/checks.py": "cosmos_framework/model/attention/flash2/checks.py", + "imaginaire/attention/flash2/functions.py": "cosmos_framework/model/attention/flash2/functions.py", + "imaginaire/attention/flash2/meta.py": "cosmos_framework/model/attention/flash2/meta.py", + "imaginaire/attention/flash2/stubs.py": "cosmos_framework/model/attention/flash2/stubs.py", + "imaginaire/attention/flash3/__init__.py": "cosmos_framework/model/attention/flash3/__init__.py", + "imaginaire/attention/flash3/checks.py": "cosmos_framework/model/attention/flash3/checks.py", + "imaginaire/attention/flash3/functions.py": "cosmos_framework/model/attention/flash3/functions.py", + "imaginaire/attention/flash3/meta.py": "cosmos_framework/model/attention/flash3/meta.py", + "imaginaire/attention/flash3/stubs.py": "cosmos_framework/model/attention/flash3/stubs.py", + "imaginaire/attention/frontend.py": "cosmos_framework/model/attention/frontend.py", + "imaginaire/attention/masks.py": "cosmos_framework/model/attention/masks.py", + "imaginaire/attention/natten/__init__.py": "cosmos_framework/model/attention/natten/__init__.py", + "imaginaire/attention/natten/checks.py": "cosmos_framework/model/attention/natten/checks.py", + "imaginaire/attention/natten/functions.py": "cosmos_framework/model/attention/natten/functions.py", + "imaginaire/attention/natten/meta.py": "cosmos_framework/model/attention/natten/meta.py", + "imaginaire/attention/natten/stubs.py": "cosmos_framework/model/attention/natten/stubs.py", + "imaginaire/attention/utils/__init__.py": "cosmos_framework/model/attention/utils/__init__.py", + "imaginaire/attention/utils/backend_filter_test.py": "cosmos_framework/model/attention/utils/backend_filter_test.py", + "imaginaire/attention/utils/determinism.py": "cosmos_framework/model/attention/utils/determinism.py", + "imaginaire/attention/utils/determinism_test.py": "cosmos_framework/model/attention/utils/determinism_test.py", + "imaginaire/attention/utils/environment.py": "cosmos_framework/model/attention/utils/environment.py", + "imaginaire/attention/utils/safe_ops/__init__.py": "cosmos_framework/model/attention/utils/safe_ops/__init__.py", + "imaginaire/attention/utils/safe_ops/functools.py": "cosmos_framework/model/attention/utils/safe_ops/functools.py", + "imaginaire/attention/utils/safe_ops/log.py": "cosmos_framework/model/attention/utils/safe_ops/log.py", + "imaginaire/attention/utils/safe_ops/safe_lru_cache_test.py": "cosmos_framework/model/attention/utils/safe_ops/safe_lru_cache_test.py", + "imaginaire/attention/utils/version.py": "cosmos_framework/model/attention/utils/version.py", + "imaginaire/attention/utils/version_test.py": "cosmos_framework/model/attention/utils/version_test.py", + "imaginaire/attention/varlen.py": "cosmos_framework/model/attention/varlen.py", + "imaginaire/auxiliary/__init__.py": "cosmos_framework/auxiliary/__init__.py", + "imaginaire/auxiliary/guardrail/__init__.py": "cosmos_framework/auxiliary/guardrail/__init__.py", + "imaginaire/auxiliary/guardrail/blocklist/__init__.py": "cosmos_framework/auxiliary/guardrail/blocklist/__init__.py", + "imaginaire/auxiliary/guardrail/blocklist/blocklist.py": "cosmos_framework/auxiliary/guardrail/blocklist/blocklist.py", + "imaginaire/auxiliary/guardrail/blocklist/blocklist_test.py": "cosmos_framework/auxiliary/guardrail/blocklist/blocklist_test.py", + "imaginaire/auxiliary/guardrail/blocklist/profile_blocklist.py": "cosmos_framework/auxiliary/guardrail/blocklist/profile_blocklist.py", + "imaginaire/auxiliary/guardrail/blocklist/utils.py": "cosmos_framework/auxiliary/guardrail/blocklist/utils.py", + "imaginaire/auxiliary/guardrail/common/__init__.py": "cosmos_framework/auxiliary/guardrail/common/__init__.py", + "imaginaire/auxiliary/guardrail/common/core.py": "cosmos_framework/auxiliary/guardrail/common/core.py", + "imaginaire/auxiliary/guardrail/common/io_utils.py": "cosmos_framework/auxiliary/guardrail/common/io_utils.py", + "imaginaire/auxiliary/guardrail/common/presets.py": "cosmos_framework/auxiliary/guardrail/common/presets.py", + "imaginaire/auxiliary/guardrail/face_blur_filter/__init__.py": "cosmos_framework/auxiliary/guardrail/face_blur_filter/__init__.py", + "imaginaire/auxiliary/guardrail/face_blur_filter/blur_utils.py": "cosmos_framework/auxiliary/guardrail/face_blur_filter/blur_utils.py", + "imaginaire/auxiliary/guardrail/face_blur_filter/face_blur_filter.py": "cosmos_framework/auxiliary/guardrail/face_blur_filter/face_blur_filter.py", + "imaginaire/auxiliary/guardrail/face_blur_filter/retinaface_utils.py": "cosmos_framework/auxiliary/guardrail/face_blur_filter/retinaface_utils.py", + "imaginaire/auxiliary/guardrail/llamaGuard3/__init__.py": "cosmos_framework/auxiliary/guardrail/llamaGuard3/__init__.py", + "imaginaire/auxiliary/guardrail/llamaGuard3/categories.py": "cosmos_framework/auxiliary/guardrail/llamaGuard3/categories.py", + "imaginaire/auxiliary/guardrail/llamaGuard3/llamaGuard3.py": "cosmos_framework/auxiliary/guardrail/llamaGuard3/llamaGuard3.py", + "imaginaire/auxiliary/guardrail/qwen3guard/__init__.py": "cosmos_framework/auxiliary/guardrail/qwen3guard/__init__.py", + "imaginaire/auxiliary/guardrail/qwen3guard/categories.py": "cosmos_framework/auxiliary/guardrail/qwen3guard/categories.py", + "imaginaire/auxiliary/guardrail/qwen3guard/qwen3guard.py": "cosmos_framework/auxiliary/guardrail/qwen3guard/qwen3guard.py", + "imaginaire/auxiliary/guardrail/video_content_safety_filter/__init__.py": "cosmos_framework/auxiliary/guardrail/video_content_safety_filter/__init__.py", + "imaginaire/auxiliary/guardrail/video_content_safety_filter/model.py": "cosmos_framework/auxiliary/guardrail/video_content_safety_filter/model.py", + "imaginaire/auxiliary/guardrail/video_content_safety_filter/video_content_safety_filter.py": "cosmos_framework/auxiliary/guardrail/video_content_safety_filter/video_content_safety_filter.py", + "imaginaire/auxiliary/guardrail/video_content_safety_filter/vision_encoder.py": "cosmos_framework/auxiliary/guardrail/video_content_safety_filter/vision_encoder.py", + "imaginaire/callbacks/every_n.py": "cosmos_framework/callbacks/every_n.py", + "imaginaire/callbacks/manual_gc.py": "cosmos_framework/callbacks/manual_gc.py", + "imaginaire/checkpointer/base.py": "cosmos_framework/checkpoint/base.py", + "imaginaire/checkpointer/dummy.py": "cosmos_framework/checkpoint/dummy.py", + "imaginaire/checkpointer/s3_filesystem.py": "cosmos_framework/checkpoint/s3_filesystem.py", + "imaginaire/config.py": "cosmos_framework/utils/config.py", + "imaginaire/configs/lr_scheduler.py": "cosmos_framework/utils/configs/lr_scheduler.py", + "imaginaire/datasets/augmentors/v3_text_transforms.py": "cosmos_framework/data/imaginaire/webdataset/augmentors/v3_text_transforms.py", + "imaginaire/datasets/webdataset/augmentors/augmentor.py": "cosmos_framework/data/imaginaire/webdataset/augmentors/augmentor.py", + "imaginaire/datasets/webdataset/augmentors/image/__init__.py": "cosmos_framework/data/imaginaire/webdataset/augmentors/image/__init__.py", + "imaginaire/datasets/webdataset/augmentors/image/cropping.py": "cosmos_framework/data/imaginaire/webdataset/augmentors/image/cropping.py", + "imaginaire/datasets/webdataset/augmentors/image/flip.py": "cosmos_framework/data/imaginaire/webdataset/augmentors/image/flip.py", + "imaginaire/datasets/webdataset/augmentors/image/misc.py": "cosmos_framework/data/imaginaire/webdataset/augmentors/image/misc.py", + "imaginaire/datasets/webdataset/augmentors/image/normalize.py": "cosmos_framework/data/imaginaire/webdataset/augmentors/image/normalize.py", + "imaginaire/datasets/webdataset/augmentors/image/padding.py": "cosmos_framework/data/imaginaire/webdataset/augmentors/image/padding.py", + "imaginaire/datasets/webdataset/augmentors/image/resize.py": "cosmos_framework/data/imaginaire/webdataset/augmentors/image/resize.py", + "imaginaire/flags.py": "cosmos_framework/utils/flags.py", + "imaginaire/flops/__init__.py": "cosmos_framework/tools/flops/__init__.py", + "imaginaire/flops/omni_mot.py": "cosmos_framework/tools/flops/omni_mot.py", + "imaginaire/flops/qwen3_vl.py": "cosmos_framework/tools/flops/qwen3_vl.py", + "imaginaire/flops/wan_vae.py": "cosmos_framework/tools/flops/wan_vae.py", + "imaginaire/functional/batch_ops.py": "cosmos_framework/utils/functional/batch_ops.py", + "imaginaire/functional/lr_scheduler.py": "cosmos_framework/utils/functional/lr_scheduler.py", + "imaginaire/functional/multi_step.py": "cosmos_framework/utils/functional/multi_step.py", + "imaginaire/functional/runge_kutta.py": "cosmos_framework/utils/functional/runge_kutta.py", + "imaginaire/lazy_config/__init__.py": "cosmos_framework/utils/lazy_config/__init__.py", + "imaginaire/lazy_config/file_io.py": "cosmos_framework/utils/lazy_config/file_io.py", + "imaginaire/lazy_config/instantiate.py": "cosmos_framework/utils/lazy_config/instantiate.py", + "imaginaire/lazy_config/lazy.py": "cosmos_framework/utils/lazy_config/lazy.py", + "imaginaire/lazy_config/lazy_call.py": "cosmos_framework/utils/lazy_config/lazy_call.py", + "imaginaire/lazy_config/omegaconf_patch.py": "cosmos_framework/utils/lazy_config/omegaconf_patch.py", + "imaginaire/lazy_config/registry.py": "cosmos_framework/utils/lazy_config/registry.py", + "imaginaire/model.py": "cosmos_framework/model/_base.py", + "imaginaire/serialization.py": "cosmos_framework/utils/serialization.py", + "imaginaire/trainer.py": "cosmos_framework/trainer/__init__.py", + "imaginaire/utils/__init__.py": "cosmos_framework/utils/__init__.py", + "imaginaire/utils/callback.py": "cosmos_framework/utils/callback.py", + "imaginaire/utils/checkpoint_db.py": "cosmos_framework/utils/checkpoint_db.py", + "imaginaire/utils/checkpointer.py": "cosmos_framework/utils/checkpointer.py", + "imaginaire/utils/cluster_env.py": "cosmos_framework/utils/cluster_env.py", + "imaginaire/utils/config_helper.py": "cosmos_framework/utils/config_helper.py", + "imaginaire/utils/context_managers.py": "cosmos_framework/utils/context_managers.py", + "imaginaire/utils/count_params.py": "cosmos_framework/utils/count_params.py", + "imaginaire/utils/device.py": "cosmos_framework/utils/device.py", + "imaginaire/utils/distributed.py": "cosmos_framework/utils/distributed.py", + "imaginaire/utils/easy_io/README.md": "cosmos_framework/utils/easy_io/README.md", + "imaginaire/utils/easy_io/__init__.py": "cosmos_framework/utils/easy_io/__init__.py", + "imaginaire/utils/easy_io/backends/__init__.py": "cosmos_framework/utils/easy_io/backends/__init__.py", + "imaginaire/utils/easy_io/backends/auto_auth.py": "cosmos_framework/utils/easy_io/backends/auto_auth.py", + "imaginaire/utils/easy_io/backends/base_backend.py": "cosmos_framework/utils/easy_io/backends/base_backend.py", + "imaginaire/utils/easy_io/backends/boto3_backend.py": "cosmos_framework/utils/easy_io/backends/boto3_backend.py", + "imaginaire/utils/easy_io/backends/boto3_client.py": "cosmos_framework/utils/easy_io/backends/boto3_client.py", + "imaginaire/utils/easy_io/backends/http_backend.py": "cosmos_framework/utils/easy_io/backends/http_backend.py", + "imaginaire/utils/easy_io/backends/local_backend.py": "cosmos_framework/utils/easy_io/backends/local_backend.py", + "imaginaire/utils/easy_io/backends/msc_backend.py": "cosmos_framework/utils/easy_io/backends/msc_backend.py", + "imaginaire/utils/easy_io/backends/registry_utils.py": "cosmos_framework/utils/easy_io/backends/registry_utils.py", + "imaginaire/utils/easy_io/easy_io.py": "cosmos_framework/utils/easy_io/easy_io.py", + "imaginaire/utils/easy_io/file_client.py": "cosmos_framework/utils/easy_io/file_client.py", + "imaginaire/utils/easy_io/handlers/__init__.py": "cosmos_framework/utils/easy_io/handlers/__init__.py", + "imaginaire/utils/easy_io/handlers/base.py": "cosmos_framework/utils/easy_io/handlers/base.py", + "imaginaire/utils/easy_io/handlers/byte_handler.py": "cosmos_framework/utils/easy_io/handlers/byte_handler.py", + "imaginaire/utils/easy_io/handlers/csv_handler.py": "cosmos_framework/utils/easy_io/handlers/csv_handler.py", + "imaginaire/utils/easy_io/handlers/gzip_handler.py": "cosmos_framework/utils/easy_io/handlers/gzip_handler.py", + "imaginaire/utils/easy_io/handlers/imageio_video_handler.py": "cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py", + "imaginaire/utils/easy_io/handlers/json_handler.py": "cosmos_framework/utils/easy_io/handlers/json_handler.py", + "imaginaire/utils/easy_io/handlers/jsonl_handler.py": "cosmos_framework/utils/easy_io/handlers/jsonl_handler.py", + "imaginaire/utils/easy_io/handlers/np_handler.py": "cosmos_framework/utils/easy_io/handlers/np_handler.py", + "imaginaire/utils/easy_io/handlers/pandas_handler.py": "cosmos_framework/utils/easy_io/handlers/pandas_handler.py", + "imaginaire/utils/easy_io/handlers/pickle_handler.py": "cosmos_framework/utils/easy_io/handlers/pickle_handler.py", + "imaginaire/utils/easy_io/handlers/pil_handler.py": "cosmos_framework/utils/easy_io/handlers/pil_handler.py", + "imaginaire/utils/easy_io/handlers/registry_utils.py": "cosmos_framework/utils/easy_io/handlers/registry_utils.py", + "imaginaire/utils/easy_io/handlers/tarfile_handler.py": "cosmos_framework/utils/easy_io/handlers/tarfile_handler.py", + "imaginaire/utils/easy_io/handlers/torch_handler.py": "cosmos_framework/utils/easy_io/handlers/torch_handler.py", + "imaginaire/utils/easy_io/handlers/torchjit_handler.py": "cosmos_framework/utils/easy_io/handlers/torchjit_handler.py", + "imaginaire/utils/easy_io/handlers/trimesh_handler.py": "cosmos_framework/utils/easy_io/handlers/trimesh_handler.py", + "imaginaire/utils/easy_io/handlers/txt_handler.py": "cosmos_framework/utils/easy_io/handlers/txt_handler.py", + "imaginaire/utils/easy_io/handlers/yaml_handler.py": "cosmos_framework/utils/easy_io/handlers/yaml_handler.py", + "imaginaire/utils/easy_io/transient_retry.py": "cosmos_framework/utils/easy_io/transient_retry.py", + "imaginaire/utils/easy_io/transient_retry_test.py": "cosmos_framework/utils/easy_io/transient_retry_test.py", + "imaginaire/utils/ema.py": "cosmos_framework/utils/ema.py", + "imaginaire/utils/env_parsers/cred_env_parser.py": "cosmos_framework/utils/env_parsers/cred_env_parser.py", + "imaginaire/utils/env_parsers/customization_env_parser.py": "cosmos_framework/utils/env_parsers/customization_env_parser.py", + "imaginaire/utils/env_parsers/env_parser.py": "cosmos_framework/utils/env_parsers/env_parser.py", + "imaginaire/utils/env_parsers/inference_env_parser.py": "cosmos_framework/utils/env_parsers/inference_env_parser.py", + "imaginaire/utils/fused_adam.py": "cosmos_framework/utils/fused_adam.py", + "imaginaire/utils/launch.py": "cosmos_framework/utils/launch.py", + "imaginaire/utils/log.py": "cosmos_framework/utils/log.py", + "imaginaire/utils/misc.py": "cosmos_framework/utils/misc.py", + "imaginaire/utils/object_store.py": "cosmos_framework/utils/object_store.py", + "imaginaire/utils/one_logger/one_logger_context_managers.py": "cosmos_framework/utils/one_logger/one_logger_context_managers.py", + "imaginaire/utils/one_logger/one_logger_global_vars.py": "cosmos_framework/utils/one_logger/one_logger_global_vars.py", + "imaginaire/utils/one_logger/one_logger_override_utils.py": "cosmos_framework/utils/one_logger/one_logger_override_utils.py", + "imaginaire/utils/one_logger/one_logger_utils.py": "cosmos_framework/utils/one_logger/one_logger_utils.py", + "imaginaire/utils/optim_instantiate.py": "cosmos_framework/utils/optim_instantiate.py", + "imaginaire/utils/profiling.py": "cosmos_framework/utils/profiling.py", + "imaginaire/utils/progress_bar.py": "cosmos_framework/utils/progress_bar.py", + "imaginaire/utils/timer.py": "cosmos_framework/utils/timer.py", + "imaginaire/utils/training_telemetry/__init__.py": "cosmos_framework/utils/training_telemetry/__init__.py", + "imaginaire/utils/training_telemetry/callback.py": "cosmos_framework/utils/training_telemetry/callback.py", + "imaginaire/utils/training_telemetry/context_managers.py": "cosmos_framework/utils/training_telemetry/context_managers.py", + "imaginaire/utils/training_telemetry/telemetry.py": "cosmos_framework/utils/training_telemetry/telemetry.py", + "imaginaire/utils/training_telemetry/utils.py": "cosmos_framework/utils/training_telemetry/utils.py", + "imaginaire/utils/validator.py": "cosmos_framework/utils/validator.py", + "imaginaire/utils/wandb_util.py": "cosmos_framework/utils/wandb_util.py", + "imaginaire/visualize/__init__.py": "cosmos_framework/tools/visualize/__init__.py", + "imaginaire/visualize/img.py": "cosmos_framework/tools/visualize/img.py", + "imaginaire/visualize/video.py": "cosmos_framework/tools/visualize/video.py", + "projects/cosmos3/cosmos3/algorithm/loss/__init__.py": "cosmos_framework/model/generator/algorithm/loss/__init__.py", + "projects/cosmos3/cosmos3/algorithm/loss/cross_entropy.py": "cosmos_framework/model/generator/algorithm/loss/cross_entropy.py", + "projects/cosmos3/cosmos3/algorithm/loss/flow_matching.py": "cosmos_framework/model/generator/algorithm/loss/flow_matching.py", + "projects/cosmos3/cosmos3/algorithm/loss/load_balancing.py": "cosmos_framework/model/generator/algorithm/loss/load_balancing.py", + "projects/cosmos3/cosmos3/algorithm/loss/time_weight.py": "cosmos_framework/model/generator/algorithm/loss/time_weight.py", + "projects/cosmos3/cosmos3/callbacks/compile_tokenizer.py": "cosmos_framework/callbacks/compile_tokenizer.py", + "projects/cosmos3/cosmos3/callbacks/data_stats.py": "cosmos_framework/callbacks/data_stats.py", + "projects/cosmos3/cosmos3/callbacks/dataloader_state.py": "cosmos_framework/callbacks/dataloader_state.py", + "projects/cosmos3/cosmos3/callbacks/device_monitor.py": "cosmos_framework/callbacks/device_monitor.py", + "projects/cosmos3/cosmos3/callbacks/dit_image_sample.py": "cosmos_framework/callbacks/dit_image_sample.py", + "projects/cosmos3/cosmos3/callbacks/every_n_draw_sample.py": "cosmos_framework/callbacks/every_n_draw_sample.py", + "projects/cosmos3/cosmos3/callbacks/expert_heatmap.py": "cosmos_framework/callbacks/expert_heatmap.py", + "projects/cosmos3/cosmos3/callbacks/grad_clip.py": "cosmos_framework/callbacks/grad_clip.py", + "projects/cosmos3/cosmos3/callbacks/heart_beat.py": "cosmos_framework/callbacks/heart_beat.py", + "projects/cosmos3/cosmos3/callbacks/hf_export.py": "cosmos_framework/callbacks/hf_export.py", + "projects/cosmos3/cosmos3/callbacks/iter_speed.py": "cosmos_framework/callbacks/iter_speed.py", + "projects/cosmos3/cosmos3/callbacks/learning_rate_logger.py": "cosmos_framework/callbacks/learning_rate_logger.py", + "projects/cosmos3/cosmos3/callbacks/load_pretrained.py": "cosmos_framework/callbacks/load_pretrained.py", + "projects/cosmos3/cosmos3/callbacks/log_tensor_shape.py": "cosmos_framework/callbacks/log_tensor_shape.py", + "projects/cosmos3/cosmos3/callbacks/mfu.py": "cosmos_framework/callbacks/mfu.py", + "projects/cosmos3/cosmos3/callbacks/moe_specialization_callback.py": "cosmos_framework/callbacks/moe_specialization_callback.py", + "projects/cosmos3/cosmos3/callbacks/moe_stability_callback.py": "cosmos_framework/callbacks/moe_stability_callback.py", + "projects/cosmos3/cosmos3/callbacks/norm_monitor.py": "cosmos_framework/callbacks/norm_monitor.py", + "projects/cosmos3/cosmos3/callbacks/ofu.py": "cosmos_framework/callbacks/ofu.py", + "projects/cosmos3/cosmos3/callbacks/param_count.py": "cosmos_framework/callbacks/param_count.py", + "projects/cosmos3/cosmos3/callbacks/sequence_packing_padding.py": "cosmos_framework/callbacks/sequence_packing_padding.py", + "projects/cosmos3/cosmos3/callbacks/sigma_loss_analysis.py": "cosmos_framework/callbacks/sigma_loss_analysis.py", + "projects/cosmos3/cosmos3/callbacks/skip_nan_step.py": "cosmos_framework/callbacks/skip_nan_step.py", + "projects/cosmos3/cosmos3/callbacks/termination_signal_checkpoint.py": "cosmos_framework/callbacks/termination_signal_checkpoint.py", + "projects/cosmos3/cosmos3/callbacks/tokens_per_sec.py": "cosmos_framework/callbacks/tokens_per_sec.py", + "projects/cosmos3/cosmos3/callbacks/training_stats.py": "cosmos_framework/callbacks/training_stats.py", + "projects/cosmos3/cosmos3/callbacks/wandb_log.py": "cosmos_framework/callbacks/wandb_log.py", + "projects/cosmos3/cosmos3/callbacks/wandb_log_eval.py": "cosmos_framework/callbacks/wandb_log_eval.py", + "projects/cosmos3/cosmos3/callbacks/wandb_vis.py": "cosmos_framework/callbacks/wandb_vis.py", + "projects/cosmos3/cosmos3/checkpointer/dcp.py": "cosmos_framework/checkpoint/dcp.py", + "projects/cosmos3/cosmos3/configs/base/__init__.py": "cosmos_framework/configs/base/__init__.py", + "projects/cosmos3/cosmos3/configs/base/defaults/__init__.py": "cosmos_framework/configs/base/defaults/__init__.py", + "projects/cosmos3/cosmos3/configs/base/defaults/activation_checkpointing.py": "cosmos_framework/configs/base/defaults/activation_checkpointing.py", + "projects/cosmos3/cosmos3/configs/base/defaults/callbacks.py": "cosmos_framework/configs/base/defaults/callbacks.py", + "projects/cosmos3/cosmos3/configs/base/defaults/checkpointer.py": "cosmos_framework/configs/base/defaults/checkpointer.py", + "projects/cosmos3/cosmos3/configs/base/defaults/compile.py": "cosmos_framework/configs/base/defaults/compile.py", + "projects/cosmos3/cosmos3/configs/base/defaults/ema.py": "cosmos_framework/configs/base/defaults/ema.py", + "projects/cosmos3/cosmos3/configs/base/defaults/model.py": "cosmos_framework/configs/base/defaults/model.py", + "projects/cosmos3/cosmos3/configs/base/defaults/model_config.py": "cosmos_framework/configs/base/defaults/model_config.py", + "projects/cosmos3/cosmos3/configs/base/defaults/optimizer.py": "cosmos_framework/configs/base/defaults/optimizer.py", + "projects/cosmos3/cosmos3/configs/base/defaults/parallelism.py": "cosmos_framework/configs/base/defaults/parallelism.py", + "projects/cosmos3/cosmos3/configs/base/defaults/reasoner.py": "cosmos_framework/configs/base/defaults/reasoner.py", + "projects/cosmos3/cosmos3/configs/base/defaults/tokenizer.py": "cosmos_framework/configs/base/defaults/tokenizer.py", + "projects/cosmos3/cosmos3/configs/base/experiment/action/__init__.py": "cosmos_framework/configs/base/experiment/action/__init__.py", + "projects/cosmos3/cosmos3/configs/base/experiment/action/posttrain_config/__init__.py": "cosmos_framework/configs/base/experiment/action/posttrain_config/__init__.py", + "projects/cosmos3/cosmos3/configs/base/experiment/action/pretrained_config/__init__.py": "cosmos_framework/configs/base/experiment/action/pretrained_config/__init__.py", + "projects/cosmos3/cosmos3/configs/base/experiment/posttrain_video/__init__.py": "cosmos_framework/configs/base/experiment/posttrain_video/__init__.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/__init__.py": "cosmos_framework/configs/base/reasoner/__init__.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/config.py": "cosmos_framework/configs/base/reasoner/config.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/defaults/__init__.py": "cosmos_framework/configs/base/reasoner/defaults/__init__.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/defaults/callbacks.py": "cosmos_framework/configs/base/reasoner/defaults/callbacks.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/defaults/config.py": "cosmos_framework/configs/base/reasoner/defaults/config.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/defaults/model.py": "cosmos_framework/configs/base/reasoner/defaults/model.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/defaults/optimizer.py": "cosmos_framework/configs/base/reasoner/defaults/optimizer.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/defaults/policy_config.py": "cosmos_framework/configs/base/reasoner/defaults/policy_config.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/defaults/vlm_policy.py": "cosmos_framework/configs/base/reasoner/defaults/vlm_policy.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/experiment/__init__.py": "cosmos_framework/configs/base/reasoner/experiment/__init__.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/experiment/utils.py": "cosmos_framework/configs/base/reasoner/experiment/utils.py", + "projects/cosmos3/cosmos3/configs/base/reasoner/freeze_config.py": "cosmos_framework/configs/base/reasoner/freeze_config.py", + "projects/cosmos3/cosmos3/datasets/action/__init__.py": "cosmos_framework/data/generator/action/__init__.py", + "projects/cosmos3/cosmos3/datasets/action/action_processing.py": "cosmos_framework/data/generator/action/action_processing.py", + "projects/cosmos3/cosmos3/datasets/action/action_spec.py": "cosmos_framework/data/generator/action/action_spec.py", + "projects/cosmos3/cosmos3/datasets/action/domain_utils.py": "cosmos_framework/data/generator/action/domain_utils.py", + "projects/cosmos3/cosmos3/datasets/action/json_formatter.py": "cosmos_framework/data/generator/action/json_formatter.py", + "projects/cosmos3/cosmos3/datasets/action/pose_utils.py": "cosmos_framework/data/generator/action/pose_utils.py", + "projects/cosmos3/cosmos3/datasets/action/pose_utils_test.py": "cosmos_framework/data/generator/action/pose_utils_test.py", + "projects/cosmos3/cosmos3/datasets/action/transforms.py": "cosmos_framework/data/generator/action/transforms.py", + "projects/cosmos3/cosmos3/datasets/action/transforms_test.py": "cosmos_framework/data/generator/action/transforms_test.py", + "projects/cosmos3/cosmos3/datasets/action/viewpoint_utils.py": "cosmos_framework/data/generator/action/viewpoint_utils.py", + "projects/cosmos3/cosmos3/datasets/augmentor_provider.py": "cosmos_framework/data/generator/augmentor_provider.py", + "projects/cosmos3/cosmos3/datasets/augmentors/__init__.py": "cosmos_framework/data/generator/augmentors/__init__.py", + "projects/cosmos3/cosmos3/datasets/augmentors/append_fps_frames_for_image.py": "cosmos_framework/data/generator/augmentors/append_fps_frames_for_image.py", + "projects/cosmos3/cosmos3/datasets/augmentors/audio_caption.py": "cosmos_framework/data/generator/augmentors/audio_caption.py", + "projects/cosmos3/cosmos3/datasets/augmentors/audio_parsing.py": "cosmos_framework/data/generator/augmentors/audio_parsing.py", + "projects/cosmos3/cosmos3/datasets/augmentors/caption_filter.py": "cosmos_framework/data/generator/augmentors/caption_filter.py", + "projects/cosmos3/cosmos3/datasets/augmentors/cropping.py": "cosmos_framework/data/generator/augmentors/cropping.py", + "projects/cosmos3/cosmos3/datasets/augmentors/duration_fps_text_timestamps.py": "cosmos_framework/data/generator/augmentors/duration_fps_text_timestamps.py", + "projects/cosmos3/cosmos3/datasets/augmentors/idle_frames_text_info.py": "cosmos_framework/data/generator/augmentors/idle_frames_text_info.py", + "projects/cosmos3/cosmos3/datasets/augmentors/image_editing_transform.py": "cosmos_framework/data/generator/augmentors/image_editing_transform.py", + "projects/cosmos3/cosmos3/datasets/augmentors/image_editing_transform_test.py": "cosmos_framework/data/generator/augmentors/image_editing_transform_test.py", + "projects/cosmos3/cosmos3/datasets/augmentors/image_resolution_filter.py": "cosmos_framework/data/generator/augmentors/image_resolution_filter.py", + "projects/cosmos3/cosmos3/datasets/augmentors/interleaved_image_transform.py": "cosmos_framework/data/generator/augmentors/interleaved_image_transform.py", + "projects/cosmos3/cosmos3/datasets/augmentors/interleaved_video_parsing.py": "cosmos_framework/data/generator/augmentors/interleaved_video_parsing.py", + "projects/cosmos3/cosmos3/datasets/augmentors/merge_datadict.py": "cosmos_framework/data/generator/augmentors/merge_datadict.py", + "projects/cosmos3/cosmos3/datasets/augmentors/multi_reference_transform.py": "cosmos_framework/data/generator/augmentors/multi_reference_transform.py", + "projects/cosmos3/cosmos3/datasets/augmentors/pkl_to_media.py": "cosmos_framework/data/generator/augmentors/pkl_to_media.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/__init__.py": "cosmos_framework/data/generator/augmentors/reasoner/__init__.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/bytes_to_media.py": "cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/filter_output_key.py": "cosmos_framework/data/generator/augmentors/reasoner/filter_output_key.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/filter_seq_length.py": "cosmos_framework/data/generator/augmentors/reasoner/filter_seq_length.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/floating_number_format.py": "cosmos_framework/data/generator/augmentors/reasoner/floating_number_format.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/format_describe_anything.py": "cosmos_framework/data/generator/augmentors/reasoner/format_describe_anything.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/nvlm_data_to_conversation.py": "cosmos_framework/data/generator/augmentors/reasoner/nvlm_data_to_conversation.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/prompt_format.py": "cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/shuffle_text_media_order.py": "cosmos_framework/data/generator/augmentors/reasoner/shuffle_text_media_order.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/timestamp.py": "cosmos_framework/data/generator/augmentors/reasoner/timestamp.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/timestamp_with_subject_tracking.py": "cosmos_framework/data/generator/augmentors/reasoner/timestamp_with_subject_tracking.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/timestamp_without_augment_message.py": "cosmos_framework/data/generator/augmentors/reasoner/timestamp_without_augment_message.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/timestamp_without_end_time.py": "cosmos_framework/data/generator/augmentors/reasoner/timestamp_without_end_time.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/tokenize_data.py": "cosmos_framework/data/generator/augmentors/reasoner/tokenize_data.py", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/user_prompt_caption_general.json": "cosmos_framework/data/generator/augmentors/reasoner/user_prompt_caption_general.json", + "projects/cosmos3/cosmos3/datasets/augmentors/reasoner/user_prompt_ocr.json": "cosmos_framework/data/generator/augmentors/reasoner/user_prompt_ocr.json", + "projects/cosmos3/cosmos3/datasets/augmentors/resolution_text_info.py": "cosmos_framework/data/generator/augmentors/resolution_text_info.py", + "projects/cosmos3/cosmos3/datasets/augmentors/sequence_plan.py": "cosmos_framework/data/generator/augmentors/sequence_plan.py", + "projects/cosmos3/cosmos3/datasets/augmentors/sound_sequence_plan.py": "cosmos_framework/data/generator/augmentors/sound_sequence_plan.py", + "projects/cosmos3/cosmos3/datasets/augmentors/text_tokenizer.py": "cosmos_framework/data/generator/augmentors/text_tokenizer.py", + "projects/cosmos3/cosmos3/datasets/augmentors/text_transforms_for_image.py": "cosmos_framework/data/generator/augmentors/text_transforms_for_image.py", + "projects/cosmos3/cosmos3/datasets/augmentors/text_transforms_for_video.py": "cosmos_framework/data/generator/augmentors/text_transforms_for_video.py", + "projects/cosmos3/cosmos3/datasets/augmentors/torchcodec_callers_test.py": "cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py", + "projects/cosmos3/cosmos3/datasets/augmentors/transfer_control_input/__init__.py": "cosmos_framework/data/generator/augmentors/transfer_control_input/__init__.py", + "projects/cosmos3/cosmos3/datasets/augmentors/transfer_control_input/blur.py": "cosmos_framework/data/generator/augmentors/transfer_control_input/blur.py", + "projects/cosmos3/cosmos3/datasets/augmentors/transfer_control_input/control_input.py": "cosmos_framework/data/generator/augmentors/transfer_control_input/control_input.py", + "projects/cosmos3/cosmos3/datasets/augmentors/transfer_control_input/fast_blur.py": "cosmos_framework/data/generator/augmentors/transfer_control_input/fast_blur.py", + "projects/cosmos3/cosmos3/datasets/augmentors/transfer_control_input/seg.py": "cosmos_framework/data/generator/augmentors/transfer_control_input/seg.py", + "projects/cosmos3/cosmos3/datasets/augmentors/transfer_control_transform.py": "cosmos_framework/data/generator/augmentors/transfer_control_transform.py", + "projects/cosmos3/cosmos3/datasets/augmentors/video_parsing.py": "cosmos_framework/data/generator/augmentors/video_parsing.py", + "projects/cosmos3/cosmos3/datasets/joint_dataloader.py": "cosmos_framework/data/generator/joint_dataloader.py", + "projects/cosmos3/cosmos3/datasets/packing_iterable_dataset.py": "cosmos_framework/data/generator/packing_iterable_dataset.py", + "projects/cosmos3/cosmos3/datasets/reasoner/video_decoder_qwen.py": "cosmos_framework/data/generator/reasoner/video_decoder_qwen.py", + "projects/cosmos3/cosmos3/datasets/recipe_database_api.py": "cosmos_framework/data/generator/recipe_database_api.py", + "projects/cosmos3/cosmos3/datasets/sound_data_utils.py": "cosmos_framework/data/generator/sound_data_utils.py", + "projects/cosmos3/cosmos3/datasets/utils.py": "cosmos_framework/data/generator/utils.py", + "projects/cosmos3/cosmos3/datasets/watchdog.py": "cosmos_framework/data/generator/watchdog.py", + "projects/cosmos3/cosmos3/diffusion/rectified_flow.py": "cosmos_framework/model/generator/diffusion/rectified_flow.py", + "projects/cosmos3/cosmos3/diffusion/samplers/__init__.py": "cosmos_framework/model/generator/diffusion/samplers/__init__.py", + "projects/cosmos3/cosmos3/diffusion/samplers/edm.py": "cosmos_framework/model/generator/diffusion/samplers/edm.py", + "projects/cosmos3/cosmos3/diffusion/samplers/fixed_step.py": "cosmos_framework/model/generator/diffusion/samplers/fixed_step.py", + "projects/cosmos3/cosmos3/diffusion/samplers/fixed_step_test.py": "cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py", + "projects/cosmos3/cosmos3/diffusion/samplers/fm_solvers_unipc.py": "cosmos_framework/model/generator/diffusion/samplers/fm_solvers_unipc.py", + "projects/cosmos3/cosmos3/diffusion/samplers/unipc.py": "cosmos_framework/model/generator/diffusion/samplers/unipc.py", + "projects/cosmos3/cosmos3/diffusion/samplers/utils.py": "cosmos_framework/model/generator/diffusion/samplers/utils.py", + "projects/cosmos3/cosmos3/diffusion/samplers/utils_test.py": "cosmos_framework/model/generator/diffusion/samplers/utils_test.py", + "projects/cosmos3/cosmos3/models/hf_model.py": "cosmos_framework/model/generator/hf_model.py", + "projects/cosmos3/cosmos3/models/mot/__init__.py": "cosmos_framework/model/generator/mot/__init__.py", + "projects/cosmos3/cosmos3/models/mot/attention.py": "cosmos_framework/model/generator/mot/attention.py", + "projects/cosmos3/cosmos3/models/mot/attention_test.py": "cosmos_framework/model/generator/mot/attention_test.py", + "projects/cosmos3/cosmos3/models/mot/cfgp_ar_test.py": "cosmos_framework/model/generator/mot/cfgp_ar_test.py", + "projects/cosmos3/cosmos3/models/mot/context_parallel_test.py": "cosmos_framework/model/generator/mot/context_parallel_test.py", + "projects/cosmos3/cosmos3/models/mot/context_parallel_utils.py": "cosmos_framework/model/generator/mot/context_parallel_utils.py", + "projects/cosmos3/cosmos3/models/mot/cosmos3_vfm_network.py": "cosmos_framework/model/generator/mot/cosmos3_vfm_network.py", + "projects/cosmos3/cosmos3/models/mot/domain_aware_linear.py": "cosmos_framework/model/generator/mot/domain_aware_linear.py", + "projects/cosmos3/cosmos3/models/mot/dot_product_attention.py": "cosmos_framework/model/generator/mot/dot_product_attention.py", + "projects/cosmos3/cosmos3/models/mot/modeling_utils.py": "cosmos_framework/model/generator/mot/modeling_utils.py", + "projects/cosmos3/cosmos3/models/mot/parallelize_unified_mot.py": "cosmos_framework/model/generator/mot/parallelize_unified_mot.py", + "projects/cosmos3/cosmos3/models/mot/parallelize_vfm_network.py": "cosmos_framework/model/generator/mot/parallelize_vfm_network.py", + "projects/cosmos3/cosmos3/models/mot/und_k_norm_example_test.py": "cosmos_framework/model/generator/mot/und_k_norm_example_test.py", + "projects/cosmos3/cosmos3/models/mot/unified_mot.py": "cosmos_framework/model/generator/mot/unified_mot.py", + "projects/cosmos3/cosmos3/models/omni_mot_model.py": "cosmos_framework/model/generator/omni_mot_model.py", + "projects/cosmos3/cosmos3/models/parallelize_vlm.py": "cosmos_framework/model/generator/parallelize_vlm.py", + "projects/cosmos3/cosmos3/models/reasoner/__init__.py": "cosmos_framework/model/generator/reasoner/__init__.py", + "projects/cosmos3/cosmos3/models/reasoner/nemotron_3_dense_vl/__init__.py": "cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/__init__.py", + "projects/cosmos3/cosmos3/models/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json": "cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json", + "projects/cosmos3/cosmos3/models/reasoner/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py": "cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configuration_nemotron_3_dense_vl.py", + "projects/cosmos3/cosmos3/models/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py": "cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/nemotron_3_dense_vl.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/__init__.py": "cosmos_framework/model/generator/reasoner/qwen3_vl/__init__.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json": "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-2B-Instruct.json", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json": "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json": "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-4B-Instruct.json", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json": "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-8B-Instruct.json", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/configs/__init__.py": "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/__init__.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/configuration_qwen3_vl.py": "cosmos_framework/model/generator/reasoner/qwen3_vl/configuration_qwen3_vl.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/qwen3_vl.py": "cosmos_framework/model/generator/reasoner/qwen3_vl/qwen3_vl.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/utils.py": "cosmos_framework/model/generator/reasoner/qwen3_vl/utils.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl/video_processing_qwen3_vl.py": "cosmos_framework/model/generator/reasoner/qwen3_vl/video_processing_qwen3_vl.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/__init__.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/__init__.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-235B-A22B-Instruct.json", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/Qwen3-VL-30B-A3B-Instruct.json", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/configs/__init__.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configs/__init__.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/configuration_qwen3_vl_moe.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/moe.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/moe_bench.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_bench.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/moe_kernels.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_kernels.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/moe_test.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py", + "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/qwen3_vl_moe.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/qwen3_vl_moe.py", + "projects/cosmos3/cosmos3/models/utils/__init__.py": "cosmos_framework/model/generator/utils/__init__.py", + "projects/cosmos3/cosmos3/models/utils/data_and_condition.py": "cosmos_framework/model/generator/utils/data_and_condition.py", + "projects/cosmos3/cosmos3/models/utils/memory.py": "cosmos_framework/model/generator/utils/memory.py", + "projects/cosmos3/cosmos3/models/utils/safetensors_loader.py": "cosmos_framework/model/generator/utils/safetensors_loader.py", + "projects/cosmos3/cosmos3/models/utils/safetensors_loader_test.py": "cosmos_framework/model/generator/utils/safetensors_loader_test.py", + "projects/cosmos3/cosmos3/models/utils/taylorseer.py": "cosmos_framework/model/generator/utils/taylorseer.py", + "projects/cosmos3/cosmos3/models/vlm_model.py": "cosmos_framework/model/generator/vlm_model.py", + "projects/cosmos3/cosmos3/processors/__init__.py": "cosmos_framework/data/generator/processors/__init__.py", + "projects/cosmos3/cosmos3/processors/base.py": "cosmos_framework/data/generator/processors/base.py", + "projects/cosmos3/cosmos3/processors/nemotron3densevl_processor.py": "cosmos_framework/data/generator/processors/nemotron3densevl_processor.py", + "projects/cosmos3/cosmos3/processors/nemotronvl_processor.py": "cosmos_framework/data/generator/processors/nemotronvl_processor.py", + "projects/cosmos3/cosmos3/processors/qwen3vl_processor.py": "cosmos_framework/data/generator/processors/qwen3vl_processor.py", + "projects/cosmos3/cosmos3/sequence_packing/__init__.py": "cosmos_framework/data/generator/sequence_packing/__init__.py", + "projects/cosmos3/cosmos3/sequence_packing/modalities.py": "cosmos_framework/data/generator/sequence_packing/modalities.py", + "projects/cosmos3/cosmos3/sequence_packing/mrope.py": "cosmos_framework/data/generator/sequence_packing/mrope.py", + "projects/cosmos3/cosmos3/sequence_packing/natten.py": "cosmos_framework/data/generator/sequence_packing/natten.py", + "projects/cosmos3/cosmos3/sequence_packing/packers.py": "cosmos_framework/data/generator/sequence_packing/packers.py", + "projects/cosmos3/cosmos3/sequence_packing/runtime.py": "cosmos_framework/data/generator/sequence_packing/runtime.py", + "projects/cosmos3/cosmos3/sequence_packing/temporal_causal.py": "cosmos_framework/data/generator/sequence_packing/temporal_causal.py", + "projects/cosmos3/cosmos3/sequence_packing/types.py": "cosmos_framework/data/generator/sequence_packing/types.py", + "projects/cosmos3/cosmos3/tokenizers/audio/__init__.py": "cosmos_framework/model/generator/tokenizers/audio/__init__.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae.py": "cosmos_framework/model/generator/tokenizers/audio/avae.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/__init__.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/__init__.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/activations.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/activations.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/alias_free_torch/__init__.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/__init__.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/alias_free_torch/act.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/act.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/alias_free_torch/filter.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/filter.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/alias_free_torch/resample.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/alias_free_torch/resample.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/bottlenecks.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/bottlenecks.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/env.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/env.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/models.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/models.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/modules.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/modules.py", + "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/modules_encodec.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/modules_encodec.py", + "projects/cosmos3/cosmos3/tokenizers/dc_ae/__init__.py": "cosmos_framework/model/generator/tokenizers/dc_ae/__init__.py", + "projects/cosmos3/cosmos3/tokenizers/dc_ae/convert_pt_to_distcp.py": "cosmos_framework/model/generator/tokenizers/dc_ae/convert_pt_to_distcp.py", + "projects/cosmos3/cosmos3/tokenizers/dc_ae/dc_ae_4x32x32.py": "cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_4x32x32.py", + "projects/cosmos3/cosmos3/tokenizers/dc_ae/dc_ae_v.py": "cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v.py", + "projects/cosmos3/cosmos3/tokenizers/dc_ae/dc_ae_v_ops.py": "cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v_ops.py", + "projects/cosmos3/cosmos3/tokenizers/dc_ae/dc_ae_v_triton_rms_norm.py": "cosmos_framework/model/generator/tokenizers/dc_ae/dc_ae_v_triton_rms_norm.py", + "projects/cosmos3/cosmos3/tokenizers/flux_vae_8x8.py": "cosmos_framework/model/generator/tokenizers/flux_vae_8x8.py", + "projects/cosmos3/cosmos3/tokenizers/interface.py": "cosmos_framework/model/generator/tokenizers/interface.py", + "projects/cosmos3/cosmos3/tokenizers/stable_diffusion_vae_8x8.py": "cosmos_framework/model/generator/tokenizers/stable_diffusion_vae_8x8.py", + "projects/cosmos3/cosmos3/tokenizers/tokenization_qwen2.py": "cosmos_framework/model/generator/tokenizers/tokenization_qwen2.py", + "projects/cosmos3/cosmos3/tokenizers/uniae/__init__.py": "cosmos_framework/model/generator/tokenizers/uniae/__init__.py", + "projects/cosmos3/cosmos3/tokenizers/uniae/frame_math.py": "cosmos_framework/model/generator/tokenizers/uniae/frame_math.py", + "projects/cosmos3/cosmos3/tokenizers/uniae/frame_math_test.py": "cosmos_framework/model/generator/tokenizers/uniae/frame_math_test.py", + "projects/cosmos3/cosmos3/tokenizers/uniae/noncausal_4x16x16.py": "cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py", + "projects/cosmos3/cosmos3/tokenizers/wan2pt1_vae_4x8x8.py": "cosmos_framework/model/generator/tokenizers/wan2pt1_vae_4x8x8.py", + "projects/cosmos3/cosmos3/tokenizers/wan2pt2_vae_4x16x16.py": "cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py", + "projects/cosmos3/cosmos3/upsampler/__init__.py": "cosmos_framework/model/generator/upsampler/__init__.py", + "projects/cosmos3/cosmos3/upsampler/prompts.py": "cosmos_framework/model/generator/upsampler/prompts.py", + "projects/cosmos3/cosmos3/utils/data_utils.py": "cosmos_framework/utils/generator/data_utils.py", + "projects/cosmos3/cosmos3/utils/dtensor_helper.py": "cosmos_framework/utils/generator/dtensor_helper.py", + "projects/cosmos3/cosmos3/utils/flash_attn.py": "cosmos_framework/utils/generator/flash_attn.py", + "projects/cosmos3/cosmos3/utils/fused_adam.py": "cosmos_framework/utils/generator/fused_adam.py", + "projects/cosmos3/cosmos3/utils/hf_attention_cosmos.py": "cosmos_framework/utils/generator/hf_attention_cosmos.py", + "projects/cosmos3/cosmos3/utils/lora.py": "cosmos_framework/utils/generator/lora.py", + "projects/cosmos3/cosmos3/utils/model_loader.py": "cosmos_framework/utils/generator/model_loader.py", + "projects/cosmos3/cosmos3/utils/model_weights_stats.py": "cosmos_framework/utils/generator/model_weights_stats.py", + "projects/cosmos3/cosmos3/utils/monkey_patch.py": "cosmos_framework/utils/generator/monkey_patch.py", + "projects/cosmos3/cosmos3/utils/optimizer.py": "cosmos_framework/utils/generator/optimizer.py", + "projects/cosmos3/cosmos3/utils/parallelism.py": "cosmos_framework/utils/generator/parallelism.py", + "projects/cosmos3/cosmos3/utils/rand_state.py": "cosmos_framework/utils/generator/rand_state.py", + "projects/cosmos3/cosmos3/utils/reasoner/__init__.py": "cosmos_framework/utils/generator/reasoner/__init__.py", + "projects/cosmos3/cosmos3/utils/reasoner/constant.py": "cosmos_framework/utils/generator/reasoner/constant.py", + "projects/cosmos3/cosmos3/utils/reasoner/create_position_ids.py": "cosmos_framework/utils/generator/reasoner/create_position_ids.py", + "projects/cosmos3/cosmos3/utils/reasoner/flop_calculator.py": "cosmos_framework/utils/generator/reasoner/flop_calculator.py", + "projects/cosmos3/cosmos3/utils/reasoner/pretrained_models_downloader.py": "cosmos_framework/utils/generator/reasoner/pretrained_models_downloader.py", + "projects/cosmos3/cosmos3/utils/video_preprocess.py": "cosmos_framework/utils/generator/video_preprocess.py", + "projects/cosmos3/tokenizer/evaluation/metric.py": "cosmos_framework/model/tokenizer/evaluation/metric.py", + "projects/cosmos3/tokenizer/evaluation/reconstruction_metrics.py": "cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py", + "projects/cosmos3/tokenizer/models/__init__.py": "cosmos_framework/model/tokenizer/models/__init__.py", + "projects/cosmos3/tokenizer/models/dense_backends.py": "cosmos_framework/model/tokenizer/models/dense_backends.py", + "projects/cosmos3/tokenizer/models/dense_runtime.py": "cosmos_framework/model/tokenizer/models/dense_runtime.py", + "projects/cosmos3/tokenizer/models/modules/__init__.py": "cosmos_framework/model/tokenizer/models/modules/__init__.py", + "projects/cosmos3/tokenizer/models/modules/attention/__init__.py": "cosmos_framework/model/tokenizer/models/modules/attention/__init__.py", + "projects/cosmos3/tokenizer/models/modules/attention/full_attn.py": "cosmos_framework/model/tokenizer/models/modules/attention/full_attn.py", + "projects/cosmos3/tokenizer/models/modules/attention/modules.py": "cosmos_framework/model/tokenizer/models/modules/attention/modules.py", + "projects/cosmos3/tokenizer/models/modules/quantizers/__init__.py": "cosmos_framework/model/tokenizer/models/modules/quantizers/__init__.py", + "projects/cosmos3/tokenizer/models/modules/quantizers/fsq.py": "cosmos_framework/model/tokenizer/models/modules/quantizers/fsq.py", + "projects/cosmos3/tokenizer/models/modules/quantizers/lfq.py": "cosmos_framework/model/tokenizer/models/modules/quantizers/lfq.py", + "projects/cosmos3/tokenizer/models/modules/quantizers/residual_vq.py": "cosmos_framework/model/tokenizer/models/modules/quantizers/residual_vq.py", + "projects/cosmos3/tokenizer/models/modules/sparse_ops.py": "cosmos_framework/model/tokenizer/models/modules/sparse_ops.py", + "projects/cosmos3/tokenizer/models/modules/sparse_tensor.py": "cosmos_framework/model/tokenizer/models/modules/sparse_tensor.py", + "projects/cosmos3/tokenizer/models/modules/transformer/__init__.py": "cosmos_framework/model/tokenizer/models/modules/transformer/__init__.py", + "projects/cosmos3/tokenizer/models/modules/transformer/blocks.py": "cosmos_framework/model/tokenizer/models/modules/transformer/blocks.py", + "projects/cosmos3/tokenizer/models/modules/transformer/modulated.py": "cosmos_framework/model/tokenizer/models/modules/transformer/modulated.py", + "projects/cosmos3/tokenizer/models/sparse_autoencoder.py": "cosmos_framework/model/tokenizer/models/sparse_autoencoder.py", + "projects/cosmos3/tokenizer/models/text_decoder.py": "cosmos_framework/model/tokenizer/models/text_decoder.py", + "projects/cosmos3/tokenizer/models/utils.py": "cosmos_framework/model/tokenizer/models/utils.py", + "projects/cosmos3/tokenizer/utils/hf.py": "cosmos_framework/model/tokenizer/utils/hf.py", + "projects/cosmos3/tokenizer/utils/vlm_prompt_format.py": "cosmos_framework/model/tokenizer/utils/vlm_prompt_format.py", + "projects/cosmos3/utils/torchcodec_video.py": "cosmos_framework/utils/generator/torchcodec_video.py", + "projects/cosmos3/vlm/configs/base/defaults/checkpointer.py": "cosmos_framework/utils/reasoner/configs_defaults/checkpointer.py", + "projects/cosmos3/vlm/processors/nemotron3densevl_processor.py": "cosmos_framework/data/reasoner/processors/nemotron3densevl_processor.py", + "projects/cosmos3/vlm/processors/nemotronvl_processor.py": "cosmos_framework/data/reasoner/processors/nemotronvl_processor.py", + "projects/cosmos3/vlm/processors/qwen3vl_processor.py": "cosmos_framework/data/reasoner/processors/qwen3vl_processor.py", + "projects/cosmos3/vlm/utils/__init__.py": "cosmos_framework/utils/reasoner/__init__.py", + "projects/cosmos3/vlm/utils/compute_flops_qwen3vl.py": "cosmos_framework/utils/reasoner/compute_flops_qwen3vl.py", + "projects/cosmos3/vlm/utils/constant.py": "cosmos_framework/utils/reasoner/constant.py", + "projects/cosmos3/vlm/utils/create_position_ids.py": "cosmos_framework/utils/reasoner/create_position_ids.py", + "projects/cosmos3/vlm/utils/dcp_checkpointer.py": "cosmos_framework/utils/reasoner/dcp_checkpointer.py", + "projects/cosmos3/vlm/utils/distributed.py": "cosmos_framework/utils/reasoner/distributed.py", + "projects/cosmos3/vlm/utils/flop_calculator.py": "cosmos_framework/utils/reasoner/flop_calculator.py", + "projects/cosmos3/vlm/utils/fused_adam.py": "cosmos_framework/utils/reasoner/fused_adam.py", + "projects/cosmos3/vlm/utils/model_wrapper.py": "cosmos_framework/utils/reasoner/model_wrapper.py", + "projects/cosmos3/vlm/utils/optimizer.py": "cosmos_framework/utils/reasoner/optimizer.py", + "projects/cosmos3/vlm/utils/planner.py": "cosmos_framework/utils/reasoner/planner.py", + "projects/cosmos3/vlm/utils/pretrained_models_downloader.py": "cosmos_framework/utils/reasoner/pretrained_models_downloader.py", + "projects/cosmos3/vlm/utils/video_preprocess.py": "cosmos_framework/utils/reasoner/video_preprocess.py" + } +} diff --git a/cosmos_framework/callbacks/every_n_draw_sample.py b/cosmos_framework/callbacks/every_n_draw_sample.py index 7d181d6f..fa87fda5 100644 --- a/cosmos_framework/callbacks/every_n_draw_sample.py +++ b/cosmos_framework/callbacks/every_n_draw_sample.py @@ -4,8 +4,9 @@ import math import os from contextlib import nullcontext +from dataclasses import dataclass from functools import partial -from typing import List, Optional +from typing import Any import numpy as np import torch @@ -23,6 +24,8 @@ from cosmos_framework.tools.visualize.video import save_img_or_video from cosmos_framework.utils.generator.data_utils import slice_data_batch +WandbImagePaths = str | dict[str, str] + def resize_image(image: torch.Tensor, size: int = 1024) -> torch.Tensor: """ @@ -49,7 +52,7 @@ def convert_to_primitive(value): return "non-primitive" # Skip non-primitive types -def pad_images_and_cat(images: List[torch.Tensor], max_w: int, max_h: int, t_crop: int = 1) -> torch.Tensor: +def pad_images_and_cat(images: list[torch.Tensor], max_w: int, max_h: int, t_crop: int = 1) -> torch.Tensor: """ Pad images to a common size and concatenate them along the batch dimension. @@ -83,6 +86,137 @@ def pad_images_and_cat(images: List[torch.Tensor], max_w: int, max_h: int, t_cro return torch.cat(padded_images, dim=0) # [total_B,C,T,max_h,max_w] (total_B = sum of batch dims) +@dataclass(frozen=True, slots=True) +class MultiviewTransferMetadata: + num_vision_items: int + sample_n_views: int + num_video_frames_per_view: int + + +@dataclass(frozen=True, slots=True) +class MultiviewTransferSampleResult: + handled: bool + image_paths: WandbImagePaths | None = None + + +def _flatten_int_metadata(value: Any) -> list[int] | None: + if value is None: + return None + if isinstance(value, torch.Tensor): + flat_value = value.detach().cpu().flatten() # [N_metadata] + return [int(item) for item in flat_value.tolist()] + if isinstance(value, (list, tuple)): + values: list[int] = [] + for item in value: + if isinstance(item, torch.Tensor): + flat_item = item.detach().cpu().flatten() # [N_metadata] + values.extend(int(v) for v in flat_item.tolist()) + elif isinstance(item, (int, np.integer)): + values.append(int(item)) + else: + return None + return values + if isinstance(value, (int, np.integer)): + return [int(value)] + return None + + +def _first_positive_metadata_value(data_batch: dict[str, Any], key: str) -> int | None: + values = _flatten_int_metadata(data_batch.get(key)) + if not values: + return None + value = values[0] + return value if value > 0 else None + + +def _get_multiview_transfer_metadata( + data_batch: dict[str, Any], + num_vision_items_per_sample: Any, +) -> MultiviewTransferMetadata | None: + if "sample_n_views" not in data_batch or "num_video_frames_per_view" not in data_batch: + return None + + num_items = _flatten_int_metadata(num_vision_items_per_sample) + if not num_items: + return None + + sample_n_views = _first_positive_metadata_value(data_batch, "sample_n_views") + num_video_frames_per_view = _first_positive_metadata_value(data_batch, "num_video_frames_per_view") + if num_items[0] < 2 or sample_n_views is None or num_video_frames_per_view is None: + return None + + return MultiviewTransferMetadata( + num_vision_items=num_items[0], + sample_n_views=sample_n_views, + num_video_frames_per_view=num_video_frames_per_view, + ) + + +def _split_multiview_tensor_by_view( + tensor: torch.Tensor, + sample_n_views: int, + num_video_frames_per_view: int, +) -> torch.Tensor | None: + if tensor.dim() == 5: + if tensor.shape[0] != 1: + return None + view_tensor = tensor[0] # [C,V*F,H,W] + elif tensor.dim() == 4: + view_tensor = tensor # [C,V*F,H,W] + else: + return None + + expected_num_frames = sample_n_views * num_video_frames_per_view + if view_tensor.shape[1] != expected_num_frames: + return None + + view_tensor = view_tensor.reshape( + view_tensor.shape[0], + sample_n_views, + num_video_frames_per_view, + view_tensor.shape[2], + view_tensor.shape[3], + ) # [C,V,F,H,W] + return view_tensor.permute(1, 0, 2, 3, 4).contiguous() # [V,C,F,H,W] + + +def _get_first_multiview_transfer_rows( + raw_data: list[torch.Tensor] | None, + metadata: MultiviewTransferMetadata, +) -> tuple[torch.Tensor, torch.Tensor] | None: + if raw_data is None or len(raw_data) < metadata.num_vision_items: + return None + + control_video = _split_multiview_tensor_by_view( + raw_data[0], + metadata.sample_n_views, + metadata.num_video_frames_per_view, + ) # [V,C,F,H,W] or None + gt_target_video = _split_multiview_tensor_by_view( + raw_data[metadata.num_vision_items - 1], + metadata.sample_n_views, + metadata.num_video_frames_per_view, + ) # [V,C,F,H,W] or None + if control_video is None or gt_target_video is None: + return None + return control_video, gt_target_video + + +def _add_wandb_image_paths( + info: dict[str, Any], + key_prefix: str, + image_paths: WandbImagePaths | None, + caption: str, +) -> None: + if image_paths is None: + return + if isinstance(image_paths, dict): + for key_suffix, image_path in image_paths.items(): + info[f"{key_prefix}_{key_suffix}"] = wandb.Image(image_path, caption=caption) + return + info[key_prefix] = wandb.Image(image_paths, caption=caption) + + class EveryNDrawSample(EveryN): """ This callback sample condition inputs from training data, run inference and save the results to wandb and s3. @@ -93,7 +227,7 @@ class EveryNDrawSample(EveryN): n_viz_sample (int, optional): for each batch, min(n_viz_sample, batch_size) samples will be saved to wandb. Defaults to 3. n_sample_to_save (int, optional): number of samples to save. The actual number of samples to save is min(n_sample_to_save, data parallel instances). Defaults to 128. num_sampling_step (int, optional): number of sampling steps. Defaults to 35. - guidance (List[float], optional): guidance scale. Defaults to [0.0, 3.0, 7.0]. + guidance (list[float], optional): guidance scale. Defaults to [0.0, 3.0, 7.0]. do_x0_prediction (bool, optional): whether to do x0 prediction. Defaults to True. n_sigmas_for_x0_prediction (int, optional): number of sigmas to use for x0 prediction. Defaults to 4. save_s3 (bool, optional): whether to save to s3. Defaults to False. @@ -109,7 +243,7 @@ def __init__( n_viz_sample: int = 2, n_sample_to_save: int = 128, num_sampling_step: int = 35, - guidance: List[float] = [0.0, 3.0, 7.0], + guidance: list[float] = [0.0, 3.0, 7.0], do_x0_prediction: bool = True, n_sigmas_for_x0_prediction: int = 4, save_s3: bool = False, @@ -119,9 +253,9 @@ def __init__( prompt_type: str = "t5_xxl", fps: int = 16, run_at_start: bool = False, - ): + ) -> None: # s3: # files: min(n_sample_to_save, data instance) # per file: min(batch_size, n_viz_sample) - # wandb: 1 file, # per file: min(batch_size, n_viz_sample) + # wandb: normal paths log one preview; multiview transfer logs one preview per selected timestamp. super().__init__(every_n, step_size, run_at_start=run_at_start) self.n_viz_sample = n_viz_sample @@ -196,7 +330,15 @@ def x0_pred(self, trainer, model, data_batch, output_batch, loss, iteration): return local_path, torch.tensor(mse_loss_list).cuda(), sigmas # [N_sigmas] @torch.no_grad() - def every_n_impl(self, trainer, model, data_batch, output_batch, loss, iteration): + def every_n_impl( + self, + trainer: Any, + model: Any, + data_batch: dict[str, Any], + output_batch: Any, + loss: Any, + iteration: int, + ) -> None: if self.is_ema: if not model.config.ema.enabled: return @@ -266,20 +408,85 @@ def every_n_impl(self, trainer, model, data_batch, output_batch, loss, iteration "sample_counter": sample_counter, } if self.do_x0_prediction: - info[f"{self.name}/{tag}_x0"] = wandb.Image(x0_img_fp, caption=f"{sample_counter}") + _add_wandb_image_paths(info, f"{self.name}/{tag}_x0", x0_img_fp, f"{sample_counter}") # convert mse_loss to a dict mse_loss = mse_loss.tolist() info.update({f"x0_pred_mse_{tag}/Sigma{sigmas[i]:0.5f}": mse_loss[i] for i in range(len(mse_loss))}) - info[f"{self.name}/{tag}_sample"] = wandb.Image(sample_img_fp, caption=f"{sample_counter}") + _add_wandb_image_paths(info, f"{self.name}/{tag}_sample", sample_img_fp, f"{sample_counter}") wandb.log( info, step=iteration, ) torch.cuda.empty_cache() + def _sample_multiview_transfer( + self, + model: Any, + data_batch: dict[str, Any], + raw_data: list[torch.Tensor] | None, + metadata: MultiviewTransferMetadata, + iteration: int, + tag: str, + ) -> MultiviewTransferSampleResult: + first_sample_rows = _get_first_multiview_transfer_rows(raw_data, metadata) + if first_sample_rows is None: + return MultiviewTransferSampleResult(handled=False) + + control_video, gt_target_video = first_sample_rows + to_show = [ + control_video.float().cpu(), # [V,C,F,H,W] + gt_target_video.float().cpu(), # [V,C,F,H,W] + ] + + generation_batch = slice_data_batch(data_batch, start=0, limit=1) + for guidance in self.guidance: + sample = model.generate_samples_from_batch( + generation_batch, + guidance=guidance, + n_sample=1, + num_steps=self.num_sampling_step, + has_negative_prompt=True if self.use_negative_prompt else False, + seed=[iteration], + ) + sample_vision = sample["vision"] + if len(sample_vision) != 1: + return MultiviewTransferSampleResult(handled=True) + assert hasattr(model, "decode") + generated_video = model.decode(sample_vision[0]) # [1,C,V*F,H,W] or [C,V*F,H,W] + generated_by_view = _split_multiview_tensor_by_view( + generated_video, + metadata.sample_n_views, + metadata.num_video_frames_per_view, + ) # [V,C,F,H,W] or None + if generated_by_view is None: + return MultiviewTransferSampleResult(handled=True) + to_show.append(generated_by_view.float().cpu()) # [V,C,F,H,W] + + if any(row.shape != to_show[0].shape for row in to_show): + return MultiviewTransferSampleResult(handled=True) + + base_fp_wo_ext = f"{tag}_ReplicateID{self.data_parallel_id:04d}_Sample_Iter{iteration:09d}" + base_fp_wo_ext = f"Iter{iteration:09d}/{base_fp_wo_ext}" + image_paths = self.run_save( + to_show, + metadata.sample_n_views, + base_fp_wo_ext, + max_columns=metadata.sample_n_views, + split_video_frames_for_wandb=True, + ) + return MultiviewTransferSampleResult(handled=True, image_paths=image_paths) + @misc.timer("EveryNDrawSample: sample") - def sample(self, trainer, model, data_batch, output_batch, loss, iteration): + def sample( + self, + trainer: Any, + model: Any, + data_batch: dict[str, Any], + output_batch: Any, + loss: Any, + iteration: int, + ) -> WandbImagePaths | None: data_batch = slice_data_batch(data_batch, start=0, limit=self.n_viz_sample) tag = "ema" if self.is_ema else "reg" @@ -303,6 +510,18 @@ def sample(self, trainer, model, data_batch, output_batch, loss, iteration): # Check if this is a multi-item vision batch (image editing) num_items = data_clean.num_vision_items_per_sample is_multi_item = num_items is not None + multiview_metadata = _get_multiview_transfer_metadata(data_batch, num_items) + if multiview_metadata is not None: + multiview_result = self._sample_multiview_transfer( + model, + data_batch, + raw_data, + multiview_metadata, + iteration, + tag, + ) + if multiview_result.handled: + return multiview_result.image_paths if is_multi_item: # Image editing: raw_data is flat [src1, tgt1, src2, tgt2, ...]. @@ -365,11 +584,19 @@ def sample(self, trainer, model, data_batch, output_batch, loss, iteration): local_path = self.run_save(to_show, batch_size, base_fp_wo_ext) return local_path - def run_save(self, to_show, batch_size, base_fp_wo_ext) -> Optional[str]: + def run_save( + self, + to_show: list[torch.Tensor], + batch_size: int, + base_fp_wo_ext: str, + max_columns: int | None = None, + split_video_frames_for_wandb: bool = False, + ) -> WandbImagePaths | None: to_show = (1.0 + torch.stack(to_show, dim=0).clamp(-1, 1)) / 2.0 # [N_rows,B,C,T,H,W] range [0,1] is_single_frame = to_show.shape[3] == 1 - n_viz_sample = min(self.n_viz_sample, batch_size) - to_show = to_show[:, :n_viz_sample] + max_columns = self.n_viz_sample if max_columns is None else max_columns + n_columns = min(max_columns, batch_size) + to_show = to_show[:, :n_columns] # ! we only save first n_sample_to_save video! video_grid = rearrange(to_show, "n b c t h w -> c t (n h) (b w)") # [C,T,N_rows*H,B*W] @@ -390,33 +617,64 @@ def run_save(self, to_show, batch_size, base_fp_wo_ext) -> Optional[str]: if self.rank == 0 and wandb.run: if is_single_frame: # image case to_show = rearrange( - to_show[:, :n_viz_sample], + to_show[:, :n_columns], "n b c t h w -> t c (n h) (b w)", ) # [1,C,N_rows*H,B*W] (t=1 for images) - image_grid = torchvision.utils.make_grid(to_show, nrow=1, padding=0, normalize=False) + image_grid = torchvision.utils.make_grid( + to_show, nrow=1, padding=0, normalize=False + ) # [C,N_rows*H,B*W] # resize so that wandb can handle it os.makedirs(os.path.dirname(local_path), exist_ok=True) - torchvision.utils.save_image(resize_image(image_grid, 1024), local_path, nrow=1, scale_each=True) + resized_image = resize_image(image_grid, 1024) # [C,H_resize,W_resize] + torchvision.utils.save_image(resized_image, local_path, nrow=1, scale_each=True) else: - to_show = to_show[:, :n_viz_sample] # [N_rows,B,C,T,H,W] + to_show = to_show[:, :n_columns] # [N_rows,B,C,T,H,W] # resize 3 frames frames so that we can display them on wandb _T = to_show.shape[3] three_frames_list = [0, _T // 2, _T - 1] - to_show = to_show[:, :, :, three_frames_list] # [N_rows,B,C,3,H,W] (3 sampled frames) log_image_size = 1024 + os.makedirs(os.path.dirname(local_path), exist_ok=True) + + if split_video_frames_for_wandb: + frame_paths: dict[str, str] = {} + for frame_name, frame_idx in zip( + ("frame_first", "frame_mid", "frame_last"), + three_frames_list, + strict=True, + ): + frame_to_show = to_show[:, :, :, frame_idx] # [N_rows,B,C,H,W] + frame_to_show = rearrange( + frame_to_show, + "n b c h w -> 1 c (n h) (b w)", + ) # [1,C,N_rows*H,B*W] + frame_image_grid = torchvision.utils.make_grid( + frame_to_show, + nrow=1, + padding=0, + normalize=False, + ) # [C,N_rows*H,B*W] + frame_path = f"{self.local_dir}/{base_fp_wo_ext}_{frame_name}_resize.jpg" + resized_frame = resize_image(frame_image_grid, log_image_size) # [C,H_resize,W_resize] + torchvision.utils.save_image(resized_frame, frame_path, nrow=1, scale_each=True) + frame_paths[frame_name] = frame_path + return frame_paths + + to_show = to_show[:, :, :, three_frames_list] # [N_rows,B,C,3,H,W] (3 sampled frames) to_show = rearrange( to_show, "n b c t h w -> 1 c (n h) (b t w)", ) # [1,C,N_rows*H,B*3*W] (t=3 sampled frames) - os.makedirs(os.path.dirname(local_path), exist_ok=True) # resize so that wandb can handle it - image_grid = torchvision.utils.make_grid(to_show, nrow=1, padding=0, normalize=False) - os.makedirs(os.path.dirname(local_path), exist_ok=True) - torchvision.utils.save_image( - resize_image(image_grid, log_image_size), local_path, nrow=1, scale_each=True - ) + image_grid = torchvision.utils.make_grid( + to_show, + nrow=1, + padding=0, + normalize=False, + ) # [C,N_rows*H,B*3*W] + resized_image = resize_image(image_grid, log_image_size) # [C,H_resize,W_resize] + torchvision.utils.save_image(resized_image, local_path, nrow=1, scale_each=True) return local_path return None diff --git a/cosmos_framework/callbacks/grad_clip.py b/cosmos_framework/callbacks/grad_clip.py index f3cb4fa3..c6a8f166 100644 --- a/cosmos_framework/callbacks/grad_clip.py +++ b/cosmos_framework/callbacks/grad_clip.py @@ -12,15 +12,17 @@ from cosmos_framework.utils.callback import Callback -@torch.compile def _fused_nan_to_num(grads: list[torch.Tensor]) -> None: """Replace NaN/Inf entries with 0.0 in every floating-point grad in-place. - The Python-level loop over ``grads`` is wrapped in ``@torch.compile`` so - Inductor can fuse the per-tensor ``nan_to_num`` ops into a single CUDA - kernel. This is NOT the ``torch._foreach_*`` API; it is fusion-via-compile - and depends on the grad list structure (length, dtypes, shapes) staying - stable across iterations so Dynamo can reuse its specialized graph. + Runs eager, NOT ``@torch.compile``. Compiling this generates a GPU-only Triton + ``nan_to_num`` kernel, which crashes whenever any grad in the list is a CPU tensor (some + parameters carry a small CPU grad): the static CUDA launcher launches it with ``stream=0`` + -> ``CUDA driver error: invalid argument``, and the standard launcher raises ``Pointer + argument cannot be accessed from Triton (cpu tensor?)``. Both were observed at 720 under + replay-TF + torch.compile (grad-clip runs at the optimizer step, after the large compiled + graphs are cached). Eager ``torch.nan_to_num`` handles CPU and CUDA grads alike; the fusion win + from compiling a handful of per-tensor ops once per step is negligible next to that fragility. """ grads = [g for g in grads if torch.is_floating_point(g)] for g in grads: diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py index d0d47ca0..5e6fd6b4 100644 --- a/cosmos_framework/configs/base/defaults/model_config.py +++ b/cosmos_framework/configs/base/defaults/model_config.py @@ -118,8 +118,8 @@ class FixedStepSamplerConfig: # Convention: exclude the final 0.0 step — FixedStepSampler appends it automatically. # Values must be descending. Using 0.999 instead of 1.0 avoids numeric edge cases at sigma=1. t_list: list[float] = [0.999, 0.75, 0.5, 0.25] - # Integrator type: "ode" (deterministic Euler) or "sde" (stochastic re-noising at each step). - sample_type: str = "ode" + # Distilled fixed-step sampling uses stochastic re-noising at each step. + sample_type: str = "sde" # Don't have any defaults and init only in config file. diff --git a/cosmos_framework/configs/base/reasoner/config.py b/cosmos_framework/configs/base/reasoner/config.py index 1a1455d6..0cb1e4db 100644 --- a/cosmos_framework/configs/base/reasoner/config.py +++ b/cosmos_framework/configs/base/reasoner/config.py @@ -7,6 +7,7 @@ from cosmos_framework.configs.base.defaults.checkpointer import register_checkpoint, register_ckpt_type from cosmos_framework.configs.base.reasoner.defaults.callbacks import register_callbacks from cosmos_framework.configs.base.reasoner.defaults.config import Config + from cosmos_framework.configs.base.reasoner.defaults.model import register_model from cosmos_framework.configs.base.reasoner.defaults.optimizer import register_optimizer, register_scheduler from cosmos_framework.configs.base.reasoner.defaults.vlm_policy import register_vlm_policy diff --git a/cosmos_framework/data/generator/augmentors/text_transforms_for_image.py b/cosmos_framework/data/generator/augmentors/text_transforms_for_image.py index 95462393..694b5964 100644 --- a/cosmos_framework/data/generator/augmentors/text_transforms_for_image.py +++ b/cosmos_framework/data/generator/augmentors/text_transforms_for_image.py @@ -9,7 +9,6 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -# COSMOS-RELEASE-END-IGNORE # For the qwen captions, we have 3 variants: short, medium, long # In addition, for synthetic data, we create prompt embeddings as well. diff --git a/cosmos_framework/data/generator/utils.py b/cosmos_framework/data/generator/utils.py index 8b04f0a9..bba5b3de 100644 --- a/cosmos_framework/data/generator/utils.py +++ b/cosmos_framework/data/generator/utils.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 +import math import re from typing import List, Tuple @@ -89,12 +90,12 @@ def get_wdinfos_w_aspect_ratio(wdinfos: list[str]) -> List[Tuple[str, str]]: return [(wdinfo, aspect_ratio.replace("_", ",")) for wdinfo, aspect_ratio in zip(wdinfos, aspect_ratios)] -def parse_frame_range_from_wdinfo(wdinfo: str) -> tuple[int, int] | None: +def parse_frame_range_from_wdinfo(wdinfo: str) -> tuple[int, int | float] | None: """ Parse frame range from wdinfo path. Args: - wdinfo: wdinfo path string containing frames_X_Y pattern + wdinfo: wdinfo path string containing a frames_X_Y pattern, where Y may be ``inf`` Returns: Tuple of (min_frames, max_frames) if found, None otherwise @@ -102,25 +103,29 @@ def parse_frame_range_from_wdinfo(wdinfo: str) -> tuple[int, int] | None: Example: >>> parse_frame_range_from_wdinfo("wdinfo/v4/tv_drama/resolution_720/aspect_ratio_16_9/frames_300_400/wdinfo.json") (300, 400) + >>> parse_frame_range_from_wdinfo("wdinfo/v4/tv_drama/resolution_720/aspect_ratio_16_9/frames_3700_inf/wdinfo.json") + (3700, inf) """ - match = re.search(r"frames_(\d+)_(\d+)", wdinfo) + match = re.search(r"frames_(\d+)_(\d+|inf)(?!\w)", wdinfo) if match: - return (int(match.group(1)), int(match.group(2))) + max_frames = math.inf if match.group(2) == "inf" else int(match.group(2)) + return (int(match.group(1)), max_frames) return None def _normalize_skip_frame_ranges( skip_frame_range: str | list[str] | None, -) -> set[tuple[int, int]]: +) -> set[tuple[int, int | float]]: """Normalize ``skip_frame_range`` into a set of (min_frames, max_frames) buckets. Args: - skip_frame_range: A single bucket string like ``"300_400"``, a list of such - strings, or None. Each string identifies the frame-range bucket - (e.g. ``frames_300_400``) that should be skipped. + skip_frame_range: A single bucket string like ``"300_400"`` or ``"3700_inf"``, + a list of such strings, or None. Each string identifies the frame-range + bucket (e.g. ``frames_300_400``) that should be skipped. Returns: - Set of (min_frames, max_frames) tuples to skip. Empty if ``skip_frame_range`` is None. + Set of (min_frames, max_frames) tuples to skip. An ``inf`` upper bound is + represented by ``math.inf``. Empty if ``skip_frame_range`` is None. """ if skip_frame_range is None: return set() @@ -128,14 +133,16 @@ def _normalize_skip_frame_ranges( if isinstance(skip_frame_range, str): skip_frame_range = [skip_frame_range] - skip_buckets: set[tuple[int, int]] = set() + skip_buckets: set[tuple[int, int | float]] = set() for bucket in skip_frame_range: - match = re.fullmatch(r"(\d+)_(\d+)", bucket.strip()) + match = re.fullmatch(r"(\d+)_(\d+|inf)", bucket.strip()) if match is None: raise ValueError( - f"Invalid skip_frame_range entry {bucket!r}. Expected the form '_', e.g. '300_400'." + f"Invalid skip_frame_range entry {bucket!r}. " + "Expected the form '_', e.g. '300_400' or '3700_inf'." ) - skip_buckets.add((int(match.group(1)), int(match.group(2)))) + max_frames = math.inf if match.group(2) == "inf" else int(match.group(2)) + skip_buckets.add((int(match.group(1)), max_frames)) return skip_buckets @@ -154,6 +161,7 @@ def filter_wdinfos_by_frame_range( based on the wdinfo's upper bound (wdinfo_max): - min_frames is EXCLUSIVE: wdinfo_max must be > min_frames - max_frames is INCLUSIVE: wdinfo_max must be <= max_frames + - an ``inf`` upper bound is treated as infinity and excluded by any finite max_frames Additionally, any wdinfo whose frame-range bucket matches an entry in ``skip_frame_range`` is excluded. @@ -163,8 +171,9 @@ def filter_wdinfos_by_frame_range( min_frames: Minimum number of frames (exclusive). If None, no lower bound. max_frames: Maximum number of frames (inclusive). If None, no upper bound. skip_frame_range: Frame-range bucket(s) to exclude, e.g. ``"300_400"`` to - drop the ``frames_300_400`` bucket. Accepts a single string or a list - of strings. If None, no bucket is skipped. + drop the ``frames_300_400`` bucket or ``"3700_inf"`` to drop the + ``frames_3700_inf`` bucket. Accepts a single string or a list of strings. + If None, no bucket is skipped. Returns: Filtered list of wdinfo paths diff --git a/cosmos_framework/data/reasoner/processors/nemotronvl_processor.py b/cosmos_framework/data/reasoner/processors/nemotronvl_processor.py index 6a4d8ce2..bcb2e2c2 100644 --- a/cosmos_framework/data/reasoner/processors/nemotronvl_processor.py +++ b/cosmos_framework/data/reasoner/processors/nemotronvl_processor.py @@ -84,15 +84,15 @@ {{- ', ' if not loop.last else '' -}} {%- endfor -%} {{- ']\n\n' -}} - + {{- 'If you decide to call any tool(s), use the following format:\n' -}} {{- '[{"name": "tool_name1", "arguments": "tool_args1"}, ' -}} {{- '{"name": "tool_name2", "arguments": "tool_args2"}]\n\n' -}} - + {{- 'The user will execute tool-calls and return responses from tool(s) in this format:\n' -}} {{- '[{"response": "tool_response1"}, ' -}} {{- '{"response": "tool_response2"}]\n\n' -}} - + {{- 'Based on the tool responses, you can call additional tools if needed, ' -}} {{- 'correct tool calls if any errors are found, or just respond to the user.' -}} {%- endif -%} diff --git a/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py b/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py index 714268e2..54472146 100644 --- a/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py +++ b/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py @@ -4,18 +4,12 @@ """Fixed-step sampler for DMD2-distilled student models. Uses an explicit, fixed sigma schedule (t_list) baked in at construction time. -Each step predicts x0 via a single velocity forward pass, then either: - - ODE: Euler step x_next = x_t + (sigma_next - sigma_cur) * v - - SDE: re-noise x0 to sigma_next with fresh noise +Each step predicts x0 via a single velocity forward pass, then re-noises x0 to +sigma_next with fresh noise. This is incompatible with multi-step solvers (UniPC, EDM) because DMD2 students are trained as one-shot denoisers at specific discrete sigmas, not as smooth score functions. - -When ``shift`` is passed at call time, the schedule is derived dynamically via -the flow-matching shift formula (same as UniPC): - sigmas = shift * s / (1 + (shift - 1) * s), s = linspace(sigma_max, sigma_min, num_steps) -In this case ``num_steps`` is required. Otherwise ``self.t_list`` is used. """ import torch @@ -27,25 +21,17 @@ class FixedStepSampler: def __init__( self, t_list: list[float], - sample_type: str = "ode", + sample_type: str = "sde", num_train_timesteps: float = 1000.0, ) -> None: assert len(t_list) >= 1, "t_list must have at least 1 entry" - assert sample_type in ("ode", "sde"), f"sample_type must be 'ode' or 'sde', got {sample_type}" + assert sample_type == "sde", f"FixedStepSampler only supports sample_type='sde', got {sample_type}" # Auto-append 0.0 if not present (convention: t_list in config excludes final step) self.t_list = t_list if t_list[-1] == 0.0 else t_list + [0.0] assert len(self.t_list) >= 2, "t_list must have at least 2 entries after appending 0.0" self.sample_type = sample_type self.num_train_timesteps = num_train_timesteps - def _build_t_list(self, num_steps: int, shift: float, device: torch.device) -> list[float]: - """Compute a shifted sigma schedule with ``num_steps`` integration steps.""" - sigma_max = 1.0 - sigma_min = 1.0 / self.num_train_timesteps - sigmas = torch.linspace(sigma_max, sigma_min, num_steps, device=device) - sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) - return sigmas.tolist() + [0.0] - def __call__( self, velocity_fn, @@ -72,13 +58,12 @@ def __call__( noise: Initial noise. Either a single ``torch.Tensor`` of shape ``(D,)`` or a ``list[torch.Tensor]`` where each element has shape ``(D,)``. - seed: RNG seed for SDE mode. Either a single ``int`` or a + seed: RNG seed. Either a single ``int`` or a ``list[int]`` with the same length as ``noise``. - num_steps: Number of denoising steps. Required when ``shift`` is - given; optional otherwise (asserted to equal - ``len(t_list) - 1`` when provided). - shift: When set, derive the sigma schedule dynamically using the - flow-matching shift formula instead of ``self.t_list``. + num_steps: Ignored. The number of denoising steps is defined by + ``self.t_list``. + shift: Ignored. Fixed-step distilled sampling always uses + ``self.t_list`` from the experiment config. condition_reference: Optional clean reference tensor(s) to preserve where ``condition_mask`` is 1. condition_mask: Optional mask tensor(s), same shape as ``noise``, @@ -111,15 +96,7 @@ def __call__( ) assert not isinstance(condition_mask, list), "condition_mask must not be a list when noise is a tensor" - if shift is not None: - assert num_steps is not None, "num_steps is required when shift is provided" - t_list = self._build_t_list(num_steps, shift, device) - else: - if num_steps is not None: - assert num_steps == len(self.t_list) - 1, ( - f"num_steps={num_steps} must match the schedule length len(t_list)-1={len(self.t_list) - 1}" - ) - t_list = self.t_list + t_list = self.t_list latent = noise @@ -139,14 +116,10 @@ def _sde_step( x0_pred = latent - sigma_cur * v_pred # [...,D] if sigma_next > 0: - if self.sample_type == "ode": - # Euler ODE step - latent_next = latent + (sigma_next - sigma_cur) * v_pred # [...,D] - else: - if seed is not None: - torch.manual_seed(seed + step_idx) - eps_fresh = torch.randn_like(x0_pred) # [...,D] - latent_next = (1.0 - sigma_next) * x0_pred + sigma_next * eps_fresh # [...,D] + if seed is not None: + torch.manual_seed(seed + step_idx) + eps_fresh = torch.randn_like(x0_pred) # [...,D] + latent_next = (1.0 - sigma_next) * x0_pred + sigma_next * eps_fresh # [...,D] else: latent_next = x0_pred # [...,D] diff --git a/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py b/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py index a6aff60e..d2d612ed 100644 --- a/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py +++ b/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py @@ -22,61 +22,6 @@ def test_no_double_append_zero(): assert sampler.t_list.count(0.0) == 1 -@pytest.mark.L0 -def test_ode_single_step(): - # 1-step ODE schedule [1.0, 0.0], v=0 => x0 = noise - 1*0 = noise, last step => x = x0_pred = noise - sampler = FixedStepSampler(t_list=[1.0, 0.0], sample_type="ode") - noise = torch.randn(1, 8) - - def zero_vel(x, _t): - return torch.zeros_like(x) - - result = sampler(zero_vel, noise) - assert torch.allclose(result, noise) - - -@pytest.mark.L0 -def test_ode_multi_step_correctness(): - # 2-step ODE: t_list=[1.0, 0.5, 0.0] - # Step 0: sigma_cur=1.0, sigma_next=0.5 - # v_pred = v0, x0_pred = noise - 1.0 * v0 - # Euler: x1 = noise + (0.5 - 1.0) * v0 = noise - 0.5 * v0 - # Step 1: sigma_cur=0.5, sigma_next=0.0 - # v_pred = v1, x0_pred = x1 - 0.5 * v1 - # last step (sigma_next==0): x = x0_pred = x1 - 0.5 * v1 - noise = torch.tensor([[1.0, 2.0, 3.0]]) - v0 = torch.tensor([[0.1, 0.2, 0.3]]) - v1 = torch.tensor([[0.4, 0.5, 0.6]]) - call_count = [0] - - def velocity_fn(x, t): - i = call_count[0] - call_count[0] += 1 - return v0 if i == 0 else v1 - - sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.0], sample_type="ode") - result = sampler(velocity_fn, noise) - - x1 = noise - 0.5 * v0 - expected = x1 - 0.5 * v1 - assert torch.allclose(result, expected, atol=1e-6) - - -@pytest.mark.L0 -def test_sde_differs_from_ode(): - sampler_ode = FixedStepSampler(t_list=[1.0, 0.5, 0.0], sample_type="ode") - sampler_sde = FixedStepSampler(t_list=[1.0, 0.5, 0.0], sample_type="sde") - noise = torch.randn(1, 16) - # Use non-zero velocity so paths diverge - - def velocity_fn(x, _t): - return torch.ones_like(x) * 0.1 - - result_ode = sampler_ode(velocity_fn, noise.clone()) - result_sde = sampler_sde(velocity_fn, noise.clone(), seed=42) - assert not torch.allclose(result_ode, result_sde) - - @pytest.mark.L0 def test_sde_reproducible_with_seed(): sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.0], sample_type="sde") @@ -110,25 +55,6 @@ def velocity_fn(x, _t): assert torch.allclose(result[0][0], condition_reference[0][0]) -@pytest.mark.L0 -def test_ode_preserves_conditioned_entries(): - sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.0], sample_type="ode") - noise = torch.tensor([5.0, 1.0]) # [N] - condition_reference = torch.tensor([5.0, 0.0]) # [N] - condition_mask = torch.tensor([1.0, 0.0]) # [N] - - def velocity_fn(x, _t): - return torch.ones_like(x) # [N] - - result = sampler( - velocity_fn, - noise, - condition_reference=condition_reference, - condition_mask=condition_mask, - ) - assert torch.allclose(result[0], condition_reference[0]) - - @pytest.mark.L0 def test_output_shape_preserved(): sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.25, 0.0]) @@ -159,6 +85,8 @@ def test_validates_single_entry_t_list(): def test_validates_sample_type(): with pytest.raises(AssertionError): FixedStepSampler(t_list=[1.0, 0.0], sample_type="invalid") + with pytest.raises(AssertionError): + FixedStepSampler(t_list=[1.0, 0.0], sample_type="ode") @pytest.mark.L0 @@ -178,90 +106,7 @@ def velocity_fn(x, _t): @pytest.mark.L0 -def test_num_steps_assertion_passes_when_correct(): - sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.0]) # 2 steps - noise = torch.randn(1, 4) - - def velocity_fn(x, _t): - return torch.zeros_like(x) - - result = sampler(velocity_fn, noise, num_steps=2) - assert result.shape == noise.shape - - -@pytest.mark.L0 -def test_num_steps_assertion_fails_when_wrong(): - sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.0]) # 2 steps - noise = torch.randn(1, 4) - - def velocity_fn(x, _t): - return torch.zeros_like(x) - - with pytest.raises(AssertionError): - sampler(velocity_fn, noise, num_steps=4) - - -@pytest.mark.L0 -def test_shift_requires_num_steps(): - sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.0]) - noise = torch.randn(1, 4) - - def velocity_fn(x, _t): - return torch.zeros_like(x) - - with pytest.raises(AssertionError, match="num_steps is required"): - sampler(velocity_fn, noise, shift=2.0) - - -@pytest.mark.L0 -def test_shift_one_gives_uniform_schedule(): - # shift=1.0 is identity: sigmas = 1*s / (1 + 0*s) = s (unchanged) - num_steps = 4 - sampler = FixedStepSampler(t_list=[0.9, 0.6, 0.3, 0.0], num_train_timesteps=1000.0) - noise = torch.zeros(1, 4) - recorded_sigmas = [] - - def velocity_fn(x, timestep): - recorded_sigmas.append(timestep.item() / 1000.0) - return torch.zeros_like(x) - - sampler(velocity_fn, noise, num_steps=num_steps, shift=1.0) - - sigma_max = 1.0 - sigma_min = 1.0 / 1000.0 - expected = torch.linspace(sigma_max, sigma_min, num_steps).tolist() - assert len(recorded_sigmas) == num_steps - for got, exp in zip(recorded_sigmas, expected): - assert abs(got - exp) < 1e-4, f"got {got}, expected {exp}" - - -@pytest.mark.L0 -def test_shift_greater_than_one_frontloads_schedule(): - # shift > 1 pushes sigmas higher (more steps near sigma_max) - num_steps = 8 - sampler = FixedStepSampler(t_list=[1.0] + [0.0] * (num_steps - 1) + [0.0]) - noise = torch.zeros(1, 4) - recorded_sigmas_uniform = [] - recorded_sigmas_shifted = [] - - def make_recorder(lst): - def velocity_fn(x, timestep): - lst.append(timestep.item() / 1000.0) - return torch.zeros_like(x) - - return velocity_fn - - sampler(make_recorder(recorded_sigmas_uniform), noise.clone(), num_steps=num_steps, shift=1.0) - sampler(make_recorder(recorded_sigmas_shifted), noise.clone(), num_steps=num_steps, shift=7.0) - - # With shift>1, all intermediate sigmas should be larger than uniform - for uniform, shifted in zip(recorded_sigmas_uniform[1:], recorded_sigmas_shifted[1:]): - assert shifted > uniform, f"expected shifted {shifted} > uniform {uniform}" - - -@pytest.mark.L0 -def test_shift_none_uses_t_list(): - # When shift is None, self.t_list is used (not a dynamic schedule) +def test_num_steps_and_shift_do_not_override_t_list(): t_list = [0.9, 0.5, 0.1, 0.0] sampler = FixedStepSampler(t_list=t_list) noise = torch.zeros(1, 4) @@ -271,7 +116,7 @@ def velocity_fn(x, timestep): recorded_sigmas.append(timestep.item() / 1000.0) return torch.zeros_like(x) - sampler(velocity_fn, noise) + sampler(velocity_fn, noise, num_steps=8, shift=7.0) expected = t_list[:-1] # last step uses sigma=0.0 but loop already excludes it assert len(recorded_sigmas) == len(expected) for got, exp in zip(recorded_sigmas, expected): diff --git a/cosmos_framework/model/generator/mot/context_parallel_test.py b/cosmos_framework/model/generator/mot/context_parallel_test.py index 3114b244..7f3dce0a 100644 --- a/cosmos_framework/model/generator/mot/context_parallel_test.py +++ b/cosmos_framework/model/generator/mot/context_parallel_test.py @@ -41,6 +41,7 @@ from cosmos_framework.utils.generator.parallelism import ParallelDims + def _broadcast_test_object(data: Any, parallel_dims: ParallelDims, iteration: int) -> Any: rank = parallel_dims.cp_rank cp_world_size = parallel_dims.cp_mesh.size() diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index d3a544c4..656364b2 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -2503,6 +2503,13 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool): return v_pred + # Run sampler for all samples at once. + sampler = sampler or self.sampler + scheduler_type = self.config.rectified_flow_inference_config.scheduler_type + if isinstance(sampler, FixedStepSampler): + num_steps = len(sampler.t_list) - 1 + shift = 0.0 + # FSDP collective-sequence alignment (sampler outer loop). See the # large block above the velocity_fn definition for the full # rationale. all_reduce on the local num_steps so every rank knows @@ -2516,9 +2523,6 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool): _max_num_steps = num_steps _extra_num_steps = _max_num_steps - num_steps - # Run sampler for all samples at once. - sampler = sampler or self.sampler - scheduler_type = self.config.rectified_flow_inference_config.scheduler_type if isinstance(sampler, FixedStepSampler): log.info(f"Using sampler: FixedStep (t_list={sampler.t_list}, sample_type={sampler.sample_type})") elif scheduler_type == "unipc": diff --git a/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py b/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py index 4eeb0657..ab5afa89 100644 --- a/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py +++ b/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py @@ -495,6 +495,5 @@ def reset(self) -> None: "SSIMMetric", "LPIPSMetric", "Rank0FIDMetric", - "FVDMetric", "calculate_psnr", ] From 8ea43183588961648a8fd8cf33866c849eeaa21d Mon Sep 17 00:00:00 2001 From: LiangHao Date: Tue, 7 Jul 2026 20:13:35 +0800 Subject: [PATCH 16/26] DROID action-policy post-training: lazy LeRobot dataset + Cosmos3-Nano recipe + numerical regression test (#86) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds **DROID action-policy post-training** on Cosmos3-Nano and a **numerical regression test** for the action-policy launches — the reference reproduction recipe for the DROID policy result. ## What's included - **Lazy DROID LeRobot dataset** — `cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py` (streaming `BaseActionLeRobotDataset`) + a rewritten `droid_lerobot_dataset.py` on top of it, plus `droid_lerobot_dataset_config.py`. Keys the versioned merged root; `use_success_only` resolves the `success/` split; eager `_register_sources()`. - **DROID Nano recipe** — `configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py` + `examples/toml/sft_config/action_policy_droid_repro.toml`: res480, `joint_pos` 8D + `use_state`, JSON action prompt (`format_prompt_as_json=True`), CPU-side color jitter. The TOML pins the GB200 reference shape — HSDP 32×8 (256 ranks), global batch 8192, lr 2e-4, 10000 iters — and trains the generation + action heads from the public Cosmos3-Nano base. - **net_ema warm-start** — `checkpoint/dcp.py`: seed `net_ema` from `net` when `net_ema` is skipped on load (fresh action heads). - **Action-policy numerical regression test** — `tests/action_policy_regression_test.py`: the action-policy analogue of `tests/launch_regression_test.py`. Deterministic 10-iter re-run of the LIBERO + DROID launches (single-node, `--deterministic`, seed 42), asserting per-arch rank-0 loss goldens with a tolerance. LIBERO golden captured on the H200 CI arch; the DROID spec skips unless its (large, out-of-CI) dataset is supplied via `DROID_ROOT`. ## Reproduction The recipe reproduces the DROID action-policy result from the public Cosmos3-Nano base. The exact training code was validated by a from-scratch 64-GPU (GB200) run to 10k iterations, and the recipe TOML now encodes that run's configuration directly. ## Companion Cookbook PR: NVIDIA/cosmos#261 — the runnable DROID/LIBERO finetune cookbook that drives this recipe. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: lfengad --- cosmos_framework/checkpoint/dcp.py | 12 +- .../configs/base/defaults/model_config.py | 2 +- .../action_policy_droid_nano.py | 14 +- .../action/datasets/action_sft_dataset.py | 19 +- .../generator/action/datasets/base_dataset.py | 42 +- .../action/datasets/cosmos3_action_lerobot.py | 1048 +++++++++++++++++ .../action/datasets/droid_lerobot_dataset.py | 734 +++++++----- .../datasets/droid_lerobot_dataset_config.py | 100 ++ .../sft_config/action_policy_droid_repro.toml | 47 +- tests/action_policy_regression_test.py | 466 ++++++++ 10 files changed, 2140 insertions(+), 344 deletions(-) create mode 100644 cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py create mode 100644 cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_config.py create mode 100644 tests/action_policy_regression_test.py diff --git a/cosmos_framework/checkpoint/dcp.py b/cosmos_framework/checkpoint/dcp.py index 43145c88..54d0d8e1 100644 --- a/cosmos_framework/checkpoint/dcp.py +++ b/cosmos_framework/checkpoint/dcp.py @@ -71,9 +71,9 @@ from cosmos_framework.checkpoint.base import AbstractCheckpointer from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter -from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.utils import callback, distributed, log, misc +from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.utils.easy_io import easy_io from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict @@ -866,6 +866,16 @@ def load( raise ValueError( f"Unexpected keys (found in checkpoint but not in model): {results.unexpected_keys}" ) + # Warm start that skipped net_ema (e.g. loading an EMA-only HF export + # with no net_ema.* keys): the EMA shadow would otherwise keep its random + # build-time generation pathway (init_moe is skipped when a checkpoint is + # present). Seed net_ema from the freshly loaded net so the EMA starts equal + # to net ("EMA warm-starts from net") instead of from random weights. + if warm_start and any("net_ema" in skip_key for skip_key in keys_to_skip_loading): + ema_worker = getattr(model, "net_ema_worker", None) + if ema_worker is not None and getattr(model, "net_ema", None) is not None: + ema_worker.copy_to(src_model=model.net, tgt_model=model.net_ema) + log.info("Warm start: re-seeded net_ema from net (net_ema was skipped on load).") elif key == "optim": log.info("- Loading the optimizer...") diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py index 5e6fd6b4..adde6bd9 100644 --- a/cosmos_framework/configs/base/defaults/model_config.py +++ b/cosmos_framework/configs/base/defaults/model_config.py @@ -5,12 +5,12 @@ import attrs -from cosmos_framework.utils.lazy_config import LazyDict from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.ema import EMAConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig from cosmos_framework.configs.base.defaults.reasoner import VLMConfig +from cosmos_framework.utils.lazy_config import LazyDict @attrs.define(slots=False) diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py index 926afea9..6ce66fb9 100644 --- a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py @@ -21,15 +21,14 @@ from hydra.core.config_store import ConfigStore -from cosmos_framework.utils.lazy_config import LazyCall as L -from cosmos_framework.utils.lazy_config import LazyDict - from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG +from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_droid_sft_dataset from cosmos_framework.data.generator.joint_dataloader import ( PackingDataLoader, RankPartitionedDataLoader, ) -from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_droid_sft_dataset +from cosmos_framework.utils.lazy_config import LazyCall as L +from cosmos_framework.utils.lazy_config import LazyDict cs = ConfigStore.instance() @@ -202,7 +201,9 @@ use_state=True, iterable_shuffle=True, # rank x worker episode-shuffle stream episode_shuffle_seed=42, - use_image_augmentation=True, # SR boost (random crop+rescale + color jitter) + # SR boost: random crop+rescale + ColorJitter, applied CPU-side in the + # DROIDLeRobotDataset image augmentor (matches i4's pipeline stage). + use_image_augmentation=True, # keep_ranges_1_0_1.json window filter (drops idle/non-task frames). Off by default; # set use_filter_dict=True + filter_dict_path to enable. use_filter_dict=False, @@ -213,6 +214,9 @@ max_action_dim="${model.config.max_action_dim}", cfg_dropout_rate=0.1, tokenizer_config="${model.config.vlm_config.tokenizer}", + # Match i4 GA (droid_lerobot_8b_ga / MR #9995): format the action + # prompt as JSON via ActionPromptJsonFormatter instead of plain text. + format_prompt_as_json=True, ), ), ), diff --git a/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py b/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py index 8a424bc9..6053259c 100644 --- a/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py @@ -19,7 +19,9 @@ from torch.utils.data import Dataset, IterableDataset, get_worker_info -from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import ( + DROIDLeRobotDataset, +) from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset from cosmos_framework.data.generator.action.transforms import ActionTransformPipeline @@ -109,13 +111,17 @@ def get_action_droid_sft_dataset( append_duration_fps_timestamps: bool = True, append_resolution_info: bool = True, append_idle_frames: bool = False, + format_prompt_as_json: bool = False, iterable_shuffle: bool = False, episode_shuffle_seed: int = 42, + use_success_only: bool = True, ) -> Dataset: """Build the DROID action SFT dataset: ``action_space='joint_pos'`` (8D) + - ``use_state`` (raw/un-normalized), concat_view, chunk_length 32.""" - dataset = DROIDLeRobotDataset( - root=root, + ``use_state`` (raw/un-normalized), concat_view, chunk_length 32. + + Reads ``root`` (a merged/versioned DROID LeRobot root) as a single flat + dataset; ``use_success_only=True`` filters to the ``success/`` split.""" + shard_kwargs = dict( fps=fps, chunk_length=chunk_length, viewpoint=viewpoint, @@ -123,10 +129,12 @@ def get_action_droid_sft_dataset( mode=mode, use_state=use_state, action_normalization=action_normalization, - use_image_augmentation=use_image_augmentation, + use_image_augmentation=use_image_augmentation, # i4: bundles random-crop+resize+ColorJitter use_filter_dict=use_filter_dict, filter_dict_path=filter_dict_path, + use_success_only=use_success_only, ) + dataset: Dataset = DROIDLeRobotDataset(root=root, **shard_kwargs) transform = ActionTransformPipeline( tokenizer_config=tokenizer_config, cfg_dropout_rate=cfg_dropout_rate, @@ -135,6 +143,7 @@ def get_action_droid_sft_dataset( append_duration_fps_timestamps=append_duration_fps_timestamps, append_resolution_info=append_resolution_info, append_idle_frames=append_idle_frames, + format_prompt_as_json=format_prompt_as_json, ) sft = ActionSFTDataset(dataset, transform, resolution) if iterable_shuffle: diff --git a/cosmos_framework/data/generator/action/datasets/base_dataset.py b/cosmos_framework/data/generator/action/datasets/base_dataset.py index e3d2ec81..ff0792ca 100644 --- a/cosmos_framework/data/generator/action/datasets/base_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/base_dataset.py @@ -74,18 +74,13 @@ def __init__( # LeRobot v2.x stores task text in a "task" column; v3.0 stores it as the # (unnamed) DataFrame index and keeps only "task_index" as a column. task_texts = tasks_df["task"] if "task" in tasks_df.columns else tasks_df.index - self._tasks = { - int(task_index): str(task) - for task, task_index in zip(task_texts, tasks_df["task_index"]) - } - self._rows = sorted( - ( - row - for path in sorted((self._root / "data").glob("chunk-*/file-*.parquet")) - for row in pq.read_table(path).to_pylist() - ), - key=lambda row: int(row["index"]), - ) + self._tasks = {int(task_index): str(task) for task, task_index in zip(task_texts, tasks_df["task_index"])} + # ``self._rows`` (the flat, index-sorted list of every frame dict) is built + # lazily on first access — see the ``_rows`` property. Materializing all + # ~18M frames as Python dicts plus a full sort costs ~13 min and tens of GB; + # subclasses that build their own compact index (e.g. DROIDLeRobotDataset) + # never touch it, so they must not pay for it at construction. + self._rows_cache: list[dict[str, Any]] | None = None @property def fps(self) -> float: @@ -140,8 +135,7 @@ def _stats_path(cls) -> Path: def load_action_stats(cls) -> dict[str, torch.Tensor]: """Return action normalization stats for this dataset as torch tensors.""" return { - key: torch.from_numpy(value).float() - for key, value in load_action_stats(str(cls._stats_path())).items() + key: torch.from_numpy(value).float() for key, value in load_action_stats(str(cls._stats_path())).items() } @abstractmethod @@ -218,5 +212,25 @@ def _build_result( **extras, } + @property + def _rows(self) -> list[dict[str, Any]]: + """Flat, index-sorted list of every frame dict, built lazily on first access. + + Only datasets that don't build their own compact index (bridge / agibot / + robomind) touch this; for them it materializes once and caches. Datasets with + a bespoke index (e.g. DROIDLeRobotDataset) never read it, so they skip the + ~13 min / tens-of-GB construction entirely. + """ + if self._rows_cache is None: + self._rows_cache = sorted( + ( + row + for path in sorted((self._root / "data").glob("chunk-*/file-*.parquet")) + for row in pq.read_table(path).to_pylist() + ), + key=lambda row: int(row["index"]), + ) + return self._rows_cache + def __len__(self) -> int: return max(0, (len(self._rows) - self._chunk_length + self._sample_stride - 1) // self._sample_stride) diff --git a/cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py b/cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py new file mode 100644 index 00000000..d74b2aee --- /dev/null +++ b/cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py @@ -0,0 +1,1048 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared LeRobot adapter utilities for Action datasets. + +These helpers centralize common behavior across Action wrappers: +- deterministic train/val episode splitting +- valid per-episode index range construction +- a reusable BaseActionLeRobotDataset class with lazy init, video formatting, + and common result building +""" + +from __future__ import annotations + +import importlib +import logging as _logging +import math +import os as _os +import random +from bisect import bisect_right +from collections import OrderedDict, defaultdict +from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from threading import Lock +from typing import Any, ClassVar + +import huggingface_hub.constants as _hf_const +import numpy as np +import torch +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from torch.utils.data import Dataset + +_hf_offline_applied = False + + +def _ensure_hf_hub_offline() -> None: + """Force HF Hub into offline mode for local-only datasets (repo_id="local"). + + Sets the ``HF_HUB_OFFLINE`` env var (for any future imports in worker + processes), patches the already-imported constant, and suppresses the + expected "Returning existing local_dir" fallback warning. + + Safe to call multiple times; only applies once per process. + """ + global _hf_offline_applied + if _hf_offline_applied: + return + if "HF_HUB_OFFLINE" not in _os.environ: + _os.environ["HF_HUB_OFFLINE"] = "1" + if not _hf_const.HF_HUB_OFFLINE: + _hf_const.HF_HUB_OFFLINE = True + _logging.getLogger("huggingface_hub._snapshot_download").setLevel(_logging.ERROR) + _hf_offline_applied = True + + +from functools import cached_property + +from cosmos_framework.data.generator.action.action_processing import ( + ActionNormalizationMethod, + ActionNormalizer, + load_action_stats, + resolve_action_normalization, +) + +# Re-export the action_spec DSL from this module so that subclass datasets +# only need a single import block (alongside ``BaseActionLeRobotDataset``). +from cosmos_framework.data.generator.action.action_spec import ( # noqa: F401 (re-export) + ActionSpec, + DimType, + Gripper, + Joint, + Pos, + Reserved, + Rot, + build_action_spec, +) +from cosmos_framework.data.generator.action.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.pose_utils import compute_idle_frames +from cosmos_framework.data.generator.action.viewpoint_utils import Viewpoint +from cosmos_framework.utils import log + + +# No-op memory-profiling shims. Profiling is disabled in cosmos-framework, so these +# keep the ported dataset's optional RSS-tracking call sites cheap and side-effect-free +# (formerly a separate memprofile stub module). +def _memprofile_enabled() -> bool: + return False + + +def _deep_size(obj, *args, **kwargs) -> int: + return 0 + + +def _fmt_mb(n, *args, **kwargs) -> str: + return "n/a" + + +def log_worker_memory_breakdown(*args, **kwargs) -> None: + return None + + +@contextmanager +def rss_tracker(*args, **kwargs): + yield + + +# --------------------------------------------------------------------------- +# LRU-capped VideoDecoderCache +# --------------------------------------------------------------------------- +_LRU_VIDEO_CACHE_MAX_SIZE: int = 64 +_LRU_DATASET_MAX_LOADED: int = 32 +ActionNormalization = ActionNormalizationMethod +_ACTION_NORMALIZATION_CHOICES: tuple[str, ...] = ("quantile", "quantile_rot", "meanstd", "minmax") + +_decoder_cache_patched = False + + +class _LRUVideoDecoderCache: + """Drop-in replacement for ``lerobot.datasets.video_utils.VideoDecoderCache`` + with LRU eviction. When the cache exceeds *max_size* entries the + least-recently-used decoder (and its file handle) is evicted. + """ + + def __init__(self, max_size: int = _LRU_VIDEO_CACHE_MAX_SIZE) -> None: + self._max_size = max_size + self._cache: OrderedDict[str, tuple[Any, Any]] = OrderedDict() + self._lock = Lock() + self._hits = 0 + self._misses = 0 + self._evictions = 0 + + def get_decoder(self, video_path: str) -> Any: + if importlib.util.find_spec("torchcodec"): # type: ignore[attr-defined] + from torchcodec.decoders import VideoDecoder + else: + raise ImportError("torchcodec is required but not available.") + + import fsspec + + video_path = str(video_path) + + with self._lock: + if video_path in self._cache: + self._cache.move_to_end(video_path) + self._hits += 1 + return self._cache[video_path][0] + + self._misses += 1 + file_handle = fsspec.open(video_path).__enter__() + decoder = VideoDecoder(file_handle, seek_mode="approximate") # type: ignore[arg-type] + self._cache[video_path] = (decoder, file_handle) + + evicted = 0 + while len(self._cache) > self._max_size: + _, (_, old_fh) = self._cache.popitem(last=False) + try: + old_fh.close() + except Exception: + pass + evicted += 1 + self._evictions += evicted + + if evicted and self._evictions % 50 <= evicted: + log.debug( + f"[VideoDecoderCache pid={_os.getpid()}] " + f"evicted={self._evictions} total, size={len(self._cache)}/{self._max_size}, " + f"hits={self._hits}, misses={self._misses}, " + f"hit_rate={100 * self._hits / max(1, self._hits + self._misses):.1f}%" + ) + + return decoder + + def clear(self) -> None: + with self._lock: + for _, file_handle in self._cache.values(): + try: + file_handle.close() + except Exception: + pass + self._cache.clear() + + def size(self) -> int: + with self._lock: + return len(self._cache) + + +def _patch_decoder_cache(max_size: int = _LRU_VIDEO_CACHE_MAX_SIZE) -> None: + """Replace the module-level ``_default_decoder_cache`` in LeRobot with an + LRU-capped version to prevent unbounded memory growth in workers.""" + global _decoder_cache_patched + if _decoder_cache_patched: + return + + import lerobot.datasets.video_utils as _vu + + lru_cache = _LRUVideoDecoderCache(max_size=max_size) + _vu._default_decoder_cache = lru_cache + _decoder_cache_patched = True + log.debug(f"Patched LeRobot VideoDecoderCache with LRU max_size={max_size}") + + +def _parallel_map( + fn: Callable[[Any], Any], + items: list[Any], + *, + max_workers: int, + label: str, +) -> list[Any]: + """Thread-pool ``map`` — returns results in input order. + + Intended for IO-bound prefetch (``LeRobotDatasetMetadata`` loads, + parquet column reads). Preserves item-order so callers can ``zip`` + with their ``indices`` / ``roots`` list. Skips the thread pool + entirely when there is 0 or 1 task — avoids per-worker + ``ThreadPoolExecutor`` setup cost and log spam under + ``shard_across_workers=True`` where each worker typically gets + only 1-2 shards. + """ + if not items: + return [] + if len(items) == 1 or max_workers <= 1: + return [fn(items[0])] if len(items) == 1 else [fn(x) for x in items] + log.info(f"{label}: {len(items)} tasks (workers={max_workers})") + with ThreadPoolExecutor(max_workers=max_workers) as ex: + return list(ex.map(fn, items)) + + +def split_episode_ids(total_episodes: int, seed: int, val_ratio: float, split: str) -> list[int]: + """Create deterministic random episode ids for train/val/full splits.""" + num_val = int(round(total_episodes * val_ratio)) + g = torch.Generator().manual_seed(seed) + episode_ids = torch.randperm(total_episodes, generator=g).tolist() + + if split == "train": + return episode_ids[num_val:] + if split == "val": + return episode_ids[:num_val] + return episode_ids + + +def build_episode_spans( + episodes: Any, + episode_ids: Sequence[int], + chunk_length: int, + sample_stride: int = 1, +) -> tuple[list[tuple[int, int, int]], int, int]: + """Build valid episode spans for LeRobot frame queries. + + Returns: + - episode spans as ``(episode_id, sample_start, valid_len)`` + - total valid sample count across selected episodes + - total raw frame count across selected episodes + """ + assert sample_stride >= 1, f"sample_stride must be >= 1, got {sample_stride}" + + dataset_from_index = list(episodes["dataset_from_index"]) + dataset_to_index = list(episodes["dataset_to_index"]) + length = list(episodes["length"]) + + spans: list[tuple[int, int, int]] = [] + valid_count = 0 + sample_count = 0 + for episode_id in episode_ids: + start = dataset_from_index[episode_id] + stop = dataset_to_index[episode_id] + raw_valid_len = stop - start - chunk_length + if raw_valid_len > 0: + valid_len = (raw_valid_len + sample_stride - 1) // sample_stride + spans.append((episode_id, start, valid_len)) + valid_count += valid_len + sample_count += int(length[episode_id]) + + return spans, valid_count, sample_count + + +def _normalize_split(split: str) -> str: + """Normalize split name to one of ``'train'``, ``'val'``, ``'full'``.""" + s = split.lower().strip() + if s in {"val", "valid", "validation", "eval", "test"}: + return "val" + if s in {"train", "full"}: + return s + raise ValueError(f"Unsupported {split=}. Use train/val/full.") + + +class BaseActionLeRobotDataset(Dataset): + """Reusable base class for Action LeRobot-backed map-style datasets. + + Subclasses typically: + 1) call ``_register_source`` to register one or more LeRobot sources + 2) implement ``__getitem__`` for dataset-specific sample parsing + 3) call ``_build_result`` to assemble the return dict + """ + + # Applied as: R_opencv = R_native @ _to_opencv + # Subclasses override in __init__; default is identity (no correction). + + # Bundled normalization stats directory. Stats are committed at + # ``<_NORMALIZERS_DIR>/__.json`` (flat + # layout matching the existing UMI files) and produced by + # ``projects/cosmos3/vfm/datasets/action/compute_action_stats.py``. + # Subclasses that need a different filename scheme can override + # :meth:`_normalizer_filename`. + _NORMALIZERS_DIR: ClassVar[Path] = Path(__file__).parent / "normalizers" + + def __init__( + self, + *, + fps: float, + chunk_length: int, + split_seed: int, + split_val_ratio: float, + split: str, + mode: str, + embodiment_type: str, + viewpoint: Viewpoint, + pose_convention: str | None = None, + rotation_format: str | None = None, + action_normalization: ActionNormalization | None = None, + tolerance_s: float = 1e-4, + max_loaded_datasets: int = _LRU_DATASET_MAX_LOADED, + skip_video_loading: bool = False, + sample_stride: int = 1, + enable_fast_init: bool = False, + fast_init_max_workers: int = 64, + min_episode_length_frames: int | None = None, + ) -> None: + super().__init__() + _ensure_hf_hub_offline() + _patch_decoder_cache() + self._memprofile = _memprofile_enabled() + + assert sample_stride >= 1, f"sample_stride must be >= 1, got {sample_stride}" + assert fast_init_max_workers >= 1, f"fast_init_max_workers must be >= 1, got {fast_init_max_workers}" + assert action_normalization is None or action_normalization in _ACTION_NORMALIZATION_CHOICES, ( + f"action_normalization must be None or one of {_ACTION_NORMALIZATION_CHOICES}, got {action_normalization!r}" + ) + + with rss_tracker(f"{self.__class__.__name__}.__init__", enabled=self._memprofile): + self._fps = fps + self._dt = 1.0 / fps + self._chunk_length = chunk_length + self._split_seed = split_seed + self._split_val_ratio = split_val_ratio + self._split = _normalize_split(split) + self._mode = mode + self._embodiment_type = embodiment_type + self._viewpoint: Viewpoint = viewpoint + self._pose_convention = pose_convention + self._rotation_format = rotation_format + self._action_normalizer: ActionNormalizer | None = None + if action_normalization is not None: + self._action_normalizer = resolve_action_normalization( + action_normalization, self._load_norm_stats(action_normalization) + ) + self._tolerance_s = tolerance_s + self._max_loaded_datasets = max_loaded_datasets + self._skip_video_loading = skip_video_loading + self._sample_stride = sample_stride + self._enable_fast_init = enable_fast_init + self._fast_init_max_workers = fast_init_max_workers + # Optional post-filter on raw episode length. When set, episodes + # whose raw frame count is below this threshold are dropped from + # ``_episode_records`` after ``build_episode_spans`` runs. Lets + # eval configs select e.g. only ``> 60s`` wall-clock episodes + # (1800 raw frames at native 30 fps) while keeping ``chunk_length`` + # (the fetch window) at a smaller value such as 900 (60 s @ fps=15). + # Subclasses that override ``_append_index_records`` are expected + # to honor this attribute themselves. + self._min_episode_length_frames: int | None = min_episode_length_frames + self._delta_timestamps: dict[str, list[float]] = {} + self._to_opencv: np.ndarray | dict[str, np.ndarray] = np.eye(3, dtype=np.float32) + + if pose_convention is None: + log.warning( + f"{self.__class__.__name__}: pose_convention is not set. " + "Consider specifying 'backward_framewise' or 'backward_anchored'." + ) + + self._datasets: list[LeRobotDataset | None] = [] + self._dataset_build_args: list[dict[str, Any] | None] = [] + self._loaded_lru: OrderedDict[int, None] = OrderedDict() + + # -- Flat index structures (populated by _append_index_records) -- + # Together these two lists form a searchable map from a flat + # global index to (dataset, row, episode, frame). One entry per + # episode span across *all* registered sources. + # + # _episode_records[i] = (ds_idx, sample_start, valid_len, episode_id) + # ds_idx – which source dataset (index into _datasets) + # sample_start – first row of this span in that dataset's table + # valid_len – number of usable frames in this span + # episode_id – the episode this span belongs to + # + # _episode_cum_ends[i] = running total of valid_len through span i + # Used for O(log N) lookup via bisect_right in _resolve_index. + self._episode_records: list[tuple[int, int, int, int]] = [] + self._episode_cum_ends: list[int] = [] + self._num_valid_indices = 0 + self._domain_id = get_domain_id(self._embodiment_type) + + # Deferred-init shard roots — a list of root paths. + # Subclasses populate this in __init__; _register_sources() + # reads _delta_timestamps and _tolerance_s from self (both + # initialised above, with _delta_timestamps overridden by + # each subclass). + # ActionUnifiedIterableDataset.assign_worker uses len() for + # round-robin shard distribution and _register_sources(indices) + # for deferred loading. When empty, shard distribution is + # skipped (every worker iterates the full dataset). + self._all_shard_roots: list[str] = [] + + # -- public properties --------------------------------------------------- + + @property + def fps(self) -> float: + return self._fps + + @property + def chunk_length(self) -> int: + return self._chunk_length + + @property + def split(self) -> str: + return self._split + + @property + def mode(self) -> str: + return self._mode + + @mode.setter + def mode(self, value: str) -> None: + self._mode = value + + @property + def domain_id(self) -> int: + return self._domain_id + + # -- source registration ------------------------------------------------- + + def _register_source( + self, + *, + delta_timestamps: dict[str, list[float]], + tolerance_s: float, + root: str | None = None, + repo_id: str = "local", + force_cache_sync: bool = False, + download_videos: bool = False, + video_backend: str | None = None, + revision: str | None = None, + dataset_label: str | None = None, + prefetched_meta: LeRobotDatasetMetadata | None = None, + ) -> LeRobotDatasetMetadata: + """Register a LeRobot dataset source lazily (metadata-only at init). + + ``prefetched_meta`` lets subclasses load metadata in a thread pool + (``LeRobotDatasetMetadata`` reads are pure I/O — ``info.json`` + + ``episodes.parquet`` + ``tasks.parquet``) and then hand the ready + object to the serial append-path below, which still manages the + order-sensitive shared state (``_datasets`` / ``_dataset_build_args`` + / ``_episode_records`` / ``_episode_cum_ends``). When ``None`` the + caller gets the original single-threaded behavior. + """ + label_str = f" [{dataset_label}]" if dataset_label else "" + cls = self.__class__.__name__ + # "local" is not a valid PEP 440 version, so LeRobot's + # is_valid_version() check skips the get_safe_version() HF API call. + if repo_id == "local" and revision is None: + revision = "local" + + with rss_tracker(f"{cls}{label_str} — metadata load", enabled=self._memprofile): + if prefetched_meta is not None: + meta = prefetched_meta + else: + meta = LeRobotDatasetMetadata( + repo_id=repo_id, + root=root, + revision=revision, + force_cache_sync=force_cache_sync, + ) + ds_idx = len(self._datasets) + self._datasets.append(None) + self._dataset_build_args.append( + { + "repo_id": repo_id, + "root": root, + "delta_timestamps": delta_timestamps, + "tolerance_s": tolerance_s, + "force_cache_sync": force_cache_sync, + "download_videos": download_videos, + "video_backend": video_backend, + "revision": revision, + } + ) + + with rss_tracker( + f"{cls}{label_str} — index records", + enabled=self._memprofile, + extras_fn=lambda: [ + f"episode_records so far: {len(self._episode_records)} entries, " + f"~{_fmt_mb(_deep_size(self._episode_records) / (1024 * 1024))}", + f"episode_cum_ends so far: {len(self._episode_cum_ends)} entries, " + f"~{_fmt_mb(_deep_size(self._episode_cum_ends) / (1024 * 1024))}", + ], + ): + self._append_index_records(meta=meta, ds_idx=ds_idx, dataset_label=dataset_label) + + return meta + + def _append_index_records( + self, + *, + meta: LeRobotDatasetMetadata, + ds_idx: int, + dataset_label: str | None = None, + ) -> None: + """Populate episode split / index records from dataset metadata.""" + episode_ids = split_episode_ids( + total_episodes=meta.total_episodes, + seed=self._split_seed, + val_ratio=self._split_val_ratio, + split=self._split, + ) + # TODO(tianweis): remove once bridge training switches to a pre-filtered mirror. + if hasattr(self, "_filter_valid_episodes"): + episode_ids = self._filter_valid_episodes(meta, episode_ids) + episode_spans, valid_count, sample_count = build_episode_spans( + episodes=meta.episodes, + episode_ids=episode_ids, + chunk_length=self._chunk_length, + sample_stride=self._sample_stride, + ) + + # Optional duration filter (see ``self._min_episode_length_frames`` + # comment in ``__init__``). Drops episodes whose raw frame count is + # below the threshold. Operates after ``build_episode_spans`` so the + # threshold is decoupled from ``chunk_length`` (which controls the + # fetch window). No-op when the attribute is None. + if self._min_episode_length_frames is not None: + length_lookup = list(meta.episodes["length"]) + before = len(episode_spans) + episode_spans = [ + (eid, ss, vl) + for (eid, ss, vl) in episode_spans + if int(length_lookup[eid]) >= self._min_episode_length_frames + ] + dropped = before - len(episode_spans) + if dropped > 0: + log.info( + f"{self.__class__.__name__}: " + f"min_episode_length_frames={self._min_episode_length_frames} " + f"dropped {dropped} / {before} chunk-eligible spans" + ) + + class_name = self.__class__.__name__ + label = f" [{dataset_label}]" if dataset_label else "" + log.info(f"{class_name}{label}: split={self._split}, num episodes={len(episode_ids)}") + if sample_count > 0: + log.info( + f"{class_name}{label}: kept {valid_count} / {sample_count} " + f"({100 * valid_count / sample_count:.2f} %) samples" + ) + + for episode_id, sample_start, valid_len in episode_spans: + self._episode_records.append((ds_idx, sample_start, valid_len, episode_id)) + self._num_valid_indices += valid_len + self._episode_cum_ends.append(self._num_valid_indices) + + # -- deferred shard registration ----------------------------------------- + + def _register_sources(self, indices: list[int] | None = None) -> None: + """Register a subset (or all) of the shard roots in ``_all_shard_roots``. + + Called by ``ActionUnifiedIterableDataset.assign_worker`` during training, + or explicitly by eval/visualization scripts after construction. + + ``_all_shard_roots`` is a list of root paths. Per-shard args that are + shared across all shards (``delta_timestamps``, ``tolerance_s``) are + taken from ``self``. Subclasses may override this for extra per-shard + setup (e.g. loading instruction segments). + + When ``enable_fast_init=True``, ``LeRobotDatasetMetadata`` (a pure-IO + read of ``info.json`` + ``episodes.parquet`` + ``tasks.parquet``) is + prefetched in a thread pool and handed to the order-sensitive + serial register loop via ``prefetched_meta=``. Shard count scales + the speedup; for single-shard datasets the two paths are + equivalent. + + Args: + indices: Which entries of ``_all_shard_roots`` to register. + ``None`` means all. + """ + if indices is None: + indices = list(range(len(self._all_shard_roots))) + if not indices: + return + + roots = [self._all_shard_roots[i] for i in indices] + + if self._enable_fast_init: + # ``_ensure_hf_hub_offline`` already ran in ``__init__`` and is + # idempotent; no need to re-invoke here. + workers = max(1, min(self._fast_init_max_workers, len(roots))) + metas: list[LeRobotDatasetMetadata | None] = _parallel_map( + lambda root: LeRobotDatasetMetadata(repo_id="local", root=root, revision="local"), + roots, + max_workers=workers, + label=f"{type(self).__name__}: LeRobotDatasetMetadata prefetch", + ) + else: + metas = [None] * len(roots) + + for root, meta in zip(roots, metas): + label = root.rsplit("/", 1)[-1] if "/" in root else root + self._register_source( + root=root, + delta_timestamps=self._delta_timestamps, + tolerance_s=self._tolerance_s, + dataset_label=label, + prefetched_meta=meta, + ) + + # -- lazy dataset access ------------------------------------------------- + + def _get_dataset(self, ds_idx: int) -> LeRobotDataset: + """Get or lazily construct the LeRobot dataset for the given source index. + + Loaded datasets are tracked with LRU ordering. When the number of + loaded datasets exceeds ``_max_loaded_datasets`` the least-recently-used + dataset is evicted (set back to ``None``) so the GC can reclaim it. + """ + ds = self._datasets[ds_idx] + if ds is not None: + self._loaded_lru.move_to_end(ds_idx) + return ds + + _ensure_hf_hub_offline() + + build_args = self._dataset_build_args[ds_idx] + if build_args is None: + raise RuntimeError(f"Missing dataset build args for dataset index {ds_idx}") + + # Evict least-recently-used datasets before loading a new one. + while len(self._loaded_lru) >= self._max_loaded_datasets: + evict_idx, _ = self._loaded_lru.popitem(last=False) + self._datasets[evict_idx] = None + + with rss_tracker( + f"[WORKER {_os.getpid()}] Lazy-loaded ds[{ds_idx}]", + enabled=self._memprofile, + extras_fn=lambda: [f"total loaded={len(self._loaded_lru)}/{len(self._datasets)}"], + ): + delta_ts = build_args["delta_timestamps"] + if self._skip_video_loading: + # Covers both LeRobot v2 (``observation.images.``) and + # v3 (``observation.image.``) video-column conventions. + delta_ts = {k: v for k, v in delta_ts.items() if not k.startswith("observation.image")} + + log.info(f"Loading shard root={build_args['root']}") + ds = LeRobotDataset( + repo_id=build_args["repo_id"], + root=build_args["root"], + delta_timestamps=delta_ts, + tolerance_s=build_args["tolerance_s"], + force_cache_sync=build_args["force_cache_sync"], + download_videos=build_args["download_videos"], + video_backend=build_args["video_backend"], + revision=build_args["revision"], + episodes=None, + ) + if self._skip_video_loading: + ds.meta.info["features"] = { + k: v for k, v in ds.meta.info["features"].items() if v.get("dtype") != "video" + } + self._datasets[ds_idx] = ds + self._loaded_lru[ds_idx] = None + + return ds + + # -- index resolution ---------------------------------------------------- + + def _resolve_index(self, idx: int) -> tuple[int, int, int, int]: + """Map a flat global index to the source dataset, row, episode, and frame. + + Multiple datasets are concatenated into a single virtual sequence. + Each episode contributes a contiguous *span* of valid frames, and + ``_episode_cum_ends[i]`` stores the running total of valid frames + through the *i*-th span. For example, with two episodes of lengths + 5 and 3 the cum-ends are ``[5, 8]``, so global index 6 falls in the + second span at offset 1. + + The lookup is O(log N) via :func:`bisect_right`. + + Returns: + dataset_idx: Which source dataset this sample belongs to. + row_idx: Row index *within* that dataset's LeRobot table. + episode_id: The episode ID for this sample. + frame_offset: Frame offset from the start of the episode span + (0-based). + + Pure index math -- no I/O or dataset access. Higher-level helpers + like :meth:`_fetch_sample` build on this. + """ + # Support negative indexing (e.g. -1 → last sample). + if idx < 0: + idx += self._num_valid_indices + if idx < 0 or idx >= self._num_valid_indices: + raise IndexError(f"{self.__class__.__name__} index {idx} out of range for size {self._num_valid_indices}") + + # _episode_cum_ends is a monotonically increasing list where entry i + # holds the cumulative number of valid frames up to and including the + # i-th episode span. bisect_right finds the first span whose + # cumulative end is strictly greater than idx, i.e. the span that + # contains idx. + # + # Example: cum_ends = [5, 8, 20] + # idx=0 -> span_idx=0 (first span, frames 0..4) + # idx=4 -> span_idx=0 + # idx=5 -> span_idx=1 (second span, frames 5..7) + # idx=8 -> span_idx=2 (third span, frames 8..19) + span_idx = bisect_right(self._episode_cum_ends, idx) + + # The global index where this span begins is the previous span's + # cumulative end (or 0 for the very first span). The frame_offset + # is how far idx is into this particular episode. + span_start = 0 if span_idx == 0 else self._episode_cum_ends[span_idx - 1] + frame_offset = idx - span_start + + # _episode_records[span_idx] stores (dataset_idx, row_start, valid_len, + # episode_id). row_start is the absolute row in the LeRobot table + # where this episode begins. With sample_stride=k, consecutive + # valid indices map to rows k apart inside the episode, so the + # effective row is row_start + frame_offset * sample_stride. + dataset_idx, row_start, _, episode_id = self._episode_records[span_idx] + row_idx = row_start + frame_offset * self._sample_stride + return dataset_idx, row_idx, episode_id, frame_offset + + def _choose_mode(self) -> str: + """Resolve the active mode for one sample request.""" + if self._mode == "joint": + return random.choice(("forward_dynamics", "inverse_dynamics", "policy")) + return self._mode + + def _fetch_sample(self, idx: int) -> tuple[str, int, int, dict[str, Any]]: + """Resolve index, pick a mode, and load the sample from the dataset. + + Returns ``(mode, dataset_idx, row_idx, sample_dict)``. + """ + mode = self._choose_mode() + dataset_idx, row_idx, _, _ = self._resolve_index(idx) + + self._getitem_count = getattr(self, "_getitem_count", 0) + 1 + profile = self._memprofile and self._getitem_count % 50 == 1 + + with rss_tracker( + f"[WORKER {_os.getpid()}] __getitem__ transient (dataset_idx={dataset_idx})", + enabled=profile, + after_fn=lambda: log_worker_memory_breakdown(self), + ): + sample = self._get_dataset(dataset_idx)[row_idx] + + if self._skip_video_loading: + sample = defaultdict(lambda: None, sample) + + return mode, dataset_idx, row_idx, sample + + # -- action normalization ------------------------------------------------ + + def _normalizer_filename(self) -> str: + """Bundled stats filename for this dataset instance. + + Default convention (matches ``compute_action_stats.py`` output): + ``[_][_].json``. + + Pose/rotation suffixes are appended only when the instance actually + has them (SE(3) pose datasets like Bridge / DROID). Joint-space + datasets — where both are ``None`` — resolve to just + ``.json``. + + Subclasses may override when the bundled filename uses a different + scheme (e.g. UMI's ``uva_umi_single_task_normalizer.json``). + """ + if not self._embodiment_type: + raise RuntimeError( + f"{self.__class__.__name__}: embodiment_type is not set; cannot resolve normalizer filename." + ) + parts = [self._embodiment_type] + if self._pose_convention: + parts.append(self._pose_convention) + if self._rotation_format: + parts.append(self._rotation_format) + return "_".join(parts) + ".json" + + def _normalizer_path(self) -> Path: + """Full path to the bundled stats JSON for this dataset.""" + return self._NORMALIZERS_DIR / self._normalizer_filename() + + def _load_norm_stats(self, action_normalization: ActionNormalization) -> dict[str, torch.Tensor]: + """Load action normalization stats for the configured normalization mode. + + Raises :class:`FileNotFoundError` if the stats file is missing. This + is intentional — silently falling back to identity normalization when + the user asked for ``quantile`` / ``quantile_rot`` / ``meanstd`` / + ``minmax`` would be a training bug. + """ + stats_key = "global_raw" if action_normalization == "quantile_rot" else "global" + raw_stats = load_action_stats(str(self._normalizer_path()), stats_key=stats_key) + return {key: torch.from_numpy(value).float() for key, value in raw_stats.items()} # dict[str,[D]] + + def get_action_normalizer( + self, + _sample: dict[str, Any] | None = None, + ) -> ActionNormalizer | None: + """Return the configured action normalizer for transform-time preprocessing.""" + return self._action_normalizer + + # -- video formatting ---------------------------------------------------- + + def _convert_video(self, video_tchw: torch.Tensor | None) -> torch.Tensor | None: + """Convert LeRobot ``(T,C,H,W)`` float video to Action ``(C,T,H,W)`` uint8. + + Args: + video_tchw: Raw floating-point video tensor in ``[0, 1]`` with + LeRobot layout, or ``None``. # [T,C,H,W] | None + + Returns: + Action-formatted video tensor, or ``None``. # [C,T,H,W] | None + """ + if self._skip_video_loading or video_tchw is None: + return None + if video_tchw.ndim != 4: + raise ValueError( + f"{self.__class__.__name__}._convert_video expected video with shape [T,C,H,W], " + f"got ndim={video_tchw.ndim}" + ) + if not torch.is_floating_point(video_tchw): + raise TypeError( + f"{self.__class__.__name__}._convert_video expected floating-point video in [0, 1], " + f"got dtype={video_tchw.dtype}" + ) + video_min = video_tchw.amin() # [] + video_max = video_tchw.amax() # [] + if video_min.item() < 0.0 or video_max.item() > 1.0: + raise ValueError( + f"{self.__class__.__name__}._convert_video expected floating-point video in [0, 1], " + f"got range=[{video_min.item():.6f}, {video_max.item():.6f}]" + ) + formatted_video = (video_tchw * 255.0).clamp(0.0, 255.0).to(torch.uint8).permute(1, 0, 2, 3) # [C,T,H,W] + return formatted_video + + # -- result building ----------------------------------------------------- + + def _build_action_spec(self) -> ActionSpec | None: + """Subclass override: declare this dataset's action layout. + + Called once per instance — the result is cached by ``self.action_spec``. + Return ``None`` to skip spec-driven idle detection; in that case + ``_compute_idle_frames`` will log a one-time warning and return + ``None`` for every sample. + """ + return None + + @cached_property + def action_spec(self) -> ActionSpec | None: + """Cached :class:`ActionSpec` from ``_build_action_spec``. + + Returns ``None`` when the subclass did not declare one; idle detection + is then skipped (with a one-time warning) until the subclass overrides + ``_build_action_spec``. + """ + return self._build_action_spec() + + @cached_property + def action_names(self) -> list[str] | None: + spec = self.action_spec + return spec.names if spec is not None else None + + # Idle-detection thresholds. Defined as **velocities** (per second) so the + # same numeric value means the same physical motion across datasets with + # different sampling rates; converted to per-frame at call time using + # ``self._fps`` via :meth:`_resolve_idle_thresholds`. + # + # Defaults: + # - ``idle_eps_t_per_sec`` = 5 mm/s (≈ 1 mm/frame at 5 Hz) + # - ``idle_eps_r_per_sec`` = 1.5°/s (geodesic, rotation-format aware) + # - ``idle_eps_g`` = 1e-2 unit gripper Δ (no fps) + # - ``idle_joint_threshold_per_sec`` = 5e-3 rad/s + # - ``idle_min_streak`` = 3 require ≥ 3 consecutive + # + # Subclasses can either override the ``*_per_sec`` attributes (preferred — + # keeps the velocity semantics) or set the corresponding ``idle_eps_*`` / + # ``idle_joint_threshold`` attribute to a non-``None`` value to bypass the + # per-fps conversion entirely (raw per-frame override). + idle_eps_t_per_sec: float = 5e-3 + idle_eps_r_per_sec: float = math.radians(1.5) + idle_eps_g: float = 1e-2 + idle_joint_threshold_per_sec: float = 5e-3 + idle_min_streak: int = 3 + + # Optional per-frame overrides. ``None`` (default) → use the ``*_per_sec`` + # attribute / fps conversion above. + idle_eps_t: float | None = None + idle_eps_r: float | None = None + idle_joint_threshold: float | None = None + + def _resolve_idle_thresholds(self) -> tuple[float, float, float, float]: + """Resolve per-frame idle thresholds for this dataset instance. + + Returns ``(eps_t, eps_r, eps_g, joint_threshold)`` in raw per-frame + units. Honours direct per-frame overrides if the subclass sets the + non-``_per_sec`` attribute; otherwise scales the ``_per_sec`` values + by ``self._fps``. + """ + fps = float(self._fps) if self._fps else 1.0 + eps_t = self.idle_eps_t if self.idle_eps_t is not None else self.idle_eps_t_per_sec / fps + eps_r = self.idle_eps_r if self.idle_eps_r is not None else self.idle_eps_r_per_sec / fps + joint_thr = ( + self.idle_joint_threshold + if self.idle_joint_threshold is not None + else self.idle_joint_threshold_per_sec / fps + ) + return float(eps_t), float(eps_r), float(self.idle_eps_g), float(joint_thr) + + def _compute_idle_frames(self, raw_action: torch.Tensor) -> torch.Tensor | None: + """Count idle frames in the *raw* (un-normalized) action chunk. + + Requires ``self.action_spec`` to be declared via ``_build_action_spec``. + Returns ``None`` when: + - ``pose_convention`` is not ``"backward_framewise"`` (TODO: extend), + - the subclass has not declared an ``ActionSpec`` (logs a one-time warning), + - the action layout does not match the declared spec. + + Detection thresholds come from the ``idle_eps_*`` class attributes + (overridable per dataset). Subclasses can also override this method + outright, or pass an explicit ``idle_frames`` integer via + ``**extras`` to :meth:`_build_result`. + """ + # TODO: currently we only support backward_framewise. Other pose + # conventions (anchored / absolute) need different idle semantics. + if self._pose_convention != "backward_framewise": + if not getattr(self, "_warned_pose_convention", False): + log.warning( + f"Dataset {self.__class__.__name__}: pose_convention=" + f"{self._pose_convention!r} is not 'backward_framewise'; " + "skipping idle-frames detection. Centralize the dataset " + "to backward_framewise to enable IdleFrames captioning." + ) + self._warned_pose_convention = True + return None + + spec = self.action_spec + if spec is None: + if not getattr(self, "_warned_no_action_spec", False): + log.warning( + f"Dataset {self.__class__.__name__} has no action spec defined; " + "skipping idle-frames detection. Override _build_action_spec() to enable it." + ) + self._warned_no_action_spec = True + return None + + eps_t, eps_r, eps_g, joint_thr = self._resolve_idle_thresholds() + try: + n = compute_idle_frames( + raw_action, + spec, + eps_t=eps_t, + eps_r=eps_r, + eps_g=eps_g, + joint_threshold=joint_thr, + min_streak=self.idle_min_streak, + ) + except (ValueError, TypeError) as e: + if not getattr(self, "_warned_action_layout", False): + log.warning( + f"Dataset {self.__class__.__name__}: action layout does " + f"not match the declared ActionSpec " + f"(action_dim={int(raw_action.shape[-1])}, " + f"spec.dim={spec.dim}); skipping idle-frames detection. " + f"Underlying error: {e}" + ) + self._warned_action_layout = True + return None + return torch.tensor(n, dtype=torch.long) + + def _build_result( + self, + *, + mode: str, + video: torch.Tensor | None, + action: torch.Tensor, + ai_caption: str, + **extras: Any, + ) -> dict[str, Any]: + """Assemble the common return dict for ``__getitem__``. + + ``video`` is expected in raw LeRobot layout before final formatting. + Subclasses may pass extra keys (e.g. ``initial_pose``) via ``**extras``. + ``idle_frames`` is auto-computed from the raw (un-normalized) ``action`` + whenever the dataset's pose/rotation conventions allow it; subclasses + can override by passing ``idle_frames`` (int or scalar tensor) via + ``**extras``. + """ + # Compute idle_frames from the raw action before normalization, unless + # the subclass has provided one explicitly via ``**extras``. + if "idle_frames" not in extras: + idle_frames = self._compute_idle_frames(action) + if idle_frames is not None: + extras = {"idle_frames": idle_frames, **extras} + + if self._skip_video_loading: + result: dict[str, Any] = {"action": action} + if "idle_frames" in extras: + result["idle_frames"] = extras["idle_frames"] + return result + formatted_video = self._convert_video(video) # [C,T,H,W] | None + return { + "ai_caption": ai_caption, + "video": formatted_video, + "action": action, + "conditioning_fps": torch.tensor(self._fps, dtype=torch.long), + "mode": mode, + "domain_id": torch.tensor(self._domain_id, dtype=torch.long), + "viewpoint": self._viewpoint, + **extras, + } + + def __len__(self) -> int: + return self._num_valid_indices + + def get_shuffle_blocks(self) -> list[tuple[int, int]]: + """Per-episode (or per kept-segment when ``use_filter_dict``) flat-index + blocks ``(start, length)`` over ``[0, len(self))``. ``ActionIterableShuffleDataset`` + shuffles the ORDER of these blocks and shards them disjointly across + ``(rank, worker)`` while keeping windows *within* a block sequential -> + decorrelates batches without random-access I/O. Derived from + ``_episode_cum_ends`` (the same monotonic cumulative index ``__getitem__`` + bisects), so blocks align exactly with flat-index addressing.""" + blocks: list[tuple[int, int]] = [] + prev = 0 + for c in self._episode_cum_ends: + c = int(c) + if c > prev: + blocks.append((prev, c - prev)) + prev = c + return blocks diff --git a/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py index 7e432ea3..64e2e12b 100644 --- a/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py @@ -1,316 +1,306 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -"""Minimal DROID LeRobot dataset for Cosmos Action v1.2 defaults.""" - -from __future__ import annotations - import json +import os import random -from pathlib import Path -from typing import Any, Literal +from typing import Any, cast import numpy as np -import pyarrow.parquet as pq import torch import torch.nn.functional as F -import torchvision.transforms as T - -from cosmos_framework.data.generator.action.action_spec import ActionSpec, Gripper, Joint, Pos, Rot, build_action_spec -from cosmos_framework.data.generator.action.datasets.base_dataset import ActionBaseDataset +import torchvision.transforms.v2 as T +from scipy.spatial.transform import Rotation as R + +from cosmos_framework.data.generator.action.datasets.cosmos3_action_lerobot import ( + ActionNormalization, + ActionSpec, + BaseActionLeRobotDataset, + Gripper, + Joint, + Pos, + Rot, + build_action_spec, + build_episode_spans, + split_episode_ids, +) +from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset_config import ( + _GRIPPER_STATE_FEATURE, + _JOINT_ACTION_FEATURE, + _JOINT_STATE_FEATURE, + ACTION_FEATURES, + HAS_MULTI_LANGUAGE_ANNOTATIONS, + IMAGE_FEATURES, + IS_FLAT_ACTION, + IS_GRIPPER_ACTION_FLIPPED, + LEROBOT_ROOTS, + STATE_FEATURES, +) from cosmos_framework.data.generator.action.pose_utils import ( + PoseConvention, build_abs_pose_from_components, + convert_rotation, pose_abs_to_rel, ) +from cosmos_framework.data.generator.action.viewpoint_utils import Viewpoint +from cosmos_framework.utils import log -PoseConvention = Literal["backward_framewise"] -Viewpoint = Literal["concat_view"] - -_IMAGE_FEATURES = { - "wrist": "observation.image.wrist_image_left", - "left": "observation.image.exterior_image_1_left", - "right": "observation.image.exterior_image_2_left", -} -_STATE_FEATURE = "observation.state.cartesian_position" -# joint_pos (8D = 7 arm joints + gripper) features, matching the internal -# DROIDLeRobotDataset(action_space="joint_pos", use_state=...). These are -# absolute joint commands/states (no normalization is applied for joint_pos, -# matching the internal canonical run which leaves action_normalization=None). -_JOINT_ACTION_FEATURE = "action.joint_position" # [7] commanded joints -_ACTION_GRIPPER_FEATURE = "action.gripper_position" # [1] commanded gripper -_JOINT_STATE_FEATURE = "observation.state.joint_positions" # [7] observed joints -_GRIPPER_STATE_FEATURE = "observation.state.gripper_position" # [1] observed gripper -# Columns whose parquet dtype is a list (need to_pylist -> stacked array). -_LIST_COLUMNS = {_STATE_FEATURE, _JOINT_ACTION_FEATURE, _JOINT_STATE_FEATURE} -_ACTION_SPACES = ("ee_pose", "joint_pos") - -# 90-degree clockwise rotation about the Z axis in the local frame. This matches -# the production DROID wrapper conversion from Franka panda_link8 to OpenCV. +_FILTER_DICT_PATH = "/scratch/fsw/portfolios/cosmos/projects/cosmos_base_training/users/haolia/workspace/droid_oss_inputs/keep_ranges_1_0_1.json" + +# 90-degree clockwise rotation about the Z axis (in local frame), converting +# DROID Franka panda_link8 orientation to the OpenCV camera convention. _DROID_TO_OPENCV: np.ndarray = np.array( - [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + [ + [0.0, -1.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ], dtype=np.float32, ) -_NORMALIZER_PATH = Path(__file__).parent.parent / "normalizer_stats/droid_lerobot_stats.json" - - -class DROIDLeRobotDataset(ActionBaseDataset): - """DROID Action dataset. - Two action layouts: - * ``action_space="ee_pose"`` (default): 10D ``[pos_delta(3), rot6d_delta(6), - gripper(1)]``, quantile-normalized (the v1.2 midtrain default). - * ``action_space="joint_pos"``: 8D ``[joint(7), gripper(1)]`` absolute joint - commands, NOT normalized, with ``use_state=True`` prepending the initial - observed joint+gripper state → ``(chunk+1, 8)`` — matching the internal - ``Cosmos3-Nano-Policy-DROID`` post-training run. - Filter dictionaries, temporal-segment validation, and image augmentation from - the production wrapper are intentionally omitted. - """ +class DROIDLeRobotDataset(BaseActionLeRobotDataset): + """ """ def __init__( self, - root: str, + root: str = "/lustre/fsw/portfolios/cosmos/projects/cosmos_base_training/cosmos3_action_datasets/droid_plus_lerobot_640x360_20260412", fps: float = 15.0, chunk_length: int = 16, - mode: str = "joint", + split_seed: int = 42, + split_val_ratio: float = 0.03, + split: str = "train", + mode: str = "policy", pose_convention: PoseConvention = "backward_framewise", - tolerance_s: float = 2e-4, + action_normalization: ActionNormalization | None = None, + tolerance_s=2e-4, viewpoint: Viewpoint = "concat_view", - action_space: str = "ee_pose", + use_success_only: bool = False, + video_mode: str | None = None, # TODO (ychao): remove + action_space: str = "midtrain", # TODO (ychao): remove use_state: bool = False, - action_normalization: str | None = "quantile", - use_image_augmentation: bool = False, use_filter_dict: bool = False, filter_dict_path: str | None = None, + enable_fast_init: bool = False, + max_num_history_actions: int = 0, + use_image_augmentation: bool = False, ) -> None: - if viewpoint != "concat_view": - raise NotImplementedError("This minimal DROID dataset only supports concat_view.") - if action_space not in _ACTION_SPACES: - raise NotImplementedError(f"action_space must be one of {_ACTION_SPACES}, got {action_space!r}.") - if use_state and action_space != "joint_pos": - raise NotImplementedError("use_state is only supported with action_space='joint_pos'.") - if use_filter_dict and not filter_dict_path: - raise ValueError("use_filter_dict=True requires filter_dict_path") - - # joint_pos uses raw joint values — disable normalization at the base level. + """ """ super().__init__( - root=root, - domain_name="droid_lerobot", fps=fps, chunk_length=chunk_length, + split_seed=split_seed, + split_val_ratio=split_val_ratio, + split=split, mode=mode, + embodiment_type="droid_lerobot", + viewpoint=viewpoint, pose_convention=pose_convention, + rotation_format="rot6d", + action_normalization=action_normalization, tolerance_s=tolerance_s, - viewpoint=viewpoint, - action_normalization=None if action_space == "joint_pos" else action_normalization, + enable_fast_init=enable_fast_init, ) - + self._use_success_only = use_success_only + self._video_mode = video_mode self._action_space = action_space - self._use_state = bool(use_state) - # Per-sample image augmentation (random crop+rescale + color jitter), applied - # to all views with shared params (temporally + cross-view consistent). Lazy-built. - self._use_image_augmentation = bool(use_image_augmentation) - self._image_augmentor: T.Compose | None = None - # Keep-ranges window filter (internal use_filter_dict): restrict training windows - # to curated active segments, dropping idle/non-task frames. Off by default; the - # keep-ranges JSON is supplied via filter_dict_path (an internal data artifact). - self._use_filter_dict = bool(use_filter_dict) - self._filter_dict_path = filter_dict_path - - # Compact, lazy frame index. Materializing every frame as a Python dict - # (``sorted(... pq.read_table(path).to_pylist() ...)``) does not scale: - # the full DROID success shard is ~18M frames, which is tens of GB of - # dicts plus an 18M-element Python sort at construction, and each - # DataLoader worker faults in its own copy. Instead we read only the - # columns the sample builder needs into contiguous numpy arrays - # (~1 GB total) -- read-only after init, so worker forks share them - # copy-on-write. - if action_space == "joint_pos": - feature_cols = [_JOINT_ACTION_FEATURE, _ACTION_GRIPPER_FEATURE, _JOINT_STATE_FEATURE, _GRIPPER_STATE_FEATURE] + self._use_state = use_state + self._use_filter_dict = use_filter_dict + self._filter_dict_path = filter_dict_path or _FILTER_DICT_PATH + self._max_num_history_actions = max_num_history_actions + self._use_image_augmentation = use_image_augmentation + if max_num_history_actions > 0 and action_space not in ("midtrain", "joint_pos"): + raise ValueError( + f"max_num_history_actions is only supported with action_space='midtrain' or 'joint_pos', got {action_space!r}" + ) + + self._is_val_temp_seg = split == "val_temp_seg" + self._to_opencv = _DROID_TO_OPENCV + + version = os.path.basename(root) + try: + lerobot_roots = LEROBOT_ROOTS[version] + self._image_features = IMAGE_FEATURES[version] + self._state_features = STATE_FEATURES[version] + self._action_features = ACTION_FEATURES[version] + self._is_flat_action = IS_FLAT_ACTION[version] + self._has_multi_language_annotations = HAS_MULTI_LANGUAGE_ANNOTATIONS[version] + self._is_gripper_action_flipped = IS_GRIPPER_ACTION_FLIPPED[version] + except KeyError as e: + raise ValueError(f"Unknown version: {version!r}. Supported: {list(LEROBOT_ROOTS.keys())}") from e + + if self._use_success_only and lerobot_roots: + lerobot_roots = [x for x in lerobot_roots if x.split("/", 1)[0] == "success"] + + self._all_shard_roots = [os.path.join(root, x) for x in lerobot_roots] if lerobot_roots else [root] + + observation_ts = [i * self._dt for i in range(0, self._chunk_length + 1)] + action_ts = [i * self._dt for i in range(0, self._chunk_length)] + if self._max_num_history_actions > 0 and self._action_space in ("midtrain", "joint_pos"): + observation_ts_ext = [i * self._dt for i in range(-self._max_num_history_actions, self._chunk_length + 1)] + action_ts_ext = [i * self._dt for i in range(-self._max_num_history_actions, self._chunk_length)] else: - feature_cols = [_STATE_FEATURE, _ACTION_GRIPPER_FEATURE] - columns = ["index", "episode_index", "task_index", "timestamp", *feature_cols] - index_parts, episode_parts, task_parts, ts_parts = [], [], [], [] - feature_parts: dict[str, list] = {c: [] for c in feature_cols} - for path in sorted((self._root / "data").glob("chunk-*/file-*.parquet")): - table = pq.read_table(path, columns=columns) - index_parts.append(table["index"].to_numpy()) - episode_parts.append(table["episode_index"].to_numpy()) - task_parts.append(table["task_index"].to_numpy()) - ts_parts.append(table["timestamp"].to_numpy()) - for c in feature_cols: - if c in _LIST_COLUMNS: - feature_parts[c].append(np.asarray(table[c].to_pylist(), dtype=np.float32)) - else: - feature_parts[c].append(np.asarray(table[c].to_numpy(), dtype=np.float32)) - order = np.argsort(np.concatenate(index_parts).astype(np.int64), kind="stable") - self._row_episode = np.concatenate(episode_parts).astype(np.int64)[order] - self._row_task = np.concatenate(task_parts).astype(np.int64)[order] - self._row_timestamp = np.concatenate(ts_parts).astype(np.float64)[order] - # Per-feature arrays keyed by parquet column name (read-only after init). - self._feat = { - c: np.concatenate(feature_parts[c], axis=0).astype(np.float32)[order] for c in feature_cols + observation_ts_ext = observation_ts + action_ts_ext = action_ts + self._delta_timestamps: dict[str, list[float]] = { + self._state_features: observation_ts_ext, + self._action_features: action_ts_ext, } - - # Group frames into episodes and keep only within-episode chunk windows. - # The global frame index is ordered by episode in LeRobot v3, so episodes - # are contiguous blocks once sorted by ``index``. The previous code sliced - # the flat row list (``rows[idx : idx + chunk + 1]``) with no boundary - # guard, so ~one chunk of samples per episode silently mixed two episodes; - # restricting to in-episode windows yields ``total - n_episodes * chunk`` - # valid samples (matching the production dataset). - assert np.all(np.diff(self._row_episode) >= 0), "episode_index is not contiguous after sorting by frame index" - ep_vals, ep_starts, ep_counts = np.unique(self._row_episode, return_index=True, return_counts=True) - self._ep_vals = ep_vals.astype(np.int64) - self._ep_starts = ep_starts.astype(np.int64) - self._valid_cum = np.cumsum(np.maximum(0, ep_counts - self._chunk_length)).astype(np.int64) - - # Keep-ranges filter: build a per-segment index over only the kept windows. - # Mirrors internal _append_index_records (use_filter_dict): the filter dict maps a - # gs:// trajectory key -> list of [start, end] frame ranges; keep windows whose start - # is in [max(start,0), min(end-chunk, valid)). Episodes absent from the dict are dropped. - if self._use_filter_dict: - with open(self._filter_dict_path) as f: - filter_dict = json.load(f) - seg_ep_pos, seg_win_start, seg_len = [], [], [] - for pos in range(len(self._ep_vals)): - valid = int(max(0, ep_counts[pos] - self._chunk_length)) - if valid <= 0: - continue - ep_id = str(self._episodes[int(self._ep_vals[pos])]["episode_id"]) - key = ( - f"gs://xembodiment_data/r2d2/r2d2-data-full/{ep_id}/recordings/" - f"MP4--gs://xembodiment_data/r2d2/r2d2-data-full/{ep_id}/trajectory.h5" - ) - ranges = filter_dict.get(key) - if ranges is None: - continue - for s, e in ranges: - ws = max(int(s), 0) - we = min(int(e) - self._chunk_length, valid) - if we - ws > 0: - seg_ep_pos.append(pos) - seg_win_start.append(ws) - seg_len.append(we - ws) - self._seg_ep_pos = np.asarray(seg_ep_pos, dtype=np.int64) - self._seg_win_start = np.asarray(seg_win_start, dtype=np.int64) - self._seg_cum = np.cumsum(seg_len).astype(np.int64) if seg_len else np.zeros(0, dtype=np.int64) - - @property - def action_dim(self) -> int: - return 8 if self._action_space == "joint_pos" else 10 - - def _action_spec(self) -> ActionSpec: + if self._viewpoint in ("wrist_view", "concat_view"): + self._delta_timestamps[self._image_features["wrist"]] = observation_ts + if self._viewpoint in ("third_person_view", "concat_view"): + self._delta_timestamps[self._image_features["left"]] = observation_ts + self._delta_timestamps[self._image_features["right"]] = observation_ts if self._action_space == "joint_pos": - return build_action_spec(Joint(n=7, label="joint"), Gripper()) - return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + self._delta_timestamps[_JOINT_ACTION_FEATURE] = action_ts + if self._use_state or self._max_num_history_actions > 0: + self._delta_timestamps[_JOINT_STATE_FEATURE] = observation_ts_ext + self._delta_timestamps[_GRIPPER_STATE_FEATURE] = observation_ts_ext + if self._use_state and self._action_space != "joint_pos": + self._delta_timestamps[_GRIPPER_STATE_FEATURE] = observation_ts - @classmethod - def _stats_path(cls) -> Path: - return _NORMALIZER_PATH - - def _window_rows(self, start: int, stop: int, episode_index: int) -> list[dict[str, Any]]: - """Reconstruct the per-frame dicts the sample builder consumes for the - half-open frame window ``[start, stop)`` from the compact column arrays. - ``start``/``stop`` are guaranteed to lie within a single episode.""" - return [ - { - "episode_index": episode_index, - "task_index": int(self._row_task[j]), - "timestamp": float(self._row_timestamp[j]), - **{c: self._feat[c][j] for c in self._feat}, - } - for j in range(start, stop) - ] - - def __getitem__(self, idx: int) -> dict[str, Any]: - mode = self._choose_mode() - idx = int(idx) - # Map the flat sample index to a within-episode frame window. if self._use_filter_dict: - seg = int(np.searchsorted(self._seg_cum, idx, side="right")) - base = int(self._seg_cum[seg - 1]) if seg > 0 else 0 - ep = int(self._seg_ep_pos[seg]) - start = int(self._ep_starts[ep]) + int(self._seg_win_start[seg]) + (idx - base) - else: - ep = int(np.searchsorted(self._valid_cum, idx, side="right")) - prev = int(self._valid_cum[ep - 1]) if ep > 0 else 0 - start = int(self._ep_starts[ep]) + (idx - prev) - episode_index = int(self._ep_vals[ep]) - episode = self._episodes[episode_index] - - observation_rows = self._window_rows(start, start + self._chunk_length + 1, episode_index) + with open(self._filter_dict_path) as f: + self._filter_dict = json.load(f) - video = self._load_concat_video(episode, observation_rows) - if self._action_space == "joint_pos": - raw_action = self._build_joint_action(observation_rows) - extras: dict[str, Any] = {} - else: - action_rows = observation_rows[: self._chunk_length] - raw_action, initial_pose = self._build_raw_action(observation_rows, action_rows) - extras = {"initial_pose": initial_pose} - task = self._tasks[int(observation_rows[0]["task_index"])] - ai_caption = random.choice(task.split(" | ")) + self._image_augmentor: T.Compose | None = None - return self._build_result( - mode=mode, - video=video, - action=raw_action, - ai_caption=ai_caption, - additional_view_description=( - "The top row is from the wrist-mounted camera. " - "The bottom row contains two horizontally concatenated third-person perspective views of the scene from opposite sides, with the robot visible." - ), - **extras, + # Eager source registration. i4 defers this to its own dataloader's + # ActionUnifiedIterableDataset.assign_worker(); cosmos-framework instead + # drives the dataset through ActionIterableShuffleDataset (block-striding, + # which needs the full flat index present in every worker), so we build + # the index at construction time here. Metadata-only (LeRobotDatasetMetadata: + # info.json + episodes.parquet + tasks.parquet); the heavy per-shard + # LeRobotDataset video readers stay lazy behind the LRU in _get_dataset. + self._register_sources() + + def _append_index_records(self, *, meta, ds_idx: int, dataset_label: str | None = None) -> None: + """ """ + if not self._use_filter_dict: + super()._append_index_records(meta=meta, ds_idx=ds_idx, dataset_label=dataset_label) + return + + episode_ids = split_episode_ids( + total_episodes=meta.total_episodes, + seed=self._split_seed, + val_ratio=self._split_val_ratio, + split=self._split, + ) + episode_spans, _, sample_count = build_episode_spans( + meta.episodes, episode_ids, self._chunk_length, sample_stride=self._sample_stride ) - def _build_joint_action(self, observation_rows: list[dict[str, Any]]) -> torch.Tensor: - """8D joint-position action ``[joint(7), gripper(1)]`` over the chunk, matching - the internal ``action_space='joint_pos'``. The window is ``chunk+1`` frames: - ``row[0]`` is the initial observed state (prepended when ``use_state``), and - ``rows[1:]`` are the ``chunk`` commanded actions. Gripper is flipped (1 - g). - No normalization is applied (internal canonical run uses raw joint values).""" - action_rows = observation_rows[1:] - joints = np.asarray([r[_JOINT_ACTION_FEATURE] for r in action_rows], dtype=np.float32) # [chunk, 7] - gripper = np.asarray([r[_ACTION_GRIPPER_FEATURE] for r in action_rows], dtype=np.float32).reshape(-1, 1) - gripper = 1.0 - gripper - action = np.concatenate([joints, gripper], axis=-1) # [chunk, 8] - if self._use_state: - init = observation_rows[0] - init_joint = np.asarray(init[_JOINT_STATE_FEATURE], dtype=np.float32) # [7] - init_gripper = np.asarray([1.0 - float(init[_GRIPPER_STATE_FEATURE])], dtype=np.float32) # [1] - initial_state = np.concatenate([init_joint, init_gripper])[None, :] # [1, 8] - action = np.concatenate([initial_state, action], axis=0) # [chunk + 1, 8] - return torch.from_numpy(action).float() - - def _load_concat_video( - self, - episode: dict[str, Any], - observation_rows: list[dict[str, Any]], - ) -> torch.Tensor: - # lerobot is a heavy, optional ("train" extra) dependency; import lazily. - from lerobot.datasets.video_utils import decode_video_frames - - timestamps = [float(row["timestamp"]) for row in observation_rows] - frames_by_view = { - name: decode_video_frames( - self._video_path(episode, video_key), - [float(episode.get(f"videos/{video_key}/from_timestamp", 0.0)) + ts for ts in timestamps], - self._tolerance_s, + class_name = self.__class__.__name__ + label = f" [{dataset_label}]" + + log.info(f"{class_name}{label}: split={self._split}, num episodes={len(episode_ids)}") + + filtered_count = 0 + for episode_id, sample_start, valid_len in episode_spans: + ep_id_str = meta.episodes[episode_id]["episode_id"] + episode_key = f"gs://xembodiment_data/r2d2/r2d2-data-full/{ep_id_str}/recordings/MP4--gs://xembodiment_data/r2d2/r2d2-data-full/{ep_id_str}/trajectory.h5" + ranges = self._filter_dict.get(episode_key) + if ranges is None: + continue + for s, e in ranges: + sub_start = max(s, 0) + sub_end = min(e - self._chunk_length, valid_len) + sub_valid_len = max(0, sub_end - sub_start) + if sub_valid_len > 0: + self._episode_records.append((ds_idx, sample_start + sub_start, sub_valid_len, episode_id)) + self._num_valid_indices += sub_valid_len + self._episode_cum_ends.append(self._num_valid_indices) + filtered_count += sub_valid_len + + if sample_count > 0: + log.info( + f"{class_name}{label}: kept {filtered_count} / {sample_count} ({100.0 * filtered_count / sample_count:.2f} %) samples" ) - for name, video_key in _IMAGE_FEATURES.items() - } - wrist = frames_by_view["wrist"] - left = frames_by_view["left"] - right = frames_by_view["right"] + def _register_sources(self, indices: list[int] | None = None) -> None: + """ """ + super()._register_sources(indices) + if self._is_val_temp_seg: + self._apply_temp_seg_filter() + + def _apply_temp_seg_filter(self) -> None: + """Replace index records with one high-scoring segment per episode. + + A segment is interesting if either: + - The gripper action changes significantly (open/close transition), or + - The gripper is closed and the end-effector position is moving. + Among qualifying segments the one with the highest score is kept. + """ + ds = self._get_dataset(0) + chunk_size = self._chunk_length + 1 + gripper_change_threshold = 0.5 + ee_movement_threshold = 0.01 + + new_records: list[tuple[int, int, int, int]] = [] + num_episodes = len(self._episode_records) + + for ds_idx, sample_start, valid_len, episode_id in self._episode_records: + end = sample_start + valid_len + self._chunk_length + num_candidates = valid_len + if num_candidates <= 0: + continue + + episode_data = ds.hf_dataset[sample_start:end] + actions = torch.tensor(np.array(episode_data[self._action_features])) # [N,action_dim] + states = torch.tensor(np.array(episode_data[self._state_features])) # [N,state_dim] + + gripper_action = actions[:, 6] if self._is_flat_action else actions # [N] + ee_pos = states[:, :3] # [N,3] + ee_disp = (ee_pos[1:] - ee_pos[:-1]).norm(dim=-1) # [N-1] + + ee_disp_windows = ee_disp.unfold(0, self._chunk_length, 1) # [num_candidates,chunk_length] + gripper_windows = gripper_action.unfold(0, chunk_size, 1) # [num_candidates,chunk_size] + + gripper_range = gripper_windows.max(dim=1).values - gripper_windows.min(dim=1).values # [num_candidates] + total_ee_movement = ee_disp_windows.sum(dim=1) # [num_candidates] + gripper_closed_ratio = (gripper_windows < 0.5).float().mean(dim=1) # [num_candidates] + + has_gripper_change = gripper_range > gripper_change_threshold + gripper_closed = gripper_closed_ratio > 0.5 + has_ee_movement = total_ee_movement > ee_movement_threshold + + scores = torch.zeros(num_candidates) # [num_candidates] + scores[has_gripper_change] = 0.5 + gripper_range[has_gripper_change] + total_ee_movement[has_gripper_change] + + closed_and_moving = gripper_closed & ~has_gripper_change & has_ee_movement + scores[closed_and_moving] = 1.0 + total_ee_movement[closed_and_moving] + + if scores.max().item() > 0: + best_offset = int(scores.argmax().item()) + new_records.append((ds_idx, sample_start + best_offset, 1, episode_id)) + + self._episode_records = new_records + self._num_valid_indices = len(new_records) + self._episode_cum_ends = list(range(1, len(new_records) + 1)) + + log.info(f"DROIDLeRobotDataset: val_temp_seg kept {len(new_records)} segments from {num_episodes} episodes") + + def _compose_multi_view(self, sample: dict[str, Any]) -> torch.Tensor: + """Compose wrist, left, and right views into a single frame. + + Layout (per frame): + ┌──────────────┐ + │ wrist │ (H, W) + ├───────┬──────┤ + │ left │ right│ (H/2, W/2) each + └───────┴──────┘ + + Left and right exterior cameras are downscaled by 2x so that they + tile to the same width as the wrist view. The output height is 3H/2. + + Returns: + Composited raw video tensor in ``(T,C,H_out,W)`` float format. + """ + wrist = sample[self._image_features["wrist"]] # [T,C,H,W] + left = sample[self._image_features["left"]] # [T,C,H_l,W_l] + right = sample[self._image_features["right"]] # [T,C,H_r,W_r] if self._use_image_augmentation: - # Random crop+rescale (spatial jitter) + color jitter, BEFORE the concat. - # All three views are stacked so one sampled set of params is applied - # uniformly across every frame and view (temporally + cross-view consistent), - # while each __getitem__ resamples. Matches the internal DROID recipe. if self._image_augmentor is None: _, _, h, w = wrist.shape self._image_augmentor = T.Compose( @@ -326,46 +316,182 @@ def _load_concat_video( _, _, h_w, w_w = wrist.shape half_h, half_w = h_w // 2, w_w // 2 - left = F.interpolate(left, size=(half_h, half_w), mode="bilinear", align_corners=False) - right = F.interpolate(right, size=(half_h, half_w), mode="bilinear", align_corners=False) - bottom = torch.cat([left, right], dim=-1) - return torch.cat([wrist, bottom], dim=-2) - def _build_raw_action( - self, - observation_rows: list[dict[str, Any]], - action_rows: list[dict[str, Any]], - ) -> tuple[torch.Tensor, torch.Tensor]: - state = np.asarray([row[_STATE_FEATURE] for row in observation_rows], dtype=np.float32) - poses_abs = build_abs_pose_from_components(state[:, 0:3], state[:, 3:6], "euler_xyz") - poses_abs[:, :3, :3] = poses_abs[:, :3, :3] @ _DROID_TO_OPENCV - - initial_pose = torch.from_numpy(poses_abs[0].copy()).float() - poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention=self._pose_convention) - gripper = np.asarray( - [row[_ACTION_GRIPPER_FEATURE] for row in action_rows], dtype=np.float32 - ).reshape(-1, 1) - gripper = 1.0 - gripper - action = np.concatenate([poses_rel[-self._chunk_length :], gripper[-self._chunk_length :]], axis=-1) - return torch.from_numpy(action).float(), initial_pose - - def __len__(self) -> int: - if self._use_filter_dict: - return int(self._seg_cum[-1]) if self._seg_cum.size else 0 - return int(self._valid_cum[-1]) if self._valid_cum.size else 0 - - def get_shuffle_blocks(self) -> list[tuple[int, int]]: - """Per-episode (or per kept-segment, when ``use_filter_dict``) flat-index blocks - ``(start, length)``. ``ActionIterableShuffleDataset`` shuffles the ORDER of these - blocks and shards them disjointly across ranks, while keeping windows *within* a - block sequential -> decorrelates batches across ranks without random-access I/O - (preserves locality + copy-on-write memory sharing across workers).""" - cum = self._seg_cum if self._use_filter_dict else self._valid_cum - blocks: list[tuple[int, int]] = [] - prev = 0 - for c in np.asarray(cum).tolist(): - c = int(c) - if c > prev: - blocks.append((prev, c - prev)) - prev = c - return blocks + left = F.interpolate(left, size=(half_h, half_w), mode="bilinear", align_corners=False) # [T,C,H/2,W/2] + right = F.interpolate(right, size=(half_h, half_w), mode="bilinear", align_corners=False) # [T,C,H/2,W/2] + bottom = torch.cat([left, right], dim=-1) # [T,C,H/2,W] + + composite = torch.cat([wrist, bottom], dim=-2) # [T,C,3H/2,W] + return composite # [T,C,3H/2,W] + + def _build_action_spec(self) -> ActionSpec: + """DROID: 10D ``[Pos, Rot6d, Gripper]`` for ``ee_pose``, + 8D ``[Joint(7), Gripper]`` for ``joint_pos``. + """ + if self._action_space == "joint_pos": + return build_action_spec(Joint(n=7, label="joint"), Gripper()) + return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + + def __getitem__(self, idx: int) -> dict[str, Any]: + """ """ + mode, _, _, sample = self._fetch_sample(idx) + + if self._has_multi_language_annotations: + tasks = sample["task"].split(" | ") + ai_caption = random.choice(tasks) + else: + ai_caption = sample["task"] + + if self._skip_video_loading: + video = None + elif self._video_mode is None: + if self._viewpoint == "concat_view": + video = self._compose_multi_view(sample) + else: + video = sample[self._image_features["wrist"]] # [T,C,H,W] + else: + if self._video_mode == "wrist": + video = sample[self._image_features["wrist"]] + if self._video_mode in ("rand_exterior", "wrist_rand_exterior"): + exterior_key = random.choice([self._image_features["left"], self._image_features["right"]]) + if self._video_mode == "rand_exterior": + video = sample[exterior_key] + else: + video = torch.cat([sample[self._image_features["wrist"]], sample[exterior_key]], dim=2) + if self._video_mode in ("wrist_left_exterior", "wrist_both_exterior"): + wrist = sample[self._image_features["wrist"]] + half_h, half_w = wrist.shape[2] // 2, wrist.shape[3] // 2 + left = F.interpolate( + sample[self._image_features["left"]], size=(half_h, half_w), mode="bilinear", align_corners=False + ) + if self._video_mode == "wrist_left_exterior": + right = torch.zeros_like(left) + if self._video_mode == "wrist_both_exterior": + right = F.interpolate( + sample[self._image_features["right"]], + size=(half_h, half_w), + mode="bilinear", + align_corners=False, + ) + video = torch.cat([wrist, torch.cat([left, right], dim=-1)], dim=-2) + + extras: dict[str, Any] = {} + + if self._action_space == "midtrain": + pose_convention = cast(PoseConvention, self._pose_convention) + state = sample[self._state_features] # [T+1, state_dim] or [H+T+1, state_dim] + poses_abs = build_abs_pose_from_components(state[:, 0:3], state[:, 3:6], "euler_xyz") + poses_abs[:, :3, :3] = poses_abs[:, :3, :3] @ self._to_opencv + initial_pose = torch.from_numpy(poses_abs[-self._chunk_length - 1].copy()).float() + poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention=pose_convention) + gripper = ( + sample[self._action_features][:, [6]] + if self._is_flat_action + else sample[self._action_features].unsqueeze(-1) + ) + if self._is_gripper_action_flipped: + gripper = 1.0 - gripper + action = torch.from_numpy( + np.concatenate([poses_rel[-self._chunk_length :], gripper[-self._chunk_length :]], axis=-1) + ).float() # [T,10] + extras["initial_pose"] = initial_pose + if self._max_num_history_actions > 0: + _, _, _, frame_offset = self._resolve_index(int(idx)) + num_available = min(self._max_num_history_actions, frame_offset * self._sample_stride) + actual_h = num_available + # with 0.5 probability, randomly sample the number of history frames + if random.random() < 0.5: + actual_h = random.randint(0, num_available) + if actual_h > 0: + hist_action_raw = torch.from_numpy( + np.concatenate( + [ + poses_rel[-self._chunk_length - actual_h : -self._chunk_length], + gripper[-self._chunk_length - actual_h : -self._chunk_length], + ], + axis=-1, + ) + ).float() # [H,D] + extras["history_action"] = hist_action_raw + if self._use_state: + initial_gripper = sample[_GRIPPER_STATE_FEATURE][0].unsqueeze(-1) + if self._is_gripper_action_flipped: + initial_gripper = 1.0 - initial_gripper + initial_rot6d = convert_rotation(poses_abs[-self._chunk_length - 1, :3, :3], "matrix", "rot6d") + initial_state = torch.from_numpy( + np.concatenate((poses_abs[-self._chunk_length - 1, :3, 3], initial_rot6d, initial_gripper), axis=-1) + ).float() + action = torch.cat([initial_state.unsqueeze(0), action], dim=0) + if self._action_space == "ee_pose_delta": + state = sample[self._state_features] + pose = np.tile(np.eye(4), (state.shape[0], 1, 1)) + pose[:, :3, :3] = R.from_euler("xyz", state[:, 3:6]).as_matrix() + pose[:, :3, 3] = state[:, 0:3] + pose_delta = np.linalg.inv(pose[0]) @ pose[1:] + gripper = sample[self._action_features].unsqueeze(-1) + if self._is_gripper_action_flipped: + gripper = 1.0 - gripper + action = torch.from_numpy( + np.concatenate((pose_delta[:, :3, 3], pose_delta[:, :3, 0], pose_delta[:, :3, 1], gripper), axis=-1) + ).float() + if self._use_state: + initial_gripper = sample[_GRIPPER_STATE_FEATURE][0].unsqueeze(-1) + if self._is_gripper_action_flipped: + initial_gripper = 1.0 - initial_gripper + initial_state = torch.from_numpy( + np.concatenate((pose[0, :3, 3], pose[0, :3, 0], pose[0, :3, 1], initial_gripper), axis=-1) + ).float() + action = torch.cat([initial_state.unsqueeze(0), action], dim=0) + if self._action_space == "joint_pos": + gripper = sample[self._action_features][-self._chunk_length :].unsqueeze(-1) + if self._is_gripper_action_flipped: + gripper = 1.0 - gripper + action = torch.cat((sample[_JOINT_ACTION_FEATURE], gripper), dim=-1).float() + if self._max_num_history_actions > 0: + _, _, _, frame_offset = self._resolve_index(int(idx)) + num_available = min(self._max_num_history_actions, frame_offset * self._sample_stride) + actual_h = num_available + if random.random() < 0.5: + actual_h = random.randint(0, num_available) + if actual_h > 0: + hist_joint = sample[_JOINT_STATE_FEATURE][ + -self._chunk_length - 1 - actual_h : -self._chunk_length - 1 + ] + hist_gripper = sample[_GRIPPER_STATE_FEATURE][ + -self._chunk_length - 1 - actual_h : -self._chunk_length - 1 + ].unsqueeze(-1) + if self._is_gripper_action_flipped: + hist_gripper = 1.0 - hist_gripper + hist_action_raw = torch.cat((hist_joint, hist_gripper), dim=-1).float() + extras["history_action"] = hist_action_raw # [H,D] + if self._use_state: + initial_gripper = sample[_GRIPPER_STATE_FEATURE][-self._chunk_length - 1].unsqueeze(-1) + if self._is_gripper_action_flipped: + initial_gripper = 1.0 - initial_gripper + initial_state = torch.cat( + (sample[_JOINT_STATE_FEATURE][-self._chunk_length - 1], initial_gripper), dim=-1 + ).float() + action = torch.cat([initial_state.unsqueeze(0), action], dim=0) + + if self._viewpoint == "concat_view" and self._video_mode in ( + None, + "wrist_left_exterior", + "wrist_both_exterior", + ): + extras["additional_view_description"] = ( + "The top row is from the wrist-mounted camera. " + "The bottom row contains two horizontally concatenated third-person perspective views of the scene from opposite sides, with the robot visible." + ) + + return self._build_result( + mode=mode, + video=video, + action=action, + ai_caption=ai_caption, + **extras, + ) + + @property + def action_dim(self) -> int: + """ """ + return 8 if self._action_space == "joint_pos" else 10 diff --git a/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_config.py b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_config.py new file mode 100644 index 00000000..e97489c5 --- /dev/null +++ b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset_config.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +_INSTITUTIONS = [ + "AUTOLab", + "CLVR", + "GuptaLab", + "ILIAD", + "IPRL", + "IRIS", + "PennPAL", + "RAD", + "RAIL", + "REAL", + "RPL", + "TRI", + "WEIRD", +] + +LEROBOT_ROOTS = { + "droid_lerobot_20260115_no_noops": None, + "droid_plus_lerobot_320x180_20260406_sharded": [f"success/{x}" for x in _INSTITUTIONS] + + [f"failure/{x}" for x in _INSTITUTIONS], + "droid_plus_lerobot_320x180_20260406": ["success", "failure"], + "droid_plus_lerobot_640x360_20260412_sharded": [f"success/{x}" for x in _INSTITUTIONS] + + [f"failure/{x}" for x in _INSTITUTIONS], + "droid_plus_lerobot_640x360_20260412": ["success", "failure"], +} + +IMAGE_FEATURES = { + "droid_lerobot_20260115_no_noops": { + "wrist": "observation.images.wrist_image_left", + "left": "observation.images.exterior_image_1_left", + "right": "observation.images.exterior_image_2_left", + }, + "droid_plus_lerobot_320x180_20260406_sharded": { + "wrist": "observation.image.wrist_image_left", + "left": "observation.image.exterior_image_1_left", + "right": "observation.image.exterior_image_2_left", + }, + "droid_plus_lerobot_320x180_20260406": { + "wrist": "observation.image.wrist_image_left", + "left": "observation.image.exterior_image_1_left", + "right": "observation.image.exterior_image_2_left", + }, + "droid_plus_lerobot_640x360_20260412_sharded": { + "wrist": "observation.image.wrist_image_left", + "left": "observation.image.exterior_image_1_left", + "right": "observation.image.exterior_image_2_left", + }, + "droid_plus_lerobot_640x360_20260412": { + "wrist": "observation.image.wrist_image_left", + "left": "observation.image.exterior_image_1_left", + "right": "observation.image.exterior_image_2_left", + }, +} + +STATE_FEATURES = { + "droid_lerobot_20260115_no_noops": "observation.state", + "droid_plus_lerobot_320x180_20260406_sharded": "observation.state.cartesian_position", + "droid_plus_lerobot_320x180_20260406": "observation.state.cartesian_position", + "droid_plus_lerobot_640x360_20260412_sharded": "observation.state.cartesian_position", + "droid_plus_lerobot_640x360_20260412": "observation.state.cartesian_position", +} + +ACTION_FEATURES = { + "droid_lerobot_20260115_no_noops": "action", + "droid_plus_lerobot_320x180_20260406_sharded": "action.gripper_position", + "droid_plus_lerobot_320x180_20260406": "action.gripper_position", + "droid_plus_lerobot_640x360_20260412_sharded": "action.gripper_position", + "droid_plus_lerobot_640x360_20260412": "action.gripper_position", +} + +IS_FLAT_ACTION = { + "droid_lerobot_20260115_no_noops": True, + "droid_plus_lerobot_320x180_20260406_sharded": False, + "droid_plus_lerobot_320x180_20260406": False, + "droid_plus_lerobot_640x360_20260412_sharded": False, + "droid_plus_lerobot_640x360_20260412": False, +} + +HAS_MULTI_LANGUAGE_ANNOTATIONS = { + "droid_lerobot_20260115_no_noops": False, + "droid_plus_lerobot_320x180_20260406_sharded": True, + "droid_plus_lerobot_320x180_20260406": True, + "droid_plus_lerobot_640x360_20260412_sharded": True, + "droid_plus_lerobot_640x360_20260412": True, +} + +IS_GRIPPER_ACTION_FLIPPED = { + "droid_lerobot_20260115_no_noops": False, + "droid_plus_lerobot_320x180_20260406_sharded": True, + "droid_plus_lerobot_320x180_20260406": True, + "droid_plus_lerobot_640x360_20260412_sharded": True, + "droid_plus_lerobot_640x360_20260412": True, +} + +_JOINT_ACTION_FEATURE = "action.joint_position" +_JOINT_STATE_FEATURE = "observation.state.joint_positions" +_GRIPPER_STATE_FEATURE = "observation.state.gripper_position" diff --git a/examples/toml/sft_config/action_policy_droid_repro.toml b/examples/toml/sft_config/action_policy_droid_repro.toml index 661f7bdb..2326d407 100644 --- a/examples/toml/sft_config/action_policy_droid_repro.toml +++ b/examples/toml/sft_config/action_policy_droid_repro.toml @@ -3,10 +3,11 @@ # ============================================================================ # DROID action-policy SFT — run config for the `action_policy_droid_nano` -# experiment. The recipe knobs (optimizer/lr, scheduler type, grad_clip, -# count-based batch, action-head skip-on-load, dataset knobs) live in the -# registered experiment; this file only sets run-level scalars (iters, ckpt -# cadence, parallelism shape, wandb, VAE path). +# experiment. Reproduces the Cosmos3-Nano-Policy-DROID reference run: +# HSDP 32x8, global batch 8192, lr 2e-4, loss_scale 10, 10000 iters. The +# remaining recipe knobs (grad_clip, count-based batch, action-head +# skip-on-load, dataset knobs) live in the registered experiment; this file +# pins the run-level scalars that define the reference reproduction. # # Env required: # DROID_ROOT=/path/to/droid_lerobot_640x360/success @@ -27,8 +28,8 @@ wandb_mode = "online" precision = "bfloat16" [model.parallelism] -data_parallel_shard_degree = 8 # 8-GPU model shard; set replicate for multi-node HSDP -data_parallel_replicate_degree = 1 +data_parallel_shard_degree = 8 # 8-GPU model shard +data_parallel_replicate_degree = 32 # HSDP 32x8 = 256 ranks (GB200 reference: 64 nodes x 4) [model.activation_checkpointing] mode = "full" @@ -37,19 +38,37 @@ save_ops_regex = ["fmha"] [model.tokenizer] vae_path = "${oc.env:WAN_VAE_PATH}" +[model.rectified_flow_training_config] +loss_scale = 10.0 # generator (diffusion) loss weight, reference value + +[optimizer] +lr = 2.0e-04 + [scheduler] -cycle_lengths = [10000] # match max_iter +cycle_lengths = [100000] # long cosine cycle: lr decays slowly across the 10000-iter run + +[dataloader_train] +max_samples_per_batch = 32 # samples packed into each per-rank batch (res480) + +[dataloader_train.dataloader] +num_workers = 16 # decode-worker throughput tuning; adjust to host CPU count +batch_size = 16 +prefetch_factor = 2 [trainer] -max_iter = 10000 -logging_iter = 50 +max_iter = 10000 +logging_iter = 50 +grad_accum_iter = 1 # global batch = max_samples 32 x (shard 8 x replicate 32) x 1 = 8192 + +[trainer.callbacks.compile_tokenizer] +enabled = true # torch.compile the Wan VAE tokenizer (speed); warms up at res480 +warmup_resolutions = ["480"] [checkpoint] load_path = "${oc.env:BASE_CHECKPOINT_PATH}" save_iter = 1000 -# max_samples_per_batch is 128 in the experiment — samples packed into each per-rank batch -# (the num_workers x prefetch_factor workers just decode in parallel to keep it fed); res480, -# reference recipe, validated multi-node on GB200. -# On lower-memory GPUs, reduce it at launch, e.g.: -# --opts dataloader_train.max_samples_per_batch=32 +# The 256-rank HSDP 32x8 shape (global batch 8192) is the GB200 reference. To fit +# fewer GPUs while keeping global batch 8192, lower data_parallel_replicate_degree +# and raise grad_accum_iter at launch — e.g. one 8-GPU node: +# --opts model.parallelism.data_parallel_replicate_degree=1 trainer.grad_accum_iter=32 diff --git a/tests/action_policy_regression_test.py b/tests/action_policy_regression_test.py new file mode 100644 index 00000000..a2206512 --- /dev/null +++ b/tests/action_policy_regression_test.py @@ -0,0 +1,466 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Numerical-stability regression test for the ACTION-policy SFT launches (DROID + LIBERO). + +The action-policy analogue of ``tests/launch_regression_test.py`` (vision/VLM): +re-runs the same ``torchrun`` command the launcher shells execute — capped to +10 iterations, ``--deterministic``, seed 42 — for each of the two action-policy +recipes and asserts that the rank-0 per-step ``Loss:`` series reproduces the +inline goldens at the bottom of this file (per GPU arch, with a tolerance). + +Both recipes post-train the public **Cosmos3-Nano** base (registered experiments +``action_policy_{libero,droid}_nano``) — the same base + Wan2.2 VAE the vision +smoke/regression tests use — so a single 4-GPU node (``data_parallel_shard_degree=4``, +``replicate_degree=1``) runs them. The recipe knobs (optimizer, action-loss +weight, dataset transforms, iterable episode-shuffle) live in the experiment; +this test only caps iters + shapes the run for a deterministic single-node +capture. + +Specs +----- +* ``action_policy_libero`` — ``examples/toml/sft_config/action_policy_libero_repro.toml``. + Data auto-downloads: the ``libero_10`` suite of ``nvidia/LIBERO_LeRobot_v3`` + (small LeRobot dir), cached across runs. Collapses the recipe's HSDP replicate + 2 -> 1 to fit one 4-GPU node. +* ``action_policy_droid`` — ``examples/toml/sft_config/action_policy_droid_repro.toml``. + The full DROID split is far too large to auto-download in CI, so this spec + requires a pre-staged LOCAL copy via ``DROID_ROOT`` and SKIPS when unset. + ``DROID_ROOT`` is the **versioned merged root** whose basename is a + ``LEROBOT_ROOTS`` version key (e.g. ``…/droid_plus_lerobot_640x360_20260412``), + NOT its ``success/`` subdir — the recipe's i4 lazy dataset appends ``success/`` + itself (``use_success_only``). The skip-check looks for + ``/success/meta/info.json``. + +Determinism notes +----------------- +The launch passes ``--deterministic`` + ``PYTHONHASHSEED=42`` and the recipes +seed the episode-shuffle stream (``episode_shuffle_seed=42``); the shard +assignment ``(rank, worker)`` and per-epoch permutation are seed-derived, so the +data order reproduces across runs for any fixed world size / ``num_workers``. +``compile_tokenizer`` is disabled for the capture (torch.compile makes the +all-rank grad-norm reduction non-bit-exact — same reason ``vision_sft_nano`` +pins ``grad_norm=None`` on H100), so ``grad_norm`` goldens are left ``None`` and +only the loss series is asserted. On GB200 the sibling vision recipe reproduces +loss bit-exact across all 10 iters under ``--deterministic``; the action recipe +shares the MoT/attention stack, so the goldens use the tight default tolerance +(loosen via ``loss_tol_bands`` if a captured tail proves noisy). + +Inputs land in the documented ``.gitignore``-d locations (``examples/data/``, +``examples/checkpoints/``), cached across runs and shared with the vision smoke +test; run output goes under the pytest tmp dir. + +Invocation (inside the training container, from the repo root, on a 4-GPU +node):: + + # LIBERO only (auto-downloads its data): + pytest -s tests/action_policy_regression_test.py --num-gpus=4 --levels=2 -o addopts= + # include DROID (pre-stage the local LeRobot split): + DROID_ROOT=/path/to/Cosmos3-DROID/success \ + pytest -s tests/action_policy_regression_test.py --num-gpus=4 --levels=2 -o addopts= + +Without ``--num-gpus``/``--levels`` (e.g. the no-GPU pre-commit CI) the tests +are not collected. + +Refreshing the goldens (after an intentional numerical change, on the target +arch):: + + COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS=1 pytest -s tests/action_policy_regression_test.py \ + --num-gpus=4 --levels=2 -o addopts= + +That prints the captured series for each spec; copy them into the matching +``_GOLDENS[]`` entry below. +""" + +from __future__ import annotations + +import os +import re +import shutil +import socket +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from cosmos_framework.inference.fixtures.args import MAX_GPUS + +REPO_ROOT = Path(__file__).resolve().parents[1] + +# Shared base artifacts — match the vision smoke test + the action launcher +# defaults so no path overrides are needed once present (all .gitignore-d, +# cached across runs). +_WAN_VAE = REPO_ROOT / "examples/checkpoints/wan22_vae/Wan2.2_VAE.pth" +_DCP_DIR = REPO_ROOT / "examples/checkpoints/Cosmos3-Nano" +# LIBERO-10 LeRobot suite (auto-downloaded). ``LIBERO_ROOT`` must point at the +# suite dir itself (contains meta/info.json). +_LIBERO_DIR = REPO_ROOT / "examples/data/LIBERO_LeRobot_v3" +_LIBERO_ROOT = _LIBERO_DIR / "libero_10" + +# rank-0 per-iteration loss from the IterSpeed callback's on_training_step_end: +# [RANK 0] Iteration 1: Hit counter: 1/50 | Loss: 0.2515 | Time: 120.42s +_VFM_LOSS_RE = re.compile(r"\[RANK\s+0\]\s+Iteration\s+\d+:\s+Hit counter:[^|]+\|\s+Loss:\s+([0-9.eE+-]+)") + +_DEFAULT_RTOL = 1e-3 +_DEFAULT_ATOL = 1e-3 + + +def _free_port() -> int: + """Return a currently-free TCP port for torchrun's rendezvous (avoids + ``EADDRINUSE`` from a hardcoded ``master_port`` when a prior run lingers).""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +def _detect_arch() -> str: + """Map ``torch.cuda.get_device_name(0)`` to a goldens key.""" + import torch + + if not torch.cuda.is_available(): + return "unknown" + name = torch.cuda.get_device_name(0).upper() + if "GB200" in name: + return "gb200" + if "H100" in name or "H200" in name: + return "h100" + return "unknown" + + +def _run(cmd: list[str], log_file: Path, extra_env: dict | None = None) -> tuple[int, str]: + """Run ``cmd`` from the repo root, tee combined output to ``log_file``. + + Streams live to stdout (so CI shows progress under ``pytest -s``) while + capturing into the log + a string. Inherits the caller's env plus + ``PYTHONPATH=.``. + """ + env = os.environ.copy() + env["PYTHONPATH"] = f".:{env.get('PYTHONPATH', '')}" + if extra_env: + env.update(extra_env) + log_file.parent.mkdir(parents=True, exist_ok=True) + captured: list[str] = [] + with log_file.open("w") as fp: + proc = subprocess.Popen( + cmd, + env=env, + cwd=str(REPO_ROOT), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + assert proc.stdout is not None + for line in proc.stdout: + sys.stdout.write(line) + sys.stdout.flush() + fp.write(line) + captured.append(line) + returncode = proc.wait() + return returncode, "".join(captured) + + +# --- input staging (shared with tests/nano_training_smoke_test.py) ----------- + + +def _ensure_wan_vae(log_dir: Path) -> None: + """Download the Wan2.2 VAE if not already present.""" + if _WAN_VAE.is_file(): + return + rc, out = _run( + [ + "uvx", + "hf@latest", + "download", + "Wan-AI/Wan2.2-TI2V-5B", + "Wan2.2_VAE.pth", + "--local-dir", + str(_WAN_VAE.parent), + "--quiet", + ], + log_dir / "download_wan_vae.log", + ) + assert rc == 0, f"Wan VAE download failed (exit {rc}):\n{out[-2000:]}" + assert _WAN_VAE.is_file(), f"Wan VAE missing at {_WAN_VAE} after download" + + +def _ensure_dcp(log_dir: Path) -> None: + """Convert Cosmos3-Nano to DCP if not already present.""" + if _DCP_DIR.is_dir() and any(_DCP_DIR.iterdir()): + return + rc, out = _run( + [ + "python", + "-m", + "cosmos_framework.scripts.convert_model_to_dcp", + "--checkpoint-path", + "Cosmos3-Nano", + "-o", + str(_DCP_DIR), + ], + log_dir / "convert_to_dcp.log", + ) + assert rc == 0, f"convert_model_to_dcp failed (exit {rc}):\n{out[-3000:]}" + assert _DCP_DIR.is_dir() and any(_DCP_DIR.iterdir()), f"DCP not written to {_DCP_DIR}" + + +def _ensure_libero(log_dir: Path) -> None: + """Download the ``libero_10`` suite of ``nvidia/LIBERO_LeRobot_v3`` if absent.""" + if (_LIBERO_ROOT / "meta" / "info.json").is_file(): + return + rc, out = _run( + [ + "uvx", + "hf@latest", + "download", + "--repo-type", + "dataset", + "nvidia/LIBERO_LeRobot_v3", + "--include", + "libero_10/**", + "--local-dir", + str(_LIBERO_DIR), + "--quiet", + ], + log_dir / "download_libero.log", + ) + assert rc == 0, f"LIBERO_LeRobot_v3 download failed (exit {rc}):\n{out[-2000:]}" + assert (_LIBERO_ROOT / "meta" / "info.json").is_file(), ( + f"LIBERO suite missing {_LIBERO_ROOT}/meta/info.json after download" + ) + + +# --- launch specs ------------------------------------------------------------ + + +# Overrides shared by both action specs: cap the run to a deterministic 10-iter +# single-node capture. Keeps ``model.config.compile.enabled`` at its recipe +# default (ON) to match launch_regression_test.py's nano spec — the loss is +# bit-exact under compile (only grad-norm is perturbed, which we don't assert). +# ``compile_tokenizer`` off just to skip its warmup. shard 4 x replicate 1 fits +# one 4-GPU node so the FSDP reduce-scatter/all-gather stays intra-node (NVLink) +# and reproduces bit-exact. A small packed batch keeps the 10 iters quick. +# +# NB (DROID on Blackwell): on gb200 the GradClip callback's compiled +# ``_fused_nan_to_num`` kernel can fail to launch through torch's static Triton +# launcher at this small shard=4 config ("CUDA driver error: invalid argument") +# — a launcher edge case, not a numeric issue (loss is identical eager). The +# H200/CI path this test targets does not hit it; the LIBERO spec is unaffected +# on either arch. +_COMMON_OVERRIDES: tuple[str, ...] = ( + "trainer.max_iter=10", + "trainer.logging_iter=1", + "trainer.seed=42", + "job.wandb_mode=disabled", + "upload_reproducible_setup=false", + "checkpoint.save_iter=999999", # no checkpoint writes during the capture + "trainer.callbacks.compile_tokenizer.enabled=false", + "model.config.parallelism.data_parallel_shard_degree=4", + "model.config.parallelism.data_parallel_replicate_degree=1", + "dataloader_train.max_samples_per_batch=8", +) + + +@dataclass(frozen=True) +class LaunchSpec: + """A single action-policy launch flow under regression.""" + + key: str + sft_toml: str + extra_hydra_args: tuple[str, ...] + requires_droid_root: bool = False + nproc_per_node: int = 4 + deterministic: bool = True + loss_rtol: float = _DEFAULT_RTOL + loss_atol: float = _DEFAULT_ATOL + # Optional tiered tolerance: each ``(count, rtol, atol)`` applies to the next + # ``count`` iters in order; counts must sum to 10. Empty => uniform default. + loss_tol_bands: tuple[tuple[int, float, float], ...] = () + + +_SPECS: dict[str, LaunchSpec] = { + "action_policy_libero": LaunchSpec( + key="action_policy_libero", + sft_toml="examples/toml/sft_config/action_policy_libero_repro.toml", + extra_hydra_args=_COMMON_OVERRIDES, + ), + "action_policy_droid": LaunchSpec( + key="action_policy_droid", + sft_toml="examples/toml/sft_config/action_policy_droid_repro.toml", + extra_hydra_args=_COMMON_OVERRIDES, + requires_droid_root=True, + ), +} + + +def _run_torchrun(spec: LaunchSpec, run_dir: Path) -> str: + """Invoke the same ``torchrun`` command the launcher shell runs; return the log text.""" + run_dir.mkdir(parents=True, exist_ok=True) + cmd = [ + "torchrun", + f"--nproc_per_node={spec.nproc_per_node}", + f"--master_port={_free_port()}", + "-m", + "cosmos_framework.scripts.train", + f"--sft-toml={spec.sft_toml}", + ] + if spec.deterministic: + cmd.append("--deterministic") + cmd += ["--", *spec.extra_hydra_args] + + rc, out = _run( + cmd, + run_dir / "training.log", + extra_env={ + "PYTHONHASHSEED": "42", + "IMAGINAIRE_OUTPUT_ROOT": str(run_dir / "output"), + "WAN_VAE_PATH": str(_WAN_VAE), + "BASE_CHECKPOINT_PATH": str(_DCP_DIR), + "LIBERO_ROOT": str(_LIBERO_ROOT), + # DROID_ROOT passes through from the caller's env (spec skips if unset). + }, + ) + if rc != 0 and "Done with training" not in out: + pytest.fail( + f"{spec.key}: torchrun failed (exit {rc}) and log lacks 'Done with training'.\nLog tail:\n{out[-3000:]}" + ) + return out + + +def _rank0_losses(text: str) -> list[float]: + """Parse the rank-0 per-iteration ``Loss:`` series (one value per step), in order.""" + vals: list[float] = [] + for m in _VFM_LOSS_RE.finditer(text): + v = float(m.group(1)) + vals.append(v) + return vals + + +# --- fixtures ---------------------------------------------------------------- + + +@pytest.fixture(scope="module", autouse=True) +def _require_4_gpus() -> None: + """Skip the module unless we can launch a 4-GPU training run here.""" + if shutil.which("torchrun") is None: + pytest.skip("torchrun not on PATH — must run inside the training container") + if shutil.which("uvx") is None: + pytest.skip("uvx not on PATH — required to download the dataset / Wan VAE") + try: + import torch + except Exception as exc: # pragma: no cover + pytest.skip(f"torch unavailable ({exc!r})") + if not torch.cuda.is_available() or torch.cuda.device_count() < 4: + pytest.skip(f"requires 4 visible CUDA devices, found {torch.cuda.device_count()}") + + +def _assert_spec_matches_goldens(spec_key: str, tmp_path: Path) -> None: + """Re-run ``spec``'s torchrun command and check the loss series against goldens.""" + spec = _SPECS[spec_key] + if spec.requires_droid_root: + droid_root = os.environ.get("DROID_ROOT", "") + # DROID_ROOT is the versioned merged root (basename is a LEROBOT_ROOTS key); + # the i4 lazy dataset appends success/ (use_success_only), so meta/info.json + # lives under /success, not at the root. + if not droid_root or not (Path(droid_root) / "success" / "meta" / "info.json").is_file(): + pytest.skip( + f"{spec.key}: set DROID_ROOT to a local versioned DROID LeRobot root " + "(basename a LEROBOT_ROOTS key, e.g. droid_plus_lerobot_640x360_20260412, " + "containing success/meta/info.json) to run this spec" + ) + + arch = _detect_arch() + + # Stage shared inputs (cached across runs); LIBERO data only for the libero spec. + _ensure_wan_vae(tmp_path) + _ensure_dcp(tmp_path) + if not spec.requires_droid_root: + _ensure_libero(tmp_path) + + out = _run_torchrun(spec, tmp_path) + assert "Done with training" in out, f"{spec.key}: training did not finish:\n{out[-3000:]}" + losses = _rank0_losses(out) + run_detail = f"\n--- {spec.key} run log (last 3000 chars) ---\n{out[-3000:]}" + assert len(losses) == 10, f"{spec.key}: expected 10 rank-0 losses, parsed {losses}{run_detail}" + assert all(v == v and abs(v) != float("inf") for v in losses), ( + f"{spec.key}: non-finite loss in series {losses}{run_detail}" + ) + + # Refresh path: print captured values for manual copy into ``_GOLDENS``. + if os.environ.get("COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS") == "1": + print(f"\n# --- goldens for arch={arch!r} key={spec.key!r} ---") + print(f'"{spec.key}": {{"loss": {losses}}},') + pytest.skip( + f"captured fresh loss series for arch={arch!r} key={spec.key!r}; copy the printed " + f"dict into _GOLDENS[{arch!r}] at the bottom of action_policy_regression_test.py, " + "then rerun without COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS to assert." + ) + + arch_goldens = _GOLDENS.get(arch) + if not arch_goldens or spec.key not in arch_goldens: + pytest.skip( + f"no goldens for arch={arch!r} key={spec.key!r} yet — capture on this hardware with " + f"COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS=1 (parsed {len(losses)} finite losses this run)" + ) + expected = arch_goldens[spec.key]["loss"] + + bands = spec.loss_tol_bands or ((10, spec.loss_rtol, spec.loss_atol),) + assert sum(c for c, _, _ in bands) == 10, f"{spec.key}: loss_tol_bands must sum to 10" + start = 0 + for count, rtol, atol in bands: + end = start + count + assert losses[start:end] == pytest.approx(expected[start:end], rel=rtol, abs=atol), ( + f"{spec.key} ({arch}): rank-0 loss[{start}:{end}] (rel={rtol}, abs={atol}) " + f"does not match goldens\n got : {losses[start:end]}\n" + f" expected: {expected[start:end]}{run_detail}" + ) + start = end + + +# --- tests ------------------------------------------------------------------- + + +if MAX_GPUS == 4: + + @pytest.mark.level(2) + @pytest.mark.gpus(4) + @pytest.mark.parametrize("spec_key", list(_SPECS), ids=list(_SPECS)) + def test_action_policy_regression(spec_key: str, tmp_path: Path) -> None: + """Deterministic 10-iter action-policy SFT; assert the rank-0 loss goldens.""" + _assert_spec_matches_goldens(spec_key, tmp_path) + + +# Goldens keyed by GPU arch then ``LaunchSpec.key``. ``loss`` is the rank-0 +# per-step series over the 10 deterministic iters. Refresh on the target arch +# with ``COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS=1`` (see the module docstring). +# The test skips (not fails) for any arch/spec without an entry, so goldens can +# land incrementally as they are captured on each arch. +# +# Captured with torch.compile ON, --deterministic, seed 42, single-node +# data_parallel_shard_degree=4 (intra-node NVLink FSDP reduction is bit-exact). +_GOLDENS: dict[str, dict[str, dict[str, list[float]]]] = { + # H200 (Hopper) CI arch. LIBERO is the primary numerical golden; the DROID + # spec needs its dataset (DROID_ROOT), so it skips unless one is provided. + "h100": { + "action_policy_libero": { + "loss": [ + 15.8107, + 15.2467, + 15.9856, + 16.5306, + 14.3738, + 16.1460, + 16.6093, + 14.8846, + 16.0632, + 16.6449, + ], + }, + }, +} + + +if __name__ == "__main__": # pragma: no cover — manual driver + sys.exit(pytest.main([__file__, "-v", "-s", "-o", "addopts="])) From 3017efb48826b69f9582ac6e3c001f50d1401067 Mon Sep 17 00:00:00 2001 From: Maosheng Liao Date: Wed, 8 Jul 2026 14:08:43 +0800 Subject: [PATCH 17/26] toml_config: add [job].upload_reproducible_setup, default False for OSS (#92) The VLM base config defaults upload_reproducible_setup=True, so a run launched through the structured-TOML flow attempts the reproducible-setup S3 upload (and wandb save_s3, which interpolate from ${upload_reproducible_setup}) by default. OSS users mostly have no S3 access and no way to turn it off from the TOML. Add a [job].upload_reproducible_setup knob (default False) that maps to the top-level config.upload_reproducible_setup: - JobConfig gains the field (default False). - PATH_REMAPS (vfm + vlm) hoist ("job","upload_reproducible_setup") to the top-level ("upload_reproducible_setup",). - load_experiment_from_toml keeps the validated model and always injects the resolved value into raw before building overrides, so the override is emitted as false even when the TOML omits it (build_hydra_overrides walks the raw dict, so an omitted field would otherwise emit nothing and the base True would win). Users opt back into S3 upload with `upload_reproducible_setup = true`. Docs (sft_config.md) + all example [job] blocks updated; adds tests/toml_config_test.py. --------- Co-authored-by: Claude Opus 4.8 (1M context) --- .../configs/toml_config/sft_config.py | 20 ++++++++++++++++++- .../configs/toml_config/toml_config_helper.py | 6 ++++++ docs/sft_config.md | 20 ++++++++++--------- 3 files changed, 36 insertions(+), 10 deletions(-) diff --git a/cosmos_framework/configs/toml_config/sft_config.py b/cosmos_framework/configs/toml_config/sft_config.py index 9ce57728..04d0efbc 100644 --- a/cosmos_framework/configs/toml_config/sft_config.py +++ b/cosmos_framework/configs/toml_config/sft_config.py @@ -83,6 +83,17 @@ class JobConfig(BaseModel): "'disabled' (no wandb at all)." ), ) + upload_reproducible_setup: bool = Field( + default=False, + description=( + "Upload the reproducible-setup bundle and wandb save_s3 artifacts " + "to S3. Defaults False (OSS: no S3 access). Set True only when S3 " + "upload is configured. Remapped out of [job] to the top-level " + "config.upload_reproducible_setup, overriding the base config value " + "(the VLM base defaults it True). Always emitted by " + "load_experiment_from_toml, so omitting it forces False." + ), + ) # ---------------------------------------------------------------- model @@ -721,7 +732,14 @@ def load_experiment_from_toml( # Validate structure against the pydantic schema (raises ValidationError on # unknown keys because of ``extra="forbid"``). - SFTExperimentConfig.model_validate(raw) + cfg = SFTExperimentConfig.model_validate(raw) + + # ``build_hydra_overrides`` walks the *raw* dict, so an omitted field emits + # no override and the base config's value wins. ``upload_reproducible_setup`` + # must instead default to False (OSS: no S3), overriding the VLM base's True. + # Inject the pydantic-resolved value back into raw so the override is always + # emitted — the default lives in exactly one place (JobConfig.Field). + raw.setdefault("job", {})["upload_reproducible_setup"] = cfg.job.upload_reproducible_setup task = raw.get("job", {}).get("task", "vfm") try: diff --git a/cosmos_framework/configs/toml_config/toml_config_helper.py b/cosmos_framework/configs/toml_config/toml_config_helper.py index 211e9d3b..4d1535c5 100644 --- a/cosmos_framework/configs/toml_config/toml_config_helper.py +++ b/cosmos_framework/configs/toml_config/toml_config_helper.py @@ -47,6 +47,9 @@ # ``model.config.`` in the Hydra tree. ``attn_implementation`` is a # VLM-only knob — skip it on VFM. Other sections pass through. "vfm": { + # [job].upload_reproducible_setup lives at the top-level config field, + # not config.job.* — hoist it out of the job section. + ("job", "upload_reproducible_setup"): ("upload_reproducible_setup",), ("model", "attn_implementation"): None, ("model", "backbone"): None, # VLM-only — VFM has no model.config.backbone # Per-caption token cap lives on the nested SFT dataset, not a top-level @@ -63,6 +66,9 @@ # tasks — so the catch-all ``("model",) -> ("model", "config")`` rule # routes them uniformly. Fields with no VLM analog map to ``None`` (skip). "vlm": { + # [job].upload_reproducible_setup lives at the top-level config field, + # not config.job.* — hoist it out of the job section. + ("job", "upload_reproducible_setup"): ("upload_reproducible_setup",), # No VLM analog — skip these leaves ("model", "max_num_tokens_after_packing"): None, ("model", "joint_attn_implementation"): None, diff --git a/docs/sft_config.md b/docs/sft_config.md index db0bb83b..241f90fe 100644 --- a/docs/sft_config.md +++ b/docs/sft_config.md @@ -70,14 +70,15 @@ The full pipeline (dataloader class, dataset wiring, model_instance LazyCall, et Run identity + meta-fields that pick the Hydra config tree to load. -| field | default | description | -| ------------ | ------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `task` | `"vfm"` | **META** — chooses which `make_config()` to call: `"vfm"` → `cosmos_framework/configs/base/config.py`, `"vlm"` → `cosmos_framework/configs/base/reasoner/config.py`. Also picks the path-remap rules in `toml_config_helper.PATH_REMAPS`. | -| `experiment` | `""` | **META** — names the Hydra experiment LazyDict registered in `ConfigStore` under `experiment/`. Resolved at load time via `experiment=` (e.g. `vision_sft_nano`). | -| `project` | `""` | W&B project (team-level bucket). Flows to `config.job.project`. | -| `group` | `""` | W&B sub-label for clustering related runs (e.g. `"sft"`). Flows to `config.job.group`. | -| `name` | `""` | W&B run name; forms part of the output dir `$IMAGINAIRE_OUTPUT_ROOT////`. Leave empty (or use `${now:%Y-%m-%d}_${now:%H-%M-%S}`) for auto-timestamped subdir. | -| `wandb_mode` | `"disabled"` | `"online"` (real-time, needs `WANDB_API_KEY`), `"offline"` (log locally, sync later via `wandb sync`), or `"disabled"`. | +| field | default | description | +| --------------------------- | ------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `task` | `"vfm"` | **META** — chooses which `make_config()` to call: `"vfm"` → `cosmos_framework/configs/base/config.py`, `"vlm"` → `cosmos_framework/configs/base/reasoner/config.py`. Also picks the path-remap rules in `toml_config_helper.PATH_REMAPS`. | +| `experiment` | `""` | **META** — names the Hydra experiment LazyDict registered in `ConfigStore` under `experiment/`. Resolved at load time via `experiment=` (e.g. `vision_sft_nano`). | +| `project` | `""` | W&B project (team-level bucket). Flows to `config.job.project`. | +| `group` | `""` | W&B sub-label for clustering related runs (e.g. `"sft"`). Flows to `config.job.group`. | +| `name` | `""` | W&B run name; forms part of the output dir `$IMAGINAIRE_OUTPUT_ROOT////`. Leave empty (or use `${now:%Y-%m-%d}_${now:%H-%M-%S}`) for auto-timestamped subdir. | +| `wandb_mode` | `"disabled"` | `"online"` (real-time, needs `WANDB_API_KEY`), `"offline"` (log locally, sync later via `wandb sync`), or `"disabled"`. | +| `upload_reproducible_setup` | `false` | Upload the reproducible-setup bundle + wandb `save_s3` artifacts to S3. **Defaults `false`** (OSS: no S3 access); set `true` only when S3 upload is configured. Remapped out of `[job]` to the top-level `config.upload_reproducible_setup`, overriding the base config value (the VLM base defaults it `true`). Always emitted by `load_experiment_from_toml`, so omitting it forces `false`. | ## `[model]` @@ -280,6 +281,7 @@ The same TOML key lands at different Hydra paths depending on `[job].task`: | TOML path | VFM (`task="vfm"`) Hydra path | VLM (`task="vlm"`) Hydra path | | ---------------------------------------------------------------------------------------- | ----------------------------------------- | ----------------------------------------- | +| `job.upload_reproducible_setup` | `upload_reproducible_setup` | `upload_reproducible_setup` | | `model.` | `model.config.` | `model.config.` | | `model.parallelism.*` | `model.config.parallelism.*` | `model.config.parallelism.*` | | `model.compile.*` | `model.config.compile.*` | `model.config.compile.*` | @@ -310,7 +312,7 @@ A few useful knobs aren't currently modeled by `SFTExperimentConfig` because the `load_experiment_from_toml(toml_path, extra_overrides)` (in `sft_config.py`) is the end-to-end loader. It: 1. Reads the TOML with `tomllib`. -2. Validates the parsed dict against `SFTExperimentConfig` (raises `ValidationError` on unknown keys). +2. Validates the parsed dict against `SFTExperimentConfig` (raises `ValidationError` on unknown keys), then injects the pydantic-resolved `job.upload_reproducible_setup` back into the raw dict so it is always emitted (defaulting to `false`) even when the TOML omits it. 3. Picks the base config from `[job].task`: `TASK_TO_BASE_CONFIG["vfm"|"vlm"]`. 4. Calls `build_hydra_overrides(raw)` to produce a `["--", "experiment=", "k.p=v", …]` list with per-task remaps applied and MISSING values filtered. `[custom]` is skipped here (it is injected verbatim in step 7, not per-leaf-remapped). 5. Appends `extra_overrides` (CLI tail) so they take precedence over the TOML. From b2237fd01a020a5a60b2f59a5c4071a042d9dc9b Mon Sep 17 00:00:00 2001 From: Dinghow Yang Date: Thu, 9 Jul 2026 14:56:38 +0800 Subject: [PATCH 18/26] convert_model_to_diffusers: export via official diffusers, drop diffusers-cosmos3 shim (#98) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What Replace the `packages/diffusers-cosmos3` shim with official diffusers (≥ 0.39.0) in the checkpoint exporter, and delete the package. ## Why - diffusers 0.39.0 ships Cosmos3 natively (`Cosmos3OmniPipeline`, `Cosmos3OmniTransformer`, `Cosmos3AVAEAudioTokenizer`); the shim was a pre-upstream stand-in. - The shim's reimplemented denoising forward has a small per-step numeric drift that compounds into visibly soft t2i images — nothing should run inference through it. - The exporter's key handling was stale against the current flat key layout (`model.`-nested `q_proj`/`vae2llm` vs flat `to_q`/`proj_in`) and could not strict-load. ## Changes - Rewrite `scripts/_convert_model_to_diffusers.py` against the official classes; keep the repo-specific extensions (action-projection export, `vision_encoder/` sidecar for the transformers/vLLM consumers). - Save `sound_tokenizer/` as a real `Cosmos3AVAEAudioTokenizer` pipeline component instead of hand-writing safetensors + patching `model_index.json`. - OSS-path fixes surfaced by e2e runs: consume the HF diffusers-layout `sound_tokenizer/` directly, load vision weights from `vision_encoder/`-subfolder sources, resolve the text config via `get_text_config()` for nested `Qwen3VLConfig`. - Deps: diffusers ≥ 0.39.0 (lock 0.37.0 → 0.39.0, safetensors 0.7.0 → 0.8.0), remove the shim from extras / uv sources, drop the now-fixed diffusers audit ignores. ## Relation to upstream `scripts/convert_cosmos3_to_diffusers.py` Upstream diffusers ships a converter for the same task, and this rewrite adopts its verified key-remap table. It cannot be used directly here: - its imports target the internal i4 module namespace (`projects.cosmos3.vfm.*`), so it does not run against this repo; - it silently drops action-projection weights (it never reads the action config, and strict-load still passes because neither side has the keys); - it does not produce the dual-purpose repo this exporter ships (no `vision_encoder/` sidecar, no top-level vLLM config / unified safetensors index). ## Verification - Key remap round-trips exactly against the real diffusers 0.39.0 `Cosmos3OmniTransformer` state-dict key set (action + sound enabled), and inverts cleanly through the native inference loader's mapping. - E2E, public checkpoint: Cosmos3-Nano → export → official `Cosmos3OmniPipeline` t2i sharp; native inference re-read of the export sharp. - E2E, real training DCP: an i4-trainer 8B DPO checkpoint (flat `net.*`/`net_ema.*` keys, t2i-only) converts and generates sharp t2i through the official pipeline (driven via the inner exporter; the CLI wrapper asserts `action_gen`/`sound_gen` and targets full omni checkpoints). ## Compat notes - Previously exported repos (`model_index: Cosmos3OmniDiffusersPipeline`) still load via an explicit `Cosmos3OmniPipeline.from_pretrained(...)`. - Loading without `cosmos_guardrail` installed requires `enable_safety_checker=False` (matches upstream semantics). --------- Co-authored-by: Claude Fable 5 --- .../scripts/_convert_model_to_diffusers.py | 302 +-- .../scripts/convert_model_to_diffusers.py | 26 +- docs/code_structure.md | 8 +- packages/diffusers-cosmos3/README.md | 1 - .../diffusers_cosmos3/__init__.py | 17 - .../diffusers_cosmos3/pipeline.py | 1124 ------------ .../diffusers_cosmos3/sequence_packing.py | 1626 ----------------- .../diffusers_cosmos3/transformer.py | 819 --------- packages/diffusers-cosmos3/pyproject.toml | 13 - .../scripts/inference_example.py | 79 - packages/diffusers-cosmos3/uv.lock | 924 ---------- pyproject.toml | 8 +- uv.lock | 61 +- 13 files changed, 204 insertions(+), 4804 deletions(-) delete mode 100644 packages/diffusers-cosmos3/README.md delete mode 100644 packages/diffusers-cosmos3/diffusers_cosmos3/__init__.py delete mode 100644 packages/diffusers-cosmos3/diffusers_cosmos3/pipeline.py delete mode 100644 packages/diffusers-cosmos3/diffusers_cosmos3/sequence_packing.py delete mode 100644 packages/diffusers-cosmos3/diffusers_cosmos3/transformer.py delete mode 100644 packages/diffusers-cosmos3/pyproject.toml delete mode 100755 packages/diffusers-cosmos3/scripts/inference_example.py delete mode 100644 packages/diffusers-cosmos3/uv.lock diff --git a/cosmos_framework/scripts/_convert_model_to_diffusers.py b/cosmos_framework/scripts/_convert_model_to_diffusers.py index e1c14016..da5549ab 100644 --- a/cosmos_framework/scripts/_convert_model_to_diffusers.py +++ b/cosmos_framework/scripts/_convert_model_to_diffusers.py @@ -11,13 +11,18 @@ import pydantic import torch from accelerate import init_empty_weights -from diffusers import AutoencoderKLWan, UniPCMultistepScheduler -from diffusers_cosmos3 import Cosmos3OmniDiffusersPipeline, Cosmos3OmniTransformer +from diffusers import ( + AutoencoderKLWan, + Cosmos3AVAEAudioTokenizer, + Cosmos3OmniPipeline, + Cosmos3OmniTransformer, + UniPCMultistepScheduler, +) from transformers import AutoConfig, AutoTokenizer from cosmos_framework.inference.model import Cosmos3OmniModel -from cosmos_framework.utils import log from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel +from cosmos_framework.utils import log DEFAULT_SOUND_TOKENIZER_CONFIG = { "model_type": "autoencoder_v2", @@ -62,14 +67,6 @@ "latent_std": None, } -SOUND_TOKENIZER_MODEL_INDEX_ENTRY = [ - "diffusers", - "Cosmos3AVAEAudioTokenizer", -] - -LEGACY_SOUND_TOKENIZER_CHECKPOINT_NAME = "model.safetensors" -DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME = "diffusion_pytorch_model.safetensors" - # Wrapper prefixes that may appear on every key of a legacy AVAE state dict # (DDP wrappers, full-model saves, training exports). Stripped iteratively # until each key reaches a recognised target prefix. @@ -80,6 +77,28 @@ # [snake1, conv1, snake2, conv2]; map sub-index → named attribute. _SOUND_TOKENIZER_RES_UNIT_INNER_NAMES = {0: "snake1", 1: "conv1", 2: "snake2", 3: "conv2"} +# The source language_model nests its transformer stack under a `model.` +# attribute (HF Qwen-style); the diffusers `Cosmos3OmniTransformer` holds +# those layers flat, so the leading `model.` prefix is stripped. The +# PackedAttentionMoT projections are renamed from the source Qwen-style +# names (`q_proj`/… plus cosmos-specific `*_moe_gen`) to the diffusers +# AttentionModuleMixin canonical names. Order matters: the `*_moe_gen` +# substrings must be substituted before the plain ones. +_ATTN_KEY_REMAP = ( + (".q_proj_moe_gen.", ".add_q_proj."), + (".k_proj_moe_gen.", ".add_k_proj."), + (".v_proj_moe_gen.", ".add_v_proj."), + (".o_proj_moe_gen.", ".to_add_out."), + (".q_norm_moe_gen.", ".norm_added_q."), + (".k_norm_moe_gen.", ".norm_added_k."), + (".q_proj.", ".to_q."), + (".k_proj.", ".to_k."), + (".v_proj.", ".to_v."), + (".o_proj.", ".to_out."), + (".q_norm.", ".norm_q."), + (".k_norm.", ".norm_k."), +) + # Legacy TimestepEmbedder stored its MLP as `nn.Sequential([Linear, SiLU, Linear])`, # so state-dict keys were `mlp.0.*` / `mlp.2.*`. The diffusers `TimestepEmbedding` # stores them as named attributes `linear_1` / `linear_2`. Index 1 (SiLU) has no @@ -93,6 +112,7 @@ DEFAULT_VISION_ENCODER_MODEL = "Qwen/Qwen3-VL-8B-Instruct" VISION_ENCODER_CHECKPOINT_PREFIX = "model.visual." +VISION_ENCODER_CHECKPOINT_SUBFOLDER = "vision_encoder" def _get_config_value(*configs, name, default=None): @@ -108,6 +128,19 @@ def _get_config_value(*configs, name, default=None): return default +def _remap_language_model_key(key: str) -> str: + """Rename a source language-model state-dict key to the diffusers + `Cosmos3OmniTransformer` layout: strip the leading `model.` prefix and + apply the attention-projection renames from `_ATTN_KEY_REMAP`. + """ + key = key.removeprefix("model.") + for old, new in _ATTN_KEY_REMAP: + if old in key: + key = key.replace(old, new) + break + return key + + def _load_sound_tokenizer_state_dict(checkpoint_path: pathlib.Path) -> dict[str, torch.Tensor]: if checkpoint_path.suffix == ".safetensors": try: @@ -154,13 +187,15 @@ def _sound_tokenizer_strip_per_key_prefixes(state_dict: dict[str, torch.Tensor]) return out -def _sound_tokenizer_filter_decoder(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Drop encoder and bottleneck keys — only the decoder is used at inference.""" - return {k: v for k, v in state_dict.items() if k.startswith("decoder.")} +def _sound_tokenizer_filter_supported_modules(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Keep only encoder/decoder keys — `Cosmos3AVAEAudioTokenizer` supports + the parameter-free `vae` bottleneck only, so bottleneck keys are dropped. + """ + return {k: v for k, v in state_dict.items() if k.startswith(("encoder.", "decoder."))} def _sound_tokenizer_infer_num_blocks(state_dict: dict[str, torch.Tensor]) -> int: - """Count the OobleckDecoderBlocks present in a flat-`Sequential` legacy + """Count the decoder blocks present in a flat-`Sequential` legacy decoder state dict by spotting `decoder.layers.{N}.layers.{M}.*` keys. """ block_indices: set[int] = set() @@ -173,7 +208,7 @@ def _sound_tokenizer_infer_num_blocks(state_dict: dict[str, torch.Tensor]) -> in def _sound_tokenizer_remap_flat_layout(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """Rewrite a legacy `decoder.layers.{N}.*` flat-`Sequential` layout to the - diffusers OobleckDecoder named-attribute layout. + diffusers Cosmos3AudioDecoder named-attribute layout. The legacy decoder is Sequential([conv1, block_0, ..., block_{N-1}, snake1, conv2]) @@ -243,7 +278,11 @@ def _sound_tokenizer_reshape_snake_params(state_dict: dict[str, torch.Tensor]) - """ out: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): - if (key.endswith(".alpha") or key.endswith(".beta")) and value.ndim == 1: + if ( + key.startswith(("encoder.", "decoder.")) + and (key.endswith(".alpha") or key.endswith(".beta")) + and value.ndim == 1 + ): value = value.unsqueeze(0).unsqueeze(-1).contiguous() out[key] = value return out @@ -252,7 +291,7 @@ def _sound_tokenizer_reshape_snake_params(state_dict: dict[str, torch.Tensor]) - def _sound_tokenizer_reapply_weight_norm(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """If a Conv layer has the folded `.weight` tensor but neither `.weight_g` nor `.weight_v`, reconstruct the pair so the resulting checkpoint can be - loaded into the weight-norm-wrapped diffusers OobleckDecoder. + loaded into the weight-norm-wrapped diffusers Cosmos3AudioDecoder. `weight_norm` with `dim=0` parameterises `weight = g * v / ||v||`. Setting `v = weight` and `g = ||weight||` along all non-zero axes is an exact @@ -262,7 +301,11 @@ def _sound_tokenizer_reapply_weight_norm(state_dict: dict[str, torch.Tensor]) -> candidate_keys = [ key for key in state_dict - if key.endswith(".weight") and any(f".{layer}." in key for layer in ("conv1", "conv2", "conv_t1")) + if key.endswith(".weight") + and ( + any(f".{layer}." in key for layer in ("conv1", "conv2", "conv_t1")) + or re.fullmatch(r"encoder\.layers\.\d+\.weight", key) + ) ] for key in candidate_keys: stem = key[: -len(".weight")] @@ -279,24 +322,17 @@ def _sound_tokenizer_reapply_weight_norm(state_dict: dict[str, torch.Tensor]) -> return out -def _remap_time_embedder_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Remap a legacy `TimestepEmbedder` state dict (with `nn.Sequential`-style - `mlp.0` / `mlp.2` keys) to the diffusers `TimestepEmbedding` layout - (`linear_1` / `linear_2`). Keys not in `_TIME_EMBEDDER_KEY_REMAP` pass - through unchanged. - """ - return {_TIME_EMBEDDER_KEY_REMAP.get(key, key): value for key, value in state_dict.items()} - - def _remap_sound_tokenizer_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """Apply the full legacy → diffusers conversion pipeline to a sound - tokenizer state dict: strip prefixes, drop non-decoder keys, remap the - flat `nn.Sequential` layout to named attributes, reshape Snake1d params, - and reconstruct `weight_g` / `weight_v` for any folded conv weights. + tokenizer state dict: strip prefixes, drop unsupported (bottleneck) keys, + remap the flat `nn.Sequential` layout to named attributes, reshape Snake1d + params, and reconstruct `weight_g` / `weight_v` for any folded conv weights. """ state_dict = _sound_tokenizer_strip_per_key_prefixes(state_dict) - state_dict = _sound_tokenizer_filter_decoder(state_dict) + state_dict = _sound_tokenizer_filter_supported_modules(state_dict) if not state_dict: + raise RuntimeError("Sound tokenizer state dict has no `encoder.*`/`decoder.*` keys after prefix stripping.") + if not any(key.startswith("decoder.") for key in state_dict): raise RuntimeError("Sound tokenizer state dict has no `decoder.*` keys after prefix stripping.") state_dict = _sound_tokenizer_remap_flat_layout(state_dict) state_dict = _sound_tokenizer_reshape_snake_params(state_dict) @@ -306,57 +342,39 @@ def _remap_sound_tokenizer_state_dict(state_dict: dict[str, torch.Tensor]) -> di return state_dict -def _load_sound_tokenizer_config(config_path: pathlib.Path | None, fallback_config_path: pathlib.Path) -> dict: - selected_config_path = config_path - if selected_config_path is None and fallback_config_path.exists(): - selected_config_path = fallback_config_path - if selected_config_path is None: +def _load_sound_tokenizer_config(config_path: pathlib.Path | None) -> dict: + if config_path is None: return dict(DEFAULT_SOUND_TOKENIZER_CONFIG) - with open(selected_config_path, encoding="utf-8") as f: + with open(config_path, encoding="utf-8") as f: return json.load(f) -def _save_sound_tokenizer( - output_dir: pathlib.Path, +def _build_sound_tokenizer( checkpoint_path: pathlib.Path, config_path: pathlib.Path | None, - remap_keys: bool = False, -) -> None: - try: - from safetensors.torch import save_file - except ImportError as exc: - raise ImportError("Saving AVAE tokenizer weights requires safetensors.") from exc - - sound_tokenizer_dir = output_dir / "sound_tokenizer" - sound_tokenizer_dir.mkdir(parents=True, exist_ok=True) - - config = _load_sound_tokenizer_config(config_path, sound_tokenizer_dir / "config.json") - with open(sound_tokenizer_dir / "config.json", "w", encoding="utf-8") as f: - json.dump(config, f, indent=4) - f.write("\n") +) -> Cosmos3AVAEAudioTokenizer: + config = _load_sound_tokenizer_config(config_path) log.info(f"Loading AVAE sound tokenizer weights from {checkpoint_path} …") - state_dict = _load_sound_tokenizer_state_dict(checkpoint_path) - if remap_keys: - log.info("Remapping AVAE sound tokenizer state dict to diffusers OobleckDecoder layout …") - state_dict = _remap_sound_tokenizer_state_dict(state_dict) - output_filename = ( - DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME if remap_keys else LEGACY_SOUND_TOKENIZER_CHECKPOINT_NAME + raw_state_dict = _load_sound_tokenizer_state_dict(checkpoint_path) + state_dict = _remap_sound_tokenizer_state_dict(raw_state_dict) + has_encoder = any(key.startswith("encoder.") for key in state_dict) + log.info( + f"Remapped {len(raw_state_dict)} → {len(state_dict)} tokenizer keys " + f"({'encoder+decoder' if has_encoder else 'decoder-only'})." ) - log.info(f"Saving AVAE sound tokenizer to {sound_tokenizer_dir / output_filename} …") - save_file(state_dict, str(sound_tokenizer_dir / output_filename), metadata={"format": "pt"}) - -def _add_sound_tokenizer_to_model_index(output_dir: pathlib.Path) -> None: - model_index_path = output_dir / "model_index.json" - if not model_index_path.exists(): - return - with open(model_index_path, encoding="utf-8") as f: - model_index = json.load(f) - model_index["sound_tokenizer"] = SOUND_TOKENIZER_MODEL_INDEX_ENTRY - with open(model_index_path, "w", encoding="utf-8") as f: - json.dump(model_index, f, indent=2) - f.write("\n") + # `Cosmos3AVAEAudioTokenizer` accepts exactly the keys of the default + # config; unknown keys in a source config JSON are ignored. + tokenizer_kwargs = {name: config.get(name, default) for name, default in DEFAULT_SOUND_TOKENIZER_CONFIG.items()} + sound_tokenizer = Cosmos3AVAEAudioTokenizer(**tokenizer_kwargs, encoder_enabled=has_encoder) + load_result = sound_tokenizer.load_state_dict(state_dict, strict=True) + if load_result.missing_keys or load_result.unexpected_keys: + raise RuntimeError( + "Cosmos3 AVAE sound tokenizer load did not match strictly: " + f"missing={load_result.missing_keys}, unexpected={load_result.unexpected_keys}." + ) + return sound_tokenizer def _checkpoint_weight_map(checkpoint_path: pathlib.Path) -> dict[str, str]: @@ -372,6 +390,20 @@ def _checkpoint_has_weight_prefix(checkpoint_path: pathlib.Path, prefix: str) -> return any(key.startswith(prefix) for key in _checkpoint_weight_map(checkpoint_path)) +def _checkpoint_vision_subfolder_files(checkpoint_path: pathlib.Path) -> dict[str, list[str]]: + """Group root-index keys stored under vision_encoder/ by shard file. + + Diffusers-layout source checkpoints keep bare Qwen3VLVisionModel keys + (`blocks.*`, `patch_embed.*`, …) in the root weight map, mapped to files + under the vision_encoder/ subfolder. + """ + files_to_keys: dict[str, list[str]] = {} + for key, filename in _checkpoint_weight_map(checkpoint_path).items(): + if filename.replace("\\", "/").startswith(f"{VISION_ENCODER_CHECKPOINT_SUBFOLDER}/"): + files_to_keys.setdefault(filename, []).append(key) + return files_to_keys + + def _load_prefixed_safetensors_state_dict(checkpoint_path: pathlib.Path, prefix: str) -> dict[str, torch.Tensor]: try: from safetensors import safe_open @@ -448,6 +480,29 @@ def _build_vision_encoder( return vision_encoder.to(dtype=dtype) +def _load_vision_subfolder_state_dict(checkpoint_path: pathlib.Path) -> dict[str, torch.Tensor]: + try: + from safetensors import safe_open + except ImportError as exc: + raise ImportError("Loading sharded safetensors vision weights requires safetensors.") from exc + + files_to_keys = _checkpoint_vision_subfolder_files(checkpoint_path) + state_dict: dict[str, torch.Tensor] = {} + for filename, keys in sorted(files_to_keys.items()): + shard_path = checkpoint_path / filename + with safe_open(str(shard_path), framework="pt", device="cpu") as shard: + for key in sorted(keys): + state_dict[key] = shard.get_tensor(key).detach().cpu().contiguous() + + if not state_dict: + raise RuntimeError( + f"No vision encoder tensors found under {VISION_ENCODER_CHECKPOINT_SUBFOLDER}/. " + "If the source checkpoint has no Qwen3-VL visual weights (e.g. a vision-generation-only " + "post-training checkpoint), pass --skip-vision-encoder to export without vision_encoder/." + ) + return state_dict + + def _load_vision_encoder( checkpoint_path: pathlib.Path, source_model, @@ -455,11 +510,14 @@ def _load_vision_encoder( dtype: torch.dtype, ): state_dict = _get_source_vision_state_dict(source_model) - if state_dict is None: + if state_dict is not None: + log.info("Extracting Qwen3-VL vision encoder weights from loaded source model …") + elif _checkpoint_has_weight_prefix(checkpoint_path, VISION_ENCODER_CHECKPOINT_PREFIX): log.info(f"Loading Qwen3-VL vision encoder weights from {checkpoint_path} …") state_dict = _load_prefixed_safetensors_state_dict(checkpoint_path, VISION_ENCODER_CHECKPOINT_PREFIX) else: - log.info("Extracting Qwen3-VL vision encoder weights from loaded source model …") + log.info(f"Loading Qwen3-VL vision encoder weights from {checkpoint_path}/vision_encoder …") + state_dict = _load_vision_subfolder_state_dict(checkpoint_path) log.info(f"Building Qwen3-VL vision encoder from {model_name_or_path} …") return _build_vision_encoder(state_dict, model_name_or_path, dtype) @@ -498,21 +556,9 @@ class Args(pydantic.BaseModel): sound_tokenizer_path: str | None = None """Optional AVAE sound tokenizer checkpoint to save under sound_tokenizer/.""" sound_tokenizer_config_path: str | None = None - """Optional AVAE config JSON to save under sound_tokenizer/config.json.""" + """Optional AVAE config JSON describing the sound tokenizer architecture.""" include_sound_tokenizer: bool = False """Require saving sound_tokenizer/ even if the source transformer is video-only.""" - remap_sound_tokenizer_keys: bool = False - """Convert the legacy AVAE state dict into the diffusers OobleckDecoder - layout (strips wrapper prefixes, drops encoder/bottleneck keys, remaps the - flat `nn.Sequential` decoder to named attributes, reshapes Snake1d - alpha/beta, and reconstructs `weight_g` / `weight_v` for folded conv - weights). When enabled, the saved file is `diffusion_pytorch_model.safetensors` - instead of `model.safetensors`.""" - remap_time_embedder_keys: bool = False - """Remap the transformer's `time_embedder` state dict from the legacy - `nn.Sequential` layout (`mlp.0.*` / `mlp.2.*`) to the diffusers - `TimestepEmbedding` layout (`linear_1.*` / `linear_2.*`). Off by default — - keys are forwarded verbatim.""" vision_encoder_model: str = DEFAULT_VISION_ENCODER_MODEL """Qwen3-VL model/config to instantiate model.visual.* weights.""" skip_vision_encoder: bool = False @@ -548,10 +594,11 @@ def convert_model_to_diffusers(args: Args) -> None: vae2llm = _tmp.net.vae2llm llm2vae = _tmp.net.llm2vae time_embedder = _tmp.net.time_embedder - lm_cfg = _tmp.net.language_model.config + # The language model may carry a nested VL config (e.g. Qwen3VLConfig); + # the text-model fields read below live on its text config. + lm_cfg = _tmp.net.language_model.config.get_text_config() net_cfg = _tmp.net.config model_cfg = _tmp.config - vlm_cfg = _tmp.net.config.vlm_config patch_latent_dim = _tmp.net.patch_latent_dim hidden_size = _tmp.net.hidden_size num_attention_heads = _tmp.net.num_heads @@ -561,17 +608,13 @@ def convert_model_to_diffusers(args: Args) -> None: latent_patch_size = _tmp.net.latent_patch_size latent_channel = _tmp.net.latent_channel timestep_scale = _tmp.net.timestep_scale - use_moe = _tmp.net.use_moe - joint_attn_implementation = net_cfg.joint_attn_implementation base_fps = int(net_cfg.base_fps) enable_fps_modulation = net_cfg.enable_fps_modulation max_action_dim = _tmp.config.max_action_dim - position_embedding_type = net_cfg.position_embedding_type unified_3d_mrope_reset_spatial_ids = _tmp.config.diffusion_expert_config.unified_3d_mrope_reset_spatial_ids unified_3d_mrope_temporal_modality_margin = ( _tmp.config.diffusion_expert_config.unified_3d_mrope_temporal_modality_margin ) - video_temporal_causal = net_cfg.video_temporal_causal action2llm = getattr(_tmp.net, "action2llm", None) llm2action = getattr(_tmp.net, "llm2action", None) action_modality_embed = getattr(_tmp.net, "action_modality_embed", None) @@ -598,9 +641,6 @@ def convert_model_to_diffusers(args: Args) -> None: if sound_dim is None and sound2llm is not None: sound_dim = sound2llm.in_features sound_latent_fps = _get_config_value(net_cfg, model_cfg, name="sound_latent_fps", default=25.0) - temporal_compression_factor_sound = _get_config_value( - net_cfg, model_cfg, name="temporal_compression_factor_sound", default=1 - ) if sound_gen: missing_sound_modules = [ name @@ -634,7 +674,9 @@ def convert_model_to_diffusers(args: Args) -> None: f"action projection weights: {missing_action_modules}." ) - has_vision_encoder_weights = _checkpoint_has_weight_prefix(checkpoint_path, VISION_ENCODER_CHECKPOINT_PREFIX) + has_vision_encoder_weights = _checkpoint_has_weight_prefix( + checkpoint_path, VISION_ENCODER_CHECKPOINT_PREFIX + ) or bool(_checkpoint_vision_subfolder_files(checkpoint_path)) vision_gen = bool( _get_config_value(net_cfg, model_cfg, name="vision_gen", default=False) or has_vision_encoder_weights ) @@ -653,66 +695,48 @@ def convert_model_to_diffusers(args: Args) -> None: attention_dropout=lm_cfg.attention_dropout, base_fps=base_fps, enable_fps_modulation=enable_fps_modulation, - freeze_und=vlm_cfg.freeze_und, head_dim=head_dim, - hidden_act=lm_cfg.hidden_act, hidden_size=hidden_size, - initializer_range=lm_cfg.initializer_range, intermediate_size=lm_cfg.intermediate_size, - joint_attn_implementation=joint_attn_implementation, latent_channel=latent_channel, latent_patch_size=latent_patch_size, action_dim=action_dim, action_gen=action_gen, - max_action_dim=max_action_dim, - max_position_embeddings=lm_cfg.max_position_embeddings, - model_type=lm_cfg.model_type, num_embodiment_domains=num_embodiment_domains, num_attention_heads=num_attention_heads, num_hidden_layers=num_hidden_layers, num_key_value_heads=num_key_value_heads, patch_latent_dim=patch_latent_dim, - position_embedding_type=position_embedding_type, - qk_norm_for_diffusion=True, - qk_norm_for_text=vlm_cfg.qk_norm_for_text, rms_norm_eps=lm_cfg.rms_norm_eps, rope_scaling=lm_cfg.rope_scaling, rope_theta=lm_cfg.rope_theta, sound_dim=sound_dim, sound_gen=sound_gen, sound_latent_fps=sound_latent_fps, - temporal_compression_factor_sound=temporal_compression_factor_sound, timestep_scale=timestep_scale, unified_3d_mrope_reset_spatial_ids=unified_3d_mrope_reset_spatial_ids, unified_3d_mrope_temporal_modality_margin=unified_3d_mrope_temporal_modality_margin, - use_cache=lm_cfg.use_cache, - use_moe=use_moe, - video_temporal_causal=video_temporal_causal, vocab_size=lm_cfg.vocab_size, ) - state_dict = language_model.state_dict() + state_dict = {_remap_language_model_key(k): v for k, v in language_model.state_dict().items()} for k, v in vae2llm.state_dict().items(): - state_dict[f"vae2llm.{k}"] = v + state_dict[f"proj_in.{k}"] = v for k, v in llm2vae.state_dict().items(): - state_dict[f"llm2vae.{k}"] = v - time_embedder_state = time_embedder.state_dict() - if args.remap_time_embedder_keys: - log.info("Remapping transformer time_embedder state dict to diffusers TimestepEmbedding layout …") - time_embedder_state = _remap_time_embedder_state_dict(time_embedder_state) - for k, v in time_embedder_state.items(): - state_dict[f"time_embedder.{k}"] = v + state_dict[f"proj_out.{k}"] = v + for k, v in time_embedder.state_dict().items(): + state_dict[f"time_embedder.{_TIME_EMBEDDER_KEY_REMAP[k]}"] = v if action_gen: for k, v in action2llm.state_dict().items(): - state_dict[f"action2llm.{k}"] = v + state_dict[f"action_proj_in.{k}"] = v for k, v in llm2action.state_dict().items(): - state_dict[f"llm2action.{k}"] = v + state_dict[f"action_proj_out.{k}"] = v state_dict["action_modality_embed"] = action_modality_embed if sound_gen: for k, v in sound2llm.state_dict().items(): - state_dict[f"sound2llm.{k}"] = v + state_dict[f"audio_proj_in.{k}"] = v for k, v in llm2sound.state_dict().items(): - state_dict[f"llm2sound.{k}"] = v - state_dict["sound_modality_embed"] = sound_modality_embed + state_dict[f"audio_proj_out.{k}"] = v + state_dict["audio_modality_embed"] = sound_modality_embed transformer.load_state_dict(state_dict, strict=True, assign=True) del ( language_model, @@ -748,6 +772,10 @@ def convert_model_to_diffusers(args: Args) -> None: "Wan-AI/Wan2.2-TI2V-5B-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16 ) + sound_tokenizer = None + if include_sound_tokenizer: + sound_tokenizer = _build_sound_tokenizer(sound_tokenizer_path, sound_tokenizer_config_path) + # Karras schedule approximating FlowUniPCMultistepScheduler with shift=5, 35 steps. # Measured from that schedule: first flow-sigma=0.9998, last flow-sigma=0.1281. # EDM sigma = flow_sigma / (1 - flow_sigma), so: @@ -762,23 +790,25 @@ def convert_model_to_diffusers(args: Args) -> None: sigma_min=0.147, ) - pipeline = Cosmos3OmniDiffusersPipeline( + # enable_safety_checker=False: constructing the checker imports + # cosmos_guardrail, which the converter environment does not ship. + # Consumers can opt back in at load time via + # `Cosmos3OmniPipeline.from_pretrained(..., enable_safety_checker=True)`. + pipeline = Cosmos3OmniPipeline( transformer=transformer, text_tokenizer=text_tokenizer, vae=diffusers_vae, scheduler=scheduler, - vision_encoder=vision_encoder, + sound_tokenizer=sound_tokenizer, + enable_safety_checker=False, ) log.info(f"Saving full pipeline to {output_dir} …") pipeline.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB") - if include_sound_tokenizer: - _save_sound_tokenizer( - output_dir, - sound_tokenizer_path, - sound_tokenizer_config_path, - remap_keys=args.remap_sound_tokenizer_keys, - ) - _add_sound_tokenizer_to_model_index(output_dir) + if vision_encoder is not None: + # Not a Cosmos3OmniPipeline component — saved as a sidecar folder for + # the transformers/vLLM consumers of the exported repository. + log.info(f"Saving Qwen3-VL vision encoder to {output_dir / 'vision_encoder'} …") + vision_encoder.save_pretrained(str(output_dir / "vision_encoder"), safe_serialization=True) else: log.info(f"Saving transformer to {output_dir} …") transformer.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB") diff --git a/cosmos_framework/scripts/convert_model_to_diffusers.py b/cosmos_framework/scripts/convert_model_to_diffusers.py index 039a2c71..98aa328a 100644 --- a/cosmos_framework/scripts/convert_model_to_diffusers.py +++ b/cosmos_framework/scripts/convert_model_to_diffusers.py @@ -40,18 +40,6 @@ class Args(pydantic.BaseModel): config_only: bool = False """If True, only save config.""" - remap_sound_tokenizer_keys: bool = True - """If True, remap the legacy AVAE sound tokenizer state dict into the - diffusers OobleckDecoder layout (and save the result as - `diffusion_pytorch_model.safetensors`). Off by default — the sound - tokenizer is written verbatim under `model.safetensors`.""" - - remap_time_embedder_keys: bool = True - """If True, remap the transformer's `time_embedder` state dict from the - legacy `nn.Sequential` layout (`mlp.0.*` / `mlp.2.*`) to the diffusers - `TimestepEmbedding` layout (`linear_1.*` / `linear_2.*`). Off by default — - keys are forwarded verbatim.""" - class SafetensorsIndexMetadata(pydantic.BaseModel): total_size: int = 0 @@ -93,9 +81,17 @@ def convert_model_to_diffusers(args: Args): sound_tokenizer_dir, sound_tokenizer_name = model_dict["config"]["sound_tokenizer"]["avae_path"].rsplit("/", 1) sound_tokenizer_checkpoint = CheckpointConfig.maybe_from_uri(f"s3://bucket/{sound_tokenizer_dir}") assert sound_tokenizer_checkpoint is not None - sound_tokenizer_path = Path(sound_tokenizer_checkpoint.hf.download()) / sound_tokenizer_name + sound_tokenizer_local = Path(sound_tokenizer_checkpoint.hf.download()) + # HF-published checkpoints ship sound_tokenizer/ in the diffusers layout + # (config.json + diffusion_pytorch_model.safetensors), which the converter + # consumes directly; fall back to the legacy AVAE file pair named by the + # model config's avae_path. + sound_tokenizer_path = sound_tokenizer_local / "diffusion_pytorch_model.safetensors" + sound_tokenizer_config_path = sound_tokenizer_local / "config.json" + if not sound_tokenizer_path.is_file(): + sound_tokenizer_path = sound_tokenizer_local / sound_tokenizer_name + sound_tokenizer_config_path = sound_tokenizer_path.with_suffix(".json") assert sound_tokenizer_path.is_file(), f"Sound tokenizer checkpoint not found: {sound_tokenizer_path}" - sound_tokenizer_config_path = sound_tokenizer_path.with_suffix(".json") assert sound_tokenizer_config_path.is_file(), f"Sound tokenizer config not found: {sound_tokenizer_config_path}" vision_encoder_model = model_dict["config"]["vlm_config"]["tokenizer"]["pretrained_model_name"] @@ -116,8 +112,6 @@ def convert_model_to_diffusers(args: Args): sound_tokenizer_path=str(sound_tokenizer_path), sound_tokenizer_config_path=str(sound_tokenizer_config_path), include_sound_tokenizer=True, - remap_sound_tokenizer_keys=args.remap_sound_tokenizer_keys, - remap_time_embedder_keys=args.remap_time_embedder_keys, vision_encoder_model=vision_encoder_model, skip_vision_encoder=False, ) diff --git a/docs/code_structure.md b/docs/code_structure.md index 3c7ecd39..25e72f1b 100644 --- a/docs/code_structure.md +++ b/docs/code_structure.md @@ -38,7 +38,7 @@ Cosmos/ │ ├── inference/ # Inference subpackage (model, args, defaults, Ray serving, common helpers, SFT experiment configs) │ └── ... # Training-infra subpackages: data, model, trainer, callbacks, checkpoint, … │ └── scripts/ # CLI entry-point scripts: train.py, _train.py, inference.py, export_model.py, … -├── packages/ # Backend shim packages: diffusers-cosmos3, transformers-cosmos3, vllm-cosmos3 +├── packages/ # Backend shim packages: transformers-cosmos3, vllm-cosmos3 ├── docs/ # User documentation (you are here) ├── docker/ # Dockerfiles for reproducible environments ├── examples/ # Runnable training / fine-tuning / inference examples @@ -49,7 +49,7 @@ Cosmos/ └── .python-version # Python version pin (used by uv) ``` -`cosmos_framework/` is the single Python package. Training infrastructure (data, model, trainer, callbacks, checkpoint, utils, …) lives in top-level subpackages; inference (Diffusers / Transformers / vLLM-friendly inference core, Ray serving, per-modality defaults, training-side experiment YAMLs) lives under `cosmos_framework/inference/`. The library-style backend shims that load Cosmos3 checkpoints into upstream ecosystems live under `packages/{diffusers,transformers,vllm}-cosmos3/`. +`cosmos_framework/` is the single Python package. Training infrastructure (data, model, trainer, callbacks, checkpoint, utils, …) lives in top-level subpackages; inference (Diffusers / Transformers / vLLM-friendly inference core, Ray serving, per-modality defaults, training-side experiment YAMLs) lives under `cosmos_framework/inference/`. The library-style backend shims that load Cosmos3 checkpoints into upstream ecosystems live under `packages/{transformers,vllm}-cosmos3/` (Diffusers needs no shim: diffusers ≥ 0.39 ships `Cosmos3OmniPipeline` natively). ## The `cosmos_framework/` Package @@ -154,7 +154,7 @@ The full inference subpackage: Training-side experiment SKUs live separately at `cosmos_framework/configs/base/experiment/sft/*.py` (see [`cosmos_framework/configs/`](#cosmos_frameworkconfigs)) — not under `inference/`. -Library-style backend shims that adapt Cosmos3 checkpoints to the Diffusers / Transformers / vLLM ecosystems live separately under `packages/{diffusers,transformers,vllm}-cosmos3/`. +Library-style backend shims that adapt Cosmos3 checkpoints to the Transformers / vLLM ecosystems live separately under `packages/{transformers,vllm}-cosmos3/`. Diffusers needs no shim: diffusers ≥ 0.39 ships `Cosmos3OmniPipeline` / `Cosmos3OmniTransformer` natively. ### `cosmos_framework/launcher/` @@ -197,7 +197,7 @@ Add new worker types as sibling subpackages — each owns its own startup, messa - `examples/` — runnable end-to-end examples; see `examples/README.md`. - `docker/` — Dockerfiles and image build helpers; see `docker/README.md`. - `cosmos_framework/scripts/` — CLI entry-point scripts (`train.py`, `inference.py`, `export_model.py`, …); invoke as `python -m cosmos_framework.scripts.`. Primary training entry point: `cosmos_framework.scripts.train` driven by a structured, pydantic-validated TOML interface (`--sft-toml=`); recipe TOMLs live under [`examples/toml/sft_config/`](../examples/toml/sft_config/) and the schema is defined in [`cosmos_framework/configs/toml_config/sft_config.py`](../cosmos_framework/configs/toml_config/sft_config.py) — see [`examples/README.md`](../examples/README.md) and [`docs/training.md`](./training.md). -- `packages/` — library-style backend shims: `packages/{diffusers,transformers,vllm}-cosmos3/`, each installable independently. +- `packages/` — library-style backend shims: `packages/{transformers,vllm}-cosmos3/`, each installable independently. ## Where to Add New Code diff --git a/packages/diffusers-cosmos3/README.md b/packages/diffusers-cosmos3/README.md deleted file mode 100644 index c97fcc6e..00000000 --- a/packages/diffusers-cosmos3/README.md +++ /dev/null @@ -1 +0,0 @@ -# Cosmos3 diffusers Plugin diff --git a/packages/diffusers-cosmos3/diffusers_cosmos3/__init__.py b/packages/diffusers-cosmos3/diffusers_cosmos3/__init__.py deleted file mode 100644 index df9abed7..00000000 --- a/packages/diffusers-cosmos3/diffusers_cosmos3/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -import diffusers - -from .pipeline import Cosmos3OmniDiffusersPipeline -from .transformer import Cosmos3OmniTransformer - -# Spoof the eventual upstream module paths so save_pretrained writes -# "diffusers" as the library in model_index.json. -Cosmos3OmniTransformer.__module__ = "diffusers.models.transformers.transformer_cosmos3" -Cosmos3OmniDiffusersPipeline.__module__ = "diffusers.pipelines.cosmos.pipeline_cosmos3_omni" - -# Loader does `getattr(importlib.import_module(library), class_name)`, -# so the class must be reachable as `diffusers.`. -diffusers.Cosmos3OmniTransformer = Cosmos3OmniTransformer -diffusers.Cosmos3OmniDiffusersPipeline = Cosmos3OmniDiffusersPipeline diff --git a/packages/diffusers-cosmos3/diffusers_cosmos3/pipeline.py b/packages/diffusers-cosmos3/diffusers_cosmos3/pipeline.py deleted file mode 100644 index d514f3e8..00000000 --- a/packages/diffusers-cosmos3/diffusers_cosmos3/pipeline.py +++ /dev/null @@ -1,1124 +0,0 @@ -# Copyright 2025 The NVIDIA Team and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import json -import math -from pathlib import Path -from typing import List, Optional, Tuple, Union - -import numpy as np -import torch -import torchvision.transforms.functional as TF -from diffusers.models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan -from diffusers.pipelines.pipeline_utils import DiffusionPipeline -from diffusers.schedulers import UniPCMultistepScheduler -from diffusers.utils.torch_utils import randn_tensor -from einops import rearrange -from tqdm import tqdm -from transformers import AutoTokenizer - -from diffusers_cosmos3.sequence_packing import ( - GenerationDataClean, - SequencePlan, - build_packed_sequence, - build_sequence_plans_from_data_batch, - get_all_seq, - pack_input_sequence, -) -from diffusers_cosmos3.transformer import ( - Cosmos3OmniTransformer, -) - -_SYSTEM_PROMPT_IMAGE = "You are a helpful assistant who will generate images from a give prompt." -_SYSTEM_PROMPT_VIDEO = "You are a helpful assistant who will generate videos from a give prompt." - - -def save_img_or_video(sample, save_fp_wo_ext, fps=24, quality=10, ffmpeg_params=None, **kwargs): - # TODO: remove this function and use diffusers-style vidoe processor - # However, it may cause numerical differences in the saved video, so we keep it for now for exact reproducibility of saved videos. - import imageio - from PIL import Image as PILImage - - assert sample.ndim == 4, "Only support 4D tensor" - - if torch.is_floating_point(sample): - sample = sample.clamp(0, 1) - else: - assert sample.dtype == torch.uint8, "Only support uint8 tensor" - sample = sample.float().div(255) - - if sample.shape[1] == 1: - save_obj = PILImage.fromarray( - rearrange((sample.cpu().float().numpy() * 255), "c 1 h w -> h w c").astype(np.uint8), - mode="RGB", - ) - save_obj.save(f"{save_fp_wo_ext}.jpg", format="JPEG", quality=85 if quality is None else quality) - else: - frames = rearrange((sample.cpu().float().numpy() * 255), "c t h w -> t h w c").astype(np.uint8) - h, w = frames.shape[1], frames.shape[2] - out_ffmpeg_params = ffmpeg_params if ffmpeg_params is not None else ["-s", f"{w}x{h}"] - imageio.mimsave( - f"{save_fp_wo_ext}.mp4", - frames, - fps=fps, - quality=quality, - macro_block_size=1, - ffmpeg_params=out_ffmpeg_params, - output_params=["-f", "mp4"], - ) - - -class DiffusersWan22VAE: - """ - Drop-in replacement for Wan2pt2VAEInterface, backed by AutoencoderKLWan. - - Bridges the following interface differences: - - 1. encode – AutoencoderKLWan returns AutoencoderKLOutput(latent_dist= - DiagonalGaussianDistribution); we extract .mode() and apply the same - (μ - mean) * inv_std normalization that WanVAE does internally. - - 2. decode – AutoencoderKLWan expects un-normalized z and returns - DecoderOutput(sample=…); we invert the normalization before calling - decode and unwrap the result to a plain tensor. - - 3. spatial/temporal_compression_factor properties – AutoencoderKLWan - stores these as config.scale_factor_spatial / scale_factor_temporal - and exposes spatial_compression_ratio (not *_factor). - - Note: AutoencoderKLWan._decode() clamps the output to [-1, 1]; - Wan2pt2VAEInterface does not. The pipeline applies .clamp(0, 1) after - decode so this difference does not affect saved videos. - - Numerical equivalence requirements (needed for bitwise-identical output - vs Wan2pt2VAEInterface): - - - No torch.amp.autocast: Wan2pt2VAEInterface constructs WanVAE with - is_amp=False, so the encoder/decoder run as pure bfloat16 with no - autocast context. Wrapping calls in autocast changes how ops such as - F.normalize accumulate internally and breaks the match. - - - mean / inv_std must be initialised directly in `dtype` (bfloat16). - WanVAE.__init__ does: - self.std = torch.tensor(std, dtype=bfloat16) - self.scale = [self.mean, 1.0 / self.std] # division in bfloat16 - Computing 1/std in float32 and then casting to bfloat16 can yield - different bit patterns, so we must perform the division in bfloat16 - from the start. - """ - - def __init__(self, vae: AutoencoderKLWan, dtype: torch.dtype = torch.bfloat16): - self.vae = vae - self.dtype = dtype - # Initialise in `dtype` so 1/std is computed in bfloat16, matching WanVAE. - mean = torch.tensor(vae.config.latents_mean, dtype=dtype) - std = torch.tensor(vae.config.latents_std, dtype=dtype) - self._mean = mean # [z_dim] - self._inv_std = 1.0 / std # [z_dim] - - @torch.no_grad() - def encode(self, x: torch.Tensor) -> torch.Tensor: - """[B,3,T,H,W] -> [B,z_dim,T//4,H//16,W//16] (normalized μ, matching Wan2pt2VAEInterface)""" - in_dtype = x.dtype - device = x.device - mean = self._mean.to(device=device, dtype=self.dtype) - inv_std = self._inv_std.to(device=device, dtype=self.dtype) - # No autocast — mirrors WanVAE(is_amp=False), pure bfloat16 forward pass. - raw_mu = self.vae.encode(x.to(self.dtype)).latent_dist.mode() - normalized = (raw_mu - mean.view(1, -1, 1, 1, 1)) * inv_std.view(1, -1, 1, 1, 1) - return normalized.to(in_dtype) - - @torch.no_grad() - def decode(self, z: torch.Tensor) -> torch.Tensor: - """[B,z_dim,T_lat,H_lat,W_lat] -> [B,3,T,H,W]""" - in_dtype = z.dtype - device = z.device - mean = self._mean.to(device=device, dtype=self.dtype) - inv_std = self._inv_std.to(device=device, dtype=self.dtype) - z_raw = z.to(self.dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1) - # No autocast — mirrors WanVAE(is_amp=False), pure bfloat16 forward pass. - out = self.vae.decode(z_raw).sample - return out.to(in_dtype) - - @property - def spatial_compression_factor(self) -> int: - return self.vae.config.scale_factor_spatial - - @property - def temporal_compression_factor(self) -> int: - return self.vae.config.scale_factor_temporal - - -class Cosmos3OmniDiffusersPipeline(DiffusionPipeline): - _optional_components = ["vision_encoder"] - model_cpu_offload_seq = "transformer" - - def __init__( - self, - transformer: Cosmos3OmniTransformer, - text_tokenizer: AutoTokenizer, - vae: AutoencoderKLWan, - scheduler: UniPCMultistepScheduler, - vision_encoder=None, - ): - super().__init__() - self.register_modules( - transformer=transformer, - text_tokenizer=text_tokenizer, - vae=vae, - scheduler=scheduler, - vision_encoder=vision_encoder, - ) - # Plain attribute (not registered): registering the wrapper would cause save_pretrained to call - # wrapper.save_pretrained(), which fails since DiffusersWan22VAE has no such method. - self.vision_tokenizer = DiffusersWan22VAE(vae) - - self.llm_special_tokens = { - "start_of_generation": text_tokenizer.convert_tokens_to_ids("<|vision_start|>"), - "end_of_generation": text_tokenizer.convert_tokens_to_ids("<|vision_end|>"), - "eos_token_id": text_tokenizer.eos_token_id, - } - - def tokenize_caption( - self, - caption: str, - is_video: bool = False, - use_system_prompt: bool = False, - ) -> list[int]: - """Tokenize a text caption into token IDs using the Qwen2 chat template. - Returns: - List of token IDs representing the full chat-formatted caption. - """ - conversations = [] - # Optionally prepend a system prompt that tells the model whether it is generating - # an image or a video. This changes the conditioning context for the LLM. - if use_system_prompt: - _system_prompt = _SYSTEM_PROMPT_VIDEO if is_video else _SYSTEM_PROMPT_IMAGE - conversations.append({"role": "system", "content": _system_prompt}) - conversations.append({"role": "user", "content": caption}) - - tokenizer_output = self.text_tokenizer.apply_chat_template( - conversations, - tokenize=True, - add_generation_prompt=True, - add_vision_id=False, - ) - return tokenizer_output - - def apply_timestep_embeds_to_noisy_tokens( - self, - packed_tokens: torch.Tensor, - packed_timestep_embeds: torch.Tensor, - noisy_frame_indexes: List[torch.Tensor], - token_shapes: list[tuple[int, ...]], - ) -> torch.Tensor: - start_noisy_index = 0 - flattened_noisy_frame_indexes = [] - for noisy_indexes_i, token_shape_i in zip(noisy_frame_indexes, token_shapes): - assert noisy_indexes_i.numel() <= token_shape_i[0] - spatial_numel_i = math.prod(token_shape_i[1:]) - spatial_indexes_i = torch.arange(spatial_numel_i, device=packed_tokens.device) - noisy_indexes_i = (noisy_indexes_i * spatial_numel_i).unsqueeze(-1).expand(-1, spatial_numel_i) - noisy_indexes_i = noisy_indexes_i.clone() + spatial_indexes_i + start_noisy_index - flattened_noisy_frame_indexes.append(noisy_indexes_i.flatten()) - start_noisy_index += math.prod(token_shape_i) - flattened_noisy_frame_indexes = torch.cat(flattened_noisy_frame_indexes, dim=0) - assert packed_tokens.dim() == 2 - assert packed_timestep_embeds.dim() == 2 - assert packed_timestep_embeds.shape[1] == packed_tokens.shape[1] - assert packed_timestep_embeds.shape[0] <= packed_tokens.shape[0] - assert flattened_noisy_frame_indexes.dim() == 1 - assert flattened_noisy_frame_indexes.shape[0] == packed_timestep_embeds.shape[0] - flattened_noisy_frame_indexes = flattened_noisy_frame_indexes.unsqueeze(-1).expand( - -1, - packed_tokens.shape[1], - ) - return packed_tokens.scatter_add( - dim=0, - index=flattened_noisy_frame_indexes, - src=packed_timestep_embeds, - ) - - def patchify_and_pack_latents( - self, - latent_patch_size: int, - latent_channel: int, - tokens_vision: torch.Tensor, - token_shapes_vision: List[Tuple[int, int, int]], - ) -> tuple[torch.Tensor, List[Tuple[int, int, int]]]: - p = latent_patch_size - packed_latent = [] - original_latent_shapes = [] - for latent, (t, h, w) in zip(tokens_vision, token_shapes_vision): - latent = latent.squeeze(0) # [C,T,H,W] - _, t_actual, h_actual, w_actual = latent.shape - original_latent_shapes.append((t_actual, h_actual, w_actual)) - h_padded = ((h_actual + p - 1) // p) * p - w_padded = ((w_actual + p - 1) // p) * p - if h_padded != h_actual or w_padded != w_actual: - padded = torch.zeros( - (latent_channel, t_actual, h_padded, w_padded), - device=latent.device, - dtype=latent.dtype, - ) - padded[:, :, :h_actual, :w_actual] = latent - latent = padded - h_patches = h_padded // p - w_patches = w_padded // p - latent = latent.reshape(latent_channel, t_actual, h_patches, p, w_patches, p) - latent = torch.einsum("cthpwq->thwpqc", latent).reshape(-1, p * p * latent_channel) - packed_latent.append(latent) - return torch.cat(packed_latent, dim=0), original_latent_shapes - - def unpatchify_and_unpack_latents( - self, - latent_patch_size: int, - latent_channel: int, - packed_mse_preds: torch.Tensor, - token_shapes_vision: List[Tuple[int, int, int]], - noisy_frame_indexes_vision: list[torch.Tensor], - original_latent_shapes: List[Tuple[int, int, int]] | None = None, - ) -> list[torch.Tensor]: - p = latent_patch_size - unpatchified_latents = [] - start_idx = 0 - for i, (t_c, h_c, w_c) in enumerate(token_shapes_vision): - if original_latent_shapes is not None: - t_orig, h_orig, w_orig = original_latent_shapes[i] - h_padded = ((h_orig + p - 1) // p) * p - w_padded = ((w_orig + p - 1) // p) * p - h_patches = h_padded // p - w_patches = w_padded // p - else: - t_orig, h_orig, w_orig = t_c, h_c * p, w_c * p - h_patches, w_patches = h_c, w_c - noisy_frame_indexes = noisy_frame_indexes_vision[i] - t_n = len(noisy_frame_indexes) - output_tensor = torch.zeros( - (latent_channel, t_c, h_orig, w_orig), - device=packed_mse_preds.device, - dtype=packed_mse_preds.dtype, - ) - num_patches = t_n * h_patches * w_patches - if num_patches > 0: - end_idx = start_idx + num_patches - latent_patches = packed_mse_preds[start_idx:end_idx] - latent_patches = latent_patches.reshape(t_n, h_patches, w_patches, p, p, latent_channel) - latent = torch.einsum("thwpqc->cthpwq", latent_patches) - latent = latent.reshape(latent_channel, t_n, h_patches * p, w_patches * p) - latent = latent[:, :, :h_orig, :w_orig] - output_tensor[:, noisy_frame_indexes] = latent - start_idx = end_idx - unpatchified_latents.append(output_tensor.unsqueeze(0)) - return unpatchified_latents - - def decode_vision( - self, - patch_latent_dim: int, - latent_patch_size: int, - latent_channel: int, - packed_seq, - last_hidden_state: torch.Tensor, - original_latent_shapes: List[Tuple[int, int, int]] | None = None, - ) -> list[torch.Tensor]: - """Decode vision predictions from last_hidden_state. Returns preds_vision list.""" - vision = packed_seq.vision - has_noisy_vision = ( - vision is not None - and vision.tokens is not None - and isinstance(vision.mse_loss_indexes, torch.Tensor) - and vision.mse_loss_indexes.numel() > 0 - ) - if not has_noisy_vision: - preds_vision = torch.zeros( - [1, patch_latent_dim], device=last_hidden_state.device, dtype=last_hidden_state.dtype - ) - preds_vision = self.transformer.proj_in(preds_vision) - preds_vision = self.transformer.proj_out(preds_vision) - if vision is not None and vision.tokens is not None: - preds_vision_list = [torch.zeros_like(tok) for tok in vision.tokens] - preds_vision_list[0] = preds_vision_list[0] + 0.0 * preds_vision.sum() - else: - preds_vision_list = [preds_vision] - else: - assert vision is not None - assert isinstance(vision.mse_loss_indexes, torch.Tensor) - assert vision.noisy_frame_indexes is not None - preds_vision = self.transformer.proj_out(last_hidden_state[vision.mse_loss_indexes]) - preds_vision_list = self.unpatchify_and_unpack_latents( - latent_patch_size, - latent_channel, - preds_vision, - token_shapes_vision=vision.token_shapes, - noisy_frame_indexes_vision=vision.noisy_frame_indexes, - original_latent_shapes=original_latent_shapes, - ) - return preds_vision_list - - def normalize_video_databatch_inplace( - self, - input_video_key: str, - data_batch: dict, - input_key: str | None = None, - device: str = "cuda", - dtype: torch.dtype = torch.bfloat16, - ) -> None: - input_key = input_video_key if input_key is None else input_key - if input_key in data_batch: - if data_batch.get("is_preprocessed", False) is True: - for i in range(len(data_batch[input_key])): - assert torch.is_floating_point(data_batch[input_key][i]) - assert torch.all((data_batch[input_key][i] >= -1.0001) & (data_batch[input_key][i] <= 1.0001)) - else: - for i in range(len(data_batch[input_key])): - item = data_batch[input_key][i] - if isinstance(item, torch.Tensor): - item = [item] - assert item[0].dtype == torch.uint8 - data_batch[input_key][i] = torch.stack(item).to(device=device, dtype=dtype) / 127.5 - 1.0 - data_batch["is_preprocessed"] = True - - def augment_image_dim_inplace( - self, - input_image_key: str, - data_batch: dict, - input_key: str | None = None, - device: str = "cuda", - dtype: torch.dtype = torch.bfloat16, - ) -> None: - input_key = input_image_key if input_key is None else input_key - if input_key in data_batch: - if data_batch.get("is_preprocessed", False) is True: - for i in range(len(data_batch[input_key])): - assert data_batch[input_key][i].shape[2] == 1 - return - else: - new_image_tensor_list = [] - for i in range(len(data_batch[input_key])): - for img_tensor in data_batch[input_key][i]: - img_tensor = rearrange(img_tensor, "c h w -> 1 c 1 h w").contiguous() - if img_tensor.dtype == torch.uint8: - img_tensor = img_tensor.to(device=device, dtype=dtype) / 127.5 - 1.0 - new_image_tensor_list.append(img_tensor) - data_batch[input_key] = new_image_tensor_list - data_batch["is_preprocessed"] = True - - def remove_padding_from_latent( - self, - spatial_compression_factor: int, - x0_tokens_vision: list[torch.Tensor], - frame_size: list[torch.Tensor], - ) -> list[torch.Tensor]: - cropped_latents = [] - for i in range(len(x0_tokens_vision)): - fs = frame_size[i] - if fs.dim() == 2: - fs = fs[0] - orig_h = int(fs[2].item()) - orig_w = int(fs[3].item()) - orig_h_latent = orig_h // spatial_compression_factor - orig_w_latent = orig_w // spatial_compression_factor - cropped_latents.append(x0_tokens_vision[i][:, :, :, :orig_h_latent, :orig_w_latent].contiguous()) - return cropped_latents - - def get_data_and_condition( - self, - input_image_key: str, - input_video_key: str, - data_batch: dict, - device: str = "cuda", - dtype: torch.dtype = torch.bfloat16, - ) -> GenerationDataClean: - assert (input_image_key in data_batch) != (input_video_key in data_batch) - is_img = input_image_key in data_batch - sample_vision_list = data_batch[input_image_key if is_img else input_video_key] - - if "num_vision_items_per_sample" not in data_batch: - has_multiple_vision_per_sample = any( - isinstance(v, (list, tuple)) and len(v) > 1 for v in sample_vision_list - ) - num_vision_items_per_sample: list[int] | None = ( - [len(v) for v in sample_vision_list] if has_multiple_vision_per_sample else None - ) - data_batch["num_vision_items_per_sample"] = num_vision_items_per_sample - if has_multiple_vision_per_sample: - media_key = input_video_key if not is_img else input_image_key - data_batch[media_key] = [item.unsqueeze(0) for sublist in sample_vision_list for item in sublist] - if data_batch[media_key][0].dtype == torch.float32 and not is_img: - data_batch["is_preprocessed"] = True - else: - num_vision_items_per_sample = data_batch["num_vision_items_per_sample"] - - batch_size = ( - len(sample_vision_list) if num_vision_items_per_sample is None else len(num_vision_items_per_sample) - ) - - self.normalize_video_databatch_inplace(input_video_key, data_batch, device=device, dtype=dtype) - self.augment_image_dim_inplace(input_image_key, data_batch, device=device, dtype=dtype) - raw_state_vision = data_batch[input_image_key if is_img else input_video_key] - x0_tokens_vision = [ - self.vision_tokenizer.encode(raw_state_vision_i).contiguous().float() - for raw_state_vision_i in raw_state_vision - ] - - frame_size = data_batch.get("image_size", None) - if frame_size is not None: - x0_tokens_vision = self.remove_padding_from_latent( - self.vision_tokenizer.spatial_compression_factor, x0_tokens_vision, frame_size - ) - - fps_raw = data_batch.get("conditioning_fps", None) - if isinstance(fps_raw, list): - fps_raw = torch.stack(fps_raw).flatten() - fps_vision = fps_raw.to(device=device, dtype=dtype) if fps_raw is not None else None - - return GenerationDataClean( - batch_size=batch_size, - is_image_batch=is_img, - raw_state_vision=raw_state_vision, - x0_tokens_vision=x0_tokens_vision, - fps_vision=fps_vision, - num_vision_items_per_sample=num_vision_items_per_sample, - ) - - def get_inference_text_tokens( - self, use_system_prompt: bool, input_caption_key: str, data_batch: dict, has_negative_prompt: bool - ) -> tuple[list[list[int]], list[list[int]]]: - cond_tokens = [ - self.tokenize_caption(c, is_video=False, use_system_prompt=use_system_prompt) - for c in data_batch[input_caption_key] - ] - if has_negative_prompt: - neg_key = "neg_" + input_caption_key - assert neg_key in data_batch, f"Negative prompt ({neg_key}) not found" - uncond_captions = data_batch[neg_key] - else: - uncond_captions = [""] * len(cond_tokens) - uncond_tokens = [ - self.tokenize_caption(c, is_video=False, use_system_prompt=use_system_prompt) for c in uncond_captions - ] - return cond_tokens, uncond_tokens - - def derive_include_end_of_generation_token(self, joint_attn_implementation: str) -> bool: - assert joint_attn_implementation in ("flex", "two_way", "three_way") - return joint_attn_implementation == "flex" - - def prepare_inference_data( - self, - use_system_prompt: bool, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - image=None, - num_frames: int = 189, - height: int = 720, - width: int = 1280, - fps: float = 24.0, - condition_frame_indexes: Optional[List[int]] = None, - noises: Optional[List[torch.Tensor]] = None, - generator: Optional[torch.Generator] = None, - input_caption_key: str = "ai_caption", - input_video_key: str = "video", - input_image_key: str = "images", - device: str = "cuda", - dtype: torch.dtype = torch.bfloat16, - ) -> tuple[ - list[SequencePlan], - GenerationDataClean, - list[list[int]], - list[list[int]], - torch.Tensor, - ]: - # Build data_batch - prompts = [prompt] if isinstance(prompt, str) else list(prompt) - batch_size = len(prompts) - is_image = num_frames == 1 - - conditioning_frames = None - if image is not None: - conditioning_frames = self._load_image_as_tensor(image, height, width) - - image_size = [ - torch.tensor([[height, width, height, width]], dtype=torch.float32, device=device) - for _ in range(batch_size) - ] - - if is_image: - img_tensor = ( - conditioning_frames.unsqueeze(0).to(device=device, dtype=dtype) - if conditioning_frames is not None - else torch.zeros(1, 3, 1, height, width, dtype=dtype, device=device) - ) - seq_plans = [ - SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=[]) - for _ in range(batch_size) - ] - data_batch = { - input_image_key: [img_tensor] * batch_size, - "image_size": image_size, - "is_preprocessed": True, - "fps": torch.full((batch_size,), float(fps), device=device), - "conditioning_fps": torch.full((batch_size,), float(fps), device=device), - "num_frames": torch.full((batch_size,), num_frames, device=device), - "sequence_plan": seq_plans, - input_caption_key: prompts, - } - else: - cond_indexes = ( - condition_frame_indexes - if condition_frame_indexes is not None - else ([0] if conditioning_frames is not None else []) - ) - if conditioning_frames is not None: - video_data = torch.zeros(1, 3, num_frames, height, width, dtype=dtype) - t_fill = min(conditioning_frames.shape[1], num_frames) - video_data[0, :, :t_fill] = conditioning_frames[:, :t_fill].to(dtype=dtype) - if t_fill < num_frames: - video_data[0, :, t_fill:] = video_data[0, :, t_fill - 1 : t_fill].expand( - -1, num_frames - t_fill, -1, -1 - ) - video_tensor = video_data.to(device=device) - else: - video_tensor = torch.zeros(1, 3, num_frames, height, width, dtype=dtype, device=device) - seq_plans = [ - SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=list(cond_indexes)) - for _ in range(batch_size) - ] - data_batch = { - input_video_key: [video_tensor] * batch_size, - "image_size": image_size, - "is_preprocessed": True, - "fps": torch.full((batch_size,), float(fps), device=device), - "conditioning_fps": torch.full((batch_size,), float(fps), device=device), - "num_frames": torch.full((batch_size,), num_frames, device=device), - "sequence_plan": seq_plans, - input_caption_key: prompts, - } - - has_negative_prompt = negative_prompt is not None - if has_negative_prompt: - neg_prompts = [negative_prompt] if isinstance(negative_prompt, str) else list(negative_prompt) - data_batch["neg_" + input_caption_key] = neg_prompts - - sequence_plans = build_sequence_plans_from_data_batch( - data_batch=data_batch, - input_video_key=input_video_key, - input_image_key=input_image_key, - ) - gen_data_clean = self.get_data_and_condition( - input_image_key, input_video_key, data_batch, device=device, dtype=dtype - ) - num_items_per_sample = gen_data_clean.num_vision_items_per_sample - cond_text_tokens, uncond_text_tokens = self.get_inference_text_tokens( - use_system_prompt, input_caption_key, data_batch, has_negative_prompt - ) - - mask_timesteps = torch.zeros((gen_data_clean.batch_size,), dtype=torch.float32) - packed_seq = pack_input_sequence( - sequence_plans=sequence_plans, - input_text_indexes=cond_text_tokens, - gen_data_clean=gen_data_clean, - input_timesteps=mask_timesteps, - special_tokens=self.llm_special_tokens, - latent_patch_size=self.transformer.config.latent_patch_size, - include_end_of_generation_token=self.derive_include_end_of_generation_token( - self.transformer.config.joint_attn_implementation - ), - position_embedding_type=self.transformer.config.position_embedding_type, - unified_3d_mrope_reset_spatial_ids=self.transformer.config.unified_3d_mrope_reset_spatial_ids, - unified_3d_mrope_temporal_modality_margin=self.transformer.config.unified_3d_mrope_temporal_modality_margin, - enable_fps_modulation=self.transformer.config.enable_fps_modulation, - base_fps=float(self.transformer.config.base_fps), - temporal_compression_factor=self.vision_tokenizer.temporal_compression_factor, - video_temporal_causal=self.transformer.config.video_temporal_causal, - action_dim=self.transformer.config.max_action_dim, - ) - - assert packed_seq.vision is not None - assert packed_seq.vision.condition_mask is not None - assert isinstance(packed_seq.vision.condition_mask, list) - assert gen_data_clean.x0_tokens_vision is not None - - noise_vision_list: list[torch.Tensor] = [] - for i, (x0_token, cond_mask) in enumerate( - zip(gen_data_clean.x0_tokens_vision, packed_seq.vision.condition_mask, strict=True) - ): - if noises is not None: - pure_noise = noises[i].to(device=device, dtype=dtype) - else: - pure_noise = randn_tensor(tuple(x0_token.shape), generator=generator, device=device, dtype=dtype) - noise_vision_list.append( - cond_mask * x0_token.to(device=device, dtype=dtype) + (1.0 - cond_mask) * pure_noise - ) - - initial_noise = torch.cat([t.reshape(-1) for t in noise_vision_list]) - - return sequence_plans, gen_data_clean, cond_text_tokens, uncond_text_tokens, initial_noise - - def encode_text( - self, - hidden_size: int, - packed_seq, - ) -> tuple[torch.Tensor, torch.dtype]: - """Embed text tokens. Returns (hidden_states [N_total, H], target_dtype).""" - packed_text_embedding = self.transformer.embed_tokens(packed_seq.text_ids) # [N_text,H] - hidden_states = packed_text_embedding.new_zeros(size=(packed_seq.sequence_length, hidden_size)) - hidden_states[packed_seq.text_indexes] = packed_text_embedding - return hidden_states, packed_text_embedding.dtype - - def encode_vision( - self, - timestep_scale: float, - latent_patch_size: int, - latent_channel: int, - packed_seq, - hidden_states: torch.Tensor, - target_dtype: torch.dtype, - fps: Optional[torch.Tensor] = None, - ) -> List[Tuple[int, int, int]] | None: - """Project vision tokens into hidden_states in-place. Returns original_latent_shapes.""" - if packed_seq.vision is None or packed_seq.vision.tokens is None: - return None - vision = packed_seq.vision - assert vision.tokens is not None - assert vision.token_shapes is not None - assert isinstance(vision.sequence_indexes, torch.Tensor) - assert isinstance(vision.timesteps, torch.Tensor) - assert isinstance(vision.mse_loss_indexes, torch.Tensor) - - packed_tokens_vision, original_latent_shapes = self.patchify_and_pack_latents( - latent_patch_size, latent_channel, vision.tokens, vision.token_shapes - ) - packed_tokens_vision = self.transformer.proj_in(packed_tokens_vision) - - if vision.mse_loss_indexes.numel() > 0: - timesteps_vision = vision.timesteps * timestep_scale - with torch.autocast("cuda", enabled=True, dtype=torch.float32): - packed_timestep_embeds_vision = self.transformer.time_embedder(timesteps_vision) - packed_timestep_embeds_vision = packed_timestep_embeds_vision.to(target_dtype) - packed_tokens_vision = self.apply_timestep_embeds_to_noisy_tokens( - packed_tokens=packed_tokens_vision, - packed_timestep_embeds=packed_timestep_embeds_vision, - noisy_frame_indexes=vision.noisy_frame_indexes, - token_shapes=vision.token_shapes, - ) - - hidden_states[vision.sequence_indexes] = packed_tokens_vision - return original_latent_shapes - - @torch.no_grad() - def run_single( - self, - packed_seq, - noise_x_vision: list[torch.Tensor], - hidden_size: int, - latent_patch_size: int, - latent_channel: int, - patch_latent_dim: int, - timestep_scale: float, - num_heads: int, - head_dim: int, - num_hidden_layers: int, - use_moe: bool, - joint_attn_implementation: str, - fps_vision: Optional[torch.Tensor], - device: str = "cuda", - dtype: torch.dtype = torch.bfloat16, - ) -> list[torch.Tensor]: - """Inlined forward pass from Cosmos3VFMNetworkSimple.forward().""" - if packed_seq.vision is not None: - packed_seq.vision.tokens = [x.to(device=device, dtype=dtype) for x in noise_x_vision] - packed_seq.to_cuda() - - # 1. Encode text - hidden_states, target_dtype = self.encode_text(hidden_size, packed_seq) - - # 2. Encode vision - original_latent_shapes = self.encode_vision( - timestep_scale, - latent_patch_size, - latent_channel, - packed_seq, - hidden_states, - target_dtype, - fps=fps_vision, - ) - - # 3. Build attention metadata - assert use_moe - assert packed_seq.attn_modes is not None - assert packed_seq.split_lens is not None - all_gen_indexes = [] - if packed_seq.vision is not None: - assert packed_seq.vision.token_shapes is not None - assert isinstance(packed_seq.vision.sequence_indexes, torch.Tensor) - all_gen_indexes.append(packed_seq.vision.sequence_indexes) - vision_sequence_indexes = torch.cat(all_gen_indexes, dim=0) if all_gen_indexes else None - - input_pack, attention_meta, _ = build_packed_sequence( - joint_attn_implementation, - packed_sequence=hidden_states, - attn_modes=packed_seq.attn_modes, - split_lens=packed_seq.split_lens, - sample_lens=packed_seq.sample_lens, - packed_und_token_indexes=packed_seq.text_indexes, - packed_gen_token_indexes=vision_sequence_indexes, - num_heads=num_heads, - is_image_batch=packed_seq.is_image_batch, - head_dim=head_dim, - num_layers=num_hidden_layers, - token_shapes=packed_seq.vision.token_shapes, - natten_parameter_list=None, - cp_world_size=1, - video_temporal_causal=False, - vision_token_shapes=packed_seq.vision.token_shapes if packed_seq.vision else None, - action_token_shapes=None, - temporal_compression_factor_vision=self.vision_tokenizer.temporal_compression_factor, - null_action_supertokens=packed_seq.null_action_supertokens, - pad_for_cuda_graphs=False, - ) - - # 4. Run transformer - packed_outputs, _ = self.transformer( - input_pack, - attention_mask=attention_meta, - position_ids=packed_seq.position_ids, - dual_kv_cache=None, - frame_idx=None, - natten_metadata_list=None, - ) - last_hidden_state = get_all_seq(packed_outputs) - - # 5. Decode vision - return self.decode_vision( - patch_latent_dim, - latent_patch_size, - latent_channel, - packed_seq, - last_hidden_state, - original_latent_shapes, - ) - - @torch.no_grad() - def get_cfg_velocity( - self, - noise_x: torch.Tensor, - timestep: torch.Tensor, - guidance: float, - gen_data_clean: GenerationDataClean, - sequence_plans: list[SequencePlan], - cond_tokens: list[list[int]], - uncond_tokens: list[list[int]], - include_eog: bool, - hidden_size: int, - latent_patch_size: int, - latent_channel: int, - patch_latent_dim: int, - timestep_scale: float, - num_heads: int, - head_dim: int, - num_hidden_layers: int, - use_moe: bool, - joint_attn_implementation: str, - skip_text_tokens_for_cfg: bool = False, - normalize_cfg: bool = False, - device: str = "cuda", - dtype: torch.dtype = torch.bfloat16, - ) -> torch.Tensor: - torch.compiler.cudagraph_mark_step_begin() - assert timestep.ndim == 2 and timestep.shape == (1, 1) - - num_items = gen_data_clean.num_vision_items_per_sample - num_vis = num_items[0] if num_items is not None else 1 - - noise_x_vision: list[torch.Tensor] = [] - offset = 0 - for j in range(num_vis): - vision_shape = gen_data_clean.x0_tokens_vision[j].shape - vision_dim = int(torch.prod(torch.tensor(vision_shape))) - noise_x_vision.append(noise_x[offset : offset + vision_dim].reshape(vision_shape)) - offset += vision_dim - - gen_data_for_packing = GenerationDataClean( - batch_size=1, - is_image_batch=gen_data_clean.is_image_batch, - raw_state_vision=gen_data_clean.raw_state_vision, - x0_tokens_vision=noise_x_vision, - fps_vision=gen_data_clean.fps_vision, - num_vision_items_per_sample=num_items, - ) - - def _run_cond(text_tokens: list[list[int]], skip_text: bool) -> torch.Tensor: - packed_seq = pack_input_sequence( - sequence_plans=sequence_plans, - input_text_indexes=text_tokens, - gen_data_clean=gen_data_for_packing, - input_timesteps=timestep.cpu(), - special_tokens=self.llm_special_tokens, - latent_patch_size=self.transformer.config.latent_patch_size, - include_end_of_generation_token=include_eog, - skip_text_tokens=skip_text, - position_embedding_type=self.transformer.config.position_embedding_type, - unified_3d_mrope_reset_spatial_ids=self.transformer.config.unified_3d_mrope_reset_spatial_ids, - unified_3d_mrope_temporal_modality_margin=self.transformer.config.unified_3d_mrope_temporal_modality_margin, - enable_fps_modulation=self.transformer.config.enable_fps_modulation, - base_fps=float(self.transformer.config.base_fps), - temporal_compression_factor=self.vision_tokenizer.temporal_compression_factor, - video_temporal_causal=self.transformer.config.video_temporal_causal, - action_dim=self.transformer.config.max_action_dim, - ) - preds = self.run_single( - packed_seq, - noise_x_vision, - hidden_size, - latent_patch_size, - latent_channel, - patch_latent_dim, - timestep_scale, - num_heads, - head_dim, - num_hidden_layers, - use_moe, - joint_attn_implementation, - fps_vision=gen_data_clean.fps_vision, - device=device, - dtype=dtype, - ) - - assert packed_seq.vision is not None - assert packed_seq.vision.condition_mask is not None - assert isinstance(packed_seq.vision.condition_mask, list) - velocity_vision = [ - pred * (1.0 - m).to(dtype=pred.dtype, device=pred.device) - if (1.0 - m).sum() > 0 - else torch.zeros_like(pred) - for pred, m in zip(preds, packed_seq.vision.condition_mask) - ] - return torch.cat([v.reshape(-1) for v in velocity_vision]) - - cond_v = _run_cond(cond_tokens, False) - uncond_v = _run_cond(uncond_tokens, skip_text_tokens_for_cfg) - - v_pred = uncond_v + guidance * (cond_v - uncond_v) - - if normalize_cfg: - v_pred = v_pred * (torch.norm(cond_v) / (torch.norm(v_pred) + 1e-8)).clamp(min=0.0, max=1.0) - return v_pred - - def _load_image_as_tensor(self, image, target_h: int, target_w: int) -> torch.Tensor: - """Load image from PIL, path, URL, or tensor; returns [3, 1, H, W] in [-1, 1].""" - from PIL import Image as PILImage - - if isinstance(image, (str, Path)): - image_str = str(image) - if image_str.startswith("http://") or image_str.startswith("https://"): - import io - import urllib.request - - with urllib.request.urlopen(image_str) as resp: - img_bytes = resp.read() - pil_img = PILImage.open(io.BytesIO(img_bytes)).convert("RGB") - else: - with open(image_str, "rb") as f: - pil_img = PILImage.open(f).convert("RGB") - img_t = torch.from_numpy(np.array(pil_img)).permute(2, 0, 1).float() - elif hasattr(image, "convert"): # PIL.Image - img_t = torch.from_numpy(np.array(image.convert("RGB"))).permute(2, 0, 1).float() - elif isinstance(image, torch.Tensor): - img_t = image.float() - if img_t.dim() == 4: - img_t = img_t.squeeze(0) - # if already normalized to [-1, 1], skip the /127.5-1 step below - if img_t.max() <= 1.1: - img_4d = img_t.unsqueeze(0) - orig_h, orig_w = img_4d.shape[2], img_4d.shape[3] - scale = max(target_w / orig_w, target_h / orig_h) - resize_h = int(math.ceil(scale * orig_h)) - resize_w = int(math.ceil(scale * orig_w)) - img_4d = TF.resize(img_4d, [resize_h, resize_w]) - img_4d = TF.center_crop(img_4d, [target_h, target_w]) - return img_4d.squeeze(0).unsqueeze(1) - else: - raise TypeError(f"Unsupported image type: {type(image)}") - - img_4d = img_t.unsqueeze(0) # [1, 3, H, W] (uint8-range [0, 255]) - orig_h, orig_w = img_4d.shape[2], img_4d.shape[3] - scale = max(target_w / orig_w, target_h / orig_h) - resize_h = int(math.ceil(scale * orig_h)) - resize_w = int(math.ceil(scale * orig_w)) - img_4d = TF.resize(img_4d, [resize_h, resize_w]) - img_4d = TF.center_crop(img_4d, [target_h, target_w]) - img_4d = img_4d / 127.5 - 1.0 # normalize after resize, matching load_conditioning_image - return img_4d.squeeze(0).unsqueeze(1) # [3, 1, H, W] - - def _resolve_defaults_and_prompts( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]], - image, - num_frames: int, - fps: float, - height: int, - width: int, - ) -> tuple[float, int, float, Union[str, List[str]], Union[str, List[str]]]: - """Load modality defaults and apply duration/resolution templates to prompts. - - Returns (guidance, num_steps, shift, formatted_prompt, formatted_negative_prompt). - """ - model_mode = "image2video" if image is not None else "text2video" - defaults = json.loads((Path(__file__).parent / f"sample_args/{model_mode}.json").read_text()) - - guidance = float(defaults["guidance"]) - num_steps = int(defaults["num_steps"]) - shift = float(defaults["shift"]) - print(f"model_mode={model_mode!r}: guidance={guidance}, num_steps={num_steps}, shift={shift}") - - duration_template = defaults.get("duration_template") - resolution_template = defaults.get("resolution_template") - negative_prompt_base = defaults.get("negative_prompt", "") - keep_metadata = defaults.get("negative_prompt_keep_metadata", False) - - def _apply_templates(text: str) -> str: - if duration_template and num_frames > 1: - text = text.rstrip(".") + ". " + duration_template.format(duration=num_frames / fps, fps=fps) - if resolution_template: - text = text.rstrip(".") + ". " + resolution_template.format(height=height, width=width) - return text - - if isinstance(prompt, str): - prompt = _apply_templates(prompt) - else: - prompt = [_apply_templates(p) for p in prompt] - - if negative_prompt is None: - negative_prompt = _apply_templates(negative_prompt_base) if keep_metadata else negative_prompt_base - - return guidance, num_steps, shift, prompt, negative_prompt - - @torch.no_grad() - def decode_latents(self, vision_list: list[torch.Tensor]) -> list[torch.Tensor]: - """Decode latents to pixel tensors of shape [C, T, H, W] in [0, 1].""" - frames = [] - for vision_latent in vision_list: - vision = self.vision_tokenizer.decode(vision_latent.cuda()) # [1, C, T, H, W] - frames.append(((1.0 + vision) / 2).clamp(0, 1).squeeze(0)) - return frames - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - image=None, - num_frames: int = 189, - height: int = 720, - width: int = 1280, - fps: float = 24.0, - condition_frame_indexes: Optional[List[int]] = None, - noises: Optional[List[torch.Tensor]] = None, - generator: Optional[torch.Generator] = None, - use_system_prompt: bool = False, - device: str = "cuda", - dtype: torch.dtype = torch.bfloat16, - output_type: str = "video", - ): - latent_patch_size = self.transformer.config.latent_patch_size - latent_channel = self.transformer.config.latent_channel - patch_latent_dim = self.transformer.config.patch_latent_dim - timestep_scale = self.transformer.config.timestep_scale - hidden_size = self.transformer.config.hidden_size - num_heads = self.transformer.config.num_attention_heads - head_dim = self.transformer.config.head_dim - num_hidden_layers = self.transformer.config.num_hidden_layers - use_moe = self.transformer.config.use_moe - joint_attn_implementation = self.transformer.config.joint_attn_implementation - - guidance, num_steps, shift, prompt, negative_prompt = self._resolve_defaults_and_prompts( - prompt, negative_prompt, image, num_frames, fps, height, width - ) - - sequence_plans, gen_data_clean, cond_tokens, uncond_tokens, initial_noise = self.prepare_inference_data( - use_system_prompt, - prompt=prompt, - negative_prompt=negative_prompt, - image=image, - num_frames=num_frames, - height=height, - width=width, - fps=fps, - condition_frame_indexes=condition_frame_indexes, - noises=noises, - generator=generator, - device=device, - dtype=dtype, - ) - - assert guidance != 1.0, "Guidance weight must be != 1.0" - - device = initial_noise.device - self.scheduler.set_timesteps(num_steps, device=device) - timesteps = self.scheduler.timesteps - # print(f"sigmas: first={self.scheduler.sigmas[0].item():.4f} last={self.scheduler.sigmas[-2].item():.4f}") - # print(f"timesteps: first={timesteps[0].item():.2f} last={timesteps[-1].item():.2f}") - # print(f"timestep_scale: {timestep_scale}") - # breakpoint() - latent = initial_noise - include_eog = self.derive_include_end_of_generation_token(joint_attn_implementation) - - # --- Denoising loop --- - print("Running generate_samples_from_batch …") - for timestep in tqdm(timesteps, desc="Denoising"): - velocity_pred = self.get_cfg_velocity( - latent, - timestep.reshape(1, 1), - guidance, - gen_data_clean, - sequence_plans, - cond_tokens, - uncond_tokens, - include_eog, - hidden_size, - latent_patch_size, - latent_channel, - patch_latent_dim, - timestep_scale, - num_heads, - head_dim, - num_hidden_layers, - use_moe, - joint_attn_implementation, - device=device, - dtype=dtype, - ) - latent = self.scheduler.step( - model_output=velocity_pred, - timestep=timestep, - sample=latent.unsqueeze(0), - return_dict=False, - )[0].squeeze(0) - - # --- Extract vision results --- - num_vision_items = gen_data_clean.num_vision_items_per_sample - n_vis = num_vision_items[0] if num_vision_items is not None else 1 - result_vision: list[torch.Tensor] = [] - offset = 0 - for j in range(n_vis): - vision_shape = gen_data_clean.x0_tokens_vision[j].shape - vision_dim = int(torch.prod(torch.tensor(vision_shape))) - if j == n_vis - 1: - result_vision.append(latent[offset : offset + vision_dim].reshape(vision_shape)) - offset += vision_dim - - if output_type == "latent": - return result_vision - return self.decode_latents(result_vision) diff --git a/packages/diffusers-cosmos3/diffusers_cosmos3/sequence_packing.py b/packages/diffusers-cosmos3/diffusers_cosmos3/sequence_packing.py deleted file mode 100644 index 658eb25f..00000000 --- a/packages/diffusers-cosmos3/diffusers_cosmos3/sequence_packing.py +++ /dev/null @@ -1,1626 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -# Copied from cosmos3._src.vfm.datasets.sequence_packing - -import math -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple - -import torch - - -@dataclass -class GenerationDataClean: - batch_size: int - is_image_batch: bool - raw_state_vision: Optional[List[torch.Tensor]] = None - x0_tokens_vision: Optional[List[torch.Tensor]] = None - fps_vision: Optional[torch.Tensor] = None - num_vision_items_per_sample: Optional[List[int]] = None - x0_tokens_action: Optional[List[torch.Tensor]] = None - fps_action: Optional[torch.Tensor] = None - action_domain_id: Optional[List[torch.Tensor]] = None - raw_action_dim: Optional[List[Optional[torch.Tensor]]] = None - x0_tokens_sound: Optional[List[torch.Tensor]] = None - fps_sound: Optional[torch.Tensor] = None - - -# "Fake" types for readability; everything is plain dict at runtime. -FactoredSequencePack = dict[str, Any] -JointSequencePack = dict[str, Any] -SequencePack = FactoredSequencePack | JointSequencePack - - -# ------------------------------------ -# Internal helpers -# ------------------------------------ - - -def _pad_to_N(N, x: torch.Tensor) -> torch.Tensor: - assert x.shape[0] <= N - padded = x.new_zeros((N, *x.shape[1:])) - padded[: x.shape[0]] = x - return padded - - -def _pad( - causal_seq: torch.Tensor, full_only_seq: torch.Tensor, max_causal_len: int, max_full_len: int -) -> tuple[torch.Tensor, torch.Tensor]: - causal_seq = _pad_to_N(max_causal_len, causal_seq) - full_only_seq = _pad_to_N(max_full_len, full_only_seq) - return causal_seq, full_only_seq - - -def _compute_mode_indices_and_offsets( - split_lens: list[int], - attn_modes: list[str], - mode: str, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor]: - indices = [] - offsets = [0] - next_offset = 0 - start = 0 - for split_len, attn_mode in zip(split_lens, attn_modes): - if attn_mode == mode: - indices.extend(range(start, start + split_len)) - next_offset += split_len - offsets.append(next_offset) - start += split_len - return ( - torch.tensor(indices, dtype=torch.int32, device=device), - torch.tensor(offsets, dtype=torch.int32, device=device), - ) - - -def _init_sequence_pack( - sample_lens: list[int], - split_lens: list[int], - attn_modes: list[str], - device: torch.device, -) -> dict: - _max_sample_len = max(sample_lens) - _max_causal_len = max((split_lens[i] for i in range(len(split_lens)) if attn_modes[i] == "causal"), default=0) - _max_full_len = max((split_lens[i] for i in range(len(split_lens)) if attn_modes[i] == "full"), default=0) - sample_lens_cu = torch.tensor([0] + sample_lens, device=device, dtype=torch.int32) - _sample_offsets = torch.cumsum(sample_lens_cu, dim=0, dtype=torch.int32) - _causal_indices, _causal_seq_offsets = _compute_mode_indices_and_offsets(split_lens, attn_modes, "causal", device) - _full_indices, _full_only_seq_offsets = _compute_mode_indices_and_offsets(split_lens, attn_modes, "full", device) - return dict( - sample_offsets=_sample_offsets, - max_sample_len=_max_sample_len, - max_causal_len=_max_causal_len, - max_full_len=_max_full_len, - _causal_indices=_causal_indices, - _full_indices=_full_indices, - _causal_seq_offsets=_causal_seq_offsets, - _full_only_seq_offsets=_full_only_seq_offsets, - _num_causal_tokens=len(_causal_indices), - _num_full_tokens=len(_full_indices), - split_lens=split_lens, - attn_modes=attn_modes, - ) - - -def _find_non_causal_text_token_idx( - attn_modes: list[str], - split_lens: list[int], - und_token_indexes: list[int], -) -> list[int]: - out = [] - full_offset = 0 - packed_idx = 0 - und_token_set = set(und_token_indexes) - for attn_mode, split_len in zip(attn_modes, split_lens): - if attn_mode == "full": - for local_idx, split_idx in enumerate(range(packed_idx, packed_idx + split_len)): - if split_idx in und_token_set: - out.append(full_offset + local_idx) - full_offset += split_len - packed_idx += split_len - return out - - -def factored_from_joint_sequence( - packed_sequence: torch.Tensor, - attn_modes: list[str], - split_lens: list[int], - sample_lens: list[int], - packed_und_token_indexes: torch.Tensor, - packed_gen_token_indexes: torch.Tensor, - is_image_batch: bool = False, - cp_world_size: int = 1, - pad_for_cuda_graphs: bool = False, -) -> FactoredSequencePack: - non_causal_text_idxs = _find_non_causal_text_token_idx(attn_modes, split_lens, packed_und_token_indexes.tolist()) - assert len(non_causal_text_idxs) == 0, "non_causal_text_idxs should be empty" - assert sum(sample_lens) == packed_sequence.shape[0] - meta = _init_sequence_pack(sample_lens, split_lens, attn_modes, packed_sequence.device) - causal_seq = packed_sequence[meta["_causal_indices"]] - full_only_seq = packed_sequence[meta["_full_indices"]] - return { - **meta, - "max_num_tokens": sum(sample_lens), - "causal_seq": causal_seq, - "full_only_seq": full_only_seq, - "is_sharded": False, - } - - -class SplitInfo: - def __init__( - self, - split_lens: list[int], - attn_modes: list[str], - sample_lens: list[int], - actual_len: int, - ): - assert sum(sample_lens) == sum(split_lens) - max_causal_len = 0 - max_full_len = 0 - for split_len, attn_mode in zip(split_lens, attn_modes): - if attn_mode == "causal": - max_causal_len = max(max_causal_len, split_len) - elif attn_mode == "full": - max_full_len = max(max_full_len, split_len) - self.max_causal_len = max_causal_len - self.max_full_len = max_full_len - self.max_sample_len = max(sample_lens) - self.split_lens = split_lens - self.attn_modes = attn_modes - self.sample_lens = sample_lens - - -def build_packed_sequence( - joint_attn_implementation: str, - *, - packed_sequence: torch.Tensor, - attn_modes: list[str], - split_lens: list[int], - sample_lens: list[int], - packed_und_token_indexes: torch.LongTensor, - packed_gen_token_indexes: torch.LongTensor, - num_heads: int, - head_dim: int, - num_layers: int, - token_shapes=None, - natten_parameter_list=None, - block_size: int = 128, - is_image_batch: bool = False, - cp_world_size: int = 1, - video_temporal_causal: bool = False, - vision_token_shapes=None, - action_token_shapes=None, - temporal_compression_factor_vision: int = 4, - null_action_supertokens: bool = False, - pad_for_cuda_graphs: bool = False, -) -> tuple[FactoredSequencePack, SplitInfo, None]: - assert joint_attn_implementation == "two_way", ( - f"Only two_way attention is supported, got {joint_attn_implementation!r}" - ) - device = packed_sequence.device - attention_meta = SplitInfo( - split_lens=split_lens, - attn_modes=attn_modes, - sample_lens=sample_lens, - actual_len=int(packed_sequence.shape[0]), - ) - input_pack = factored_from_joint_sequence( - packed_sequence=packed_sequence, - attn_modes=attn_modes, - split_lens=split_lens, - sample_lens=sample_lens, - packed_und_token_indexes=packed_und_token_indexes.to(device), - packed_gen_token_indexes=packed_gen_token_indexes.to(device), - is_image_batch=is_image_batch, - cp_world_size=cp_world_size, - pad_for_cuda_graphs=pad_for_cuda_graphs, - ) - input_pack.pop("split_lens", None) - input_pack.pop("attn_modes", None) - return input_pack, attention_meta, None - - -def _ensure_core_metadata(pack: SequencePack) -> None: - required = [ - "sample_offsets", - "max_sample_len", - "max_causal_len", - "max_full_len", - "_causal_indices", - "_full_indices", - "_causal_seq_offsets", - "_full_only_seq_offsets", - "is_sharded", - ] - for key in required: - if key not in pack: - raise KeyError(f"Missing required pack field: {key}") - - -def from_mode_splits( - causal_seq: torch.Tensor, - full_only_seq: torch.Tensor, - orig: FactoredSequencePack | JointSequencePack, - is_sharded: bool | None = None, -): - """ - Create a new sequence pack from two mode splits. - Args: - causal_seq (torch.Tensor): The causal sequence. - full_only_seq (torch.Tensor): The full-only sequence. - orig (FactoredSequencePack | JointSequencePack): The metadata source to copy from. - is_sharded (bool | None): If True, create a local pack for context parallel. - If None, inherits from orig. - """ - _ensure_core_metadata(orig) - if is_sharded is None: - is_sharded = orig.get("is_sharded", False) - - if "packed_sequence" in orig: - all_len = int(orig["_causal_indices"].shape[0] + orig["_full_indices"].shape[0]) - packed_sequence = causal_seq.new_zeros((all_len, *causal_seq.shape[1:])) # [seq_len,D] - packed_sequence[orig["_causal_indices"]] = causal_seq - packed_sequence[orig["_full_indices"]] = full_only_seq - return from_joint(packed_sequence, orig) - else: - out = dict(orig) - out["causal_seq"] = causal_seq - out["full_only_seq"] = full_only_seq - out["is_sharded"] = is_sharded - return out - - -# ------------------------------------ -# Public API -# ------------------------------------ - - -def zeros_like(orig: FactoredSequencePack | JointSequencePack, shape: Tuple[int, ...] | torch.Size | None = None): - """ - Create a new sequence pack with the same metadata as the original, but with all tokens set to zero. - Args: - orig (FactoredSequencePack | JointSequencePack): The original sequence pack to copy metadata from. - shape (Tuple[int, ...] | torch.Size | None): The shape of the new sequence pack. If None, the shape will be the same as the original. - """ - _ensure_core_metadata(orig) - if "packed_sequence" in orig: - if shape is None: - shape_ = orig["packed_sequence"].shape - else: - assert len(shape) >= 1 and shape[0] == -1 - shape_ = (orig["packed_sequence"].shape[0],) + tuple(shape)[1:] - packed_sequence = torch.zeros( - shape_, device=orig["packed_sequence"].device, dtype=orig["packed_sequence"].dtype - ) # [seq_len,D] - return from_joint(packed_sequence, orig) - else: - if shape is None: - shape_causal = orig["causal_seq"].shape - shape_full = orig["full_only_seq"].shape - else: - assert len(shape) >= 1 and shape[0] == -1 - shape_causal = (orig["causal_seq"].shape[0],) + tuple(shape)[1:] - shape_full = (orig["full_only_seq"].shape[0],) + tuple(shape)[1:] - causal_seq = torch.zeros( - shape_causal, device=orig["causal_seq"].device, dtype=orig["causal_seq"].dtype - ) # [N_causal_tokens,D] - full_only_seq = torch.zeros( - shape_full, device=orig["full_only_seq"].device, dtype=orig["full_only_seq"].dtype - ) # [N_full_tokens,D] - return from_mode_splits(causal_seq, full_only_seq, orig) - - -def from_joint(packed_sequence: torch.Tensor, metadata_source: FactoredSequencePack | JointSequencePack): - """ - Create a new sequence pack from a packed sequence and another sequence pack with the same metadata. - Args: - packed_sequence (torch.Tensor): Tensor containing all tokens in the batch of sequences. - metadata_source (FactoredSequencePack | JointSequencePack): The metadata source to copy from. - """ - _ensure_core_metadata(metadata_source) - if "packed_sequence" in metadata_source: - out = dict(metadata_source) - out["packed_sequence"] = packed_sequence - return out - else: - if metadata_source["is_sharded"]: - # Use sharded sequences as is when is_sharded is True (used in Context Parallel) - causal_seq = packed_sequence[: len(metadata_source["causal_seq"])] # [N_causal_tokens,D] - full_only_seq = packed_sequence[len(metadata_source["causal_seq"]) :] # [N_full_tokens,D] - else: - causal_seq = packed_sequence[metadata_source["_causal_indices"]] # [N_causal_tokens,D] - full_only_seq = packed_sequence[metadata_source["_full_indices"]] # [N_full_tokens,D] - causal_seq, full_only_seq = _pad( - causal_seq, - full_only_seq, - max_causal_len=metadata_source["causal_seq"].shape[0], - max_full_len=metadata_source["full_only_seq"].shape[0], - ) - - return from_mode_splits(causal_seq, full_only_seq, metadata_source) - - -def from_und_gen_splits(und_seq: torch.Tensor, gen_seq: torch.Tensor, orig: FactoredSequencePack | JointSequencePack): - """ - Create a new sequence pack from two und/gen splits. - Args: - und_seq (torch.Tensor): The understanding sequence. - gen_seq (torch.Tensor): The generating sequence. - orig (FactoredSequencePack | JointSequencePack): The metadata source to copy from. - """ - # If we have a joint pack (single packed_sequence), place by und/gen indexes. - if "packed_sequence" in orig and "packed_und_token_indexes" in orig and "packed_gen_token_indexes" in orig: - all_len = int(und_seq.shape[0] + gen_seq.shape[0]) - packed_sequence = und_seq.new_zeros((all_len, *und_seq.shape[1:])) # [seq_len,D] - packed_sequence[orig["packed_und_token_indexes"]] = und_seq - packed_sequence[orig["packed_gen_token_indexes"]] = gen_seq - return from_joint(packed_sequence, orig) - # Otherwise, treat und/gen as mode splits (und == causal; gen == full). - return from_mode_splits(und_seq, gen_seq, orig) - - -def get_und_seq(pack: SequencePack) -> torch.Tensor: - """ - Get all understanding tokens in a sequence pack in a single tensor. - - Args: - pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the understanding sequence from. - Returns: - torch.Tensor: All understanding tokens concatenated over all sequences in the batch. - """ - if "causal_seq" in pack: - return pack["causal_seq"] - if "packed_sequence" in pack and "packed_und_token_indexes" in pack: - return pack["packed_sequence"][pack["packed_und_token_indexes"]] - raise KeyError("Cannot derive und_seq from provided pack") - - -def set_und_seq(pack: SequencePack, value: torch.Tensor) -> None: - """ - Override the understanding tokens in a sequence pack. - The order of tokens passed in must correspond to the order of tokens returned by get_und_seq. - - Args: - pack (FactoredSequencePack | JointSequencePack): The sequence pack to set the understanding sequence in. - value (torch.Tensor): The understanding sequence to set. - """ - if "packed_sequence" in pack and "packed_und_token_indexes" in pack: - pack["packed_sequence"][pack["packed_und_token_indexes"]] = value - elif "causal_seq" in pack: - pack["causal_seq"] = value - else: - raise KeyError("Cannot set und_seq from provided pack") - - -def get_gen_seq(pack: SequencePack) -> torch.Tensor: - """ - Get all generating tokens in a sequence pack in a single tensor. - Args: - pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the generating sequence from. - Returns: - torch.Tensor: All generating tokens concatenated over all sequences in the batch. - """ - if "full_only_seq" in pack: - return pack["full_only_seq"] - if "packed_sequence" in pack and "packed_gen_token_indexes" in pack: - return pack["packed_sequence"][pack["packed_gen_token_indexes"]] - raise KeyError("Cannot derive gen_seq from provided pack") - - -def set_gen_seq(pack: SequencePack, value: torch.Tensor) -> None: - """ - Override the generating tokens in a sequence pack. - The order of tokens passed in must correspond to the order of tokens returned by get_gen_seq. - Args: - pack (FactoredSequencePack | JointSequencePack): The sequence pack to set the generating sequence in. - value (torch.Tensor): The generating sequence to set. - """ - if "packed_sequence" in pack and "packed_gen_token_indexes" in pack: - pack["packed_sequence"][pack["packed_gen_token_indexes"]] = value - elif "full_only_seq" in pack: - pack["full_only_seq"] = value - else: - raise KeyError("Cannot set gen_seq from provided pack") - - -def get_device_and_dtype(pack: SequencePack) -> Tuple[torch.device, torch.dtype]: - """ - Get the device and dtype of a sequence pack. - Args: - pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the device and dtype from. - Returns: - Tuple[torch.device, torch.dtype]: The device and dtype of the sequence pack. - """ - if "packed_sequence" in pack: - return pack["packed_sequence"].device, pack["packed_sequence"].dtype - if "causal_seq" in pack and "full_only_seq" in pack: - return pack["causal_seq"].device, pack["causal_seq"].dtype - raise KeyError("Cannot derive device and dtype from provided pack") - - -def get_all_seq(pack: SequencePack) -> torch.Tensor: - """ - Get all tokens in a sequence pack in a single tensor. - Args: - pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the all sequence from. - Returns: - torch.Tensor: All tokens concatenated over all sequences in the batch. - """ - if "all_seq" in pack: - return pack["all_seq"] - if "packed_sequence" in pack: - return pack["packed_sequence"] - if "causal_seq" in pack and "full_only_seq" in pack: - _ensure_core_metadata(pack) - if pack["is_sharded"]: - assert False, "get_all_seq is not supported in context parallel sharded mode" - else: - out = pack["causal_seq"].new_zeros( - int(pack["_causal_indices"].shape[0] + pack["_full_indices"].shape[0]), *pack["causal_seq"].shape[1:] - ) # [seq_len,D] - if pack["causal_seq"].shape[0] > 0: - out[pack["_causal_indices"]] = pack["causal_seq"][: pack["_causal_indices"].shape[0]] - if pack["full_only_seq"].shape[0] > 0: - out[pack["_full_indices"]] = pack["full_only_seq"][: pack["_full_indices"].shape[0]] - return out - raise KeyError("Cannot derive all_seq from provided pack") - - -def get_causal_seq(pack: SequencePack) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Get the causal sequence and its offsets in a sequence pack. - Args: - pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the causal sequence from. - Returns: - Tuple[torch.Tensor, torch.Tensor]: The concatenated causal sub-sequences and the starting offset for each sub-sequence. - """ - _ensure_core_metadata(pack) - if "causal_seq" in pack: - return pack["causal_seq"], pack["_causal_seq_offsets"] - assert "packed_sequence" in pack - return pack["packed_sequence"][pack["_causal_indices"]], pack["_causal_seq_offsets"] - - -def get_full_only_seq(pack: SequencePack) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Get the full-only sequence and its offsets in a sequence pack. - Args: - pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the full-only sequence from. - Returns: - Tuple[torch.Tensor, torch.Tensor]: The concatenated full-only sub-sequences and the starting offset for each sub-sequence. - """ - _ensure_core_metadata(pack) - if "full_only_seq" in pack: - return pack["full_only_seq"], pack["_full_only_seq_offsets"] - assert "packed_sequence" in pack - return pack["packed_sequence"][pack["_full_indices"]], pack["_full_only_seq_offsets"] - - -# ============================================================================ -# 3D mRoPE position ID utilities -# Copied from cosmos3._src.vfm.models.mot.unified_3dmrope_utils -# ============================================================================ - - -def get_3d_mrope_ids_text_tokens( - num_tokens: int, - temporal_offset: int | float, - use_float_positions: bool = False, -) -> tuple[torch.Tensor, int | float]: - """Generate 3D mRoPE position IDs for text tokens. - - For text tokens, all three axes (temporal, height, width) share the same - monotonically increasing position IDs, starting from ``temporal_offset``. - - Args: - num_tokens: Number of text tokens. - temporal_offset: Current temporal offset to start from. - use_float_positions: If True, generate float position IDs. - - Returns: - Tuple of position IDs tensor of shape (3, num_tokens) and updated temporal offset. - """ - if use_float_positions: - ids = torch.arange(num_tokens, dtype=torch.float32) + temporal_offset - else: - ids = torch.arange(num_tokens, dtype=torch.long) + int(temporal_offset) - - mrope_ids = ids.unsqueeze(0).expand(3, -1).contiguous() # [3,num_tokens] - next_temporal_offset = temporal_offset + num_tokens - return mrope_ids, next_temporal_offset - - -def get_3d_mrope_ids_vae_tokens( - grid_t: int, - grid_h: int, - grid_w: int, - temporal_offset: int | float, - reset_spatial_indices: bool = True, - fps: float | None = None, - base_fps: float = 24.0, - temporal_compression_factor: int = 4, - base_temporal_compression_factor: int | None = None, - start_frame_offset: int = 0, -) -> tuple[torch.Tensor, int | float]: - """Generate 3D mRoPE position IDs for VAE vision tokens (image/video latents). - - Args: - grid_t: Number of temporal frames in the latent grid. - grid_h: Height of the latent grid (after patchification). - grid_w: Width of the latent grid (after patchification). - temporal_offset: Current temporal offset. - reset_spatial_indices: If True, spatial indices start from 0 for each vision segment. - fps: Frames per second. If None, FPS modulation is disabled. - base_fps: Base FPS for normalization. - temporal_compression_factor: VAE temporal compression factor. - base_temporal_compression_factor: Base temporal compression factor. - start_frame_offset: Offset added to frame indices before FPS scaling. - - Returns: - Tuple of position IDs tensor of shape (3, grid_t * grid_h * grid_w) and updated offset. - """ - fps_modulation_enabled = fps is not None and grid_t > 1 - effective_base_tcf = ( - base_temporal_compression_factor - if base_temporal_compression_factor is not None - else temporal_compression_factor - ) - - if fps_modulation_enabled: - tps = fps / temporal_compression_factor - base_tps = base_fps / effective_base_tcf - frame_indices = torch.arange(grid_t, dtype=torch.float32) - scaled_t = (frame_indices + start_frame_offset) / tps * base_tps + temporal_offset - t_index = scaled_t.view(-1, 1).expand(-1, grid_h * grid_w).flatten() - t_dtype = torch.float32 - else: - t_index = ( - torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() - + int(temporal_offset) - + start_frame_offset - ) - t_dtype = torch.long - - h_index = torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten() - w_index = torch.arange(grid_w, dtype=torch.long).view(1, 1, -1).expand(grid_t, grid_h, -1).flatten() - - if not reset_spatial_indices: - spatial_offset = int(temporal_offset) - h_index = h_index + spatial_offset - w_index = w_index + spatial_offset - - if fps_modulation_enabled: - mrope_ids = torch.stack([t_index, h_index.to(torch.float32), w_index.to(torch.float32)], dim=0) - else: - mrope_ids = torch.stack([t_index, h_index, w_index], dim=0) - - max_position = mrope_ids.max().item() - next_temporal_offset = math.ceil(max_position) + 1 - return mrope_ids, next_temporal_offset - - -# ============================================================================ -# Data structures for sequence packing -# Copied from cosmos3._src.vfm.datasets.sequence_packing -# ============================================================================ - - -@dataclass -class ModalityData: - """Unified container for a single generation modality's data. - - This dataclass serves dual purposes: - 1. During packing: Acts as a builder, accumulating data in lists - 2. After finalize(): Holds finalized tensors ready for model consumption - """ - - sequence_indexes: list[int] | torch.Tensor = field(default_factory=list) - timesteps: list[float] | torch.Tensor = field(default_factory=list) - mse_loss_indexes: list[int] | torch.Tensor = field(default_factory=list) - token_shapes: list = field(default_factory=list) - - tokens: list[torch.Tensor] = field(default_factory=list) - condition_mask: list[torch.Tensor] = field(default_factory=list) - noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list) - domain_id: list[torch.Tensor] = field(default_factory=list) - raw_action_dim: list[torch.Tensor | None] | None = field(default_factory=list) - - def to_cuda(self) -> None: - if isinstance(self.sequence_indexes, torch.Tensor): - self.sequence_indexes = self.sequence_indexes.cuda() - if isinstance(self.timesteps, torch.Tensor): - self.timesteps = self.timesteps.cuda() - if isinstance(self.mse_loss_indexes, torch.Tensor): - self.mse_loss_indexes = self.mse_loss_indexes.cuda() - self.tokens = [token.cuda() for token in self.tokens] - self.condition_mask = [cm.cuda() for cm in self.condition_mask] - self.noisy_frame_indexes = [ni.cuda() for ni in self.noisy_frame_indexes] - self.domain_id = [d.cuda() for d in self.domain_id] - if self.raw_action_dim is not None: - self.raw_action_dim = [d.cuda() if d is not None else None for d in self.raw_action_dim] - - -@dataclass -class PackedSequence: - """Unified sequence container - works as builder during packing and final output.""" - - # Sequence structure - sample_lens: list[int] = field(default_factory=list) - split_lens: list[int] = field(default_factory=list) - attn_modes: list[str] = field(default_factory=list) - is_image_batch: bool = False - sequence_length: int = 0 - - # Build-time tracking - curr: int = 0 - - # Text modality (list during build, tensor after finalize) - text_ids: list[int] | torch.Tensor = field(default_factory=list) - text_indexes: list[int] | torch.Tensor = field(default_factory=list) - position_ids: list[int] | torch.Tensor = field(default_factory=list) - - # Loss computation - Cross Entropy (text) - label_ids: list[int] | torch.Tensor | None = field(default_factory=list) - ce_loss_indexes: list[int] | torch.Tensor | None = field(default_factory=list) - ce_loss_weights: list[float] | torch.Tensor | None = field(default_factory=list) - - # Build-time mRoPE tracking - _use_mrope: bool = False - _mrope_temporal_offset: int | float = 0 - _mrope_reset_spatial: bool = True - - # Temporal causal - null_action_supertokens: bool = False - - # Generation modalities - vision: ModalityData | None = None - action: ModalityData | None = None - sound: ModalityData | None = None - - def finalize( - self, - gen_data_clean: "GenerationDataClean", - ) -> "PackedSequence": - """Convert all lists to tensors and compute derived values.""" - sequence_length = sum(self.sample_lens) - sample_lens = self.sample_lens.copy() - split_lens = self.split_lens.copy() - attn_modes = self.attn_modes.copy() - - label_ids: torch.Tensor | None = None - ce_loss_indexes: torch.Tensor | None = None - ce_loss_weights: torch.Tensor | None = None - if self.label_ids and len(self.label_ids) > 0: - label_ids = torch.tensor(self.label_ids) - ce_loss_indexes = torch.tensor(self.ce_loss_indexes) - ce_loss_weights = torch.tensor(self.ce_loss_weights) - - vision: ModalityData | None = None - if self.vision is not None and len(self.vision.sequence_indexes) > 0: - vision = ModalityData( - sequence_indexes=torch.tensor(self.vision.sequence_indexes, dtype=torch.long), - timesteps=torch.tensor(self.vision.timesteps), - mse_loss_indexes=torch.tensor(self.vision.mse_loss_indexes, dtype=torch.long), - token_shapes=list(self.vision.token_shapes), - tokens=self.vision.tokens, - condition_mask=list(self.vision.condition_mask), - noisy_frame_indexes=list(self.vision.noisy_frame_indexes), - ) - - action: ModalityData | None = None - if self.action is not None and len(self.action.sequence_indexes) > 0: - action = ModalityData( - sequence_indexes=torch.tensor(self.action.sequence_indexes, dtype=torch.long), - timesteps=torch.tensor(self.action.timesteps), - mse_loss_indexes=torch.tensor(self.action.mse_loss_indexes, dtype=torch.long), - token_shapes=list(self.action.token_shapes), - tokens=self.action.tokens, - condition_mask=list(self.action.condition_mask), - noisy_frame_indexes=list(self.action.noisy_frame_indexes), - domain_id=( - gen_data_clean.action_domain_id - if gen_data_clean.action_domain_id is not None - else [torch.zeros(1, dtype=torch.long)] * len(self.action.token_shapes) - ), - raw_action_dim=gen_data_clean.raw_action_dim, - ) - - sound: ModalityData | None = None - if self.sound is not None and len(self.sound.sequence_indexes) > 0: - sound = ModalityData( - sequence_indexes=torch.tensor(self.sound.sequence_indexes, dtype=torch.long), - timesteps=torch.tensor(self.sound.timesteps), - mse_loss_indexes=torch.tensor(self.sound.mse_loss_indexes, dtype=torch.long), - token_shapes=list(self.sound.token_shapes), - tokens=self.sound.tokens, - condition_mask=list(self.sound.condition_mask), - noisy_frame_indexes=list(self.sound.noisy_frame_indexes), - ) - - if self._use_mrope and len(self.position_ids) > 0 and isinstance(self.position_ids[0], torch.Tensor): - mrope_tensors: list[torch.Tensor] = self.position_ids # type: ignore[assignment] - position_ids = torch.cat(mrope_tensors, dim=1) # [3,actual_seq_len] - else: - position_ids = torch.tensor(self.position_ids) # [seq_len] - - return PackedSequence( - sequence_length=sequence_length, - sample_lens=sample_lens, - split_lens=split_lens, - attn_modes=attn_modes, - is_image_batch=gen_data_clean.is_image_batch, - text_ids=torch.tensor(self.text_ids, dtype=torch.long), - text_indexes=torch.tensor(self.text_indexes, dtype=torch.long), - position_ids=position_ids, - label_ids=label_ids, - ce_loss_indexes=ce_loss_indexes, - ce_loss_weights=ce_loss_weights, - vision=vision, - action=action, - sound=sound, - null_action_supertokens=self.null_action_supertokens, - ) - - def to_cuda(self) -> None: - if isinstance(self.text_ids, torch.Tensor): - self.text_ids = self.text_ids.cuda() - if isinstance(self.text_indexes, torch.Tensor): - self.text_indexes = self.text_indexes.cuda() - if isinstance(self.position_ids, torch.Tensor): - self.position_ids = self.position_ids.cuda() - if isinstance(self.label_ids, torch.Tensor): - self.label_ids = self.label_ids.cuda() - if isinstance(self.ce_loss_indexes, torch.Tensor): - self.ce_loss_indexes = self.ce_loss_indexes.cuda() - if isinstance(self.ce_loss_weights, torch.Tensor): - self.ce_loss_weights = self.ce_loss_weights.cuda() - if self.vision is not None: - self.vision.to_cuda() - if self.action is not None: - self.action.to_cuda() - if self.sound is not None: - self.sound.to_cuda() - - -@dataclass -class SequencePlan: - """Plan describing which modalities are present in a sample.""" - - has_text: bool - has_vision: bool = False - condition_frame_indexes_vision: list[int] = field(default_factory=list) - has_action: bool = False - condition_frame_indexes_action: list[int] = field(default_factory=list) - has_sound: bool = False - condition_frame_indexes_sound: list[int] = field(default_factory=list) - - def as_dict(self) -> dict: - return { - "has_text": self.has_text, - "has_vision": self.has_vision, - "has_action": self.has_action, - "has_sound": self.has_sound, - "condition_frame_indexes_vision": self.condition_frame_indexes_vision, - "condition_frame_indexes_action": self.condition_frame_indexes_action, - "condition_frame_indexes_sound": self.condition_frame_indexes_sound, - } - - -# ============================================================================ -# Helper functions for packing sequences -# ============================================================================ - - -def _pack_text_tokens( - packed_seq: PackedSequence, - text_ids: List[int], - special_tokens: Dict[str, int], - curr_rope_id: int, - has_generation: bool, - use_float_positions: bool = False, -) -> Tuple[int, int, int]: - """Pack text tokens into the sequence.""" - assert isinstance(packed_seq.text_ids, list), "PackedSequence must be in build mode" - assert isinstance(packed_seq.text_indexes, list) - assert isinstance(packed_seq.position_ids, list) - assert isinstance(packed_seq.label_ids, list) - assert isinstance(packed_seq.ce_loss_indexes, list) - assert isinstance(packed_seq.ce_loss_weights, list) - - curr = packed_seq.curr - - if "bos_token_id" in special_tokens: - shifted_text_ids = [special_tokens["bos_token_id"]] + text_ids - else: - shifted_text_ids = text_ids - - split_len = 0 - - packed_seq.text_ids.extend(shifted_text_ids) - packed_seq.text_indexes.extend(range(curr, curr + len(shifted_text_ids))) - - packed_seq.ce_loss_indexes.extend(range(curr, curr + len(shifted_text_ids))) - packed_seq.ce_loss_weights.extend([1.0] * len(shifted_text_ids)) - packed_seq.label_ids.extend(text_ids[1:] + [special_tokens["eos_token_id"]]) - - curr += len(shifted_text_ids) - split_len += len(shifted_text_ids) - - packed_seq.text_ids.append(special_tokens["eos_token_id"]) - packed_seq.text_indexes.append(curr) - curr += 1 - split_len += 1 - - if has_generation: - packed_seq.text_ids.append(special_tokens["start_of_generation"]) - packed_seq.text_indexes.append(curr) - curr += 1 - split_len += 1 - - if packed_seq._use_mrope: - text_mrope_ids, packed_seq._mrope_temporal_offset = get_3d_mrope_ids_text_tokens( - num_tokens=split_len, - temporal_offset=packed_seq._mrope_temporal_offset, - use_float_positions=use_float_positions, - ) - packed_seq.position_ids.append(text_mrope_ids) - else: - packed_seq.position_ids.extend(range(curr_rope_id, curr_rope_id + split_len)) - packed_seq.attn_modes.append("causal") - packed_seq.split_lens.append(split_len) - - packed_seq.curr = curr - return curr_rope_id + split_len, split_len, split_len - - -def _pack_vision_tokens( - packed_seq: PackedSequence, - input_vision_tokens: torch.Tensor, - condition_frame_indexes_vision: list[int], - input_timestep: float | torch.Tensor, - curr_rope_id: int, - latent_patch_size: int = 1, - vision_fps: float | None = None, - enable_fps_modulation: bool = False, - base_fps: float = 24.0, - temporal_compression_factor: int = 4, -) -> int: - """Pack vision tokens into the sequence.""" - assert isinstance(packed_seq.position_ids, list), "PackedSequence must be in build mode" - - curr = packed_seq.curr - vision_split_len = 0 - - if packed_seq.vision is None: - packed_seq.vision = ModalityData() - - assert isinstance(packed_seq.vision.sequence_indexes, list) - assert isinstance(packed_seq.vision.mse_loss_indexes, list) - assert isinstance(packed_seq.vision.timesteps, list) - assert isinstance(packed_seq.vision.tokens, list) - - _, _, latent_t, latent_h, latent_w = input_vision_tokens.shape - if latent_patch_size < 1: - raise ValueError(f"latent_patch_size must be >= 1, got {latent_patch_size}") - patch_h = math.ceil(latent_h / latent_patch_size) - patch_w = math.ceil(latent_w / latent_patch_size) - packed_seq.vision.token_shapes.append((latent_t, patch_h, patch_w)) - packed_seq.vision.tokens.append(input_vision_tokens) - - num_vision_tokens = latent_t * patch_h * patch_w - packed_seq.vision.sequence_indexes.extend(range(curr, curr + num_vision_tokens)) - - condition_set = {idx for idx in condition_frame_indexes_vision if 0 <= idx < latent_t} - assert isinstance(packed_seq.vision.condition_mask, list) - - vision_condition_mask = torch.zeros( - (latent_t, 1, 1), device=input_vision_tokens.device, dtype=input_vision_tokens.dtype - ) - for frame_idx in condition_set: - vision_condition_mask[frame_idx, 0, 0] = 1.0 - packed_seq.vision.condition_mask.append(vision_condition_mask) - - vision_noisy_frame_indexes = torch.tensor( - [idx for idx in range(latent_t) if idx not in condition_set], - device=input_vision_tokens.device, - dtype=torch.long, - ) - assert isinstance(packed_seq.vision.noisy_frame_indexes, list) - packed_seq.vision.noisy_frame_indexes.append(vision_noisy_frame_indexes) - - frame_token_stride = patch_h * patch_w - for frame_idx in range(latent_t): - if frame_idx in condition_set: - continue - frame_start = curr + frame_idx * frame_token_stride - frame_end = frame_start + frame_token_stride - packed_seq.vision.mse_loss_indexes.extend(range(frame_start, frame_end)) - if isinstance(input_timestep, torch.Tensor): - frame_ts = input_timestep[frame_idx].item() - else: - frame_ts = input_timestep - packed_seq.vision.timesteps.extend([frame_ts] * frame_token_stride) - - curr += num_vision_tokens - vision_split_len += num_vision_tokens - - if packed_seq._use_mrope: - effective_fps = vision_fps if enable_fps_modulation else None - vision_mrope_ids, packed_seq._mrope_temporal_offset = get_3d_mrope_ids_vae_tokens( - grid_t=latent_t, - grid_h=patch_h, - grid_w=patch_w, - temporal_offset=packed_seq._mrope_temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, - fps=effective_fps, - base_fps=base_fps, - temporal_compression_factor=temporal_compression_factor, - ) - packed_seq.position_ids.append(vision_mrope_ids) - else: - packed_seq.position_ids.extend([curr_rope_id] * vision_split_len) - - packed_seq.curr = curr - return vision_split_len - - -def _pack_action_tokens( - packed_seq: PackedSequence, - input_action_tokens: torch.Tensor, - condition_frame_indexes_action: list[int], - input_timestep: float, - curr_rope_id: int, - action_temporal_offset: int | float = 0, - enable_fps_modulation: bool = False, - base_fps: float = 24.0, - action_fps: float | None = None, - base_temporal_compression_factor: int | None = None, -) -> int: - """Pack action tokens into the sequence.""" - assert isinstance(packed_seq.position_ids, list), "PackedSequence must be in build mode" - - curr = packed_seq.curr - action_split_len = input_action_tokens.shape[0] - - if packed_seq.action is None: - packed_seq.action = ModalityData() - - assert isinstance(packed_seq.action.sequence_indexes, list) - assert isinstance(packed_seq.action.mse_loss_indexes, list) - assert isinstance(packed_seq.action.timesteps, list) - assert isinstance(packed_seq.action.tokens, list) - - action_indexes = list(range(curr, curr + action_split_len)) - packed_seq.action.sequence_indexes.extend(action_indexes) - packed_seq.action.token_shapes.append((action_split_len,)) - packed_seq.action.tokens.append(input_action_tokens) - - condition_set = {idx for idx in condition_frame_indexes_action if 0 <= idx < action_split_len} - assert isinstance(packed_seq.action.condition_mask, list) - - action_condition_mask = torch.zeros( - (action_split_len, 1), device=input_action_tokens.device, dtype=input_action_tokens.dtype - ) - for frame_idx in condition_set: - action_condition_mask[frame_idx, 0] = 1.0 - packed_seq.action.condition_mask.append(action_condition_mask) - - action_noisy_frame_indexes = torch.tensor( - [idx for idx in range(action_split_len) if idx not in condition_set], - device=input_action_tokens.device, - dtype=torch.long, - ) - assert isinstance(packed_seq.action.noisy_frame_indexes, list) - packed_seq.action.noisy_frame_indexes.append(action_noisy_frame_indexes) - - frame_token_stride = 1 - for frame_idx in range(action_split_len): - if frame_idx in condition_set: - continue - frame_start = curr + frame_idx * frame_token_stride - frame_end = frame_start + frame_token_stride - packed_seq.action.mse_loss_indexes.extend(range(frame_start, frame_end)) - packed_seq.action.timesteps.extend([input_timestep] * frame_token_stride) - - if packed_seq._use_mrope: - effective_fps = action_fps if enable_fps_modulation else None - action_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( - grid_t=action_split_len, - grid_h=1, - grid_w=1, - temporal_offset=action_temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, - fps=effective_fps, - base_fps=base_fps, - temporal_compression_factor=1, - base_temporal_compression_factor=base_temporal_compression_factor, - start_frame_offset=1, - ) - packed_seq.position_ids.append(action_mrope_ids) - else: - packed_seq.position_ids.extend([curr_rope_id] * action_split_len) - - packed_seq.curr = curr + action_split_len - return action_split_len - - -def _pack_sound_tokens( - packed_seq: PackedSequence, - input_sound_tokens: torch.Tensor, - condition_frame_indexes_sound: list[int], - input_timestep: float, - curr_rope_id: int, - sound_temporal_offset: int | float = 0, - enable_fps_modulation: bool = False, - base_fps: float = 24.0, - sound_fps: float | None = None, -) -> int: - """Pack sound/audio tokens into the sequence.""" - assert isinstance(packed_seq.position_ids, list), "PackedSequence must be in build mode" - - curr = packed_seq.curr - _, sound_split_len = input_sound_tokens.shape - - if packed_seq.sound is None: - packed_seq.sound = ModalityData() - - assert isinstance(packed_seq.sound.sequence_indexes, list) - assert isinstance(packed_seq.sound.mse_loss_indexes, list) - assert isinstance(packed_seq.sound.timesteps, list) - assert isinstance(packed_seq.sound.tokens, list) - - packed_seq.sound.token_shapes.append((sound_split_len, 1, 1)) - packed_seq.sound.sequence_indexes.extend(range(curr, curr + sound_split_len)) - packed_seq.sound.tokens.append(input_sound_tokens) - - condition_set = {idx for idx in condition_frame_indexes_sound if 0 <= idx < sound_split_len} - assert isinstance(packed_seq.sound.condition_mask, list) - - sound_condition_mask = torch.zeros( - (sound_split_len, 1), device=input_sound_tokens.device, dtype=input_sound_tokens.dtype - ) - for frame_idx in condition_set: - sound_condition_mask[frame_idx, 0] = 1.0 - packed_seq.sound.condition_mask.append(sound_condition_mask) - - sound_noisy_frame_indexes = torch.tensor( - [idx for idx in range(sound_split_len) if idx not in condition_set], - device=input_sound_tokens.device, - dtype=torch.long, - ) - assert isinstance(packed_seq.sound.noisy_frame_indexes, list) - packed_seq.sound.noisy_frame_indexes.append(sound_noisy_frame_indexes) - - for frame_idx in range(sound_split_len): - if frame_idx in condition_set: - continue - frame_start = curr + frame_idx - frame_end = frame_start + 1 - packed_seq.sound.mse_loss_indexes.extend(range(frame_start, frame_end)) - packed_seq.sound.timesteps.extend([input_timestep]) - - if packed_seq._use_mrope: - effective_fps = sound_fps if enable_fps_modulation else None - sound_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( - grid_t=sound_split_len, - grid_h=1, - grid_w=1, - temporal_offset=sound_temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, - fps=effective_fps, - base_fps=base_fps, - temporal_compression_factor=1, - start_frame_offset=0, - ) - packed_seq.position_ids.append(sound_mrope_ids) - else: - packed_seq.position_ids.extend([curr_rope_id] * sound_split_len) - - packed_seq.curr = curr + sound_split_len - return sound_split_len - - -def _pack_supertokens_temporal_causal( - packed_seq: "PackedSequence", - input_vision_tokens: torch.Tensor, - input_action_tokens: torch.Tensor | None, - condition_frame_indexes_vision: list[int], - input_timestep: float | torch.Tensor, - curr_rope_id: int, - latent_patch_size: int, - temporal_compression_factor: int, - action_dim: int, - vision_fps: float | None = None, - action_fps: float | None = None, - enable_fps_modulation: bool = False, - base_fps: float = 24.0, -) -> tuple[int, bool]: - """Pack vision and action tokens in interleaved supertoken order for temporal causal attention. - - Buffer layout: [action_t0, vision_t0, action_t1, vision_t1, ..., action_{T-1}, vision_{T-1}] - """ - assert isinstance(packed_seq.position_ids, list), "PackedSequence must be in build mode" - - _, _, latent_t, latent_h, latent_w = input_vision_tokens.shape - patch_h = math.ceil(latent_h / latent_patch_size) - patch_w = math.ceil(latent_w / latent_patch_size) - tcf = temporal_compression_factor - patches_per_frame = patch_h * patch_w - supertoken_len = tcf + patches_per_frame - - if packed_seq.vision is None: - packed_seq.vision = ModalityData() - if packed_seq.action is None: - packed_seq.action = ModalityData() - - assert isinstance(packed_seq.vision.sequence_indexes, list) - assert isinstance(packed_seq.vision.mse_loss_indexes, list) - assert isinstance(packed_seq.vision.timesteps, list) - assert isinstance(packed_seq.vision.tokens, list) - assert isinstance(packed_seq.vision.condition_mask, list) - assert isinstance(packed_seq.action.sequence_indexes, list) - assert isinstance(packed_seq.action.mse_loss_indexes, list) - assert isinstance(packed_seq.action.timesteps, list) - assert isinstance(packed_seq.action.tokens, list) - assert isinstance(packed_seq.action.condition_mask, list) - - device = input_vision_tokens.device - dtype = input_vision_tokens.dtype - - null_tokens = torch.zeros(tcf, action_dim, device=device, dtype=dtype) - if input_action_tokens is not None: - if input_action_tokens.dim() == 3: - real_actions = input_action_tokens.squeeze(0) - else: - real_actions = input_action_tokens - if latent_t == 1: - all_action_tokens = real_actions - else: - all_action_tokens = torch.cat([null_tokens, real_actions], dim=0) - else: - all_action_tokens = torch.zeros(latent_t * tcf, action_dim, device=device, dtype=dtype) - - null_action_flag = not (latent_t == 1 and input_action_tokens is not None) - - packed_seq.vision.token_shapes.append((latent_t, patch_h, patch_w)) - packed_seq.vision.tokens.append(input_vision_tokens) - - condition_set_vision = {idx for idx in condition_frame_indexes_vision if 0 <= idx < latent_t} - vision_condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=dtype) - for fidx in condition_set_vision: - vision_condition_mask[fidx, 0, 0] = 1.0 - packed_seq.vision.condition_mask.append(vision_condition_mask) - - vision_noisy_frame_indexes = torch.tensor( - [idx for idx in range(latent_t) if idx not in condition_set_vision], - device=device, - dtype=torch.long, - ) - packed_seq.vision.noisy_frame_indexes.append(vision_noisy_frame_indexes) - - packed_seq.action.token_shapes.append((latent_t * tcf,)) - packed_seq.action.tokens.append(all_action_tokens) - - action_condition_mask = torch.ones((latent_t * tcf, 1), device=device, dtype=dtype) - packed_seq.action.condition_mask.append(action_condition_mask) - - curr = packed_seq.curr - total_split_len = 0 - - if packed_seq._use_mrope: - temporal_offset = packed_seq._mrope_temporal_offset - effective_action_fps = action_fps if enable_fps_modulation else None - effective_vision_fps = vision_fps if enable_fps_modulation else None - - fps_active = effective_action_fps is not None - t_dtype = torch.float32 if fps_active else torch.long - t_offset = float(temporal_offset) if fps_active else int(temporal_offset) - null_t = torch.full((tcf,), t_offset, dtype=t_dtype) - null_hw = torch.zeros(tcf, dtype=t_dtype) - null_ids = torch.stack([null_t, null_hw, null_hw]) # [3,tcf] - - def _real_action_ids(n_frames: int, start_frame_offset: int) -> torch.Tensor: - flat, _ = get_3d_mrope_ids_vae_tokens( - grid_t=n_frames * tcf, - grid_h=1, - grid_w=1, - temporal_offset=temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, - fps=effective_action_fps, - base_fps=base_fps, - temporal_compression_factor=1, - base_temporal_compression_factor=tcf, - start_frame_offset=start_frame_offset, - ) - return flat.reshape(3, n_frames, tcf) - - if latent_t > 1: - null_ids_3d = null_ids.reshape(3, 1, tcf) - real_ids_3d = _real_action_ids(latent_t - 1, start_frame_offset=1) - action_ids_3d = torch.cat([null_ids_3d, real_ids_3d], dim=1) - elif input_action_tokens is None: - action_ids_3d = null_ids.reshape(3, 1, tcf) - else: - action_ids_3d = _real_action_ids(1, start_frame_offset=0) - - vision_ids_flat, new_offset = get_3d_mrope_ids_vae_tokens( - grid_t=latent_t, - grid_h=patch_h, - grid_w=patch_w, - temporal_offset=temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, - fps=effective_vision_fps, - base_fps=base_fps, - temporal_compression_factor=tcf, - ) - vision_ids_3d = vision_ids_flat.reshape(3, latent_t, patches_per_frame) - - interleaved_ids = torch.cat([action_ids_3d, vision_ids_3d], dim=2).reshape(3, latent_t * supertoken_len) - packed_seq.position_ids.append(interleaved_ids) - packed_seq._mrope_temporal_offset = new_offset - - for frame_t in range(latent_t): - action_indexes = list(range(curr, curr + tcf)) - packed_seq.action.sequence_indexes.extend(action_indexes) - curr += tcf - total_split_len += tcf - - if not packed_seq._use_mrope: - packed_seq.position_ids.extend([curr_rope_id] * tcf) - - frame_indexes = list(range(curr, curr + patches_per_frame)) - packed_seq.vision.sequence_indexes.extend(frame_indexes) - curr += patches_per_frame - total_split_len += patches_per_frame - - if not packed_seq._use_mrope: - packed_seq.position_ids.extend([curr_rope_id] * patches_per_frame) - - if frame_t not in condition_set_vision: - packed_seq.vision.mse_loss_indexes.extend(frame_indexes) - frame_ts = input_timestep[frame_t].item() if isinstance(input_timestep, torch.Tensor) else input_timestep - packed_seq.vision.timesteps.extend([frame_ts] * patches_per_frame) - - packed_seq.curr = curr - return total_split_len, null_action_flag - - -# ============================================================================ -# Main packing functions -# ============================================================================ - - -def pack_input_sequence( - sequence_plans: list[SequencePlan], - input_text_indexes: list[list[int]], - gen_data_clean: GenerationDataClean, - input_timesteps: torch.Tensor, - special_tokens: dict[str, int], - max_num_tokens: int | None = None, - latent_patch_size: int = 1, - skip_text_tokens: bool = False, - include_end_of_generation_token: bool = False, - position_embedding_type: str = "3d_rope", - unified_3d_mrope_reset_spatial_ids: bool = True, - unified_3d_mrope_temporal_modality_margin: int = 0, - enable_fps_modulation: bool = False, - base_fps: float = 24.0, - temporal_compression_factor: int = 4, - video_temporal_causal: bool = False, - action_dim: int = 32, - initial_mrope_temporal_offset: int | float = 0, -) -> PackedSequence: - """Pack a sequence of input strings and VAE latents into a packed tensor format. - - Args: - sequence_plans: List of SequencePlan items describing which modalities are present. - input_text_indexes: List of text token ID sequences. - gen_data_clean: GenerationDataClean containing vision, action, and sound tensors. - input_timesteps: Diffusion timesteps for each sample. Shape (B,) or (B, 1). - special_tokens: Dictionary containing special token IDs. - max_num_tokens: Maximum number of tokens (unused, kept for API compatibility). - latent_patch_size: Patch size used by the network to pack latents. - skip_text_tokens: If True, skip packing text tokens. - include_end_of_generation_token: If True, append end-of-generation token. - position_embedding_type: Position embedding type for vision tokens. - unified_3d_mrope_reset_spatial_ids: If True, spatial indices start from 0 per segment. - unified_3d_mrope_temporal_modality_margin: Temporal margin between text and vision. - enable_fps_modulation: If True, scale temporal position IDs based on video FPS. - base_fps: Base FPS for normalization. - temporal_compression_factor: VAE temporal compression factor. - video_temporal_causal: If True, pack vision and action as interleaved supertokens. - action_dim: Action token dimension for temporal causal packing. - initial_mrope_temporal_offset: Initial temporal offset for AR inference. - - Returns: - PackedSequence containing all packed tensors and metadata. - """ - del max_num_tokens - - assert special_tokens is not None, "Special tokens must be provided" - assert isinstance(input_timesteps, torch.Tensor), "input_timesteps must be a tensor" - if input_timesteps.is_cuda: - raise ValueError("input_timesteps must be on CPU, not CUDA") - if isinstance(input_text_indexes, torch.Tensor): - raise ValueError("input_text_tokens must be a list, not a tensor") - - packed_seq = PackedSequence() - packed_seq._use_mrope = position_embedding_type == "unified_3d_mrope" - packed_seq._mrope_reset_spatial = unified_3d_mrope_reset_spatial_ids - - idx_text = 0 - idx_vision = 0 - idx_action = 0 - idx_sound = 0 - null_action_flags: list[bool] = [] - - if not skip_text_tokens: - for plan in sequence_plans: - assert plan.has_text, "All sequence plans must have has_text=True when skip_text_tokens=False" - - for sample_idx, sequence_plan in enumerate(sequence_plans): - curr_rope_id = 0 - sample_len = 0 - - packed_seq._mrope_temporal_offset = initial_mrope_temporal_offset - - _ts = input_timesteps[sample_idx] - input_timestep = _ts.item() if _ts.numel() == 1 else _ts - - if sequence_plan.has_text and not skip_text_tokens: - text_ids = input_text_indexes[idx_text] - idx_text += 1 - - has_generation_for_sample = sequence_plan.has_vision or sequence_plan.has_action or sequence_plan.has_sound - curr_rope_id, _, text_sample_len = _pack_text_tokens( - packed_seq, - text_ids, - special_tokens, - curr_rope_id, - has_generation=has_generation_for_sample, - use_float_positions=enable_fps_modulation, - ) - sample_len += text_sample_len - packed_seq._mrope_temporal_offset += unified_3d_mrope_temporal_modality_margin - - vision_start_temporal_offset = packed_seq._mrope_temporal_offset - - if video_temporal_causal and sequence_plan.has_vision: - assert position_embedding_type == "unified_3d_mrope", ( - "video_temporal_causal=True requires position_embedding_type='unified_3d_mrope'" - ) - input_vision_tokens = gen_data_clean.x0_tokens_vision[idx_vision] - idx_vision += 1 - - vision_fps = None - if ( - enable_fps_modulation - and gen_data_clean.fps_vision is not None - and idx_vision - 1 < len(gen_data_clean.fps_vision) - ): - vision_fps = float(gen_data_clean.fps_vision[idx_vision - 1].item()) - - input_action_tokens_tc: torch.Tensor | None = None - action_fps_tc: float | None = None - if sequence_plan.has_action: - input_action_tokens_tc = gen_data_clean.x0_tokens_action[idx_action] - if ( - enable_fps_modulation - and gen_data_clean.fps_action is not None - and idx_action < len(gen_data_clean.fps_action) - ): - action_fps_tc = float(gen_data_clean.fps_action[idx_action].item()) - idx_action += 1 - - supertoken_split_len, null_flag = _pack_supertokens_temporal_causal( - packed_seq=packed_seq, - input_vision_tokens=input_vision_tokens, - input_action_tokens=input_action_tokens_tc, - condition_frame_indexes_vision=sequence_plan.condition_frame_indexes_vision, - input_timestep=input_timestep, - curr_rope_id=curr_rope_id, - latent_patch_size=latent_patch_size, - temporal_compression_factor=temporal_compression_factor, - action_dim=action_dim, - vision_fps=vision_fps, - action_fps=action_fps_tc, - enable_fps_modulation=enable_fps_modulation, - base_fps=base_fps, - ) - null_action_flags.append(null_flag) - sample_len += supertoken_split_len - vision_split_len = supertoken_split_len - action_split_len = 0 - - else: - if sequence_plan.has_vision: - num_vis = ( - gen_data_clean.num_vision_items_per_sample[sample_idx] - if gen_data_clean.num_vision_items_per_sample is not None - else 1 - ) - - vision_split_len = 0 - for item_idx in range(num_vis): - input_vision_tokens = gen_data_clean.x0_tokens_vision[idx_vision] - - vision_fps: float | None = None - if ( - enable_fps_modulation - and gen_data_clean.fps_vision is not None - and idx_vision < len(gen_data_clean.fps_vision) - ): - vision_fps = float(gen_data_clean.fps_vision[idx_vision].item()) - - idx_vision += 1 - - if num_vis > 1 and item_idx < num_vis - 1: - latent_t = input_vision_tokens.shape[2] - item_condition_frames = list(range(latent_t)) - else: - item_condition_frames = sequence_plan.condition_frame_indexes_vision - - item_split_len = _pack_vision_tokens( - packed_seq=packed_seq, - input_vision_tokens=input_vision_tokens, - condition_frame_indexes_vision=item_condition_frames, - input_timestep=input_timestep, - curr_rope_id=curr_rope_id, - latent_patch_size=latent_patch_size, - vision_fps=vision_fps, - enable_fps_modulation=enable_fps_modulation, - base_fps=base_fps, - temporal_compression_factor=temporal_compression_factor, - ) - vision_split_len += item_split_len - sample_len += vision_split_len - - else: - vision_split_len = 0 - - if sequence_plan.has_action: - input_action_tokens = gen_data_clean.x0_tokens_action[idx_action] - - action_fps: float | None = None - if ( - enable_fps_modulation - and gen_data_clean.fps_action is not None - and idx_action < len(gen_data_clean.fps_action) - ): - action_fps = float(gen_data_clean.fps_action[idx_action].item()) - - idx_action += 1 - - action_split_len = _pack_action_tokens( - packed_seq=packed_seq, - input_action_tokens=input_action_tokens, - condition_frame_indexes_action=sequence_plan.condition_frame_indexes_action, - input_timestep=input_timestep, - curr_rope_id=curr_rope_id, - action_temporal_offset=vision_start_temporal_offset, - enable_fps_modulation=enable_fps_modulation, - base_fps=base_fps, - action_fps=action_fps, - base_temporal_compression_factor=temporal_compression_factor, - ) - sample_len += action_split_len - else: - action_split_len = 0 - - if sequence_plan.has_sound: - input_sound_tokens = gen_data_clean.x0_tokens_sound[idx_sound] - - sound_fps: float | None = None - if ( - enable_fps_modulation - and gen_data_clean.fps_sound is not None - and idx_sound < len(gen_data_clean.fps_sound) - ): - sound_fps = float(gen_data_clean.fps_sound[idx_sound].item()) - - idx_sound += 1 - - sound_split_len = _pack_sound_tokens( - packed_seq=packed_seq, - input_sound_tokens=input_sound_tokens, - condition_frame_indexes_sound=sequence_plan.condition_frame_indexes_sound, - input_timestep=input_timestep, - curr_rope_id=curr_rope_id, - sound_temporal_offset=vision_start_temporal_offset, - enable_fps_modulation=enable_fps_modulation, - base_fps=base_fps, - sound_fps=sound_fps, - ) - sample_len += sound_split_len - else: - sound_split_len = 0 - - eov_len = 0 - has_any_generation = sequence_plan.has_vision or sequence_plan.has_action or sequence_plan.has_sound - if include_end_of_generation_token and has_any_generation: - assert isinstance(packed_seq.text_ids, list) - assert isinstance(packed_seq.text_indexes, list) - assert isinstance(packed_seq.position_ids, list) - - packed_seq.text_ids.append(special_tokens["end_of_generation"]) - packed_seq.text_indexes.append(packed_seq.curr) - - if packed_seq._use_mrope: - eov_dtype = torch.float32 if enable_fps_modulation else torch.long - eov_mrope_ids = torch.full((3, 1), packed_seq._mrope_temporal_offset, dtype=eov_dtype) - packed_seq.position_ids.append(eov_mrope_ids) # type: ignore[arg-type] - packed_seq._mrope_temporal_offset += 1 - else: - packed_seq.position_ids.append(curr_rope_id) # type: ignore[arg-type] - - packed_seq.curr += 1 - eov_len = 1 - sample_len += 1 - - combined_split_len = vision_split_len + action_split_len + sound_split_len + eov_len - packed_seq.attn_modes.append("full") - packed_seq.split_lens.append(combined_split_len) - packed_seq.sample_lens.append(sample_len) - - if null_action_flags: - assert len(set(null_action_flags)) == 1, ( - f"Inconsistent null_action_supertokens across samples: {null_action_flags}." - ) - packed_seq.null_action_supertokens = null_action_flags[0] - - return packed_seq.finalize(gen_data_clean=gen_data_clean) - - -def build_sequence_plans_from_data_batch( - data_batch: dict, - input_video_key, - input_image_key: str, -) -> list[SequencePlan]: - """Build or retrieve sequence plans from a data batch dictionary. - - This function extracts sequence plans from the data batch if they exist, - otherwise creates default SequencePlan objects for each sample in the batch. - - Args: - data_batch: Dictionary containing the data batch from the dataloader. - input_video_key: Key for video tensors in the batch. - input_image_key: Key for image tensors in the batch. - - Returns: - List of SequencePlan objects, one per sample in the batch. - """ - # NOTE: this function is ONLY intended for backward compatibility. - # For new modalities, please generate the sequence_plan in the dataset class. - - if "sequence_plan" in data_batch: - return data_batch["sequence_plan"] - - assert "action" not in data_batch or data_batch["action"] is None, "Action data SHOULD have sequence_plans!" - assert "sound" not in data_batch or data_batch["sound"] is None, "Sound data SHOULD have sequence_plans!" - - batch_size = 0 - for key in [input_video_key, input_image_key]: - if key in data_batch: - val = data_batch[key] - if isinstance(val, torch.Tensor): - batch_size = val.shape[0] - break - elif isinstance(val, list): - batch_size = len(val) - break - - if batch_size == 0: - raise ValueError( - f"Cannot determine batch size from data_batch. Expected {input_video_key}, {input_image_key}, or similar key." - ) - - return [ - SequencePlan( - has_text=True, - has_vision=True, - condition_frame_indexes_vision=[], - ) - for _ in range(batch_size) - ] diff --git a/packages/diffusers-cosmos3/diffusers_cosmos3/transformer.py b/packages/diffusers-cosmos3/diffusers_cosmos3/transformer.py deleted file mode 100644 index 7b432712..00000000 --- a/packages/diffusers-cosmos3/diffusers_cosmos3/transformer.py +++ /dev/null @@ -1,819 +0,0 @@ -# Copyright 2025 The NVIDIA Team and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import math -from typing import Optional, Tuple - -import torch -import torch.nn as nn -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.models.attention_dispatch import dispatch_attention_fn -from diffusers.models.modeling_utils import ModelMixin -from transformers.activations import ACT2FN -from transformers.integrations import use_kernel_forward_from_hub -from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update -from transformers.models.qwen3_vl.modeling_qwen3_vl import apply_rotary_pos_emb - -from diffusers_cosmos3.sequence_packing import ( - FactoredSequencePack, - from_joint, - from_mode_splits, - from_und_gen_splits, - get_all_seq, - get_causal_seq, - get_device_and_dtype, - get_full_only_seq, - get_gen_seq, - get_und_seq, - set_gen_seq, - set_und_seq, - zeros_like, -) - - -def _pack_to_batch(tokens: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int) -> torch.Tensor: - """Unpack (total_tokens, heads, dim) → (batch, max_seqlen, heads, dim).""" - batch = cu_seqlens.shape[0] - 1 - cu = cu_seqlens.tolist() - out = tokens.new_zeros(batch, max_seqlen, *tokens.shape[1:]) - for i in range(batch): - n = cu[i + 1] - cu[i] - out[i, :n] = tokens[cu[i] : cu[i + 1]] - return out - - -def _batch_to_pack(batched: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: - """Repack (batch, max_seqlen, heads, dim) → (total_tokens, heads, dim).""" - cu = cu_seqlens.tolist() - return torch.cat([batched[i, : cu[i + 1] - cu[i]] for i in range(len(cu) - 1)], dim=0) - - -def _kv_padding_mask(cu_seqlens: torch.Tensor, max_seqlen: int, dtype: torch.dtype, device: torch.device): - """Float mask (batch, 1, 1, max_seqlen) with -inf at padding positions, or None if uniform.""" - batch = cu_seqlens.shape[0] - 1 - cu = cu_seqlens.tolist() - mask = torch.zeros(batch, 1, 1, max_seqlen, dtype=dtype, device=device) - for i in range(batch): - kl = cu[i + 1] - cu[i] - if kl < max_seqlen: - mask[i, 0, 0, kl:] = float("-inf") - return None if (mask == 0).all() else mask - - -class CosmosAttnProcessor3_0: - """ - Packed two-way attention processor for Cosmos3. Implements separate causal - (understanding) and full (generation) attention pathways via dispatch_attention_fn. - """ - - def __call__( - self, - packed_query_states: FactoredSequencePack, - packed_key_states: FactoredSequencePack, - packed_value_states: FactoredSequencePack, - ) -> FactoredSequencePack: - causal_q, causal_offsets = get_causal_seq(packed_query_states) - causal_k, _ = get_causal_seq(packed_key_states) - causal_v, _ = get_causal_seq(packed_value_states) - full_q, full_offsets = get_full_only_seq(packed_query_states) - sample_offsets = packed_query_states["sample_offsets"] - max_causal = packed_query_states["max_causal_len"] - max_full = packed_query_states["max_full_len"] - max_sample = packed_query_states["max_sample_len"] - - # Causal (understanding) self-attention - causal_out = dispatch_attention_fn( - _pack_to_batch(causal_q, causal_offsets, max_causal), - _pack_to_batch(causal_k, causal_offsets, max_causal), - _pack_to_batch(causal_v, causal_offsets, max_causal), - is_causal=True, - enable_gqa=True, - ) - causal_out = _batch_to_pack(causal_out, causal_offsets).flatten(-2, -1) - - # Full (generation) cross-attention: Q = gen tokens, K/V = all tokens - all_k = get_all_seq(packed_key_states) - all_v = get_all_seq(packed_value_states) - full_out = dispatch_attention_fn( - _pack_to_batch(full_q, full_offsets, max_full), - _pack_to_batch(all_k, sample_offsets, max_sample), - _pack_to_batch(all_v, sample_offsets, max_sample), - attn_mask=_kv_padding_mask(sample_offsets, max_sample, causal_q.dtype, causal_q.device), - is_causal=False, - enable_gqa=True, - ) - full_out = _batch_to_pack(full_out, full_offsets).flatten(-2, -1) - - return from_mode_splits(causal_out, full_out, packed_query_states) - - -class TimestepEmbedder(nn.Module): - """Embeds scalar timesteps into vector representations.""" - - def __init__(self, hidden_size, frequency_embedding_size=256): - super().__init__() - self.linear_1 = nn.Linear(frequency_embedding_size, hidden_size, bias=True) - self.act = nn.SiLU() - self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True) - self.frequency_embedding_size = frequency_embedding_size - self.hidden_size = hidden_size - - def _init_weights(self): - std = 1.0 / math.sqrt(self.frequency_embedding_size) - torch.nn.init.trunc_normal_(self.mlp[0].weight, std=std, a=-3 * std, b=3 * std) - torch.nn.init.zeros_(self.mlp[0].bias) - - std = 1.0 / math.sqrt(self.hidden_size) - torch.nn.init.trunc_normal_(self.mlp[2].weight, std=std, a=-3 * std, b=3 * std) - torch.nn.init.zeros_(self.mlp[2].bias) - - @staticmethod - def timestep_embedding(t, dim, max_period=10000): - half = dim // 2 - freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( - device=t.device - ) - args = t[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) - return embedding - - def forward(self, t): - t_freq = self.timestep_embedding(t, self.frequency_embedding_size) - return self.linear_2(self.act(self.linear_1(t_freq))) - - -class DomainAwareLinear(nn.Module): - """Linear projection with one weight/bias pair per action embodiment domain.""" - - def __init__(self, input_size: int, output_size: int, num_domains: int) -> None: - super().__init__() - self.input_size = int(input_size) - self.output_size = int(output_size) - self.num_domains = int(num_domains) - self.fc = nn.Embedding(self.num_domains, self.output_size * self.input_size) - self.bias = nn.Embedding(self.num_domains, self.output_size) - nn.init.xavier_uniform_(self.fc.weight) - nn.init.zeros_(self.bias.weight) - - def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: - if domain_id.ndim == 0: - domain_id = domain_id.unsqueeze(0) - domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1) - if x.shape[0] != domain_id.shape[0]: - raise ValueError( - "Cosmos3 action domain_id batch size must match action tokens: " - f"tokens={x.shape[0]}, domain_id={domain_id.shape[0]}." - ) - if torch.any((domain_id < 0) | (domain_id >= self.num_domains)): - raise ValueError(f"Cosmos3 action domain_id must be in [0, {self.num_domains}), got {domain_id.tolist()}.") - - weight = self.fc(domain_id).view(domain_id.shape[0], self.input_size, self.output_size) - bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size) - if x.ndim == 2: - return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias - if x.ndim == 3: - return torch.bmm(x, weight) + bias.unsqueeze(1) - raise ValueError(f"Cosmos3 DomainAwareLinear expected rank-2 or rank-3 input, got {tuple(x.shape)}.") - - -class LayerTypes: - def __init__(self, is_moe: bool): - self.is_moe = is_moe - if is_moe: # TODO: moe is not yet tested - self.mlp = Qwen3VLMoeTextMLP - self.rms_norm = Qwen3VLMoeTextRMSNorm - self.rotary_embedding = Qwen3VLMoeTextRotaryEmbedding - else: - self.mlp = Cosmos3VLTextMLP - self.rms_norm = Cosmos3VLTextRMSNorm - self.rotary_embedding = Cosmos3VLTextRotaryEmbedding - - -class Cosmos3VLTextRotaryEmbedding(nn.Module): - def __init__(self, config): - super().__init__() - if hasattr(config, "rope_scaling") and config.rope_scaling is not None: - self.rope_type = config.rope_scaling.get("rope_type", "default") - else: - self.rope_type = "default" - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings - - self.config = config - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - - self.mrope_section = ( - config.rope_scaling.get("mrope_section", [24, 20, 20]) if config.rope_scaling is not None else [24, 20, 20] - ) - - def init_weights(self, buffer_device: torch.device | None = None) -> None: - inv_freq, self.attention_scaling = self.rope_init_fn(self.config, buffer_device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - - def apply_interleaved_mrope(self, freqs, mrope_section): - """Apply interleaved MRoPE to 3D rotary embeddings. - Reorganizes frequency layout from chunked [TTT...HHH...WWW] to - interleaved [THTHWHTHW...TT], preserving frequency continuity. - args: - x: (3, bs, seq_len, head_dim // 2) - mrope_section: (3,) - returns: - x_t: (bs, seq_len, head_dim // 2) - """ - freqs_t = freqs[0] # just overwrite the first dimension T - for dim, offset in enumerate((1, 2), start=1): # H, W - length = mrope_section[dim] * 3 - idx = slice(offset, length, 3) - freqs_t[..., idx] = freqs[dim, ..., idx] - return freqs_t - - @torch.no_grad() - @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) - def forward(self, x, position_ids): - assert self.inv_freq.dtype == torch.float32, f"inv_freq must be float32, but got {self.inv_freq.dtype}" - - # In contrast to other models, Cosmos3Omni has different position ids for the grids - # So we expand the inv_freq to shape (3, ...) - if position_ids.ndim == 2: - position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) # [3,B,N] - inv_freq_expanded = ( - self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1).to(x.device) - ) # [3,B,head_dim//2,1] - position_ids_expanded = position_ids[:, :, None, :].float() # [3,B,1,N] - - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) # [3,B,N,head_dim//2] - freqs = self.apply_interleaved_mrope(freqs, self.mrope_section) # [B,N,head_dim//2] - emb = torch.cat((freqs, freqs), dim=-1) # [B,N,head_dim] - cos = emb.cos() * self.attention_scaling # [B,N,head_dim] - sin = emb.sin() * self.attention_scaling # [B,N,head_dim] - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) # each: [B,N,head_dim] - - -class Cosmos3VLTextRMSNorm(nn.Module): - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: - """ - Cosmos3VLTextRMSNorm is equivalent to T5LayerNorm - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) - - def extra_repr(self) -> str: - return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" - - -class Cosmos3VLTextMLP(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj - - -class Cosmos3VLTextAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config, layer_idx: int): - super().__init__() - self.config = config - self.layer_idx = layer_idx - self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads - self.scaling = self.head_dim**-0.5 - self.attention_dropout = config.attention_dropout - self.is_causal = True - - self.to_q = nn.Linear( - config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias - ) - self.to_k = nn.Linear( - config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias - ) - self.to_v = nn.Linear( - config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias - ) - self.to_out = nn.Linear( - config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias - ) - self.norm_q = Cosmos3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) # unlike olmo, only on the head dim! - self.norm_k = Cosmos3VLTextRMSNorm( - self.head_dim, eps=config.rms_norm_eps - ) # thus post norm_q does not need reshape - - -class PackedAttentionMoT(Cosmos3VLTextAttention): - """ - Dual-pathway packed attention for Qwen3VL MoT (Dense version). - Implements understanding and generation pathways with separate projections. - - Note that this implementation is used for both Qwen3VL and Qwen3VL-MoE variants, - even though it derives from the dense version of Qwen3VLTextAttention. - """ - - def __init__(self, config, layer_idx: int, layer_types: LayerTypes): - super().__init__(config, layer_idx) - - # Add missing attributes for MoT compatibility - self.hidden_size = config.hidden_size - self.num_attention_heads = config.num_attention_heads - self.num_key_value_heads = config.num_key_value_heads - self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads - self.scaling = self.head_dim**-0.5 - self.attention_dropout = config.attention_dropout - - # Generation pathway projections (separate from understanding pathway) - # Qwen3VL already has query/key norms built in, so we add generation versions - self.norm_added_q = layer_types.rms_norm(self.head_dim, eps=config.rms_norm_eps) - self.norm_added_k = layer_types.rms_norm(self.head_dim, eps=config.rms_norm_eps) - - # Generation pathway linear projections - self.add_q_proj = nn.Linear( - self.hidden_size, self.num_attention_heads * self.head_dim, bias=config.attention_bias - ) - self.add_k_proj = nn.Linear( - self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias - ) - self.add_v_proj = nn.Linear( - self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias - ) - self.to_add_out = nn.Linear( - self.num_attention_heads * self.head_dim, self.hidden_size, bias=config.attention_bias - ) - self.dispatch_attention_fn = CosmosAttnProcessor3_0() - self.cp_mesh = None - - def forward( - self, - pack: FactoredSequencePack, - attention_mask, - packed_position_embeddings: Tuple[FactoredSequencePack, FactoredSequencePack], - dual_kv_cache=None, - natten_metadata: dict | None = None, - ) -> FactoredSequencePack: - """Forward pass with optional KV cache for autoregressive generation. - - This method is used for frame 0 where we store K/V for both und and gen tokens. - For frame 1+, forward_with_kv_cache() is used instead (optimized path). - - Args: - pack: Packed sequence with und/gen tokens - attention_mask: Attention mask (BlockMask or SplitInfo) - packed_position_embeddings: RoPE embeddings (cos, sin) - dual_kv_cache: Optional dual KV cache for AR generation (frame 0). - """ - - q_und_in = self.to_q(get_und_seq(pack)) # [N_und,num_heads*head_dim] - q_gen_in = self.add_q_proj(get_gen_seq(pack)) # [N_gen,num_heads*head_dim] - - k_und_in = self.to_k(get_und_seq(pack)) # [N_und,num_kv_heads*head_dim] - k_gen_in = self.add_k_proj(get_gen_seq(pack)) # [N_gen,num_kv_heads*head_dim] - - v_und_in = self.to_v(get_und_seq(pack)) # [N_und,num_kv_heads*head_dim] - v_gen_in = self.add_v_proj(get_gen_seq(pack)) # [N_gen,num_kv_heads*head_dim] - - q_und = q_und_in.view(-1, self.num_attention_heads, self.head_dim) # [N_und,num_heads,head_dim] - k_und = k_und_in.view(-1, self.num_key_value_heads, self.head_dim) # [N_und,num_kv_heads,head_dim] - v_und = v_und_in.view(-1, self.num_key_value_heads, self.head_dim) # [N_und,num_kv_heads,head_dim] - - q_gen = q_gen_in.view(-1, self.num_attention_heads, self.head_dim) # [N_gen,num_heads,head_dim] - k_gen = k_gen_in.view(-1, self.num_key_value_heads, self.head_dim) # [N_gen,num_kv_heads,head_dim] - v_gen = v_gen_in.view(-1, self.num_key_value_heads, self.head_dim) # [N_gen,num_kv_heads,head_dim] - - q_und = self.norm_q(q_und) # [N_und,num_heads,head_dim] - k_und = self.norm_k(k_und) # [N_und,num_kv_heads,head_dim] - - q_gen = self.norm_added_q(q_gen) # [N_gen,num_heads,head_dim] - k_gen = self.norm_added_k(k_gen) # [N_gen,num_kv_heads,head_dim] - - if self.config.freeze_und: - q_und = q_und.detach() - k_und = k_und.detach() - v_und = v_und.detach() - - # Attempted port: Apply RoPE (BAGEL qwen-2.5) - # Note: Position embeddings are now pre-squeezed at model level - packed_cos = packed_position_embeddings[0] - packed_sin = packed_position_embeddings[1] - - q_und_, k_und_ = apply_rotary_pos_emb( - q_und, - k_und, - get_und_seq(packed_cos), - get_und_seq(packed_sin), - unsqueeze_dim=1, - ) # q_und_: [N_und,num_heads,head_dim], k_und_: [N_und,num_kv_heads,head_dim] - q_gen_, k_gen_ = apply_rotary_pos_emb( - q_gen, - k_gen, - get_gen_seq(packed_cos), - get_gen_seq(packed_sin), - unsqueeze_dim=1, - ) # q_gen_: [N_gen,num_heads,head_dim], k_gen_: [N_gen,num_kv_heads,head_dim] - - # === KV CACHE INTEGRATION FOR AUTOREGRESSIVE GENERATION === - # Frame 0: Store und and gen K/V (no fetching) - # Apply cache after RoPE (cached keys already have positional info) - # CP path: storage happens inside context_parallel_attention() after all-to-all, - # so tensors are stored head-sharded [1,S,H/cp,D]. - # Non-CP path: store here as [1,S,H,D] for fetch_kv() dim=1 compat. - if dual_kv_cache is not None and self.cp_mesh is None: - und_len = pack["_num_causal_tokens"] - gen_len = pack["_num_full_tokens"] - if not dual_kv_cache.und_cache.is_initialized: - dual_kv_cache.und_cache.store( - k_und_[:und_len].unsqueeze(0), v_und[:und_len].unsqueeze(0) - ) # [1,S_und,H,D] - dual_kv_cache.gen_cache.store_kv( - k_gen_[:gen_len].unsqueeze(0), v_gen[:gen_len].unsqueeze(0), frame_idx=0 - ) # [1,S_gen,H,D] - - packed_query_states_ = from_und_gen_splits(q_und_, q_gen_, pack) # [N_und+N_gen,num_heads,head_dim] - packed_key_states_ = from_und_gen_splits(k_und_, k_gen_, pack) # [N_und+N_gen,num_kv_heads,head_dim] - packed_value_states_ = from_und_gen_splits(v_und, v_gen, pack) # [N_und+N_gen,num_kv_heads,head_dim] - - # CP: pass dual_kv_cache so context_parallel_attention() stores head-sharded K/V - dispatch_kwargs: dict = {} - if self.cp_mesh is not None and dual_kv_cache is not None: - dispatch_kwargs["dual_kv_cache"] = dual_kv_cache - dispatch_kwargs["frame_idx"] = 0 - - packed_attn_output = self.dispatch_attention_fn( - packed_query_states_, - packed_key_states_, - packed_value_states_, - ) - - # Apply projections directly to get final results - und_seq = self.to_out(get_und_seq(packed_attn_output)) # [N_und,hidden_size] - gen_seq = self.to_add_out(get_gen_seq(packed_attn_output)) # [N_gen,hidden_size] - return from_und_gen_splits(und_seq, gen_seq, pack) # [N_und+N_gen,hidden_size] - - -class Cosmos3VLTextMoTDecoderLayer(nn.Module): - """ - Qwen3VL text MoT (Mixture of Tokens) decoder layer. - Features dual-pathway attention for understanding vs generation. - - This is used for both Dense and MoE models. - """ - - def __init__( - self, - config, - layer_idx: int, - layer_types: LayerTypes, - ): - super().__init__() - self.hidden_size = config.hidden_size - self.freeze_und = config.freeze_und - self.self_attn = PackedAttentionMoT(config, layer_idx, layer_types) - - # TODO: Qwen3VLMoeTextSparseMoeBlock not supported yet - self.mlp = layer_types.mlp(config) - self.mlp_moe_gen = layer_types.mlp(config) - - self.input_layernorm = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) - self.input_layernorm_moe_gen = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm_moe_gen = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) - - def forward( - self, - input: FactoredSequencePack, - attention_mask, - packed_position_embeddings: Tuple[FactoredSequencePack, FactoredSequencePack], - dual_kv_cache: None = None, - frame_idx: Optional[int] = None, - natten_metadata: dict | None = None, - ) -> FactoredSequencePack: - """Training forward pass with MoT routing - Attempted port from qwen2_mot - - Args: - input: Packed sequence with und/gen tokens - attention_mask: Attention mask - packed_position_embeddings: RoPE embeddings (cos, sin) - dual_kv_cache: Optional dual KV cache for AR generation - frame_idx: Current frame index (default: None, treated as 0) - """ - - # Handle None frame_idx as 0 - if frame_idx is None: - frame_idx = 0 - - # TODO: support gen_only = True and AR generation - gen_only = False - # if dual_kv_cache is not None and isinstance(dual_kv_cache, DualKVCache): - # gen_only = frame_idx > 0 and dual_kv_cache.und_cache.is_initialized - - # Pre-Attention layernorm - pack_norm_out = from_und_gen_splits( - self.input_layernorm(get_und_seq(input)), # [N_und,hidden_size] - self.input_layernorm_moe_gen(get_gen_seq(input)), # [N_gen,hidden_size] - input, - ) # [N_und+N_gen,hidden_size] - - # STANDARD PATH: Process both und and gen tokens (frame 0) - pack_attn_out = self.self_attn( - pack_norm_out, - attention_mask, - packed_position_embeddings, - dual_kv_cache, - natten_metadata=natten_metadata, - ) - residual_und = get_und_seq(input) + get_und_seq(pack_attn_out) # [N_und,hidden_size] - residual_gen = get_gen_seq(input) + get_gen_seq(pack_attn_out) # [N_gen,hidden_size] - - # STANDARD PATH: Process both und and gen tokens - ln_out_und = self.post_attention_layernorm(residual_und) # [N_und,hidden_size] - ln_out_gen = self.post_attention_layernorm_moe_gen(residual_gen) # [N_gen,hidden_size] - - # UNPAD MLP INPUT =============== - # NOTE: This is only need for the MoE auxiliary loss computation and to avoid - # artificial expert inbalance due to routing padding tokens. - gen_len = pack_attn_out["_num_full_tokens"] - und_len = pack_attn_out["_num_causal_tokens"] - ln_out_und_unpadded = ln_out_und[:und_len] # [N_und_unpadded,hidden_size] - ln_out_gen_unpadded = ln_out_gen[:gen_len] # [N_gen_unpadded,hidden_size] - - mlp_out_und_unpadded = self.mlp(ln_out_und_unpadded) # [N_und_unpadded,hidden_size] - mlp_out_gen_unpadded = self.mlp_moe_gen(ln_out_gen_unpadded) # [N_gen_unpadded,hidden_size] - - # PAD MLP OUTPUT =============== - mlp_out_und = torch.cat([mlp_out_und_unpadded, ln_out_und[und_len:]], dim=0) # [N_und,hidden_size] - mlp_out_gen = torch.cat([mlp_out_gen_unpadded, ln_out_gen[gen_len:]], dim=0) # [N_gen,hidden_size] - - mlp_out_und_seq = residual_und + mlp_out_und # [N_und,hidden_size] - mlp_out_gen_seq = residual_gen + mlp_out_gen # [N_gen,hidden_size] - - return from_und_gen_splits(mlp_out_und_seq, mlp_out_gen_seq, input) - - -class Cosmos3OmniTransformer(ModelMixin, ConfigMixin): - @register_to_config - def __init__( - self, - attention_bias: bool = False, - attention_dropout: float = 0.0, - dtype: str = "bfloat16", - freeze_und: bool = False, - head_dim: int = 128, - hidden_act: str = "silu", - hidden_size: int = 4096, - initializer_range: float = 0.02, - intermediate_size: int = 12288, - base_fps: int = 24, - enable_fps_modulation: bool = True, - joint_attn_implementation: str = "two_way", - latent_channel: int = 48, - action_dim: int | None = None, - action_gen: bool = False, - max_action_dim: int = 32, - num_embodiment_domains: int = 32, - position_embedding_type: str = "unified_3d_mrope", - unified_3d_mrope_reset_spatial_ids: bool = True, - unified_3d_mrope_temporal_modality_margin: int = 15000, - video_temporal_causal: bool = False, - latent_patch_size: int = 2, - max_position_embeddings: int = 262144, - model_type: str = "qwen3_vl_text", - num_attention_heads: int = 32, - num_hidden_layers: int = 36, - num_key_value_heads: int = 8, - patch_latent_dim: int = 192, - qk_norm: bool = False, - qk_norm_for_diffusion: bool = True, - qk_norm_for_text: bool = True, - rms_norm_eps: float = 1e-6, - rope_scaling: dict | None = None, - rope_theta: float = 5000000.0, - sound_dim: int | None = None, - sound_gen: bool = False, - sound_latent_fps: float = 25.0, - temporal_compression_factor_sound: int = 1, - timestep_scale: float = 0.001, - use_cache: bool = True, - use_moe: bool = True, - vocab_size: int = 151936, - ): - super().__init__() - - if rope_scaling is None: - rope_scaling = {"mrope_interleaved": True, "mrope_section": [24, 20, 20], "rope_type": "default"} - self.register_to_config(rope_scaling=rope_scaling) - - layer_types = LayerTypes(is_moe=False) - self.embed_tokens = nn.Embedding(self.config.vocab_size, self.config.hidden_size) - self.layers = nn.ModuleList( - [ - Cosmos3VLTextMoTDecoderLayer(self.config, layer_idx, layer_types) - for layer_idx in range(self.config.num_hidden_layers) - ] - ) - # Understanding pathway final norm - self.norm = layer_types.rms_norm(self.config.hidden_size, eps=self.config.rms_norm_eps) - # Generation pathway final norm - self.norm_moe_gen = layer_types.rms_norm(self.config.hidden_size, eps=self.config.rms_norm_eps) - self.rotary_emb = Cosmos3VLTextRotaryEmbedding(config=self.config) - self.vocab_size = vocab_size - self.action_gen = action_gen - self.action_dim = int(max_action_dim if action_dim is None else action_dim) - self.num_embodiment_domains = int(num_embodiment_domains) - self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False) - self.proj_in = nn.Linear(patch_latent_dim, hidden_size, bias=True) - self.proj_out = nn.Linear(hidden_size, patch_latent_dim, bias=True) - self.time_embedder = TimestepEmbedder(hidden_size) - if action_gen: - self.action_proj_in = DomainAwareLinear(self.action_dim, hidden_size, self.num_embodiment_domains) - self.action_proj_out = DomainAwareLinear(hidden_size, self.action_dim, self.num_embodiment_domains) - self.action_modality_embed = nn.Parameter(torch.zeros(hidden_size)) - if sound_gen: - if sound_dim is None: - raise ValueError("`sound_dim` must be provided when `sound_gen=True`.") - self.audio_proj_in = nn.Linear(sound_dim, hidden_size, bias=True) - self.audio_proj_out = nn.Linear(hidden_size, sound_dim, bias=True) - self.audio_modality_embed = nn.Parameter(torch.zeros(hidden_size)) - - @classmethod - def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): - model = super().from_pretrained(pretrained_model_name_or_path, **kwargs) - # inv_freq is a non-persistent buffer absent from the saved state_dict. - # Initialize it on CPU; it will move to the correct device with .to() / .cuda(). - model.rotary_emb.init_weights(buffer_device=None) - return model - - def forward( - self, - pack: FactoredSequencePack, - attention_mask, - position_ids: torch.Tensor, - dual_kv_cache: None = None, - frame_idx: Optional[int] = None, - natten_metadata_list: list | None = None, - ) -> Tuple[FactoredSequencePack, None]: - """Training forward pass - simplified to match qwen3_mot. - - Returns: - (outputs, None) — the None placeholder mirrors the (packed_outputs, lbl_metadata) - tuple returned by the original language_model so callers can unpack both. - """ - # Handle None frame_idx as 0 - if frame_idx is None: - frame_idx = 0 - - # Create position embeddings (Qwen3 style) - squeeze once at model level - # tensor below is only used for its dtype and device - device, dtype = get_device_and_dtype(pack) - _meta_tensor = torch.tensor([], dtype=dtype, device=device) # [0] - cos, sin = self.rotary_emb( - _meta_tensor, - position_ids=position_ids.unsqueeze(0) if position_ids.ndim == 1 else position_ids.unsqueeze(1), - ) # if ndim == 2, then the mrope position_ids is (3, seq_len), we need to put batch dimension in the middle to make it compatible with the rotary_emb - # cos, sin: [1,N,head_dim] (1D pos_ids) or [3,1,N,head_dim] (mrope pos_ids) - cos = cos.squeeze(0) # [N,head_dim] or [3,N,head_dim] - sin = sin.squeeze(0) # [N,head_dim] or [3,N,head_dim] - position_embeddings = ( - from_joint(cos, pack), - from_joint(sin, pack), - ) - - # TODO: Add lbl_metadata_all (we don't need it at inference) - hidden_states = pack - - for i, decoder_layer in enumerate(self.layers): - hidden_states = decoder_layer( - hidden_states, - attention_mask, - position_embeddings, - dual_kv_cache[i] if dual_kv_cache is not None else None, - frame_idx, - natten_metadata=None if natten_metadata_list is None else natten_metadata_list[i], - ) - - outputs = zeros_like(hidden_states) # [N_und+N_gen,hidden_size] - set_und_seq(outputs, self.norm(get_und_seq(hidden_states))) # [N_und,hidden_size] - set_gen_seq(outputs, self.norm_moe_gen(get_gen_seq(hidden_states))) # [N_gen,hidden_size] - return outputs, None - - -@use_kernel_forward_from_hub("RMSNorm") -class Qwen3VLMoeTextRMSNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-6): - """ - Qwen3VLMoeTextRMSNorm is equivalent to T5LayerNorm - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states): - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) - - def extra_repr(self): - return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" - - -class Qwen3VLMoeTextMLP(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj - - -class Qwen3VLMoeTextRotaryEmbedding(nn.Module): - def __init__(self, config): - super().__init__() - if hasattr(config, "rope_scaling") and config.rope_scaling is not None: - self.rope_type = config.rope_scaling.get("rope_type", "default") - else: - self.rope_type = "default" - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings - self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20]) - - self.config = config - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - - def init_weights(self, buffer_device: torch.device | None = None) -> None: - inv_freq, self.attention_scaling = self.rope_init_fn(self.config, buffer_device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - - def apply_interleaved_mrope(self, freqs, mrope_section): - """Apply interleaved MRoPE to 3D rotary embeddings. - Reorganizes frequency layout from chunked [TTT...HHH...WWW] to - interleaved [THTHWHTHW...TT], preserving frequency continuity. - args: - x: (3, bs, seq_len, head_dim // 2) - mrope_section: (3,) - returns: - x_t: (bs, seq_len, head_dim // 2) - """ - freqs_t = freqs[0] # just overwrite the first dimension T - for dim, offset in enumerate((1, 2), start=1): # H, W - length = mrope_section[dim] * 3 - idx = slice(offset, length, 3) - freqs_t[..., idx] = freqs[dim, ..., idx] - return freqs_t - - @torch.no_grad() - @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) - def forward(self, x, position_ids): - assert self.inv_freq.dtype == torch.float32, f"inv_freq must be float32, but got {self.inv_freq.dtype}" - - # In contrast to other models, Qwen3VLMoe has different position ids for the grids - # So we expand the inv_freq to shape (3, ...) - if position_ids.ndim == 2: - position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) # [3,B,N] - inv_freq_expanded = ( - self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) - ) # [3,B,head_dim//2,1] - position_ids_expanded = position_ids[:, :, None, :].float() # [3,B,1,N] - - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) # [3,B,N,head_dim//2] - freqs = self.apply_interleaved_mrope(freqs, self.mrope_section) # [B,N,head_dim//2] - emb = torch.cat((freqs, freqs), dim=-1) # [B,N,head_dim] - cos = emb.cos() * self.attention_scaling # [B,N,head_dim] - sin = emb.sin() * self.attention_scaling # [B,N,head_dim] - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) diff --git a/packages/diffusers-cosmos3/pyproject.toml b/packages/diffusers-cosmos3/pyproject.toml deleted file mode 100644 index edf6500d..00000000 --- a/packages/diffusers-cosmos3/pyproject.toml +++ /dev/null @@ -1,13 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[project] -name = "diffusers-cosmos3" -version = "0.1.0" -description = "diffusers plugin that loads Cosmos3 checkpoints" -requires-python = ">=3.10" -dependencies = ["diffusers>=0.37.0"] diff --git a/packages/diffusers-cosmos3/scripts/inference_example.py b/packages/diffusers-cosmos3/scripts/inference_example.py deleted file mode 100755 index ec145319..00000000 --- a/packages/diffusers-cosmos3/scripts/inference_example.py +++ /dev/null @@ -1,79 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: OpenMDW-1.1 -""" -Load a Cosmos3 diffusers pipeline from a converted checkpoint and run inference. - -CUDA_VISIBLE_DEVICES=0 python inference_cosmos3.py \ - --pipeline-path converted/cosmos3-nano-pipeline \ - --input inputs/omni/i2v.json - -The pipeline must have been produced by convert_cosmos3_to_diffusers.py -with --save-pipeline. -""" - -import argparse -import json -import pathlib - -import torch -from diffusers_cosmos3 import Cosmos3OmniDiffusersPipeline -from diffusers_cosmos3.pipeline import save_img_or_video - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--pipeline-path", - default="converted/cosmos3-nano-pipeline", - help="Path to directory saved by cosmos_convert_cosmos_to_diffusers.py --save-pipeline.", - ) - parser.add_argument( - "--input", - default="inputs/omni/i2v.json", - help="Path to JSON input file with 'prompt' and optional 'vision_path'.", - ) - parser.add_argument("--output", default=".", help="Directory to save generated video/image files.") - parser.add_argument("--height", type=int, default=720) - parser.add_argument("--width", type=int, default=1280) - parser.add_argument("--num-frames", type=int, default=189) - args = parser.parse_args() - - pipeline_path = pathlib.Path(args.pipeline_path) - print(f"Loading pipeline from {pipeline_path} …") - pipeline = Cosmos3OmniDiffusersPipeline.from_pretrained( - str(pipeline_path), - torch_dtype=torch.bfloat16, - device_map="cuda", - ) - print("Pipeline loaded successfully.") - - # --- Load JSON input --- - input_path = pathlib.Path(args.input) - print(f"Loading input from {input_path} …") - with open(input_path) as f: - input_data = json.load(f) - prompt = input_data["prompt"] - vision_path = input_data.get("vision_path", None) - - output_dir = pathlib.Path(args.output) - output_dir.mkdir(parents=True, exist_ok=True) - - result_vision = pipeline( - prompt=prompt, - image=vision_path, - num_frames=args.num_frames, - height=args.height, - width=args.width, - output_type="latent", - ) - - result_frames = pipeline.decode_latents(result_vision) - for i, frames in enumerate(result_frames): - save_path = str(output_dir / f"sample-{i}") - save_img_or_video(frames, save_path) - print(f"Saved: {save_path}.mp4") - - -if __name__ == "__main__": - main() diff --git a/packages/diffusers-cosmos3/uv.lock b/packages/diffusers-cosmos3/uv.lock deleted file mode 100644 index ef5f4413..00000000 --- a/packages/diffusers-cosmos3/uv.lock +++ /dev/null @@ -1,924 +0,0 @@ -version = 1 -revision = 3 -requires-python = ">=3.10" -resolution-markers = [ - "python_full_version >= '3.11'", - "python_full_version < '3.11'", -] - -[[package]] -name = "annotated-doc" -version = "0.0.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/57/ba/046ceea27344560984e26a590f90bc7f4a75b06701f653222458922b558c/annotated_doc-0.0.4.tar.gz", hash = "sha256:fbcda96e87e9c92ad167c2e53839e57503ecfda18804ea28102353485033faa4", size = 7288, upload-time = "2025-11-10T22:07:42.062Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/d3/26bf1008eb3d2daa8ef4cacc7f3bfdc11818d111f7e2d0201bc6e3b49d45/annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320", size = 5303, upload-time = "2025-11-10T22:07:40.673Z" }, -] - -[[package]] -name = "anyio" -version = "4.13.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, - { name = "idna" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/19/14/2c5dd9f512b66549ae92767a9c7b330ae88e1932ca57876909410251fe13/anyio-4.13.0.tar.gz", hash = "sha256:334b70e641fd2221c1505b3890c69882fe4a2df910cba14d97019b90b24439dc", size = 231622, upload-time = "2026-03-24T12:59:09.671Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, -] - -[[package]] -name = "certifi" -version = "2026.5.20" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f3/ce/ee2ecad540810a79593028e88299baeae54d346cc7a0d94b6199988b89b1/certifi-2026.5.20.tar.gz", hash = "sha256:69dea482ab64caa7b9f6aba1c6bf48bb6a5448d1c0f1b17ab42ad8c763a5344d", size = 135422, upload-time = "2026-05-20T11:46:50.073Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/59/8c/57e832b7af6d7c5abe66eb3fbe3a3a32f4d11ea23a1aa7131371035be991/certifi-2026.5.20-py3-none-any.whl", hash = "sha256:3c52e209ba0a4ad7aebe60436a4ab349c39e1e602e8c134221e546902ad25897", size = 134134, upload-time = "2026-05-20T11:46:48.578Z" }, -] - -[[package]] -name = "charset-normalizer" -version = "3.4.7" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e7/a1/67fe25fac3c7642725500a3f6cfe5821ad557c3abb11c9d20d12c7008d3e/charset_normalizer-3.4.7.tar.gz", hash = "sha256:ae89db9e5f98a11a4bf50407d4363e7b09b31e55bc117b4f7d80aab97ba009e5", size = 144271, upload-time = "2026-04-02T09:28:39.342Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/26/08/0f303cb0b529e456bb116f2d50565a482694fbb94340bf56d44677e7ed03/charset_normalizer-3.4.7-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:cdd68a1fb318e290a2077696b7eb7a21a49163c455979c639bf5a5dcdc46617d", size = 315182, upload-time = "2026-04-02T09:25:40.673Z" }, - { url = "https://files.pythonhosted.org/packages/24/47/b192933e94b546f1b1fe4df9cc1f84fcdbf2359f8d1081d46dd029b50207/charset_normalizer-3.4.7-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e17b8d5d6a8c47c85e68ca8379def1303fd360c3e22093a807cd34a71cd082b8", size = 209329, upload-time = "2026-04-02T09:25:42.354Z" }, - { url = "https://files.pythonhosted.org/packages/c2/b4/01fa81c5ca6141024d89a8fc15968002b71da7f825dd14113207113fabbd/charset_normalizer-3.4.7-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:511ef87c8aec0783e08ac18565a16d435372bc1ac25a91e6ac7f5ef2b0bff790", size = 231230, upload-time = "2026-04-02T09:25:44.281Z" }, - { url = "https://files.pythonhosted.org/packages/20/f7/7b991776844dfa058017e600e6e55ff01984a063290ca5622c0b63162f68/charset_normalizer-3.4.7-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:007d05ec7321d12a40227aae9e2bc6dca73f3cb21058999a1df9e193555a9dcc", size = 225890, upload-time = "2026-04-02T09:25:45.475Z" }, - { url = "https://files.pythonhosted.org/packages/20/e7/bed0024a0f4ab0c8a9c64d4445f39b30c99bd1acd228291959e3de664247/charset_normalizer-3.4.7-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cf29836da5119f3c8a8a70667b0ef5fdca3bb12f80fd06487cfa575b3909b393", size = 216930, upload-time = "2026-04-02T09:25:46.58Z" }, - { url = "https://files.pythonhosted.org/packages/e2/ab/b18f0ab31cdd7b3ddb8bb76c4a414aeb8160c9810fdf1bc62f269a539d87/charset_normalizer-3.4.7-cp310-cp310-manylinux_2_31_armv7l.whl", hash = "sha256:12d8baf840cc7889b37c7c770f478adea7adce3dcb3944d02ec87508e2dcf153", size = 202109, upload-time = "2026-04-02T09:25:48.031Z" }, - { url = "https://files.pythonhosted.org/packages/82/e5/7e9440768a06dfb3075936490cb82dbf0ee20a133bf0dd8551fa096914ec/charset_normalizer-3.4.7-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d560742f3c0d62afaccf9f41fe485ed69bd7661a241f86a3ef0f0fb8b1a397af", size = 214684, upload-time = "2026-04-02T09:25:49.245Z" }, - { url = "https://files.pythonhosted.org/packages/71/94/8c61d8da9f062fdf457c80acfa25060ec22bf1d34bbeaca4350f13bcfd07/charset_normalizer-3.4.7-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:b14b2d9dac08e28bb8046a1a0434b1750eb221c8f5b87a68f4fa11a6f97b5e34", size = 212785, upload-time = "2026-04-02T09:25:50.671Z" }, - { url = "https://files.pythonhosted.org/packages/66/cd/6e9889c648e72c0ab2e5967528bb83508f354d706637bc7097190c874e13/charset_normalizer-3.4.7-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:bc17a677b21b3502a21f66a8cc64f5bfad4df8a0b8434d661666f8ce90ac3af1", size = 203055, upload-time = "2026-04-02T09:25:51.802Z" }, - { url = "https://files.pythonhosted.org/packages/92/2e/7a951d6a08aefb7eb8e1b54cdfb580b1365afdd9dd484dc4bee9e5d8f258/charset_normalizer-3.4.7-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:750e02e074872a3fad7f233b47734166440af3cdea0add3e95163110816d6752", size = 232502, upload-time = "2026-04-02T09:25:53.388Z" }, - { url = "https://files.pythonhosted.org/packages/58/d5/abcf2d83bf8e0a1286df55cd0dc1d49af0da4282aa77e986df343e7de124/charset_normalizer-3.4.7-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:4e5163c14bffd570ef2affbfdd77bba66383890797df43dc8b4cc7d6f500bf53", size = 214295, upload-time = "2026-04-02T09:25:54.765Z" }, - { url = "https://files.pythonhosted.org/packages/47/3a/7d4cd7ed54be99973a0dc176032cba5cb1f258082c31fa6df35cff46acfc/charset_normalizer-3.4.7-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:6ed74185b2db44f41ef35fd1617c5888e59792da9bbc9190d6c7300617182616", size = 227145, upload-time = "2026-04-02T09:25:55.904Z" }, - { url = "https://files.pythonhosted.org/packages/1d/98/3a45bf8247889cf28262ebd3d0872edff11565b2a1e3064ccb132db3fbb0/charset_normalizer-3.4.7-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:94e1885b270625a9a828c9793b4d52a64445299baa1fea5a173bf1d3dd9a1a5a", size = 218884, upload-time = "2026-04-02T09:25:57.074Z" }, - { url = "https://files.pythonhosted.org/packages/ad/80/2e8b7f8915ed5c9ef13aa828d82738e33888c485b65ebf744d615040c7ea/charset_normalizer-3.4.7-cp310-cp310-win32.whl", hash = "sha256:6785f414ae0f3c733c437e0f3929197934f526d19dfaa75e18fdb4f94c6fb374", size = 148343, upload-time = "2026-04-02T09:25:58.199Z" }, - { url = "https://files.pythonhosted.org/packages/35/1b/3b8c8c77184af465ee9ad88b5aea46ea6b2e1f7b9dc9502891e37af21e30/charset_normalizer-3.4.7-cp310-cp310-win_amd64.whl", hash = "sha256:6696b7688f54f5af4462118f0bfa7c1621eeb87154f77fa04b9295ce7a8f2943", size = 159174, upload-time = "2026-04-02T09:25:59.322Z" }, - { url = "https://files.pythonhosted.org/packages/be/c1/feb40dca40dbb21e0a908801782d9288c64fc8d8e562c2098e9994c8c21b/charset_normalizer-3.4.7-cp310-cp310-win_arm64.whl", hash = "sha256:66671f93accb62ed07da56613636f3641f1a12c13046ce91ffc923721f23c008", size = 147805, upload-time = "2026-04-02T09:26:00.756Z" }, - { url = "https://files.pythonhosted.org/packages/c2/d7/b5b7020a0565c2e9fa8c09f4b5fa6232feb326b8c20081ccded47ea368fd/charset_normalizer-3.4.7-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:7641bb8895e77f921102f72833904dcd9901df5d6d72a2ab8f31d04b7e51e4e7", size = 309705, upload-time = "2026-04-02T09:26:02.191Z" }, - { url = "https://files.pythonhosted.org/packages/5a/53/58c29116c340e5456724ecd2fff4196d236b98f3da97b404bc5e51ac3493/charset_normalizer-3.4.7-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:202389074300232baeb53ae2569a60901f7efadd4245cf3a3bf0617d60b439d7", size = 206419, upload-time = "2026-04-02T09:26:03.583Z" }, - { url = "https://files.pythonhosted.org/packages/b2/02/e8146dc6591a37a00e5144c63f29fb7c97a734ea8a111190783c0e60ab63/charset_normalizer-3.4.7-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:30b8d1d8c52a48c2c5690e152c169b673487a2a58de1ec7393196753063fcd5e", size = 227901, upload-time = "2026-04-02T09:26:04.738Z" }, - { url = "https://files.pythonhosted.org/packages/fb/73/77486c4cd58f1267bf17db420e930c9afa1b3be3fe8c8b8ebbebc9624359/charset_normalizer-3.4.7-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:532bc9bf33a68613fd7d65e4b1c71a6a38d7d42604ecf239c77392e9b4e8998c", size = 222742, upload-time = "2026-04-02T09:26:06.36Z" }, - { url = "https://files.pythonhosted.org/packages/a1/fa/f74eb381a7d94ded44739e9d94de18dc5edc9c17fb8c11f0a6890696c0a9/charset_normalizer-3.4.7-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2fe249cb4651fd12605b7288b24751d8bfd46d35f12a20b1ba33dea122e690df", size = 214061, upload-time = "2026-04-02T09:26:08.347Z" }, - { url = "https://files.pythonhosted.org/packages/dc/92/42bd3cefcf7687253fb86694b45f37b733c97f59af3724f356fa92b8c344/charset_normalizer-3.4.7-cp311-cp311-manylinux_2_31_armv7l.whl", hash = "sha256:65bcd23054beab4d166035cabbc868a09c1a49d1efe458fe8e4361215df40265", size = 199239, upload-time = "2026-04-02T09:26:09.823Z" }, - { url = "https://files.pythonhosted.org/packages/4c/3d/069e7184e2aa3b3cddc700e3dd267413dc259854adc3380421c805c6a17d/charset_normalizer-3.4.7-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:08e721811161356f97b4059a9ba7bafb23ea5ee2255402c42881c214e173c6b4", size = 210173, upload-time = "2026-04-02T09:26:10.953Z" }, - { url = "https://files.pythonhosted.org/packages/62/51/9d56feb5f2e7074c46f93e0ebdbe61f0848ee246e2f0d89f8e20b89ebb8f/charset_normalizer-3.4.7-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:e060d01aec0a910bdccb8be71faf34e7799ce36950f8294c8bf612cba65a2c9e", size = 209841, upload-time = "2026-04-02T09:26:12.142Z" }, - { url = "https://files.pythonhosted.org/packages/d2/59/893d8f99cc4c837dda1fe2f1139079703deb9f321aabcb032355de13b6c7/charset_normalizer-3.4.7-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:38c0109396c4cfc574d502df99742a45c72c08eff0a36158b6f04000043dbf38", size = 200304, upload-time = "2026-04-02T09:26:13.711Z" }, - { url = "https://files.pythonhosted.org/packages/7d/1d/ee6f3be3464247578d1ed5c46de545ccc3d3ff933695395c402c21fa6b77/charset_normalizer-3.4.7-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:1c2a768fdd44ee4a9339a9b0b130049139b8ce3c01d2ce09f67f5a68048d477c", size = 229455, upload-time = "2026-04-02T09:26:14.941Z" }, - { url = "https://files.pythonhosted.org/packages/54/bb/8fb0a946296ea96a488928bdce8ef99023998c48e4713af533e9bb98ef07/charset_normalizer-3.4.7-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:1a87ca9d5df6fe460483d9a5bbf2b18f620cbed41b432e2bddb686228282d10b", size = 210036, upload-time = "2026-04-02T09:26:16.478Z" }, - { url = "https://files.pythonhosted.org/packages/9a/bc/015b2387f913749f82afd4fcba07846d05b6d784dd16123cb66860e0237d/charset_normalizer-3.4.7-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:d635aab80466bc95771bb78d5370e74d36d1fe31467b6b29b8b57b2a3cd7d22c", size = 224739, upload-time = "2026-04-02T09:26:17.751Z" }, - { url = "https://files.pythonhosted.org/packages/17/ab/63133691f56baae417493cba6b7c641571a2130eb7bceba6773367ab9ec5/charset_normalizer-3.4.7-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:ae196f021b5e7c78e918242d217db021ed2a6ace2bc6ae94c0fc596221c7f58d", size = 216277, upload-time = "2026-04-02T09:26:18.981Z" }, - { url = "https://files.pythonhosted.org/packages/06/6d/3be70e827977f20db77c12a97e6a9f973631a45b8d186c084527e53e77a4/charset_normalizer-3.4.7-cp311-cp311-win32.whl", hash = "sha256:adb2597b428735679446b46c8badf467b4ca5f5056aae4d51a19f9570301b1ad", size = 147819, upload-time = "2026-04-02T09:26:20.295Z" }, - { url = "https://files.pythonhosted.org/packages/20/d9/5f67790f06b735d7c7637171bbfd89882ad67201891b7275e51116ed8207/charset_normalizer-3.4.7-cp311-cp311-win_amd64.whl", hash = "sha256:8e385e4267ab76874ae30db04c627faaaf0b509e1ccc11a95b3fc3e83f855c00", size = 159281, upload-time = "2026-04-02T09:26:21.74Z" }, - { url = "https://files.pythonhosted.org/packages/ca/83/6413f36c5a34afead88ce6f66684d943d91f233d76dd083798f9602b75ae/charset_normalizer-3.4.7-cp311-cp311-win_arm64.whl", hash = "sha256:d4a48e5b3c2a489fae013b7589308a40146ee081f6f509e047e0e096084ceca1", size = 147843, upload-time = "2026-04-02T09:26:22.901Z" }, - { url = "https://files.pythonhosted.org/packages/0c/eb/4fc8d0a7110eb5fc9cc161723a34a8a6c200ce3b4fbf681bc86feee22308/charset_normalizer-3.4.7-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:eca9705049ad3c7345d574e3510665cb2cf844c2f2dcfe675332677f081cbd46", size = 311328, upload-time = "2026-04-02T09:26:24.331Z" }, - { url = "https://files.pythonhosted.org/packages/f8/e3/0fadc706008ac9d7b9b5be6dc767c05f9d3e5df51744ce4cc9605de7b9f4/charset_normalizer-3.4.7-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6178f72c5508bfc5fd446a5905e698c6212932f25bcdd4b47a757a50605a90e2", size = 208061, upload-time = "2026-04-02T09:26:25.568Z" }, - { url = "https://files.pythonhosted.org/packages/42/f0/3dd1045c47f4a4604df85ec18ad093912ae1344ac706993aff91d38773a2/charset_normalizer-3.4.7-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e1421b502d83040e6d7fb2fb18dff63957f720da3d77b2fbd3187ceb63755d7b", size = 229031, upload-time = "2026-04-02T09:26:26.865Z" }, - { url = "https://files.pythonhosted.org/packages/dc/67/675a46eb016118a2fbde5a277a5d15f4f69d5f3f5f338e5ee2f8948fcf43/charset_normalizer-3.4.7-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:edac0f1ab77644605be2cbba52e6b7f630731fc42b34cb0f634be1a6eface56a", size = 225239, upload-time = "2026-04-02T09:26:28.044Z" }, - { url = "https://files.pythonhosted.org/packages/4b/f8/d0118a2f5f23b02cd166fa385c60f9b0d4f9194f574e2b31cef350ad7223/charset_normalizer-3.4.7-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5649fd1c7bade02f320a462fdefd0b4bd3ce036065836d4f42e0de958038e116", size = 216589, upload-time = "2026-04-02T09:26:29.239Z" }, - { url = "https://files.pythonhosted.org/packages/b1/f1/6d2b0b261b6c4ceef0fcb0d17a01cc5bc53586c2d4796fa04b5c540bc13d/charset_normalizer-3.4.7-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:203104ed3e428044fd943bc4bf45fa73c0730391f9621e37fe39ecf477b128cb", size = 202733, upload-time = "2026-04-02T09:26:30.5Z" }, - { url = "https://files.pythonhosted.org/packages/6f/c0/7b1f943f7e87cc3db9626ba17807d042c38645f0a1d4415c7a14afb5591f/charset_normalizer-3.4.7-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:298930cec56029e05497a76988377cbd7457ba864beeea92ad7e844fe74cd1f1", size = 212652, upload-time = "2026-04-02T09:26:31.709Z" }, - { url = "https://files.pythonhosted.org/packages/38/dd/5a9ab159fe45c6e72079398f277b7d2b523e7f716acc489726115a910097/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:708838739abf24b2ceb208d0e22403dd018faeef86ddac04319a62ae884c4f15", size = 211229, upload-time = "2026-04-02T09:26:33.282Z" }, - { url = "https://files.pythonhosted.org/packages/d5/ff/531a1cad5ca855d1c1a8b69cb71abfd6d85c0291580146fda7c82857caa1/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:0f7eb884681e3938906ed0434f20c63046eacd0111c4ba96f27b76084cd679f5", size = 203552, upload-time = "2026-04-02T09:26:34.845Z" }, - { url = "https://files.pythonhosted.org/packages/c1/4c/a5fb52d528a8ca41f7598cb619409ece30a169fbdf9cdce592e53b46c3a6/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:4dc1e73c36828f982bfe79fadf5919923f8a6f4df2860804db9a98c48824ce8d", size = 230806, upload-time = "2026-04-02T09:26:36.152Z" }, - { url = "https://files.pythonhosted.org/packages/59/7a/071feed8124111a32b316b33ae4de83d36923039ef8cf48120266844285b/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:aed52fea0513bac0ccde438c188c8a471c4e0f457c2dd20cdbf6ea7a450046c7", size = 212316, upload-time = "2026-04-02T09:26:37.672Z" }, - { url = "https://files.pythonhosted.org/packages/fd/35/f7dba3994312d7ba508e041eaac39a36b120f32d4c8662b8814dab876431/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:fea24543955a6a729c45a73fe90e08c743f0b3334bbf3201e6c4bc1b0c7fa464", size = 227274, upload-time = "2026-04-02T09:26:38.93Z" }, - { url = "https://files.pythonhosted.org/packages/8a/2d/a572df5c9204ab7688ec1edc895a73ebded3b023bb07364710b05dd1c9be/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bb6d88045545b26da47aa879dd4a89a71d1dce0f0e549b1abcb31dfe4a8eac49", size = 218468, upload-time = "2026-04-02T09:26:40.17Z" }, - { url = "https://files.pythonhosted.org/packages/86/eb/890922a8b03a568ca2f336c36585a4713c55d4d67bf0f0c78924be6315ca/charset_normalizer-3.4.7-cp312-cp312-win32.whl", hash = "sha256:2257141f39fe65a3fdf38aeccae4b953e5f3b3324f4ff0daf9f15b8518666a2c", size = 148460, upload-time = "2026-04-02T09:26:41.416Z" }, - { url = "https://files.pythonhosted.org/packages/35/d9/0e7dffa06c5ab081f75b1b786f0aefc88365825dfcd0ac544bdb7b2b6853/charset_normalizer-3.4.7-cp312-cp312-win_amd64.whl", hash = "sha256:5ed6ab538499c8644b8a3e18debabcd7ce684f3fa91cf867521a7a0279cab2d6", size = 159330, upload-time = "2026-04-02T09:26:42.554Z" }, - { url = "https://files.pythonhosted.org/packages/9e/5d/481bcc2a7c88ea6b0878c299547843b2521ccbc40980cb406267088bc701/charset_normalizer-3.4.7-cp312-cp312-win_arm64.whl", hash = "sha256:56be790f86bfb2c98fb742ce566dfb4816e5a83384616ab59c49e0604d49c51d", size = 147828, upload-time = "2026-04-02T09:26:44.075Z" }, - { url = "https://files.pythonhosted.org/packages/c1/3b/66777e39d3ae1ddc77ee606be4ec6d8cbd4c801f65e5a1b6f2b11b8346dd/charset_normalizer-3.4.7-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:f496c9c3cc02230093d8330875c4c3cdfc3b73612a5fd921c65d39cbcef08063", size = 309627, upload-time = "2026-04-02T09:26:45.198Z" }, - { url = "https://files.pythonhosted.org/packages/2e/4e/b7f84e617b4854ade48a1b7915c8ccfadeba444d2a18c291f696e37f0d3b/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ea948db76d31190bf08bd371623927ee1339d5f2a0b4b1b4a4439a65298703c", size = 207008, upload-time = "2026-04-02T09:26:46.824Z" }, - { url = "https://files.pythonhosted.org/packages/c4/bb/ec73c0257c9e11b268f018f068f5d00aa0ef8c8b09f7753ebd5f2880e248/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a277ab8928b9f299723bc1a2dabb1265911b1a76341f90a510368ca44ad9ab66", size = 228303, upload-time = "2026-04-02T09:26:48.397Z" }, - { url = "https://files.pythonhosted.org/packages/85/fb/32d1f5033484494619f701e719429c69b766bfc4dbc61aa9e9c8c166528b/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3bec022aec2c514d9cf199522a802bd007cd588ab17ab2525f20f9c34d067c18", size = 224282, upload-time = "2026-04-02T09:26:49.684Z" }, - { url = "https://files.pythonhosted.org/packages/fa/07/330e3a0dda4c404d6da83b327270906e9654a24f6c546dc886a0eb0ffb23/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e044c39e41b92c845bc815e5ae4230804e8e7bc29e399b0437d64222d92809dd", size = 215595, upload-time = "2026-04-02T09:26:50.915Z" }, - { url = "https://files.pythonhosted.org/packages/e3/7c/fc890655786e423f02556e0216d4b8c6bcb6bdfa890160dc66bf52dee468/charset_normalizer-3.4.7-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:f495a1652cf3fbab2eb0639776dad966c2fb874d79d87ca07f9d5f059b8bd215", size = 201986, upload-time = "2026-04-02T09:26:52.197Z" }, - { url = "https://files.pythonhosted.org/packages/d8/97/bfb18b3db2aed3b90cf54dc292ad79fdd5ad65c4eae454099475cbeadd0d/charset_normalizer-3.4.7-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e712b419df8ba5e42b226c510472b37bd57b38e897d3eca5e8cfd410a29fa859", size = 211711, upload-time = "2026-04-02T09:26:53.49Z" }, - { url = "https://files.pythonhosted.org/packages/6f/a5/a581c13798546a7fd557c82614a5c65a13df2157e9ad6373166d2a3e645d/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7804338df6fcc08105c7745f1502ba68d900f45fd770d5bdd5288ddccb8a42d8", size = 210036, upload-time = "2026-04-02T09:26:54.975Z" }, - { url = "https://files.pythonhosted.org/packages/8c/bf/b3ab5bcb478e4193d517644b0fb2bf5497fbceeaa7a1bc0f4d5b50953861/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:481551899c856c704d58119b5025793fa6730adda3571971af568f66d2424bb5", size = 202998, upload-time = "2026-04-02T09:26:56.303Z" }, - { url = "https://files.pythonhosted.org/packages/e7/4e/23efd79b65d314fa320ec6017b4b5834d5c12a58ba4610aa353af2e2f577/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:f59099f9b66f0d7145115e6f80dd8b1d847176df89b234a5a6b3f00437aa0832", size = 230056, upload-time = "2026-04-02T09:26:57.554Z" }, - { url = "https://files.pythonhosted.org/packages/b9/9f/1e1941bc3f0e01df116e68dc37a55c4d249df5e6fa77f008841aef68264f/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:f59ad4c0e8f6bba240a9bb85504faa1ab438237199d4cce5f622761507b8f6a6", size = 211537, upload-time = "2026-04-02T09:26:58.843Z" }, - { url = "https://files.pythonhosted.org/packages/80/0f/088cbb3020d44428964a6c97fe1edfb1b9550396bf6d278330281e8b709c/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:3dedcc22d73ec993f42055eff4fcfed9318d1eeb9a6606c55892a26964964e48", size = 226176, upload-time = "2026-04-02T09:27:00.437Z" }, - { url = "https://files.pythonhosted.org/packages/6a/9f/130394f9bbe06f4f63e22641d32fc9b202b7e251c9aef4db044324dac493/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:64f02c6841d7d83f832cd97ccf8eb8a906d06eb95d5276069175c696b024b60a", size = 217723, upload-time = "2026-04-02T09:27:02.021Z" }, - { url = "https://files.pythonhosted.org/packages/73/55/c469897448a06e49f8fa03f6caae97074fde823f432a98f979cc42b90e69/charset_normalizer-3.4.7-cp313-cp313-win32.whl", hash = "sha256:4042d5c8f957e15221d423ba781e85d553722fc4113f523f2feb7b188cc34c5e", size = 148085, upload-time = "2026-04-02T09:27:03.192Z" }, - { url = "https://files.pythonhosted.org/packages/5d/78/1b74c5bbb3f99b77a1715c91b3e0b5bdb6fe302d95ace4f5b1bec37b0167/charset_normalizer-3.4.7-cp313-cp313-win_amd64.whl", hash = "sha256:3946fa46a0cf3e4c8cb1cc52f56bb536310d34f25f01ca9b6c16afa767dab110", size = 158819, upload-time = "2026-04-02T09:27:04.454Z" }, - { url = "https://files.pythonhosted.org/packages/68/86/46bd42279d323deb8687c4a5a811fd548cb7d1de10cf6535d099877a9a9f/charset_normalizer-3.4.7-cp313-cp313-win_arm64.whl", hash = "sha256:80d04837f55fc81da168b98de4f4b797ef007fc8a79ab71c6ec9bc4dd662b15b", size = 147915, upload-time = "2026-04-02T09:27:05.971Z" }, - { url = "https://files.pythonhosted.org/packages/97/c8/c67cb8c70e19ef1960b97b22ed2a1567711de46c4ddf19799923adc836c2/charset_normalizer-3.4.7-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:c36c333c39be2dbca264d7803333c896ab8fa7d4d6f0ab7edb7dfd7aea6e98c0", size = 309234, upload-time = "2026-04-02T09:27:07.194Z" }, - { url = "https://files.pythonhosted.org/packages/99/85/c091fdee33f20de70d6c8b522743b6f831a2f1cd3ff86de4c6a827c48a76/charset_normalizer-3.4.7-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1c2aed2e5e41f24ea8ef1590b8e848a79b56f3a5564a65ceec43c9d692dc7d8a", size = 208042, upload-time = "2026-04-02T09:27:08.749Z" }, - { url = "https://files.pythonhosted.org/packages/87/1c/ab2ce611b984d2fd5d86a5a8a19c1ae26acac6bad967da4967562c75114d/charset_normalizer-3.4.7-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:54523e136b8948060c0fa0bc7b1b50c32c186f2fceee897a495406bb6e311d2b", size = 228706, upload-time = "2026-04-02T09:27:09.951Z" }, - { url = "https://files.pythonhosted.org/packages/a8/29/2b1d2cb00bf085f59d29eb773ce58ec2d325430f8c216804a0a5cd83cbca/charset_normalizer-3.4.7-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:715479b9a2802ecac752a3b0efa2b0b60285cf962ee38414211abdfccc233b41", size = 224727, upload-time = "2026-04-02T09:27:11.175Z" }, - { url = "https://files.pythonhosted.org/packages/47/5c/032c2d5a07fe4d4855fea851209cca2b6f03ebeb6d4e3afdb3358386a684/charset_normalizer-3.4.7-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bd6c2a1c7573c64738d716488d2cdd3c00e340e4835707d8fdb8dc1a66ef164e", size = 215882, upload-time = "2026-04-02T09:27:12.446Z" }, - { url = "https://files.pythonhosted.org/packages/2c/c2/356065d5a8b78ed04499cae5f339f091946a6a74f91e03476c33f0ab7100/charset_normalizer-3.4.7-cp314-cp314-manylinux_2_31_armv7l.whl", hash = "sha256:c45e9440fb78f8ddabcf714b68f936737a121355bf59f3907f4e17721b9d1aae", size = 200860, upload-time = "2026-04-02T09:27:13.721Z" }, - { url = "https://files.pythonhosted.org/packages/0c/cd/a32a84217ced5039f53b29f460962abb2d4420def55afabe45b1c3c7483d/charset_normalizer-3.4.7-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:3534e7dcbdcf757da6b85a0bbf5b6868786d5982dd959b065e65481644817a18", size = 211564, upload-time = "2026-04-02T09:27:15.272Z" }, - { url = "https://files.pythonhosted.org/packages/44/86/58e6f13ce26cc3b8f4a36b94a0f22ae2f00a72534520f4ae6857c4b81f89/charset_normalizer-3.4.7-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:e8ac484bf18ce6975760921bb6148041faa8fef0547200386ea0b52b5d27bf7b", size = 211276, upload-time = "2026-04-02T09:27:16.834Z" }, - { url = "https://files.pythonhosted.org/packages/8f/fe/d17c32dc72e17e155e06883efa84514ca375f8a528ba2546bee73fc4df81/charset_normalizer-3.4.7-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:a5fe03b42827c13cdccd08e6c0247b6a6d4b5e3cdc53fd1749f5896adcdc2356", size = 201238, upload-time = "2026-04-02T09:27:18.229Z" }, - { url = "https://files.pythonhosted.org/packages/6a/29/f33daa50b06525a237451cdb6c69da366c381a3dadcd833fa5676bc468b3/charset_normalizer-3.4.7-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:2d6eb928e13016cea4f1f21d1e10c1cebd5a421bc57ddf5b1142ae3f86824fab", size = 230189, upload-time = "2026-04-02T09:27:19.445Z" }, - { url = "https://files.pythonhosted.org/packages/b6/6e/52c84015394a6a0bdcd435210a7e944c5f94ea1055f5cc5d56c5fe368e7b/charset_normalizer-3.4.7-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:e74327fb75de8986940def6e8dee4f127cc9752bee7355bb323cc5b2659b6d46", size = 211352, upload-time = "2026-04-02T09:27:20.79Z" }, - { url = "https://files.pythonhosted.org/packages/8c/d7/4353be581b373033fb9198bf1da3cf8f09c1082561e8e922aa7b39bf9fe8/charset_normalizer-3.4.7-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:d6038d37043bced98a66e68d3aa2b6a35505dc01328cd65217cefe82f25def44", size = 227024, upload-time = "2026-04-02T09:27:22.063Z" }, - { url = "https://files.pythonhosted.org/packages/30/45/99d18aa925bd1740098ccd3060e238e21115fffbfdcb8f3ece837d0ace6c/charset_normalizer-3.4.7-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:7579e913a5339fb8fa133f6bbcfd8e6749696206cf05acdbdca71a1b436d8e72", size = 217869, upload-time = "2026-04-02T09:27:23.486Z" }, - { url = "https://files.pythonhosted.org/packages/5c/05/5ee478aa53f4bb7996482153d4bfe1b89e0f087f0ab6b294fcf92d595873/charset_normalizer-3.4.7-cp314-cp314-win32.whl", hash = "sha256:5b77459df20e08151cd6f8b9ef8ef1f961ef73d85c21a555c7eed5b79410ec10", size = 148541, upload-time = "2026-04-02T09:27:25.146Z" }, - { url = "https://files.pythonhosted.org/packages/48/77/72dcb0921b2ce86420b2d79d454c7022bf5be40202a2a07906b9f2a35c97/charset_normalizer-3.4.7-cp314-cp314-win_amd64.whl", hash = "sha256:92a0a01ead5e668468e952e4238cccd7c537364eb7d851ab144ab6627dbbe12f", size = 159634, upload-time = "2026-04-02T09:27:26.642Z" }, - { url = "https://files.pythonhosted.org/packages/c6/a3/c2369911cd72f02386e4e340770f6e158c7980267da16af8f668217abaa0/charset_normalizer-3.4.7-cp314-cp314-win_arm64.whl", hash = "sha256:67f6279d125ca0046a7fd386d01b311c6363844deac3e5b069b514ba3e63c246", size = 148384, upload-time = "2026-04-02T09:27:28.271Z" }, - { url = "https://files.pythonhosted.org/packages/94/09/7e8a7f73d24dba1f0035fbbf014d2c36828fc1bf9c88f84093e57d315935/charset_normalizer-3.4.7-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:effc3f449787117233702311a1b7d8f59cba9ced946ba727bdc329ec69028e24", size = 330133, upload-time = "2026-04-02T09:27:29.474Z" }, - { url = "https://files.pythonhosted.org/packages/8d/da/96975ddb11f8e977f706f45cddd8540fd8242f71ecdb5d18a80723dcf62c/charset_normalizer-3.4.7-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fbccdc05410c9ee21bbf16a35f4c1d16123dcdeb8a1d38f33654fa21d0234f79", size = 216257, upload-time = "2026-04-02T09:27:30.793Z" }, - { url = "https://files.pythonhosted.org/packages/e5/e8/1d63bf8ef2d388e95c64b2098f45f84758f6d102a087552da1485912637b/charset_normalizer-3.4.7-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:733784b6d6def852c814bce5f318d25da2ee65dd4839a0718641c696e09a2960", size = 234851, upload-time = "2026-04-02T09:27:32.44Z" }, - { url = "https://files.pythonhosted.org/packages/9b/40/e5ff04233e70da2681fa43969ad6f66ca5611d7e669be0246c4c7aaf6dc8/charset_normalizer-3.4.7-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a89c23ef8d2c6b27fd200a42aa4ac72786e7c60d40efdc76e6011260b6e949c4", size = 233393, upload-time = "2026-04-02T09:27:34.03Z" }, - { url = "https://files.pythonhosted.org/packages/be/c1/06c6c49d5a5450f76899992f1ee40b41d076aee9279b49cf9974d2f313d5/charset_normalizer-3.4.7-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6c114670c45346afedc0d947faf3c7f701051d2518b943679c8ff88befe14f8e", size = 223251, upload-time = "2026-04-02T09:27:35.369Z" }, - { url = "https://files.pythonhosted.org/packages/2b/9f/f2ff16fb050946169e3e1f82134d107e5d4ae72647ec8a1b1446c148480f/charset_normalizer-3.4.7-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:a180c5e59792af262bf263b21a3c49353f25945d8d9f70628e73de370d55e1e1", size = 206609, upload-time = "2026-04-02T09:27:36.661Z" }, - { url = "https://files.pythonhosted.org/packages/69/d5/a527c0cd8d64d2eab7459784fb4169a0ac76e5a6fc5237337982fd61347e/charset_normalizer-3.4.7-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:3c9a494bc5ec77d43cea229c4f6db1e4d8fe7e1bbffa8b6f0f0032430ff8ab44", size = 220014, upload-time = "2026-04-02T09:27:38.019Z" }, - { url = "https://files.pythonhosted.org/packages/7e/80/8a7b8104a3e203074dc9aa2c613d4b726c0e136bad1cc734594b02867972/charset_normalizer-3.4.7-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8d828b6667a32a728a1ad1d93957cdf37489c57b97ae6c4de2860fa749b8fc1e", size = 218979, upload-time = "2026-04-02T09:27:39.37Z" }, - { url = "https://files.pythonhosted.org/packages/02/9a/b759b503d507f375b2b5c153e4d2ee0a75aa215b7f2489cf314f4541f2c0/charset_normalizer-3.4.7-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:cf1493cd8607bec4d8a7b9b004e699fcf8f9103a9284cc94962cb73d20f9d4a3", size = 209238, upload-time = "2026-04-02T09:27:40.722Z" }, - { url = "https://files.pythonhosted.org/packages/c2/4e/0f3f5d47b86bdb79256e7290b26ac847a2832d9a4033f7eb2cd4bcf4bb5b/charset_normalizer-3.4.7-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:0c96c3b819b5c3e9e165495db84d41914d6894d55181d2d108cc1a69bfc9cce0", size = 236110, upload-time = "2026-04-02T09:27:42.33Z" }, - { url = "https://files.pythonhosted.org/packages/96/23/bce28734eb3ed2c91dcf93abeb8a5cf393a7b2749725030bb630e554fdd8/charset_normalizer-3.4.7-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:752a45dc4a6934060b3b0dab47e04edc3326575f82be64bc4fc293914566503e", size = 219824, upload-time = "2026-04-02T09:27:43.924Z" }, - { url = "https://files.pythonhosted.org/packages/2c/6f/6e897c6984cc4d41af319b077f2f600fc8214eb2fe2d6bcb79141b882400/charset_normalizer-3.4.7-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:8778f0c7a52e56f75d12dae53ae320fae900a8b9b4164b981b9c5ce059cd1fcb", size = 233103, upload-time = "2026-04-02T09:27:45.348Z" }, - { url = "https://files.pythonhosted.org/packages/76/22/ef7bd0fe480a0ae9b656189ec00744b60933f68b4f42a7bb06589f6f576a/charset_normalizer-3.4.7-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ce3412fbe1e31eb81ea42f4169ed94861c56e643189e1e75f0041f3fe7020abe", size = 225194, upload-time = "2026-04-02T09:27:46.706Z" }, - { url = "https://files.pythonhosted.org/packages/c5/a7/0e0ab3e0b5bc1219bd80a6a0d4d72ca74d9250cb2382b7c699c147e06017/charset_normalizer-3.4.7-cp314-cp314t-win32.whl", hash = "sha256:c03a41a8784091e67a39648f70c5f97b5b6a37f216896d44d2cdcb82615339a0", size = 159827, upload-time = "2026-04-02T09:27:48.053Z" }, - { url = "https://files.pythonhosted.org/packages/7a/1d/29d32e0fb40864b1f878c7f5a0b343ae676c6e2b271a2d55cc3a152391da/charset_normalizer-3.4.7-cp314-cp314t-win_amd64.whl", hash = "sha256:03853ed82eeebbce3c2abfdbc98c96dc205f32a79627688ac9a27370ea61a49c", size = 174168, upload-time = "2026-04-02T09:27:49.795Z" }, - { url = "https://files.pythonhosted.org/packages/de/32/d92444ad05c7a6e41fb2036749777c163baf7a0301a040cb672d6b2b1ae9/charset_normalizer-3.4.7-cp314-cp314t-win_arm64.whl", hash = "sha256:c35abb8bfff0185efac5878da64c45dafd2b37fb0383add1be155a763c1f083d", size = 153018, upload-time = "2026-04-02T09:27:51.116Z" }, - { url = "https://files.pythonhosted.org/packages/db/8f/61959034484a4a7c527811f4721e75d02d653a35afb0b6054474d8185d4c/charset_normalizer-3.4.7-py3-none-any.whl", hash = "sha256:3dce51d0f5e7951f8bb4900c257dad282f49190fdbebecd4ba99bcc41fef404d", size = 61958, upload-time = "2026-04-02T09:28:37.794Z" }, -] - -[[package]] -name = "click" -version = "8.4.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/9b/98/518d8e5081007684232226f475082b30087d0f585e8457db087298259f49/click-8.4.1.tar.gz", hash = "sha256:918b5633eddf6b41c32d4f454bf0de810065c74e3f7dbf8ee5452f8be88d3e96", size = 353007, upload-time = "2026-05-22T04:08:37.769Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c7/0d/67e5b4109ea4a837e80daa87c2c696711955e40449a97e8926672534def2/click-8.4.1-py3-none-any.whl", hash = "sha256:482be17c6991b8c19c5429a1e995d9b0efdbb63172824c41f99965dc0ade8ec2", size = 116639, upload-time = "2026-05-22T04:08:35.26Z" }, -] - -[[package]] -name = "colorama" -version = "0.4.6" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, -] - -[[package]] -name = "diffusers" -version = "0.37.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "filelock" }, - { name = "httpx" }, - { name = "huggingface-hub" }, - { name = "importlib-metadata" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "pillow" }, - { name = "regex" }, - { name = "requests" }, - { name = "safetensors" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/46/5c/f4c2eb8d481fe8784a7e2331fbaab820079c06676185fa6d2177b386d590/diffusers-0.37.1.tar.gz", hash = "sha256:2346c21f77f835f273b7aacbaada1c34a596a3a2cc6ddc99d149efcd0ec298fa", size = 4135139, upload-time = "2026-03-25T08:04:04.515Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9c/dd/51c38785ce5e1c287b5ad17ba550edaaaffce0deb0da4857019c6700fbaf/diffusers-0.37.1-py3-none-any.whl", hash = "sha256:0537c0b28cb53cf39d6195489bcf8f833986df556c10f5e28ab7427b86fc8b90", size = 5001536, upload-time = "2026-03-25T08:04:02.385Z" }, -] - -[[package]] -name = "diffusers-cosmos3" -version = "0.1.0" -source = { editable = "." } -dependencies = [ - { name = "diffusers" }, -] - -[package.metadata] -requires-dist = [{ name = "diffusers", specifier = ">=0.37.0" }] - -[[package]] -name = "exceptiongroup" -version = "1.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598", size = 16740, upload-time = "2025-11-21T23:01:53.443Z" }, -] - -[[package]] -name = "filelock" -version = "3.29.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b5/fe/997687a931ab51049acce6fa1f23e8f01216374ea81374ddee763c493db5/filelock-3.29.0.tar.gz", hash = "sha256:69974355e960702e789734cb4871f884ea6fe50bd8404051a3530bc07809cf90", size = 57571, upload-time = "2026-04-19T15:39:10.068Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/81/47/dd9a212ef6e343a6857485ffe25bba537304f1913bdbed446a23f7f592e1/filelock-3.29.0-py3-none-any.whl", hash = "sha256:96f5f6344709aa1572bbf631c640e4ebeeb519e08da902c39a001882f30ac258", size = 39812, upload-time = "2026-04-19T15:39:08.752Z" }, -] - -[[package]] -name = "fsspec" -version = "2026.4.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d5/8d/1c51c094345df128ca4a990d633fe1a0ff28726c9e6b3c41ba65087bba1d/fsspec-2026.4.0.tar.gz", hash = "sha256:301d8ac70ae90ef3ad05dcf94d6c3754a097f9b5fe4667d2787aa359ec7df7e4", size = 312760, upload-time = "2026-04-29T20:42:38.635Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d5/0c/043d5e551459da400957a1395e0febbf771446ff34291afcbe3d8be2a279/fsspec-2026.4.0-py3-none-any.whl", hash = "sha256:11ef7bb35dab8a394fde6e608221d5cf3e8499401c249bebaeaad760a1a8dec2", size = 203402, upload-time = "2026-04-29T20:42:36.842Z" }, -] - -[[package]] -name = "h11" -version = "0.16.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, -] - -[[package]] -name = "hf-xet" -version = "1.5.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/74/d8/5c06fc76461418326a7decf8367480c35be11a41fd938633929c60a9ec6b/hf_xet-1.5.0.tar.gz", hash = "sha256:e0fb0a34d9f406eed88233e829a67ec016bec5af19e480eac65a233ea289a948", size = 837196, upload-time = "2026-05-06T06:18:15.583Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/68/9b/6912c99070915a4f28119e3c5b52a9abd1eec0ad5cb293b8c967a0c6f5a2/hf_xet-1.5.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:7d70fe2ce97b9db73b9c9b9c81fe3693640aec83416a966c446afea54acfae3c", size = 4023383, upload-time = "2026-05-06T06:17:53.947Z" }, - { url = "https://files.pythonhosted.org/packages/0f/6d/9563cfde59b5d8128a9c7ec972a087f4c782e4f7bac5a85234edfd5d5e49/hf_xet-1.5.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:73a0dae8c71de3b0633a45c73f4a4a5ed09e94b43441d82981a781d4f12baa42", size = 3792751, upload-time = "2026-05-06T06:17:51.791Z" }, - { url = "https://files.pythonhosted.org/packages/07/a5/ed5a0cf35b49a0571af5a8f53416dad1877a718c021c9937c3a53cb45781/hf_xet-1.5.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a60290ec57e9b71767fba7c3645ddafdd0759974b540441510c629c6db6db24a", size = 4456058, upload-time = "2026-05-06T06:17:40.735Z" }, - { url = "https://files.pythonhosted.org/packages/60/fb/3ae8bf2a7a37a4197d0195d7247fd25b3952e15cb8a599e285dfaa6f52b3/hf_xet-1.5.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e5de0f6deada0dada870bb376a11bcd1f08abf3a968a6d118f33e72d1b1eb480", size = 4250783, upload-time = "2026-05-06T06:17:38.412Z" }, - { url = "https://files.pythonhosted.org/packages/a2/9b/8bae40d4d91525085137196e84eb0ed49cf65b5e96e5c3ecdadd8bd0fac2/hf_xet-1.5.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c799d49f1a5544a0ef7591c0ee75e0d6b93d6f56dc7a4979f59f7518d2872216", size = 4445594, upload-time = "2026-05-06T06:18:04.219Z" }, - { url = "https://files.pythonhosted.org/packages/13/59/c74efbbd4e8728172b2cc72a2bc014d2947a4b7bdced932fbd3f5da1a4e5/hf_xet-1.5.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:2baea1b0b989e5c152fe81425f7745ddc8901280ba3d97c98d8cdece7b706c60", size = 4663995, upload-time = "2026-05-06T06:18:06.1Z" }, - { url = "https://files.pythonhosted.org/packages/73/32/8e1e0410af64cda9b139d1dcebdc993a8ff9c8c7c0e2696ae356d75ccc0d/hf_xet-1.5.0-cp313-cp313t-win_amd64.whl", hash = "sha256:526345b3ed45f374f6317349df489167606736c876241ba984105afe7fd4839d", size = 3966608, upload-time = "2026-05-06T06:18:19.74Z" }, - { url = "https://files.pythonhosted.org/packages/fc/34/a8febc8f4edbea8b3e21b02ebc8b628679b84ba7e45cde624a7736b51500/hf_xet-1.5.0-cp313-cp313t-win_arm64.whl", hash = "sha256:786d28e2eb8315d5035544b9d137b4a842d600c434bb91bf7d0d953cce906ad4", size = 3796946, upload-time = "2026-05-06T06:18:17.568Z" }, - { url = "https://files.pythonhosted.org/packages/2a/20/8fc8996afe5815fa1a6be8e9e5c02f24500f409d599e905800d498a4e14d/hf_xet-1.5.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:872d5601e6deea30d15865ede55d29eac6daf5a534ab417b99b6ef6b076dd96c", size = 4023495, upload-time = "2026-05-06T06:18:01.94Z" }, - { url = "https://files.pythonhosted.org/packages/32/6a/93d84463c00cecb561a7508aa6303e35ee2894294eac14245526924415fe/hf_xet-1.5.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:9929561f5abf4581c8ea79587881dfef6b8abb2a0d8a51915936fc2a614f4e73", size = 3792731, upload-time = "2026-05-06T06:18:00.021Z" }, - { url = "https://files.pythonhosted.org/packages/9d/5a/8ec8e0c863b382d00b3c2e2af6ded6b06371be617144a625903a6d562f4b/hf_xet-1.5.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f7b7bbae318e583a86fb21e5a4a175d6721d628a2874f4bd022d0e660c32a682", size = 4456738, upload-time = "2026-05-06T06:17:49.574Z" }, - { url = "https://files.pythonhosted.org/packages/c5/ca/f7effa1a67717da2bcc6b6c28f71c6ca648c77acaec4e2c32f40cbe16d85/hf_xet-1.5.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:cf7b2dc6f31a4ea754bb50f74cde482dcf5d366d184076d8530b9872787f3761", size = 4251622, upload-time = "2026-05-06T06:17:47.096Z" }, - { url = "https://files.pythonhosted.org/packages/65/f2/19247dba3e231cf77dec59ddfb878f00057635ff773d099c9b59d37812c3/hf_xet-1.5.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8dbcbab554c9ef158ef2c991545c3e970ddd8cc7acdcd0a78c5a41095dab4ded", size = 4445667, upload-time = "2026-05-06T06:18:11.983Z" }, - { url = "https://files.pythonhosted.org/packages/7f/64/6f116801a3bcfb6f59f5c251f48cadc47ea54026441c4a385079286a94fa/hf_xet-1.5.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5906bf7718d3636dc13402914736abe723492cb730f744834f5f5b67d3a12702", size = 4664619, upload-time = "2026-05-06T06:18:13.771Z" }, - { url = "https://files.pythonhosted.org/packages/5c/e8/069542d37946ed08669b127e1496fa99e78196d71de8d41eda5e9f1b7a58/hf_xet-1.5.0-cp314-cp314t-win_amd64.whl", hash = "sha256:5f3dc2248fc01cc0a00cd392ab497f1ca373fcbc7e3f2da1f452480b384e839e", size = 3966802, upload-time = "2026-05-06T06:18:28.162Z" }, - { url = "https://files.pythonhosted.org/packages/f9/91/fc6fdec27b14d04e88c386ac0a0129732b53fa23f7c4a78f4b83a039c567/hf_xet-1.5.0-cp314-cp314t-win_arm64.whl", hash = "sha256:b285cea1b5bab46b758772716ba8d6854a1a0310fed1c249d678a8b38601e5a0", size = 3797168, upload-time = "2026-05-06T06:18:26.287Z" }, - { url = "https://files.pythonhosted.org/packages/3d/fb/69ff198a82cae7eb1a69fb84d93b3a3e4816564d76817fe541ddc96874eb/hf_xet-1.5.0-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:dad0dc84e941b8ba3c860659fe1fdc35c049d47cce293f003287757e971a8f56", size = 4030814, upload-time = "2026-05-06T06:17:57.933Z" }, - { url = "https://files.pythonhosted.org/packages/9b/ff/edcc2b40162bef3ff78e14ab637e5f3b89243d6aee72f5949d3bb6a5af83/hf_xet-1.5.0-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:fd6e5a9b0fdac4ed03ed45ef79254a655b1aaab514a02202617fbf643f5fdf7a", size = 3798444, upload-time = "2026-05-06T06:17:55.79Z" }, - { url = "https://files.pythonhosted.org/packages/49/4d/103f76b04310e5e57656696cc184690d20c466af0bca3ca88f8c8ea5d4f3/hf_xet-1.5.0-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3531b1823a0e6d77d80f9ed15ca0e00f0d115094f8ac033d5cae88f4564cc949", size = 4465986, upload-time = "2026-05-06T06:17:44.886Z" }, - { url = "https://files.pythonhosted.org/packages/c4/a2/546f47f464737b3edbab6f8ddb57f2599b93d2cbb66f06abb475ccb48651/hf_xet-1.5.0-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:9a0ee58cd18d5ea799f7ed11290bbccbe56bdd8b1d97ca74b9cc49a3945d7a3b", size = 4259865, upload-time = "2026-05-06T06:17:42.639Z" }, - { url = "https://files.pythonhosted.org/packages/95/7f/1be593c1f28613be2e196473481cd81bfc5910795e30a34e8f744f6cac4f/hf_xet-1.5.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1e60df5a42e9bed8628b6416af2cba4cba57ae9f02de226a06b020d98e1aab18", size = 4459835, upload-time = "2026-05-06T06:18:08.026Z" }, - { url = "https://files.pythonhosted.org/packages/aa/b2/703569fc881f3284487e68cda7b42179978480da3c438042a6bbbb4a671c/hf_xet-1.5.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:4b35549ce62601b84da4ff9b24d970032ace3d4430f52d91bcbb26c901d6c690", size = 4672414, upload-time = "2026-05-06T06:18:09.864Z" }, - { url = "https://files.pythonhosted.org/packages/af/37/1b6def445c567286b50aa3b33828158e135b1be44938dde59f11382a500c/hf_xet-1.5.0-cp37-abi3-win_amd64.whl", hash = "sha256:2806c7c17b4d23f8d88f7c4814f838c3b6150773fe339c20af23e1cfaf2797e4", size = 3977238, upload-time = "2026-05-06T06:18:23.621Z" }, - { url = "https://files.pythonhosted.org/packages/62/94/3b66b148778ee100dcfd69c2ca22b57b41b44d3063ceec934f209e9184ce/hf_xet-1.5.0-cp37-abi3-win_arm64.whl", hash = "sha256:b6c9df403040248c76d808d3e047d64db2d923bae593eb244c41e425cf6cd7be", size = 3806916, upload-time = "2026-05-06T06:18:21.7Z" }, -] - -[[package]] -name = "httpcore" -version = "1.0.9" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "certifi" }, - { name = "h11" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, -] - -[[package]] -name = "httpx" -version = "0.28.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "certifi" }, - { name = "httpcore" }, - { name = "idna" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, -] - -[[package]] -name = "huggingface-hub" -version = "1.17.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "filelock" }, - { name = "fsspec" }, - { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, - { name = "httpx" }, - { name = "packaging" }, - { name = "pyyaml" }, - { name = "tqdm" }, - { name = "typer" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/bd/65/9826515abb600b5722bcf53f8b4a2fb58340b1f8bfcaee19f83561c13a44/huggingface_hub-1.17.0.tar.gz", hash = "sha256:fad842b6763ef70ebc3919665b1b9273645203185400a7d6c5eddc2323cc3435", size = 797082, upload-time = "2026-05-28T15:12:13.347Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/02/28/d7cef5e477b855c25d415b8f57e5bc7347c7a90cad3acf1725d0c92ca294/huggingface_hub-1.17.0-py3-none-any.whl", hash = "sha256:3b8156d23118e87f6a587648bfbc04f04a12a757ccb4ed298b35c4ae638bf24c", size = 671546, upload-time = "2026-05-28T15:12:11.441Z" }, -] - -[[package]] -name = "idna" -version = "3.17" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b9/28/99c51f664567218d824af024c0251650fb27e4ca066df188dab0769c5b91/idna-3.17.tar.gz", hash = "sha256:5eb0cb53bc467c12eadcf6de83163ad8527cec9416f44b9b61b19caedad2b87f", size = 196048, upload-time = "2026-05-28T14:32:38.55Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/de/a7/f76514cc40ad6234098ecdebda08732d75964776c51a42845b7da10649e2/idna-3.17-py3-none-any.whl", hash = "sha256:466e48829084efe2548012b855df21540b96f2e20e51bd124c851536556a592c", size = 65316, upload-time = "2026-05-28T14:32:37.035Z" }, -] - -[[package]] -name = "importlib-metadata" -version = "9.0.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "zipp" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/38/3d/2d244233ac4f76e38533cfcb2991c9eb4c7bf688ae0a036d30725b8faafe/importlib_metadata-9.0.0-py3-none-any.whl", hash = "sha256:2d21d1cc5a017bd0559e36150c21c830ab1dc304dedd1b7ea85d20f45ef3edd7", size = 27789, upload-time = "2026-03-20T06:42:55.665Z" }, -] - -[[package]] -name = "markdown-it-py" -version = "4.2.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "mdurl" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/06/ff/7841249c247aa650a76b9ee4bbaeae59370dc8bfd2f6c01f3630c35eb134/markdown_it_py-4.2.0.tar.gz", hash = "sha256:04a21681d6fbb623de53f6f364d352309d4094dd4194040a10fd51833e418d49", size = 82454, upload-time = "2026-05-07T12:08:28.36Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" }, -] - -[[package]] -name = "mdurl" -version = "0.1.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, -] - -[[package]] -name = "numpy" -version = "2.2.6" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version < '3.11'", -] -sdist = { url = "https://files.pythonhosted.org/packages/76/21/7d2a95e4bba9dc13d043ee156a356c0a8f0c6309dff6b21b4d71a073b8a8/numpy-2.2.6.tar.gz", hash = "sha256:e29554e2bef54a90aa5cc07da6ce955accb83f21ab5de01a62c8478897b264fd", size = 20276440, upload-time = "2025-05-17T22:38:04.611Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9a/3e/ed6db5be21ce87955c0cbd3009f2803f59fa08df21b5df06862e2d8e2bdd/numpy-2.2.6-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b412caa66f72040e6d268491a59f2c43bf03eb6c96dd8f0307829feb7fa2b6fb", size = 21165245, upload-time = "2025-05-17T21:27:58.555Z" }, - { url = "https://files.pythonhosted.org/packages/22/c2/4b9221495b2a132cc9d2eb862e21d42a009f5a60e45fc44b00118c174bff/numpy-2.2.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8e41fd67c52b86603a91c1a505ebaef50b3314de0213461c7a6e99c9a3beff90", size = 14360048, upload-time = "2025-05-17T21:28:21.406Z" }, - { url = "https://files.pythonhosted.org/packages/fd/77/dc2fcfc66943c6410e2bf598062f5959372735ffda175b39906d54f02349/numpy-2.2.6-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:37e990a01ae6ec7fe7fa1c26c55ecb672dd98b19c3d0e1d1f326fa13cb38d163", size = 5340542, upload-time = "2025-05-17T21:28:30.931Z" }, - { url = "https://files.pythonhosted.org/packages/7a/4f/1cb5fdc353a5f5cc7feb692db9b8ec2c3d6405453f982435efc52561df58/numpy-2.2.6-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:5a6429d4be8ca66d889b7cf70f536a397dc45ba6faeb5f8c5427935d9592e9cf", size = 6878301, upload-time = "2025-05-17T21:28:41.613Z" }, - { url = "https://files.pythonhosted.org/packages/eb/17/96a3acd228cec142fcb8723bd3cc39c2a474f7dcf0a5d16731980bcafa95/numpy-2.2.6-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:efd28d4e9cd7d7a8d39074a4d44c63eda73401580c5c76acda2ce969e0a38e83", size = 14297320, upload-time = "2025-05-17T21:29:02.78Z" }, - { url = "https://files.pythonhosted.org/packages/b4/63/3de6a34ad7ad6646ac7d2f55ebc6ad439dbbf9c4370017c50cf403fb19b5/numpy-2.2.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc7b73d02efb0e18c000e9ad8b83480dfcd5dfd11065997ed4c6747470ae8915", size = 16801050, upload-time = "2025-05-17T21:29:27.675Z" }, - { url = "https://files.pythonhosted.org/packages/07/b6/89d837eddef52b3d0cec5c6ba0456c1bf1b9ef6a6672fc2b7873c3ec4e2e/numpy-2.2.6-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:74d4531beb257d2c3f4b261bfb0fc09e0f9ebb8842d82a7b4209415896adc680", size = 15807034, upload-time = "2025-05-17T21:29:51.102Z" }, - { url = "https://files.pythonhosted.org/packages/01/c8/dc6ae86e3c61cfec1f178e5c9f7858584049b6093f843bca541f94120920/numpy-2.2.6-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8fc377d995680230e83241d8a96def29f204b5782f371c532579b4f20607a289", size = 18614185, upload-time = "2025-05-17T21:30:18.703Z" }, - { url = "https://files.pythonhosted.org/packages/5b/c5/0064b1b7e7c89137b471ccec1fd2282fceaae0ab3a9550f2568782d80357/numpy-2.2.6-cp310-cp310-win32.whl", hash = "sha256:b093dd74e50a8cba3e873868d9e93a85b78e0daf2e98c6797566ad8044e8363d", size = 6527149, upload-time = "2025-05-17T21:30:29.788Z" }, - { url = "https://files.pythonhosted.org/packages/a3/dd/4b822569d6b96c39d1215dbae0582fd99954dcbcf0c1a13c61783feaca3f/numpy-2.2.6-cp310-cp310-win_amd64.whl", hash = "sha256:f0fd6321b839904e15c46e0d257fdd101dd7f530fe03fd6359c1ea63738703f3", size = 12904620, upload-time = "2025-05-17T21:30:48.994Z" }, - { url = "https://files.pythonhosted.org/packages/da/a8/4f83e2aa666a9fbf56d6118faaaf5f1974d456b1823fda0a176eff722839/numpy-2.2.6-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f9f1adb22318e121c5c69a09142811a201ef17ab257a1e66ca3025065b7f53ae", size = 21176963, upload-time = "2025-05-17T21:31:19.36Z" }, - { url = "https://files.pythonhosted.org/packages/b3/2b/64e1affc7972decb74c9e29e5649fac940514910960ba25cd9af4488b66c/numpy-2.2.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c820a93b0255bc360f53eca31a0e676fd1101f673dda8da93454a12e23fc5f7a", size = 14406743, upload-time = "2025-05-17T21:31:41.087Z" }, - { url = "https://files.pythonhosted.org/packages/4a/9f/0121e375000b5e50ffdd8b25bf78d8e1a5aa4cca3f185d41265198c7b834/numpy-2.2.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:3d70692235e759f260c3d837193090014aebdf026dfd167834bcba43e30c2a42", size = 5352616, upload-time = "2025-05-17T21:31:50.072Z" }, - { url = "https://files.pythonhosted.org/packages/31/0d/b48c405c91693635fbe2dcd7bc84a33a602add5f63286e024d3b6741411c/numpy-2.2.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:481b49095335f8eed42e39e8041327c05b0f6f4780488f61286ed3c01368d491", size = 6889579, upload-time = "2025-05-17T21:32:01.712Z" }, - { url = "https://files.pythonhosted.org/packages/52/b8/7f0554d49b565d0171eab6e99001846882000883998e7b7d9f0d98b1f934/numpy-2.2.6-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b64d8d4d17135e00c8e346e0a738deb17e754230d7e0810ac5012750bbd85a5a", size = 14312005, upload-time = "2025-05-17T21:32:23.332Z" }, - { url = "https://files.pythonhosted.org/packages/b3/dd/2238b898e51bd6d389b7389ffb20d7f4c10066d80351187ec8e303a5a475/numpy-2.2.6-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba10f8411898fc418a521833e014a77d3ca01c15b0c6cdcce6a0d2897e6dbbdf", size = 16821570, upload-time = "2025-05-17T21:32:47.991Z" }, - { url = "https://files.pythonhosted.org/packages/83/6c/44d0325722cf644f191042bf47eedad61c1e6df2432ed65cbe28509d404e/numpy-2.2.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:bd48227a919f1bafbdda0583705e547892342c26fb127219d60a5c36882609d1", size = 15818548, upload-time = "2025-05-17T21:33:11.728Z" }, - { url = "https://files.pythonhosted.org/packages/ae/9d/81e8216030ce66be25279098789b665d49ff19eef08bfa8cb96d4957f422/numpy-2.2.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9551a499bf125c1d4f9e250377c1ee2eddd02e01eac6644c080162c0c51778ab", size = 18620521, upload-time = "2025-05-17T21:33:39.139Z" }, - { url = "https://files.pythonhosted.org/packages/6a/fd/e19617b9530b031db51b0926eed5345ce8ddc669bb3bc0044b23e275ebe8/numpy-2.2.6-cp311-cp311-win32.whl", hash = "sha256:0678000bb9ac1475cd454c6b8c799206af8107e310843532b04d49649c717a47", size = 6525866, upload-time = "2025-05-17T21:33:50.273Z" }, - { url = "https://files.pythonhosted.org/packages/31/0a/f354fb7176b81747d870f7991dc763e157a934c717b67b58456bc63da3df/numpy-2.2.6-cp311-cp311-win_amd64.whl", hash = "sha256:e8213002e427c69c45a52bbd94163084025f533a55a59d6f9c5b820774ef3303", size = 12907455, upload-time = "2025-05-17T21:34:09.135Z" }, - { url = "https://files.pythonhosted.org/packages/82/5d/c00588b6cf18e1da539b45d3598d3557084990dcc4331960c15ee776ee41/numpy-2.2.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:41c5a21f4a04fa86436124d388f6ed60a9343a6f767fced1a8a71c3fbca038ff", size = 20875348, upload-time = "2025-05-17T21:34:39.648Z" }, - { url = "https://files.pythonhosted.org/packages/66/ee/560deadcdde6c2f90200450d5938f63a34b37e27ebff162810f716f6a230/numpy-2.2.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:de749064336d37e340f640b05f24e9e3dd678c57318c7289d222a8a2f543e90c", size = 14119362, upload-time = "2025-05-17T21:35:01.241Z" }, - { url = "https://files.pythonhosted.org/packages/3c/65/4baa99f1c53b30adf0acd9a5519078871ddde8d2339dc5a7fde80d9d87da/numpy-2.2.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:894b3a42502226a1cac872f840030665f33326fc3dac8e57c607905773cdcde3", size = 5084103, upload-time = "2025-05-17T21:35:10.622Z" }, - { url = "https://files.pythonhosted.org/packages/cc/89/e5a34c071a0570cc40c9a54eb472d113eea6d002e9ae12bb3a8407fb912e/numpy-2.2.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:71594f7c51a18e728451bb50cc60a3ce4e6538822731b2933209a1f3614e9282", size = 6625382, upload-time = "2025-05-17T21:35:21.414Z" }, - { url = "https://files.pythonhosted.org/packages/f8/35/8c80729f1ff76b3921d5c9487c7ac3de9b2a103b1cd05e905b3090513510/numpy-2.2.6-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2618db89be1b4e05f7a1a847a9c1c0abd63e63a1607d892dd54668dd92faf87", size = 14018462, upload-time = "2025-05-17T21:35:42.174Z" }, - { url = "https://files.pythonhosted.org/packages/8c/3d/1e1db36cfd41f895d266b103df00ca5b3cbe965184df824dec5c08c6b803/numpy-2.2.6-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd83c01228a688733f1ded5201c678f0c53ecc1006ffbc404db9f7a899ac6249", size = 16527618, upload-time = "2025-05-17T21:36:06.711Z" }, - { url = "https://files.pythonhosted.org/packages/61/c6/03ed30992602c85aa3cd95b9070a514f8b3c33e31124694438d88809ae36/numpy-2.2.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:37c0ca431f82cd5fa716eca9506aefcabc247fb27ba69c5062a6d3ade8cf8f49", size = 15505511, upload-time = "2025-05-17T21:36:29.965Z" }, - { url = "https://files.pythonhosted.org/packages/b7/25/5761d832a81df431e260719ec45de696414266613c9ee268394dd5ad8236/numpy-2.2.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fe27749d33bb772c80dcd84ae7e8df2adc920ae8297400dabec45f0dedb3f6de", size = 18313783, upload-time = "2025-05-17T21:36:56.883Z" }, - { url = "https://files.pythonhosted.org/packages/57/0a/72d5a3527c5ebffcd47bde9162c39fae1f90138c961e5296491ce778e682/numpy-2.2.6-cp312-cp312-win32.whl", hash = "sha256:4eeaae00d789f66c7a25ac5f34b71a7035bb474e679f410e5e1a94deb24cf2d4", size = 6246506, upload-time = "2025-05-17T21:37:07.368Z" }, - { url = "https://files.pythonhosted.org/packages/36/fa/8c9210162ca1b88529ab76b41ba02d433fd54fecaf6feb70ef9f124683f1/numpy-2.2.6-cp312-cp312-win_amd64.whl", hash = "sha256:c1f9540be57940698ed329904db803cf7a402f3fc200bfe599334c9bd84a40b2", size = 12614190, upload-time = "2025-05-17T21:37:26.213Z" }, - { url = "https://files.pythonhosted.org/packages/f9/5c/6657823f4f594f72b5471f1db1ab12e26e890bb2e41897522d134d2a3e81/numpy-2.2.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0811bb762109d9708cca4d0b13c4f67146e3c3b7cf8d34018c722adb2d957c84", size = 20867828, upload-time = "2025-05-17T21:37:56.699Z" }, - { url = "https://files.pythonhosted.org/packages/dc/9e/14520dc3dadf3c803473bd07e9b2bd1b69bc583cb2497b47000fed2fa92f/numpy-2.2.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:287cc3162b6f01463ccd86be154f284d0893d2b3ed7292439ea97eafa8170e0b", size = 14143006, upload-time = "2025-05-17T21:38:18.291Z" }, - { url = "https://files.pythonhosted.org/packages/4f/06/7e96c57d90bebdce9918412087fc22ca9851cceaf5567a45c1f404480e9e/numpy-2.2.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:f1372f041402e37e5e633e586f62aa53de2eac8d98cbfb822806ce4bbefcb74d", size = 5076765, upload-time = "2025-05-17T21:38:27.319Z" }, - { url = "https://files.pythonhosted.org/packages/73/ed/63d920c23b4289fdac96ddbdd6132e9427790977d5457cd132f18e76eae0/numpy-2.2.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:55a4d33fa519660d69614a9fad433be87e5252f4b03850642f88993f7b2ca566", size = 6617736, upload-time = "2025-05-17T21:38:38.141Z" }, - { url = "https://files.pythonhosted.org/packages/85/c5/e19c8f99d83fd377ec8c7e0cf627a8049746da54afc24ef0a0cb73d5dfb5/numpy-2.2.6-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f92729c95468a2f4f15e9bb94c432a9229d0d50de67304399627a943201baa2f", size = 14010719, upload-time = "2025-05-17T21:38:58.433Z" }, - { url = "https://files.pythonhosted.org/packages/19/49/4df9123aafa7b539317bf6d342cb6d227e49f7a35b99c287a6109b13dd93/numpy-2.2.6-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1bc23a79bfabc5d056d106f9befb8d50c31ced2fbc70eedb8155aec74a45798f", size = 16526072, upload-time = "2025-05-17T21:39:22.638Z" }, - { url = "https://files.pythonhosted.org/packages/b2/6c/04b5f47f4f32f7c2b0e7260442a8cbcf8168b0e1a41ff1495da42f42a14f/numpy-2.2.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e3143e4451880bed956e706a3220b4e5cf6172ef05fcc397f6f36a550b1dd868", size = 15503213, upload-time = "2025-05-17T21:39:45.865Z" }, - { url = "https://files.pythonhosted.org/packages/17/0a/5cd92e352c1307640d5b6fec1b2ffb06cd0dabe7d7b8227f97933d378422/numpy-2.2.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b4f13750ce79751586ae2eb824ba7e1e8dba64784086c98cdbbcc6a42112ce0d", size = 18316632, upload-time = "2025-05-17T21:40:13.331Z" }, - { url = "https://files.pythonhosted.org/packages/f0/3b/5cba2b1d88760ef86596ad0f3d484b1cbff7c115ae2429678465057c5155/numpy-2.2.6-cp313-cp313-win32.whl", hash = "sha256:5beb72339d9d4fa36522fc63802f469b13cdbe4fdab4a288f0c441b74272ebfd", size = 6244532, upload-time = "2025-05-17T21:43:46.099Z" }, - { url = "https://files.pythonhosted.org/packages/cb/3b/d58c12eafcb298d4e6d0d40216866ab15f59e55d148a5658bb3132311fcf/numpy-2.2.6-cp313-cp313-win_amd64.whl", hash = "sha256:b0544343a702fa80c95ad5d3d608ea3599dd54d4632df855e4c8d24eb6ecfa1c", size = 12610885, upload-time = "2025-05-17T21:44:05.145Z" }, - { url = "https://files.pythonhosted.org/packages/6b/9e/4bf918b818e516322db999ac25d00c75788ddfd2d2ade4fa66f1f38097e1/numpy-2.2.6-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0bca768cd85ae743b2affdc762d617eddf3bcf8724435498a1e80132d04879e6", size = 20963467, upload-time = "2025-05-17T21:40:44Z" }, - { url = "https://files.pythonhosted.org/packages/61/66/d2de6b291507517ff2e438e13ff7b1e2cdbdb7cb40b3ed475377aece69f9/numpy-2.2.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:fc0c5673685c508a142ca65209b4e79ed6740a4ed6b2267dbba90f34b0b3cfda", size = 14225144, upload-time = "2025-05-17T21:41:05.695Z" }, - { url = "https://files.pythonhosted.org/packages/e4/25/480387655407ead912e28ba3a820bc69af9adf13bcbe40b299d454ec011f/numpy-2.2.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:5bd4fc3ac8926b3819797a7c0e2631eb889b4118a9898c84f585a54d475b7e40", size = 5200217, upload-time = "2025-05-17T21:41:15.903Z" }, - { url = "https://files.pythonhosted.org/packages/aa/4a/6e313b5108f53dcbf3aca0c0f3e9c92f4c10ce57a0a721851f9785872895/numpy-2.2.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:fee4236c876c4e8369388054d02d0e9bb84821feb1a64dd59e137e6511a551f8", size = 6712014, upload-time = "2025-05-17T21:41:27.321Z" }, - { url = "https://files.pythonhosted.org/packages/b7/30/172c2d5c4be71fdf476e9de553443cf8e25feddbe185e0bd88b096915bcc/numpy-2.2.6-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e1dda9c7e08dc141e0247a5b8f49cf05984955246a327d4c48bda16821947b2f", size = 14077935, upload-time = "2025-05-17T21:41:49.738Z" }, - { url = "https://files.pythonhosted.org/packages/12/fb/9e743f8d4e4d3c710902cf87af3512082ae3d43b945d5d16563f26ec251d/numpy-2.2.6-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f447e6acb680fd307f40d3da4852208af94afdfab89cf850986c3ca00562f4fa", size = 16600122, upload-time = "2025-05-17T21:42:14.046Z" }, - { url = "https://files.pythonhosted.org/packages/12/75/ee20da0e58d3a66f204f38916757e01e33a9737d0b22373b3eb5a27358f9/numpy-2.2.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:389d771b1623ec92636b0786bc4ae56abafad4a4c513d36a55dce14bd9ce8571", size = 15586143, upload-time = "2025-05-17T21:42:37.464Z" }, - { url = "https://files.pythonhosted.org/packages/76/95/bef5b37f29fc5e739947e9ce5179ad402875633308504a52d188302319c8/numpy-2.2.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8e9ace4a37db23421249ed236fdcdd457d671e25146786dfc96835cd951aa7c1", size = 18385260, upload-time = "2025-05-17T21:43:05.189Z" }, - { url = "https://files.pythonhosted.org/packages/09/04/f2f83279d287407cf36a7a8053a5abe7be3622a4363337338f2585e4afda/numpy-2.2.6-cp313-cp313t-win32.whl", hash = "sha256:038613e9fb8c72b0a41f025a7e4c3f0b7a1b5d768ece4796b674c8f3fe13efff", size = 6377225, upload-time = "2025-05-17T21:43:16.254Z" }, - { url = "https://files.pythonhosted.org/packages/67/0e/35082d13c09c02c011cf21570543d202ad929d961c02a147493cb0c2bdf5/numpy-2.2.6-cp313-cp313t-win_amd64.whl", hash = "sha256:6031dd6dfecc0cf9f668681a37648373bddd6421fff6c66ec1624eed0180ee06", size = 12771374, upload-time = "2025-05-17T21:43:35.479Z" }, - { url = "https://files.pythonhosted.org/packages/9e/3b/d94a75f4dbf1ef5d321523ecac21ef23a3cd2ac8b78ae2aac40873590229/numpy-2.2.6-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:0b605b275d7bd0c640cad4e5d30fa701a8d59302e127e5f79138ad62762c3e3d", size = 21040391, upload-time = "2025-05-17T21:44:35.948Z" }, - { url = "https://files.pythonhosted.org/packages/17/f4/09b2fa1b58f0fb4f7c7963a1649c64c4d315752240377ed74d9cd878f7b5/numpy-2.2.6-pp310-pypy310_pp73-macosx_14_0_x86_64.whl", hash = "sha256:7befc596a7dc9da8a337f79802ee8adb30a552a94f792b9c9d18c840055907db", size = 6786754, upload-time = "2025-05-17T21:44:47.446Z" }, - { url = "https://files.pythonhosted.org/packages/af/30/feba75f143bdc868a1cc3f44ccfa6c4b9ec522b36458e738cd00f67b573f/numpy-2.2.6-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce47521a4754c8f4593837384bd3424880629f718d87c5d44f8ed763edd63543", size = 16643476, upload-time = "2025-05-17T21:45:11.871Z" }, - { url = "https://files.pythonhosted.org/packages/37/48/ac2a9584402fb6c0cd5b5d1a91dcf176b15760130dd386bbafdbfe3640bf/numpy-2.2.6-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:d042d24c90c41b54fd506da306759e06e568864df8ec17ccc17e9e884634fd00", size = 12812666, upload-time = "2025-05-17T21:45:31.426Z" }, -] - -[[package]] -name = "numpy" -version = "2.4.6" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.11'", -] -sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b3/49/ec46835a70be8fa6446c495126ac84fdb28cb2558e1620ffb87a10c8b64c/numpy-2.4.6-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0280e0356c0829a18d9de1cb7eee50ec22ca639878d7240307ca0943d73cd2c4", size = 16969194, upload-time = "2026-05-18T23:33:13.503Z" }, - { url = "https://files.pythonhosted.org/packages/0e/0d/f5957185c0ee2f3e12f78715aa9e3b353fd83633316c8532b38faa37e3f6/numpy-2.4.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:110f8b71aacb688ec69062bb7f6938a0f8acb01b7c1c4beb453c65b6d234584d", size = 14964111, upload-time = "2026-05-18T23:33:17.795Z" }, - { url = "https://files.pythonhosted.org/packages/ad/40/40a40ee0ddf7ceb782c49af278894b686e586d65d8c1889c8b5da01a3d7d/numpy-2.4.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:4cfe66903cc32a9921a6733d96b19bb6abf310397581bbad89c228f5abaf0ee8", size = 5469159, upload-time = "2026-05-18T23:33:20.654Z" }, - { url = "https://files.pythonhosted.org/packages/63/13/f9a8046535cb21deae82f8d03de9617e08882d274fad2539630761888228/numpy-2.4.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:8155154c7c691289fe18f510b5d4657c68c67989f293f0535a91360392ff6538", size = 6798936, upload-time = "2026-05-18T23:33:22.987Z" }, - { url = "https://files.pythonhosted.org/packages/33/a8/6fa8c1a345a8c85dbb21932c447bee07c30a2c2a3f31e369c0a84b300147/numpy-2.4.6-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ab0a9c4ffb1a6d95ef519fe4247dba8eb6b18ad93999f76b7f657039acabd47", size = 15966692, upload-time = "2026-05-18T23:33:26.62Z" }, - { url = "https://files.pythonhosted.org/packages/02/03/74fe2a4cb3817d94d86402f2506554130a2f01414e299b5a843e5a8a957f/numpy-2.4.6-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:89cd468399cfd2504718f0ba50e410dca55a170b61a02ad92bb18c8a65186e93", size = 16918164, upload-time = "2026-05-18T23:33:29.955Z" }, - { url = "https://files.pythonhosted.org/packages/c5/80/3615be3313f7e7696609bc194b9f0101da809df79e859bdb84e0cd043f46/numpy-2.4.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:c2d37ab77531417474168eb79d6d80b14f821a966818505d03013d0833edb7a8", size = 17322877, upload-time = "2026-05-18T23:33:34.724Z" }, - { url = "https://files.pythonhosted.org/packages/ca/ac/a691e0fe2675e370d0e08ff905adc49a1c8830e8cae03efe4477e92cd55d/numpy-2.4.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f407cb6b8e9d6d8c626bc73c945db1706035af8fd632295547bf1c9e46d092d6", size = 18651487, upload-time = "2026-05-18T23:33:38.217Z" }, - { url = "https://files.pythonhosted.org/packages/15/a7/9bc1cd626d7bf6869bfedf27b91b6ab5dd607758bf8e959d6fa80c6a59cb/numpy-2.4.6-cp311-cp311-win32.whl", hash = "sha256:ddea102b48f9e339f3948bf22040944184627a30fdf7f858667673b9c5f033c8", size = 6233945, upload-time = "2026-05-18T23:33:41.331Z" }, - { url = "https://files.pythonhosted.org/packages/c5/31/7fc6239c12bce7e931463251cca4426c465e1876ba3cc785402ef4dd8f4e/numpy-2.4.6-cp311-cp311-win_amd64.whl", hash = "sha256:1e254a00cdf42b1e4d5b3d68d33af63268d41340d8885df2ab6470f2e1500147", size = 12608406, upload-time = "2026-05-18T23:33:44.131Z" }, - { url = "https://files.pythonhosted.org/packages/27/83/140f85a466595a16382996a1bf06b2b54bcd597488921b0c9daaeeda72af/numpy-2.4.6-cp311-cp311-win_arm64.whl", hash = "sha256:ed9749eef4cbd126da3dc1d6bcb3a57f5eb7ac6a6484146bdbf743f552dfc577", size = 10479528, upload-time = "2026-05-18T23:33:50.725Z" }, - { url = "https://files.pythonhosted.org/packages/95/2a/3d7b5ac8aac24feaf9ad7ed58f45b0bbc06d37e4338ae84c9f2298b570f9/numpy-2.4.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:001fbb8e08d942dd57599e781f2472269ee7f2755fae407b4f67b2f0b17da3f1", size = 16689119, upload-time = "2026-05-18T23:33:54.065Z" }, - { url = "https://files.pythonhosted.org/packages/ea/12/92c4c131527599e8288d6918e888d88726f84d805d784b771f32408aeaef/numpy-2.4.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ebfb099f8dcf083deef3ac1ca4c1503f387cf76296fcb3816b66f5ecb5f54fdb", size = 14699246, upload-time = "2026-05-18T23:33:57.621Z" }, - { url = "https://files.pythonhosted.org/packages/ad/fe/c0a6b7b2ca128a8fb228575147073b660656734b8ebe4d76c8fd748dcc79/numpy-2.4.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:3213d622a0283a39a93d188f3cf72b26862df52fbb4ca3697f51705016523d41", size = 5204410, upload-time = "2026-05-18T23:34:00.302Z" }, - { url = "https://files.pythonhosted.org/packages/f3/d4/9770d14ba719432bb90a421bfd443872ed0f70f7264b64bec12ea363d5fd/numpy-2.4.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:357cc07a6d7b0b182ff02249616a03742827ebb1277546b5c7cd7f7620a45698", size = 6551240, upload-time = "2026-05-18T23:34:02.852Z" }, - { url = "https://files.pythonhosted.org/packages/c9/c6/50a46a6205feba2343f1d6d17438107c5dc491ed1c736e6ea68689fd906b/numpy-2.4.6-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5f9fb9157b4ce2971008323afe46053787b526ef624fea915b261468a8421a0f", size = 15671012, upload-time = "2026-05-18T23:34:05.485Z" }, - { url = "https://files.pythonhosted.org/packages/99/60/14115e6364fa676c5397c2ad3004e527e9aa487abf5d0706ec81bbd08529/numpy-2.4.6-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:90f9849678c75fe7afa2d348ac842c168b0a4d3d61919687216dfc547976d853", size = 16645538, upload-time = "2026-05-18T23:34:09.265Z" }, - { url = "https://files.pythonhosted.org/packages/ae/c5/693cbe59e57db94d2231fa519ca3978dc9e19da5a8f088588f5c6e947ff2/numpy-2.4.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c1a2af6c6ef86344a6b0db6b97834208bf598db514f2b155042439b62605601a", size = 17020706, upload-time = "2026-05-18T23:34:13.053Z" }, - { url = "https://files.pythonhosted.org/packages/ef/fc/85b7c4eff9b4966ade25c2273cf7e7012e92366c032058653934b37de044/numpy-2.4.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e5805d5a22fd19c8ccff10a9561f9df94436b0545619ea579db2d3c35294bce2", size = 18368541, upload-time = "2026-05-18T23:34:17.024Z" }, - { url = "https://files.pythonhosted.org/packages/f6/81/e1b27545deedce7f4a0b348618c6b62d74e36a4dc9ccd42f3eb2f85eee32/numpy-2.4.6-cp312-cp312-win32.whl", hash = "sha256:e3eeb0aabd6bd5ce64faae67e9935203a6991b4bc2a485a767fbafb2c5125f45", size = 5962825, upload-time = "2026-05-18T23:34:20.3Z" }, - { url = "https://files.pythonhosted.org/packages/ab/ca/feab00bd44aa5fe1ad2c18f08b4d3bb92e26484b0b1d1443897809ed528c/numpy-2.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:d8e8286dd7cea7895157318d1b91cdacac64c479f3cbc8dce548331728484751", size = 12321687, upload-time = "2026-05-18T23:34:23.095Z" }, - { url = "https://files.pythonhosted.org/packages/63/cf/5a6d34850a39d1093558564f77ee8e8e0bee5061151b8f05a55711001ec7/numpy-2.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:4081eb135ac24158bd51cdfbef16f1c64df7063b1143f24731387137c092bec8", size = 10221482, upload-time = "2026-05-18T23:34:25.876Z" }, - { url = "https://files.pythonhosted.org/packages/fb/82/bdab26d7438c6791ca31b7c024ca37c1eab8b726ba236129005cd4a06e45/numpy-2.4.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:511dbaf848decaaaf4b4ca48032619fb3138710c4bf7da7617765edad1ef96b0", size = 16684648, upload-time = "2026-05-18T23:34:29.41Z" }, - { url = "https://files.pythonhosted.org/packages/1b/30/a80189bcc7f5e4258b3fbc3968d909d1756f54d023299ecc39ad6fdb9ef8/numpy-2.4.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bf162abab1c1a736333192707cef898e735a5ca00f38f27eeedf44b39d9e85eb", size = 14693902, upload-time = "2026-05-18T23:34:33.013Z" }, - { url = "https://files.pythonhosted.org/packages/97/12/70b5d0d7c15e1ebb8a6a84a8caa1d19e181d84fb58bb6d70aca29099dec1/numpy-2.4.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:043191bfa8eab18c776647b62723ac9dddece59743b13f49b2016094129c2b3f", size = 5198992, upload-time = "2026-05-18T23:34:36.132Z" }, - { url = "https://files.pythonhosted.org/packages/ba/8c/ebd2a8f8a83541f8d38cc5667e8c2b69cecfd30da6e45693e8158857d44b/numpy-2.4.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:6180d8b35af935aed8ece3a85e0a43f87393ae0ac87c8d2c8bd2c993f7270ef3", size = 6546944, upload-time = "2026-05-18T23:34:38.484Z" }, - { url = "https://files.pythonhosted.org/packages/bb/c5/7b863a97a91671a0338f4253bd3b5a3d3852f0692dae91711c9f4a10e787/numpy-2.4.6-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:72fbe16c6fac95aedf5937fa873445cec2110be35d8a4e9433d7501fd98dae6b", size = 15669392, upload-time = "2026-05-18T23:34:41.257Z" }, - { url = "https://files.pythonhosted.org/packages/a5/9d/3584b9984ca4c047aea75214ce1a4c4c73d849bd71b604264b7f5653f8a8/numpy-2.4.6-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a7830bab239b79cda9c08c2da014761cafb48da6150e1da17ac06283f43b6089", size = 16633220, upload-time = "2026-05-18T23:34:45.075Z" }, - { url = "https://files.pythonhosted.org/packages/05/ae/7c67fba23bd98caec7c99261f3a16072ade14813486b0282cb29846de832/numpy-2.4.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ef4aea96ce4d3b074422cb4f2f64e216bf9e213004bb58ecfdf50ea02ea8eb9a", size = 17020800, upload-time = "2026-05-18T23:34:49.065Z" }, - { url = "https://files.pythonhosted.org/packages/d9/5d/3b6725cb31d983c5e66916f5d36f6d7e5521129e4c4404d64f918292a5b6/numpy-2.4.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dfa20cc6ca228e6b155b11da03825975ce66aea520985dbbddf0f2a5a495c605", size = 18357600, upload-time = "2026-05-18T23:34:52.709Z" }, - { url = "https://files.pythonhosted.org/packages/f7/da/2ccc6c2fe8898dee01d90c75c5f5f914a23daf99e3e0f59516a08760c8b5/numpy-2.4.6-cp313-cp313-win32.whl", hash = "sha256:56b39e5e0622a09a25bf5baf62f4bcf0cb8a41ae6e2819cf49bbc5a74c083f91", size = 5961134, upload-time = "2026-05-18T23:34:55.618Z" }, - { url = "https://files.pythonhosted.org/packages/b5/cd/9cc4dc876fb065d5c220aae4d5e14826b2715331bb7618ce1fb07a679d99/numpy-2.4.6-cp313-cp313-win_amd64.whl", hash = "sha256:c4fc99836233ea196540b17ab0983aff60ed07941751930f5f4d05bc3b3b7359", size = 12318598, upload-time = "2026-05-18T23:34:58.928Z" }, - { url = "https://files.pythonhosted.org/packages/39/1e/c0bcba1f8694116485fe28fd1be698c278fcda4141c5b0e53a2aed8b12a8/numpy-2.4.6-cp313-cp313-win_arm64.whl", hash = "sha256:a7c711e21628b52034bb5ab8d1bce291f752fcc5e92accc615778acee1ff4778", size = 10222272, upload-time = "2026-05-18T23:35:02.167Z" }, - { url = "https://files.pythonhosted.org/packages/63/6d/cc5619247c8f4204e507f5883528372e4ac4bb189e579fb859a12e480b1f/numpy-2.4.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:112b06a867b235ef466ed3508ddf0238050df9c727cafb5301ac385b899189a1", size = 14821197, upload-time = "2026-05-18T23:35:05.468Z" }, - { url = "https://files.pythonhosted.org/packages/00/58/f1c39161c87d9e9bed660f1ed4bafc0e403d5ec9650b6dd77aead07d489b/numpy-2.4.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:eaf7fa2de5c0be8ae6ff8e9bea2ccd725e980541244521d8d4b5f3354a27babe", size = 5326287, upload-time = "2026-05-18T23:35:08.693Z" }, - { url = "https://files.pythonhosted.org/packages/af/57/3917ab0fd97f271a8694513581b8a36c655f111c446852c302f04ccdb6fc/numpy-2.4.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:7265a2f3d436e54ef9f2b52b5c937e6be778781bd97a590319d7348f1c1ca997", size = 6646763, upload-time = "2026-05-18T23:35:11.459Z" }, - { url = "https://files.pythonhosted.org/packages/eb/0f/037e64c494b67581ae18193d770adef354c41f3f2c8ebf865602d949bf8f/numpy-2.4.6-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f74a575920ab21fe304421a3fc28793d82e299cae9eccb37084e9fc7f3617c20", size = 15728070, upload-time = "2026-05-18T23:35:14.79Z" }, - { url = "https://files.pythonhosted.org/packages/21/a6/5d2bae9c9542eb4df16dc9c46dc79c186e9bad53805dfa5399a6023c6db0/numpy-2.4.6-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ede83e07a75dd06bc501566c1eca2afc0d61677c1472ac9ad93fdee6e638a48d", size = 16681752, upload-time = "2026-05-18T23:35:18.836Z" }, - { url = "https://files.pythonhosted.org/packages/92/14/23d1dfb410ae362cd59ce53e936b1513d545eb40db3949ced632e19a459e/numpy-2.4.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:68bb27509ac1b9a3443094260f6326150663b06abe40b73a2f81160623da5b67", size = 17086024, upload-time = "2026-05-18T23:35:22.52Z" }, - { url = "https://files.pythonhosted.org/packages/4b/6e/23595a2c642cdf3bc567877064bdd7f91c8b0038a4453cf2daf7248eafe9/numpy-2.4.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:a0df0043bdb289bde1f62da130d20df23d58b45429f752bc7a8fc5325a225ecd", size = 18403398, upload-time = "2026-05-18T23:35:26.398Z" }, - { url = "https://files.pythonhosted.org/packages/8a/90/0ac3bc947217e66dec77e7cbc6a1979d1af70b6461b82f620d3bccd5e4c8/numpy-2.4.6-cp313-cp313t-win32.whl", hash = "sha256:29a287e0cf63ff528da061de6b9f64a4618da591ca1046aafc54062e40ca7eab", size = 6084971, upload-time = "2026-05-18T23:35:29.387Z" }, - { url = "https://files.pythonhosted.org/packages/77/71/5673e351671a1d2bd6063b91b44f70c0affea7d1516fa7a6572941ba4aa1/numpy-2.4.6-cp313-cp313t-win_amd64.whl", hash = "sha256:25c692919ac5a01f170a3bfcd62d745b24fd095c353d50812637d6fcab442e75", size = 12458532, upload-time = "2026-05-18T23:35:32.175Z" }, - { url = "https://files.pythonhosted.org/packages/3f/88/19d3503c5046e688f049274b27a3ef3d771152fa80d3ba3d01a3dff61abe/numpy-2.4.6-cp313-cp313t-win_arm64.whl", hash = "sha256:1e978ec1e8bd0e0e4de6bb75de9d30cbb74db6b6a2bb727618613703ca0167dd", size = 10291881, upload-time = "2026-05-18T23:35:35.465Z" }, - { url = "https://files.pythonhosted.org/packages/f8/91/3ab2044d05fd16d343c5ac2e69b127f1b2854040dd20b193257c78028bd3/numpy-2.4.6-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:06ca2f61ec4385a07a6977c55ba998a4466c123642b4a32694d3128fce18c079", size = 16683458, upload-time = "2026-05-18T23:35:38.353Z" }, - { url = "https://files.pythonhosted.org/packages/8e/62/764ce66fa4147ae6d73071a3abf804ffe606f174618697c571acdf26a7c9/numpy-2.4.6-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:38efbc8de75c7a0fc1ac190162d892787f3f47b57cc291231aafee36b80982b7", size = 14704559, upload-time = "2026-05-18T23:35:42.14Z" }, - { url = "https://files.pythonhosted.org/packages/60/61/23f27c172f022e04025b7dc2367f4d63c1a398120607ec896228649a6f48/numpy-2.4.6-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d581b735e177fdcdce6fed8e7e8880a3fb6ee4e3653a3ac6af01c6f4c03effc5", size = 5209716, upload-time = "2026-05-18T23:35:45.377Z" }, - { url = "https://files.pythonhosted.org/packages/03/71/21cf70dc6ea3e3acb95fc53a265b2fc248b981f0194ceb5b475271b8809d/numpy-2.4.6-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:0a041d3d761dc3c35cc56ce0351506a02bcbc25f7b169f652435141a17db9096", size = 6543947, upload-time = "2026-05-18T23:35:47.926Z" }, - { url = "https://files.pythonhosted.org/packages/d5/91/64288395ee1799bd2e0b04a305dce9666da90c961e1f3fe982a05ee1c036/numpy-2.4.6-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:40fdc1ae7125e518ea98e53e69a4ebc27e1fd50510c47b7ea130cf21e5e1d42b", size = 15685197, upload-time = "2026-05-18T23:35:50.863Z" }, - { url = "https://files.pythonhosted.org/packages/f3/eb/ebffaa97dc55502df69584a8f0dcf07f69a3e0b3e2323670a2722db9aa39/numpy-2.4.6-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a2c306dea656c12c68f51f4cea133cbe78ca7435eb28c735eac1d3ebe73be6e8", size = 16638245, upload-time = "2026-05-18T23:35:54.752Z" }, - { url = "https://files.pythonhosted.org/packages/b8/0b/54f9da33128d7e350fab89c7455902eeae70349ee52bddb448dc4a576f45/numpy-2.4.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:33111801a01c12a8a1e3721f0a9232f8cfc8ae2c6b7098167e6f623c6073f402", size = 17036587, upload-time = "2026-05-18T23:35:58.355Z" }, - { url = "https://files.pythonhosted.org/packages/b6/f0/fdebc1052db1cc37c64beb22072d67cd6d1c71adca1299f53dec2b5e20d3/numpy-2.4.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ae506e6902902557576a26ff33eda8695e7ecb3cb36c3b573a0765dee114ebdb", size = 18363226, upload-time = "2026-05-18T23:36:02.845Z" }, - { url = "https://files.pythonhosted.org/packages/aa/b4/298628d98c72b57e57f7165ae6a481a1deaf6f3c28262a6e4c739c275930/numpy-2.4.6-cp314-cp314-win32.whl", hash = "sha256:aaf159caa35993cb1f56fb9b8e4610d35758e7ca005412eb1daa856a78c9c4b1", size = 6010196, upload-time = "2026-05-18T23:36:05.92Z" }, - { url = "https://files.pythonhosted.org/packages/df/ac/46de6dda46478f7942f839e094970be2d4a861e005c4b3bf07c92e291a09/numpy-2.4.6-cp314-cp314-win_amd64.whl", hash = "sha256:b507f5c4c1d508876d1819b6bf9a49d365b96320b5d4993426b33a23ca4b8261", size = 12450334, upload-time = "2026-05-18T23:36:09.107Z" }, - { url = "https://files.pythonhosted.org/packages/78/92/b8b798ac784102c0da830d2257d59358e3d3d90d1e2b3f2575dad976c5cf/numpy-2.4.6-cp314-cp314-win_arm64.whl", hash = "sha256:6f41ae150c4e32db4f3310cdaf64b1593a03dbabe29eec77fc9b50fe64061df6", size = 10495678, upload-time = "2026-05-18T23:36:12.766Z" }, - { url = "https://files.pythonhosted.org/packages/30/34/ec28d1aa8115971537c01469ab2011ee96827930f0a124de1000cc2a7ed7/numpy-2.4.6-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ece3d2cfe132e7d51f44a832b303895e6f2d499c5e74dfbdb06ee246147a304a", size = 14823672, upload-time = "2026-05-18T23:36:16.473Z" }, - { url = "https://files.pythonhosted.org/packages/16/bd/f6d1fede4e54e8042a7ff97bb495510f3c220f94bcd9e8b228e87c92cc0d/numpy-2.4.6-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:e3e5193ef5a3dc73bceee50f7fdc2c90dbb76c42df8d8fae3d1067a583df579e", size = 5328731, upload-time = "2026-05-18T23:36:19.767Z" }, - { url = "https://files.pythonhosted.org/packages/f4/f0/e105b9e2fd728a9910103884decd6951d9dd73896b914a98d9a231de02ee/numpy-2.4.6-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:17f9ade344e7d9b464a084d69bcf18fc691cb1db67c62ed80820bf4926d78f0e", size = 6649805, upload-time = "2026-05-18T23:36:22.266Z" }, - { url = "https://files.pythonhosted.org/packages/82/dd/1206a7ca6ab15e3f02069707ca96222e202af681bb73756da7527f3cb837/numpy-2.4.6-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9cd5ffd25db4e7ba6a375693b3fc0fc1791ec636c17db3720da19bde7180ec43", size = 15730496, upload-time = "2026-05-18T23:36:25.713Z" }, - { url = "https://files.pythonhosted.org/packages/51/e7/38d3ea825dcab85a591734decb2f6c67caa7c8367d374df1a1c3842f9b07/numpy-2.4.6-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7d92c3819208a60205a12a245c91ad70cb0a85336659b19b834205573ac8456e", size = 16679616, upload-time = "2026-05-18T23:36:29.652Z" }, - { url = "https://files.pythonhosted.org/packages/93/b7/caabfdf53edf663e0b4eb74d7d405d83baef09eb5e83bcd32d601d72b93e/numpy-2.4.6-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e85b752a1e912b70eaad4fafbd4d1238007ab221de2009b9a2f5ae7461239895", size = 17085145, upload-time = "2026-05-18T23:36:33.449Z" }, - { url = "https://files.pythonhosted.org/packages/f9/45/68d7c33a6bcf3e5aa3bdbd57a367e6f615286dfd6482f97e8ffeb734306e/numpy-2.4.6-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:29cb7f67d10b479ff07c17d33e39f78c07f71c40ef30d63c153d340e96cd3fb4", size = 18403813, upload-time = "2026-05-18T23:36:37.369Z" }, - { url = "https://files.pythonhosted.org/packages/9c/50/0753655aa844c99cd9e018aacf76f130f1bd81d881bb74bc0aef5d73a8ba/numpy-2.4.6-cp314-cp314t-win32.whl", hash = "sha256:260a5d70215b61ab4fadf5c7baacd64821842975eea312125ed3c39a6391b063", size = 6156982, upload-time = "2026-05-18T23:36:40.817Z" }, - { url = "https://files.pythonhosted.org/packages/b2/d4/7c67becf668f973cb490cec3e98dfd799d866f9c989a54d355672cfa0db6/numpy-2.4.6-cp314-cp314t-win_amd64.whl", hash = "sha256:81a1cca95ed5bb92aa8b10dd2cdc9a0d3853a50fad926c28b5d7e8ea54389627", size = 12638908, upload-time = "2026-05-18T23:36:43.996Z" }, - { url = "https://files.pythonhosted.org/packages/43/bb/e1c71a4295b1b1d1393d50dbb4f2a36283c6859d9d3892e84f00ec5a91d5/numpy-2.4.6-cp314-cp314t-win_arm64.whl", hash = "sha256:0c9136e14ed34a9e343a31c533d78a9813a69a3148332bce5e9821cb2f996e66", size = 10565867, upload-time = "2026-05-18T23:36:47.114Z" }, - { url = "https://files.pythonhosted.org/packages/de/12/b422cc84439adc0d00de605bf4a308890ae5c26f2c71fbd73e5d08fbb0dd/numpy-2.4.6-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:55cced7c52e981362f708ad635198e97a752dfba412cc03c23bbf3bd8d5cd662", size = 16847511, upload-time = "2026-05-18T23:36:50.673Z" }, - { url = "https://files.pythonhosted.org/packages/44/53/f481bef68011740f8849418d82db07230e825013f31f4eef5ba5b805316a/numpy-2.4.6-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:d6da64deb6b8ed903e7560180a92f2d804ee1ba5eeb849ac2748b8c1aba1f6d7", size = 14889064, upload-time = "2026-05-18T23:36:53.879Z" }, - { url = "https://files.pythonhosted.org/packages/7f/57/42ed575c10ced8af951d426bc4e1f8aff16fd851db33f067036215a7f860/numpy-2.4.6-pp311-pypy311_pp73-macosx_14_0_arm64.whl", hash = "sha256:68a5124b13fa6cc2086764a20005d30bc0548146f7f5322f02fce212ca14317f", size = 5394157, upload-time = "2026-05-18T23:36:57.194Z" }, - { url = "https://files.pythonhosted.org/packages/6a/ef/f66cc724fcc36c1e364c67f51ae9146090b8b584f27d58b97fdae3edd737/numpy-2.4.6-pp311-pypy311_pp73-macosx_14_0_x86_64.whl", hash = "sha256:948424b06129ce883307e8cff868c31396d8dc7630a59c61d70d98dbe70f222c", size = 6708728, upload-time = "2026-05-18T23:36:59.575Z" }, - { url = "https://files.pythonhosted.org/packages/1a/9c/c531f2293b91265d8b48e9b329f54fdd7ffae73cb4134ea10cca4237e9cc/numpy-2.4.6-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5dbbdb29840ca3d91ee0fece42fc29278886d908280bfec0a5846c6f901a3eb0", size = 15798374, upload-time = "2026-05-18T23:37:02.674Z" }, - { url = "https://files.pythonhosted.org/packages/1a/b0/413077f6b1153ed3cba361401c6783bbad6114804a000cc22eb71c13e190/numpy-2.4.6-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8ad03c0965fb3c692200e74d458ca28c1dbb4ce96f9a479a8aa041ad5fabca02", size = 16747286, upload-time = "2026-05-18T23:37:06.327Z" }, - { url = "https://files.pythonhosted.org/packages/15/ce/e5ec180bc41812edcd8daeb8639d205622c0e8c02259d8ab25a0201b3c2a/numpy-2.4.6-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:2803abfebfc990042cd494d8ce2d5f82e9d847af6d35ec486923aa19dbad5e73", size = 12504263, upload-time = "2026-05-18T23:37:09.715Z" }, -] - -[[package]] -name = "packaging" -version = "26.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134, upload-time = "2026-04-24T20:15:23.917Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" }, -] - -[[package]] -name = "pillow" -version = "12.2.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/8c/21/c2bcdd5906101a30244eaffc1b6e6ce71a31bd0742a01eb89e660ebfac2d/pillow-12.2.0.tar.gz", hash = "sha256:a830b1a40919539d07806aa58e1b114df53ddd43213d9c8b75847eee6c0182b5", size = 46987819, upload-time = "2026-04-01T14:46:17.687Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3a/aa/d0b28e1c811cd4d5f5c2bfe2e022292bd255ae5744a3b9ac7d6c8f72dd75/pillow-12.2.0-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:a4e8f36e677d3336f35089648c8955c51c6d386a13cf6ee9c189c5f5bd713a9f", size = 5354355, upload-time = "2026-04-01T14:42:15.402Z" }, - { url = "https://files.pythonhosted.org/packages/27/8e/1d5b39b8ae2bd7650d0c7b6abb9602d16043ead9ebbfef4bc4047454da2a/pillow-12.2.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2e589959f10d9824d39b350472b92f0ce3b443c0a3442ebf41c40cb8361c5b97", size = 4695871, upload-time = "2026-04-01T14:42:18.234Z" }, - { url = "https://files.pythonhosted.org/packages/f0/c5/dcb7a6ca6b7d3be41a76958e90018d56c8462166b3ef223150360850c8da/pillow-12.2.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a52edc8bfff4429aaabdf4d9ee0daadbbf8562364f940937b941f87a4290f5ff", size = 6269734, upload-time = "2026-04-01T14:42:20.608Z" }, - { url = "https://files.pythonhosted.org/packages/ea/f1/aa1bb13b2f4eba914e9637893c73f2af8e48d7d4023b9d3750d4c5eb2d0c/pillow-12.2.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:975385f4776fafde056abb318f612ef6285b10a1f12b8570f3647ad0d74b48ec", size = 8076080, upload-time = "2026-04-01T14:42:23.095Z" }, - { url = "https://files.pythonhosted.org/packages/a1/2a/8c79d6a53169937784604a8ae8d77e45888c41537f7f6f65ed1f407fe66d/pillow-12.2.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bd9c0c7a0c681a347b3194c500cb1e6ca9cab053ea4d82a5cf45b6b754560136", size = 6382236, upload-time = "2026-04-01T14:42:25.82Z" }, - { url = "https://files.pythonhosted.org/packages/b5/42/bbcb6051030e1e421d103ce7a8ecadf837aa2f39b8f82ef1a8d37c3d4ebc/pillow-12.2.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:88d387ff40b3ff7c274947ed3125dedf5262ec6919d83946753b5f3d7c67ea4c", size = 7070220, upload-time = "2026-04-01T14:42:28.68Z" }, - { url = "https://files.pythonhosted.org/packages/3f/e1/c2a7d6dd8cfa6b231227da096fd2d58754bab3603b9d73bf609d3c18b64f/pillow-12.2.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:51c4167c34b0d8ba05b547a3bb23578d0ba17b80a5593f93bd8ecb123dd336a3", size = 6493124, upload-time = "2026-04-01T14:42:31.579Z" }, - { url = "https://files.pythonhosted.org/packages/5f/41/7c8617da5d32e1d2f026e509484fdb6f3ad7efaef1749a0c1928adbb099e/pillow-12.2.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:34c0d99ecccea270c04882cb3b86e7b57296079c9a4aff88cb3b33563d95afaa", size = 7194324, upload-time = "2026-04-01T14:42:34.615Z" }, - { url = "https://files.pythonhosted.org/packages/2d/de/a777627e19fd6d62f84070ee1521adde5eeda4855b5cf60fe0b149118bca/pillow-12.2.0-cp310-cp310-win32.whl", hash = "sha256:b85f66ae9eb53e860a873b858b789217ba505e5e405a24b85c0464822fe88032", size = 6376363, upload-time = "2026-04-01T14:42:37.19Z" }, - { url = "https://files.pythonhosted.org/packages/e7/34/fc4cb5204896465842767b96d250c08410f01f2f28afc43b257de842eed5/pillow-12.2.0-cp310-cp310-win_amd64.whl", hash = "sha256:673aa32138f3e7531ccdbca7b3901dba9b70940a19ccecc6a37c77d5fdeb05b5", size = 7083523, upload-time = "2026-04-01T14:42:39.62Z" }, - { url = "https://files.pythonhosted.org/packages/2d/a0/32852d36bc7709f14dc3f64f929a275e958ad8c19a6deba9610d458e28b3/pillow-12.2.0-cp310-cp310-win_arm64.whl", hash = "sha256:3e080565d8d7c671db5802eedfb438e5565ffa40115216eabb8cd52d0ecce024", size = 2463318, upload-time = "2026-04-01T14:42:42.063Z" }, - { url = "https://files.pythonhosted.org/packages/68/e1/748f5663efe6edcfc4e74b2b93edfb9b8b99b67f21a854c3ae416500a2d9/pillow-12.2.0-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:8be29e59487a79f173507c30ddf57e733a357f67881430449bb32614075a40ab", size = 5354347, upload-time = "2026-04-01T14:42:44.255Z" }, - { url = "https://files.pythonhosted.org/packages/47/a1/d5ff69e747374c33a3b53b9f98cca7889fce1fd03d79cdc4e1bccc6c5a87/pillow-12.2.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:71cde9a1e1551df7d34a25462fc60325e8a11a82cc2e2f54578e5e9a1e153d65", size = 4695873, upload-time = "2026-04-01T14:42:46.452Z" }, - { url = "https://files.pythonhosted.org/packages/df/21/e3fbdf54408a973c7f7f89a23b2cb97a7ef30c61ab4142af31eee6aebc88/pillow-12.2.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f490f9368b6fc026f021db16d7ec2fbf7d89e2edb42e8ec09d2c60505f5729c7", size = 6280168, upload-time = "2026-04-01T14:42:49.228Z" }, - { url = "https://files.pythonhosted.org/packages/d3/f1/00b7278c7dd52b17ad4329153748f87b6756ec195ff786c2bdf12518337d/pillow-12.2.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8bd7903a5f2a4545f6fd5935c90058b89d30045568985a71c79f5fd6edf9b91e", size = 8088188, upload-time = "2026-04-01T14:42:51.735Z" }, - { url = "https://files.pythonhosted.org/packages/ad/cf/220a5994ef1b10e70e85748b75649d77d506499352be135a4989c957b701/pillow-12.2.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3997232e10d2920a68d25191392e3a4487d8183039e1c74c2297f00ed1c50705", size = 6394401, upload-time = "2026-04-01T14:42:54.343Z" }, - { url = "https://files.pythonhosted.org/packages/e9/bd/e51a61b1054f09437acfbc2ff9106c30d1eb76bc1453d428399946781253/pillow-12.2.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e74473c875d78b8e9d5da2a70f7099549f9eb37ded4e2f6a463e60125bccd176", size = 7079655, upload-time = "2026-04-01T14:42:56.954Z" }, - { url = "https://files.pythonhosted.org/packages/6b/3d/45132c57d5fb4b5744567c3817026480ac7fc3ce5d4c47902bc0e7f6f853/pillow-12.2.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:56a3f9c60a13133a98ecff6197af34d7824de9b7b38c3654861a725c970c197b", size = 6503105, upload-time = "2026-04-01T14:42:59.847Z" }, - { url = "https://files.pythonhosted.org/packages/7d/2e/9df2fc1e82097b1df3dce58dc43286aa01068e918c07574711fcc53e6fb4/pillow-12.2.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:90e6f81de50ad6b534cab6e5aef77ff6e37722b2f5d908686f4a5c9eba17a909", size = 7203402, upload-time = "2026-04-01T14:43:02.664Z" }, - { url = "https://files.pythonhosted.org/packages/bd/2e/2941e42858ebb67e50ae741473de81c2984e6eff7b397017623c676e2e8d/pillow-12.2.0-cp311-cp311-win32.whl", hash = "sha256:8c984051042858021a54926eb597d6ee3012393ce9c181814115df4c60b9a808", size = 6378149, upload-time = "2026-04-01T14:43:05.274Z" }, - { url = "https://files.pythonhosted.org/packages/69/42/836b6f3cd7f3e5fa10a1f1a5420447c17966044c8fbf589cc0452d5502db/pillow-12.2.0-cp311-cp311-win_amd64.whl", hash = "sha256:6e6b2a0c538fc200b38ff9eb6628228b77908c319a005815f2dde585a0664b60", size = 7082626, upload-time = "2026-04-01T14:43:08.557Z" }, - { url = "https://files.pythonhosted.org/packages/c2/88/549194b5d6f1f494b485e493edc6693c0a16f4ada488e5bd974ed1f42fad/pillow-12.2.0-cp311-cp311-win_arm64.whl", hash = "sha256:9a8a34cc89c67a65ea7437ce257cea81a9dad65b29805f3ecee8c8fe8ff25ffe", size = 2463531, upload-time = "2026-04-01T14:43:10.743Z" }, - { url = "https://files.pythonhosted.org/packages/58/be/7482c8a5ebebbc6470b3eb791812fff7d5e0216c2be3827b30b8bb6603ed/pillow-12.2.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2d192a155bbcec180f8564f693e6fd9bccff5a7af9b32e2e4bf8c9c69dbad6b5", size = 5308279, upload-time = "2026-04-01T14:43:13.246Z" }, - { url = "https://files.pythonhosted.org/packages/d8/95/0a351b9289c2b5cbde0bacd4a83ebc44023e835490a727b2a3bd60ddc0f4/pillow-12.2.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f3f40b3c5a968281fd507d519e444c35f0ff171237f4fdde090dd60699458421", size = 4695490, upload-time = "2026-04-01T14:43:15.584Z" }, - { url = "https://files.pythonhosted.org/packages/de/af/4e8e6869cbed569d43c416fad3dc4ecb944cb5d9492defaed89ddd6fe871/pillow-12.2.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:03e7e372d5240cc23e9f07deca4d775c0817bffc641b01e9c3af208dbd300987", size = 6284462, upload-time = "2026-04-01T14:43:18.268Z" }, - { url = "https://files.pythonhosted.org/packages/e9/9e/c05e19657fd57841e476be1ab46c4d501bffbadbafdc31a6d665f8b737b6/pillow-12.2.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b86024e52a1b269467a802258c25521e6d742349d760728092e1bc2d135b4d76", size = 8094744, upload-time = "2026-04-01T14:43:20.716Z" }, - { url = "https://files.pythonhosted.org/packages/2b/54/1789c455ed10176066b6e7e6da1b01e50e36f94ba584dc68d9eebfe9156d/pillow-12.2.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7371b48c4fa448d20d2714c9a1f775a81155050d383333e0a6c15b1123dda005", size = 6398371, upload-time = "2026-04-01T14:43:23.443Z" }, - { url = "https://files.pythonhosted.org/packages/43/e3/fdc657359e919462369869f1c9f0e973f353f9a9ee295a39b1fea8ee1a77/pillow-12.2.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:62f5409336adb0663b7caa0da5c7d9e7bdbaae9ce761d34669420c2a801b2780", size = 7087215, upload-time = "2026-04-01T14:43:26.758Z" }, - { url = "https://files.pythonhosted.org/packages/8b/f8/2f6825e441d5b1959d2ca5adec984210f1ec086435b0ed5f52c19b3b8a6e/pillow-12.2.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:01afa7cf67f74f09523699b4e88c73fb55c13346d212a59a2db1f86b0a63e8c5", size = 6509783, upload-time = "2026-04-01T14:43:29.56Z" }, - { url = "https://files.pythonhosted.org/packages/67/f9/029a27095ad20f854f9dba026b3ea6428548316e057e6fc3545409e86651/pillow-12.2.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fc3d34d4a8fbec3e88a79b92e5465e0f9b842b628675850d860b8bd300b159f5", size = 7212112, upload-time = "2026-04-01T14:43:32.091Z" }, - { url = "https://files.pythonhosted.org/packages/be/42/025cfe05d1be22dbfdb4f264fe9de1ccda83f66e4fc3aac94748e784af04/pillow-12.2.0-cp312-cp312-win32.whl", hash = "sha256:58f62cc0f00fd29e64b29f4fd923ffdb3859c9f9e6105bfc37ba1d08994e8940", size = 6378489, upload-time = "2026-04-01T14:43:34.601Z" }, - { url = "https://files.pythonhosted.org/packages/5d/7b/25a221d2c761c6a8ae21bfa3874988ff2583e19cf8a27bf2fee358df7942/pillow-12.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:7f84204dee22a783350679a0333981df803dac21a0190d706a50475e361c93f5", size = 7084129, upload-time = "2026-04-01T14:43:37.213Z" }, - { url = "https://files.pythonhosted.org/packages/10/e1/542a474affab20fd4a0f1836cb234e8493519da6b76899e30bcc5d990b8b/pillow-12.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:af73337013e0b3b46f175e79492d96845b16126ddf79c438d7ea7ff27783a414", size = 2463612, upload-time = "2026-04-01T14:43:39.421Z" }, - { url = "https://files.pythonhosted.org/packages/4a/01/53d10cf0dbad820a8db274d259a37ba50b88b24768ddccec07355382d5ad/pillow-12.2.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:8297651f5b5679c19968abefd6bb84d95fe30ef712eb1b2d9b2d31ca61267f4c", size = 4100837, upload-time = "2026-04-01T14:43:41.506Z" }, - { url = "https://files.pythonhosted.org/packages/0f/98/f3a6657ecb698c937f6c76ee564882945f29b79bad496abcba0e84659ec5/pillow-12.2.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:50d8520da2a6ce0af445fa6d648c4273c3eeefbc32d7ce049f22e8b5c3daecc2", size = 4176528, upload-time = "2026-04-01T14:43:43.773Z" }, - { url = "https://files.pythonhosted.org/packages/69/bc/8986948f05e3ea490b8442ea1c1d4d990b24a7e43d8a51b2c7d8b1dced36/pillow-12.2.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:766cef22385fa1091258ad7e6216792b156dc16d8d3fa607e7545b2b72061f1c", size = 3640401, upload-time = "2026-04-01T14:43:45.87Z" }, - { url = "https://files.pythonhosted.org/packages/34/46/6c717baadcd62bc8ed51d238d521ab651eaa74838291bda1f86fe1f864c9/pillow-12.2.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:5d2fd0fa6b5d9d1de415060363433f28da8b1526c1c129020435e186794b3795", size = 5308094, upload-time = "2026-04-01T14:43:48.438Z" }, - { url = "https://files.pythonhosted.org/packages/71/43/905a14a8b17fdb1ccb58d282454490662d2cb89a6bfec26af6d3520da5ec/pillow-12.2.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:56b25336f502b6ed02e889f4ece894a72612fe885889a6e8c4c80239ff6e5f5f", size = 4695402, upload-time = "2026-04-01T14:43:51.292Z" }, - { url = "https://files.pythonhosted.org/packages/73/dd/42107efcb777b16fa0393317eac58f5b5cf30e8392e266e76e51cff28c3d/pillow-12.2.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f1c943e96e85df3d3478f7b691f229887e143f81fedab9b20205349ab04d73ed", size = 6280005, upload-time = "2026-04-01T14:43:54.242Z" }, - { url = "https://files.pythonhosted.org/packages/a8/68/b93e09e5e8549019e61acf49f65b1a8530765a7f812c77a7461bca7e4494/pillow-12.2.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:03f6fab9219220f041c74aeaa2939ff0062bd5c364ba9ce037197f4c6d498cd9", size = 8090669, upload-time = "2026-04-01T14:43:57.335Z" }, - { url = "https://files.pythonhosted.org/packages/4b/6e/3ccb54ce8ec4ddd1accd2d89004308b7b0b21c4ac3d20fa70af4760a4330/pillow-12.2.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5cdfebd752ec52bf5bb4e35d9c64b40826bc5b40a13df7c3cda20a2c03a0f5ed", size = 6395194, upload-time = "2026-04-01T14:43:59.864Z" }, - { url = "https://files.pythonhosted.org/packages/67/ee/21d4e8536afd1a328f01b359b4d3997b291ffd35a237c877b331c1c3b71c/pillow-12.2.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:eedf4b74eda2b5a4b2b2fb4c006d6295df3bf29e459e198c90ea48e130dc75c3", size = 7082423, upload-time = "2026-04-01T14:44:02.74Z" }, - { url = "https://files.pythonhosted.org/packages/78/5f/e9f86ab0146464e8c133fe85df987ed9e77e08b29d8d35f9f9f4d6f917ba/pillow-12.2.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:00a2865911330191c0b818c59103b58a5e697cae67042366970a6b6f1b20b7f9", size = 6505667, upload-time = "2026-04-01T14:44:05.381Z" }, - { url = "https://files.pythonhosted.org/packages/ed/1e/409007f56a2fdce61584fd3acbc2bbc259857d555196cedcadc68c015c82/pillow-12.2.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1e1757442ed87f4912397c6d35a0db6a7b52592156014706f17658ff58bbf795", size = 7208580, upload-time = "2026-04-01T14:44:08.39Z" }, - { url = "https://files.pythonhosted.org/packages/23/c4/7349421080b12fb35414607b8871e9534546c128a11965fd4a7002ccfbee/pillow-12.2.0-cp313-cp313-win32.whl", hash = "sha256:144748b3af2d1b358d41286056d0003f47cb339b8c43a9ea42f5fea4d8c66b6e", size = 6375896, upload-time = "2026-04-01T14:44:11.197Z" }, - { url = "https://files.pythonhosted.org/packages/3f/82/8a3739a5e470b3c6cbb1d21d315800d8e16bff503d1f16b03a4ec3212786/pillow-12.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:390ede346628ccc626e5730107cde16c42d3836b89662a115a921f28440e6a3b", size = 7081266, upload-time = "2026-04-01T14:44:13.947Z" }, - { url = "https://files.pythonhosted.org/packages/c3/25/f968f618a062574294592f668218f8af564830ccebdd1fa6200f598e65c5/pillow-12.2.0-cp313-cp313-win_arm64.whl", hash = "sha256:8023abc91fba39036dbce14a7d6535632f99c0b857807cbbbf21ecc9f4717f06", size = 2463508, upload-time = "2026-04-01T14:44:16.312Z" }, - { url = "https://files.pythonhosted.org/packages/4d/a4/b342930964e3cb4dce5038ae34b0eab4653334995336cd486c5a8c25a00c/pillow-12.2.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:042db20a421b9bafecc4b84a8b6e444686bd9d836c7fd24542db3e7df7baad9b", size = 5309927, upload-time = "2026-04-01T14:44:18.89Z" }, - { url = "https://files.pythonhosted.org/packages/9f/de/23198e0a65a9cf06123f5435a5d95cea62a635697f8f03d134d3f3a96151/pillow-12.2.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:dd025009355c926a84a612fecf58bb315a3f6814b17ead51a8e48d3823d9087f", size = 4698624, upload-time = "2026-04-01T14:44:21.115Z" }, - { url = "https://files.pythonhosted.org/packages/01/a6/1265e977f17d93ea37aa28aa81bad4fa597933879fac2520d24e021c8da3/pillow-12.2.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:88ddbc66737e277852913bd1e07c150cc7bb124539f94c4e2df5344494e0a612", size = 6321252, upload-time = "2026-04-01T14:44:23.663Z" }, - { url = "https://files.pythonhosted.org/packages/3c/83/5982eb4a285967baa70340320be9f88e57665a387e3a53a7f0db8231a0cd/pillow-12.2.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d362d1878f00c142b7e1a16e6e5e780f02be8195123f164edf7eddd911eefe7c", size = 8126550, upload-time = "2026-04-01T14:44:26.772Z" }, - { url = "https://files.pythonhosted.org/packages/4e/48/6ffc514adce69f6050d0753b1a18fd920fce8cac87620d5a31231b04bfc5/pillow-12.2.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2c727a6d53cb0018aadd8018c2b938376af27914a68a492f59dfcaca650d5eea", size = 6433114, upload-time = "2026-04-01T14:44:29.615Z" }, - { url = "https://files.pythonhosted.org/packages/36/a3/f9a77144231fb8d40ee27107b4463e205fa4677e2ca2548e14da5cf18dce/pillow-12.2.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:efd8c21c98c5cc60653bcb311bef2ce0401642b7ce9d09e03a7da87c878289d4", size = 7115667, upload-time = "2026-04-01T14:44:32.773Z" }, - { url = "https://files.pythonhosted.org/packages/c1/fc/ac4ee3041e7d5a565e1c4fd72a113f03b6394cc72ab7089d27608f8aaccb/pillow-12.2.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9f08483a632889536b8139663db60f6724bfcb443c96f1b18855860d7d5c0fd4", size = 6538966, upload-time = "2026-04-01T14:44:35.252Z" }, - { url = "https://files.pythonhosted.org/packages/c0/a8/27fb307055087f3668f6d0a8ccb636e7431d56ed0750e07a60547b1e083e/pillow-12.2.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:dac8d77255a37e81a2efcbd1fc05f1c15ee82200e6c240d7e127e25e365c39ea", size = 7238241, upload-time = "2026-04-01T14:44:37.875Z" }, - { url = "https://files.pythonhosted.org/packages/ad/4b/926ab182c07fccae9fcb120043464e1ff1564775ec8864f21a0ebce6ac25/pillow-12.2.0-cp313-cp313t-win32.whl", hash = "sha256:ee3120ae9dff32f121610bb08e4313be87e03efeadfc6c0d18f89127e24d0c24", size = 6379592, upload-time = "2026-04-01T14:44:40.336Z" }, - { url = "https://files.pythonhosted.org/packages/c2/c4/f9e476451a098181b30050cc4c9a3556b64c02cf6497ea421ac047e89e4b/pillow-12.2.0-cp313-cp313t-win_amd64.whl", hash = "sha256:325ca0528c6788d2a6c3d40e3568639398137346c3d6e66bb61db96b96511c98", size = 7085542, upload-time = "2026-04-01T14:44:43.251Z" }, - { url = "https://files.pythonhosted.org/packages/00/a4/285f12aeacbe2d6dc36c407dfbbe9e96d4a80b0fb710a337f6d2ad978c75/pillow-12.2.0-cp313-cp313t-win_arm64.whl", hash = "sha256:2e5a76d03a6c6dcef67edabda7a52494afa4035021a79c8558e14af25313d453", size = 2465765, upload-time = "2026-04-01T14:44:45.996Z" }, - { url = "https://files.pythonhosted.org/packages/bf/98/4595daa2365416a86cb0d495248a393dfc84e96d62ad080c8546256cb9c0/pillow-12.2.0-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:3adc9215e8be0448ed6e814966ecf3d9952f0ea40eb14e89a102b87f450660d8", size = 4100848, upload-time = "2026-04-01T14:44:48.48Z" }, - { url = "https://files.pythonhosted.org/packages/0b/79/40184d464cf89f6663e18dfcf7ca21aae2491fff1a16127681bf1fa9b8cf/pillow-12.2.0-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:6a9adfc6d24b10f89588096364cc726174118c62130c817c2837c60cf08a392b", size = 4176515, upload-time = "2026-04-01T14:44:51.353Z" }, - { url = "https://files.pythonhosted.org/packages/b0/63/703f86fd4c422a9cf722833670f4f71418fb116b2853ff7da722ea43f184/pillow-12.2.0-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:6a6e67ea2e6feda684ed370f9a1c52e7a243631c025ba42149a2cc5934dec295", size = 3640159, upload-time = "2026-04-01T14:44:53.588Z" }, - { url = "https://files.pythonhosted.org/packages/71/e0/fb22f797187d0be2270f83500aab851536101b254bfa1eae10795709d283/pillow-12.2.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:2bb4a8d594eacdfc59d9e5ad972aa8afdd48d584ffd5f13a937a664c3e7db0ed", size = 5312185, upload-time = "2026-04-01T14:44:56.039Z" }, - { url = "https://files.pythonhosted.org/packages/ba/8c/1a9e46228571de18f8e28f16fabdfc20212a5d019f3e3303452b3f0a580d/pillow-12.2.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:80b2da48193b2f33ed0c32c38140f9d3186583ce7d516526d462645fd98660ae", size = 4695386, upload-time = "2026-04-01T14:44:58.663Z" }, - { url = "https://files.pythonhosted.org/packages/70/62/98f6b7f0c88b9addd0e87c217ded307b36be024d4ff8869a812b241d1345/pillow-12.2.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:22db17c68434de69d8ecfc2fe821569195c0c373b25cccb9cbdacf2c6e53c601", size = 6280384, upload-time = "2026-04-01T14:45:01.5Z" }, - { url = "https://files.pythonhosted.org/packages/5e/03/688747d2e91cfbe0e64f316cd2e8005698f76ada3130d0194664174fa5de/pillow-12.2.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7b14cc0106cd9aecda615dd6903840a058b4700fcb817687d0ee4fc8b6e389be", size = 8091599, upload-time = "2026-04-01T14:45:04.5Z" }, - { url = "https://files.pythonhosted.org/packages/f6/35/577e22b936fcdd66537329b33af0b4ccfefaeabd8aec04b266528cddb33c/pillow-12.2.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8cbeb542b2ebc6fcdacabf8aca8c1a97c9b3ad3927d46b8723f9d4f033288a0f", size = 6396021, upload-time = "2026-04-01T14:45:07.117Z" }, - { url = "https://files.pythonhosted.org/packages/11/8d/d2532ad2a603ca2b93ad9f5135732124e57811d0168155852f37fbce2458/pillow-12.2.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4bfd07bc812fbd20395212969e41931001fd59eb55a60658b0e5710872e95286", size = 7083360, upload-time = "2026-04-01T14:45:09.763Z" }, - { url = "https://files.pythonhosted.org/packages/5e/26/d325f9f56c7e039034897e7380e9cc202b1e368bfd04d4cbe6a441f02885/pillow-12.2.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:9aba9a17b623ef750a4d11b742cbafffeb48a869821252b30ee21b5e91392c50", size = 6507628, upload-time = "2026-04-01T14:45:12.378Z" }, - { url = "https://files.pythonhosted.org/packages/5f/f7/769d5632ffb0988f1c5e7660b3e731e30f7f8ec4318e94d0a5d674eb65a4/pillow-12.2.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:deede7c263feb25dba4e82ea23058a235dcc2fe1f6021025dc71f2b618e26104", size = 7209321, upload-time = "2026-04-01T14:45:15.122Z" }, - { url = "https://files.pythonhosted.org/packages/6a/7a/c253e3c645cd47f1aceea6a8bacdba9991bf45bb7dfe927f7c893e89c93c/pillow-12.2.0-cp314-cp314-win32.whl", hash = "sha256:632ff19b2778e43162304d50da0181ce24ac5bb8180122cbe1bf4673428328c7", size = 6479723, upload-time = "2026-04-01T14:45:17.797Z" }, - { url = "https://files.pythonhosted.org/packages/cd/8b/601e6566b957ca50e28725cb6c355c59c2c8609751efbecd980db44e0349/pillow-12.2.0-cp314-cp314-win_amd64.whl", hash = "sha256:4e6c62e9d237e9b65fac06857d511e90d8461a32adcc1b9065ea0c0fa3a28150", size = 7217400, upload-time = "2026-04-01T14:45:20.529Z" }, - { url = "https://files.pythonhosted.org/packages/d6/94/220e46c73065c3e2951bb91c11a1fb636c8c9ad427ac3ce7d7f3359b9b2f/pillow-12.2.0-cp314-cp314-win_arm64.whl", hash = "sha256:b1c1fbd8a5a1af3412a0810d060a78b5136ec0836c8a4ef9aa11807f2a22f4e1", size = 2554835, upload-time = "2026-04-01T14:45:23.162Z" }, - { url = "https://files.pythonhosted.org/packages/b6/ab/1b426a3974cb0e7da5c29ccff4807871d48110933a57207b5a676cccc155/pillow-12.2.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:57850958fe9c751670e49b2cecf6294acc99e562531f4bd317fa5ddee2068463", size = 5314225, upload-time = "2026-04-01T14:45:25.637Z" }, - { url = "https://files.pythonhosted.org/packages/19/1e/dce46f371be2438eecfee2a1960ee2a243bbe5e961890146d2dee1ff0f12/pillow-12.2.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:d5d38f1411c0ed9f97bcb49b7bd59b6b7c314e0e27420e34d99d844b9ce3b6f3", size = 4698541, upload-time = "2026-04-01T14:45:28.355Z" }, - { url = "https://files.pythonhosted.org/packages/55/c3/7fbecf70adb3a0c33b77a300dc52e424dc22ad8cdc06557a2e49523b703d/pillow-12.2.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5c0a9f29ca8e79f09de89293f82fc9b0270bb4af1d58bc98f540cc4aedf03166", size = 6322251, upload-time = "2026-04-01T14:45:30.924Z" }, - { url = "https://files.pythonhosted.org/packages/1c/3c/7fbc17cfb7e4fe0ef1642e0abc17fc6c94c9f7a16be41498e12e2ba60408/pillow-12.2.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1610dd6c61621ae1cf811bef44d77e149ce3f7b95afe66a4512f8c59f25d9ebe", size = 8127807, upload-time = "2026-04-01T14:45:33.908Z" }, - { url = "https://files.pythonhosted.org/packages/ff/c3/a8ae14d6defd2e448493ff512fae903b1e9bd40b72efb6ec55ce0048c8ce/pillow-12.2.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a34329707af4f73cf1782a36cd2289c0368880654a2c11f027bcee9052d35dd", size = 6433935, upload-time = "2026-04-01T14:45:36.623Z" }, - { url = "https://files.pythonhosted.org/packages/6e/32/2880fb3a074847ac159d8f902cb43278a61e85f681661e7419e6596803ed/pillow-12.2.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8e9c4f5b3c546fa3458a29ab22646c1c6c787ea8f5ef51300e5a60300736905e", size = 7116720, upload-time = "2026-04-01T14:45:39.258Z" }, - { url = "https://files.pythonhosted.org/packages/46/87/495cc9c30e0129501643f24d320076f4cc54f718341df18cc70ec94c44e1/pillow-12.2.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:fb043ee2f06b41473269765c2feae53fc2e2fbf96e5e22ca94fb5ad677856f06", size = 6540498, upload-time = "2026-04-01T14:45:41.879Z" }, - { url = "https://files.pythonhosted.org/packages/18/53/773f5edca692009d883a72211b60fdaf8871cbef075eaa9d577f0a2f989e/pillow-12.2.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:f278f034eb75b4e8a13a54a876cc4a5ab39173d2cdd93a638e1b467fc545ac43", size = 7239413, upload-time = "2026-04-01T14:45:44.705Z" }, - { url = "https://files.pythonhosted.org/packages/c9/e4/4b64a97d71b2a83158134abbb2f5bd3f8a2ea691361282f010998f339ec7/pillow-12.2.0-cp314-cp314t-win32.whl", hash = "sha256:6bb77b2dcb06b20f9f4b4a8454caa581cd4dd0643a08bacf821216a16d9c8354", size = 6482084, upload-time = "2026-04-01T14:45:47.568Z" }, - { url = "https://files.pythonhosted.org/packages/ba/13/306d275efd3a3453f72114b7431c877d10b1154014c1ebbedd067770d629/pillow-12.2.0-cp314-cp314t-win_amd64.whl", hash = "sha256:6562ace0d3fb5f20ed7290f1f929cae41b25ae29528f2af1722966a0a02e2aa1", size = 7225152, upload-time = "2026-04-01T14:45:50.032Z" }, - { url = "https://files.pythonhosted.org/packages/ff/6e/cf826fae916b8658848d7b9f38d88da6396895c676e8086fc0988073aaf8/pillow-12.2.0-cp314-cp314t-win_arm64.whl", hash = "sha256:aa88ccfe4e32d362816319ed727a004423aab09c5cea43c01a4b435643fa34eb", size = 2556579, upload-time = "2026-04-01T14:45:52.529Z" }, - { url = "https://files.pythonhosted.org/packages/4e/b7/2437044fb910f499610356d1352e3423753c98e34f915252aafecc64889f/pillow-12.2.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:0538bd5e05efec03ae613fd89c4ce0368ecd2ba239cc25b9f9be7ed426b0af1f", size = 5273969, upload-time = "2026-04-01T14:45:55.538Z" }, - { url = "https://files.pythonhosted.org/packages/f6/f4/8316e31de11b780f4ac08ef3654a75555e624a98db1056ecb2122d008d5a/pillow-12.2.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:394167b21da716608eac917c60aa9b969421b5dcbbe02ae7f013e7b85811c69d", size = 4659674, upload-time = "2026-04-01T14:45:58.093Z" }, - { url = "https://files.pythonhosted.org/packages/d4/37/664fca7201f8bb2aa1d20e2c3d5564a62e6ae5111741966c8319ca802361/pillow-12.2.0-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5d04bfa02cc2d23b497d1e90a0f927070043f6cbf303e738300532379a4b4e0f", size = 5288479, upload-time = "2026-04-01T14:46:01.141Z" }, - { url = "https://files.pythonhosted.org/packages/49/62/5b0ed78fce87346be7a5cfcfaaad91f6a1f98c26f86bdbafa2066c647ef6/pillow-12.2.0-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0c838a5125cee37e68edec915651521191cef1e6aa336b855f495766e77a366e", size = 7032230, upload-time = "2026-04-01T14:46:03.874Z" }, - { url = "https://files.pythonhosted.org/packages/c3/28/ec0fc38107fc32536908034e990c47914c57cd7c5a3ece4d8d8f7ffd7e27/pillow-12.2.0-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a6c9fa44005fa37a91ebfc95d081e8079757d2e904b27103f4f5fa6f0bf78c0", size = 5355404, upload-time = "2026-04-01T14:46:06.33Z" }, - { url = "https://files.pythonhosted.org/packages/5e/8b/51b0eddcfa2180d60e41f06bd6d0a62202b20b59c68f5a132e615b75aecf/pillow-12.2.0-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:25373b66e0dd5905ed63fa3cae13c82fbddf3079f2c8bf15c6fb6a35586324c1", size = 6002215, upload-time = "2026-04-01T14:46:08.83Z" }, - { url = "https://files.pythonhosted.org/packages/bc/60/5382c03e1970de634027cee8e1b7d39776b778b81812aaf45b694dfe9e28/pillow-12.2.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:bfa9c230d2fe991bed5318a5f119bd6780cda2915cca595393649fc118ab895e", size = 7080946, upload-time = "2026-04-01T14:46:11.734Z" }, -] - -[[package]] -name = "pygments" -version = "2.20.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, -] - -[[package]] -name = "pyyaml" -version = "6.0.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f4/a0/39350dd17dd6d6c6507025c0e53aef67a9293a6d37d3511f23ea510d5800/pyyaml-6.0.3-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:214ed4befebe12df36bcc8bc2b64b396ca31be9304b8f59e25c11cf94a4c033b", size = 184227, upload-time = "2025-09-25T21:31:46.04Z" }, - { url = "https://files.pythonhosted.org/packages/05/14/52d505b5c59ce73244f59c7a50ecf47093ce4765f116cdb98286a71eeca2/pyyaml-6.0.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:02ea2dfa234451bbb8772601d7b8e426c2bfa197136796224e50e35a78777956", size = 174019, upload-time = "2025-09-25T21:31:47.706Z" }, - { url = "https://files.pythonhosted.org/packages/43/f7/0e6a5ae5599c838c696adb4e6330a59f463265bfa1e116cfd1fbb0abaaae/pyyaml-6.0.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b30236e45cf30d2b8e7b3e85881719e98507abed1011bf463a8fa23e9c3e98a8", size = 740646, upload-time = "2025-09-25T21:31:49.21Z" }, - { url = "https://files.pythonhosted.org/packages/2f/3a/61b9db1d28f00f8fd0ae760459a5c4bf1b941baf714e207b6eb0657d2578/pyyaml-6.0.3-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:66291b10affd76d76f54fad28e22e51719ef9ba22b29e1d7d03d6777a9174198", size = 840793, upload-time = "2025-09-25T21:31:50.735Z" }, - { url = "https://files.pythonhosted.org/packages/7a/1e/7acc4f0e74c4b3d9531e24739e0ab832a5edf40e64fbae1a9c01941cabd7/pyyaml-6.0.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9c7708761fccb9397fe64bbc0395abcae8c4bf7b0eac081e12b809bf47700d0b", size = 770293, upload-time = "2025-09-25T21:31:51.828Z" }, - { url = "https://files.pythonhosted.org/packages/8b/ef/abd085f06853af0cd59fa5f913d61a8eab65d7639ff2a658d18a25d6a89d/pyyaml-6.0.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:418cf3f2111bc80e0933b2cd8cd04f286338bb88bdc7bc8e6dd775ebde60b5e0", size = 732872, upload-time = "2025-09-25T21:31:53.282Z" }, - { url = "https://files.pythonhosted.org/packages/1f/15/2bc9c8faf6450a8b3c9fc5448ed869c599c0a74ba2669772b1f3a0040180/pyyaml-6.0.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:5e0b74767e5f8c593e8c9b5912019159ed0533c70051e9cce3e8b6aa699fcd69", size = 758828, upload-time = "2025-09-25T21:31:54.807Z" }, - { url = "https://files.pythonhosted.org/packages/a3/00/531e92e88c00f4333ce359e50c19b8d1de9fe8d581b1534e35ccfbc5f393/pyyaml-6.0.3-cp310-cp310-win32.whl", hash = "sha256:28c8d926f98f432f88adc23edf2e6d4921ac26fb084b028c733d01868d19007e", size = 142415, upload-time = "2025-09-25T21:31:55.885Z" }, - { url = "https://files.pythonhosted.org/packages/2a/fa/926c003379b19fca39dd4634818b00dec6c62d87faf628d1394e137354d4/pyyaml-6.0.3-cp310-cp310-win_amd64.whl", hash = "sha256:bdb2c67c6c1390b63c6ff89f210c8fd09d9a1217a465701eac7316313c915e4c", size = 158561, upload-time = "2025-09-25T21:31:57.406Z" }, - { url = "https://files.pythonhosted.org/packages/6d/16/a95b6757765b7b031c9374925bb718d55e0a9ba8a1b6a12d25962ea44347/pyyaml-6.0.3-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:44edc647873928551a01e7a563d7452ccdebee747728c1080d881d68af7b997e", size = 185826, upload-time = "2025-09-25T21:31:58.655Z" }, - { url = "https://files.pythonhosted.org/packages/16/19/13de8e4377ed53079ee996e1ab0a9c33ec2faf808a4647b7b4c0d46dd239/pyyaml-6.0.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:652cb6edd41e718550aad172851962662ff2681490a8a711af6a4d288dd96824", size = 175577, upload-time = "2025-09-25T21:32:00.088Z" }, - { url = "https://files.pythonhosted.org/packages/0c/62/d2eb46264d4b157dae1275b573017abec435397aa59cbcdab6fc978a8af4/pyyaml-6.0.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:10892704fc220243f5305762e276552a0395f7beb4dbf9b14ec8fd43b57f126c", size = 775556, upload-time = "2025-09-25T21:32:01.31Z" }, - { url = "https://files.pythonhosted.org/packages/10/cb/16c3f2cf3266edd25aaa00d6c4350381c8b012ed6f5276675b9eba8d9ff4/pyyaml-6.0.3-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:850774a7879607d3a6f50d36d04f00ee69e7fc816450e5f7e58d7f17f1ae5c00", size = 882114, upload-time = "2025-09-25T21:32:03.376Z" }, - { url = "https://files.pythonhosted.org/packages/71/60/917329f640924b18ff085ab889a11c763e0b573da888e8404ff486657602/pyyaml-6.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b8bb0864c5a28024fac8a632c443c87c5aa6f215c0b126c449ae1a150412f31d", size = 806638, upload-time = "2025-09-25T21:32:04.553Z" }, - { url = "https://files.pythonhosted.org/packages/dd/6f/529b0f316a9fd167281a6c3826b5583e6192dba792dd55e3203d3f8e655a/pyyaml-6.0.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:1d37d57ad971609cf3c53ba6a7e365e40660e3be0e5175fa9f2365a379d6095a", size = 767463, upload-time = "2025-09-25T21:32:06.152Z" }, - { url = "https://files.pythonhosted.org/packages/f2/6a/b627b4e0c1dd03718543519ffb2f1deea4a1e6d42fbab8021936a4d22589/pyyaml-6.0.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:37503bfbfc9d2c40b344d06b2199cf0e96e97957ab1c1b546fd4f87e53e5d3e4", size = 794986, upload-time = "2025-09-25T21:32:07.367Z" }, - { url = "https://files.pythonhosted.org/packages/45/91/47a6e1c42d9ee337c4839208f30d9f09caa9f720ec7582917b264defc875/pyyaml-6.0.3-cp311-cp311-win32.whl", hash = "sha256:8098f252adfa6c80ab48096053f512f2321f0b998f98150cea9bd23d83e1467b", size = 142543, upload-time = "2025-09-25T21:32:08.95Z" }, - { url = "https://files.pythonhosted.org/packages/da/e3/ea007450a105ae919a72393cb06f122f288ef60bba2dc64b26e2646fa315/pyyaml-6.0.3-cp311-cp311-win_amd64.whl", hash = "sha256:9f3bfb4965eb874431221a3ff3fdcddc7e74e3b07799e0e84ca4a0f867d449bf", size = 158763, upload-time = "2025-09-25T21:32:09.96Z" }, - { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, - { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, - { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, - { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, - { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, - { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, - { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, - { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, - { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, - { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, - { url = "https://files.pythonhosted.org/packages/d1/11/0fd08f8192109f7169db964b5707a2f1e8b745d4e239b784a5a1dd80d1db/pyyaml-6.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8da9669d359f02c0b91ccc01cac4a67f16afec0dac22c2ad09f46bee0697eba8", size = 181669, upload-time = "2025-09-25T21:32:23.673Z" }, - { url = "https://files.pythonhosted.org/packages/b1/16/95309993f1d3748cd644e02e38b75d50cbc0d9561d21f390a76242ce073f/pyyaml-6.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2283a07e2c21a2aa78d9c4442724ec1eb15f5e42a723b99cb3d822d48f5f7ad1", size = 173252, upload-time = "2025-09-25T21:32:25.149Z" }, - { url = "https://files.pythonhosted.org/packages/50/31/b20f376d3f810b9b2371e72ef5adb33879b25edb7a6d072cb7ca0c486398/pyyaml-6.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c", size = 767081, upload-time = "2025-09-25T21:32:26.575Z" }, - { url = "https://files.pythonhosted.org/packages/49/1e/a55ca81e949270d5d4432fbbd19dfea5321eda7c41a849d443dc92fd1ff7/pyyaml-6.0.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a33284e20b78bd4a18c8c2282d549d10bc8408a2a7ff57653c0cf0b9be0afce5", size = 841159, upload-time = "2025-09-25T21:32:27.727Z" }, - { url = "https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6", size = 801626, upload-time = "2025-09-25T21:32:28.878Z" }, - { url = "https://files.pythonhosted.org/packages/f9/11/ba845c23988798f40e52ba45f34849aa8a1f2d4af4b798588010792ebad6/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6", size = 753613, upload-time = "2025-09-25T21:32:30.178Z" }, - { url = "https://files.pythonhosted.org/packages/3d/e0/7966e1a7bfc0a45bf0a7fb6b98ea03fc9b8d84fa7f2229e9659680b69ee3/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be", size = 794115, upload-time = "2025-09-25T21:32:31.353Z" }, - { url = "https://files.pythonhosted.org/packages/de/94/980b50a6531b3019e45ddeada0626d45fa85cbe22300844a7983285bed3b/pyyaml-6.0.3-cp313-cp313-win32.whl", hash = "sha256:d0eae10f8159e8fdad514efdc92d74fd8d682c933a6dd088030f3834bc8e6b26", size = 137427, upload-time = "2025-09-25T21:32:32.58Z" }, - { url = "https://files.pythonhosted.org/packages/97/c9/39d5b874e8b28845e4ec2202b5da735d0199dbe5b8fb85f91398814a9a46/pyyaml-6.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:79005a0d97d5ddabfeeea4cf676af11e647e41d81c9a7722a193022accdb6b7c", size = 154090, upload-time = "2025-09-25T21:32:33.659Z" }, - { url = "https://files.pythonhosted.org/packages/73/e8/2bdf3ca2090f68bb3d75b44da7bbc71843b19c9f2b9cb9b0f4ab7a5a4329/pyyaml-6.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb", size = 140246, upload-time = "2025-09-25T21:32:34.663Z" }, - { url = "https://files.pythonhosted.org/packages/9d/8c/f4bd7f6465179953d3ac9bc44ac1a8a3e6122cf8ada906b4f96c60172d43/pyyaml-6.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:8d1fab6bb153a416f9aeb4b8763bc0f22a5586065f86f7664fc23339fc1c1fac", size = 181814, upload-time = "2025-09-25T21:32:35.712Z" }, - { url = "https://files.pythonhosted.org/packages/bd/9c/4d95bb87eb2063d20db7b60faa3840c1b18025517ae857371c4dd55a6b3a/pyyaml-6.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:34d5fcd24b8445fadc33f9cf348c1047101756fd760b4dacb5c3e99755703310", size = 173809, upload-time = "2025-09-25T21:32:36.789Z" }, - { url = "https://files.pythonhosted.org/packages/92/b5/47e807c2623074914e29dabd16cbbdd4bf5e9b2db9f8090fa64411fc5382/pyyaml-6.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:501a031947e3a9025ed4405a168e6ef5ae3126c59f90ce0cd6f2bfc477be31b7", size = 766454, upload-time = "2025-09-25T21:32:37.966Z" }, - { url = "https://files.pythonhosted.org/packages/02/9e/e5e9b168be58564121efb3de6859c452fccde0ab093d8438905899a3a483/pyyaml-6.0.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b3bc83488de33889877a0f2543ade9f70c67d66d9ebb4ac959502e12de895788", size = 836355, upload-time = "2025-09-25T21:32:39.178Z" }, - { url = "https://files.pythonhosted.org/packages/88/f9/16491d7ed2a919954993e48aa941b200f38040928474c9e85ea9e64222c3/pyyaml-6.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c458b6d084f9b935061bc36216e8a69a7e293a2f1e68bf956dcd9e6cbcd143f5", size = 794175, upload-time = "2025-09-25T21:32:40.865Z" }, - { url = "https://files.pythonhosted.org/packages/dd/3f/5989debef34dc6397317802b527dbbafb2b4760878a53d4166579111411e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7c6610def4f163542a622a73fb39f534f8c101d690126992300bf3207eab9764", size = 755228, upload-time = "2025-09-25T21:32:42.084Z" }, - { url = "https://files.pythonhosted.org/packages/d7/ce/af88a49043cd2e265be63d083fc75b27b6ed062f5f9fd6cdc223ad62f03e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5190d403f121660ce8d1d2c1bb2ef1bd05b5f68533fc5c2ea899bd15f4399b35", size = 789194, upload-time = "2025-09-25T21:32:43.362Z" }, - { url = "https://files.pythonhosted.org/packages/23/20/bb6982b26a40bb43951265ba29d4c246ef0ff59c9fdcdf0ed04e0687de4d/pyyaml-6.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:4a2e8cebe2ff6ab7d1050ecd59c25d4c8bd7e6f400f5f82b96557ac0abafd0ac", size = 156429, upload-time = "2025-09-25T21:32:57.844Z" }, - { url = "https://files.pythonhosted.org/packages/f4/f4/a4541072bb9422c8a883ab55255f918fa378ecf083f5b85e87fc2b4eda1b/pyyaml-6.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:93dda82c9c22deb0a405ea4dc5f2d0cda384168e466364dec6255b293923b2f3", size = 143912, upload-time = "2025-09-25T21:32:59.247Z" }, - { url = "https://files.pythonhosted.org/packages/7c/f9/07dd09ae774e4616edf6cda684ee78f97777bdd15847253637a6f052a62f/pyyaml-6.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:02893d100e99e03eda1c8fd5c441d8c60103fd175728e23e431db1b589cf5ab3", size = 189108, upload-time = "2025-09-25T21:32:44.377Z" }, - { url = "https://files.pythonhosted.org/packages/4e/78/8d08c9fb7ce09ad8c38ad533c1191cf27f7ae1effe5bb9400a46d9437fcf/pyyaml-6.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c1ff362665ae507275af2853520967820d9124984e0f7466736aea23d8611fba", size = 183641, upload-time = "2025-09-25T21:32:45.407Z" }, - { url = "https://files.pythonhosted.org/packages/7b/5b/3babb19104a46945cf816d047db2788bcaf8c94527a805610b0289a01c6b/pyyaml-6.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6adc77889b628398debc7b65c073bcb99c4a0237b248cacaf3fe8a557563ef6c", size = 831901, upload-time = "2025-09-25T21:32:48.83Z" }, - { url = "https://files.pythonhosted.org/packages/8b/cc/dff0684d8dc44da4d22a13f35f073d558c268780ce3c6ba1b87055bb0b87/pyyaml-6.0.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a80cb027f6b349846a3bf6d73b5e95e782175e52f22108cfa17876aaeff93702", size = 861132, upload-time = "2025-09-25T21:32:50.149Z" }, - { url = "https://files.pythonhosted.org/packages/b1/5e/f77dc6b9036943e285ba76b49e118d9ea929885becb0a29ba8a7c75e29fe/pyyaml-6.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00c4bdeba853cc34e7dd471f16b4114f4162dc03e6b7afcc2128711f0eca823c", size = 839261, upload-time = "2025-09-25T21:32:51.808Z" }, - { url = "https://files.pythonhosted.org/packages/ce/88/a9db1376aa2a228197c58b37302f284b5617f56a5d959fd1763fb1675ce6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:66e1674c3ef6f541c35191caae2d429b967b99e02040f5ba928632d9a7f0f065", size = 805272, upload-time = "2025-09-25T21:32:52.941Z" }, - { url = "https://files.pythonhosted.org/packages/da/92/1446574745d74df0c92e6aa4a7b0b3130706a4142b2d1a5869f2eaa423c6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:16249ee61e95f858e83976573de0f5b2893b3677ba71c9dd36b9cf8be9ac6d65", size = 829923, upload-time = "2025-09-25T21:32:54.537Z" }, - { url = "https://files.pythonhosted.org/packages/f0/7a/1c7270340330e575b92f397352af856a8c06f230aa3e76f86b39d01b416a/pyyaml-6.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4ad1906908f2f5ae4e5a8ddfce73c320c2a1429ec52eafd27138b7f1cbe341c9", size = 174062, upload-time = "2025-09-25T21:32:55.767Z" }, - { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, -] - -[[package]] -name = "regex" -version = "2026.5.9" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/dc/0e/49aee608ad09480e7fd276898c99ec6192985fa331abe4eb3a986094490b/regex-2026.5.9.tar.gz", hash = "sha256:a8234aa23ec39894bfe4a3f1b85616a7032481964a13ac6fc9f10de4f6fca270", size = 416074, upload-time = "2026-05-09T23:15:19.37Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fe/ed/0ad2c8edf634918eb4484365d3819fa7bd7f58daf807fe7fb21812c316e5/regex-2026.5.9-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:a9e1328e17c84c1a5d22ec9f785ecef4a967fab9a42b6a8dc3bcbebd0a0c9e44", size = 489438, upload-time = "2026-05-09T23:11:29.374Z" }, - { url = "https://files.pythonhosted.org/packages/89/a9/4ed972ad263963b860b7c3e86e0e1bcc791def47b43b8c8efe57e710f139/regex-2026.5.9-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:bfe1ce50cbfb569d74e1e4337da6468961f31dbea55fd85aa5de59c0947a805a", size = 291270, upload-time = "2026-05-09T23:11:33.254Z" }, - { url = "https://files.pythonhosted.org/packages/16/81/075930d9fa28c4ea1f53398dd015ee7c882f623539759113cda1257f4b82/regex-2026.5.9-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:15ee42209947f4ca045412eae98416317238163618ace2a8e54f99586a466733", size = 289198, upload-time = "2026-05-09T23:11:35.769Z" }, - { url = "https://files.pythonhosted.org/packages/d4/c8/5cdfbf0b5dc6599e1b6131eff43262e5275d4ec3469ce10216061659aadb/regex-2026.5.9-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b4bb445ff3f725f59df8f6014edb547ee928ec7023a774f6a39a3f953038cbb2", size = 784765, upload-time = "2026-05-09T23:11:37.689Z" }, - { url = "https://files.pythonhosted.org/packages/cd/ca/ae5fd6edc59b7f84b904b31d6ec39a860cbcecd10f64bd5a062ca83a4864/regex-2026.5.9-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:446ddd671e43ab535810c4b21cff7104945c701d4a14d1e6d1cd6f4e445a8bea", size = 852115, upload-time = "2026-05-09T23:11:39.973Z" }, - { url = "https://files.pythonhosted.org/packages/f6/ce/a91cf555afb51f3b74a182e24ba073b91ea7bb64592fc4b315c111bb19fd/regex-2026.5.9-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:7b92817338591505f282cf3864c145244b1edcf5381d237038df955001091538", size = 899503, upload-time = "2026-05-09T23:11:42.48Z" }, - { url = "https://files.pythonhosted.org/packages/55/7f/725a0a2b245a4cf0c4bab29d0e97c74285d94136a65d1b55a6459a583502/regex-2026.5.9-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6b8a143aca6c39b446ea8092cde25cc8fe9304d4f5fecfbc1a9dbb0282703c2", size = 794093, upload-time = "2026-05-09T23:11:44.681Z" }, - { url = "https://files.pythonhosted.org/packages/e3/2a/996efbd59ce6b5d4a09e3af6180ceb62af171f4a9a6fb557d2f0ae0d462b/regex-2026.5.9-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0f03aa6898aaaac4592479821df16e68e8d0e29e903e65d8f2dfb2f19028a989", size = 786234, upload-time = "2026-05-09T23:11:46.882Z" }, - { url = "https://files.pythonhosted.org/packages/4b/0a/8731e8b8806174c9cdd5903f80a14990331c1f42fc4209b540952e9e010d/regex-2026.5.9-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:ed457d8e98ae812ed7732bef7bf78de78e834eae0372a74e23ca90ef21d910f9", size = 769895, upload-time = "2026-05-09T23:11:49.324Z" }, - { url = "https://files.pythonhosted.org/packages/9a/0b/932473194bd563f342a412ae2ffbbd6da608306a2bc4e99249a41c2b0b92/regex-2026.5.9-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:71b61c5bfe1c806332defc42ad6c780b3c55f661986d7f40283a3a88274b4c00", size = 774991, upload-time = "2026-05-09T23:11:51.261Z" }, - { url = "https://files.pythonhosted.org/packages/98/80/9523d196010031df25f7177ee0a467efbee436324038e5d99def17a57515/regex-2026.5.9-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:3b1e39888c5e0c7d92cea4fc777396c4a90363b05de75d02eb459a4752200808", size = 848790, upload-time = "2026-05-09T23:11:53.232Z" }, - { url = "https://files.pythonhosted.org/packages/3c/07/56987b35e89edf47e4a38cf2845aeee476bfa688a6bdbd3e820cda461dc1/regex-2026.5.9-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:6ba42b2e7e7f46cf68cc6a5ca36fa07959f9bbd9c6bdcc47b6ee76549a590248", size = 757679, upload-time = "2026-05-09T23:11:55.82Z" }, - { url = "https://files.pythonhosted.org/packages/04/2a/ff713fff0c566507c06a4ce2dc0ae8e7eeebc88811a95fc81cf1e7d534dd/regex-2026.5.9-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:c010eb8caca74bdb40c07498d7ece26b4428fd3f04aa8a72c9ac6f79e8faaac6", size = 837116, upload-time = "2026-05-09T23:11:57.934Z" }, - { url = "https://files.pythonhosted.org/packages/77/90/df6d982b03e3614785c6937ba51b57f6733d97d2ee1c9bc7531dbfab3a54/regex-2026.5.9-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:a6a563446a41adc451393dc6b8e6ad87979efaee3c8738690a8d1b08ebead1b4", size = 782081, upload-time = "2026-05-09T23:11:59.607Z" }, - { url = "https://files.pythonhosted.org/packages/c7/8a/4e88a5f7c3e98489aac4dd23142723d907b2a595b4a6abcbacabefeded09/regex-2026.5.9-cp310-cp310-win32.whl", hash = "sha256:954cc214c04663ee6d266fc61739cad83054683048de65c5bd1d640ad28098ac", size = 266247, upload-time = "2026-05-09T23:12:01.116Z" }, - { url = "https://files.pythonhosted.org/packages/6a/40/4b224cb0582b2dca1786726e6cdabe26abbf757d7f6718332f186da155d2/regex-2026.5.9-cp310-cp310-win_amd64.whl", hash = "sha256:b310768746dd314ea6e2ff4cc89ef215426813396ff4e94ee8e6f7096c8b6e03", size = 278416, upload-time = "2026-05-09T23:12:03.2Z" }, - { url = "https://files.pythonhosted.org/packages/12/4d/014fbe803204cab0947ee428f09f658a29632053dde1d3c6176bb4f0fd4c/regex-2026.5.9-cp310-cp310-win_arm64.whl", hash = "sha256:19c16ceb4a267a8789e25733e583983eeab9f0f8664e66b0bd1c5d21f14c2d4b", size = 270413, upload-time = "2026-05-09T23:12:04.649Z" }, - { url = "https://files.pythonhosted.org/packages/c2/dc/c1f2df4027e82fc54b5a473e4b250f5139faca49a0fbe29a48668d228f34/regex-2026.5.9-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ccf5249114cc3e772ecdd88a98a86eca0fd74c61ce32a94743758c083fc05d48", size = 489445, upload-time = "2026-05-09T23:12:06.111Z" }, - { url = "https://files.pythonhosted.org/packages/03/d2/59f01110660081cce9c0bc30ebd0b5ee250dacf658e3248ed92f01e0e8ee/regex-2026.5.9-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:46f1326ca6e65b0879d23ca302c0f2415aad42ff0309b9c818e7949fe19a41d8", size = 291271, upload-time = "2026-05-09T23:12:07.731Z" }, - { url = "https://files.pythonhosted.org/packages/58/b6/14b2c84ff90ddb370c81d27503f4a0fcf071496416f4855f6cc8c5d81c35/regex-2026.5.9-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:ef31cbfe458e21c6122ba8150ff060e0c7789ed0d26eb423f25472584920b555", size = 289212, upload-time = "2026-05-09T23:12:09.266Z" }, - { url = "https://files.pythonhosted.org/packages/03/d0/4db86529117320de0c84afd90e70bb47434625875e34fcef9d8c127c5b16/regex-2026.5.9-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:992604d02e6d9c6d786c24a706a71ecffe1020fc1ef264044474cd81fa2c3919", size = 792310, upload-time = "2026-05-09T23:12:11.416Z" }, - { url = "https://files.pythonhosted.org/packages/07/78/fe4800cd322f862ecffd2d553409b20d80650e5ed71b9d178f853d020b82/regex-2026.5.9-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c9411dd64ca95477225734a93dfc8583b51916b8d5942f99d6cac21e09965451", size = 861721, upload-time = "2026-05-09T23:12:13.681Z" }, - { url = "https://files.pythonhosted.org/packages/b5/d0/b3618a895dd8feb897c61bb2954edd265e1767d82a01d53065d5871127a3/regex-2026.5.9-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3dd4a3ff360dfb836fecdb93a4598f9d6e2ac81e3e397125145c6221bf58cf4c", size = 906460, upload-time = "2026-05-09T23:12:15.443Z" }, - { url = "https://files.pythonhosted.org/packages/33/6f/1481597e859ef19508b345eec4afd1416ed6e6b459c75a64026ef193aecf/regex-2026.5.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2a661a7d270a61f7cf460caee8b9fa2d5ef9e5c681234bcb9e0fe14f488e7dfc", size = 799843, upload-time = "2026-05-09T23:12:16.892Z" }, - { url = "https://files.pythonhosted.org/packages/73/59/955734c803f59108deccba3597ae440c76b62a652733c0006e6243758420/regex-2026.5.9-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f079e50a0d3cc3cd5091fa9ff45869a2e6b2cd35895731edafb0327901a8d86d", size = 773610, upload-time = "2026-05-09T23:12:19.127Z" }, - { url = "https://files.pythonhosted.org/packages/68/8f/70c04a236d651c81881dac42ef8538bddda6121434509d0a22d9e601503b/regex-2026.5.9-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:4ebe8f0b5ec5a5024dc4a4c59f444c4e9afc5f2abdbb8962065b75d27fb971f9", size = 781645, upload-time = "2026-05-09T23:12:20.806Z" }, - { url = "https://files.pythonhosted.org/packages/1d/96/05c7434d88185e5d27fe54aeb74df86bd77cd79f52f0b4eae54faa8fea70/regex-2026.5.9-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:97cf3bc1b7d7d2306772ec07366c80d9df00ff79e79cea32898883a646d2fae2", size = 854473, upload-time = "2026-05-09T23:12:22.465Z" }, - { url = "https://files.pythonhosted.org/packages/4e/c1/6e3d8202d981f3117004bf341ee74893ba4ba8a9fbaf4b94615846550a08/regex-2026.5.9-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:0f9eede6a5cbdc02d4978090186390936e1776a7d1359b21e41014c609880bcf", size = 763311, upload-time = "2026-05-09T23:12:24.351Z" }, - { url = "https://files.pythonhosted.org/packages/93/c7/e7737f1526b3fb32bd4c337fd6c71c3ebb5c8296fc34d11197e0955d2e35/regex-2026.5.9-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:01f0f5f55f4b64dacec85dc116d3c05fd23ad3ff037bbc73a2085775953c2611", size = 844593, upload-time = "2026-05-09T23:12:26.341Z" }, - { url = "https://files.pythonhosted.org/packages/a5/27/0daffb1a535bb39f422c3d200f4ab023c71110ad66a32b366bee708baba0/regex-2026.5.9-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1268eddd8486dc561d08eee1156e40aa3a8fe10f4bdec8fa653b455fcbffd12c", size = 789167, upload-time = "2026-05-09T23:12:27.975Z" }, - { url = "https://files.pythonhosted.org/packages/ce/fc/294fe4fac4f2ed67207b17471815870c1c45b3a489e08e0ac96daea16ef6/regex-2026.5.9-cp311-cp311-win32.whl", hash = "sha256:8676474c07469d6f33dd1085ca2cd45f65785f32518f2b20e36d9953ca07f994", size = 266249, upload-time = "2026-05-09T23:12:30.141Z" }, - { url = "https://files.pythonhosted.org/packages/d0/b0/8dce459f6245bcf8f6e9f23ac9569f1a0f15c131cc0745e82b43226204cf/regex-2026.5.9-cp311-cp311-win_amd64.whl", hash = "sha256:246de9d60aa3f8538b519834dd95cbf276ea263d6a7bd5a3666dc3fa0230505b", size = 278423, upload-time = "2026-05-09T23:12:31.676Z" }, - { url = "https://files.pythonhosted.org/packages/db/8d/f9aeff6ad63a3ef720386f2907e6d34a35a510a6e498ebad28b0fb3f6ab6/regex-2026.5.9-cp311-cp311-win_arm64.whl", hash = "sha256:d726ca3f0d76969bf1e8e477d160d3d666bbf999f6860bd314889e5345782046", size = 270420, upload-time = "2026-05-09T23:12:33.194Z" }, - { url = "https://files.pythonhosted.org/packages/50/9b/6550044bc44e17c84d312c031c2ec42fbdb6a4ec4e29093be3a172d08772/regex-2026.5.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57eeeb05db7979413dec5438f2db21d7ecbba787cde7a711df1a6f6df672aa06", size = 490451, upload-time = "2026-05-09T23:12:34.72Z" }, - { url = "https://files.pythonhosted.org/packages/1e/95/fc7ba4303b5a0f92446a12ee6778ef2c6c799233f5060042a31bf390cfe9/regex-2026.5.9-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:398c521292f4c7fb807001dcd54694d3a1fcafc179a36ad9cc56f98df85930b6", size = 292112, upload-time = "2026-05-09T23:12:36.285Z" }, - { url = "https://files.pythonhosted.org/packages/54/4b/ee27938d1b2c443e89a9a10e00d2d19aa5ee300cd3d61140644e93bb083e/regex-2026.5.9-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f7a7c26137296beba7784de6eba69c6a93a63ccebc385e4962fe67e267a91225", size = 289599, upload-time = "2026-05-09T23:12:38.089Z" }, - { url = "https://files.pythonhosted.org/packages/d8/dd/ba103dc19614e25f3880800ca67ce093d6e21b325d72b8383c7bf906e9fa/regex-2026.5.9-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6441cc660d76107934a09c22167200839a0e89604a6297f78a974e66e931d2c0", size = 796732, upload-time = "2026-05-09T23:12:40.062Z" }, - { url = "https://files.pythonhosted.org/packages/cf/e7/f035b4fd858b050b0080bf302968dc0f59ba34e391872d54936758e6844e/regex-2026.5.9-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:91328f1c23d47595ca3ef0a7557fa129c5a23404b775c770697d2f35b33e0107", size = 865440, upload-time = "2026-05-09T23:12:42.059Z" }, - { url = "https://files.pythonhosted.org/packages/0a/51/8cd301ecc899aea28124357f729f4272f44de7806fc7ca02490bfbe253e8/regex-2026.5.9-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:93a7860539414dddaefba2b40f8771765ae17949d4c7182b876ce429e11a8309", size = 912329, upload-time = "2026-05-09T23:12:44.373Z" }, - { url = "https://files.pythonhosted.org/packages/cc/1e/3fbe2fa1e8cebd62f3bb7d3321cff1640aca2e240b51d9bd624aad949260/regex-2026.5.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dd2810d22146b6d838acc5ec15602cb6b47920aa4e33015df3868eedfd20bab8", size = 801239, upload-time = "2026-05-09T23:12:46.268Z" }, - { url = "https://files.pythonhosted.org/packages/17/2f/6f6008682bf2cf98040a0d3153a8e557b6ab728d7713d045cee4ce544ab8/regex-2026.5.9-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:daff2bdbaf1d23e52fdff7c0b7bc2048b68f978df6a4d107ac981f94caef2e66", size = 777054, upload-time = "2026-05-09T23:12:48.051Z" }, - { url = "https://files.pythonhosted.org/packages/19/2b/eee0d20a6842ba04df4b8847a920b57ef56853f14ef85405473e586b605a/regex-2026.5.9-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4eeb011098fcb77af513dcef521a3dbecbf8849b1e38940759d293b7a93f5026", size = 785098, upload-time = "2026-05-09T23:12:49.851Z" }, - { url = "https://files.pythonhosted.org/packages/4a/98/6fc1e6410feefb92159edaed5041992bfe390e8d26c721865434acbca558/regex-2026.5.9-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:ea9c8ecfa1b73c73b626534d6626e5340d429630943672b8480724f44e84b962", size = 860095, upload-time = "2026-05-09T23:12:51.666Z" }, - { url = "https://files.pythonhosted.org/packages/18/a3/bd855e0f2cb1a978ecf6fa6bb69632dd9c3f6ea3b81cde62fde14c9daec7/regex-2026.5.9-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:cd2846168eb9ee3c513902bc8225409cb1caab31d04728b145171fa1625d9621", size = 765762, upload-time = "2026-05-09T23:12:53.413Z" }, - { url = "https://files.pythonhosted.org/packages/dc/66/0ae8c092e60b14c79d24f8e0b7f0aea5bfbffdcab00b5483d13404d3c3a5/regex-2026.5.9-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:39617fb0cde9c0e6306dc70e3bfc096f3da793219879f7ae7aa341a69fbdcf6d", size = 852100, upload-time = "2026-05-09T23:12:55.256Z" }, - { url = "https://files.pythonhosted.org/packages/21/de/8dfde60fc1b21c946a893ba273403b72617edb261370cb1087099a83f088/regex-2026.5.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fd03c4f0e33280d15cae17159b899245d6b7c53d21def19b263b39655061f5ce", size = 789479, upload-time = "2026-05-09T23:12:57.573Z" }, - { url = "https://files.pythonhosted.org/packages/c3/1c/bdcc98f9a4af4fdd166c74941174619ccff4726d3ce32faa8e9a2ecd38dd/regex-2026.5.9-cp312-cp312-win32.whl", hash = "sha256:164eba9b755ea6f244b0d881196fbc1fac09714e9782c9e2732b813142033c8e", size = 266699, upload-time = "2026-05-09T23:12:59.14Z" }, - { url = "https://files.pythonhosted.org/packages/78/87/240d36864f9e48ace85f72e79ced97ceb7f27ce87739a947dcb834b4e6bc/regex-2026.5.9-cp312-cp312-win_amd64.whl", hash = "sha256:86f40a5d6444db30a125c9c9177e6b25dad981cbc37451fd838f145e6edac92e", size = 277783, upload-time = "2026-05-09T23:13:00.789Z" }, - { url = "https://files.pythonhosted.org/packages/4f/b5/7b30f312b0669dff5beebe5b0989dc2d1a312b1a44fab852199c387a5b96/regex-2026.5.9-cp312-cp312-win_arm64.whl", hash = "sha256:96f5f58b54a063d7ea9dca08e1cf57bfe10499c4d579ee672da284f57f5f0070", size = 270513, upload-time = "2026-05-09T23:13:02.426Z" }, - { url = "https://files.pythonhosted.org/packages/aa/da/797e91ecec6f84135da778ddce78c20e0af5d2a15c26f87a81bc3eadb6db/regex-2026.5.9-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:d626b84406444b165fc0ba981604edea39f0588ff1f92baa23fe50799ea9afdb", size = 490303, upload-time = "2026-05-09T23:13:04.382Z" }, - { url = "https://files.pythonhosted.org/packages/44/da/bf30abaaa737b58f4a4b8c4a03659e02fd92092c822e0197ed9e0daab917/regex-2026.5.9-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d7bdc0ab8f3dd7e1b4f9ab88634e13374669db86bb3c72e8292f07ae313f539f", size = 292019, upload-time = "2026-05-09T23:13:06.022Z" }, - { url = "https://files.pythonhosted.org/packages/2d/e7/d0eaf5713828417b9e5648cf81fa9bacd4961f6ab98c380c2034f8716e35/regex-2026.5.9-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a8820737949116ffff55fe18f9fc644530063ba6ebfcb8314239416e78f1347c", size = 289468, upload-time = "2026-05-09T23:13:08.214Z" }, - { url = "https://files.pythonhosted.org/packages/d3/9b/b3fdd62b003baa1a9b593cd8c8699c9651c2e80cc21a5c715707983c42d7/regex-2026.5.9-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aa0fbdbac82cb3e4450d0ccde7d7a35607f4cb2dd9fba4b8b69bfaf8c9fa6aed", size = 796749, upload-time = "2026-05-09T23:13:10.573Z" }, - { url = "https://files.pythonhosted.org/packages/d4/30/66ab84588765f5b4b271a9ca09ef7ce2b87caa95176ec3d2ad65d7bc4902/regex-2026.5.9-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:57e8915c7986aa33d25e4d3629cef711cd2863f2961b10409f0c04cb8b7d9020", size = 865445, upload-time = "2026-05-09T23:13:12.523Z" }, - { url = "https://files.pythonhosted.org/packages/1a/89/f05169e8588aac365f35ffc7f3bc3184f095ef4cfded7cfaa3c7fd5dbd89/regex-2026.5.9-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:508f56a89ba9cb26e4168cbc37dbd60a28d82430a9e18ad1d25fe0883c314ca2", size = 912322, upload-time = "2026-05-09T23:13:14.281Z" }, - { url = "https://files.pythonhosted.org/packages/30/e1/c93444052cf41581f3c884ab3fb5823daf0992f11cd4388d4275ca610558/regex-2026.5.9-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b6d189041f15691cfa2b6c4290448ec221244d225b3f5fe9e7771b34ffcdf6e2", size = 801269, upload-time = "2026-05-09T23:13:16.569Z" }, - { url = "https://files.pythonhosted.org/packages/50/fe/0cf96b882f540e62e8b9956599798203d599c44cf4c77917ca27400ff69b/regex-2026.5.9-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e82db382b44d0111b22601c509c89f64434816c9e0eef9d1989cda8cc6ff1c04", size = 777085, upload-time = "2026-05-09T23:13:18.675Z" }, - { url = "https://files.pythonhosted.org/packages/23/5c/d78d4924e7fc875557b9e9b768423925fdfaac5549d06da7810019a9bd26/regex-2026.5.9-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2acfb48634f64996b57f90f39afa692ff362162722581921fe92239a59960f3c", size = 785153, upload-time = "2026-05-09T23:13:20.525Z" }, - { url = "https://files.pythonhosted.org/packages/bf/e0/5214774090e7b4524dcea3e3c4aa74141d43043f8beb49c1599db1c8b53a/regex-2026.5.9-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:d29eebfc9525db68cad3c97eedd7f754fa265aa5cd0cf4f863b2421e1b48fc9f", size = 860164, upload-time = "2026-05-09T23:13:22.263Z" }, - { url = "https://files.pythonhosted.org/packages/6e/e1/4a57a83350319b1271f0d7a249b8672513ed928b237a741631270de6caea/regex-2026.5.9-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:debb893095e944091c16e641a6e33c1b0f4cb61ab945ec5afbf53ce7068834d8", size = 765731, upload-time = "2026-05-09T23:13:24.277Z" }, - { url = "https://files.pythonhosted.org/packages/12/f4/499e74a20c156fc75836ee04a72a38d1a063978f600937f9760467beb1b0/regex-2026.5.9-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:d659eee77986549c9ea45b861c7567e44d6287c3dc9a4565478853f7b9fe2ff6", size = 852062, upload-time = "2026-05-09T23:13:26.125Z" }, - { url = "https://files.pythonhosted.org/packages/5b/92/7eebc0d0a01e78629695f342ba17e0deaff8fb45e79cc0d7b98287da6e3e/regex-2026.5.9-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:2efa205e6d98b24d1f3ab395c11aa15cdf10935bca283d0285e0499c284fba21", size = 789577, upload-time = "2026-05-09T23:13:27.814Z" }, - { url = "https://files.pythonhosted.org/packages/05/a4/018e71f7d2ad48c1ebe6d3ae0026f9b7cb4802fd15c7cc02fdf724355102/regex-2026.5.9-cp313-cp313-win32.whl", hash = "sha256:f3844f134e834076677dd369976e9f5068679fcb8e50102fdf6b7ac96a3ec127", size = 266691, upload-time = "2026-05-09T23:13:29.549Z" }, - { url = "https://files.pythonhosted.org/packages/e6/1d/861a93719fb9ee7dbfc3761b3797b7a3e112a5d42c6129459d2d741be9b5/regex-2026.5.9-cp313-cp313-win_amd64.whl", hash = "sha256:3527bb4942d2c14552155406cdedd906567456821848aed1cb4933a391bf5eca", size = 277747, upload-time = "2026-05-09T23:13:31.859Z" }, - { url = "https://files.pythonhosted.org/packages/d9/c6/0a2436ae4da1ba76e51cb98943c6838a9a721faa40ebe2dce07694ae34e3/regex-2026.5.9-cp313-cp313-win_arm64.whl", hash = "sha256:56a33f191f17d8c417f99945ebdc1e691d3af9605d86ec68c7e54a57e3e17af6", size = 270500, upload-time = "2026-05-09T23:13:33.525Z" }, - { url = "https://files.pythonhosted.org/packages/e8/e9/d21346f7b60ed58789371358ed66b09d00f832e1bd7c06e55d9da5679882/regex-2026.5.9-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:01f28d868834624c934b8d2e0aa1c8341337e37831f4a012f18a5afcba4cbaf3", size = 494172, upload-time = "2026-05-09T23:13:35.935Z" }, - { url = "https://files.pythonhosted.org/packages/c4/43/fd1177a2032037c681baecdb3422ee4e1424aec4e4f470ef47793d325274/regex-2026.5.9-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:48036f6374aaa79eb3b754ec29c61d1c6b1606749d705a13f8854fa2539671f6", size = 293952, upload-time = "2026-05-09T23:13:38.307Z" }, - { url = "https://files.pythonhosted.org/packages/f2/7d/9fbf919768368d3f8a4f6c692cf2aa61e482b2b81ec6a298ace4cbf02480/regex-2026.5.9-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:b96350aa424e79d4fd6b567b344dcbe2b2d6bfc48dfe7717587e1fa6d43da6ff", size = 292314, upload-time = "2026-05-09T23:13:40.353Z" }, - { url = "https://files.pythonhosted.org/packages/e2/6c/e41bfeecb589716843e7c4df09ba46ff2a42961457afece19059d85caeef/regex-2026.5.9-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8f3af7a4903c5c04a11a196a5aa75cdd7dd3f8508132f9fb3259d9f5908e3b88", size = 811681, upload-time = "2026-05-09T23:13:42.543Z" }, - { url = "https://files.pythonhosted.org/packages/87/83/a5c1c525fba0aa656e88ad0face0b1829788ef4c2fb6b26df58aa1151b84/regex-2026.5.9-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7e87577720152d2caae19fe2baaf1f8d5ca12091e9e229f03915c37d1e4b9178", size = 871135, upload-time = "2026-05-09T23:13:44.326Z" }, - { url = "https://files.pythonhosted.org/packages/18/d4/80882e799e440dd878b0979cbebf8fa4d54624a332c83037c7a701649e3f/regex-2026.5.9-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c8b9b9d294cfea3cd19c718ade7cc93492b2c4991abd9a68d0b3477ae6d8e100", size = 917265, upload-time = "2026-05-09T23:13:47.295Z" }, - { url = "https://files.pythonhosted.org/packages/ae/ff/8db60211e2286e396aad7dc7725356c502bff0901ea05bd6cdc2e1a042b9/regex-2026.5.9-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:728d8bfd28a8845c8b6bc5dc7ce010453d206396786c0765c2740cb65f37791e", size = 816311, upload-time = "2026-05-09T23:13:49.885Z" }, - { url = "https://files.pythonhosted.org/packages/4c/47/742ef579c61730f8d268e5cf1f9ce0e37e2ea041ad0f5644724f2378e463/regex-2026.5.9-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7e30b874d341fac767d7df5a0870540541c2c054b80cfaac116e8d367a8a7ff2", size = 785498, upload-time = "2026-05-09T23:13:52.25Z" }, - { url = "https://files.pythonhosted.org/packages/7f/ab/cb0999802dcb0fb95b1ab005e8d4163d8afdd67efc2cb6b6630ac13f8cb1/regex-2026.5.9-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:fd190e88a895a8901325fad284a3f74ea52b1da8525b76cc811fa9b1edf0ce2b", size = 801348, upload-time = "2026-05-09T23:13:54.127Z" }, - { url = "https://files.pythonhosted.org/packages/7d/62/8ca59a24c55bc34d166eefaf3717bd77772f329fdbf984d86581e0a3571c/regex-2026.5.9-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:8e76e8161ad00694cfce6767d5dea860c6391ac5b83e5c3a39661e696f11fc7e", size = 866493, upload-time = "2026-05-09T23:13:56.067Z" }, - { url = "https://files.pythonhosted.org/packages/8d/3d/30f2ae62cef3278bb5bb821f467277a55fb73f01032cf85997e15e8289a8/regex-2026.5.9-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:ddda5340e6c01a293027dd46232fa79eaff1b48058ce7a98f572b6445b088041", size = 772811, upload-time = "2026-05-09T23:13:57.867Z" }, - { url = "https://files.pythonhosted.org/packages/d8/ae/7d2089bcd78ad0c0161bc684339df50032acb438a7bd3305e7ddb1193cec/regex-2026.5.9-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:205109e96b3cf5adf8f4cd62bedde9487feb282b9497a3535451e5a24cd706a0", size = 856584, upload-time = "2026-05-09T23:13:59.679Z" }, - { url = "https://files.pythonhosted.org/packages/a9/29/92ff47f75990131ea4f24ba17819e5a9d141e10819807e09addd73409af6/regex-2026.5.9-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:dfbe4579b9f08036aa7d101d1835437a20783574ac66327e6b29b4018a138081", size = 803453, upload-time = "2026-05-09T23:14:01.978Z" }, - { url = "https://files.pythonhosted.org/packages/04/99/eff29f1037dcab36702c9ee5d6858cf1ce2336ea8ea2987f64245b99ea5e/regex-2026.5.9-cp313-cp313t-win32.whl", hash = "sha256:ed2c9e8068b614c574d8d30e543d617cf5379b0535d46f97ef00e904745a08b5", size = 269951, upload-time = "2026-05-09T23:14:03.661Z" }, - { url = "https://files.pythonhosted.org/packages/0e/9d/8870b8981d27b22cda77bb26a5ac7ebfa9c7d9e0dea195a834a82380e748/regex-2026.5.9-cp313-cp313t-win_amd64.whl", hash = "sha256:b46b0f094dc1d3b90356c85a0bd2c9bafc4a6a190b9d6f8ddd5a033b6e088ed4", size = 281240, upload-time = "2026-05-09T23:14:05.56Z" }, - { url = "https://files.pythonhosted.org/packages/72/b1/3379415e8f135c13ac551353397cc4fe97b4978f3cac73c5fcbcded548b8/regex-2026.5.9-cp313-cp313t-win_arm64.whl", hash = "sha256:872acc074bd29ffc9913ecdfedf6ea77502312ca44a4aa0d3779089c6069d8de", size = 272383, upload-time = "2026-05-09T23:14:07.843Z" }, - { url = "https://files.pythonhosted.org/packages/13/3e/9c3cd292d8808b3645a2ce517e200179b6d0e903f176300bd8b542e14de5/regex-2026.5.9-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:1bd7587a2948b4085195d5a3374eaf4a425dc3e55784c038175355ecf3bbbf8a", size = 490376, upload-time = "2026-05-09T23:14:09.64Z" }, - { url = "https://files.pythonhosted.org/packages/60/70/d43ee8a2ca0a8b68d167f21658b85520ac0574617c7f320367c5047f7556/regex-2026.5.9-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:dea2e88e1cce4522496cce630e11e67b98b7076620bc4336c3f674bc21a375f4", size = 291964, upload-time = "2026-05-09T23:14:11.424Z" }, - { url = "https://files.pythonhosted.org/packages/21/91/9d50b433828d8e74196904e168a43abf1e6e88b2a15d47ed742456720c37/regex-2026.5.9-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:2099f7e7ff7b6aa3192312650a56e91cc091e49d50b04e4f6f8b6e28b3b27f1c", size = 289682, upload-time = "2026-05-09T23:14:13.123Z" }, - { url = "https://files.pythonhosted.org/packages/3e/d2/b835e3cafbb9d977736912436259ff551d60919f7d7b3d37d46659c63564/regex-2026.5.9-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ecd353045824e4477562a2ac718c25799cdaaa41f7aa925a806a8a3e6848a5b9", size = 796996, upload-time = "2026-05-09T23:14:14.923Z" }, - { url = "https://files.pythonhosted.org/packages/2c/a6/9f992d00019166b9de01c546dd4549bc679f2a68df11b877740b0760b7c2/regex-2026.5.9-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:65c8c8c37377794bd5b2f3ebe51919042bf17aec802e23c833d89782ed0c78af", size = 866089, upload-time = "2026-05-09T23:14:17.757Z" }, - { url = "https://files.pythonhosted.org/packages/e0/08/4d32af657e049b19cb62b02e46e38fe1518797bfb2203ee93a510b21b0dc/regex-2026.5.9-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5b73ab8afcf66c622db143d1c6fda4e58e4d537ee4f125229ad47b1ab80f34c0", size = 911530, upload-time = "2026-05-09T23:14:20.353Z" }, - { url = "https://files.pythonhosted.org/packages/d9/27/2af43dd1dc201d1fecefda64a45f4ad0995855b92724f795a777b402ee69/regex-2026.5.9-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0de5cf193997384ed2ca6f1cd4f78055b255d93d82d5a8cd6ba0d11c10b167e4", size = 800643, upload-time = "2026-05-09T23:14:22.265Z" }, - { url = "https://files.pythonhosted.org/packages/a4/dd/23a249047013b5321d4a60c4d2437462086f601b061776a525e5fba2a59f/regex-2026.5.9-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d641a8c9a61618047796d572a39a79b26167b0411d2c3031937b2fe2d081e2cf", size = 777223, upload-time = "2026-05-09T23:14:24.179Z" }, - { url = "https://files.pythonhosted.org/packages/94/6a/e85ed9538cd19586d0465076a4578a12e093ce776d15f3f8ce92733a8dd6/regex-2026.5.9-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:24b2355ef5cc9aa5b8f07d17704face1c166fdcc2290fa7bd6e6c925655a8346", size = 785760, upload-time = "2026-05-09T23:14:26.065Z" }, - { url = "https://files.pythonhosted.org/packages/2a/c4/f25473209438638e947c55f9156fd8f236f74169229028cc99116380868e/regex-2026.5.9-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:a24852d3c29ad9e47593593d8a247c44ccc3d0548ef12c822d6ed0810affe676", size = 860891, upload-time = "2026-05-09T23:14:28.17Z" }, - { url = "https://files.pythonhosted.org/packages/f9/f7/f4f86e3c74419c37370e91f150ae0c2ef7d34b2e0e4cdd5da046a02e4022/regex-2026.5.9-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:916714069da19329ef7de197dcbc77bb3104145c7c2c864dbfbe318f46b88b14", size = 765891, upload-time = "2026-05-09T23:14:30.06Z" }, - { url = "https://files.pythonhosted.org/packages/26/70/704d8e13765939146b1cd0ef4e2feb71d7929727d2290f026eed10095955/regex-2026.5.9-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:fa411799ca8da32a8d38d020a88faa5b6f91657d284761352940ecf9f7c3bbdd", size = 851380, upload-time = "2026-05-09T23:14:32.123Z" }, - { url = "https://files.pythonhosted.org/packages/26/29/1a13582a8460038edc38e49f64ceb0dd7c60f5caba77571f4bf6601965d9/regex-2026.5.9-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:1e6da47d679b7010ef27556b6e0f99771b744936db1792a10ceac6547ae1503e", size = 789350, upload-time = "2026-05-09T23:14:34.799Z" }, - { url = "https://files.pythonhosted.org/packages/73/56/3dcafe34fc72e271d62ad9a291801e88a1457bb251c132f15fcc2e5aad1a/regex-2026.5.9-cp314-cp314-win32.whl", hash = "sha256:98bd73080e8756255137e1bd3f3f00295bbc5aa383c0e0f973920e9134d7c4ad", size = 272130, upload-time = "2026-05-09T23:14:36.729Z" }, - { url = "https://files.pythonhosted.org/packages/d0/9c/02eebf0be95efe416c664db7fb8b6b05b7a0b06a7544f2884f2558b0526f/regex-2026.5.9-cp314-cp314-win_amd64.whl", hash = "sha256:ff8d372ac2acdc048d1c19916f27ee61bc5722728458ba6ca5052f2c72d51763", size = 280999, upload-time = "2026-05-09T23:14:39.126Z" }, - { url = "https://files.pythonhosted.org/packages/70/5a/1dd1abee76cb7a846a0bcf42fdc87e5720c3c33c24f3e37814310a513d9f/regex-2026.5.9-cp314-cp314-win_arm64.whl", hash = "sha256:e1d93bf647916292e8edcec150c07ddf3dc50179ccaf770c04a7f9e452155372", size = 273500, upload-time = "2026-05-09T23:14:41.059Z" }, - { url = "https://files.pythonhosted.org/packages/86/c1/c5f619b0057a7965cb78ec559c1d7a45ce8c99a35bea95483d64959a93d9/regex-2026.5.9-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:83d0ee4a57d1c87cb549e195ec300b8f0ec3a82eba66d835e4e2ed8634fe4499", size = 494269, upload-time = "2026-05-09T23:14:42.869Z" }, - { url = "https://files.pythonhosted.org/packages/05/2c/5d01f1aee33de4bbe60c8452945bfc8477ca7c5ae4450f6bfe711036cb36/regex-2026.5.9-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:d3d7eb5c9a7f6df82ed3cfac9beb93882a5cbcb5b8b157b56cb2b3b276574ac1", size = 293954, upload-time = "2026-05-09T23:14:44.822Z" }, - { url = "https://files.pythonhosted.org/packages/7a/fe/e8988b2ae2108c6ef71bd4aa8d87fbe257976dd0810e826cd75f701c68b6/regex-2026.5.9-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:075160bf16658e16d35233300b8453aac25de4cbea808d22348b6979668e924d", size = 292405, upload-time = "2026-05-09T23:14:47.211Z" }, - { url = "https://files.pythonhosted.org/packages/79/34/d2b0937faa7859263f7f0a3c6b103a1296306be6952dc173d0154e9a2f49/regex-2026.5.9-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:45375819235558a4ff1c4971dc32881f022613abdb180128f5cb4768c1765a1c", size = 811855, upload-time = "2026-05-09T23:14:49.21Z" }, - { url = "https://files.pythonhosted.org/packages/80/fe/daf53a47457a8486db66c66c01ceb9c2303eecee3f87197f1e77eb1a736d/regex-2026.5.9-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ead4b163ac30a29574510cd4b3e2e985ac5290c05fc7095557d6a5f403fc31b5", size = 871189, upload-time = "2026-05-09T23:14:51.555Z" }, - { url = "https://files.pythonhosted.org/packages/1c/75/058fc4470cbfbf57d800aff1a0022b929a3f9fa553ee10a0cdf2070eb31f/regex-2026.5.9-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8c6e4218fbdfbcd4f6c19efca40930d24a621bf4b48cb76bc6640543bd28ef20", size = 917485, upload-time = "2026-05-09T23:14:53.633Z" }, - { url = "https://files.pythonhosted.org/packages/88/e7/179cfda3a28bc843b5c6cfe7f79f23489c791ed95f151083803660878432/regex-2026.5.9-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6351571c8a42b505eb555c0dc47d740d0fb66977dc142919eea6f4325b7c56a0", size = 816369, upload-time = "2026-05-09T23:14:56.198Z" }, - { url = "https://files.pythonhosted.org/packages/41/90/6f0cc422071688266d344fca8462d787cba0a2c144acb25721f9a61ec265/regex-2026.5.9-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:002205cafd2a9e78c6290c7d1df277bf3277b3b7a30e0b4bb0dac2e2e3f7cb2d", size = 785869, upload-time = "2026-05-09T23:14:58.602Z" }, - { url = "https://files.pythonhosted.org/packages/02/67/a31f1760f09c27b251ef39e9beb541f462cf977381d067faa764c2c0e393/regex-2026.5.9-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8abd33fef90b2a9efac5557d6033ca82d1195ed3a15fea5af15ba7b463c6a63b", size = 801427, upload-time = "2026-05-09T23:15:00.642Z" }, - { url = "https://files.pythonhosted.org/packages/e3/c4/1a80654597b6bc1e1ea0494824c31200e8a956abe290afae9b19a166a148/regex-2026.5.9-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:31037c82eccb44b7ea2e9e221d7c01429430e989a1f4b91ea5a855f6017b509a", size = 866482, upload-time = "2026-05-09T23:15:03.384Z" }, - { url = "https://files.pythonhosted.org/packages/d1/11/960724e06482c08466ff5611e242e86f80062949cdf6b4b9cc317b9dd93d/regex-2026.5.9-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:5604dfd046dc37eca90250fc3be938b076c8059fa772ac0ed6f499b0f0fb0415", size = 773022, upload-time = "2026-05-09T23:15:05.625Z" }, - { url = "https://files.pythonhosted.org/packages/50/a8/a9979c3e7918280e93159ebcab5ef1a65116dd4f3bd6091be0eae4a126e8/regex-2026.5.9-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:0e1b1b4e496afbb24f4a62aba855ee4f88f25578927697b340702e48c9ee6bc2", size = 856642, upload-time = "2026-05-09T23:15:07.966Z" }, - { url = "https://files.pythonhosted.org/packages/fe/d4/a9b732f2f0072c0ab12227483abb24fffcb9f73f8a2b203df0a6d0434735/regex-2026.5.9-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:be3372b9df6ddecff6486d37e19095a7b4973137caf5512407a89f4455361f41", size = 803552, upload-time = "2026-05-09T23:15:10.215Z" }, - { url = "https://files.pythonhosted.org/packages/d5/fe/1b3113817447a1d4155e4ac76d2e072f42c0bcba2f43fa8a0e756ea2cd91/regex-2026.5.9-cp314-cp314t-win32.whl", hash = "sha256:3ddd90103f9e5c471c49c7852ecc1fe27c7e45eb99e977aefe7caa4e779f4f58", size = 275746, upload-time = "2026-05-09T23:15:12.609Z" }, - { url = "https://files.pythonhosted.org/packages/92/73/93d42045302636c91f2e5ef588b65b84b01428f28ec77de256b1dfdfbe5c/regex-2026.5.9-cp314-cp314t-win_amd64.whl", hash = "sha256:ca518ed29c46eecba6010b15f1b9a479314d2de409536e71b6a13aa04e3b8a77", size = 285685, upload-time = "2026-05-09T23:15:15.086Z" }, - { url = "https://files.pythonhosted.org/packages/da/80/35b4c33c804a165a7f55289afda3ea9e3eb6d15800341a2d66455c0f1f30/regex-2026.5.9-cp314-cp314t-win_arm64.whl", hash = "sha256:5e41809d2683fcde7d5a8c87a6567ba1fb1ce0de9f31bff578de00a4b2d76daa", size = 275713, upload-time = "2026-05-09T23:15:16.98Z" }, -] - -[[package]] -name = "requests" -version = "2.34.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "certifi" }, - { name = "charset-normalizer" }, - { name = "idna" }, - { name = "urllib3" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ac/c3/e2a2b89f2d3e2179abd6d00ebd70bff6273f37fb3e0cc209f48b39d00cbf/requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed", size = 142856, upload-time = "2026-05-14T19:25:27.735Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, -] - -[[package]] -name = "rich" -version = "15.0.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "markdown-it-py" }, - { name = "pygments" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/0722ca900cc807c13a6a0c696dacf35430f72e0ec571c4275d2371fca3e9/rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36", size = 230680, upload-time = "2026-04-12T08:24:00.75Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, -] - -[[package]] -name = "safetensors" -version = "0.7.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/29/9c/6e74567782559a63bd040a236edca26fd71bc7ba88de2ef35d75df3bca5e/safetensors-0.7.0.tar.gz", hash = "sha256:07663963b67e8bd9f0b8ad15bb9163606cd27cc5a1b96235a50d8369803b96b0", size = 200878, upload-time = "2025-11-19T15:18:43.199Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fa/47/aef6c06649039accf914afef490268e1067ed82be62bcfa5b7e886ad15e8/safetensors-0.7.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c82f4d474cf725255d9e6acf17252991c3c8aac038d6ef363a4bf8be2f6db517", size = 467781, upload-time = "2025-11-19T15:18:35.84Z" }, - { url = "https://files.pythonhosted.org/packages/e8/00/374c0c068e30cd31f1e1b46b4b5738168ec79e7689ca82ee93ddfea05109/safetensors-0.7.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:94fd4858284736bb67a897a41608b5b0c2496c9bdb3bf2af1fa3409127f20d57", size = 447058, upload-time = "2025-11-19T15:18:34.416Z" }, - { url = "https://files.pythonhosted.org/packages/f1/06/578ffed52c2296f93d7fd2d844cabfa92be51a587c38c8afbb8ae449ca89/safetensors-0.7.0-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e07d91d0c92a31200f25351f4acb2bc6aff7f48094e13ebb1d0fb995b54b6542", size = 491748, upload-time = "2025-11-19T15:18:09.79Z" }, - { url = "https://files.pythonhosted.org/packages/ae/33/1debbbb70e4791dde185edb9413d1fe01619255abb64b300157d7f15dddd/safetensors-0.7.0-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8469155f4cb518bafb4acf4865e8bb9d6804110d2d9bdcaa78564b9fd841e104", size = 503881, upload-time = "2025-11-19T15:18:16.145Z" }, - { url = "https://files.pythonhosted.org/packages/8e/1c/40c2ca924d60792c3be509833df711b553c60effbd91da6f5284a83f7122/safetensors-0.7.0-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:54bef08bf00a2bff599982f6b08e8770e09cc012d7bba00783fc7ea38f1fb37d", size = 623463, upload-time = "2025-11-19T15:18:21.11Z" }, - { url = "https://files.pythonhosted.org/packages/9b/3a/13784a9364bd43b0d61eef4bea2845039bc2030458b16594a1bd787ae26e/safetensors-0.7.0-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:42cb091236206bb2016d245c377ed383aa7f78691748f3bb6ee1bfa51ae2ce6a", size = 532855, upload-time = "2025-11-19T15:18:25.719Z" }, - { url = "https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dac7252938f0696ddea46f5e855dd3138444e82236e3be475f54929f0c510d48", size = 507152, upload-time = "2025-11-19T15:18:33.023Z" }, - { url = "https://files.pythonhosted.org/packages/3c/a8/4b45e4e059270d17af60359713ffd83f97900d45a6afa73aaa0d737d48b6/safetensors-0.7.0-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1d060c70284127fa805085d8f10fbd0962792aed71879d00864acda69dbab981", size = 541856, upload-time = "2025-11-19T15:18:31.075Z" }, - { url = "https://files.pythonhosted.org/packages/06/87/d26d8407c44175d8ae164a95b5a62707fcc445f3c0c56108e37d98070a3d/safetensors-0.7.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:cdab83a366799fa730f90a4ebb563e494f28e9e92c4819e556152ad55e43591b", size = 674060, upload-time = "2025-11-19T15:18:37.211Z" }, - { url = "https://files.pythonhosted.org/packages/11/f5/57644a2ff08dc6325816ba7217e5095f17269dada2554b658442c66aed51/safetensors-0.7.0-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:672132907fcad9f2aedcb705b2d7b3b93354a2aec1b2f706c4db852abe338f85", size = 771715, upload-time = "2025-11-19T15:18:38.689Z" }, - { url = "https://files.pythonhosted.org/packages/86/31/17883e13a814bd278ae6e266b13282a01049b0c81341da7fd0e3e71a80a3/safetensors-0.7.0-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:5d72abdb8a4d56d4020713724ba81dac065fedb7f3667151c4a637f1d3fb26c0", size = 714377, upload-time = "2025-11-19T15:18:40.162Z" }, - { url = "https://files.pythonhosted.org/packages/4a/d8/0c8a7dc9b41dcac53c4cbf9df2b9c83e0e0097203de8b37a712b345c0be5/safetensors-0.7.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b0f6d66c1c538d5a94a73aa9ddca8ccc4227e6c9ff555322ea40bdd142391dd4", size = 677368, upload-time = "2025-11-19T15:18:41.627Z" }, - { url = "https://files.pythonhosted.org/packages/05/e5/cb4b713c8a93469e3c5be7c3f8d77d307e65fe89673e731f5c2bfd0a9237/safetensors-0.7.0-cp38-abi3-win32.whl", hash = "sha256:c74af94bf3ac15ac4d0f2a7c7b4663a15f8c2ab15ed0fc7531ca61d0835eccba", size = 326423, upload-time = "2025-11-19T15:18:45.74Z" }, - { url = "https://files.pythonhosted.org/packages/5d/e6/ec8471c8072382cb91233ba7267fd931219753bb43814cbc71757bfd4dab/safetensors-0.7.0-cp38-abi3-win_amd64.whl", hash = "sha256:d1239932053f56f3456f32eb9625590cc7582e905021f94636202a864d470755", size = 341380, upload-time = "2025-11-19T15:18:44.427Z" }, - { url = "https://files.pythonhosted.org/packages/a7/6a/4d08d89a6fcbe905c5ae68b8b34f0791850882fc19782d0d02c65abbdf3b/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f4729811a6640d019a4b7ba8638ee2fd21fa5ca8c7e7bdf0fed62068fcaac737", size = 492430, upload-time = "2025-11-19T15:18:11.884Z" }, - { url = "https://files.pythonhosted.org/packages/dd/29/59ed8152b30f72c42d00d241e58eaca558ae9dbfa5695206e2e0f54c7063/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:12f49080303fa6bb424b362149a12949dfbbf1e06811a88f2307276b0c131afd", size = 503977, upload-time = "2025-11-19T15:18:17.523Z" }, - { url = "https://files.pythonhosted.org/packages/d3/0b/4811bfec67fa260e791369b16dab105e4bae82686120554cc484064e22b4/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0071bffba4150c2f46cae1432d31995d77acfd9f8db598b5d1a2ce67e8440ad2", size = 623890, upload-time = "2025-11-19T15:18:22.666Z" }, - { url = "https://files.pythonhosted.org/packages/58/5b/632a58724221ef03d78ab65062e82a1010e1bef8e8e0b9d7c6d7b8044841/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:473b32699f4200e69801bf5abf93f1a4ecd432a70984df164fc22ccf39c4a6f3", size = 531885, upload-time = "2025-11-19T15:18:27.146Z" }, -] - -[[package]] -name = "shellingham" -version = "1.5.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, -] - -[[package]] -name = "tqdm" -version = "4.67.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/09/a9/6ba95a270c6f1fbcd8dac228323f2777d886cb206987444e4bce66338dd4/tqdm-4.67.3.tar.gz", hash = "sha256:7d825f03f89244ef73f1d4ce193cb1774a8179fd96f31d7e1dcde62092b960bb", size = 169598, upload-time = "2026-02-03T17:35:53.048Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl", hash = "sha256:ee1e4c0e59148062281c49d80b25b67771a127c85fc9676d3be5f243206826bf", size = 78374, upload-time = "2026-02-03T17:35:50.982Z" }, -] - -[[package]] -name = "typer" -version = "0.25.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "annotated-doc" }, - { name = "click" }, - { name = "rich" }, - { name = "shellingham" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e4/51/9aed62104cea109b820bbd6c14245af756112017d309da813ef107d42e7e/typer-0.25.1.tar.gz", hash = "sha256:9616eb8853a09ffeabab1698952f33c6f29ffdbceb4eaeecf571880e8d7664cc", size = 122276, upload-time = "2026-04-30T19:32:16.964Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/f9/2b3ff4e56e5fa7debfaf9eb135d0da96f3e9a1d5b27222223c7296336e5f/typer-0.25.1-py3-none-any.whl", hash = "sha256:75caa44ed46a03fb2dab8808753ffacdbfea88495e74c85a28c5eefcf5f39c89", size = 58409, upload-time = "2026-04-30T19:32:18.271Z" }, -] - -[[package]] -name = "typing-extensions" -version = "4.15.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, -] - -[[package]] -name = "urllib3" -version = "2.7.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, -] - -[[package]] -name = "zipp" -version = "4.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b9/d8/eab98a517c14134c0b2eb4e2387bc5f457334293ec5d2dd3857ec2966802/zipp-4.1.0.tar.gz", hash = "sha256:4cb57381f544315db7688e976e922a2b18cdb513d21cc194eb42232ba2a3e602", size = 26214, upload-time = "2026-05-18T20:08:57.967Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3a/13/547360d81e6d88d58492968ffda9f9542854f11310ee556fef14260cc886/zipp-4.1.0-py3-none-any.whl", hash = "sha256:25ad4e16390cd314347dd8f1de67a2ac538ae658ed4ab9db16029c07c188e97f", size = 10238, upload-time = "2026-05-18T20:08:57.045Z" }, -] diff --git a/pyproject.toml b/pyproject.toml index 2130b4f9..ed9d36b7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "accelerate", "av", "cattrs", - "diffusers", + "diffusers>=0.39.0", "einops", "hydra-core", "imageio-ffmpeg", @@ -72,7 +72,6 @@ train = [ "boto3", "botocore", "datasets", - "diffusers-cosmos3", "dists-pytorch", "einx", "fastparquet", @@ -298,7 +297,6 @@ required-environments = [ ] [tool.uv.sources] -diffusers-cosmos3 = { path = "packages/diffusers-cosmos3", editable = true } flash-attn = { index = "cosmos"} flash-attn-3-nv = { index = "cosmos" } lerobot = { git = "https://github.com/mli0603/lerobot.git" } @@ -336,10 +334,6 @@ ignore = [ "PYSEC-2025-217", # flash-attn: flash-attn deserialization isn't used in cosmos3 package "GHSA-7g5w-pq96-8c5w", - # Need diffusers 0.38 - "GHSA-7wx4-6vff-v64p", - "GHSA-98h9-4798-4q5v", - "PYSEC-2026-41", # torch, no fix "PYSEC-2026-139", # Need vllm 0.20 diff --git a/uv.lock b/uv.lock index 81fa99bd..5dc44665 100644 --- a/uv.lock +++ b/uv.lock @@ -1789,7 +1789,6 @@ train = [ { name = "boto3" }, { name = "botocore" }, { name = "datasets" }, - { name = "diffusers-cosmos3" }, { name = "dists-pytorch" }, { name = "einx" }, { name = "fastparquet" }, @@ -2018,8 +2017,7 @@ requires-dist = [ { name = "botocore", marker = "extra == 'train'" }, { name = "cattrs" }, { name = "datasets", marker = "extra == 'train'" }, - { name = "diffusers" }, - { name = "diffusers-cosmos3", marker = "extra == 'train'", editable = "packages/diffusers-cosmos3" }, + { name = "diffusers", specifier = ">=0.39.0" }, { name = "dists-pytorch", marker = "extra == 'train'" }, { name = "einops" }, { name = "einx", marker = "extra == 'train'" }, @@ -2820,7 +2818,7 @@ wheels = [ [[package]] name = "diffusers" -version = "0.37.0" +version = "0.39.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filelock" }, @@ -2833,22 +2831,11 @@ dependencies = [ { name = "requests" }, { name = "safetensors" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/58/3b/01d0ff800b811c5ad8bba682f4c6abf1d7071cd81464c01724333fefb7ba/diffusers-0.37.0.tar.gz", hash = "sha256:408789af73898585f525afd07ca72b3955affea4216a669558e9f59b5b1fe704", size = 4141136, upload-time = "2026-03-05T14:58:39.704Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1a/81/6095237b86a3116c4789f28c4435d5296c00c0fc74ffde99008fd6b3a36c/diffusers-0.39.0.tar.gz", hash = "sha256:14bb1d98c85a0e463d734c99aaa73b480a7bc9bad22af30fbf730ef8f09c1d67", size = 4651240, upload-time = "2026-07-03T08:48:47.904Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f7/55/586a3a2b9c95f371c9c3cb048c3cac15aedcce8d6d53ebd6bbc46860722d/diffusers-0.37.0-py3-none-any.whl", hash = "sha256:7eab74bf896974250b5e1027cae813aba1004f02d97c9b44891b83713386aa08", size = 5000449, upload-time = "2026-03-05T14:58:37.361Z" }, -] - -[[package]] -name = "diffusers-cosmos3" -version = "0.1.0" -source = { editable = "packages/diffusers-cosmos3" } -dependencies = [ - { name = "diffusers" }, + { url = "https://files.pythonhosted.org/packages/3f/3f/7469c46e9d22307ea686bab687d70e6bf328722952f9d10339f5e913e608/diffusers-0.39.0-py3-none-any.whl", hash = "sha256:912aca51b5787365110806e984d5555735bf8a461073bb8459029d0bca7870ef", size = 5631176, upload-time = "2026-07-03T08:48:45.337Z" }, ] -[package.metadata] -requires-dist = [{ name = "diffusers", specifier = ">=0.37.0" }] - [[package]] name = "dill" version = "0.4.0" @@ -11377,28 +11364,26 @@ wheels = [ [[package]] name = "safetensors" -version = "0.7.0" +version = "0.8.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/29/9c/6e74567782559a63bd040a236edca26fd71bc7ba88de2ef35d75df3bca5e/safetensors-0.7.0.tar.gz", hash = "sha256:07663963b67e8bd9f0b8ad15bb9163606cd27cc5a1b96235a50d8369803b96b0", size = 200878, upload-time = "2025-11-19T15:18:43.199Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fa/47/aef6c06649039accf914afef490268e1067ed82be62bcfa5b7e886ad15e8/safetensors-0.7.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c82f4d474cf725255d9e6acf17252991c3c8aac038d6ef363a4bf8be2f6db517", size = 467781, upload-time = "2025-11-19T15:18:35.84Z" }, - { url = "https://files.pythonhosted.org/packages/e8/00/374c0c068e30cd31f1e1b46b4b5738168ec79e7689ca82ee93ddfea05109/safetensors-0.7.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:94fd4858284736bb67a897a41608b5b0c2496c9bdb3bf2af1fa3409127f20d57", size = 447058, upload-time = "2025-11-19T15:18:34.416Z" }, - { url = "https://files.pythonhosted.org/packages/f1/06/578ffed52c2296f93d7fd2d844cabfa92be51a587c38c8afbb8ae449ca89/safetensors-0.7.0-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e07d91d0c92a31200f25351f4acb2bc6aff7f48094e13ebb1d0fb995b54b6542", size = 491748, upload-time = "2025-11-19T15:18:09.79Z" }, - { url = "https://files.pythonhosted.org/packages/ae/33/1debbbb70e4791dde185edb9413d1fe01619255abb64b300157d7f15dddd/safetensors-0.7.0-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8469155f4cb518bafb4acf4865e8bb9d6804110d2d9bdcaa78564b9fd841e104", size = 503881, upload-time = "2025-11-19T15:18:16.145Z" }, - { url = "https://files.pythonhosted.org/packages/8e/1c/40c2ca924d60792c3be509833df711b553c60effbd91da6f5284a83f7122/safetensors-0.7.0-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:54bef08bf00a2bff599982f6b08e8770e09cc012d7bba00783fc7ea38f1fb37d", size = 623463, upload-time = "2025-11-19T15:18:21.11Z" }, - { url = "https://files.pythonhosted.org/packages/9b/3a/13784a9364bd43b0d61eef4bea2845039bc2030458b16594a1bd787ae26e/safetensors-0.7.0-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:42cb091236206bb2016d245c377ed383aa7f78691748f3bb6ee1bfa51ae2ce6a", size = 532855, upload-time = "2025-11-19T15:18:25.719Z" }, - { url = "https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dac7252938f0696ddea46f5e855dd3138444e82236e3be475f54929f0c510d48", size = 507152, upload-time = "2025-11-19T15:18:33.023Z" }, - { url = "https://files.pythonhosted.org/packages/3c/a8/4b45e4e059270d17af60359713ffd83f97900d45a6afa73aaa0d737d48b6/safetensors-0.7.0-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1d060c70284127fa805085d8f10fbd0962792aed71879d00864acda69dbab981", size = 541856, upload-time = "2025-11-19T15:18:31.075Z" }, - { url = "https://files.pythonhosted.org/packages/06/87/d26d8407c44175d8ae164a95b5a62707fcc445f3c0c56108e37d98070a3d/safetensors-0.7.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:cdab83a366799fa730f90a4ebb563e494f28e9e92c4819e556152ad55e43591b", size = 674060, upload-time = "2025-11-19T15:18:37.211Z" }, - { url = "https://files.pythonhosted.org/packages/11/f5/57644a2ff08dc6325816ba7217e5095f17269dada2554b658442c66aed51/safetensors-0.7.0-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:672132907fcad9f2aedcb705b2d7b3b93354a2aec1b2f706c4db852abe338f85", size = 771715, upload-time = "2025-11-19T15:18:38.689Z" }, - { url = "https://files.pythonhosted.org/packages/86/31/17883e13a814bd278ae6e266b13282a01049b0c81341da7fd0e3e71a80a3/safetensors-0.7.0-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:5d72abdb8a4d56d4020713724ba81dac065fedb7f3667151c4a637f1d3fb26c0", size = 714377, upload-time = "2025-11-19T15:18:40.162Z" }, - { url = "https://files.pythonhosted.org/packages/4a/d8/0c8a7dc9b41dcac53c4cbf9df2b9c83e0e0097203de8b37a712b345c0be5/safetensors-0.7.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b0f6d66c1c538d5a94a73aa9ddca8ccc4227e6c9ff555322ea40bdd142391dd4", size = 677368, upload-time = "2025-11-19T15:18:41.627Z" }, - { url = "https://files.pythonhosted.org/packages/05/e5/cb4b713c8a93469e3c5be7c3f8d77d307e65fe89673e731f5c2bfd0a9237/safetensors-0.7.0-cp38-abi3-win32.whl", hash = "sha256:c74af94bf3ac15ac4d0f2a7c7b4663a15f8c2ab15ed0fc7531ca61d0835eccba", size = 326423, upload-time = "2025-11-19T15:18:45.74Z" }, - { url = "https://files.pythonhosted.org/packages/5d/e6/ec8471c8072382cb91233ba7267fd931219753bb43814cbc71757bfd4dab/safetensors-0.7.0-cp38-abi3-win_amd64.whl", hash = "sha256:d1239932053f56f3456f32eb9625590cc7582e905021f94636202a864d470755", size = 341380, upload-time = "2025-11-19T15:18:44.427Z" }, - { url = "https://files.pythonhosted.org/packages/a7/6a/4d08d89a6fcbe905c5ae68b8b34f0791850882fc19782d0d02c65abbdf3b/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f4729811a6640d019a4b7ba8638ee2fd21fa5ca8c7e7bdf0fed62068fcaac737", size = 492430, upload-time = "2025-11-19T15:18:11.884Z" }, - { url = "https://files.pythonhosted.org/packages/dd/29/59ed8152b30f72c42d00d241e58eaca558ae9dbfa5695206e2e0f54c7063/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:12f49080303fa6bb424b362149a12949dfbbf1e06811a88f2307276b0c131afd", size = 503977, upload-time = "2025-11-19T15:18:17.523Z" }, - { url = "https://files.pythonhosted.org/packages/d3/0b/4811bfec67fa260e791369b16dab105e4bae82686120554cc484064e22b4/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0071bffba4150c2f46cae1432d31995d77acfd9f8db598b5d1a2ce67e8440ad2", size = 623890, upload-time = "2025-11-19T15:18:22.666Z" }, - { url = "https://files.pythonhosted.org/packages/58/5b/632a58724221ef03d78ab65062e82a1010e1bef8e8e0b9d7c6d7b8044841/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:473b32699f4200e69801bf5abf93f1a4ecd432a70984df164fc22ccf39c4a6f3", size = 531885, upload-time = "2025-11-19T15:18:27.146Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/a0/f718cda65b05407d228f97602cf60dca269c979867aa5beb25410de26cd3/safetensors-0.8.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c554f85858e05226d3c2828e32395e677434685d6d94594a41643361c5e837f0", size = 473568, upload-time = "2026-06-09T07:52:18.829Z" }, + { url = "https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c80201d22cbf405b80647a60ada77bba06c8fba2da2743ba1e89cdcc39a81f25", size = 484562, upload-time = "2026-06-09T07:52:17.518Z" }, + { url = "https://files.pythonhosted.org/packages/09/7d/65a7de0af421317bb36a067241e4235fff194eed60b961ed6d3f59a3fc60/safetensors-0.8.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a46e5ff292c356d6991e60942ba7f79817682d3a2cef0702136448cb9c4d235", size = 502844, upload-time = "2026-06-09T07:52:07.624Z" }, + { url = "https://files.pythonhosted.org/packages/91/4f/3175c9d75634e0e0dda0082794193521035edd7c70a6f212bf33ca06ddf4/safetensors-0.8.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4124502b78f03534117c848f87a39b8f31e577b15eff423bf8bfb95f2a8c30d0", size = 511823, upload-time = "2026-06-09T07:52:09.565Z" }, + { url = "https://files.pythonhosted.org/packages/20/87/846c289e7aa2299eff406335717cf43ce8777194ece8aad75772e0411615/safetensors-0.8.0-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7bc0a787ba8a35be368ee3574edfa2b1ad389eebd0a72e482ae275490e3f6c98", size = 633461, upload-time = "2026-06-09T07:52:11.128Z" }, + { url = "https://files.pythonhosted.org/packages/76/22/8d64d9df2c45d5ded401df889d0ad90882804ca172d79ec4f0df8f727fe0/safetensors-0.8.0-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:040070828e36dc8e122178bbbd5830ff9e97920affb84cbe0f46442497bed358", size = 545148, upload-time = "2026-06-09T07:52:13.603Z" }, + { url = "https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd6f3f93c9a0a7cc2788ee63fb763353d4bd2e89b0751bc78fcf7dda00bea774", size = 516040, upload-time = "2026-06-09T07:52:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/46/fb/cdaed17ceb2948784fd9c36b6fd3e951b608547cea81a48e8ee6f8cfdfcb/safetensors-0.8.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:fcdd41ec4628fee5799f807c73c353629130fbd942aa23d83c623dd6c9d52d78", size = 513832, upload-time = "2026-06-09T07:52:12.37Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/1e15de264dcc3b77943d2d0c56a95809956883b1c2d6d585c792523f180b/safetensors-0.8.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:8e9f537aa183a38ace122d27303dcd986b26bd2a7591f9181d7f0c396f4677ca", size = 559930, upload-time = "2026-06-09T07:52:14.743Z" }, + { url = "https://files.pythonhosted.org/packages/2a/43/bf38443278eab4b1be1fce2931e2b012ad9cb7df52ada751d0aab8f7659a/safetensors-0.8.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87eec7ffed2b809f05a398a8becb7d013f19f7837cd15d9748580d6cf30dbaf4", size = 678670, upload-time = "2026-06-09T07:52:20.032Z" }, + { url = "https://files.pythonhosted.org/packages/72/e3/68cd3fa5b48488e84add63e04cb12f3bc28ae4638c06d4508c6e88823d0e/safetensors-0.8.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4a95ae2b05d7726d751da4ebf626a2ca782b706e101bd894c95bc2450b1cffcc", size = 786679, upload-time = "2026-06-09T07:52:21.322Z" }, + { url = "https://files.pythonhosted.org/packages/29/4b/1c19c509d56e01f4fbb3d0a2e597450f6cc04d1d56cf52defb0a62dfd715/safetensors-0.8.0-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:3ae091f16662658bdc019a4ff6cb4c085bb7d725eb5978b183ffd265863b6d2d", size = 765683, upload-time = "2026-06-09T07:52:22.594Z" }, + { url = "https://files.pythonhosted.org/packages/27/43/41c1621732edd934d868a00d1b891584c892a7b62a9aab82ea5a0a5623ee/safetensors-0.8.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8e080062fcde23be189565e1c3305d16751a218ecf9412c8601e64204eb6f846", size = 722361, upload-time = "2026-06-09T07:52:23.924Z" }, + { url = "https://files.pythonhosted.org/packages/8e/3f/73ccf82579412b4a71c4ca673f10b5f1f888d7cf5af7fe24f27d30307be4/safetensors-0.8.0-cp310-abi3-win32.whl", hash = "sha256:2ddf52eac562eda224f99acfa7889d02968c1fd59a5b011ae7d8137c37e9c02d", size = 342401, upload-time = "2026-06-09T07:52:28.895Z" }, + { url = "https://files.pythonhosted.org/packages/1b/6d/3fba214c1e5e0f69991677ec3bc17023f0421776975e1de0c682dca475e2/safetensors-0.8.0-cp310-abi3-win_amd64.whl", hash = "sha256:096ec1a98435df7beb08853bb5aa9081a84f23d0adc67ed1a0a10550f608373f", size = 355540, upload-time = "2026-06-09T07:52:27.832Z" }, + { url = "https://files.pythonhosted.org/packages/8d/fc/7eedc3510d97878876e32774eebbeb61c43f148a96e915c84229a3e967aa/safetensors-0.8.0-cp310-abi3-win_arm64.whl", hash = "sha256:f7838e5135a406ad3e02efdcb8cf2e5397d368b0154537c4fec682dbc544d452", size = 340500, upload-time = "2026-06-09T07:52:26.745Z" }, ] [[package]] From 6d84808e7c796cc7b68aa54b0c5eab2f3107e8ae Mon Sep 17 00:00:00 2001 From: Maosheng Liao Date: Thu, 9 Jul 2026 19:05:55 +0800 Subject: [PATCH 19/26] Support reasoner video input. (#25) Co-authored-by: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/args.py | 5 +- cosmos_framework/inference/args_test.py | 13 ++++ .../defaults/reasoner/sample_args.json | 3 +- cosmos_framework/inference/inference.py | 69 ++++++++++++++++--- cosmos_framework/inference/inference_test.py | 30 +++++++- docs/inference.md | 9 +++ inputs/reasoner/reasoner_video.json | 5 ++ 7 files changed, 120 insertions(+), 14 deletions(-) create mode 100644 inputs/reasoner/reasoner_video.json diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 914d0807..4ef9092c 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -647,11 +647,12 @@ class ReasonerDataArgs(ArgsBase): top_p: _ReasonerTopP | None = None repetition_penalty: _ReasonerRepetitionPenalty | None = None presence_penalty: float | None = None + video_fps: pydantic.PositiveFloat | None = None class ReasonerDataOverrides(OverridesBase): """Reasoner overrides for ``model_mode='reasoner'``. ``vision_path`` (if set) is - used as the conditioning image; the VLM processor handles preprocessing.""" + used as the conditioning image or video; the VLM processor handles preprocessing.""" max_new_tokens: pydantic.PositiveInt | None = None """Maximum number of new tokens to generate per prompt.""" @@ -667,6 +668,8 @@ class ReasonerDataOverrides(OverridesBase): """CTRL/HF-style multiplicative repetition penalty (>0). ``1.0`` is identity.""" presence_penalty: float | None = None """Additive presence penalty (any sign). ``0.0`` is identity.""" + video_fps: pydantic.PositiveFloat | None = None + """Frames per second to sample from a video vision_path. None -> decoder default (2.0).""" def _build_reasoner_data(self, model_config: "OmniMoTModelConfig", sample_meta: SampleMeta): if not sample_meta.model_mode.is_reasoner: diff --git a/cosmos_framework/inference/args_test.py b/cosmos_framework/inference/args_test.py index eb5002b1..b3f9b0e7 100644 --- a/cosmos_framework/inference/args_test.py +++ b/cosmos_framework/inference/args_test.py @@ -209,6 +209,7 @@ def test_build_sound_data_rejects_model_without_sound_gen(): def test_audio_image2video_conditions_image_and_sound(tmp_path: Path): import omegaconf + from cosmos_framework.inference.common.config import structure_config setup_args = OmniSetupOverrides( @@ -234,3 +235,15 @@ def test_audio_image2video_conditions_image_and_sound(tmp_path: Path): assert args.condition_frame_indexes_vision == [0] assert args.enable_sound is True assert Path(args.sound_path).name == "clip.wav" + + +def test_reasoner_video_fps_defaults_none(): + from cosmos_framework.inference.args import ReasonerDataOverrides + + assert ReasonerDataOverrides().video_fps is None + + +def test_reasoner_video_fps_accepts_positive_float(): + from cosmos_framework.inference.args import ReasonerDataOverrides + + assert ReasonerDataOverrides(video_fps=2.0).video_fps == 2.0 diff --git a/cosmos_framework/inference/defaults/reasoner/sample_args.json b/cosmos_framework/inference/defaults/reasoner/sample_args.json index cc539916..e7a25adf 100644 --- a/cosmos_framework/inference/defaults/reasoner/sample_args.json +++ b/cosmos_framework/inference/defaults/reasoner/sample_args.json @@ -6,5 +6,6 @@ "top_k": null, "top_p": null, "repetition_penalty": 1.0, - "presence_penalty": 0.0 + "presence_penalty": 0.0, + "video_fps": null } diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index 53ee1a57..a7f338ca 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -14,11 +14,15 @@ import cattrs.preconf.json import safetensors.torch import torch +import torchvision.io from PIL import Image +from qwen_vl_utils.vision_process import smart_nframes from torch.utils._pytree import tree_map_only from torch.utils.data import Dataset from typing_extensions import Self +from cosmos_framework.configs.base.defaults.compile import CompileConfig +from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig from cosmos_framework.inference.args import ( ModelMode, NegativeMetadataMode, @@ -26,6 +30,7 @@ OmniSetupArgs, ) from cosmos_framework.inference.common.args import ( + VIDEO_EXTENSIONS, CheckpointType, ConfigFileType, ParallelismArgs, @@ -46,13 +51,11 @@ pil_to_conditioning_frames, resize_pil_image, ) -from cosmos_framework.utils import log -from cosmos_framework.tools.visualize.video import save_img_or_video -from cosmos_framework.configs.base.defaults.compile import CompileConfig -from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import _SYSTEM_PROMPT_IMAGE_EDITING from cosmos_framework.model.generator.upsampler.prompts import is_upsampled_prompt +from cosmos_framework.tools.visualize.video import save_img_or_video +from cosmos_framework.utils import log if TYPE_CHECKING: from cosmos_framework.configs.base.defaults.model_config import OmniMoTModelConfig @@ -463,14 +466,44 @@ def _get_prompt_sample_data(sample_args: OmniSampleArgs, model: OmniMoTModel, *, return out +def _decode_reasoner_video(vision_path: str, video_fps: float | None) -> dict[str, Any]: + """Decode a local video file into the frame-list payload the Qwen3-VL processor expects. + + Returns ``{"frames": [PIL.Image, ...], "fps": float}``. Uses the same + ``torchvision.io.read_video`` decode the rest of the inference path relies on + (no ``decord`` dependency), then uniformly samples frames toward ``video_fps`` + (default 2.0) via Qwen's ``smart_nframes``. The repo ``Qwen3VLProcessor`` runs + with ``do_sample_frames=False``, so it consumes this pre-sampled frame list + as-is and handles its own per-frame resize.""" + frames, _, info = torchvision.io.read_video(str(vision_path), pts_unit="sec") # [T,H,W,C] uint8 + total_frames = int(frames.shape[0]) + if total_frames == 0: + raise ValueError(f"Decoded zero frames from reasoner video: {vision_path}") + src_fps = float(info.get("video_fps") or 0.0) or 1.0 + target_fps = video_fps if video_fps is not None else 2.0 + nframes = smart_nframes({"fps": target_fps}, total_frames=total_frames, video_fps=src_fps) + idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() + pil_frames = [Image.fromarray(frames[i].numpy()) for i in idx] + sample_fps = nframes / total_frames * src_fps + return {"frames": pil_frames, "fps": sample_fps} + + def _get_reasoner_sample_data(sample_args: OmniSampleArgs, model: OmniMoTModel) -> dict[str, Any]: - """Sample batch for reasoner text generation: prompt + optional conditioning image.""" + """Sample batch for reasoner text generation: prompt + optional conditioning image or video.""" image: Image.Image | None = None + video: dict[str, Any] | None = None if sample_args.vision_path is not None: - image = Image.open(sample_args.vision_path).convert("RGB") + if Path(sample_args.vision_path).suffix.lower() in VIDEO_EXTENSIONS: + video = _decode_reasoner_video(str(sample_args.vision_path), sample_args.video_fps) + else: + image = Image.open(sample_args.vision_path).convert("RGB") + # Both keys are emitted for every sample (``None`` when absent) so the batch + # builder can positionally align them and the three-way homogeneity check in + # ``_generate_reasoner_batch`` reliably detects an image/video/text mix. return { model.input_caption_key: [sample_args.prompt], "reasoner_images": [image], + "reasoner_videos": [video], } @@ -1684,13 +1717,28 @@ def _generate_reasoner_batch( prompts: list[str] = data_batch[self.model.input_caption_key] raw_images: list[Image.Image | None] = data_batch["reasoner_images"] - n_set = sum(img is not None for img in raw_images) - if 0 < n_set < len(raw_images): + raw_videos: list[dict[str, Any] | None] | None = data_batch.get("reasoner_videos") + + n_img = sum(img is not None for img in raw_images) + n_vid = sum(v is not None for v in (raw_videos or [])) + if n_img and n_vid: + raise ValueError( + "Reasoner batch mixes image- and video-conditioned samples. Split into separate batches." + ) + if 0 < n_img < len(raw_images): raise ValueError( "Reasoner batch mixes image-conditioned and text-only samples " - f"({n_set}/{len(raw_images)} have vision_path). Split into separate batches." + f"({n_img}/{len(raw_images)} have an image vision_path). Split into separate batches." + ) + if raw_videos is not None and 0 < n_vid < len(raw_videos): + raise ValueError( + "Reasoner batch mixes video-conditioned and text-only samples " + f"({n_vid}/{len(raw_videos)} have a video vision_path). Split into separate batches." ) - images: list[Image.Image] | None = cast(list[Image.Image], raw_images) if n_set == len(raw_images) else None + images: list[Image.Image] | None = cast(list[Image.Image], raw_images) if n_img == len(raw_images) else None + videos: list[dict[str, Any]] | None = ( + cast(list[dict[str, Any]], raw_videos) if raw_videos is not None and n_vid == len(raw_videos) else None + ) try: with sync_distributed_errors(): @@ -1715,6 +1763,7 @@ def _generate_reasoner_batch( prompts, max_new_tokens=sample_args_list[0].max_new_tokens, images=images, + videos=videos, do_sample=sample_args_list[0].do_sample, temperature=sample_args_list[0].temperature, top_k=sample_args_list[0].top_k, diff --git a/cosmos_framework/inference/inference_test.py b/cosmos_framework/inference/inference_test.py index d1f501b8..da10591e 100644 --- a/cosmos_framework/inference/inference_test.py +++ b/cosmos_framework/inference/inference_test.py @@ -169,6 +169,7 @@ def _make_reasoner_sample_args(**overrides: Any) -> SimpleNamespace: model_mode=ModelMode.REASONER, prompt="Describe a robotic arm.", vision_path=None, + video_fps=None, max_new_tokens=8, do_sample=False, temperature=1.0, @@ -189,7 +190,11 @@ def test_get_sample_data_reasoner_text_only() -> None: out = inference.get_sample_data(sample_args, model, device="cpu") - assert out == {"caption": ["Describe a robotic arm."], "reasoner_images": [None]} + assert out == { + "caption": ["Describe a robotic arm."], + "reasoner_images": [None], + "reasoner_videos": [None], + } @pytest.mark.L0 @@ -205,13 +210,34 @@ def test_get_sample_data_reasoner_with_image(tmp_path: Path) -> None: out = inference.get_sample_data(sample_args, model, device="cpu") - assert list(out) == ["caption", "reasoner_images"] + assert list(out) == ["caption", "reasoner_images", "reasoner_videos"] assert out["caption"] == ["Describe a robotic arm."] + assert out["reasoner_videos"] == [None] assert len(out["reasoner_images"]) == 1 assert out["reasoner_images"][0].size == (8, 8) assert out["reasoner_images"][0].mode == "RGB" +@pytest.mark.L0 +def test_get_sample_data_reasoner_with_video(monkeypatch: pytest.MonkeyPatch) -> None: + """A video ``vision_path`` routes through ``_decode_reasoner_video`` into ``reasoner_videos``. + + The decoder is monkeypatched (real decode needs torchvision + an actual clip); + this asserts the routing/contract, not the decode itself.""" + from cosmos_framework.inference import inference + + decoded = {"frames": ["F0", "F1"], "fps": 2.0} + monkeypatch.setattr(inference, "_decode_reasoner_video", lambda path, fps: decoded) + model = SimpleNamespace(input_caption_key="caption") + sample_args = _make_reasoner_sample_args(vision_path="/tmp/clip.mp4", video_fps=2.0) + + out = inference.get_sample_data(sample_args, model, device="cpu") + + assert out["caption"] == ["Describe a robotic arm."] + assert out["reasoner_videos"] == [decoded] + assert out["reasoner_images"] == [None] + + @pytest.mark.L0 def test_reasoner_defaults_json_round_trip() -> None: import json as _json diff --git a/docs/inference.md b/docs/inference.md index 15801476..cad7c44a 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -20,6 +20,7 @@ ______________________________________________________________________ - [Text](#text) - [Vision (Image/Video)](#vision-imagevideo) - [Action](#action) + - [Reasoner](#reasoner) - [Custom Defaults](#custom-defaults) - [Guardrails](#guardrails) - [Troubleshooting](#troubleshooting) @@ -213,6 +214,14 @@ The action output is written to `sample_outputs.json`. See the [Modes](#modes) table above for the action mode inputs/outputs and example files. +### Reasoner + +`model_mode=reasoner` generates text (written to `reasoner_text.txt`) from a prompt and an optional `vision_path`. The `vision_path` may point to an **image** (`.jpg`/`.png`/…) or a **video** (`.mp4`): a video is decoded and uniformly sampled into frames that condition the reasoner. + +- `video_fps`: frames per second to sample from the video (default: the decoder's default of 2.0). + +Examples: [`inputs/reasoner/reasoner.json`](../inputs/reasoner/reasoner.json) (text), [`inputs/reasoner/reasoner_image.json`](../inputs/reasoner/reasoner_image.json) (image), [`inputs/reasoner/reasoner_video.json`](../inputs/reasoner/reasoner_video.json) (video). + ### Custom Defaults To use your own default values instead of the built-in presets, pass a JSON file via the `defaults_file` field in your sample arguments: diff --git a/inputs/reasoner/reasoner_video.json b/inputs/reasoner/reasoner_video.json new file mode 100644 index 00000000..1c683088 --- /dev/null +++ b/inputs/reasoner/reasoner_video.json @@ -0,0 +1,5 @@ +{ + "model_mode": "reasoner", + "prompt": "Describe what is happening in this video in one sentence.", + "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/vision/robot_pouring.mp4" +} From 3d9c0878fd0dde76eac98161aed0493d85a036fd Mon Sep 17 00:00:00 2001 From: Xuanmeng Zhang Date: Fri, 10 Jul 2026 14:37:08 +0800 Subject: [PATCH 20/26] Add Cosmos3-Super VideoPhy-2 Reasoner SFT recipe (#101) - Add Cosmos3-Super VideoPhy-2 Reasoner SFT recipe. It fully fine-tunes` Qwen3-VL-32B` using FSDP full sharding and initializes from the Cosmos3-Super language model merged with the 32B visual tower. - The new experiment adds a launcher and TOML recipe. - Use a learning rate of `1e-6 `for both Nano and Super. - Set `lr_multipliers={"model.visual": 1.0} `to override the reasoner default of `0.1` and train the visual projector at the base learning rate. The previous `mm_projector` and `merger` keys matched no `Qwen3-VL` parameters because the projector is named `model.visual.merger`, so the intended `20x` multiplier was never applied. - Document the Super recipe in `docs/training.md `and `examples/README.md`. --- .../reasoner/experiment/videophy2_sft_nano.py | 24 ++++- .../experiment/videophy2_sft_super.py | 65 +++++++++++ docs/training.md | 36 +++++++ examples/README.md | 1 + examples/launch_sft_videophy2_super.sh | 55 ++++++++++ .../toml/sft_config/videophy2_sft_nano.toml | 3 + .../toml/sft_config/videophy2_sft_super.toml | 101 ++++++++++++++++++ 7 files changed, 284 insertions(+), 1 deletion(-) create mode 100644 cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_super.py create mode 100755 examples/launch_sft_videophy2_super.sh create mode 100644 examples/toml/sft_config/videophy2_sft_super.toml diff --git a/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py index 1405ac7f..1123e5fb 100644 --- a/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py +++ b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py @@ -160,11 +160,33 @@ def _dl(dataset_key, split, num_workers, persistent_workers=False, pin_memory=Fa grad_accum_iter=8, ), optimizer=dict( + # Uniform 1e-6 across both tiers. Nano is stable at 1e-5 too, but the 32B + # super run spikes transiently (loss->6.6 @ iter 23) at 1e-5 before + # recovering to the same endpoint; 1e-6 reaches the identical final loss + # with no instability, so both tiers ship at 1e-6. (Super deepcopies this + # block, so setting it here covers both.) lr=1e-6, fused=True, weight_decay=0.05, betas=[0.9, 0.999], - lr_multipliers={"mm_projector": 20.0, "merger": 20.0}, + # NOTE: the ViT is FROZEN (freeze_vision_encoder=True below), so the only + # trainable "model.visual.*" params are the projector (model.visual.merger + # + deepstack_merger_list). This key therefore sets the PROJECTOR LR only, + # to 1.0x optimizer.lr -- giving a uniform LR across projector + LLM + + # lm_head. It is NOT touching the frozen ViT backbone. + # + # Why the key is "model.visual" (broad) and not "model.visual.merger": + # the reasoner default optimizer (configs/base/reasoner/defaults/optimizer.py) + # ships lr_multipliers={"model.visual": 0.1}. Hydra merges dicts base-first + # (so that key stays PINNED FIRST) and _build_params_with_metadata + # (utils/generator/optimizer.py) is first-substring-match-wins -- so a + # narrower "model.visual.merger" key is shadowed by "model.visual" and never + # matches (verified: the merger still resolves to 0.1x). Overriding the SAME + # "model.visual" key to 1.0 is the only way to cancel the 0.1x default; with + # the ViT frozen that reaches only the merger. (To give an UNFROZEN ViT its + # own LR you'd likewise override this "model.visual" value -- a narrower key + # can't out-prioritize it.) + lr_multipliers={"model.visual": 1.0}, ), scheduler=dict( warm_up_steps=[5], diff --git a/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_super.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_super.py new file mode 100644 index 00000000..83d5d7e2 --- /dev/null +++ b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_super.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""VideoPhy-2 SFT recipe (Cosmos3-Super tier): Qwen3-VL-32B full fine-tune. + +Super-tier counterpart of ``videophy2_sft_nano``. Reuses that recipe's VideoPhy-2 +``LocalSFTDataset`` + ``CosmosDataLoader`` dataflow verbatim (imported + deepcopied, +the same variant idiom ``llava_ov_vlm.py`` uses) and changes only what the larger +backbone needs: + + * ``vlm_policy`` ``qwen3_vl_8b_instruct`` -> ``qwen3_vl_32b_instruct`` + (``Qwen/Qwen3-VL-32B-Instruct``) — the visual tower + config the Cosmos3-Super + Reasoner is built on, mirroring how the nano recipe rides the 8B tower. + * FSDP full-shard across every rank (``data_parallel_shard_degree=-1``) so the 32B + weights + optimizer state fit, instead of the nano recipe's fixed dp=8. This is + the same super-tier sharding switch ``vision_sft_super`` makes, and lets the recipe + run unchanged on a 4-GPU (e.g. GB200x4) or 8-GPU allocation. + +Still a full fine-tune (no LoRA): the freeze config is inherited from the nano recipe +(vision encoder frozen, LM + mm_projector trained). + +Launch via ``examples/launch_sft_videophy2_super.sh`` after +``prepare_videophy2_from_hf`` populates ``$VIDEOPHYSICS_ROOT``, supplying the merged +Cosmos3-Super Reasoner checkpoint through ``VLM_SAFETENSORS_PATH`` (see the launch shell). +""" + +from __future__ import annotations + +import copy + +from hydra.core.config_store import ConfigStore + +# Importing the nano module registers `videophy2_sft_nano` and pulls in the shared +# VideoPhy-2 dataflow helpers; we clone its LazyDict rather than re-declaring them. +from cosmos_framework.configs.base.reasoner.experiment.videophy2_sft_nano import videophy2_sft_nano + +cs = ConfigStore.instance() + + +videophy2_sft_super = copy.deepcopy(videophy2_sft_nano) + +# Backbone: nano 8B -> super 32B. The vlm_policy override lives in the Hydra +# `defaults` list; rewrite that one entry's value in place, leaving the rest of the +# recipe (checkpoint backend, callbacks, dataflow) untouched. +for _default in videophy2_sft_super["defaults"]: + if not isinstance(_default, str) and "override /vlm_policy" in _default: + _default["override /vlm_policy"] = "qwen3_vl_32b_instruct" + +# 32B full fine-tune: shard model + optimizer state across every rank (FSDP full +# shard, auto-sized from WORLD_SIZE) instead of the nano recipe's fixed dp_shard=8, +# so the recipe fits on 4- or 8-GPU nodes. +# examples/toml/sft_config/videophy2_sft_super.toml is authoritative at launch and +# repeats these; keeping them here lets `experiment=videophy2_sft_super` run standalone. +videophy2_sft_super.model.config.parallelism.data_parallel_shard_degree = -1 +videophy2_sft_super.model.config.parallelism.data_parallel_replicate_degree = 1 + + +for _item in [videophy2_sft_super]: + experiment_name = [name.lower() for name, value in globals().items() if value is _item][0] + if "job" not in _item: + _item["job"] = dict(name=experiment_name + "_${now:%Y-%m-%d}_${now:%H-%M-%S}") + else: + _item["job"]["name"] = experiment_name + "_${now:%Y-%m-%d}_${now:%H-%M-%S}" + + cs.store(group="experiment", package="_global_", name=experiment_name, node=_item) diff --git a/docs/training.md b/docs/training.md index 5b11e5ff..056ceb0a 100644 --- a/docs/training.md +++ b/docs/training.md @@ -115,6 +115,21 @@ python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ +
Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Super) + +Super-tier counterpart of the VideoPhy-2 recipe above: same 1–5 physical-plausibility scoring on [videophysics/videophy2_train](https://huggingface.co/datasets/videophysics/videophy2_train), but a full fine-tune of the 32B backbone. `[job].task = "vlm"`. Bootstraps from `Cosmos3-Super`'s language-model weights merged onto the public Qwen3-VL-32B-Instruct visual tower; the merged HF directory is consumed via `[model.backbone].safetensors_path` (plumbed by `VLM_SAFETENSORS_PATH`). Full-shard FSDP across all ranks, so it runs on a 4-GPU (e.g. GB200x4) or 8-GPU allocation. + +Launch shell: `examples/launch_sft_videophy2_super.sh` + +```shell +# Step 1 (data): same as the nano recipe — materialize the public HF dataset into +# the canonical local layout (videophy2_{train,val}/{meta.json, media/, text/}). +python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ + --out_root examples/data/videophysics --split both +``` + +
+ ## Step 2 — Prepare checkpoint Convert the base checkpoint to [PyTorch Distributed Checkpoint (DCP)](https://pytorch.org/docs/stable/distributed.checkpoint.html) format. `cosmos_framework.scripts.convert_model_to_dcp` ships in the unified `cosmos_framework/` package, so this step runs from the repo root (with the environment activated per [Setup](./setup.md)). @@ -144,6 +159,16 @@ python -m cosmos_framework.scripts.convert_model_to_vlm_safetensors \ -o examples/checkpoints/Cosmos3-Nano-VLM ``` +**Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Super):** Same converter, but merge the `Cosmos3-Super` LM onto the **32B** Qwen3-VL visual tower via `--vlm-model-name Qwen/Qwen3-VL-32B-Instruct`. + +```shell +# Step 2 (Super VLM checkpoint): merge Cosmos3-Super LM onto the Qwen3-VL-32B visual tower. +python -m cosmos_framework.scripts.convert_model_to_vlm_safetensors \ + --checkpoint-path Cosmos3-Super \ + --vlm-model-name Qwen/Qwen3-VL-32B-Instruct \ + -o examples/checkpoints/Cosmos3-Super-VLM +``` + ## Step 3 — Run training **Weights & Biases (optional):** every recipe TOML defaults to `job.wandb_mode = "disabled"`. To log a run to W&B, flip that field to `"online"` in the TOML and export `WANDB_API_KEY` in your environment before launching. @@ -165,6 +190,7 @@ Each launcher's default paths come from the `DATASET_PATH` + `BASE_CHECKPOINT_PA | `launch_sft_vision_super.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Super` | | `launch_sft_llava_ov.sh` | Reasoner SFT | (none; dataset streams from HF Hub) | (none; backbone fetched at startup, or set `VLM_SAFETENSORS_PATH`) | | `launch_sft_videophy2_nano.sh` | Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; set `VLM_SAFETENSORS_PATH` env) | +| `launch_sft_videophy2_super.sh`| Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; set `VLM_SAFETENSORS_PATH` env — Cosmos3-Super-VLM merge) | `WAN_VAE_PATH` defaults to `examples/checkpoints/wan22_vae/Wan2.2_VAE.pth` for every non-reasoner recipe. @@ -177,6 +203,16 @@ export VLM_SAFETENSORS_PATH=$PWD/examples/checkpoints/Cosmos3-Nano-VLM bash examples/launch_sft_videophy2_nano.sh ``` +**Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Super):** + +```shell +# Same as the nano recipe, but point VLM_SAFETENSORS_PATH at the Cosmos3-Super +# merge from Step 2. On a 4-GPU node (e.g. GB200x4) also set NPROC_PER_NODE=4. +export VIDEOPHYSICS_ROOT=$PWD/examples/data/videophysics +export VLM_SAFETENSORS_PATH=$PWD/examples/checkpoints/Cosmos3-Super-VLM +bash examples/launch_sft_videophy2_super.sh +``` + #### Overriding the defaults If you'd rather put data or checkpoints on a different filesystem (e.g. a faster SSD or shared mount), download to your chosen path in Step 1 / convert the DCP to your chosen path in Step 2, then export the matching env var(s) before launching: diff --git a/examples/README.md b/examples/README.md index 0e51140f..08c219a8 100644 --- a/examples/README.md +++ b/examples/README.md @@ -16,3 +16,4 @@ This directory contains: | Vision SFT LoRA (Cosmos3-Super) | `launch_sft_vision_super.sh` | | Reasoner Alignment SFT | `launch_sft_llava_ov.sh` | | Reasoner Alignment SFT (Cosmos3-Nano) | `launch_sft_videophy2_nano.sh` | +| Reasoner Alignment SFT (Cosmos3-Super) | `launch_sft_videophy2_super.sh` | diff --git a/examples/launch_sft_videophy2_super.sh b/examples/launch_sft_videophy2_super.sh new file mode 100755 index 00000000..785b896b --- /dev/null +++ b/examples/launch_sft_videophy2_super.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# Structured-TOML launch for videophy2_sft_super (VLM dialog SFT on VideoPhy-2 +# via CosmosDataLoader, Cosmos3-Super tier / Qwen3-VL-32B full fine-tune). Drives +# cosmos_framework.scripts.train against +# examples/toml/sft_config/videophy2_sft_super.toml. +# +# [job].task = "vlm" — picks cosmos_framework/configs/base/reasoner/config.py as the base config. +# +# Required env: +# VIDEOPHYSICS_ROOT dir containing videophy2_train/ and videophy2_val/ +# (each with meta.json + media/ + text/). Populate via +# `python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf`. +# +# Optional env: +# HF_TOKEN for gated Qwen3-VL-32B-Instruct downloads. +# VLM_SAFETENSORS_PATH local directory of pre-converted Qwen3-VL safetensors +# (e.g. Cosmos3-Super LM merged with the Qwen3-VL-32B visual +# tower via `cosmos_framework.scripts.convert_model_to_vlm_safetensors +# --checkpoint-path Cosmos3-Super --vlm-model-name Qwen/Qwen3-VL-32B-Instruct`). +# When set, plumbed to backbone.safetensors_path via a +# tail override. When unset, the framework falls back +# to the public Qwen/Qwen3-VL-32B-Instruct HF snapshot. +# NPROC_PER_NODE torchrun GPUs per node; default 8. Set 4 on a GB200x4 node. +# +# Usage (8-GPU allocation, inside the training container, from the repo root): +# VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_super.sh +# # on a 4-GPU node (e.g. GB200x4): +# NPROC_PER_NODE=4 VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_super.sh + +TOML_FILE="examples/toml/sft_config/videophy2_sft_super.toml" + +# Super-variant allocator tweak: expandable_segments so the 32B backbone fits +# without OOM during compile/decode. (Unlike launch_sft_vision_super.sh we do NOT +# clear LD_LIBRARY_PATH — this reasoner recipe decodes VideoPhy-2 clips with +# torchcodec, which dlopen()s the CUDA NPP + FFmpeg libs off LD_LIBRARY_PATH; the +# nano videophy2 launcher leaves it untouched for the same reason.) +export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" + +TAIL_OVERRIDES=( + ${EXTRA_TAIL_OVERRIDES:-} +) + +# When VLM_SAFETENSORS_PATH is set, plumb it to backbone.safetensors_path so the +# framework loads weights from the local snapshot (e.g. a Cosmos3-Super LM merged +# with the Qwen3-VL-32B visual tower via +# `cosmos_framework.scripts.convert_model_to_vlm_safetensors`) while keeping the +# public HF model_name for tokenizer/architecture discovery. +if [[ -n "${VLM_SAFETENSORS_PATH:-}" ]]; then + TAIL_OVERRIDES+=("model.config.policy.backbone.safetensors_path=$VLM_SAFETENSORS_PATH") +fi + +source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh" diff --git a/examples/toml/sft_config/videophy2_sft_nano.toml b/examples/toml/sft_config/videophy2_sft_nano.toml index 35b8e785..81f43964 100644 --- a/examples/toml/sft_config/videophy2_sft_nano.toml +++ b/examples/toml/sft_config/videophy2_sft_nano.toml @@ -55,6 +55,9 @@ eps = 1.0e-8 fused = true lr = 1.0e-6 weight_decay = 0.1 +# NOTE: per-part LR multipliers live in the experiment (videophy2_sft_nano.py), +# NOT here -- the toml->hydra override path parses a dict key's "." as nesting, +# so a dotted key like "model.visual" can't be set from toml. Edit the experiment. [scheduler] cycle_lengths = [50] diff --git a/examples/toml/sft_config/videophy2_sft_super.toml b/examples/toml/sft_config/videophy2_sft_super.toml new file mode 100644 index 00000000..9029e246 --- /dev/null +++ b/examples/toml/sft_config/videophy2_sft_super.toml @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# videophy2_sft_super — VLM dialog SFT on VideoPhy-2 via CosmosDataLoader, +# Cosmos3-Super tier (Qwen3-VL-32B full fine-tune). +# Base config = cosmos_framework/configs/base/reasoner/config.py (selected by [job].task="vlm"). +# +# Super-tier counterpart of videophy2_sft_nano: same dataset + dataflow, but the +# 32B backbone (Qwen/Qwen3-VL-32B-Instruct) and FSDP full-shard across every rank +# (data_parallel_shard_degree=-1, auto from WORLD_SIZE) so it fits on a 4-GPU +# (e.g. GB200x4) or 8-GPU allocation. Still a full fine-tune (no LoRA). +# +# Dataset prep: +# python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ +# --out_root $VIDEOPHYSICS_ROOT --split train # and again with --split val +# +# Required env at launch: VIDEOPHYSICS_ROOT (read by the experiment Python). +# Supply the merged Cosmos3-Super Reasoner checkpoint via VLM_SAFETENSORS_PATH. +# +# Example launch: +# bash examples/launch_sft_videophy2_super.sh + +[job] +task = "vlm" +experiment = "videophy2_sft_super" +project = "cosmos3" +group = "vlm_videophy2_sft" +name = "videophy2_sft_super" +wandb_mode = "disabled" + +[model] +attn_implementation = "cosmos" +precision = "bfloat16" # was [model.parallelism].precision + +[model.backbone] +model_name = "Qwen/Qwen3-VL-32B-Instruct" + +[model.ema] +enabled = false +rate = 0.1 +iteration_shift = 0 + +[model.parallelism] +data_parallel_shard_degree = -1 # super: FSDP full shard, auto from WORLD_SIZE (4- or 8-GPU) +data_parallel_replicate_degree = 1 +context_parallel_shard_degree = 1 # raise to 2 (needs even GPU count) if the 32B run OOMs +cfg_parallel_shard_degree = 1 + +[model.compile] +enabled = false # was [model.parallelism].use_torch_compile +compile_dynamic = true + +[model.activation_checkpointing] +mode = "full" +save_ops_regex = ["fmha"] +preserve_rng_state = true +determinism_check = "default" + +[optimizer] +betas = [0.9, 0.95] +eps = 1.0e-8 +fused = true +lr = 1.0e-6 +weight_decay = 0.1 +# NOTE: per-part LR multipliers live in the experiment (videophy2_sft_nano.py), +# NOT here -- the toml->hydra override path parses a dict key's "." as nesting, +# so a dotted key like "model.visual" can't be set from toml. Edit the experiment. + +[scheduler] +cycle_lengths = [50] +f_max = [1.0] +f_min = [0.1] +f_start = [0.05] +verbosity_interval = 0 +warm_up_steps = [5] + +[trainer] +distributed_parallelism = "fsdp" +grad_accum_iter = 8 +logging_iter = 1 +max_iter = 50 + +[trainer.callbacks.compile_tokenizer] +compile_after_iterations = 3 +enabled = false + +[trainer.callbacks.grad_clip] +clip_norm = 1.0 +force_finite = false + +[checkpoint] +keys_to_skip_loading = [] +load_path = "???" +save_iter = 100 + +[dataloader_train] +# Routed by PATH_REMAPS["vlm"] onto the CosmosDataLoader's nested PoolPackingBatcher: +# max_samples_per_batch -> dataloader_train.batcher.max_batch_size +# max_sequence_length -> dataloader_train.batcher.max_tokens +max_samples_per_batch = 1 +max_sequence_length = 16000 From 1e139fb641ad8bc9671fae527bcc61fe3a2a4ad4 Mon Sep 17 00:00:00 2001 From: lfengad Date: Sat, 11 Jul 2026 19:34:30 +0800 Subject: [PATCH 21/26] feat(transfer-sft): control-conditioned SFT dataset + curator support (#104) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports Cosmos Transfer (control-conditioned) SFT post-training to the OSS layout (from i4 MR !10217). ## Changes - `curator_to_sft_jsonl.py`: add `--control-type` (edge/blur/depth/seg) and `--control-path-root`. Emits `control_type` on every row and `control_path` for precomputed depth/seg; clips with no matching precomputed file are dropped with reason `missing_control_path`. Summary records the control config. - `transfer_sft_dataset.py` (new): `TransferSFTDataset(SFTDataset)` — computes edge/blur on-the-fly or loads precomputed depth/seg control videos, returns `video=[control, target]` with a shared-temporal-position `SequencePlan`; plus `get_transfer_sft_dataset()` loader. - Tests: 10 new cases covering edge/blur/depth control paths and the converter end-to-end. ## Verification - Live import of `transfer_sft_dataset` in the i4 container: OK (base class = `SFTDataset`). - Curator converter tests: **33 passed**. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 (1M context) --- .../local_datasets/transfer_sft_dataset.py | 394 ++++++++++++++++++ .../scripts/curator_to_sft_jsonl.py | 122 +++++- .../scripts/curator_to_sft_jsonl_test.py | 141 +++++++ 3 files changed, 650 insertions(+), 7 deletions(-) create mode 100644 cosmos_framework/data/generator/local_datasets/transfer_sft_dataset.py diff --git a/cosmos_framework/data/generator/local_datasets/transfer_sft_dataset.py b/cosmos_framework/data/generator/local_datasets/transfer_sft_dataset.py new file mode 100644 index 00000000..20f926c7 --- /dev/null +++ b/cosmos_framework/data/generator/local_datasets/transfer_sft_dataset.py @@ -0,0 +1,394 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""Transfer SFT dataset: (control, target) pairs for Cosmos Transfer post-training. + +Subclasses SFTDataset — all frame loading, rank partitioning, and iteration +logic is inherited. Only ``process_one_sample()`` is overridden to add the +control signal and repackage the output as ``video=[control, target]``. + +Supported control types +----------------------- +- ``edge``: Canny edge map, computed on-the-fly. Recommended starting point. +- ``blur``: Gaussian blur, computed on-the-fly. +- ``depth``: Precomputed depth video at ``metadata["control_path"]``. +- ``seg``: Precomputed segmentation video at ``metadata["control_path"]``. + +Workflow +-------- +1. Run cosmos-curator with ``--upload-clip-info-in-chunks``. +2. Convert to transfer JSONL:: + + python -m cosmos_framework.scripts.curator_to_sft_jsonl \\ + --curator-output outputs/curator_split/ \\ + -o outputs/transfer_sft.jsonl \\ + --control-type edge + +3. Wire into an experiment config inheriting from exp302 (video) or exp301 (image):: + + from cosmos_framework.data.generator.local_datasets.transfer_sft_dataset import ( + get_transfer_sft_dataset, + ) + dataset=L(get_transfer_sft_dataset)( + jsonl_paths="s3://your-bucket/transfer_sft.jsonl", + resolution="480", + num_video_frames=61, + control_type="edge", + tokenizer_config=..., + ) +""" + +import gzip +import io +import json +import os +import random +import tempfile +from pathlib import Path +from typing import Any, Optional + +import boto3 +import cv2 +import numpy as np +import torch + +from cosmos_framework.data.generator.local_datasets.helper import ( + client_config, + download_from_s3, + ffmpeg_decode_video, + get_aspect_ratio, + parse_s3_url, +) +from cosmos_framework.data.generator.local_datasets.sft_dataset import SFTDataset +from cosmos_framework.data.generator.sequence_packing import SequencePlan +from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO +from cosmos_framework.utils import log +from cosmos_framework.utils.flags import INTERNAL + +_ON_THE_FLY_CONTROLS = frozenset({"edge", "blur"}) +_PRECOMPUTED_CONTROLS = frozenset({"depth", "seg"}) + +# Canny threshold presets matching AddControlInputEdge. +_EDGE_THRESHOLDS = [ + (20, 50), # very_low + (50, 100), # low + (100, 200), # medium + (200, 300), # high + (300, 400), # very_high +] + + +def _compute_edge(frames: np.ndarray) -> np.ndarray: + """Canny edge map from uint8 [T, H, W, 3] RGB → uint8 [T, H, W, 3].""" + low, high = random.choice(_EDGE_THRESHOLDS) + edges = np.zeros_like(frames) + for t in range(frames.shape[0]): + gray = cv2.cvtColor(frames[t], cv2.COLOR_RGB2GRAY) + edge = cv2.Canny(gray, low, high) + edges[t] = np.stack([edge, edge, edge], axis=-1) + return edges + + +def _compute_blur(frames: np.ndarray) -> np.ndarray: + """Gaussian blur from uint8 [T, H, W, 3] → uint8 [T, H, W, 3].""" + ksize = random.choice([11, 21, 31, 41]) + blurred = np.zeros_like(frames) + for t in range(frames.shape[0]): + blurred[t] = cv2.GaussianBlur(frames[t], (ksize, ksize), 0) + return blurred + + +def _normalize(frames: np.ndarray) -> torch.Tensor: + """uint8 [T, H, W, 3] → float32 [3, T, H, W] in [-1, 1].""" + x = torch.from_numpy(np.ascontiguousarray(frames)).float() # [T, H, W, 3] + return (x.permute(3, 0, 1, 2) / 255.0 - 0.5) / 0.5 # [3, T, H, W] + + +class TransferSFTDataset(SFTDataset): + """SFT dataset for Cosmos Transfer post-training. + + Inherits all frame loading, rank partitioning, and iteration from + ``SFTDataset``. Overrides ``process_one_sample()`` to compute or load a + control signal and return ``video=[control, target]`` in the format + expected by the Transfer training loop. + """ + + def __init__( + self, + *args, + control_type: str = "edge", + **kwargs, + ): + assert control_type in _ON_THE_FLY_CONTROLS | _PRECOMPUTED_CONTROLS, f"Unknown control_type={control_type!r}" + super().__init__(*args, **kwargs) + self.control_type = control_type + + def _load_precomputed_control( + self, + control_path: str, + start_frame: int, + end_frame: int, + temporal_interval: int, + resize_hw: tuple[int, int], + crop_y: int, + crop_x: int, + target_h: int, + target_w: int, + ) -> np.ndarray: + """Decode a precomputed control video and extract the matching frame window.""" + ctrl_bytes = download_from_s3(self.s3_client, control_path) + if ctrl_bytes is None: + raise RuntimeError(f"Failed to read control video: {control_path}") + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=True) as tmp: + tmp.write(ctrl_bytes) + tmp.flush() + frames = list(ffmpeg_decode_video(tmp.name, scale_hw=resize_hw, num_threads=2)) + ctrl = np.stack(frames, axis=0) # [T_full, H, W, 3] + chunk = [] + for idx in range(ctrl.shape[0]): + if idx < start_frame: + continue + if idx > end_frame: + break + if (idx - start_frame) % temporal_interval == 0: + chunk.append(ctrl[idx]) + if not chunk: + return np.empty((0, target_h, target_w, 3), dtype=np.uint8) + arr = np.stack(chunk, axis=0) + return arr[:, crop_y : crop_y + target_h, crop_x : crop_x + target_w] + + def process_one_sample(self, metadata: dict) -> dict | None: + sample = super().process_one_sample(metadata) + if sample is None: + return None + + # Parent returns video as uint8 [3, T, H, W]. Convert to [T, H, W, 3] + # numpy for on-the-fly control computation, then normalize both streams. + video_uint8 = sample["video"] # [3, T, H, W] uint8 + target_np = video_uint8.permute(1, 2, 3, 0).numpy() # [T, H, W, 3] uint8 + T, H, W = target_np.shape[:3] + + if self.control_type == "edge": + control_np = _compute_edge(target_np) + elif self.control_type == "blur": + control_np = _compute_blur(target_np) + else: + ctrl_path = metadata.get("control_path") + if not ctrl_path: + log.warning(f"No control_path for {metadata['uuid']}, skipping") + return None + # Re-derive the resize / crop parameters to align the control video. + input_w, input_h = metadata["width"], metadata["height"] + target_w_out, target_h_out = self.output_sizes[metadata["aspect_ratio"]] + resize_ratio = max(target_w_out / input_w, target_h_out / input_h) + resize_h = round(input_h * resize_ratio) + resize_w = round(input_w * resize_ratio) + crop_y = round((resize_h - target_h_out) / 2) + crop_x = round((resize_w - target_w_out) / 2) + # Derive end_frame from T (post-truncation) rather than sample["frame_end"] + # (pre-truncation). The parent truncates to target_t = (N-1)//tcf*tcf+1 + # frames, which may be less than N when (num_video_frames-1) % tcf != 0. + # Using frame_end directly would request more frames than T from the + # control video, causing the shape-mismatch check below to drop every + # depth/seg sample silently for non-standard frame counts. + effective_end_frame = sample["frame_start"] + (T - 1) * sample["num_multiplier"] + try: + control_np = self._load_precomputed_control( + ctrl_path, + start_frame=sample["frame_start"], + end_frame=effective_end_frame, + temporal_interval=sample["num_multiplier"], + resize_hw=(resize_h, resize_w), + crop_y=crop_y, + crop_x=crop_x, + target_h=target_h_out, + target_w=target_w_out, + ) + except Exception as exc: + log.warning(f"Failed to load control for {metadata['uuid']}: {exc}") + return None + if control_np.shape[0] == 0 or control_np.shape[0] != T: + log.warning(f"Control/target frame count mismatch for {metadata['uuid']}: {control_np.shape[0]} vs {T}") + return None + + control_tensor = _normalize(control_np) # [3, T, H, W] float32 [-1, 1] + target_tensor = _normalize(target_np) # [3, T, H, W] float32 [-1, 1] + + sample["video"] = [control_tensor, target_tensor] + sample["dataset_name"] = "video_transfer" + sample["selected_caption_type"] = "transfer_caption" + + # Always provide a SequencePlan with shared temporal positions. + cond_indexes = sample["sequence_plan"].condition_frame_indexes_vision if "sequence_plan" in sample else [] + sample["sequence_plan"] = SequencePlan( + has_text=True, + has_vision=True, + condition_frame_indexes_vision=cond_indexes, + share_vision_temporal_positions=True, + ) + + # Two vision streams → wrap image_size as a list. + sample["image_size"] = [sample["image_size"], sample["image_size"]] + + return sample + + +def _load_transfer_metadata( + s3_client: Any, + jsonl_url: str, + min_frames: int, + uuid_prefix: str = "", + min_short_edge: int = 0, + control_type: str = "edge", +) -> list[dict]: + """Load transfer SFT metadata, passing control_path through for depth/seg.""" + log.info(f"Loading transfer SFT metadata from {jsonl_url}", rank0_only=False) + metadata_list: list[dict] = [] + num_raw = 0 + + with io.BytesIO() as buffer: + if jsonl_url.startswith("s3://"): + bucket, key = parse_s3_url(jsonl_url) + s3_client.download_fileobj(Bucket=bucket, Key=key, Fileobj=buffer) + else: + path = Path(jsonl_url).absolute() + jsonl_url = str(path) + buffer.write(path.read_bytes()) + buffer.seek(0) + + line_iter = gzip.open(buffer, "rb") if jsonl_url.endswith(".gz") else buffer + for line in line_iter: + num_raw += 1 + record = json.loads(line.decode("utf-8")) + uuid = f"{uuid_prefix}{record['uuid']}" if uuid_prefix else record["uuid"] + + if record["duration"] > 61.0: + continue + if min_short_edge > 0 and min(record["width"], record["height"]) < min_short_edge: + continue + if control_type in _PRECOMPUTED_CONTROLS and not record.get("control_path"): + continue + + kept_windows = [ + w for w in (record.get("t2w_windows") or []) if w["end_frame"] - w["start_frame"] + 1 >= min_frames + ] + if not kept_windows: + continue + + vision_path = record["vision_path"] + if "://" not in vision_path and not vision_path.startswith("/"): + vision_path = f"{os.path.dirname(jsonl_url)}/{vision_path}" + + entry: dict[str, Any] = { + "uuid": uuid, + "vision_path": vision_path, + "width": record["width"], + "height": record["height"], + "nb_frames": record.get("nb_frames"), + "framerate": record.get("framerate"), + "aspect_ratio": get_aspect_ratio(record["width"], record["height"]), + "t2w_windows": kept_windows, + } + if record.get("control_path"): + ctrl_path = record["control_path"] + if "://" not in ctrl_path and not ctrl_path.startswith("/"): + ctrl_path = f"{os.path.dirname(jsonl_url)}/{ctrl_path}" + entry["control_path"] = ctrl_path + metadata_list.append(entry) + + log.info( + f"Transfer SFT: kept {len(metadata_list)} / {num_raw} records from {jsonl_url}", + rank0_only=False, + ) + return metadata_list + + +def get_transfer_sft_dataset( + jsonl_paths: str | list[str], + resolution: str = "480", + num_video_frames: int = 61, + control_type: str = "edge", + temporal_interval_mode: str = "entire_chunk", + frame_selection_mode: str = "center", + tokenizer_config: Optional[Any] = None, + cfg_dropout_rate: float = 0.1, + use_system_prompt: bool = False, + append_duration_fps_timestamps: bool = True, + append_resolution_info: bool = True, + min_short_edge: int = 0, + conditioning_config: dict[int, float] | None = None, + temporal_compression_factor: int = 4, + **kwargs, +) -> TransferSFTDataset: + """Create a TransferSFTDataset from one or more JSONL files. + + Args: + jsonl_paths: Local path(s) or ``s3://`` URI(s) to transfer SFT JSONL(s) + produced by ``curator_to_sft_jsonl.py --control-type ``. + Multiple files are concatenated; uuids are prefixed to avoid collisions. + resolution: Output resolution bucket, e.g. ``"480"`` or ``"720"``. + num_video_frames: Frames per clip (≥ 61 to pass the converter hard filter). + control_type: ``"edge"`` | ``"blur"`` | ``"depth"`` | ``"seg"``. + Edge and blur are computed on-the-fly; depth and seg require + ``control_path`` entries in the JSONL. + temporal_interval_mode: ``"force_one"`` | ``"max_30fps"`` | ``"entire_chunk"``. + frame_selection_mode: ``"center"`` | ``"first"`` | ``"random"``. + tokenizer_config: Tokenizer config (same as ``get_sft_dataset``). + cfg_dropout_rate: Caption dropout rate for CFG training. + use_system_prompt: Include system prompt in tokenization. + append_duration_fps_timestamps: Append duration/FPS text to captions. + append_resolution_info: Append resolution text to captions. + min_short_edge: Drop videos whose shortest edge is below this value. + conditioning_config: I2V conditioning frame distribution, e.g. + ``{0: 0.7, 1: 0.2, 2: 0.1}`` (70% unconditioned). + temporal_compression_factor: VAE temporal compression factor (default 4). + """ + log.info(f"get_transfer_sft_dataset: ignoring unknown kwargs: {list(kwargs)}") + assert resolution in VIDEO_RES_SIZE_INFO, f"Unknown resolution {resolution!r}; known: {sorted(VIDEO_RES_SIZE_INFO)}" + + if isinstance(jsonl_paths, str): + jsonl_paths = [jsonl_paths] + + if INTERNAL: + with open("credentials/gcs.secret") as f: + credentials = json.load(f) + else: + credentials = {} + + s3_client = boto3.client("s3", **credentials, config=client_config) + + metadata_list: list[dict] = [] + for idx, jsonl_url in enumerate(jsonl_paths): + prefix = f"{idx}/" if len(jsonl_paths) > 1 else "" + metadata_list.extend( + _load_transfer_metadata( + s3_client, + jsonl_url, + min_frames=61, + uuid_prefix=prefix, + min_short_edge=min_short_edge, + control_type=control_type, + ) + ) + + total_windows = sum(len(m["t2w_windows"]) for m in metadata_list) + log.info( + f"Transfer SFT: {len(metadata_list)} videos, {total_windows} windows, " + f"control_type={control_type}, resolution={resolution}" + ) + + return TransferSFTDataset( + metadata=metadata_list, + num_video_frames=num_video_frames, + resolution=resolution, + s3_credentials=credentials, + control_type=control_type, + temporal_interval_mode=temporal_interval_mode, + frame_selection_mode=frame_selection_mode, + tokenizer_config=tokenizer_config, + cfg_dropout_rate=cfg_dropout_rate, + use_system_prompt=use_system_prompt, + append_duration_fps_timestamps=append_duration_fps_timestamps, + append_resolution_info=append_resolution_info, + conditioning_config=conditioning_config, + temporal_compression_factor=temporal_compression_factor, + ) diff --git a/cosmos_framework/scripts/curator_to_sft_jsonl.py b/cosmos_framework/scripts/curator_to_sft_jsonl.py index e70d0469..0ae07573 100644 --- a/cosmos_framework/scripts/curator_to_sft_jsonl.py +++ b/cosmos_framework/scripts/curator_to_sft_jsonl.py @@ -14,19 +14,35 @@ silently at train time (so dataset counts match), and writes a sidecar ``.summary.json`` with per-reason drop counts. +Transfer (control-conditioned) training +---------------------------------------- +Pass ``--control-type`` to produce a JSONL suitable for ``TransferSFTDataset`` +(``transfer_sft_dataset.py``) in addition to the standard fields above. Two +extra fields are written per row: + +- ``control_type`` — the control signal to use (``edge``, ``blur``, ``depth``, ``seg``). +- ``control_path`` — path to the precomputed control video; only present for + ``depth`` and ``seg`` (edge/blur are computed on-the-fly from ``vision_path``). + Usage ----- + # Standard T2V/I2V/V2V SFT python -m cosmos_framework.scripts.curator_to_sft_jsonl \\ --curator-output outputs/curator_split/ \\ - -o outputs/curator_split/cosmos3_sft.jsonl + -o outputs/curator_split/cosmos3_sft.jsonl \\ + --caption-model qwen --enhanced-caption-model qwen_lm - # With explicit caption-model preference (e.g. when curator ran with - # both qwen captioning and qwen_lm enhancement): + # Transfer SFT — edge control (no precomputed files needed) python -m cosmos_framework.scripts.curator_to_sft_jsonl \\ --curator-output outputs/curator_split/ \\ - -o outputs/curator_split/cosmos3_sft.jsonl \\ - --caption-model qwen \\ - --enhanced-caption-model qwen_lm + -o outputs/curator_split/cosmos3_transfer_sft.jsonl \\ + --control-type edge + + # Transfer SFT — depth control (requires precomputed depth videos) + python -m cosmos_framework.scripts.curator_to_sft_jsonl \\ + --curator-output outputs/curator_split/ \\ + -o outputs/curator_split/cosmos3_transfer_sft.jsonl \\ + --control-type depth --control-path-root outputs/depth_maps/ Curator must have been run with ``--upload-clip-info-in-chunks`` so that ``metas_jsonl/v0/`` is populated; otherwise the converter has no input. @@ -38,7 +54,7 @@ from collections import Counter from collections.abc import Iterator from pathlib import Path -from typing import Annotated, Any +from typing import Annotated, Any, Literal import tyro @@ -52,6 +68,11 @@ _CAPTION_SUFFIX = "_caption" _ENHANCED_SUFFIX = "_enhanced_caption" +# Control types for transfer training. +_ON_THE_FLY_CONTROL_TYPES = frozenset({"edge", "blur"}) # computed from vision_path at train time +_PRECOMPUTED_CONTROL_TYPES = frozenset({"depth", "seg"}) # require a separate control video +_ALL_CONTROL_TYPES = _ON_THE_FLY_CONTROL_TYPES | _PRECOMPUTED_CONTROL_TYPES + def _relativize_vision_path(vision_path: str, output_jsonl: Path) -> str: """Rewrite ``vision_path`` relative to the output JSONL's directory. @@ -133,6 +154,39 @@ def _resolve_window_caption( return None +def _find_control_path( + clip_uuid: str, + clip_location: str, + control_type: str, + control_path_root: Path | None, +) -> str | None: + """Locate a precomputed control video for a clip. + + Search order: + 1. ``{control_path_root}/{uuid}.mp4`` + 2. ``{control_path_root}/{uuid}/{uuid}.mp4`` + 3. ``{control_path_root}/{uuid}_{control_type}.mp4`` + 4. ``{clip_dir}/{uuid}_{control_type}.mp4`` (sibling convention) + 5. ``{clip_dir}/{uuid}_{control_type}.mkv`` + """ + if control_path_root is not None: + for candidate in [ + control_path_root / f"{clip_uuid}.mp4", + control_path_root / clip_uuid / f"{clip_uuid}.mp4", + control_path_root / f"{clip_uuid}_{control_type}.mp4", + ]: + if candidate.is_file(): + return str(candidate) + + clip_dir = Path(clip_location).parent + for suffix in (f"_{control_type}.mp4", f"_{control_type}.mkv"): + sibling = clip_dir / f"{clip_uuid}{suffix}" + if sibling.is_file(): + return str(sibling) + + return None + + def _compute_duration(record: dict[str, Any]) -> float | None: """Derive clip duration in seconds from a curator row. @@ -162,10 +216,17 @@ def _build_sft_row( min_window_frames: int, max_duration_s: float, temporal_interval: int, + control_type: str | None = None, + control_path_root: Path | None = None, ) -> tuple[dict[str, Any] | None, str | None]: """Translate a curator metas_jsonl row into an SFT row, or report a drop reason. Returns ``(row, None)`` for kept records and ``(None, reason)`` for drops. + + When ``control_type`` is set, two extra fields are written: + - ``control_type`` — always present. + - ``control_path`` — only for ``depth``/``seg``; clips without a matching + precomputed file are dropped with reason ``missing_control_path``. """ clip_uuid = record.get("span_uuid") clip_location = record.get("clip_location") @@ -184,6 +245,13 @@ def _build_sft_row( if min_short_edge > 0 and min(int(width), int(height)) < min_short_edge: return None, "short_edge_too_small" + # For precomputed control types, require a matching control video. + control_path: str | None = None + if control_type in _PRECOMPUTED_CONTROL_TYPES: + control_path = _find_control_path(str(clip_uuid), str(clip_location), control_type, control_path_root) + if control_path is None: + return None, "missing_control_path" + windows = record.get("windows") or [] t2w_windows: list[dict[str, Any]] = [] for window in windows: @@ -223,6 +291,10 @@ def _build_sft_row( "vision_path": str(clip_location), "t2w_windows": t2w_windows, } + if control_type is not None: + row["control_type"] = control_type + if control_path is not None: + row["control_path"] = control_path return row, None @@ -266,8 +338,38 @@ def main( # noqa: PLR0913 int, tyro.conf.arg(help="temporal_interval to record on every t2w_window."), ] = DEFAULT_TEMPORAL_INTERVAL, + control_type: Annotated[ + Literal["edge", "blur", "depth", "seg"] | None, + tyro.conf.arg( + help=( + "Produce a transfer SFT JSONL by adding 'control_type' (and optionally " + "'control_path') to each row. 'edge' and 'blur' are computed on-the-fly " + "at train time — no precomputed files needed. 'depth' and 'seg' require " + "precomputed control videos; pass --control-path-root to locate them. " + "Omit to produce a standard T2V/I2V/V2V SFT JSONL (default behaviour)." + ), + ), + ] = None, + control_path_root: Annotated[ + Path | None, + tyro.conf.arg( + help=( + "Directory containing precomputed control videos for --control-type depth/seg. " + "Searched as {root}/{uuid}.mp4 or {root}/{uuid}/{uuid}.mp4. " + "Not needed for edge or blur." + ), + ), + ] = None, ) -> None: """Build an SFT JSONL from a curator splitting-pipeline output directory.""" + if control_type in _PRECOMPUTED_CONTROL_TYPES and control_path_root is None: + print( + f"WARNING: --control-type={control_type!r} requires precomputed control files. " + "Pass --control-path-root to locate them; clips without a matching file will be " + "dropped with reason 'missing_control_path'.", + file=sys.stderr, + ) + jsonl_files = list(_iter_metas_jsonl_files(curator_output)) if not jsonl_files: print( @@ -302,11 +404,15 @@ def main( # noqa: PLR0913 min_window_frames=min_window_frames, max_duration_s=max_duration_s, temporal_interval=temporal_interval, + control_type=control_type, + control_path_root=control_path_root, ) if row is None: drops[reason or "unknown"] += 1 continue row["vision_path"] = _relativize_vision_path(row["vision_path"], output) + if "control_path" in row: + row["control_path"] = _relativize_vision_path(row["control_path"], output) kept_rows.append(row) output.parent.mkdir(parents=True, exist_ok=True) @@ -330,6 +436,8 @@ def main( # noqa: PLR0913 "caption_model": caption_model, "enhanced_caption_model": enhanced_caption_model, "temporal_interval": temporal_interval, + "control_type": control_type, + "control_path_root": str(control_path_root) if control_path_root else None, } summary_path = output.with_suffix(output.suffix + ".summary.json") summary_path.write_text(json.dumps(summary, indent=2) + "\n") diff --git a/cosmos_framework/scripts/curator_to_sft_jsonl_test.py b/cosmos_framework/scripts/curator_to_sft_jsonl_test.py index 16845e0e..0c0c4fc6 100644 --- a/cosmos_framework/scripts/curator_to_sft_jsonl_test.py +++ b/cosmos_framework/scripts/curator_to_sft_jsonl_test.py @@ -13,6 +13,7 @@ MAX_VIDEO_DURATION_S, MIN_WINDOW_FRAMES, _build_sft_row, + _find_control_path, _relativize_vision_path, _resolve_window_caption, main, @@ -581,3 +582,143 @@ def test_main_exits_when_no_metas_jsonl_found(tmp_path: Path) -> None: temporal_interval=DEFAULT_TEMPORAL_INTERVAL, ) assert exc_info.value.code == 1 + + +# --------------------------------------------------------------------------- +# Transfer (control-conditioned) tests +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_build_sft_row_edge_control_adds_control_type_no_control_path() -> None: + """edge is computed on-the-fly: row gets control_type but no control_path.""" + record = _make_record() + row, reason = _build_sft_row(record, **_default_kwargs(), control_type="edge") + assert reason is None + assert row is not None + assert row["control_type"] == "edge" + assert "control_path" not in row + + +@pytest.mark.L0 +def test_build_sft_row_blur_control_adds_control_type_no_control_path() -> None: + """blur is computed on-the-fly: row gets control_type but no control_path.""" + record = _make_record() + row, reason = _build_sft_row(record, **_default_kwargs(), control_type="blur") + assert reason is None + assert row is not None + assert row["control_type"] == "blur" + assert "control_path" not in row + + +@pytest.mark.L0 +def test_build_sft_row_depth_without_control_path_root_drops_missing_control_path() -> None: + """depth requires a precomputed file; clip is dropped when none can be found.""" + record = _make_record(clip_location="/out/clips/clip-uuid-0.mp4") + row, reason = _build_sft_row(record, **_default_kwargs(), control_type="depth") + assert row is None + assert reason == "missing_control_path" + + +@pytest.mark.L0 +def test_build_sft_row_depth_with_existing_control_file_emits_control_path(tmp_path: Path) -> None: + """depth with a matching {root}/{uuid}.mp4 is kept and carries control_path.""" + control_root = tmp_path / "depth_maps" + control_root.mkdir() + uuid = "clip-uuid-0" + control_file = control_root / f"{uuid}.mp4" + control_file.touch() + + record = _make_record(span_uuid=uuid, clip_location=str(tmp_path / "clips" / f"{uuid}.mp4")) + row, reason = _build_sft_row(record, **_default_kwargs(), control_type="depth", control_path_root=control_root) + assert reason is None + assert row is not None + assert row["control_type"] == "depth" + assert row["control_path"] == str(control_file) + + +@pytest.mark.L0 +def test_find_control_path_finds_flat_uuid_mp4_at_root(tmp_path: Path) -> None: + uuid = "abc-123" + control_root = tmp_path / "depth" + control_root.mkdir() + expected = control_root / f"{uuid}.mp4" + expected.touch() + result = _find_control_path(uuid, str(tmp_path / "clips" / f"{uuid}.mp4"), "depth", control_root) + assert result == str(expected) + + +@pytest.mark.L0 +def test_find_control_path_finds_sibling_when_no_root(tmp_path: Path) -> None: + uuid = "abc-123" + clip_dir = tmp_path / "clips" + clip_dir.mkdir() + sibling = clip_dir / f"{uuid}_depth.mp4" + sibling.touch() + result = _find_control_path(uuid, str(clip_dir / f"{uuid}.mp4"), "depth", None) + assert result == str(sibling) + + +@pytest.mark.L0 +def test_find_control_path_returns_none_when_no_file_exists(tmp_path: Path) -> None: + result = _find_control_path("no-such-uuid", str(tmp_path / "clips" / "no-such-uuid.mp4"), "depth", None) + assert result is None + + +@pytest.mark.L0 +def test_main_edge_control_emits_control_type_field(tmp_path: Path) -> None: + """End-to-end: --control-type edge adds 'control_type' to every output row.""" + curator_output = tmp_path / "curator_out" + metas_dir = curator_output / "metas_jsonl" / "v0" + metas_dir.mkdir(parents=True) + (metas_dir / "shard_0.jsonl").write_text(json.dumps(_make_record()) + "\n") + + output = curator_output / "cosmos3_transfer_sft.jsonl" + main( + curator_output=curator_output, + output=output, + caption_model="qwen", + enhanced_caption_model=None, + min_short_edge=0, + min_window_frames=MIN_WINDOW_FRAMES, + max_duration_s=MAX_VIDEO_DURATION_S, + temporal_interval=DEFAULT_TEMPORAL_INTERVAL, + control_type="edge", + ) + + rows = [json.loads(line) for line in output.read_text().splitlines() if line] + assert len(rows) == 1 + assert rows[0]["control_type"] == "edge" + assert "control_path" not in rows[0] + + summary = json.loads(output.with_suffix(output.suffix + ".summary.json").read_text()) + assert summary["control_type"] == "edge" + assert summary["control_path_root"] is None + + +@pytest.mark.L0 +def test_main_depth_control_drops_clips_without_precomputed_files(tmp_path: Path) -> None: + """--control-type depth drops clips for which no precomputed file can be found.""" + curator_output = tmp_path / "curator_out" + metas_dir = curator_output / "metas_jsonl" / "v0" + metas_dir.mkdir(parents=True) + record = _make_record(clip_location=str(tmp_path / "clips" / "clip-uuid-0.mp4")) + (metas_dir / "shard_0.jsonl").write_text(json.dumps(record) + "\n") + + output = curator_output / "cosmos3_transfer_sft.jsonl" + with pytest.raises(SystemExit) as exc_info: + main( + curator_output=curator_output, + output=output, + caption_model=None, + enhanced_caption_model=None, + min_short_edge=0, + min_window_frames=MIN_WINDOW_FRAMES, + max_duration_s=MAX_VIDEO_DURATION_S, + temporal_interval=DEFAULT_TEMPORAL_INTERVAL, + control_type="depth", + ) + assert exc_info.value.code == 1 + + summary = json.loads(output.with_suffix(output.suffix + ".summary.json").read_text()) + assert summary["drops_by_reason"].get("missing_control_path", 0) == 1 From 0fa3ba479a7c1c5be28f81fe084b9ea544d697c9 Mon Sep 17 00:00:00 2001 From: yy-code-nv Date: Sat, 11 Jul 2026 20:07:03 +0800 Subject: [PATCH 22/26] Release 2026-07-11 (i4 7392c4b5) (#105) Automated release from i4. _source_commit: `7392c4b59ab7c1f0a94e59fbc4a75a0f3684b66f-dirty` _dest_commit (base): `3d9c0878fd0dde76eac98161aed0493d85a036fd` Co-authored-by: lfengad --- .file_mapping.json | 17 +- .../callbacks/every_n_draw_sample.py | 104 +- cosmos_framework/checkpoint/dcp.py | 12 +- .../configs/base/defaults/checkpointer.py | 1 + .../configs/base/defaults/model_config.py | 2 +- .../configs/base/defaults/optimizer.py | 71 + .../configs/base/defaults/reasoner.py | 32 + .../configs/base/defaults/tokenizer.py | 1 + .../data/generator/action/domain_utils.py | 2 + cosmos_framework/data/generator/utils.py | 9 +- cosmos_framework/model/attention/backends.py | 11 +- .../attention/benchmarks/benchmark_fmha.py | 638 +++++++++ .../model/attention/cudnn/__init__.py | 42 +- .../model/attention/cudnn/checks.py | 110 +- .../model/attention/cudnn/functions.py | 279 ++-- .../model/attention/cudnn/meta.py | 41 +- cosmos_framework/model/attention/frontend.py | 3 +- .../model/attention/utils/__init__.py | 9 + .../model/generator/mot/attention.py | 58 +- .../model/generator/omni_mot_model.py | 154 ++- .../tokenizers/uniae/noncausal_4x16x16.py | 309 +++-- .../tokenizers/wan2pt2_vae_4x16x16.py | 5 +- .../model/tokenizer/checkpoint_io.py | 318 +++++ .../model/tokenizer/evaluation/metric.py | 54 +- .../evaluation/reconstruction_metrics.py | 90 -- .../model/tokenizer/models/architecture.py | 53 + .../model/tokenizer/models/dense_runtime.py | 56 +- .../tokenizer/models/modules/__init__.py | 26 - .../models/modules/attention/full_attn.py | 9 +- .../models/modules/attention/modules.py | 15 +- .../tokenizer/models/modules/sparse_tensor.py | 34 +- .../tokenizer/models/sparse_autoencoder.py | 19 +- .../model/tokenizer/models/text_decoder.py | 421 ++++-- .../model/tokenizer/models/utils.py | 13 +- .../model/tokenizer/utils/tensors.py | 63 + cosmos_framework/tools/flops/qwen3_vl.py | 21 +- cosmos_framework/tools/flops/test_qwen3_vl.py | 158 +++ .../easy_io/handlers/imageio_video_handler.py | 18 +- .../handlers/imageio_video_handler_test.py | 57 + .../utils/generator/aux_optimizer_utils.py | 316 +++++ .../utils/generator/dion2_with_aux_adamw.py | 1149 +++++++++++++++++ .../utils/generator/muon_with_aux_adamw.py | 835 ++++++++++++ cosmos_framework/utils/generator/optimizer.py | 85 +- .../utils/generator/torchcodec_video.py | 16 +- 44 files changed, 5091 insertions(+), 645 deletions(-) create mode 100644 cosmos_framework/model/attention/benchmarks/benchmark_fmha.py create mode 100644 cosmos_framework/model/tokenizer/checkpoint_io.py create mode 100644 cosmos_framework/model/tokenizer/models/architecture.py create mode 100644 cosmos_framework/model/tokenizer/utils/tensors.py create mode 100644 cosmos_framework/tools/flops/test_qwen3_vl.py create mode 100644 cosmos_framework/utils/easy_io/handlers/imageio_video_handler_test.py create mode 100644 cosmos_framework/utils/generator/aux_optimizer_utils.py create mode 100644 cosmos_framework/utils/generator/dion2_with_aux_adamw.py create mode 100644 cosmos_framework/utils/generator/muon_with_aux_adamw.py diff --git a/.file_mapping.json b/.file_mapping.json index 503fa929..5b0e65ba 100644 --- a/.file_mapping.json +++ b/.file_mapping.json @@ -1,16 +1,15 @@ { - "_source_commit": "c2c7152e59e4090a61822dc22a1691e7afb44702-dirty", - "_dest_commit": "463ac142b47c637098a7e86d0631252f82b65418", - "_generated_at": "2026-07-07T05:54:17Z", + "_source_commit": "7392c4b59ab7c1f0a94e59fbc4a75a0f3684b66f-dirty", + "_dest_commit": "3d9c0878fd0dde76eac98161aed0493d85a036fd", + "_generated_at": "2026-07-11T05:50:38Z", "files": { "imaginaire/__init__.py": "cosmos_framework/__init__.py", "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", "imaginaire/attention/backends.py": "cosmos_framework/model/attention/backends.py", - "imaginaire/attention/benchmarks/benchmark_fmha_head_sharding.py": "cosmos_framework/model/attention/benchmarks/benchmark_fmha_head_sharding.py", + "imaginaire/attention/benchmarks/benchmark_fmha.py": "cosmos_framework/model/attention/benchmarks/benchmark_fmha.py", "imaginaire/attention/checks.py": "cosmos_framework/model/attention/checks.py", "imaginaire/attention/cudnn/__init__.py": "cosmos_framework/model/attention/cudnn/__init__.py", "imaginaire/attention/cudnn/checks.py": "cosmos_framework/model/attention/cudnn/checks.py", - "imaginaire/attention/cudnn/cudnn_forward.py": "cosmos_framework/model/attention/cudnn/cudnn_forward.py", "imaginaire/attention/cudnn/functions.py": "cosmos_framework/model/attention/cudnn/functions.py", "imaginaire/attention/cudnn/meta.py": "cosmos_framework/model/attention/cudnn/meta.py", "imaginaire/attention/cudnn/stubs.py": "cosmos_framework/model/attention/cudnn/stubs.py", @@ -88,6 +87,7 @@ "imaginaire/flops/__init__.py": "cosmos_framework/tools/flops/__init__.py", "imaginaire/flops/omni_mot.py": "cosmos_framework/tools/flops/omni_mot.py", "imaginaire/flops/qwen3_vl.py": "cosmos_framework/tools/flops/qwen3_vl.py", + "imaginaire/flops/test_qwen3_vl.py": "cosmos_framework/tools/flops/test_qwen3_vl.py", "imaginaire/flops/wan_vae.py": "cosmos_framework/tools/flops/wan_vae.py", "imaginaire/functional/batch_ops.py": "cosmos_framework/utils/functional/batch_ops.py", "imaginaire/functional/lr_scheduler.py": "cosmos_framework/utils/functional/lr_scheduler.py", @@ -132,6 +132,7 @@ "imaginaire/utils/easy_io/handlers/csv_handler.py": "cosmos_framework/utils/easy_io/handlers/csv_handler.py", "imaginaire/utils/easy_io/handlers/gzip_handler.py": "cosmos_framework/utils/easy_io/handlers/gzip_handler.py", "imaginaire/utils/easy_io/handlers/imageio_video_handler.py": "cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py", + "imaginaire/utils/easy_io/handlers/imageio_video_handler_test.py": "cosmos_framework/utils/easy_io/handlers/imageio_video_handler_test.py", "imaginaire/utils/easy_io/handlers/json_handler.py": "cosmos_framework/utils/easy_io/handlers/json_handler.py", "imaginaire/utils/easy_io/handlers/jsonl_handler.py": "cosmos_framework/utils/easy_io/handlers/jsonl_handler.py", "imaginaire/utils/easy_io/handlers/np_handler.py": "cosmos_framework/utils/easy_io/handlers/np_handler.py", @@ -405,7 +406,9 @@ "projects/cosmos3/cosmos3/tokenizers/wan2pt2_vae_4x16x16.py": "cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py", "projects/cosmos3/cosmos3/upsampler/__init__.py": "cosmos_framework/model/generator/upsampler/__init__.py", "projects/cosmos3/cosmos3/upsampler/prompts.py": "cosmos_framework/model/generator/upsampler/prompts.py", + "projects/cosmos3/cosmos3/utils/aux_optimizer_utils.py": "cosmos_framework/utils/generator/aux_optimizer_utils.py", "projects/cosmos3/cosmos3/utils/data_utils.py": "cosmos_framework/utils/generator/data_utils.py", + "projects/cosmos3/cosmos3/utils/dion2_with_aux_adamw.py": "cosmos_framework/utils/generator/dion2_with_aux_adamw.py", "projects/cosmos3/cosmos3/utils/dtensor_helper.py": "cosmos_framework/utils/generator/dtensor_helper.py", "projects/cosmos3/cosmos3/utils/flash_attn.py": "cosmos_framework/utils/generator/flash_attn.py", "projects/cosmos3/cosmos3/utils/fused_adam.py": "cosmos_framework/utils/generator/fused_adam.py", @@ -414,6 +417,7 @@ "projects/cosmos3/cosmos3/utils/model_loader.py": "cosmos_framework/utils/generator/model_loader.py", "projects/cosmos3/cosmos3/utils/model_weights_stats.py": "cosmos_framework/utils/generator/model_weights_stats.py", "projects/cosmos3/cosmos3/utils/monkey_patch.py": "cosmos_framework/utils/generator/monkey_patch.py", + "projects/cosmos3/cosmos3/utils/muon_with_aux_adamw.py": "cosmos_framework/utils/generator/muon_with_aux_adamw.py", "projects/cosmos3/cosmos3/utils/optimizer.py": "cosmos_framework/utils/generator/optimizer.py", "projects/cosmos3/cosmos3/utils/parallelism.py": "cosmos_framework/utils/generator/parallelism.py", "projects/cosmos3/cosmos3/utils/rand_state.py": "cosmos_framework/utils/generator/rand_state.py", @@ -423,9 +427,11 @@ "projects/cosmos3/cosmos3/utils/reasoner/flop_calculator.py": "cosmos_framework/utils/generator/reasoner/flop_calculator.py", "projects/cosmos3/cosmos3/utils/reasoner/pretrained_models_downloader.py": "cosmos_framework/utils/generator/reasoner/pretrained_models_downloader.py", "projects/cosmos3/cosmos3/utils/video_preprocess.py": "cosmos_framework/utils/generator/video_preprocess.py", + "projects/cosmos3/tokenizer/checkpoint_io.py": "cosmos_framework/model/tokenizer/checkpoint_io.py", "projects/cosmos3/tokenizer/evaluation/metric.py": "cosmos_framework/model/tokenizer/evaluation/metric.py", "projects/cosmos3/tokenizer/evaluation/reconstruction_metrics.py": "cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py", "projects/cosmos3/tokenizer/models/__init__.py": "cosmos_framework/model/tokenizer/models/__init__.py", + "projects/cosmos3/tokenizer/models/architecture.py": "cosmos_framework/model/tokenizer/models/architecture.py", "projects/cosmos3/tokenizer/models/dense_backends.py": "cosmos_framework/model/tokenizer/models/dense_backends.py", "projects/cosmos3/tokenizer/models/dense_runtime.py": "cosmos_framework/model/tokenizer/models/dense_runtime.py", "projects/cosmos3/tokenizer/models/modules/__init__.py": "cosmos_framework/model/tokenizer/models/modules/__init__.py", @@ -445,6 +451,7 @@ "projects/cosmos3/tokenizer/models/text_decoder.py": "cosmos_framework/model/tokenizer/models/text_decoder.py", "projects/cosmos3/tokenizer/models/utils.py": "cosmos_framework/model/tokenizer/models/utils.py", "projects/cosmos3/tokenizer/utils/hf.py": "cosmos_framework/model/tokenizer/utils/hf.py", + "projects/cosmos3/tokenizer/utils/tensors.py": "cosmos_framework/model/tokenizer/utils/tensors.py", "projects/cosmos3/tokenizer/utils/vlm_prompt_format.py": "cosmos_framework/model/tokenizer/utils/vlm_prompt_format.py", "projects/cosmos3/utils/torchcodec_video.py": "cosmos_framework/utils/generator/torchcodec_video.py", "projects/cosmos3/vlm/configs/base/defaults/checkpointer.py": "cosmos_framework/utils/reasoner/configs_defaults/checkpointer.py", diff --git a/cosmos_framework/callbacks/every_n_draw_sample.py b/cosmos_framework/callbacks/every_n_draw_sample.py index fa87fda5..015c5f7f 100644 --- a/cosmos_framework/callbacks/every_n_draw_sample.py +++ b/cosmos_framework/callbacks/every_n_draw_sample.py @@ -217,6 +217,91 @@ def _add_wandb_image_paths( info[key_prefix] = wandb.Image(image_paths, caption=caption) +def _pixel_tensor_to_5d(t: torch.Tensor) -> torch.Tensor: + """Ensure a pixel tensor has shape (B, C, T, H, W) for the visualization grid. + + Handles (C, H, W), (B, C, H, W), and (B, C, T, H, W) inputs. + """ + if t.ndim == 3: + return t.unsqueeze(0).unsqueeze(2) # (C,H,W) -> (1,C,1,H,W) + if t.ndim == 4: + return t.unsqueeze(2) # (B,C,H,W) -> (B,C,1,H,W) + return t + + +def _resize_5d_to_width(img5d: torch.Tensor, target_width: int) -> torch.Tensor: + """Resize a single-frame (1, C, 1, H, W) tensor so its width is exactly ``target_width``. + + Height is scaled proportionally to preserve the overall aspect ratio. Assumes a + single temporal frame (T == 1). + """ + h, w = img5d.shape[-2], img5d.shape[-1] + if w == target_width: + return img5d + new_h = max(1, round(h * target_width / w)) + chw = img5d[0, :, 0] # (C,H,W) + resized = torchvision_F.resize(chw, [new_h, target_width], antialias=True) # (C,new_h,target_width) + return resized.unsqueeze(0).unsqueeze(2) # (1,C,1,new_h,target_width) + + +def _resize_pad_to_square(img5d: torch.Tensor, cell: int) -> torch.Tensor: + """Resize a single-frame (1, C, 1, H, W) tensor to fit inside a ``cell`` x ``cell`` square. + + Aspect ratio is preserved (the image is scaled so its longer side equals ``cell``), then + the result is center-padded with zeros to exactly ``cell`` x ``cell``. Assumes T == 1. + """ + h, w = img5d.shape[-2], img5d.shape[-1] + scale = cell / max(h, w) + new_h = max(1, round(h * scale)) + new_w = max(1, round(w * scale)) + chw = img5d[0, :, 0] # (C,H,W) + resized = torchvision_F.resize(chw, [new_h, new_w], antialias=True) # (C,new_h,new_w) + pad_h = cell - new_h + pad_w = cell - new_w + top = pad_h // 2 + left = pad_w // 2 + padded = torch.nn.functional.pad(resized, (left, pad_w - left, top, pad_h - top)) # (C,cell,cell) + return padded.unsqueeze(0).unsqueeze(2) # (1,C,1,cell,cell) + + +def _build_reference_grid(references: list[torch.Tensor], target_width: int) -> torch.Tensor: + """Tile reference images into a compact near-square grid sized to ``target_width``. + + All references are arranged in a ``rows`` x ``cols`` grid (``cols = ceil(sqrt(n))``), each in + an aspect-preserved square cell, then the whole grid is resized so its width equals + ``target_width``. This keeps the condition montage aligned with the generated-image column + width (better space utilization) instead of one overly-long row. + + Args: + references: list of single-frame pixel tensors (each (1,C,1,H,W) or (C,H,W)), in [-1, 1]. + target_width: width of the generated/target image for this sample. + + Returns: + A (1, C, 1, H_grid, target_width) tensor. + """ + if not references: + raise ValueError("Expected at least one reference image to build the condition grid.") + + refs = [_pixel_tensor_to_5d(r) for r in references] # list[(1,C,1,H,W)] + n = len(refs) + cols = math.ceil(math.sqrt(n)) + rows = math.ceil(n / cols) + cell = max(1, target_width // cols) + + cells = [_resize_pad_to_square(r, cell) for r in refs] # list[(1,C,1,cell,cell)] + blank = torch.zeros_like(cells[0]) + while len(cells) < rows * cols: + cells.append(blank) + + row_imgs = [] + for r in range(rows): + row = torch.cat(cells[r * cols : (r + 1) * cols], dim=-1) # (1,C,1,cell,cols*cell) + row_imgs.append(row) + grid = torch.cat(row_imgs, dim=-2) # (1,C,1,rows*cell,cols*cell) + + return _resize_5d_to_width(grid, target_width) # (1,C,1,H_grid,target_width) + + class EveryNDrawSample(EveryN): """ This callback sample condition inputs from training data, run inference and save the results to wandb and s3. @@ -531,11 +616,20 @@ def sample( vis_offset = 0 for sample_idx in range(data_clean.batch_size): n_vis = num_items[sample_idx] - # First item(s) are condition, last item is generation target - # but we need to support multiple conditions per sample in the future. Current code - # can handle this without throwing an error. - condition_images.append(raw_data[vis_offset]) # source image (1, C, 1, H, W) - gt_target_images.append(raw_data[vis_offset + n_vis - 1]) # target image (1, C, 1, H, W) + # First item(s) are condition references, last item is the generation target. + refs = raw_data[vis_offset : vis_offset + n_vis - 1] # all condition items + target = raw_data[vis_offset + n_vis - 1] # target image (1, C, 1, H, W) + # Multi-reference generation (>1 single-frame image references): tile every + # reference into a compact grid resized to the target width, so all references + # are visible without blowing up the row width. For video editing/transfer + # (T > 1) keep the existing behavior (first item = condition) so those tasks + # render exactly as before and stay consistent with the t_crop frame cropping. + refs_are_images = all(r.shape[-3] == 1 for r in refs) and target.shape[-3] == 1 + if refs_are_images and len(refs) > 1: + condition_images.append(_build_reference_grid(refs, target.shape[-1])) + else: + condition_images.append(raw_data[vis_offset]) # source image (1, C, 1, H, W) / video clip + gt_target_images.append(target) vis_offset += n_vis # Use target images for max_w/max_h/t_crop (generated samples match target size) diff --git a/cosmos_framework/checkpoint/dcp.py b/cosmos_framework/checkpoint/dcp.py index 54d0d8e1..43145c88 100644 --- a/cosmos_framework/checkpoint/dcp.py +++ b/cosmos_framework/checkpoint/dcp.py @@ -71,9 +71,9 @@ from cosmos_framework.checkpoint.base import AbstractCheckpointer from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter +from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.utils import callback, distributed, log, misc -from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.utils.easy_io import easy_io from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict @@ -866,16 +866,6 @@ def load( raise ValueError( f"Unexpected keys (found in checkpoint but not in model): {results.unexpected_keys}" ) - # Warm start that skipped net_ema (e.g. loading an EMA-only HF export - # with no net_ema.* keys): the EMA shadow would otherwise keep its random - # build-time generation pathway (init_moe is skipped when a checkpoint is - # present). Seed net_ema from the freshly loaded net so the EMA starts equal - # to net ("EMA warm-starts from net") instead of from random weights. - if warm_start and any("net_ema" in skip_key for skip_key in keys_to_skip_loading): - ema_worker = getattr(model, "net_ema_worker", None) - if ema_worker is not None and getattr(model, "net_ema", None) is not None: - ema_worker.copy_to(src_model=model.net, tgt_model=model.net_ema) - log.info("Warm start: re-seeded net_ema from net (net_ema was skipped on load).") elif key == "optim": log.info("- Loading the optimizer...") diff --git a/cosmos_framework/configs/base/defaults/checkpointer.py b/cosmos_framework/configs/base/defaults/checkpointer.py index 9502070e..151f6e61 100644 --- a/cosmos_framework/configs/base/defaults/checkpointer.py +++ b/cosmos_framework/configs/base/defaults/checkpointer.py @@ -98,6 +98,7 @@ strict_resume=True, enable_gcs_patch_in_boto3=True, dcp_async_mode_enabled=True, + dcp_load_dedup=True, ) CHECKPOINT_S3_EAST2 = CheckpointConfig( diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py index adde6bd9..5e6fd6b4 100644 --- a/cosmos_framework/configs/base/defaults/model_config.py +++ b/cosmos_framework/configs/base/defaults/model_config.py @@ -5,12 +5,12 @@ import attrs +from cosmos_framework.utils.lazy_config import LazyDict from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.ema import EMAConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig from cosmos_framework.configs.base.defaults.reasoner import VLMConfig -from cosmos_framework.utils.lazy_config import LazyDict @attrs.define(slots=False) diff --git a/cosmos_framework/configs/base/defaults/optimizer.py b/cosmos_framework/configs/base/defaults/optimizer.py index dc56906c..a619b2fe 100644 --- a/cosmos_framework/configs/base/defaults/optimizer.py +++ b/cosmos_framework/configs/base/defaults/optimizer.py @@ -31,6 +31,57 @@ disable_weight_decay_for_1d_params=False, ) +# Muon / Dion2 share the standard factory knobs (keys_to_select, lr_multipliers, +# disable_weight_decay_for_1d_params) plus their own orthogonalization +# hyperparameters. ``fused`` is required by the factory; the AdamW side is fused +# by construction and ``capturable`` / ``master_weights`` are forced on. +MUON_OPTIMIZER_KWARGS: dict[str, Any] = dict( + # Base learning rate. Muon scales matrix params by muon_lr_scale*sqrt(max(A,B)); + # the AdamW side and the per-param-group lr_multipliers use it directly. + lr=1e-4, + weight_decay=0.1, + adam_betas=[0.9, 0.99], + eps=1e-8, + fused=True, + keys_to_select=[], + lr_multipliers={}, + disable_weight_decay_for_1d_params=False, + # Name substrings for stacked MoE expert params ([E, M, N]) to orthogonalize + # per expert slice. Empty = experts stay on AdamW (no behavior change). + # e.g. ["gate_up_proj", "down_proj"] for grouped-MM MoE experts. + expert_param_keywords=[], + # Muon-specific. + muon_momentum=0.95, + muon_lr_scale=0.2, + ns_steps=5, + nesterov=True, + use_distributed=True, +) + +DION2_OPTIMIZER_KWARGS: dict[str, Any] = dict( + lr=1e-4, + weight_decay=0.1, + adam_betas=[0.9, 0.99], + eps=1e-8, + fused=True, + keys_to_select=[], + lr_multipliers={}, + disable_weight_decay_for_1d_params=False, + # Name substrings for stacked MoE expert params ([E, M, N]) to orthogonalize + # per expert slice. Empty = experts stay on AdamW (no behavior change). + # e.g. ["gate_up_proj", "down_proj"] for grouped-MM MoE experts. + expert_param_keywords=[], + # Muon/Dion2-specific. + muon_momentum=0.95, + muon_lr_scale=0.2, + ns_steps=5, + nesterov=True, + use_distributed=True, + # Dion2-specific: submatrix selection fraction and error-feedback decay. + fraction=1.0, + ef_decay=0.95, +) + LAMBDACOSINE_KWARGS: dict[str, Any] = dict( warm_up_steps=[2000], cycle_lengths=[100000], @@ -64,6 +115,26 @@ def register_optimizers(optimizer_kwargs: dict[str, Any]) -> None: **optimizer_kwargs, ), ) + cs.store( + group="optimizer", + package="optimizer", + name="muonwithauxadamw", + node=L(build_optimizer)( + model=PLACEHOLDER, + optimizer_type="MuonWithAuxAdamW", + **MUON_OPTIMIZER_KWARGS, + ), + ) + cs.store( + group="optimizer", + package="optimizer", + name="dion2withauxadamw", + node=L(build_optimizer)( + model=PLACEHOLDER, + optimizer_type="Dion2WithAuxAdamW", + **DION2_OPTIMIZER_KWARGS, + ), + ) def register_schedulers(lambdacosine_kwargs: dict[str, Any]) -> None: diff --git a/cosmos_framework/configs/base/defaults/reasoner.py b/cosmos_framework/configs/base/defaults/reasoner.py index bee81b63..09e92efe 100644 --- a/cosmos_framework/configs/base/defaults/reasoner.py +++ b/cosmos_framework/configs/base/defaults/reasoner.py @@ -785,6 +785,32 @@ class VLMConfig: ), ) +# Same as 590c1c0 but with use_und_k_norm_for_gen=True: normalises K_und before +# it is used as a key in the gen→und cross-attention path (the qk-norm fix for the +# generator). Adds a freshly-initialised k_norm_und_for_gen RMSNorm. +Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0_UndKNorm: VLMConfig = VLMConfig( + model_name="nvidia/Cosmos3-Edge-Reasoner", + model_instance=L(Nemotron3DenseVLTextForCausalLM)( + config=L(create_vlm_config)( + base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( + json_file="cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json" + ), + qk_norm_for_text=False, + use_und_k_norm_for_gen=True, + ), + ), + tokenizer=L(build_processor_lazy)( + tokenizer_type="nvidia/Cosmos3-Edge-Reasoner", + config_variant="gcp", + ), + pretrained_weights=PretrainedWeightsConfig( + backbone_path="s3://bucket0/cosmos3/pretrained/huggingface/nvidia/Cosmos3-Edge-Reasoner-590c1c0/", + credentials_path="credentials/gcp_checkpoint.secret", + enable_gcs_patch_in_boto3=True, + checkpoint_format="nemotron_3_dense_vl", + ), +) + # Same as 9b4c028 but with use_und_k_norm_for_gen=True: normalises K_und before # it is used as a key in the gen→und cross-attention path. Cosmos3EdgeReasoner_VLM_GCP_Config_9b4c028_UndKNorm: VLMConfig = VLMConfig( @@ -987,3 +1013,9 @@ def register_vlm(): name="cosmos3_edge_reasoner_vlm_gcp_590c1c0", node=Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0, ) + cs.store( + group="vlm_config", + package="model.config.vlm_config", + name="cosmos3_edge_reasoner_vlm_gcp_590c1c0_und_k_norm", + node=Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0_UndKNorm, + ) diff --git a/cosmos_framework/configs/base/defaults/tokenizer.py b/cosmos_framework/configs/base/defaults/tokenizer.py index 04c08873..55e79016 100644 --- a/cosmos_framework/configs/base/defaults/tokenizer.py +++ b/cosmos_framework/configs/base/defaults/tokenizer.py @@ -128,6 +128,7 @@ vae_path=PRETRAINED_TOKENIZER_UNIAE_4X16X16_C48_T16TO160_MIXP_FPS_MIX_ENCODER_NONCAUSAL_DECODER_NONCAUSAL_NOGAN_S3_NEMOTRON2B_VAE_PTH, spatial_compression_factor=16, temporal_compression_factor=4, + pad_frames=1, pixel_trim=True, causal=False, ) diff --git a/cosmos_framework/data/generator/action/domain_utils.py b/cosmos_framework/data/generator/action/domain_utils.py index 6f433f73..c75f7f38 100644 --- a/cosmos_framework/data/generator/action/domain_utils.py +++ b/cosmos_framework/data/generator/action/domain_utils.py @@ -20,6 +20,7 @@ "agibotworld": 15, "embodiment_c_gripper": 15, "embodiment_c_gripper_ext": 15, + "xdof_yam": 16, "fractal": 20, } @@ -38,6 +39,7 @@ "agibotworld": 29, "embodiment_c_gripper": 29, "embodiment_c_gripper_ext": 29, + "xdof_yam": 20, "fractal": 10, # NOTE: ``libero`` (7/10/13 depending on ``rotation_space``) and ``hand_pose`` # (variable with ``keypoint_option`` and ``rotation_format``) are absent diff --git a/cosmos_framework/data/generator/utils.py b/cosmos_framework/data/generator/utils.py index bba5b3de..ecfcc0fe 100644 --- a/cosmos_framework/data/generator/utils.py +++ b/cosmos_framework/data/generator/utils.py @@ -98,7 +98,10 @@ def parse_frame_range_from_wdinfo(wdinfo: str) -> tuple[int, int | float] | None wdinfo: wdinfo path string containing a frames_X_Y pattern, where Y may be ``inf`` Returns: - Tuple of (min_frames, max_frames) if found, None otherwise + Tuple of (min_frames, max_frames) if found, where max_frames is + ``math.inf`` for unbounded ``frames_X_inf`` and ``frames_gt_X`` buckets. + ``frames_gt_X`` is treated as starting at X. Returns None if no frame + range is present. Example: >>> parse_frame_range_from_wdinfo("wdinfo/v4/tv_drama/resolution_720/aspect_ratio_16_9/frames_300_400/wdinfo.json") @@ -110,6 +113,10 @@ def parse_frame_range_from_wdinfo(wdinfo: str) -> tuple[int, int | float] | None if match: max_frames = math.inf if match.group(2) == "inf" else int(match.group(2)) return (int(match.group(1)), max_frames) + match = re.search(r"frames_gt_(\d+)", wdinfo) + if match: + # Wdinfo generation uses frames_gt_X for the bucket that starts at X, so treat it as >= X. + return (int(match.group(1)), math.inf) return None diff --git a/cosmos_framework/model/attention/backends.py b/cosmos_framework/model/attention/backends.py index 4236b35b..8d5f83fb 100644 --- a/cosmos_framework/model/attention/backends.py +++ b/cosmos_framework/model/attention/backends.py @@ -10,6 +10,7 @@ import torch +from cosmos_framework.model.attention.cudnn.checks import cudnn_attention_check from cosmos_framework.model.attention.flash2.checks import flash2_attention_check from cosmos_framework.model.attention.flash3.checks import flash3_attention_check from cosmos_framework.model.attention.masks import CausalType @@ -22,8 +23,8 @@ from cosmos_framework.model.attention.utils.safe_ops import log from cosmos_framework.model.attention.utils.safe_ops.functools import lru_cache - BACKEND_CHECK_MAP = { + "cudnn": cudnn_attention_check, "natten": natten_attention_check, "flash2": flash2_attention_check, "flash3": flash3_attention_check, @@ -131,17 +132,25 @@ def get_backend_list(arch_tag: int) -> list[str]: if arch_tag == 90: default_backends = [ "flash3", + "cudnn", "natten", "flash2", ] elif arch_tag in [100, 103]: default_backends = [ + "cudnn", "natten", "flash2", ] + elif arch_tag in [110, 120, 121]: + default_backends = [ + "cudnn", + "natten", + ] elif arch_tag >= 80: default_backends = [ "flash2", + "cudnn", "natten", ] else: diff --git a/cosmos_framework/model/attention/benchmarks/benchmark_fmha.py b/cosmos_framework/model/attention/benchmarks/benchmark_fmha.py new file mode 100644 index 00000000..bc68b1c5 --- /dev/null +++ b/cosmos_framework/model/attention/benchmarks/benchmark_fmha.py @@ -0,0 +1,638 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +""" +Benchmark the unified ``cosmos_framework.model.attention.attention`` frontend across a +user-specified list of backends, for a single sequence-length (and optional +context-parallel shard shape). + +Dense / batched attention ONLY: this benchmark deliberately has no varlen +(sequence-packed) support. Inputs are ``[B, S, H, D]``. + +This single script combines two previously-separate benchmarks: + +1. Backend sweep (the former ``benchmark_fmha``): compare backends at a single query + length (``--seqlen``), with optional cross-attention (``--kv-seqlen``), causal + masking, backward-pass timing, MLA-style value head dim (``--head-dim-v``), and + per-backend rejection diagnostics (``--debug``). + +2. Context-parallel per-rank shape study : + isolate the per-rank attention kernel shape a CP strategy would produce, via + ``--shard-mode`` + ``--cp-size``. ``head`` divides Q/KV heads by ``cp_size`` and + keeps the full KV length; ``sequence`` keeps the heads and splits the KV length + across ``cp_size`` shards (Q is intentionally NOT divided — each rank computes a + partial result over a local KV shard). There is intentionally no all-to-all, + all-gather, or all-reduce in the timed region: the intent is to study local kernel + scalability. The shard transform is applied to the ``(q_len, kv_len)`` shape. With + ``--shard-mode none`` (default) no sharding is applied. + +The script runs single-process by default, but also runs under ``torchrun`` so a real +multi-rank CP shape can be measured (``sequence`` mode is rank-dependent). Metrics +reported per row: median / min latency, achieved TFLOP/s (on the LOCAL shape), tokens/s, +and peak allocated memory. Only rank 0 prints. + +By default this is self-attention (``q_len == kv_len``). Cross-attention with a +different KV length can be benchmarked via ``--kv-seqlen``, but only when ``--causal`` +is NOT set: causal masking here uses ``CausalType.DontCare``, which requires +``q_len == kv_len``. When ``--causal`` is set, ``kv_len`` always equals ``q_len``. + +Examples: + + # Backend sweep, self-attention, causal, forward-only: + python -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --backends natten cudnn \ + --seqlen 3093 \ + --batch 1 --num-q-heads 16 --num-kv-heads 8 --head-dim 128 \ + --dtype bf16 --causal + + # Cross-attention (non-causal), distinct q_len and kv_len: + python -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --backends natten cudnn \ + --seqlen 4096 --kv-seqlen 1024 \ + --num-q-heads 16 --num-kv-heads 16 --head-dim 128 --dtype bf16 + + # One-GPU mock of a CP=4 per-rank head-sharded shape: + python -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --backends natten cudnn \ + --seqlen 396 --kv-seqlen 177771 \ + --num-q-heads 32 --num-kv-heads 8 --head-dim 128 \ + --shard-mode head --cp-size 4 --compile + + # True four-GPU CP head-sharded run: + torchrun --standalone --nproc_per_node=4 \ + -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --seqlen 396 --kv-seqlen 177771 \ + --num-q-heads 32 --num-kv-heads 8 --head-dim 128 \ + --shard-mode head --cp-size 4 + + # One-GPU mock of a CP=4 per-rank sequence-sharded (split-KV) shape. + # KV length is divided by cp_size while heads are kept; --causal is not allowed: + python -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --backends natten cudnn \ + --seqlen 396 --kv-seqlen 177771 \ + --num-q-heads 32 --num-kv-heads 8 --head-dim 128 \ + --shard-mode sequence --cp-size 4 + + # True four-GPU CP sequence-sharded run (each rank holds a different KV shard): + torchrun --standalone --nproc_per_node=4 \ + -m cosmos_framework.model.attention.benchmarks.benchmark_fmha \ + --seqlen 396 --kv-seqlen 177771 \ + --num-q-heads 32 --num-kv-heads 8 --head-dim 128 \ + --shard-mode sequence --cp-size 4 + +Pass ``--debug`` to report the concrete reason an incompatible backend was rejected +(rather than the generic "incompatible" message) and to print DEBUG-level +backend-selection logs. Pass ``--json`` to emit one JSON object per result row instead +of the formatted table. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch import Tensor + +from cosmos_framework.model.attention import attention +from cosmos_framework.model.attention.backends import choose_backend +from cosmos_framework.model.attention.masks import CausalType +from cosmos_framework.model.attention.utils import get_arch_tag + +DTYPE_MAP: dict[str, torch.dtype] = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} + + +@dataclass(frozen=True) +class LocalAttentionShape: + """Per-rank attention shape after applying the CP shard transform to a global pair.""" + + q_len: int + kv_len: int + num_q_heads: int + num_kv_heads: int + sequence_shard_index: int + sequence_shard_start: int + + +def local_shape_for_pair( + q_len: int, + kv_len: int, + num_q_heads: int, + num_kv_heads: int, + shard_mode: str, + cp_size: int, + rank: int, +) -> LocalAttentionShape: + """Transform a global ``(q_len, kv_len, heads)`` pair into the LOCAL per-rank shape. + + - ``none``: identity (cp_size must be 1). + - ``head``: divide Q and KV heads by ``cp_size``, keep the full KV length. All ranks + see the same local shape (only head counts shrink). + - ``sequence``: keep the heads, split the KV length across ``cp_size`` shards; the + local KV length depends on ``rank % cp_size`` (remainder distributed to the low + shards). Q is intentionally not divided. + """ + if shard_mode == "none": + if cp_size != 1: + raise ValueError(f"--shard-mode none requires --cp-size 1, got {cp_size}") + return LocalAttentionShape(q_len, kv_len, num_q_heads, num_kv_heads, 0, 0) + + if shard_mode == "head": + if num_q_heads % cp_size != 0: + raise ValueError(f"num_q_heads={num_q_heads} must be divisible by cp_size={cp_size}") + if num_kv_heads % cp_size != 0: + raise ValueError(f"num_kv_heads={num_kv_heads} must be divisible by cp_size={cp_size}") + local_q_heads = num_q_heads // cp_size + local_kv_heads = num_kv_heads // cp_size + if local_q_heads % local_kv_heads != 0: + raise ValueError(f"local_q_heads={local_q_heads} must be divisible by local_kv_heads={local_kv_heads}") + return LocalAttentionShape(q_len, kv_len, local_q_heads, local_kv_heads, 0, 0) + + if shard_mode == "sequence": + shard_index = rank % cp_size + base_kv_len = kv_len // cp_size + remainder = kv_len % cp_size + local_kv_len = base_kv_len + int(shard_index < remainder) + if local_kv_len <= 0: + raise ValueError( + f"sequence-sharded local_kv_len must be positive, got {local_kv_len}; " + f"kv_len={kv_len}, cp_size={cp_size}, shard_index={shard_index}" + ) + sequence_shard_start = shard_index * base_kv_len + min(shard_index, remainder) + return LocalAttentionShape(q_len, local_kv_len, num_q_heads, num_kv_heads, shard_index, sequence_shard_start) + + raise ValueError(f"Unsupported shard_mode={shard_mode!r}") + + +def attention_flops( + batch: int, + q_len: int, + kv_len: int, + num_q_heads: int, + head_dim: int, + head_dim_v: int, + is_causal: bool, + include_backward: bool, +) -> float: + """ + Approximate attention FLOPs (supports q_len != kv_len for cross-attention). + + Forward is 2 GEMMs (QK^T and P@V); each multiply-add counts as 2 FLOPs. + Backward is ~2x the forward cost. Causal masking roughly halves the work + (only applicable when q_len == kv_len). + """ + qk = 2.0 * batch * num_q_heads * q_len * kv_len * head_dim # [B,Hq,Sq,Skv] scores + pv = 2.0 * batch * num_q_heads * q_len * kv_len * head_dim_v # [B,Hq,Sq,Dv] output + flops = qk + pv + if is_causal: + flops *= 0.5 + if include_backward: + flops *= 3.0 # fwd (1x) + bwd (~2x) + return flops + + +def make_inputs( + batch: int, + q_len: int, + kv_len: int, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + head_dim_v: int, + dtype: torch.dtype, + device: torch.device, + requires_grad: bool, + generator: torch.Generator | None, +) -> tuple[Tensor, Tensor, Tensor]: + """Create heads-last dense QKV tensors (q_len may differ from kv_len).""" + q = torch.randn( + batch, q_len, num_q_heads, head_dim, dtype=dtype, device=device, generator=generator + ).requires_grad_(requires_grad) # [B,Sq,Hq,D] + k = torch.randn( + batch, kv_len, num_kv_heads, head_dim, dtype=dtype, device=device, generator=generator + ).requires_grad_(requires_grad) # [B,Skv,Hkv,D] + v = torch.randn( + batch, kv_len, num_kv_heads, head_dim_v, dtype=dtype, device=device, generator=generator + ).requires_grad_(requires_grad) # [B,Skv,Hkv,Dv] + return q, k, v + + +def compatibility_reason( + q: Tensor, # [B,S,Hq,D] + k: Tensor, # [B,S,Hkv,D] + v: Tensor, # [B,S,Hkv,Dv] + backend: str | None, + is_causal: bool, +) -> str | None: + """ + Surface the concrete reason a backend was rejected for this use case / device. + + The frontend runs backend selection with ``raise_error=False`` and only reports a + generic "incompatible" message, swallowing the underlying reason. Here we re-run + ``choose_backend`` with ``raise_error=True`` (a distinct lru_cache key, so it does not + hit the frontend's cached result) to capture the specific rejection reason. + + Returns the reason string, or None if the backend is compatible. + """ + causal_type = CausalType.DontCare if is_causal else None + try: + choose_backend( + query_shape=q.shape, + key_shape=k.shape, + value_shape=v.shape, + dtype=q.dtype, + device=q.device, + requires_grad=q.requires_grad or k.requires_grad or v.requires_grad, + is_causal=is_causal, + causal_type=causal_type, + is_varlen=False, + deterministic=False, + backend=backend, + raise_error=True, + ) + return None + except Exception as e: # noqa: BLE001 - surface the concrete rejection reason + return f"{type(e).__name__}: {e}" + + +def run_once( + q: Tensor, # [B,S,Hq,D] + k: Tensor, # [B,S,Hkv,D] + v: Tensor, # [B,S,Hkv,Dv] + backend: str | None, + is_causal: bool, + include_backward: bool, +) -> Tensor: + """Run a single forward (and optional backward) attention call. + + Returns the attention output. The caller MUST keep/consume the returned tensor: + under ``torch.compile`` an unused output lets Inductor dead-code-eliminate the whole + attention (yielding absurd, ~PFLOP/s "timings"), so the output is returned as a graph + output to keep the kernel live. + """ + causal_type = CausalType.DontCare if is_causal else None + out = attention( + query=q, + key=k, + value=v, + is_causal=is_causal, + causal_type=causal_type, + backend=backend, + return_lse=False, + ) # [B,S,Hq,Dv] + if include_backward: + grad = torch.ones_like(out) # [B,S,Hq,Dv] + out.backward(grad) + return out + + +def benchmark_backend( + q: Tensor, # [B,S,Hq,D] + k: Tensor, # [B,S,Hkv,D] + v: Tensor, # [B,S,Hkv,Dv] + backend: str | None, + is_causal: bool, + include_backward: bool, + warmup: int, + iters: int, + compile_fn: bool, + shard_mode: str, + cp_size: int, + debug: bool, +) -> tuple[list[float], int, str | None]: + """ + Time a single backend on a single (already-sharded) input shape. + + Returns ``(per-iter latencies in ms, peak allocated bytes, error message)``. + If the backend is incompatible or errors out, latencies is empty, peak bytes is 0, + and the error message describes why. When ``debug`` is set, an incompatible backend's + concrete rejection reason is reported instead of the generic frontend message. + + The forward (+ optional backward) is wrapped in an NVTX range so the timed region is + easy to find in Nsight Systems. Backward requires grad, so ``inference_mode`` is only + used for the forward-only path. ``compile_fn`` compiles the call before timing. + """ + if debug: + reason = compatibility_reason(q, k, v, backend, is_causal) + if reason is not None: + return [], 0, reason + + def _call() -> Tensor: + return run_once(q, k, v, backend, is_causal, include_backward) + + call = torch.compile(_call, fullgraph=not include_backward) if compile_fn else _call + + # inference_mode is incompatible with autograd; only use it for forward-only timing. + ctx = torch.inference_mode() if not include_backward else torch.enable_grad() + + if dist.is_initialized(): + dist.barrier() + + # Sink holding the latest output so Inductor cannot dead-code-eliminate the attention + # (an unused compiled output would be optimized away, giving impossible TFLOP/s). + sink: Tensor | None = None + try: + with ctx: + for _ in range(warmup): + sink = call() + torch.cuda.synchronize() + + torch.cuda.reset_peak_memory_stats(q.device) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + latencies_ms: list[float] = [] + torch.cuda.nvtx.range_push(f"fmha.{shard_mode}.cp{cp_size}.{backend or 'auto'}") + for _ in range(iters): + start.record() + sink = call() + end.record() + torch.cuda.synchronize() + latencies_ms.append(start.elapsed_time(end)) # ms + torch.cuda.nvtx.range_pop() + # Force a host-side read of the result so the compiled output is genuinely + # materialized (belt-and-suspenders against dead-code elimination). + if sink is not None: + _ = sink.sum().item() + except Exception as e: # noqa: BLE001 - report and skip incompatible backends + return [], 0, f"{type(e).__name__}: {e}" + + peak_bytes = torch.cuda.max_memory_allocated(q.device) + return latencies_ms, peak_bytes, None + + +def format_row(cells: list[str], widths: list[int]) -> str: + return " ".join(cell.ljust(w) for cell, w in zip(cells, widths)) + + +def init_distributed() -> tuple[int, int, int]: + """Read torchrun env (defaults to single process) and init NCCL when world_size > 1.""" + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + if world_size > 1 and not dist.is_initialized(): + dist.init_process_group("nccl") + return rank, local_rank, world_size + + +def resolve_seqlen_pair(seqlen: int, kv_seqlen: int | None, causal: bool) -> tuple[int, int]: + """Resolve the single (q_len, kv_len) shape. Causal masking here uses + CausalType.DontCare, which requires q_len == kv_len, so a distinct KV length is only + allowed when not causal.""" + if kv_seqlen is None: + return (seqlen, seqlen) + if causal: + raise SystemExit( + "--kv-seqlen is only valid without --causal: causal masking here uses CausalType.DontCare, " + "which requires q_len == kv_len. Drop --kv-seqlen, or drop --causal." + ) + return (seqlen, kv_seqlen) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Benchmark cosmos_framework.model.attention (dense / non-varlen only) across backends, " + "sequence lengths, and optional context-parallel shard shapes." + ) + parser.add_argument( + "--backends", + nargs="+", + default=["auto"], + help="Backends to benchmark. Use 'auto' to let the frontend choose. " + "Named choices: flash2, flash3, natten, cudnn.", + ) + parser.add_argument("--seqlen", type=int, default=4096, help="Query sequence length.") + parser.add_argument( + "--kv-seqlen", + type=int, + default=None, + help="KV sequence length for cross-attention (q_len != kv_len). Only allowed when " + "--causal is NOT set. If omitted, kv_len == q_len (self-attention).", + ) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--num-q-heads", type=int, default=16) + parser.add_argument("--num-kv-heads", type=int, default=8) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument( + "--head-dim-v", type=int, default=None, help="Value head dim (defaults to --head-dim; set differently for MLA)." + ) + parser.add_argument("--dtype", choices=list(DTYPE_MAP), default="bf16") + parser.add_argument("--causal", action="store_true", help="Enable causal masking (CausalType.DontCare).") + parser.add_argument("--backward", action="store_true", help="Include a backward pass (training) in timing.") + parser.add_argument( + "--shard-mode", + choices=("none", "head", "sequence"), + default="none", + help="Context-parallel per-rank shape to simulate. none: no sharding; " + "head: divide Q/KV heads by --cp-size (full KV length); " + "sequence: keep heads and split KV length across --cp-size shards.", + ) + parser.add_argument( + "--cp-size", type=int, default=1, help="Context-parallel sharding factor (>1 needs --shard-mode)." + ) + parser.add_argument("--compile", action="store_true", help="torch.compile the attention call before benchmarking.") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--seed", type=int, default=1234, help="Random seed for input generation.") + parser.add_argument("--device", default="cuda") + parser.add_argument("--json", action="store_true", help="Emit one JSON object per result row instead of a table.") + parser.add_argument( + "--debug", + action="store_true", + help="Report the concrete reason a backend is rejected (instead of the generic " + "'incompatible' message) and enable DEBUG-level backend-selection logs.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + if args.debug: + # Raise the logger to DEBUG so the per-backend selection diagnostics (emitted by the + # attention checks via log.debug) are printed, then re-initialize the stdout sink. + from cosmos_framework.utils import log + + log.LEVEL = "DEBUG" + log.init_loguru_stdout() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA is not available; this benchmark requires a GPU.") + + if args.cp_size < 1: + raise SystemExit(f"--cp-size must be >= 1, got {args.cp_size}") + if args.shard_mode == "sequence" and args.causal: + raise SystemExit( + "--shard-mode sequence is incompatible with --causal: a KV-split shard holds a " + "partial key range, so a causal mask over the local shard is not well-defined here." + ) + + rank, local_rank, world_size = init_distributed() + is_main = rank == 0 + device = torch.device("cuda", local_rank) if args.device == "cuda" else torch.device(args.device) + dtype = DTYPE_MAP[args.dtype] + head_dim_v = args.head_dim_v if args.head_dim_v is not None else args.head_dim + backends: list[str | None] = [None if b == "auto" else b for b in args.backends] + generator = torch.Generator(device=device).manual_seed(args.seed) + + q_len, kv_len = resolve_seqlen_pair(args.seqlen, args.kv_seqlen, args.causal) + local = local_shape_for_pair( + q_len=q_len, + kv_len=kv_len, + num_q_heads=args.num_q_heads, + num_kv_heads=args.num_kv_heads, + shard_mode=args.shard_mode, + cp_size=args.cp_size, + rank=rank, + ) + seq_label = str(local.q_len) if local.q_len == local.kv_len else f"{local.q_len}/{local.kv_len}" + heads_label = f"{local.num_q_heads}/{local.num_kv_heads}" + shard_label = "-" if args.shard_mode == "none" else f"{args.shard_mode}/{args.cp_size}" + + if is_main and not args.json: + arch_tag = get_arch_tag(device) + print(f"Device: {torch.cuda.get_device_name(device)} (arch_tag={arch_tag}, sm_{arch_tag})") + print( + f"Config: batch={args.batch} Hq={args.num_q_heads} Hkv={args.num_kv_heads} " + f"D={args.head_dim} Dv={head_dim_v} dtype={args.dtype} causal={args.causal} " + f"backward={args.backward} compile={args.compile}" + ) + print( + f"Shard: mode={args.shard_mode} cp_size={args.cp_size} world_size={world_size} | " + f"Timing: warmup={args.warmup} iters={args.iters}\n" + ) + headers = [ + "backend", + "seq(q/kv)", + "heads(q/kv)", + "shard", + "median(ms)", + "min(ms)", + "TFLOP/s", + "tok/s", + "peakMB", + "status", + ] + widths = [10, 15, 12, 10, 11, 10, 9, 12, 9, 40] + print(format_row(headers, widths)) + print(format_row(["-" * w for w in widths], widths)) + else: + widths = [10, 15, 12, 10, 11, 10, 9, 12, 9, 40] + + for backend in backends: + backend_label = backend if backend is not None else "auto" + + q, k, v = make_inputs( + batch=args.batch, + q_len=local.q_len, + kv_len=local.kv_len, + num_q_heads=local.num_q_heads, + num_kv_heads=local.num_kv_heads, + head_dim=args.head_dim, + head_dim_v=head_dim_v, + dtype=dtype, + device=device, + requires_grad=args.backward, + generator=generator, + ) + + latencies_ms, peak_bytes, error = benchmark_backend( + q, + k, + v, + backend=backend, + is_causal=args.causal, + include_backward=args.backward, + warmup=args.warmup, + iters=args.iters, + compile_fn=args.compile, + shard_mode=args.shard_mode, + cp_size=args.cp_size, + debug=args.debug, + ) + + if not is_main: + continue + + if error is not None: + if args.json: + print(json.dumps({"backend": backend_label, "seq": seq_label, "status": error})) + else: + print( + format_row( + [backend_label, seq_label, heads_label, shard_label, "-", "-", "-", "-", "-", error], widths + ) + ) + continue + + median_ms = statistics.median(latencies_ms) + min_ms = min(latencies_ms) + flops = attention_flops( + batch=args.batch, + q_len=local.q_len, + kv_len=local.kv_len, + num_q_heads=local.num_q_heads, + head_dim=args.head_dim, + head_dim_v=head_dim_v, + is_causal=args.causal, + include_backward=args.backward, + ) + tflops = flops / (median_ms * 1e-3) / 1e12 + tokens_per_s = (args.batch * local.q_len) / (median_ms * 1e-3) + peak_mb = peak_bytes / (1024**2) + + if args.json: + print( + json.dumps( + { + "backend": backend_label, + "q_len": local.q_len, + "kv_len": local.kv_len, + "num_q_heads": local.num_q_heads, + "num_kv_heads": local.num_kv_heads, + "shard_mode": args.shard_mode, + "cp_size": args.cp_size, + "median_ms": median_ms, + "min_ms": min_ms, + "tflops": tflops, + "tokens_per_s": tokens_per_s, + "peak_mb": peak_mb, + "status": "ok", + } + ) + ) + else: + print( + format_row( + [ + backend_label, + seq_label, + heads_label, + shard_label, + f"{median_ms:.4f}", + f"{min_ms:.4f}", + f"{tflops:.1f}", + f"{tokens_per_s:.0f}", + f"{peak_mb:.1f}", + "ok", + ], + widths, + ) + ) + + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/cosmos_framework/model/attention/cudnn/__init__.py b/cosmos_framework/model/attention/cudnn/__init__.py index ca3697b9..c585d320 100644 --- a/cosmos_framework/model/attention/cudnn/__init__.py +++ b/cosmos_framework/model/attention/cudnn/__init__.py @@ -11,12 +11,10 @@ import torch from cosmos_framework.model.attention.utils.safe_ops import log -from cosmos_framework.model.attention.utils.version import version_at_least -CUDNN_DISALLOWED = True - -CUDNN_MIN_BACKEND_VERSION = 91300 -CUDNN_MIN_FRONTEND_VERSION = "1.14.0" +# Minimum cuDNN runtime version, in ``torch.backends.cudnn.version()`` encoding +# (major * 10000 + minor * 100 + patch). 92200 == cuDNN 9.22.0. +CUDNN_MIN_BACKEND_VERSION = 92200 def cudnn_supported() -> bool: @@ -24,36 +22,34 @@ def cudnn_supported() -> bool: Returns whether cuDNN Attention is supported in this environment. Requirements are: * Presence of CUDA Runtime (via PyTorch) - * Presence of cuDNN and its Python frontend, meeting minimum version requirements + * Presence of the cuDNN runtime that ships with / is linked by PyTorch, meeting the minimum + version requirement - This check guards imports / dependencies on the cuDNN package. + The backend runs cuDNN attention through PyTorch's own SDPA cuDNN dispatch + (``F.scaled_dot_product_attention`` / ``aten._scaled_dot_product_cudnn_attention``), so it no + longer depends on the standalone cuDNN Python frontend package -- only on the cuDNN that PyTorch + itself uses. """ if not torch.cuda.is_available(): log.debug("cuDNN Attention is not supported because PyTorch did not detect CUDA runtime.") return False - try: - import cudnn - - except ImportError: - log.debug("cuDNN Attention is not supported because the frontend Python package was not found.") - return False - except Exception as e: - log.debug(f"cuDNN Attention is not supported because importing the frontend Python package failed: {e}") + if not torch.backends.cudnn.is_available(): + log.debug("cuDNN Attention is not supported because PyTorch reports cuDNN is unavailable.") return False - if cudnn.backend_version() < CUDNN_MIN_BACKEND_VERSION: + backend_version = torch.backends.cudnn.version() + if backend_version is None or backend_version < CUDNN_MIN_BACKEND_VERSION: log.debug( - "cuDNN Attention is not supported due to insufficient cuDNN backend version " - f"{cudnn.backend_version()=}, expected at least {CUDNN_MIN_BACKEND_VERSION=}." + "cuDNN Attention is not supported due to insufficient cuDNN runtime version " + f"{backend_version=}, expected at least {CUDNN_MIN_BACKEND_VERSION=}." ) return False - if not version_at_least(cudnn.__version__, CUDNN_MIN_FRONTEND_VERSION): - log.debug( - "cuDNN Attention is not supported due to insufficient cuDNN frontend version " - f"{cudnn.__version__}, expected at least {CUDNN_MIN_FRONTEND_VERSION}." - ) + # The cuDNN SDPA ATen op is the mechanism this backend relies on; if it is missing (unexpected on + # a CUDA build), decline rather than fail later at call time. + if not hasattr(torch.ops.aten, "_scaled_dot_product_cudnn_attention"): + log.debug("cuDNN Attention is not supported because torch lacks the cuDNN SDPA ATen op.") return False return True diff --git a/cosmos_framework/model/attention/cudnn/checks.py b/cosmos_framework/model/attention/cudnn/checks.py index ffc5f96a..e6e941af 100644 --- a/cosmos_framework/model/attention/cudnn/checks.py +++ b/cosmos_framework/model/attention/cudnn/checks.py @@ -13,10 +13,83 @@ import torch from cosmos_framework.model.attention.checks import attention_param_checks, attention_tensor_checks -from cosmos_framework.model.attention.cudnn import CUDNN_DISALLOWED, CUDNN_SUPPORTED -from cosmos_framework.model.attention.cudnn.meta import get_bwd_dtypes, get_fwd_dtypes +from cosmos_framework.model.attention.cudnn import CUDNN_SUPPORTED +from cosmos_framework.model.attention.cudnn.meta import CUDNN_HEAD_DIM_ALIGNMENT, get_bwd_dtypes, get_fwd_dtypes, get_max_head_dim from cosmos_framework.model.attention.masks import CausalType -from cosmos_framework.model.attention.utils import get_arch_tag, is_torch_compiling, log_or_raise_error +from cosmos_framework.model.attention.utils import get_arch_tag, log_or_raise_error + + +def cudnn_sdpa_eligible( + query_shape: torch.Size, + key_shape: torch.Size, + value_shape: torch.Size, + arch_tag: int, + raise_error: bool = False, +) -> bool: + """ + Mirror the cuDNN SDPA eligibility constraints that the generic ``attention_tensor_checks`` does + not cover, so unsupported shapes are rejected here (and the chooser can fall back to another + backend such as NATTEN) instead of failing deep inside the raw + ``torch.ops.aten._scaled_dot_product_cudnn_attention`` operator. + + PyTorch's cuDNN eligibility logic additionally rejects: + * a key/value sequence length of 1, and + * head dims (Q/K and V) that are not a multiple of ``CUDNN_HEAD_DIM_ALIGNMENT`` (8) or that + exceed the architecture-specific maximum (see ``get_max_head_dim``, typically 128). + + This helper is intentionally tensorless and takes ``arch_tag`` directly (rather than a device) + so the eligibility rules can be unit-tested for a specific architecture (e.g. SM120) without + requiring that physical GPU. + + Parameters: + query_shape (torch.Size): Shape of 4-D query tensor (`[batch, seqlen, heads, head_dim]`). + + key_shape (torch.Size): Shape of 4-D key tensor (`[batch, seqlen_kv, heads_kv, head_dim]`). + + value_shape (torch.Size): Shape of 4-D value tensor (`[batch, seqlen_kv, heads_kv, head_dim_v]`). + + arch_tag (int): Arch tag for the target CUDA device (see ``get_arch_tag``). Example: 90 for + Hopper, 120 for workstation/consumer Blackwell. + + raise_error (bool): whether to raise an error if any checks fail, instead of just returning + False. Default is False. + + Returns: + success (bool): whether the shapes satisfy cuDNN's SDPA eligibility constraints. + """ + target_fn = partial(log_or_raise_error, raise_error=raise_error) + + # cuDNN's fused attention rejects a KV sequence length of 1. Layout is [batch, seqlen, heads, + # head_dim], so the KV sequence length is index 1 of the key shape. + seqlen_kv = key_shape[1] + if seqlen_kv == 1: + target_fn( + f"cuDNN Attention does not support a key/value sequence length of 1, got {seqlen_kv=}.", + exception=ValueError, + ) + return False + + max_head_dim = get_max_head_dim(arch_tag) + # QK and V head dims are the last dimension of each tensor. Q/K head-dim equality and (absent + # MLA) Q/V head-dim equality are already enforced by attention_tensor_checks, but we validate + # both explicitly for precise, actionable error messages. + for name, head_dim in (("QK", query_shape[-1]), ("V", value_shape[-1])): + if head_dim % CUDNN_HEAD_DIM_ALIGNMENT != 0: + target_fn( + f"cuDNN Attention requires the {name} head dim to be a multiple of " + f"{CUDNN_HEAD_DIM_ALIGNMENT}, got {head_dim=}.", + exception=ValueError, + ) + return False + if head_dim > max_head_dim: + target_fn( + f"cuDNN Attention on this architecture ({arch_tag=}) supports a maximum {name} head " + f"dim of {max_head_dim}, got {head_dim=}.", + exception=ValueError, + ) + return False + + return True def cudnn_attention_check( @@ -76,21 +149,14 @@ def cudnn_attention_check( ) return False - if CUDNN_DISALLOWED: - target_fn("cuDNN backend is disabled. Please choose another backend.", exception=RuntimeError) - return False - if deterministic: target_fn("cuDNN Attention does not support deterministic mode.", exception=RuntimeError) return False - if is_torch_compiling(): - target_fn( - "cuDNN backend does not support torch.compile yet.", - exception=RuntimeError, - ) - return False - + # cuDNN Attention supports both the forward (inference) and backward (training) passes: it runs on + # native, differentiable PyTorch ops (F.scaled_dot_product_attention / the cuDNN SDPA ATen op), so + # autograd handles the backward pass. attention_tensor_checks validates the backward dtype below + # (see get_bwd_dtypes) whenever operands require grad. arch_tag = get_arch_tag(device) fwd_dtypes = get_fwd_dtypes(arch_tag) bwd_dtypes = get_bwd_dtypes(arch_tag) @@ -103,13 +169,27 @@ def cudnn_attention_check( supported_dtypes_forward=fwd_dtypes, supported_dtypes_backward=bwd_dtypes, supports_mla=False, - supports_gqa_mqa=False, + supports_gqa_mqa=True, raise_error=raise_error, backend_name="cuDNN Attention", ): target_fn("cuDNN does not support the given inputs.", exception=RuntimeError) return False + # Mirror PyTorch's cuDNN SDPA eligibility constraints that attention_tensor_checks does not + # cover (KV sequence length 1, head-dim alignment, and the architecture-specific head-dim + # maximum). cuDNN is first in the Blackwell backend order, so without this the chooser would + # accept these shapes and then fail inside the raw cuDNN ATen operator instead of falling back + # to another backend (e.g. NATTEN). + if not cudnn_sdpa_eligible( + query_shape=query_shape, + key_shape=key_shape, + value_shape=value_shape, + arch_tag=arch_tag, + raise_error=raise_error, + ): + return False + if is_varlen: target_fn("Varlen for cuDNN Attention is not integrated yet.", exception=RuntimeError) return False diff --git a/cosmos_framework/model/attention/cudnn/functions.py b/cosmos_framework/model/attention/cudnn/functions.py index 6d06c59c..5a2d943b 100644 --- a/cosmos_framework/model/attention/cudnn/functions.py +++ b/cosmos_framework/model/attention/cudnn/functions.py @@ -7,147 +7,57 @@ cuDNN Backend: intermediate APIs Only safe to import when CUDNN_SUPPORTED is True. -""" -import time -from functools import partial +This backend runs cuDNN attention through PyTorch's *own* cuDNN SDPA dispatch rather than the +standalone cuDNN Python frontend (``import cudnn`` + ``cudnn.pygraph``). It always uses +``torch.ops.aten._scaled_dot_product_cudnn_attention`` -- the exact ATen op that +``torch.nn.functional.scaled_dot_product_attention`` lowers to for the cuDNN backend -- and simply +discards the logsumexp when ``return_lse=False``. Unlike SDPA's public API, this op also exposes the +logsumexp statistics needed for the ``return_lse=True`` path. + +It is a first-class, meta-registered PyTorch op, so it traces cleanly under +``torch.compile(fullgraph=True)``. This deliberately replaces the previous cuDNN-frontend graph +builder, which had to be hidden behind an opaque custom op (Dynamo cannot trace the frontend's +sourceless enum objects) and required a special TorchInductor flag to run on Thor -- a flag that +regressed other GPUs. +""" import torch from torch import Tensor -from torch.amp import custom_bwd, custom_fwd -from torch.autograd import Function from cosmos_framework.model.attention.checks import assert_universal_tensor_checks from cosmos_framework.model.attention.cudnn.checks import cudnn_attention_check -from cosmos_framework.model.attention.cudnn.cudnn_forward import ( - cudnn_sdpa_fwd_generate_op, - cudnn_sdpa_fwd_generate_operands, - cudnn_sdpa_fwd_post_process, -) from cosmos_framework.model.attention.masks import CausalType -from cosmos_framework.model.attention.utils.safe_ops import log - -amp_fwd = partial(custom_fwd, device_type="cuda") -amp_bwd = partial(custom_bwd, device_type="cuda") - - -CUDNN_PADDING_REQUIRED = False - - -class CudnnAttentionAutogradFn(Function): - @staticmethod - @amp_fwd - def forward( - ctx, - query: Tensor, - key: Tensor, - value: Tensor, - num_heads: int, - is_causal: bool, - scale: float, - ) -> tuple[Tensor, Tensor]: - query = query.contiguous() - key = key.contiguous() - value = value.contiguous() - - seqlen_Q = None - seqlen_KV = None - padding_Q = 0 - padding_KV = 0 - - if CUDNN_PADDING_REQUIRED: - Q_multiplier = 256 - KV_multiplier = 256 - - if query.shape[1] % Q_multiplier != 0: - seqlen_Q = query.shape[1] - padding_Q = Q_multiplier - (seqlen_Q % Q_multiplier) - - old_shape = query.shape - query = torch.nn.functional.pad(query, (0, 0, 0, 0, 0, padding_Q), "constant", 0) - log.debug(f"cuDNN Attention: padded query from {old_shape} to {query.shape}.") - - if key.shape[1] % KV_multiplier != 0: - seqlen_KV = key.shape[1] - padding_KV = KV_multiplier - (seqlen_KV % KV_multiplier) - - old_shape = key.shape - key = torch.nn.functional.pad(key, (0, 0, 0, 0, 0, padding_KV), "constant", 0) - value = torch.nn.functional.pad(value, (0, 0, 0, 0, 0, padding_KV), "constant", 0) - log.debug(f"cuDNN Attention: padded KV from {old_shape} to {key.shape}.") - - # Transform operands to cuDNN-compatible layouts, make output tensors - (q_cudnn_layout, k_cudnn_layout, v_cudnn_layout, output_cudnn_layout, lse_cudnn_layout) = ( - cudnn_sdpa_fwd_generate_operands(q=query, k=key, v=value, num_heads=num_heads, return_lse=True) - ) - - # Construct graph - assert q_cudnn_layout.device == k_cudnn_layout.device == v_cudnn_layout.device == output_cudnn_layout.device - assert q_cudnn_layout.dtype == k_cudnn_layout.dtype == v_cudnn_layout.dtype == output_cudnn_layout.dtype - cudnn_graph_gen_start = time.time() * 1e3 - cudnn_sdpa = cudnn_sdpa_fwd_generate_op( - dtype=q_cudnn_layout.dtype, - device=q_cudnn_layout.device, - q_shape=q_cudnn_layout.shape, - q_stride=q_cudnn_layout.stride(), - k_shape=k_cudnn_layout.shape, - k_stride=k_cudnn_layout.stride(), - v_shape=v_cudnn_layout.shape, - v_stride=v_cudnn_layout.stride(), - output_shape=output_cudnn_layout.shape, - output_stride=output_cudnn_layout.stride(), - lse_shape=None if lse_cudnn_layout is None else lse_cudnn_layout.shape, - lse_stride=None if lse_cudnn_layout is None else lse_cudnn_layout.stride(), - is_causal=is_causal, - attn_scale=scale, - seqlen_Q=seqlen_Q, - seqlen_KV=seqlen_KV, - ) - cudnn_graph_gen_time = time.time() * 1e3 - cudnn_graph_gen_start - log.debug(f"cuDNN Attention forward graph generation took {cudnn_graph_gen_time:.1f} ms.") - - # Execute graph - cudnn_sdpa( - q=q_cudnn_layout, - k=k_cudnn_layout, - v=v_cudnn_layout, - output=output_cudnn_layout, - lse=lse_cudnn_layout, - ) - - # Transform outputs back to torch contiguous layouts - output, logsumexp = cudnn_sdpa_fwd_post_process( - output_cudnn_layout=output_cudnn_layout, - lse_cudnn_layout=lse_cudnn_layout, - ) - - ctx.save_for_backward(q_cudnn_layout, k_cudnn_layout, v_cudnn_layout, lse_cudnn_layout, output_cudnn_layout) - ctx.num_heads = num_heads - ctx.scale = scale - - if padding_Q > 0: - old_shape = output.shape - output = output[:, :seqlen_Q, :, :] - logsumexp = logsumexp[:, :seqlen_Q, :, :] - assert output.shape[1] == seqlen_Q - assert logsumexp.shape[1] == seqlen_Q - log.debug(f"cuDNN Attention: unpadded output from {old_shape} to {output.shape}.") - - return output, logsumexp - - @staticmethod - @amp_bwd - def backward( - ctx, grad_out: Tensor, grad_lse: Tensor - ) -> tuple[ - Tensor, - Tensor, - Tensor, - None, - None, - None, - ]: - raise NotImplementedError() + + +def _cudnn_sdpa_with_lse(q: Tensor, k: Tensor, v: Tensor, is_causal: bool, scale: float) -> tuple[Tensor, Tensor]: + """cuDNN SDPA returning logsumexp, via torch's native cuDNN ATen op. + + ``torch.ops.aten._scaled_dot_product_cudnn_attention`` is the op ``F.scaled_dot_product_attention`` + dispatches to for the cuDNN backend; unlike the public API it exposes the logsumexp statistics. The + op is autograd-aware, so backward is provided automatically (no explicit autograd wrapper here). + The op requires equal Q/KV head counts, so the caller must expand K/V heads to match Q beforehand + (GQA/MQA). Inputs are heads-first ``[B, H, S, D]``; returns output ``[B, H, S, Dv]`` and logsumexp + ``[B, H, S]`` (float32). Note some torch versions return logsumexp with a trailing singleton + dim (``[B, H, S, 1]``); callers must normalize the rank. + + Only positional args ``(query, key, value, attn_bias, compute_log_sumexp, dropout_p, is_causal, + return_debug_mask)`` plus the keyword-only ``scale`` are used; these indices/names have been + stable across torch versions, and only the first two outputs (output, logsumexp) are consumed. + """ + results = torch.ops.aten._scaled_dot_product_cudnn_attention( + q, + k, + v, + None, # attn_bias + True, # compute_log_sumexp + 0.0, # dropout_p + is_causal, + False, # return_debug_mask + scale=scale, + ) + output, logsumexp = results[0], results[1] # [B,H,S,Dv], [B,H,S] or [B,H,S,1] + return output, logsumexp def cudnn_attention( @@ -167,7 +77,7 @@ def cudnn_attention( ) -> Tensor | tuple[Tensor, Tensor]: """ Runs cuDNN Attention on given operands (Q, K, V) with the heads-last contiguous layout - (`[batch, seqlen, heads, head_dim]`). + (`[batch, seqlen, heads, head_dim]`), dispatched through PyTorch's built-in cuDNN SDPA. Parameters: query (Tensor): 4-D query tensor, with the heads-last contiguous layout @@ -182,46 +92,46 @@ def cudnn_attention( is_causal (bool): whether or not causal masking is enabled. Default is False. causal_type (CausalType): causal masking mode. Choices: `CausalType.TopLeft`, - `CausalType.BottomRight`. Required when `is_causal = True`. + `CausalType.DontCare`. Required when `is_causal = True`. cuDNN SDPA's causal masking is + top-left aligned. scale (float | None): Dot product scale (attention scale). Defaults to head_dim ** -0.5. - cumulative_seqlen_Q (Tensor | None): (varlen) Optional 1-D tensor with size `batch + 1` - indicating the cumulative sum of number of query tokens in each batch, with an - additional 0 element in the beginning. Must be passed together with - `cumulative_seqlen_KV` and `max_seqlen_{Q,KV}`. + cumulative_seqlen_Q (Tensor | None): (varlen) Not supported by this backend. - cumulative_seqlen_KV (Tensor | None): (varlen) Optional 1-D tensor with size `batch + 1` - indicating the cumulative sum of number of key/value tokens in each batch, with an - additional 0 element in the beginning. Must be passed together with - `cumulative_seqlen_Q` and `max_seqlen_{Q,KV}`. + cumulative_seqlen_KV (Tensor | None): (varlen) Not supported by this backend. - max_seqlen_Q (int | None): (varlen) Optional integer indicating the maximum query - sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}` - and `max_seqlen_KV`. + max_seqlen_Q (int | None): (varlen) Not supported by this backend. - max_seqlen_KV (int | None): (varlen) Optional integer indicating the maximum key/value - sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}` - and `max_seqlen_Q`. + max_seqlen_KV (int | None): (varlen) Not supported by this backend. Other Parameters: return_lse (bool): Whether to return the logsumexp values. Default is False. - backend_kwargs (dict | None): Key-value pair for passing arguments specific to cuDNN's - attention operator, if any. + backend_kwargs (dict | None): Key-value pair for passing backend-specific arguments. Only + ``deterministic`` is recognized (and must be False); any other key raises an error. - deterministic (bool): Deterministic backward pass required. + deterministic (bool): Deterministic backward pass required. Not supported by this backend. Returns: output (Tensor): 4-D output tensor, with the heads-last contiguous layout (`[batch, seqlen, heads, head_dim_v]`). - logsumexp (Tensor): logsumexp tensor, with the heads-last contiguous layout - (`[batch, seqlen, heads, 1]`). Only returned when return_lse is True. + logsumexp (Tensor): logsumexp tensor, with the heads-last layout + (`[batch, seqlen, heads]`). Only returned when return_lse is True. + NOTE: not guaranteed to be contiguous (it is a transposed view) and must not be + made contiguous, so its results stay correct when merged via `merge_attentions`. """ is_varlen = cumulative_seqlen_Q is not None assert_universal_tensor_checks(query, key, value) + + backend_kwargs = backend_kwargs.copy() if backend_kwargs is not None else {} + # Determinism in backend_kwargs supersedes primary flag, if set to True + if "deterministic" in backend_kwargs: + deterministic = deterministic or backend_kwargs["deterministic"] + del backend_kwargs["deterministic"] + assert cudnn_attention_check( query_shape=query.shape, key_shape=key.shape, @@ -236,21 +146,52 @@ def cudnn_attention( raise_error=True, ) - assert not is_varlen # cudnn_attention_check should prevent this assertion failing - - num_heads = query.shape[-2] - scale = scale if scale is not None else query.shape[-1] ** -0.5 + # cudnn_attention_check should prevent this assertion failing (varlen is not integrated). + assert not is_varlen - output, lse = CudnnAttentionAutogradFn.apply( - query, - key, - value, - num_heads, - is_causal, - scale, - ) + # cuDNN SDPA takes no extra operator arguments; reject anything unrecognized instead of silently + # ignoring it. + if backend_kwargs: + raise ValueError(f"cuDNN Attention backend received unsupported backend_kwargs: {sorted(backend_kwargs)}.") - if return_lse: - return output, lse + scale = scale if scale is not None else query.shape[-1] ** -0.5 - return output + heads = query.shape[-2] + heads_kv = key.shape[-2] + assert heads % heads_kv == 0 + group_size = heads // heads_kv + + # The cuDNN ATen op expects the heads-first layout: [B,H,S,D]. + q = query.transpose(1, 2) # [B,H,S_q,D] + k = key.transpose(1, 2) # [B,H_kv,S_kv,D] + v = value.transpose(1, 2) # [B,H_kv,S_kv,D_v] + + # Always take the LSE-capable ATen path (the same op SDPA lowers to for cuDNN) and simply drop + # the logsumexp when it isn't requested. This keeps a single code path for both cases. Tradeoff: + # the raw ATen op requires equal Q/KV head counts, so GQA/MQA must materialize expanded K/V here, + # whereas the public SDPA API (enable_gqa=True) can handle grouped heads without materializing. + if group_size > 1: + k = k.repeat_interleave(group_size, dim=1, output_size=heads) # [B,H,S_kv,D] + v = v.repeat_interleave(group_size, dim=1, output_size=heads) # [B,H,S_kv,D_v] + + output, logsumexp = _cudnn_sdpa_with_lse(q, k, v, is_causal, scale) # [B,H,S_q,D_v], [B,H,S_q(,1)] + + assert output.dim() == 4 + output = output.transpose(1, 2).contiguous() # [B,S_q,H,D_v] + + if not return_lse: + return output + + # The cuDNN ATen op returns logsumexp with a trailing singleton dim on some torch + # versions ([B,H,S_q,1]); drop it so LSE is rank-3, matching every other backend and + # what `merge_attentions` requires. Guarded on rank so we neither crash on torch builds + # that already return [B,H,S_q] nor squeeze the heads dim when S_q/H is 1. + if logsumexp.dim() == 4: + logsumexp = logsumexp.squeeze(-1) # [B,H,S_q] + + # NOTE: Do NOT call .contiguous() on LSE. Attention merging's backward pass requires the + # output and LSE tensors passed into `merge_attentions` to share the same (transposed) + # data layout; forcing contiguity here makes that backward pass incorrect (see flash2). + logsumexp = logsumexp.transpose(1, 2) # [B,S_q,H] + assert logsumexp.dim() == 3 + return output, logsumexp diff --git a/cosmos_framework/model/attention/cudnn/meta.py b/cosmos_framework/model/attention/cudnn/meta.py index 5cd15344..c4e2092f 100644 --- a/cosmos_framework/model/attention/cudnn/meta.py +++ b/cosmos_framework/model/attention/cudnn/meta.py @@ -13,6 +13,38 @@ from cosmos_framework.model.attention.utils.safe_ops import log +# cuDNN's fused attention (SDPA) requires every head dim (Q/K and V) to be a multiple of 8. This +# mirrors PyTorch's own cuDNN SDPA eligibility logic (see check_cudnn_head_dim in sdp_utils.cpp). +CUDNN_HEAD_DIM_ALIGNMENT = 8 + + +def get_max_head_dim(arch_tag: int) -> int: + """ + Returns the maximum head dim cuDNN's fused attention (SDPA) supports for the given arch. + + PyTorch's cuDNN SDPA eligibility logic enforces an architecture-specific head-dim maximum + (typically 128). We stay conservative: when unsure we prefer the lower bound so the chooser + falls back to another backend (e.g. NATTEN) rather than accepting a shape that then fails + inside the raw ``torch.ops.aten._scaled_dot_product_cudnn_attention`` operator. + + Parameters: + arch_tag (int): Arch tag for the current CUDA device. Example: 80 for A100, 90 for H100, + 120 for workstation/consumer Blackwell (e.g. RTX PRO 6000). + + Returns: + max_head_dim (int): the maximum supported head dim, or 0 if the arch is unsupported. + """ + if arch_tag < 80: + log.debug("cuDNN Attention is not supported because compute capability is below the minimum (8.0).") + return 0 + + # Hopper's cuDNN fused attention supports head dims up to 256 (FP16/BF16). Every other + # currently-supported arch (Ampere/Ada and the Blackwell family, including workstation/consumer + # sm_120/121) is held to the conservative 128 maximum until larger dims are verified. + if arch_tag == 90: + return 256 + return 128 + def get_fwd_dtypes(arch_tag: int) -> list[torch.dtype]: """ @@ -30,6 +62,7 @@ def get_fwd_dtypes(arch_tag: int) -> list[torch.dtype]: log.debug("cuDNN Attention is not supported because compute capability is below the minimum (8.0).") return [] + # As of version 91400 FP8 inference via the python frontend does not seem to work. log.debug(f"cuDNN Attention only supports FP16 and BF16 for {arch_tag=}.") return [torch.float16, torch.bfloat16] @@ -46,4 +79,10 @@ def get_bwd_dtypes(arch_tag: int) -> list[torch.dtype]: """ - return [] + if arch_tag < 80: + log.debug("cuDNN Attention is not supported because compute capability is below the minimum (8.0).") + return [] + + # cuDNN SDPA backward supports the same FP16/BF16 precisions as the forward pass. + log.debug(f"cuDNN Attention backward supports FP16 and BF16 for {arch_tag=}.") + return [torch.float16, torch.bfloat16] diff --git a/cosmos_framework/model/attention/frontend.py b/cosmos_framework/model/attention/frontend.py index 018932a8..dfc93863 100644 --- a/cosmos_framework/model/attention/frontend.py +++ b/cosmos_framework/model/attention/frontend.py @@ -21,6 +21,7 @@ universal_tensor_checks, varlen_tensor_checks, ) +from cosmos_framework.model.attention.cudnn import cudnn_attention from cosmos_framework.model.attention.flash2 import flash2_attention from cosmos_framework.model.attention.flash3 import flash3_attention from cosmos_framework.model.attention.masks import CausalType @@ -28,9 +29,9 @@ from cosmos_framework.model.attention.utils.environment import filter_attention_merge_backends from cosmos_framework.model.attention.utils.safe_ops import log - # Map backend names to their frontend attention API BACKEND_MAP = { + "cudnn": cudnn_attention, "natten": natten_attention, "flash2": flash2_attention, "flash3": flash3_attention, diff --git a/cosmos_framework/model/attention/utils/__init__.py b/cosmos_framework/model/attention/utils/__init__.py index 67eaaae5..e7aece57 100644 --- a/cosmos_framework/model/attention/utils/__init__.py +++ b/cosmos_framework/model/attention/utils/__init__.py @@ -60,6 +60,14 @@ def is_blackwell_dc(device: torch.device | None = None) -> bool: return get_arch_tag(device) in [100, 103] +def is_blackwell(device: torch.device | None = None) -> bool: + """True for any Blackwell GPU: data-center (sm_100/103), Thor (sm_110, DRIVE/Jetson Thor), + or workstation/consumer (sm_120/121, e.g. RTX PRO 6000 / RTX 50-series). This matches the + non-data-center grouping ``[110, 120, 121]`` in ``get_backend_list``. Use ``is_blackwell_dc`` + when a check must be restricted to the data-center parts (e.g. kernels only qualified there).""" + return get_arch_tag(device) in [100, 103, 110, 120, 121] + + __all__ = [ "get_arch_tag", "log_or_raise_error", @@ -68,6 +76,7 @@ def is_blackwell_dc(device: torch.device | None = None) -> bool: "is_fp8", "is_hopper", "is_blackwell_dc", + "is_blackwell", "is_torch_compiling", "torch_deterministic_mode", ] diff --git a/cosmos_framework/model/generator/mot/attention.py b/cosmos_framework/model/generator/mot/attention.py index f0d87660..7eb7fa6c 100644 --- a/cosmos_framework/model/generator/mot/attention.py +++ b/cosmos_framework/model/generator/mot/attention.py @@ -111,34 +111,74 @@ def two_way_attention( causal_v, _ = get_causal_seq(packed_value_states) full_q, full_q_offsets = get_full_only_seq(packed_query_states) + # NOTE: we can only use the don't care causal mask when we know seqlen_Q == seqlen_KV. + # Since this is a varlen use case, we would need to statically check all Q and KV offsets + # are the same. + # We don't want to launch a kernel just to perform this check and slow down our model, and + # we definitely don't want to complicate the sequence_packing code so that it performs a + # static check when creating the packed sequence and metadata. Instead, we just rely + # on causal_q_offsets and causal_k_offsets being the same tensor. + use_dont_care_mask = causal_q_offsets is causal_k_offsets + sample_offsets = packed_query_states["sample_offsets"] - use_dont_care_mask = causal_q_offsets is causal_k_offsets + # Number of packed samples. sample_offsets has shape [num_samples + 1], so this is a static + # (metadata) shape read: it does not launch a kernel or force a device sync. + num_samples = sample_offsets.shape[0] - 1 + + # With a single sample there is exactly one sequence in the pack, so the varlen (sequence-packed) + # metadata is redundant and we can call the dense attention API instead (no cumulative/max seqlen + # args). This remains correct in the presence of trailing padding: for causal self-attention the + # mask never lets a real query attend to padded keys (padding is appended after all real tokens), + # get_all_seq returns unpadded KV for the full path, and any padded query rows are independent of + # the real rows and simply discarded downstream. + # + # The dense path is gated to forward-only (inference) execution via + # torch.is_grad_enabled(), which is False under torch.no_grad()/torch.inference_mode() + # and True during training. This avoids branching on the sample count during training, + # where batch composition varies between single- and multi-sample packs; keeping a single + # code path there prevents torch.compile from specializing on both shapes and incurring + # the associated recompilation overhead. + use_dense = num_samples == 1 and not torch.is_grad_enabled() + + if use_dense: + causal_varlen_kwargs = {} + else: + causal_varlen_kwargs = dict( + cumulative_seqlen_Q=causal_q_offsets, + cumulative_seqlen_KV=causal_k_offsets, + max_seqlen_Q=packed_query_states["max_causal_len"], + max_seqlen_KV=packed_query_states["max_causal_len"], + ) # NOTE: cosmos_framework attention is BSHD in, BSHD out causal_res = attention( causal_q.unsqueeze(0), # [1,N_und,heads,head_dim] causal_k.unsqueeze(0), # [1,N_und,heads,head_dim] causal_v.unsqueeze(0), # [1,N_und,heads,head_dim] - cumulative_seqlen_Q=causal_q_offsets, - cumulative_seqlen_KV=causal_k_offsets, - max_seqlen_Q=packed_query_states["max_causal_len"], - max_seqlen_KV=packed_query_states["max_causal_len"], is_causal=True, causal_type=CausalType.DontCare if use_dont_care_mask else CausalType.TopLeft, + **causal_varlen_kwargs, ) # [1,N_und,heads,head_dim] # [1,N_und,heads,head_dim] -> [N_und,heads,head_dim] -> [N_und,heads*head_dim] causal_out = causal_res.squeeze(0).flatten(-2, -1) # type: ignore # [N_und,heads*head_dim] + if use_dense: + full_varlen_kwargs = {} + else: + full_varlen_kwargs = dict( + cumulative_seqlen_Q=full_q_offsets, + cumulative_seqlen_KV=sample_offsets, + max_seqlen_Q=packed_query_states["max_full_len"], + max_seqlen_KV=packed_query_states["max_sample_len"], + ) + full_res = attention( full_q.unsqueeze(0), # [1,N_full,heads,head_dim] get_all_seq(packed_key_normalized).unsqueeze(0), # [1,N_all,heads,head_dim] normed und K for gen get_all_seq(packed_value_states).unsqueeze(0), # [1,N_all,heads,head_dim] - cumulative_seqlen_Q=full_q_offsets, - cumulative_seqlen_KV=sample_offsets, - max_seqlen_Q=packed_query_states["max_full_len"], - max_seqlen_KV=packed_query_states["max_sample_len"], + **full_varlen_kwargs, ) # [1,N_full,heads,head_dim] # [1,N_full,heads,head_dim] -> [N_full,heads,head_dim] -> [N_full,heads*head_dim] diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index 656364b2..6fed5dee 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -392,13 +392,9 @@ def set_up_memory(self) -> None: def set_up_parallelism(self) -> None: """Set up the fsdp for the model.""" - if not torch.distributed.is_initialized(): - self.parallel_dims = None - return - self.parallel_dims = ParallelDims( enable_inference_mode=self.config.parallelism.enable_inference_mode, - world_size=torch.distributed.get_world_size(), + world_size=torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1, dp_shard=self.config.parallelism.data_parallel_shard_degree, cfgp=self.config.parallelism.cfg_parallel_shard_degree, cp=self.config.parallelism.context_parallel_shard_degree, @@ -1666,8 +1662,13 @@ def _prepare_inference_data( ) # 2. Get data and condition (same as training) - # This encodes vision to x0_tokens - gen_data_clean = self.get_data_and_condition(data_batch) + # This encodes vision to x0_tokens. Pass each sample's vision conditioning + # frame indexes so a causal tokenizer can skip encoding pixel frames that only + # feed generated (non-conditioned) latent positions (e.g. policy / + # forward-dynamics condition latent frame 0 only, so only the first pixel frame + # is encoded instead of the whole clip). + vision_condition_indexes = [plan.condition_frame_indexes_vision for plan in sequence_plans] + gen_data_clean = self.get_data_and_condition(data_batch, vision_condition_indexes=vision_condition_indexes) num_items_per_sample = gen_data_clean.num_vision_items_per_sample # None for standard T2I/T2V @@ -2846,11 +2847,139 @@ def validation_step(self, data_batch: dict[str, torch.Tensor], iteration: int): def forward(self, xt, t): pass - def get_data_and_condition(self, data_batch: dict[str, torch.Tensor], iteration: int = 1) -> GenerationDataClean: + def _encode_vision_x0_tokens( + self, + raw_state_vision: list[torch.Tensor], + num_vision_items_per_sample: list[int] | None, + vision_condition_indexes: list[list[int]] | None, + ) -> list[torch.Tensor]: + """Encode vision items into x0 latent tokens, optionally skipping unused frames. + + Default behavior (``vision_condition_indexes is None``) encodes every pixel + frame of every vision item. This is the path used during training and by any + caller that does not opt in, and it is byte-for-byte the original encode loop. + + Inference optimization: when ``vision_condition_indexes`` is provided (one + conditioning-frame-index list per sample, from each ``SequencePlan``) and the + vision tokenizer is temporally causal, only the pixel-frame prefix required to + reconstruct the highest conditioned latent frame is encoded. The remaining + latent frames are generated (pure-noise) positions that the ``condition_mask`` + blend discards, so they are zero-filled instead of encoded. Under a causal VAE + the kept latents are identical to the corresponding frames of a full-clip + encode, while the encode processes far fewer frames (e.g. 1 pixel frame instead + of the full clip for policy / forward-dynamics modes that only condition latent + frame 0). Inverse-dynamics conditions every latent frame, so it keeps the full + encode automatically. + + The optimization falls back to a full encode for multi-vision samples, + non-causal tokenizers, samples with no conditioning frames, and any item whose + full clip is already the minimal prefix. + """ + + def _full_encode(state: torch.Tensor) -> torch.Tensor: + return self.encode(state).contiguous().float() + + # Only opt in when a caller supplied per-sample conditioning indexes, the + # samples map 1:1 to vision items (no multi-vision flattening), and the + # tokenizer is causal (so a pixel prefix reproduces the leading latents). + optimization_applicable = ( + vision_condition_indexes is not None + and num_vision_items_per_sample is None + and self.tokenizer_vision_gen is not None + and self.tokenizer_vision_gen.is_causal + and len(vision_condition_indexes) == len(raw_state_vision) + ) + if not optimization_applicable: + return [_full_encode(raw_state_vision_i) for raw_state_vision_i in raw_state_vision] + + # The prefix-encode optimization zero-fills the generated (pure-noise) latent + # frames that the condition_mask blend later discards. That is only valid when + # no gradient is required: torch.is_grad_enabled() is False under both + # torch.inference_mode() and torch.no_grad(), so this asserts we are on an + # inference path and never silently corrupts a training forward. + if torch.is_grad_enabled(): + raise ValueError( + "prefix-encode optimization is inference-only, but grad is enabled " + "(expected torch.inference_mode()/torch.no_grad())" + ) + + tokenizer = self.tokenizer_vision_gen + assert vision_condition_indexes is not None # narrowed by optimization_applicable above + x0_tokens_vision: list[torch.Tensor] = [] + for raw_state_vision_i, condition_indexes in zip(raw_state_vision, vision_condition_indexes, strict=True): + # The temporal axis is the first of the trailing (T, H, W) dims for both the + # [B, C, T, H, W] and [C, T, H, W] layouts the tokenizer accepts, and encode + # preserves rank so the same index locates T in the latent output. + temporal_dim = raw_state_vision_i.ndim - 3 + num_pixel_frames = int(raw_state_vision_i.shape[temporal_dim]) + num_latent_frames = tokenizer.get_latent_num_frames(num_pixel_frames) + + valid_condition = [idx for idx in condition_indexes if 0 <= idx < num_latent_frames] + if not valid_condition: + # No conditioning frames to preserve (e.g. fully generated T2V item): + # keep the full encode so behavior is unchanged for non-action items. + x0_tokens_vision.append(_full_encode(raw_state_vision_i)) + continue + + needed_latent_frames = max(valid_condition) + 1 + needed_pixel_frames = tokenizer.get_pixel_num_frames(needed_latent_frames) + assert needed_pixel_frames <= num_pixel_frames, ( + f"needed_pixel_frames ({needed_pixel_frames}) cannot be greater than " + f"num_pixel_frames ({num_pixel_frames})" + ) + if needed_pixel_frames == num_pixel_frames: + # The conditioning already spans (nearly) the whole clip; no work saved. + x0_tokens_vision.append(_full_encode(raw_state_vision_i)) + continue + + prefix = raw_state_vision_i.narrow(temporal_dim, 0, needed_pixel_frames) + prefix_latent = _full_encode(prefix) # leading latent frames, causally exact + + # The scatter below assumes the pixel->latent round trip is exact, i.e. that + # encoding ``needed_pixel_frames`` yields exactly ``needed_latent_frames`` + # latents. All current causal tokenizers satisfy this, but guard it: an + # overshoot would make the narrow below request more frames than full_latent + # holds (confusing runtime error), and an undershoot would silently zero-fill + # the highest conditioning frame and corrupt the condition_mask blend with no + # error at all. + num_prefix_latent_frames = prefix_latent.shape[temporal_dim] + assert num_prefix_latent_frames == needed_latent_frames, ( + f"causal tokenizer pixel<->latent round trip is not exact: encoding " + f"{needed_pixel_frames} pixel frames (from get_pixel_num_frames({needed_latent_frames})) " + f"produced {num_prefix_latent_frames} latent frames, expected {needed_latent_frames}" + ) + + # Scatter the encoded prefix into a full-length latent; the unconditioned + # tail frames stay zero since the condition_mask blend replaces them with noise. + full_shape = list(prefix_latent.shape) + full_shape[temporal_dim] = num_latent_frames + full_latent = prefix_latent.new_zeros(full_shape) + full_latent.narrow(temporal_dim, 0, prefix_latent.shape[temporal_dim]).copy_(prefix_latent) + x0_tokens_vision.append(full_latent) + + return x0_tokens_vision + + def get_data_and_condition( + self, + data_batch: dict[str, torch.Tensor], + iteration: int = 1, + vision_condition_indexes: list[list[int]] | None = None, + ) -> GenerationDataClean: """ - Get raw data of different modalities from databatch - Tokenize into corresponding latents - Load other conditioning information if any (fps, etc.) + + Args: + data_batch: Raw data batch from the dataloader. + iteration: Current iteration (only used for time-based logging during training). + vision_condition_indexes: Optional per-sample list of conditioning latent + frame indexes (one list per sample, taken from each ``SequencePlan``). + When provided (inference only), it enables the causal-VAE prefix-encode + optimization in ``_encode_vision_x0_tokens`` that skips encoding pixel + frames which only feed generated (non-conditioned) latent positions. + Leaving it ``None`` (the default, used during training) encodes every + frame — the original behavior. """ # Detect whether any sample has multiple vision items (e.g. image editing). # If so, track the count per sample before all vision items from this batch are flattened into a list. @@ -2898,13 +3027,16 @@ def get_data_and_condition(self, data_batch: dict[str, torch.Tensor], iteration: if log_enc_time: timer = Timer(unit="s") timer.start() + # Vision (image/video) raw state and tokenized latent state self._normalize_video_databatch_inplace(data_batch) self._augment_image_dim_inplace(data_batch) # converts each image tensor to (1, C, 1, H, W) raw_state_vision = data_batch[self.input_image_key if is_image_batch else self.input_video_key] - x0_tokens_vision = [ - self.encode(raw_state_vision_i).contiguous().float() for raw_state_vision_i in raw_state_vision - ] + x0_tokens_vision = self._encode_vision_x0_tokens( + raw_state_vision, + num_vision_items_per_sample, + vision_condition_indexes, + ) frame_size = data_batch.get("image_size", None) if frame_size is not None: diff --git a/cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py b/cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py index a9e0dac7..c65788bb 100644 --- a/cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py +++ b/cosmos_framework/model/generator/tokenizers/uniae/noncausal_4x16x16.py @@ -17,7 +17,10 @@ recon = vae.decode(latents) # [B, 48, T_p, H//16, W//16] -> [B, 3, 4*T_p, H, W] """ +import os from collections.abc import Mapping, Sequence +from numbers import Real +from typing import Any import torch @@ -32,52 +35,174 @@ normalize_resolution_int_mapping, ) from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution +from cosmos_framework.model.tokenizer.checkpoint_io import ( + DCP_MODEL_LOAD_INFO_KEY, + DCPModelLoadInfo, + load_dcp_model_checkpoint, + load_torch_checkpoint, +) +from cosmos_framework.model.tokenizer.models.architecture import get_siglip2_so400m_common_arch from cosmos_framework.model.tokenizer.models.dense_runtime import DenseAutoencoderRuntime from cosmos_framework.model.tokenizer.models.sparse_autoencoder import AutoencoderKL -# S3 architecture config (avoids importing configs/base which pulls in loss deps) _S3_ARCH = dict( - patch_size=(4, 16, 16), - in_channels=3072, - out_channels=3072, - # Encoder - encoder_model_channels=1152, - encoder_num_blocks=27, - encoder_num_heads=16, - encoder_mlp_channels=4304, - encoder_pe_mode="joint", - encoder_qk_rms_norm=False, - encoder_use_bias=True, - encoder_use_rms_norm=False, - # Decoder - decoder_model_channels=1152, - decoder_num_blocks=27, - decoder_num_heads=16, - decoder_mlp_channels=4304, - decoder_pe_mode="joint", - decoder_qk_rms_norm=True, - decoder_use_bias=False, - decoder_use_rms_norm=True, - # Common settings - use_decoder=True, - quantizer_type="rq", - quantizer_codebook_size=65536, - quantizer_num_codebooks=1, - quantizer_chunk_size=1, - use_vf_loss=False, - freeze_encoder=False, - pretrained_model_name="google/siglip2-so400m-patch16-naflex", - concat_latent=None, - random_num_sample_frames_batch_sizes=[8, 12, 16, 20, 24], - inference_num_sample_frames_batch_size=16, - inference_num_sample_frames_stride=16, - inference_kv_cache_size=0, + **get_siglip2_so400m_common_arch(), use_quantizer=False, use_dual_latent=False, use_text_alignment=False, use_post_text_alignment=False, ) +_IGNORED_LEGACY_CHECKPOINT_PREFIXES = ( + "ema.", + "loss_fn.", + "text_decoder_wrapper.", + "tokenizer.", +) +_LEGACY_CHECKPOINT_SUFFIXES = (".pt", ".pth", ".ckpt") + + +def _is_legacy_checkpoint_file(checkpoint_path: str) -> bool: + """Distinguish torch checkpoint files from DCP component directories.""" + is_remote = checkpoint_path.startswith("s3://") + if not is_remote and os.path.isdir(checkpoint_path): + return False + return checkpoint_path.lower().endswith(_LEGACY_CHECKPOINT_SUFFIXES) or ( + not is_remote and os.path.isfile(checkpoint_path) + ) + + +def _derive_legacy_latent_norm_path(checkpoint_path: str) -> str: + """Derive the latent-stat sidecar for a legacy checkpoint file.""" + checkpoint_stem, checkpoint_suffix = os.path.splitext(checkpoint_path) + if checkpoint_suffix.lower() in _LEGACY_CHECKPOINT_SUFFIXES: + return checkpoint_stem + "_latent_norm.pt" + return checkpoint_path + "_latent_norm.pt" + + +def _coerce_latent_norm_vector( + values: object, + *, + field_name: str, + z_dim: int, + source: str, + strictly_positive: bool = False, +) -> torch.Tensor: + """Convert one tensor/JSON vector to validated CPU float64 statistics.""" + if isinstance(values, torch.Tensor): + if values.dtype == torch.bool: + raise ValueError(f"Latent-normalization {field_name} in {source} must contain real numbers, not booleans.") + if values.is_complex(): + raise ValueError( + f"Latent-normalization {field_name} in {source} must contain real numbers, not complex values." + ) + vector = values.detach().to(device="cpu", dtype=torch.float64) # [C_candidate] + elif isinstance(values, Sequence) and not isinstance(values, (str, bytes)): + if any(isinstance(value, bool) or not isinstance(value, Real) for value in values): + raise ValueError(f"Latent-normalization {field_name} in {source} must contain only real numbers.") + vector = torch.tensor(list(values), dtype=torch.float64) # [C_candidate] + else: + raise TypeError( + f"Latent-normalization {field_name} in {source} must be a tensor or numeric sequence, " + f"got {type(values).__name__}." + ) + + if vector.ndim != 1 or vector.shape[0] != z_dim: + raise ValueError( + f"Latent-normalization {field_name} in {source} must have shape ({z_dim},), got {tuple(vector.shape)}." + ) + finite_mask = torch.isfinite(vector) # [C] + if not bool(finite_mask.all().item()): + raise ValueError(f"Latent-normalization {field_name} in {source} contains non-finite values.") + if strictly_positive: + positive_mask = vector > 0 # [C] + if not bool(positive_mask.all().item()): + raise ValueError(f"Latent-normalization {field_name} in {source} must be strictly positive.") + return vector + + +def _load_latent_norm_stats( + norm_path: str, + *, + backend_args: dict[str, str] | None, + z_dim: int, + dtype: torch.dtype, + device: str, +) -> tuple[torch.Tensor, torch.Tensor]: + """Load latent statistics and return validated runtime mean and inverse standard deviation.""" + if norm_path.lower().endswith(".json"): + norm_stats = easy_io.load(norm_path, backend_args=backend_args) + else: + norm_stats = easy_io.load(norm_path, backend_args=backend_args, map_location="cpu", weights_only=True) + if not isinstance(norm_stats, Mapping): + raise TypeError( + f"Latent-normalization sidecar {norm_path} must contain a mapping, got {type(norm_stats).__name__}." + ) + + stats_z_dim = norm_stats.get("z_dim") + if stats_z_dim is not None and (isinstance(stats_z_dim, bool) or not isinstance(stats_z_dim, int)): + raise ValueError(f"Latent-normalization z_dim in {norm_path} must be an integer, got {stats_z_dim!r}.") + if isinstance(stats_z_dim, int) and stats_z_dim != z_dim: + raise ValueError(f"Latent-normalization z_dim in {norm_path} is {stats_z_dim}, expected {z_dim}.") + if "mean" not in norm_stats or "std" not in norm_stats: + raise ValueError(f"Latent-normalization sidecar {norm_path} must contain mean and std entries.") + + mean_cpu = _coerce_latent_norm_vector( + norm_stats["mean"], + field_name="mean", + z_dim=z_dim, + source=norm_path, + ) # [C] + std_cpu = _coerce_latent_norm_vector( + norm_stats["std"], + field_name="std", + z_dim=z_dim, + source=norm_path, + strictly_positive=True, + ) # [C] + mean = mean_cpu.to(dtype=dtype, device=device) # [C] + std = std_cpu.to(dtype=dtype, device=device) # [C] + mean_finite_mask = torch.isfinite(mean) # [C] + std_finite_positive_mask = torch.isfinite(std) & (std > 0) # [C] + if not bool(mean_finite_mask.all().item()): + raise ValueError(f"Latent-normalization mean in {norm_path} is non-finite after conversion to {dtype}.") + if not bool(std_finite_positive_mask.all().item()): + raise ValueError( + f"Latent-normalization std in {norm_path} must remain finite and positive after conversion to {dtype}." + ) + inv_std = std.reciprocal() # [C] + inv_std_finite_mask = torch.isfinite(inv_std) # [C] + if not bool(inv_std_finite_mask.all().item()): + raise ValueError(f"Latent-normalization inverse std in {norm_path} is non-finite after conversion to {dtype}.") + return mean, inv_std + + +def _extract_visual_tokenizer_state_dict(model_state: Mapping[str, Any]) -> dict[str, Any]: + """Extract bare AutoencoderKL state from a full TokenizerModel checkpoint.""" + network_state = { + key.removeprefix("network."): value for key, value in model_state.items() if key.startswith("network.") + } + return network_state if network_state else dict(model_state) + + +def _get_dcp_unexpected_visual_keys(load_info: DCPModelLoadInfo) -> list[str]: + """Return unexpected keys owned by the tokenizer network from a DCP report.""" + return [key.removeprefix("network.") for key in load_info.unexpected_checkpoint_keys if key.startswith("network.")] + + +def _validate_uniae_checkpoint_keys(missing: list[str], unexpected: list[str], checkpoint_path: str) -> None: + """Reject promoted checkpoints that do not fully cover the visual tokenizer.""" + if missing: + raise RuntimeError( + f"UniAE checkpoint {checkpoint_path} is missing {len(missing)} visual model keys (sample: {missing[:5]})." + ) + invalid_unexpected = [key for key in unexpected if not key.startswith(_IGNORED_LEGACY_CHECKPOINT_PREFIXES)] + if invalid_unexpected: + raise RuntimeError( + f"UniAE checkpoint {checkpoint_path} has {len(invalid_unexpected)} unexpected visual model keys " + f"(sample: {invalid_unexpected[:5]})." + ) + class UniAEVAE: """UniAE S3 VAE wrapper for diffusion training. @@ -95,6 +220,7 @@ def __init__( z_dim: int = 48, vae_pth: str = "", object_store_credential_path_pretrained: str = "", + latent_norm_path: str = "", dtype: torch.dtype = torch.bfloat16, device: str = "cuda", backend: str = "batched", @@ -102,7 +228,9 @@ def __init__( pixel_trim: bool = True, chunk_size: int | Mapping[str, int] = 16, encode_chunk_batch_size: int | Mapping[str, int] = 1, - ): + ) -> None: + if torch.device(device).type == "meta": + raise ValueError("UniAEVAE requires a concrete CPU or CUDA device; device='meta' is not supported.") chunk_size = normalize_resolution_int_mapping(chunk_size, name="chunk_size") if any(chunk_frames % 4 != 0 for chunk_frames in chunk_size.values()): raise ValueError("chunk_size values must be multiples of 4.") @@ -135,11 +263,14 @@ def __init__( if not vae_pth: raise ValueError("vae_pth must be provided to load latent normalization stats") vae_pth_str = str(vae_pth) - # Derive stats path: strip .pt suffix, append _latent_norm.pt - if vae_pth_str.endswith(".pt"): - norm_pth = vae_pth_str[:-3] + "_latent_norm.pt" + is_legacy_checkpoint = _is_legacy_checkpoint_file(vae_pth_str) + # Legacy files use a paired _latent_norm.pt sidecar. + if latent_norm_path: + norm_pth = latent_norm_path + elif is_legacy_checkpoint: + norm_pth = _derive_legacy_latent_norm_path(vae_pth_str) else: - norm_pth = vae_pth_str + "_latent_norm.pt" + raise ValueError("latent_norm_path is required when vae_pth points to a DCP checkpoint directory.") if norm_pth.startswith("s3://"): norm_backend_args = { "backend": "s3", @@ -147,13 +278,17 @@ def __init__( } else: norm_backend_args = None - norm_stats = easy_io.load(norm_pth, backend_args=norm_backend_args, map_location="cpu", weights_only=False) - mean = norm_stats["mean"].to(dtype=dtype, device=device) - std = norm_stats["std"].to(dtype=dtype, device=device) - self._latent_mean = mean.view(1, z_dim, 1, 1, 1) - self._latent_inv_std = (1.0 / std).view(1, z_dim, 1, 1, 1) - - # make compatible with meta device + mean, inv_std = _load_latent_norm_stats( + norm_pth, + backend_args=norm_backend_args, + z_dim=z_dim, + dtype=dtype, + device=device, + ) # [C],[C] + self._latent_mean = mean.view(1, z_dim, 1, 1, 1) # [1,C,1,1,1] + self._latent_inv_std = inv_std.view(1, z_dim, 1, 1, 1) # [1,C,1,1,1] + + # Construct the autoencoder on the requested concrete device. autoencoder = AutoencoderKL( **_S3_ARCH, latent_channels=z_dim, @@ -164,27 +299,44 @@ def __init__( # Load checkpoint if vae_pth and get_rank() == 0: - if str(vae_pth).startswith("s3://"): - backend_args = {"backend": "s3", "s3_credential_path": object_store_credential_path_pretrained} + dcp_unexpected_keys: list[str] = [] + if is_legacy_checkpoint: + if str(vae_pth).startswith("s3://"): + backend_args = { + "backend": "s3", + "s3_credential_path": object_store_credential_path_pretrained, + } + else: + backend_args = None + state_dict = load_torch_checkpoint( + vae_pth, + backend_args=backend_args, + map_location="cpu", + ) + if "model" in state_dict: + model_state = state_dict["model"] + elif "state_dict" in state_dict: + model_state = state_dict["state_dict"] + else: + model_state = state_dict + model_state = _extract_visual_tokenizer_state_dict(model_state) else: - backend_args = None - state_dict = easy_io.load(vae_pth, backend_args=backend_args, map_location="cpu", weights_only=False) - if "model" in state_dict: + state_dict = load_dcp_model_checkpoint( + autoencoder, + vae_pth, + checkpoint_key_prefix="network.", + s3_credential=object_store_credential_path_pretrained, + ) model_state = state_dict["model"] - elif "state_dict" in state_dict: - model_state = state_dict["state_dict"] - else: - model_state = state_dict - # Checkpoint may be saved from a wrapper with a 'network.' prefix — strip it. - if any(k.startswith("network.") for k in model_state): - model_state = { - k[len("network.") :] if k.startswith("network.") else k: v for k, v in model_state.items() - } + dcp_load_info = state_dict[DCP_MODEL_LOAD_INFO_KEY] + if not isinstance(dcp_load_info, DCPModelLoadInfo): + raise TypeError(f"Invalid DCP model load info: {type(dcp_load_info).__name__}") + dcp_unexpected_keys = _get_dcp_unexpected_visual_keys(dcp_load_info) missing, unexpected = autoencoder.load_state_dict(model_state, strict=False) - if missing: - log.warning(f"Missing keys: {len(missing)} (e.g., {missing[:3]})") - if unexpected: - log.warning(f"Unexpected keys: {len(unexpected)} (e.g., {unexpected[:3]})") + all_unexpected = [*unexpected, *dcp_unexpected_keys] + _validate_uniae_checkpoint_keys(missing, all_unexpected, str(vae_pth)) + if all_unexpected: + log.info(f"Ignored {len(all_unexpected)} non-visual checkpoint keys (e.g., {all_unexpected[:3]}).") log.info(f"Loaded checkpoint from {vae_pth}") elif vae_pth: autoencoder.to_empty(device=device) @@ -364,16 +516,25 @@ def __init__( bucket_name: str = "", object_store_credential_path_pretrained: str = "", vae_path: str = "", + latent_norm_path: str = "", encode_chunk_frames: int | Mapping[str, int] = 16, encode_chunk_batch_size: int | Mapping[str, int] = 1, spatial_compression_factor: int = 16, temporal_compression_factor: int = 4, - pad_frames: int = 0, + pad_frames: int = 1, pixel_trim: bool = True, backend: str = "batched_with_padding", causal: bool = False, - ): - super().__init__(object_store_credential_path_pretrained) + ) -> None: + if spatial_compression_factor != 16: + raise ValueError( + f"UniAEVAEInterface requires spatial_compression_factor=16, got {spatial_compression_factor}." + ) + if temporal_compression_factor != 4: + raise ValueError( + f"UniAEVAEInterface requires temporal_compression_factor=4, got {temporal_compression_factor}." + ) + super().__init__(object_store_credential_path_pretrained or None) self._causal = causal assert not self._causal, "UniAEVAEInterface is a non-causal tokenizer; causal must be False." self._spatial_compression_factor = spatial_compression_factor @@ -400,12 +561,16 @@ def __init__( self.use_streaming_encode = False vae_full_path = vae_path + latent_norm_full_path = latent_norm_path if bucket_name and not vae_path.startswith("s3://"): vae_full_path = f"s3://{bucket_name}/{vae_path}" + if bucket_name and latent_norm_path and not latent_norm_path.startswith("s3://"): + latent_norm_full_path = f"s3://{bucket_name}/{latent_norm_path}" self.vae = UniAEVAE( vae_pth=vae_full_path, object_store_credential_path_pretrained=object_store_credential_path_pretrained, + latent_norm_path=latent_norm_full_path, pad_frames=pad_frames, pixel_trim=pixel_trim, backend=backend, @@ -414,7 +579,7 @@ def __init__( ) self.is_compiled = False - def reset_dtype(self): + def reset_dtype(self) -> None: pass def encode(self, state: torch.Tensor) -> torch.Tensor: @@ -483,11 +648,11 @@ def get_latent_temporal_positions( ) @property - def spatial_compression_factor(self): + def spatial_compression_factor(self) -> int: return self._spatial_compression_factor @property - def temporal_compression_factor(self): + def temporal_compression_factor(self) -> int: return self._temporal_compression_factor @property diff --git a/cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py b/cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py index fea6320d..0117a13c 100644 --- a/cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py +++ b/cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py @@ -1285,7 +1285,10 @@ def _collect_warmup_shapes( raise ValueError(f"Aspect ratio {aspect_ratio} not found in resolution {res_key}") res_dict = {aspect_ratio: res_dict[aspect_ratio]} - for H, W in res_dict.values(): + # ``VIDEO_RES_SIZE_INFO`` stores each bucket as ``(W, H)`` (e.g. the + # landscape ``"4,3": (736, 544)``), so unpack width-first to match the + # runtime encode key, which is derived from the actual ``(H, W)`` tensor. + for W, H in res_dict.values(): if isinstance(tokenizer.encode_chunk_frames, Mapping): chunk_frames = tokenizer.encode_chunk_frames[res_key] else: diff --git a/cosmos_framework/model/tokenizer/checkpoint_io.py b/cosmos_framework/model/tokenizer/checkpoint_io.py new file mode 100644 index 00000000..1f2079a7 --- /dev/null +++ b/cosmos_framework/model/tokenizer/checkpoint_io.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Safe checkpoint loading helpers for tokenizer code.""" + +from __future__ import annotations + +import io +import os +import re +from collections import defaultdict +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from omegaconf.base import ContainerMetadata, Metadata +from omegaconf.dictconfig import DictConfig +from omegaconf.listconfig import ListConfig +from omegaconf.nodes import AnyNode + +from cosmos_framework.utils.easy_io import easy_io + +TorchMapLocation = str | torch.device | Callable[[Any, str], Any] | dict[str, str] | None +DCP_MODEL_LOAD_INFO_KEY = "dcp_model_load_info" +_RQ_CODEBOOK_STAGE_KEY_PATTERN = re.compile(r"(?:^|\.)quantizer\.codebooks\.(\d+)(?:\.|$)") + + +@dataclass(frozen=True) +class DCPModelLoadInfo: + """Model-key coverage discovered from one DCP component's metadata.""" + + loaded_model_keys: frozenset[str] + missing_model_keys: frozenset[str] + unexpected_checkpoint_keys: frozenset[str] + + +def infer_rq_codebook_depth(state_keys: Iterable[str], *, source: str) -> int | None: + """Infer a contiguous residual-quantizer depth from state-dict keys.""" + stage_indices = { + int(match.group(1)) + for key in state_keys + if (match := _RQ_CODEBOOK_STAGE_KEY_PATTERN.search(str(key))) is not None + } + if not stage_indices: + return None + + inferred_depth = max(stage_indices) + 1 + expected_indices = set(range(inferred_depth)) + if stage_indices != expected_indices: + raise ValueError( + f"Residual-quantizer stages in {source} are not contiguous from zero: " + f"found {sorted(stage_indices)}, expected {sorted(expected_indices)}." + ) + return inferred_depth + + +def get_model_state_keys(model: Any) -> set[str]: + """Return model-state names without constructing a mapping for torch modules.""" + if isinstance(model, nn.Module): + parameter_keys = {name for name, _param in model.named_parameters(remove_duplicate=False)} + buffer_keys = {name for name, _buffer in model.named_buffers(remove_duplicate=False)} + return parameter_keys | buffer_keys + + state_dict = getattr(model, "state_dict", None) + return set(state_dict()) if callable(state_dict) else set() + + +def validate_rq_checkpoint_depth( + model: nn.Module, + checkpoint_keys: Iterable[str], + *, + checkpoint_path: str, + model_keys: Iterable[str] | None = None, +) -> None: + """Reject RQ checkpoints whose persisted depth disagrees with the instantiated model.""" + resolved_model_keys = get_model_state_keys(model) if model_keys is None else model_keys + model_depth = infer_rq_codebook_depth(resolved_model_keys, source="the instantiated model") + checkpoint_depth = infer_rq_codebook_depth( + checkpoint_keys, + source=f"checkpoint {checkpoint_path}", + ) + if model_depth is None or checkpoint_depth is None or model_depth == checkpoint_depth: + return + + raise ValueError( + f"Residual-quantizer depth mismatch for checkpoint {checkpoint_path}: " + f"the instantiated model has {model_depth} stage(s), but the checkpoint contains {checkpoint_depth}. " + "Historical UniAE configs may record quantizer_num_codebooks=1 even though their checkpoints contain " + "four stages; migrate the config to the checkpoint depth before loading." + ) + + +def _build_s3_backend_args(checkpoint_path: str, s3_credential: str | None) -> dict[str, str] | None: + """Build easy_io backend args for an S3 checkpoint when explicit credentials are provided.""" + if not checkpoint_path.startswith("s3://") or s3_credential is None: + return None + return { + "backend": "s3", + "s3_credential_path": s3_credential, + } + + +def _safe_checkpoint_globals() -> list[Any]: + """Return trusted non-tensor classes present in tokenizer training checkpoints.""" + return [ + Any, + AnyNode, + ContainerMetadata, + DictConfig, + ListConfig, + Metadata, + defaultdict, + dict, + int, + list, + ] + + +def load_torch_checkpoint( + checkpoint_path: str | Path, + *, + map_location: TorchMapLocation = "cpu", + s3_credential: str | None = None, + backend_args: dict[str, Any] | None = None, + backend_key: str | None = None, +) -> Any: + """Load a torch checkpoint through easy_io with safe tensor-only unpickling.""" + checkpoint_path_str = str(checkpoint_path) + resolved_backend_args = ( + backend_args if backend_args is not None else _build_s3_backend_args(checkpoint_path_str, s3_credential) + ) + with torch.serialization.safe_globals(_safe_checkpoint_globals()): + return easy_io.load( + checkpoint_path_str, + backend_args=resolved_backend_args, + backend_key=backend_key, + map_location=map_location, + weights_only=True, + ) + + +def load_torch_checkpoint_from_bytes( + checkpoint_bytes: bytes, + *, + map_location: TorchMapLocation = "cpu", +) -> Any: + """Load a torch checkpoint from bytes with safe tensor-only unpickling.""" + with io.BytesIO(checkpoint_bytes) as checkpoint_buffer: + with torch.serialization.safe_globals(_safe_checkpoint_globals()): + return torch.load(checkpoint_buffer, map_location=map_location, weights_only=True) + + +def load_torch_checkpoint_from_easy_io_backend( + backend: Any, + checkpoint_path: str, + *, + map_location: TorchMapLocation = "cpu", +) -> Any: + """Load one checkpoint through an already-configured easy_io backend.""" + return load_torch_checkpoint_from_bytes( + backend.get(filepath=checkpoint_path), + map_location=map_location, + ) + + +def normalize_dcp_model_checkpoint_path(checkpoint_path: str | Path) -> str: + """Map a trainer iteration directory or URI to its DCP model component.""" + checkpoint_path_str = str(checkpoint_path).rstrip("/") + model_component_path = os.path.join(checkpoint_path_str, "model") + if os.path.isdir(model_component_path): + return model_component_path + if checkpoint_path_str.startswith("s3://") and re.search(r"/iter_\d{9}$", checkpoint_path_str): + return model_component_path + return checkpoint_path_str + + +def _cpu_state_placeholder(value: Any) -> Any: + """Create a CPU DCP target without duplicating source-device storage.""" + if isinstance(value, torch.Tensor): + return torch.empty_like(value, device="cpu") # [...] + return value + + +def _add_dcp_ema_template_tensors( + model: nn.Module, + model_state: dict[str, Any], + metadata_keys: set[str], +) -> None: + """Request EMA buffers from DCP when checkpoint metadata says they exist.""" + from cosmos_framework.utils.ema import get_buffer_name + + for name, param in model.named_parameters(): + ema_key = f"ema.{get_buffer_name(name)}" + if ema_key not in metadata_keys or ema_key in model_state: + continue + model_state[ema_key] = torch.empty_like(param.detach(), dtype=torch.float32, device="cpu") # [...] + + +def load_dcp_model_checkpoint( + model: nn.Module, + checkpoint_path: str | Path, + *, + include_ema: bool = False, + allow_partial: bool = False, + checkpoint_key_prefix: str = "", + s3_credential: str | None = None, +) -> dict[str, Any]: + """Load model state from a local or S3 torch.distributed.checkpoint component.""" + from torch.distributed.checkpoint import load + from torch.distributed.checkpoint.filesystem import FileSystemReader + + checkpoint_path_str = normalize_dcp_model_checkpoint_path(checkpoint_path) + + def _build_storage_reader() -> Any: + if checkpoint_path_str.startswith("s3://"): + if not s3_credential: + raise ValueError("s3_credential is required to load a remote DCP checkpoint.") + from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader + + return S3StorageReader( + credential_path=s3_credential, + path=checkpoint_path_str, + ) + return FileSystemReader(checkpoint_path_str) + + storage_reader = _build_storage_reader() + state_dict_metadata = storage_reader.read_metadata().state_dict_metadata + metadata_keys = set(state_dict_metadata) + if include_ema and checkpoint_key_prefix: + raise ValueError("checkpoint_key_prefix cannot be combined with include_ema.") + + current_model_state = model.state_dict() + checkpoint_key_by_model_key = { + model_key: f"{checkpoint_key_prefix}{model_key}" for model_key in current_model_state + } + expected_checkpoint_keys = set(checkpoint_key_by_model_key.values()) + wrapped_matches = sum(f"model.{key}" in metadata_keys for key in expected_checkpoint_keys) + raw_matches = sum(key in metadata_keys for key in expected_checkpoint_keys) + if wrapped_matches == 0 and raw_matches == 0: + raise ValueError( + f"DCP checkpoint {checkpoint_path_str} does not contain tensors matching the requested model " + f"(checkpoint sample: {sorted(metadata_keys)[:5]}; model sample: {sorted(expected_checkpoint_keys)[:5]})." + ) + + metadata_prefix = "model." if wrapped_matches >= raw_matches else "" + component_metadata = { + key.removeprefix(metadata_prefix): value + for key, value in state_dict_metadata.items() + if not metadata_prefix or key.startswith(metadata_prefix) + } + validate_rq_checkpoint_depth( + model, + component_metadata, + checkpoint_path=checkpoint_path_str, + model_keys=current_model_state, + ) + checkpoint_model_state: dict[str, Any] = {} + loaded_model_keys: set[str] = set() + missing_model_keys: set[str] = set() + shape_mismatches: list[tuple[str, tuple[int, ...], tuple[int, ...]]] = [] + for model_key, current_value in current_model_state.items(): + checkpoint_key = checkpoint_key_by_model_key[model_key] + checkpoint_metadata = component_metadata.get(checkpoint_key) + if checkpoint_metadata is None: + missing_model_keys.add(model_key) + continue + checkpoint_shape = getattr(checkpoint_metadata, "size", None) + if isinstance(current_value, torch.Tensor) and checkpoint_shape is not None: + current_shape = tuple(current_value.shape) + saved_shape = tuple(checkpoint_shape) + if current_shape != saved_shape: + shape_mismatches.append((model_key, saved_shape, current_shape)) + continue + checkpoint_model_state[checkpoint_key] = _cpu_state_placeholder(current_value) + loaded_model_keys.add(model_key) + + if shape_mismatches: + raise ValueError( + f"DCP checkpoint {checkpoint_path_str} has {len(shape_mismatches)} model shape mismatches " + f"(sample: {shape_mismatches[:5]})." + ) + if missing_model_keys and not allow_partial: + raise ValueError( + f"DCP checkpoint {checkpoint_path_str} is missing {len(missing_model_keys)} model keys " + f"(sample: {sorted(missing_model_keys)[:10]})." + ) + + if include_ema: + _add_dcp_ema_template_tensors( + model, + checkpoint_model_state, + set(component_metadata), + ) + requested_checkpoint_keys = set(checkpoint_model_state) + unexpected_checkpoint_keys = set(component_metadata) - requested_checkpoint_keys + checkpoint_state = {"model": checkpoint_model_state} if metadata_prefix else checkpoint_model_state + load(checkpoint_state, storage_reader=storage_reader, no_dist=True) + loaded_model_state = { + model_key: checkpoint_model_state[checkpoint_key_by_model_key[model_key]] for model_key in loaded_model_keys + } + loaded_model_state.update( + { + key: value + for key, value in checkpoint_model_state.items() + if key.startswith("ema.") and key not in loaded_model_state + } + ) + return { + "model": loaded_model_state, + DCP_MODEL_LOAD_INFO_KEY: DCPModelLoadInfo( + loaded_model_keys=frozenset(loaded_model_keys), + missing_model_keys=frozenset(missing_model_keys), + unexpected_checkpoint_keys=frozenset(unexpected_checkpoint_keys), + ), + } diff --git a/cosmos_framework/model/tokenizer/evaluation/metric.py b/cosmos_framework/model/tokenizer/evaluation/metric.py index 955b6582..6659dde6 100644 --- a/cosmos_framework/model/tokenizer/evaluation/metric.py +++ b/cosmos_framework/model/tokenizer/evaluation/metric.py @@ -18,6 +18,8 @@ import torch from loguru import logger as logging +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs + def compute_psnr( original: torch.Tensor, @@ -157,7 +159,7 @@ def compute_codebook_usage( all_indices = [] for i, g in enumerate(gathered): all_indices.append(g[: sizes[i].item()]) - flat_indices = torch.cat(all_indices) + flat_indices = cat_with_bounded_inputs(all_indices, dim=0) # Compute code histogram histogram = torch.bincount(flat_indices, minlength=num_codes).float() @@ -318,34 +320,70 @@ def reset(self) -> None: if self._metric is not None: self._metric.reset() + @staticmethod + def _flatten_video_batch(images: torch.Tensor, layout: str = "auto") -> torch.Tensor: + """Flatten image/video tensors into a 4D image batch. + + Args: + images: Tensor shaped [B, C, H, W], [B, C, T, H, W], or [B, T, C, H, W]. + layout: Layout for 5D tensors. Use "bcthw", "btchw", or "auto". + + Returns: + Tensor shaped [B*, C, H, W]. + """ + if images.ndim == 4: + return images + if images.ndim != 5: + raise ValueError(f"FIDComputer.update expected a 4D or 5D tensor, got shape {tuple(images.shape)}") + + normalized_layout = layout.lower() + if normalized_layout == "auto": + channels_at_dim1 = images.shape[1] in (1, 3) + channels_at_dim2 = images.shape[2] in (1, 3) + if channels_at_dim1 == channels_at_dim2: + raise ValueError( + "Ambiguous 5D FID image layout. Pass layout='bcthw' for [B, C, T, H, W] " + "or layout='btchw' for [B, T, C, H, W]." + ) + normalized_layout = "bcthw" if channels_at_dim1 else "btchw" + + if normalized_layout == "bcthw": + images = images.permute(0, 2, 1, 3, 4).contiguous() # [B, T, C, H, W] + elif normalized_layout == "btchw": + images = images.contiguous() # [B, T, C, H, W] + else: + raise ValueError(f"Unsupported FID video layout '{layout}'. Expected 'auto', 'bcthw', or 'btchw'.") + + return images.reshape(-1, *images.shape[2:]) # [B*T, C, H, W] + @torch.no_grad() def update( self, images: torch.Tensor, real: bool = True, + layout: str = "auto", ) -> None: """Update with a batch of images. Separate calls for real and fake images. Args: - images: Images in [0, 1] range, shape (B, C, H, W). + images: Images in [0, 1] range, shaped [B, C, H, W], [B, C, T, H, W], or [B, T, C, H, W]. real: Whether these are real images (True) or fake/reconstructed (False). + layout: Layout for 5D tensors. Use "bcthw", "btchw", or "auto". """ if not self._ensure_initialized(): return - # Handle video tensors (B, C, T, H, W) -> (B*T, C, H, W) - if images.ndim == 5: - images = images.reshape(-1, *images.shape[-3:]) + images = self._flatten_video_batch(images, layout=layout) # [B*, C, H, W] if self._normalize and images.dtype == torch.uint8: - images = images.float() / 255.0 + images = images.float() / 255.0 # [B*, C, H, W] if self._normalize: - images = images.float().to(self.device) + images = images.float().to(self.device) # [B*, C, H, W] else: - images = images.to(self.device) + images = images.to(self.device) # [B*, C, H, W] # Update with autocast disabled for numerical stability with self._autocast_context(): diff --git a/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py b/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py index ab5afa89..732c20eb 100644 --- a/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py +++ b/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py @@ -398,95 +398,6 @@ def calculate_psnr( return psnr.mean() -class Rank0FIDMetric(nn.Module): - """FID metric that runs only on rank 0 to avoid distributed sync issues. - - Uses torchmetrics FrechetInceptionDistance internally but only computes - on rank 0's data to avoid NCCL collective operation mismatches caused by - torchmetrics/torch-fidelity's internal distributed synchronization. - - Note: FID is computed only on rank 0's portion of the data (1/world_size), - which may be less representative than full dataset FID, but avoids - distributed synchronization issues. - - Usage: - fid = Rank0FIDMetric(rank=rank).to(device) - - # During evaluation loop (only rank 0 updates) - for batch in dataloader: - fid.update(real_images, fake_images) - - # Compute FID (only rank 0 has valid result) - if rank == 0: - fid_value = fid.compute() - """ - - def __init__(self, rank: int = 0, feature_dim: int = 2048) -> None: - super().__init__() - self.rank = rank - self.feature_dim = feature_dim - self._fid_metric = None - - # Only initialize FID metric on rank 0 - if self.rank == 0: - try: - from torchmetrics.image.fid import FrechetInceptionDistance - - # normalize=True means input is [0, 1] float, not uint8 - self._fid_metric = FrechetInceptionDistance( - feature=feature_dim, - normalize=True, - sync_on_compute=False, - dist_sync_on_step=False, - ) - except ImportError: - pass - - @torch.no_grad() - def update(self, real_images: torch.Tensor, fake_images: torch.Tensor) -> None: - """Update FID statistics with a batch of real and fake images. - - Only updates on rank 0. - - Args: - real_images: Real images in [0, 1] range, shape (B, C, H, W) or (B, T, C, H, W) - fake_images: Fake/reconstructed images in [0, 1] range - """ - if self.rank != 0 or self._fid_metric is None: - return - - # Handle video format by flattening batch and time dimensions - if real_images.dim() == 5: # (B, T, C, H, W) - real_images = real_images.reshape(-1, *real_images.shape[2:]) - fake_images = fake_images.reshape(-1, *fake_images.shape[2:]) - - # Move metric to same device as images - device = real_images.device - self._fid_metric = self._fid_metric.to(device) - - # torchmetrics FID update - self._fid_metric.update(real_images, real=True) - self._fid_metric.update(fake_images, real=False) - - def compute(self) -> torch.Tensor: - """Compute FID from accumulated statistics. - - Only valid on rank 0. - - Returns: - FID value as a scalar tensor (inf if not rank 0 or metric unavailable) - """ - if self.rank != 0 or self._fid_metric is None: - return torch.tensor(float("inf")) - - return self._fid_metric.compute() - - def reset(self) -> None: - """Reset accumulated statistics.""" - if self._fid_metric is not None: - self._fid_metric.reset() - - __all__ = [ @@ -494,6 +405,5 @@ def reset(self) -> None: "PSNRMetric", "SSIMMetric", "LPIPSMetric", - "Rank0FIDMetric", "calculate_psnr", ] diff --git a/cosmos_framework/model/tokenizer/models/architecture.py b/cosmos_framework/model/tokenizer/models/architecture.py new file mode 100644 index 00000000..f70542b6 --- /dev/null +++ b/cosmos_framework/model/tokenizer/models/architecture.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared UniAE architecture specifications used by training and inference.""" + +from __future__ import annotations + +import copy +from typing import Any + +_SIGLIP2_SO400M_COMMON_ARCH: dict[str, Any] = { + "patch_size": (4, 16, 16), + "in_channels": 3072, + "out_channels": 3072, + "encoder_model_channels": 1152, + "encoder_num_blocks": 27, + "encoder_num_heads": 16, + "encoder_mlp_channels": 4304, + "encoder_pe_mode": "joint", + "encoder_qk_rms_norm": False, + "encoder_use_bias": True, + "encoder_use_rms_norm": False, + "decoder_model_channels": 1152, + "decoder_num_blocks": 27, + "decoder_num_heads": 16, + "decoder_mlp_channels": 4304, + "decoder_pe_mode": "joint", + "decoder_qk_rms_norm": True, + "decoder_use_bias": False, + "decoder_use_rms_norm": True, + "use_decoder": True, + "quantizer_type": "rq", + "quantizer_codebook_size": 65536, + # Production checkpoints use four residual-quantization stages. + "quantizer_num_codebooks": 4, + "quantizer_chunk_size": 1, + "use_vf_loss": False, + "freeze_encoder": False, + "pretrained_model_name": "google/siglip2-so400m-patch16-naflex", + "concat_latent": None, + "random_num_sample_frames_batch_sizes": [8, 12, 16, 20, 24], + "inference_num_sample_frames_batch_size": 16, + "inference_num_sample_frames_stride": 16, + "inference_kv_cache_size": 0, +} + + +def get_siglip2_so400m_common_arch() -> dict[str, Any]: + """Return an isolated copy of the shared SigLIP2-SO400M UniAE architecture.""" + return copy.deepcopy(_SIGLIP2_SO400M_COMMON_ARCH) + + +__all__ = ["get_siglip2_so400m_common_arch"] diff --git a/cosmos_framework/model/tokenizer/models/dense_runtime.py b/cosmos_framework/model/tokenizer/models/dense_runtime.py index 93820284..57c08b2d 100644 --- a/cosmos_framework/model/tokenizer/models/dense_runtime.py +++ b/cosmos_framework/model/tokenizer/models/dense_runtime.py @@ -22,6 +22,7 @@ ) from cosmos_framework.model.tokenizer.models.modules.transformer.blocks import LearnedPositionEmbedder from cosmos_framework.model.tokenizer.models.sparse_autoencoder import AutoencoderKL, SparseTransformerBase +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs @dataclass(frozen=True) @@ -50,55 +51,6 @@ class DenseGridMetadata: DenseGridMetadataKey = tuple[str, int, int, int, int, str, str] -class DenseDiagonalGaussianDistribution: - """Diagonal Gaussian posterior for dense channels-last latent tensors.""" - - def __init__(self, parameters: torch.Tensor, deterministic: bool = False) -> None: - """Initialize the dense posterior from `[mean, logvar]` moments.""" - if parameters.ndim not in (4, 5): - raise ValueError( - "DenseDiagonalGaussianDistribution expects 4D/5D channels-last moments, " - f"got shape {tuple(parameters.shape)}." - ) - self.original_dtype = parameters.dtype - self.parameters = parameters.to(torch.float32) - self.mean, self.logvar = torch.chunk(self.parameters, 2, dim=-1) - self.deterministic = deterministic - self.std = torch.exp(0.5 * self.logvar) - self.var = torch.exp(self.logvar) - - if self.deterministic: - self.var = self.std = torch.zeros_like( - self.mean, - device=self.parameters.device, - dtype=self.parameters.dtype, - ) - - def sample(self) -> torch.Tensor: - """Sample a dense channels-last latent tensor.""" - sample = torch.randn_like(self.mean) - return (self.mean + self.std * sample).to(self.original_dtype) - - def kl(self, other: "DenseDiagonalGaussianDistribution" | None = None) -> torch.Tensor: - """Compute KL divergence per latent token, matching sparse scaling.""" - reduce_dims = (-1,) - if self.deterministic: - num_tokens = math.prod(self.mean.shape[:-1]) - return torch.zeros(num_tokens, device=self.parameters.device, dtype=self.parameters.dtype) - if other is None: - kl = 0.5 * torch.sum(torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, dim=reduce_dims) - else: - kl = 0.5 * torch.sum( - torch.pow(self.mean - other.mean, 2) / other.var - + self.var / other.var - - 1.0 - - self.logvar - + other.logvar, - dim=reduce_dims, - ) - return kl.reshape(-1) - - class DenseAutoencoderRuntime(nn.Module): """Dense frozen-runtime wrapper around an existing sparse autoencoder. @@ -406,7 +358,7 @@ def _encode_padded_chunks(padded_chunks: list[torch.Tensor]) -> list[torch.Tenso encoded = self._encode_video_chunk(padded_chunks[0], pad_to=pad_to) return [_trim_boundary_latents(encoded)] - batched_video = torch.cat(padded_chunks, dim=0) # [B*G,t_pad,H,W,3] + batched_video = cat_with_bounded_inputs(padded_chunks, dim=0) # [B*G,t_pad,H,W,3] encoded = self._encode_video_chunk(batched_video, pad_to=pad_to) # [B*G,T_lat,Hp,Wp,2C] per_video_batch = padded_chunks[0].shape[0] return list(_trim_boundary_latents(encoded).split(per_video_batch, dim=0)) @@ -449,7 +401,7 @@ def _encode_padded_chunks(padded_chunks: list[torch.Tensor]) -> list[torch.Tenso if pending_full_chunks: encoded_chunks.extend(_encode_padded_chunks(pending_full_chunks)) - return torch.cat(encoded_chunks, dim=1) + return cat_with_bounded_inputs(encoded_chunks, dim=1) def decode( self, @@ -511,7 +463,7 @@ def decode( if trim_pixel and not is_image: decoded_chunk = decoded_chunk[:, pad_frames:-pad_frames] decoded_chunks.append(decoded_chunk) - return torch.cat(decoded_chunks, dim=1) + return cat_with_bounded_inputs(decoded_chunks, dim=1) def _metadata_cache_key( self, diff --git a/cosmos_framework/model/tokenizer/models/modules/__init__.py b/cosmos_framework/model/tokenizer/models/modules/__init__.py index aae53fc5..ec173616 100644 --- a/cosmos_framework/model/tokenizer/models/modules/__init__.py +++ b/cosmos_framework/model/tokenizer/models/modules/__init__.py @@ -14,7 +14,6 @@ import importlib import os -from typing import Literal from loguru import logger as logging @@ -51,31 +50,6 @@ def _init_from_env() -> None: _init_from_env() -def set_backend(backend: Literal["pytorch", "spconv", "torchsparse"]) -> None: - """Set the sparse tensor backend. - - Args: - backend: Backend to use. - - "pytorch": Pure PyTorch implementation (default, no external dependencies) - - "spconv": Uses spconv.pytorch.SparseConvTensor - - "torchsparse": Uses torchsparse.SparseTensor - """ - global BACKEND - if backend not in _VALID_BACKENDS: - raise ValueError(f"Invalid backend: {backend}. Must be one of {_VALID_BACKENDS}") - BACKEND = backend - - -def set_debug(debug: bool) -> None: - """Enable or disable debug mode. - - Args: - debug: Whether to enable debug mode. - """ - global DEBUG - DEBUG = debug - - # Lazy loading attribute mapping _ATTRIBUTES = { # SparseTensor and operations diff --git a/cosmos_framework/model/tokenizer/models/modules/attention/full_attn.py b/cosmos_framework/model/tokenizer/models/modules/attention/full_attn.py index 78b077f6..06e59eb4 100644 --- a/cosmos_framework/model/tokenizer/models/modules/attention/full_attn.py +++ b/cosmos_framework/model/tokenizer/models/modules/attention/full_attn.py @@ -17,6 +17,7 @@ from cosmos_framework.model.attention.frontend import attention as i4_attention from cosmos_framework.model.attention.varlen import generate_varlen_parameters +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs if TYPE_CHECKING: from cosmos_framework.model.tokenizer.models.modules.sparse_tensor import SparseTensor @@ -191,12 +192,12 @@ def _pack_sparse_temporal_causal_qkv( return empty_q, empty_k, empty_v, [], [], empty_idx return ( - torch.cat(q_feat_chunks, dim=0), - torch.cat(k_feat_chunks, dim=0), - torch.cat(v_feat_chunks, dim=0), + cat_with_bounded_inputs(q_feat_chunks, dim=0), + cat_with_bounded_inputs(k_feat_chunks, dim=0), + cat_with_bounded_inputs(v_feat_chunks, dim=0), q_seqlen, kv_seqlen, - torch.cat(q_index_chunks, dim=0), + cat_with_bounded_inputs(q_index_chunks, dim=0), ) diff --git a/cosmos_framework/model/tokenizer/models/modules/attention/modules.py b/cosmos_framework/model/tokenizer/models/modules/attention/modules.py index fd18129f..4ebcbc2e 100644 --- a/cosmos_framework/model/tokenizer/models/modules/attention/modules.py +++ b/cosmos_framework/model/tokenizer/models/modules/attention/modules.py @@ -24,6 +24,7 @@ sparse_scaled_dot_product_attention, tensor_varlen_scaled_dot_product_attention, ) +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs if TYPE_CHECKING: from cosmos_framework.model.tokenizer.models.modules.sparse_tensor import SparseTensor @@ -89,7 +90,7 @@ def _manual_segmented_scaled_dot_product_attention( ) if out_chunks: - output = torch.cat(out_chunks, dim=0) # [Tq,H,D] + output = cat_with_bounded_inputs(out_chunks, dim=0) # [Tq,H,D] else: output = q.new_empty((0, q.shape[1], v.shape[-1])) # [0,H,D] captured_attn = attn_chunks if store_attn else None @@ -927,7 +928,7 @@ def _concat_temporal_freqs( if not combined_parts: return current_freqs_cis.new_empty((0,) + current_freqs_cis.shape[1:]) - return torch.cat(combined_parts, dim=0) + return cat_with_bounded_inputs(combined_parts, dim=0) combined_parts = [] current_batch_size = current.shape[0] @@ -966,7 +967,7 @@ def _concat_temporal_freqs( if not combined_parts: return current_freqs_cis.new_empty((0,) + current_freqs_cis.shape[1:]) - return torch.cat(combined_parts, dim=0) + return cat_with_bounded_inputs(combined_parts, dim=0) def _concat_temporal_sparse(self, cached: "SparseTensor", current: "SparseTensor") -> "SparseTensor": """Concatenate cached and current sparse tensors along temporal dimension. @@ -1044,9 +1045,9 @@ def _update_kv_cache( offset += batch_len if cached_coords_parts: - selected_coords = torch.cat(cached_coords_parts, dim=0) - cached_k_feats = torch.cat(cached_k_parts, dim=0) - cached_v_feats = torch.cat(cached_v_parts, dim=0) + selected_coords = cat_with_bounded_inputs(cached_coords_parts, dim=0) + cached_k_feats = cat_with_bounded_inputs(cached_k_parts, dim=0) + cached_v_feats = cat_with_bounded_inputs(cached_v_parts, dim=0) else: selected_coords = torch.empty( (0, k.coords.shape[1]), @@ -1087,7 +1088,7 @@ def _update_kv_cache( kv_cache["cached_v"] = cached_v if k_freqs_cis is not None: if cached_freq_parts: - cached_k_freqs_cis = torch.cat(cached_freq_parts, dim=0) + cached_k_freqs_cis = cat_with_bounded_inputs(cached_freq_parts, dim=0) else: cached_k_freqs_cis = k_freqs_cis.new_empty((0,) + k_freqs_cis.shape[1:]) if kv_cache_detach: diff --git a/cosmos_framework/model/tokenizer/models/modules/sparse_tensor.py b/cosmos_framework/model/tokenizer/models/modules/sparse_tensor.py index 4198dd30..268bb695 100644 --- a/cosmos_framework/model/tokenizer/models/modules/sparse_tensor.py +++ b/cosmos_framework/model/tokenizer/models/modules/sparse_tensor.py @@ -17,6 +17,7 @@ import torch from cosmos_framework.model.tokenizer.models.modules import BACKEND, DEBUG +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs SparseTensorData = None # Lazy import @@ -472,8 +473,8 @@ def _select_contiguous_batch_ranges( ) return SparseTensor( - feats=torch.cat(selected_feats_parts, dim=0), - coords=torch.cat(selected_coords_parts, dim=0), + feats=cat_with_bounded_inputs(selected_feats_parts, dim=0), + coords=cat_with_bounded_inputs(selected_coords_parts, dim=0), shape=torch.Size([self.shape[0]] + list(self.feats.shape[1:])), layout=new_layout, has_special_tokens=has_special_tokens, @@ -906,8 +907,8 @@ def __getitem__(self, idx): # Concatenate all coordinates and features (including empty ones) if len(coords) > 0: - coords = torch.cat(coords, dim=0).contiguous() - feats = torch.cat(feats, dim=0).contiguous() + coords = cat_with_bounded_inputs(coords, dim=0).contiguous() + feats = cat_with_bounded_inputs(feats, dim=0).contiguous() else: # Handle case where idx is empty coords = torch.empty( @@ -1431,8 +1432,8 @@ def concat_temporal_at_batch_start(self, other: "SparseTensor") -> "SparseTensor if not has_cached_tokens: return self - combined_coords = torch.cat(combined_coords_parts, dim=0) - combined_feats = torch.cat(combined_feats_parts, dim=0) + combined_coords = cat_with_bounded_inputs(combined_coords_parts, dim=0) + combined_feats = cat_with_bounded_inputs(combined_feats_parts, dim=0) new_shape = torch.Size([current_batch_size] + list(self.feats.shape[1:])) result = SparseTensor( feats=combined_feats, @@ -1499,8 +1500,8 @@ def concat_temporal_at_batch_start(self, other: "SparseTensor") -> "SparseTensor if not has_cached_tokens: return self - combined_coords = torch.cat(combined_coords_parts, dim=0) - combined_feats = torch.cat(combined_feats_parts, dim=0) + combined_coords = cat_with_bounded_inputs(combined_coords_parts, dim=0) + combined_feats = cat_with_bounded_inputs(combined_feats_parts, dim=0) new_shape = torch.Size([current_batch_size] + list(self.feats.shape[1:])) result = SparseTensor( feats=combined_feats, @@ -1631,14 +1632,14 @@ def sparse_cat(inputs: list[SparseTensor], dim: int = 0) -> SparseTensor: coords.append(input.coords.clone()) coords[-1][:, 0] += start start += input.shape[0] - coords = torch.cat(coords, dim=0) - feats = torch.cat([input.feats for input in inputs], dim=0) + coords = cat_with_bounded_inputs(coords, dim=0) + feats = cat_with_bounded_inputs([input.feats for input in inputs], dim=0) output = SparseTensor( coords=coords, feats=feats, ) else: - feats = torch.cat([input.feats for input in inputs], dim=dim) + feats = cat_with_bounded_inputs([input.feats for input in inputs], dim=dim) output = inputs[0].replace(feats) return output @@ -1735,9 +1736,6 @@ def reconstruct_from_temporal_slices( all_coords.append(restored_slice.coords) all_feats.append(restored_slice.feats) - combined_coords = torch.cat(all_coords, dim=0) - combined_feats = torch.cat(all_feats, dim=0) - if target_coords is not None: fast_path = _reconstruct_in_target_batch_order( restored_slices, @@ -1746,6 +1744,8 @@ def reconstruct_from_temporal_slices( ) if fast_path is not None: return fast_path + combined_coords = cat_with_bounded_inputs(all_coords, dim=0) # [N_available, 5] + combined_feats = cat_with_bounded_inputs(all_feats, dim=0) # [N_available, *F] return _match_target_coordinates( combined_coords, combined_feats, @@ -1753,6 +1753,8 @@ def reconstruct_from_temporal_slices( has_special_tokens=has_special_tokens, ) else: + combined_coords = cat_with_bounded_inputs(all_coords, dim=0) # [N_available, 5] + combined_feats = cat_with_bounded_inputs(all_feats, dim=0) # [N_available, *F] sort_key = (combined_coords[:, 0].long() << 20) + combined_coords[:, 1].long() sorted_indices = torch.argsort(sort_key) @@ -1852,12 +1854,12 @@ def _reconstruct_in_target_batch_order( has_special_tokens=False, ) - combined_coords = torch.cat(combined_coords_parts, dim=0) + combined_coords = cat_with_bounded_inputs(combined_coords_parts, dim=0) # [N_target, 5] if combined_coords.shape != target_coords.shape or not torch.equal(combined_coords, target_coords): return None return SparseTensor( - feats=torch.cat(combined_feats_parts, dim=0), + feats=cat_with_bounded_inputs(combined_feats_parts, dim=0), # [N_target, *F] coords=target_coords, shape=torch.Size([batch_size] + list(restored_slices[0].feats.shape[1:])), layout=new_layout, diff --git a/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py b/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py index f69884b9..40b79fd0 100644 --- a/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py +++ b/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py @@ -58,6 +58,7 @@ sparse_to_batched_tensor, sparse_to_img_list, ) +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs, stack_with_bounded_inputs # ============================================================================= # Helper Functions @@ -146,7 +147,9 @@ def _sparse_tensor_to_dense_batch_tokens(sparse_tensor: SparseTensor) -> tuple[t dense_feats = sparse_tensor.feats.reshape(batch_size, seq_len, *feature_shape) return dense_feats, True - dense_feats = torch.stack([sparse_tensor.feats[batch_slice] for batch_slice in sparse_tensor.layout], dim=0) + dense_feats = stack_with_bounded_inputs( + [sparse_tensor.feats[batch_slice] for batch_slice in sparse_tensor.layout], dim=0 + ) # [B,S,*F] return dense_feats, False @@ -154,7 +157,7 @@ def _dense_batch_tokens_to_flat_features(dense_feats: torch.Tensor, used_reshape """Restore dense `[B, S, D]` features to the sparse flat token order.""" if used_reshape_fast_path: return dense_feats.reshape(-1, *dense_feats.shape[2:]) - return torch.cat(list(dense_feats.unbind(dim=0)), dim=0) + return cat_with_bounded_inputs(dense_feats.unbind(dim=0), dim=0) def _crop_temporal_slices_to_ownership( @@ -1208,6 +1211,7 @@ class AutoencoderKLConfig: text_decoder_model_name: str | None = None # e.g., "Qwen/Qwen3-0.6B" or local path text_decoder_family: str = "qwen3" text_decoder_gradient_checkpointing: bool = True + text_decoder_packed_attention_backend: Literal["sdpa", "natten"] = "sdpa" encoder_use_checkpoint: bool | None = None decoder_use_checkpoint: bool | None = None encoder_dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled" @@ -1291,6 +1295,7 @@ def __init__( text_decoder_model_name: str | None = None, text_decoder_family: str = "qwen3", text_decoder_gradient_checkpointing: bool = True, + text_decoder_packed_attention_backend: Literal["sdpa", "natten"] = "sdpa", encoder_use_checkpoint: bool | None = None, decoder_use_checkpoint: bool | None = None, encoder_dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled", @@ -1311,7 +1316,7 @@ def __init__( task_random_num_sample_frames_batch_sizes: dict[str, list[int]] | None = None, use_dual_latent: bool = False, use_checkpoint: bool = True, - ): + ) -> None: super().__init__() self.use_checkpoint = use_checkpoint self.encoder_use_checkpoint = use_checkpoint if encoder_use_checkpoint is None else encoder_use_checkpoint @@ -1325,6 +1330,7 @@ def __init__( self.use_post_text_decoder = use_post_text_decoder self.spatial_pool_size = spatial_pool_size self.text_decoder_family = text_decoder_family + self.text_decoder_packed_attention_backend: Literal["sdpa", "natten"] = text_decoder_packed_attention_backend self.decoder_temporal_mode = decoder_temporal_mode self.decoder_temporal_query_latent_steps = decoder_temporal_query_latent_steps self.decoder_temporal_cache_latent_steps = decoder_temporal_cache_latent_steps @@ -1435,7 +1441,7 @@ def __init__( elif self.quantizer_type == "rq": self.quantizer = RQBottleneck( latent_shape=(16, 16, latent_channels), - code_shape=(16, 16, 4), + code_shape=(16, 16, self.quantizer_num_codebooks), n_embed=self.quantizer_codebook_size, decay=0.99, shared_codebook=True, @@ -1534,6 +1540,7 @@ def __init__( image_hidden_size=encoder_model_channels, spatial_pool_size=spatial_pool_size, gradient_checkpointing=text_decoder_gradient_checkpointing, + packed_attention_backend=text_decoder_packed_attention_backend, family_spec=get_text_decoder_family_spec( family=text_decoder_family, model_name=text_decoder_model_name, @@ -1548,7 +1555,7 @@ def decode_text( image_feats: "SparseTensor", image_patch_indices: torch.Tensor, segment_ids: torch.Tensor | None = None, - ): + ) -> tuple[Any, int]: """Decode text from image features using the configured text decoder. Passes encoder output (x_no_proj) through spatial merger + causal LM. @@ -1921,7 +1928,7 @@ def decode( logging.warning(f"Decoded shapes are not the same: {[x.shape for x in decoded_list]}") decoded = decoded_list else: - decoded = torch.stack(decoded_list, dim=0) + decoded = stack_with_bounded_inputs(decoded_list, dim=0) # [B,T,C,H,W] decoded = rearrange(decoded, "b t c h w -> b t h w c") if decoded.shape[1] == 1: decoded = decoded.squeeze(1) diff --git a/cosmos_framework/model/tokenizer/models/text_decoder.py b/cosmos_framework/model/tokenizer/models/text_decoder.py index b9f296c1..3cddc135 100644 --- a/cosmos_framework/model/tokenizer/models/text_decoder.py +++ b/cosmos_framework/model/tokenizer/models/text_decoder.py @@ -23,7 +23,6 @@ """ import os -import re from dataclasses import dataclass from types import SimpleNamespace from typing import Any @@ -33,11 +32,15 @@ import torch.nn.functional as F from loguru import logger as logging +from cosmos_framework.model.attention import attention as i4_attention +from cosmos_framework.model.attention.masks import CausalType +from cosmos_framework.model.attention.natten import NATTEN_SUPPORTED, get_natten_version from cosmos_framework.model.tokenizer.utils.hf import ( load_auto_tokenizer_from_cache, prepare_nemotron_tokenizer_snapshot, resolve_hf_snapshot_path, ) +from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs, stack_with_bounded_inputs from cosmos_framework.model.tokenizer.utils.vlm_prompt_format import densevl_add_vision_id_text QWEN3_PAD_TOKEN_ID = 151643 @@ -58,6 +61,9 @@ NEMOTRON_2B_THINK_START_TOKEN = "" NEMOTRON_2B_THINK_END_TOKEN = "" TEXT_DECODER_ATTN_IMPLEMENTATION_ENV = "TOKENIZER_TEXT_DECODER_ATTN_IMPLEMENTATION" +PACKED_ATTENTION_BACKEND_SDPA = "sdpa" +PACKED_ATTENTION_BACKEND_NATTEN = "natten" +PACKED_ATTENTION_BACKENDS = frozenset({PACKED_ATTENTION_BACKEND_SDPA, PACKED_ATTENTION_BACKEND_NATTEN}) VQA_THINKING_MODE_OFF = "off" VQA_THINKING_MODE_ON = "on" VQA_THINKING_MODE_RAW = "raw" @@ -88,6 +94,128 @@ def _append_vqa_reasoning_suffix(question: str, reasoning_suffix: str) -> str: return f"{question.rstrip()}{reasoning_suffix}" +def _rotate_half_for_nemotron(hidden_states: torch.Tensor) -> torch.Tensor: + """Rotate the two RoPE feature halves using Nemotron's convention.""" + midpoint = hidden_states.shape[-1] // 2 + first_half = hidden_states[..., :midpoint] # [B,H,T,D/2] + second_half = hidden_states[..., midpoint:] # [B,H,T,D/2] + rotated = torch.cat((-second_half, first_half), dim=-1) # [B,H,T,D] + return rotated + + +def _apply_nemotron_rotary_embeddings( + query_states: torch.Tensor, + key_states: torch.Tensor, + cosine: torch.Tensor, + sine: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply the remote Nemotron model's RoPE formulation to Q and K.""" + rotated_query = _rotate_half_for_nemotron(query_states) # [B,Hq,T,D] + rotated_key = _rotate_half_for_nemotron(key_states) # [B,Hkv,T,D] + query_states_ = query_states * cosine + rotated_query * sine # [B,Hq,T,D] + key_states_ = key_states * cosine + rotated_key * sine # [B,Hkv,T,D] + return query_states_, key_states_ + + +def _forward_nemotron_natten_attention( + attention_module: nn.Module, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + cumulative_seqlen: torch.Tensor, + max_seqlen: int, + rotary_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """Run one Nemotron self-attention block with NATTEN causal varlen attention.""" + batch_size, sequence_length, _ = hidden_states.shape + if batch_size != 1: + raise ValueError(f"NATTEN varlen requires packed batch size 1, got {batch_size}.") + + query_states = attention_module.q_proj(hidden_states) # [1,T,Hq*D] + query_states = query_states.view( + batch_size, + sequence_length, + attention_module.num_heads, + attention_module.head_dim, + ).transpose(1, 2) # [1,Hq,T,D] + key_states = attention_module.k_proj(hidden_states) # [1,T,Hkv*D] + key_states = key_states.view( + batch_size, + sequence_length, + attention_module.num_key_value_heads, + attention_module.head_dim, + ).transpose(1, 2) # [1,Hkv,T,D] + value_states = attention_module.v_proj(hidden_states) # [1,T,Hkv*D] + value_states = value_states.view( + batch_size, + sequence_length, + attention_module.num_key_value_heads, + attention_module.head_dim, + ).transpose(1, 2) # [1,Hkv,T,D] + + if rotary_embeddings is None: + normalized_position_ids = position_ids.reshape(batch_size, sequence_length).long() # [1,T] + cosine, sine = attention_module.rotary_emb( + value_states, + normalized_position_ids, + ) # each [1,1,T,D] + else: + cosine, sine = rotary_embeddings + query_states_, key_states_ = _apply_nemotron_rotary_embeddings( + query_states, + key_states, + cosine, + sine, + ) # [1,Hq,T,D], [1,Hkv,T,D] + + query_states_ = query_states_.transpose(1, 2).contiguous() # [1,T,Hq,D] + key_states_ = key_states_.transpose(1, 2).contiguous() # [1,T,Hkv,D] + value_states = value_states.transpose(1, 2).contiguous() # [1,T,Hkv,D] + attention_output = i4_attention( + query=query_states_, + key=key_states_, + value=value_states, + is_causal=True, + causal_type=CausalType.DontCare, + cumulative_seqlen_Q=cumulative_seqlen, + cumulative_seqlen_KV=cumulative_seqlen, + max_seqlen_Q=max_seqlen, + max_seqlen_KV=max_seqlen, + backend=PACKED_ATTENTION_BACKEND_NATTEN, + ) + if not isinstance(attention_output, torch.Tensor): + raise TypeError("NATTEN attention unexpectedly returned log-sum-exp output.") + attention_output = attention_output.reshape(batch_size, sequence_length, attention_module.hidden_size) # [1,T,D] + projected_output = attention_module.o_proj(attention_output) # [1,T,D] + return projected_output + + +def _forward_nemotron_natten_decoder_layer( + decoder_layer: nn.Module, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + cumulative_seqlen: torch.Tensor, + max_seqlen: int, + rotary_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """Run one remote-code Nemotron layer while replacing only self-attention.""" + residual = hidden_states # [1,T,D] + normalized_hidden_states = decoder_layer.input_layernorm(hidden_states) # [1,T,D] + attention_output = _forward_nemotron_natten_attention( + decoder_layer.self_attn, + normalized_hidden_states, + position_ids, + cumulative_seqlen, + max_seqlen, + rotary_embeddings, + ) # [1,T,D] + hidden_states = residual + attention_output # [1,T,D] + + residual = hidden_states # [1,T,D] + normalized_hidden_states = decoder_layer.post_attention_layernorm(hidden_states) # [1,T,D] + mlp_output = decoder_layer.mlp(normalized_hidden_states) # [1,T,D] + return residual + mlp_output # [1,T,D] + + @dataclass(frozen=True) class TextDecoderFamilySpec: """Static configuration for one supported text decoder family.""" @@ -179,22 +307,6 @@ def _repair_nemotron_rotary_buffers(module: nn.Module) -> int: return repaired_count -def _keep_only_first_image_placeholder( - text: str, - image_placeholder: str = "", -) -> str: - """Keep only the first image placeholder occurrence in a prompt string.""" - first_index = text.find(image_placeholder) - if first_index == -1: - return text - - before_first = text[:first_index] - after_first = text[first_index + len(image_placeholder) :] - escaped_placeholder = re.escape(image_placeholder) - after_first_cleaned = re.sub(escaped_placeholder + r"\s*", "", after_first) - return before_first + image_placeholder + after_first_cleaned - - def infer_text_decoder_family(model_name: str) -> str: """Infer the decoder family from a HuggingFace model name.""" model_name_lower = model_name.lower() @@ -437,8 +549,8 @@ def forward( new_layout.append(slice(current_offset, current_offset + num_merged)) current_offset += num_merged - merged_feats = torch.cat(merged_feats_list, dim=0) - merged_coords = torch.cat(merged_coords_list, dim=0) + merged_feats = cat_with_bounded_inputs(merged_feats_list, dim=0) + merged_coords = cat_with_bounded_inputs(merged_coords_list, dim=0) # MLP projection: merged_hidden_size -> out_hidden_size merged_feats = self.linear_fc2(self.act_fn(self.linear_fc1(merged_feats))) @@ -463,7 +575,8 @@ def __init__( dtype: torch.dtype | None = None, attn_implementation: str = "flash_attention_2", family_spec: TextDecoderFamilySpec | None = None, - ): + packed_attention_backend: str = PACKED_ATTENTION_BACKEND_SDPA, + ) -> None: super().__init__() if dtype is None: @@ -472,8 +585,28 @@ def __init__( from transformers import AutoModelForCausalLM self.spec = family_spec or get_text_decoder_family_spec(model_name=model_name) + if packed_attention_backend not in PACKED_ATTENTION_BACKENDS: + raise ValueError( + f"Unsupported packed_attention_backend={packed_attention_backend!r}; " + f"expected one of {sorted(PACKED_ATTENTION_BACKENDS)}." + ) + if packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN and self.spec.family != NEMOTRON_2B_SPEC.family: + raise ValueError("NATTEN packed attention currently supports only the Nemotron 2B text decoder.") + if ( + packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN + and torch.cuda.is_available() + and not NATTEN_SUPPORTED + ): + raise RuntimeError("NATTEN packed attention was requested but NATTEN is unavailable or unsupported.") + self.packed_attention_backend: str = packed_attention_backend self._model_name: str = model_name or self.spec.default_model_name resolved_attn_implementation = _resolve_attn_implementation(attn_implementation) + if packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN: + # Packed training bypasses the remote decoder attention. Keep all + # non-packed calls on the explicit dependency-free fallback. + resolved_attn_implementation = "eager" + if NATTEN_SUPPORTED: + logging.info(f"Using NATTEN {get_natten_version()} for packed text decoder attention") logging.info(f"Loading text decoder: {self._model_name}") logging.info(f"Using text decoder attention backend: {resolved_attn_implementation}") hf_cache_dir = os.environ.get("HF_HOME") @@ -592,7 +725,8 @@ def __init__( logging.info( f"TextDecoderWrapper ready: family={self.spec.family}, hidden_size={hidden_size}, " f"layers={len(self.text_decoder.model.layers)}, " - f"gradient_checkpointing={gradient_checkpointing_enabled}, use_cache=False" + f"gradient_checkpointing={gradient_checkpointing_enabled}, " + f"packed_attention_backend={self.packed_attention_backend}, use_cache=False" ) def _get_eos_token_ids(self) -> tuple[int, ...]: @@ -636,6 +770,77 @@ def _embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: """Embed token IDs using the model's input embedding module.""" return self.text_decoder.get_input_embeddings()(input_ids) + def _forward_nemotron_natten_from_embeddings( + self, + text_embeds: torch.Tensor, + position_ids: torch.Tensor, + cumulative_seqlen: torch.Tensor, + max_seqlen: int, + ) -> SimpleNamespace: + """Run packed Nemotron layers with NATTEN varlen self-attention.""" + if self.spec.family != NEMOTRON_2B_SPEC.family: + raise ValueError("NATTEN packed attention currently supports only the Nemotron 2B text decoder.") + if text_embeds.shape[0] != 1: + raise ValueError(f"NATTEN varlen requires packed batch size 1, got {text_embeds.shape[0]}.") + if cumulative_seqlen.dtype != torch.int32: + raise ValueError(f"NATTEN cumulative sequence lengths must be int32, got {cumulative_seqlen.dtype}.") + + hidden_states = text_embeds # [1,T,D] + decoder_layers = self.text_decoder.model.layers + if not decoder_layers: + raise ValueError("NATTEN packed attention requires at least one Nemotron decoder layer.") + first_attention = decoder_layers[0].self_attn + normalized_position_ids = position_ids.reshape(hidden_states.shape[0], hidden_states.shape[1]).long() # [1,T] + rotary_reference = hidden_states.new_empty( + ( + hidden_states.shape[0], + first_attention.num_key_value_heads, + 0, + first_attention.head_dim, + ) + ) # [1,Hkv,0,D] + rotary_embeddings = first_attention.rotary_emb( + rotary_reference, + normalized_position_ids, + ) # each [1,1,T,D] + use_manual_gradient_checkpointing = ( + self._manual_gradient_checkpointing_enabled and self.training and hidden_states.requires_grad + ) + for decoder_layer in decoder_layers: + if use_manual_gradient_checkpointing: + + def _layer_forward( + layer_hidden_states: torch.Tensor, + layer_module: nn.Module = decoder_layer, + ) -> torch.Tensor: + return _forward_nemotron_natten_decoder_layer( + layer_module, + layer_hidden_states, + position_ids, + cumulative_seqlen, + max_seqlen, + rotary_embeddings, + ) + + hidden_states = torch.utils.checkpoint.checkpoint( + _layer_forward, + hidden_states, + preserve_rng_state=True, + use_reentrant=False, + ) # [1,T,D] + else: + hidden_states = _forward_nemotron_natten_decoder_layer( + decoder_layer, + hidden_states, + position_ids, + cumulative_seqlen, + max_seqlen, + rotary_embeddings, + ) # [1,T,D] + + hidden_states = self.text_decoder.model.norm(hidden_states) # [1,T,D] + return SimpleNamespace(last_hidden_state=hidden_states, past_key_values=None) + def _forward_from_embeddings( self, text_embeds: torch.Tensor, @@ -644,8 +849,30 @@ def _forward_from_embeddings( cache_position: torch.Tensor | None = None, use_cache: bool = False, past_key_values: Any | None = None, + packed_cumulative_seqlen: torch.Tensor | None = None, + packed_max_seqlen: int | None = None, ) -> SimpleNamespace: """Run the LM from caller-provided token embeddings.""" + has_packed_cumulative_seqlen = packed_cumulative_seqlen is not None + has_packed_max_seqlen = packed_max_seqlen is not None + if has_packed_cumulative_seqlen != has_packed_max_seqlen: + raise ValueError("packed_cumulative_seqlen and packed_max_seqlen must be provided together.") + if has_packed_cumulative_seqlen: + if self.packed_attention_backend != PACKED_ATTENTION_BACKEND_NATTEN: + raise ValueError("Packed varlen metadata was provided without the NATTEN packed attention backend.") + if attention_mask is not None or use_cache or past_key_values is not None: + raise ValueError("NATTEN packed training does not support masks or KV-cache inputs.") + if position_ids is None: + raise ValueError("NATTEN packed attention requires per-token position IDs.") + assert packed_cumulative_seqlen is not None + assert packed_max_seqlen is not None + return self._forward_nemotron_natten_from_embeddings( + text_embeds=text_embeds, + position_ids=position_ids, + cumulative_seqlen=packed_cumulative_seqlen, + max_seqlen=packed_max_seqlen, + ) + if self.spec.supports_inputs_embeds_forward: model_kwargs: dict[str, Any] = { "inputs_embeds": text_embeds, @@ -798,10 +1025,11 @@ def _encode(text: str) -> list[int]: image_placeholder = "" num_visuals = len(image_token_counts_by_visual) normalized_question = str(question).strip() + placeholder_count = normalized_question.count(image_placeholder) if num_visuals == 1: - normalized_question = _keep_only_first_image_placeholder(normalized_question, image_placeholder) + if placeholder_count > 1: + raise ValueError(f"VQA prompt has {placeholder_count} image placeholders for one visual input.") else: - placeholder_count = normalized_question.count(image_placeholder) if placeholder_count == 0: visual_placeholders = "\n".join(image_placeholder for _ in image_token_counts_by_visual) normalized_question = f"{visual_placeholders}\n{normalized_question}".strip() @@ -868,10 +1096,9 @@ def forward( ) -> tuple[torch.Tensor, int]: """Forward pass: inject image features into text and run causal LM. - When segment_ids are provided (packed sequences), the packed [1, S] sequence - is split into per-segment batch elements [num_segments, max_seg_len]. - HuggingFace flash attention with unpadding then produces cu_seqlens per batch - element, giving segment-isolated causal attention. + When segment_ids are provided, SDPA reshapes the packed sequence into a + padded per-segment batch, while NATTEN retains the packed layout and uses + cumulative sequence offsets. Both paths provide segment-isolated causal attention. Position IDs reset per segment for correct RoPE. Args: @@ -879,14 +1106,15 @@ def forward( image_feats_tensor: [N, encoder_dim] raw encoder output features. image_coords: [N, 4 or 5] spatial coordinates. image_patch_indices: [N_pooled] flat indices into [B*S] sequence where - merged image features should be inserted. Padded entries are -1. + merged image features should be inserted. The host collator validates + and removes padding sentinels before device transfer. image_layout: Optional batch layout slices for the image features. segment_ids: [B, S] segment IDs for packed sequences. Values >= 0 indicate valid segments, -1 indicates padding. When provided, enables segment-isolated attention and per-segment position IDs. Returns: - lm_logits: [B, S, vocab_size] next-token prediction logits. + Dense [B, S, vocab_size] logits. num_pooled_tokens: Number of image tokens after spatial merging. """ image_features = image_feats_tensor @@ -917,27 +1145,46 @@ def forward( text_embeds = self._embed_input_ids(input_ids) # [B, S, d] B, S, d = text_embeds.shape - # Zero out vision placeholder positions, then insert real image features - vision_mask = (input_ids != self.vision_token_id).to(dtype=text_embeds.dtype) - text_embeds = text_embeds * vision_mask[:, :, None] - - if num_pooled_tokens > 0 and len(image_patch_indices) > 0: - valid_indices = image_patch_indices[image_patch_indices >= 0] - max_idx = B * S - 1 - valid_indices = valid_indices[valid_indices <= max_idx] - num_to_insert = min(len(valid_indices), num_pooled_tokens) + # The host collator validates exact placeholder/index coverage and removes + # padding sentinels before tensors move to the accelerator. + # Zero out vision placeholder positions, then insert real image features. + vision_mask = (input_ids != self.vision_token_id).to(dtype=text_embeds.dtype) # [B,S] + text_embeds = text_embeds * vision_mask[:, :, None] # [B,S,d] - if num_to_insert > 0: - text_embeds_flat = text_embeds.reshape(B * S, d) - insert_idx = valid_indices[:num_to_insert].to(device=text_embeds.device, dtype=torch.long) - text_embeds_flat[insert_idx] = image_features[:num_to_insert].to(text_embeds.dtype) - text_embeds = text_embeds_flat.reshape(B, S, d) + if image_patch_indices.ndim != 1: + raise ValueError( + f"Image patch indices must be one-dimensional, got shape {tuple(image_patch_indices.shape)}." + ) + if ( + image_patch_indices.dtype == torch.bool + or image_patch_indices.is_floating_point() + or image_patch_indices.is_complex() + ): + raise TypeError(f"Image patch indices must use an integer dtype, got {image_patch_indices.dtype}.") + if image_patch_indices.numel() != num_pooled_tokens: + raise ValueError( + "Image feature/index count mismatch: " + f"got {num_pooled_tokens} pooled features and {image_patch_indices.numel()} indices." + ) + if image_patch_indices.numel() > 0: + insert_idx = image_patch_indices.to(device=input_ids.device, dtype=torch.long) # [N] + text_embeds_flat = text_embeds.reshape(B * S, d) # [B*S,d] + text_insert_idx = insert_idx.to(device=text_embeds.device) # [N] + text_embeds_flat[text_insert_idx] = image_features.to(text_embeds.dtype) # [N,d] + text_embeds = text_embeds_flat.reshape(B, S, d) # [B,S,d] # Segment-aware forward pass for packed sequences if segment_ids is not None: - lm_logits = self._forward_packed(text_embeds, input_ids, segment_ids) + lm_logits = self._forward_packed( + text_embeds, + input_ids, + segment_ids, + ) else: - lm_logits = self._forward_standard(text_embeds, input_ids=input_ids) + lm_logits = self._forward_standard( + text_embeds, + input_ids=input_ids, + ) return lm_logits, num_pooled_tokens @@ -992,9 +1239,9 @@ def _forward_packed( ) -> torch.Tensor: """Segment-isolated forward pass for packed sequences. - Splits packed [B, S] into per-segment batch [num_segments, max_seg_len]. - Flash attention with unpadding naturally isolates batch elements via cu_seqlens. - Position IDs reset per segment for correct RoPE. + SDPA reshapes packed [B, S] into padded per-segment batches. NATTEN keeps + valid tokens contiguous and uses cumulative sequence offsets. Position IDs + reset per segment for correct RoPE in both paths. Args: text_embeds: [B, S, d] embeddings with image features already injected. @@ -1002,14 +1249,14 @@ def _forward_packed( segment_ids: [B, S] segment IDs (>=0 for valid, -1 for padding). Returns: - lm_logits: [B, S, vocab_size] reassembled in original packed layout. + Dense [B, S, vocab_size] logits. """ B, S, d = text_embeds.shape device = text_embeds.device dtype = text_embeds.dtype # Process each batch element (typically B=1 for packed sequences) - all_logits = [] + all_logits: list[torch.Tensor] = [] for b in range(B): seg_ids = segment_ids[b] # [S] embeds = text_embeds[b] # [S, d] @@ -1035,42 +1282,54 @@ def _forward_packed( "Segment IDs must be contiguous within packed text decoder inputs" ) num_segments = seg_lengths.numel() - max_seg_len = int(seg_lengths.max().item()) + + segment_starts = torch.cumsum(seg_lengths, dim=0) - seg_lengths # [R] + segment_positions = torch.arange(valid_indices.numel(), device=device, dtype=torch.long) # [V] + segment_positions = segment_positions - torch.repeat_interleave(segment_starts, seg_lengths) # [V] + + if self.packed_attention_backend == PACKED_ATTENTION_BACKEND_NATTEN: + segment_lengths_int32 = seg_lengths.to(dtype=torch.int32) # [R] + zero_offset = torch.zeros(1, device=device, dtype=torch.int32) # [1] + cumulative_lengths = segment_lengths_int32.cumsum(dim=0, dtype=torch.int32) # [R] + cumulative_seqlen = torch.cat([zero_offset, cumulative_lengths], dim=0) # [R+1] + max_seqlen = int(seg_lengths.max().item()) + packed_embeds = valid_embeds.unsqueeze(0) # [1,V,D] + packed_position_ids = segment_positions.unsqueeze(0) # [1,V] + outputs = self._forward_from_embeddings( + text_embeds=packed_embeds, + attention_mask=None, + position_ids=packed_position_ids, + use_cache=False, + packed_cumulative_seqlen=cumulative_seqlen, + packed_max_seqlen=max_seqlen, + ) + valid_logits = self._lm_head(outputs.last_hidden_state) # [1,V,Vocab] + packed_logits = valid_logits.new_zeros((S, valid_logits.shape[-1])) # [S,Vocab] + packed_logits[valid_indices] = valid_logits[0] # [V,Vocab] + all_logits.append(packed_logits) + continue segment_rows = torch.repeat_interleave( torch.arange(num_segments, device=device, dtype=torch.long), seg_lengths, - ) - segment_starts = torch.cumsum(seg_lengths, dim=0) - seg_lengths - segment_positions = torch.arange(valid_indices.numel(), device=device, dtype=torch.long) - segment_positions = segment_positions - torch.repeat_interleave(segment_starts, seg_lengths) - - # Build per-segment batch: [num_segments, max_seg_len, d] - batch_embeds = torch.zeros(num_segments, max_seg_len, d, device=device, dtype=dtype) - batch_embeds[segment_rows, segment_positions] = valid_embeds - - position_template = torch.arange(max_seg_len, device=device, dtype=torch.long) - batch_mask = (position_template.unsqueeze(0) < seg_lengths.unsqueeze(1)).to(dtype=torch.long) - batch_pos = position_template.unsqueeze(0).expand(num_segments, -1) * batch_mask + ) # [V] + max_seg_len = int(seg_lengths.max().item()) + batch_embeds = torch.zeros(num_segments, max_seg_len, d, device=device, dtype=dtype) # [R,T,D] + batch_embeds[segment_rows, segment_positions] = valid_embeds # [V,D] - # Forward through the text decoder with one batch element per segment, - # which isolates attention across packed examples. - cache_position = position_template + position_template = torch.arange(max_seg_len, device=device, dtype=torch.long) # [T] + batch_mask = (position_template.unsqueeze(0) < seg_lengths.unsqueeze(1)).to(dtype=torch.long) # [R,T] + batch_pos = position_template.unsqueeze(0).expand(num_segments, -1) * batch_mask # [R,T] outputs = self._forward_from_embeddings( text_embeds=batch_embeds, attention_mask=batch_mask, position_ids=batch_pos, - cache_position=cache_position, + cache_position=position_template, use_cache=False, ) - seg_logits = self._lm_head(outputs.last_hidden_state) - # seg_logits: [num_segments, max_seg_len, vocab_size] - - # Reassemble into original packed layout [S, vocab_size] - V = seg_logits.shape[-1] - packed_logits = torch.zeros(S, V, device=device, dtype=seg_logits.dtype) - packed_logits[valid_indices] = seg_logits[segment_rows, segment_positions] - + seg_logits = self._lm_head(outputs.last_hidden_state) # [R,T,Vocab] + packed_logits = seg_logits.new_zeros((S, seg_logits.shape[-1])) # [S,Vocab] + packed_logits[valid_indices] = seg_logits[segment_rows, segment_positions] # [V,Vocab] all_logits.append(packed_logits) # Guard: empty batch (all samples skipped by collate) -> empty logits tensor. @@ -1078,7 +1337,7 @@ def _forward_packed( V = self.lm_config.vocab_size return torch.zeros(0, S, V, device=device, dtype=dtype) # Stack batch: [B, S, vocab_size] - return torch.stack(all_logits, dim=0) + return stack_with_bounded_inputs(all_logits, dim=0) # [B,S,Vocab] def _decode_generation_result( self, @@ -1197,7 +1456,8 @@ def _generate_from_prefix_embeddings( if len(generated_tokens) == 0: return input_ids - return torch.cat([input_ids, torch.cat(generated_tokens, dim=1)], dim=1) + generated_ids = cat_with_bounded_inputs(generated_tokens, dim=1) # [B,T_generated] + return torch.cat([input_ids, generated_ids], dim=1) @torch.no_grad() def generate_caption( @@ -1366,6 +1626,7 @@ def _encode(s: str) -> list[int]: return tok.encode(s, add_special_tokens=False) answer_decode_prefix_ids: list[int] | None = None + answer_suppress_token_ids: tuple[int, ...] = () if self.spec.family == QWEN3_SPEC.family: system_turn = ( [QWEN3_IM_START_TOKEN_ID] @@ -1383,6 +1644,7 @@ def _encode(s: str) -> list[int]: # Empty think block signals Qwen3 non-thinking mode. no_think = [QWEN3_THINK_START_TOKEN_ID] + _encode("\n\n") + [QWEN3_THINK_END_TOKEN_ID] + _encode("\n\n") asst_prefix += no_think + answer_suppress_token_ids = self.spec.suppress_token_ids eos_token_ids = self._get_eos_token_ids() elif self.spec.family == NEMOTRON_2B_SPEC.family: im_start_id = _get_required_token_id(tok, NEMOTRON_2B_IM_START_TOKEN) @@ -1398,6 +1660,7 @@ def _encode(s: str) -> list[int]: asst_prefix = [im_start_id] + _encode("assistant\n") if resolved_thinking_mode == VQA_THINKING_MODE_OFF: asst_prefix += [think_id, end_think_id] + answer_suppress_token_ids = (think_id, end_think_id) elif resolved_thinking_mode == VQA_THINKING_MODE_ON: answer_decode_prefix_ids = [think_id] + _encode("\n") asst_prefix += answer_decode_prefix_ids @@ -1436,9 +1699,7 @@ def _encode(s: str) -> list[int]: text_embeds=text_embeds, temperature=temperature, eos_token_ids=eos_token_ids, - suppress_token_ids=( - self.spec.suppress_token_ids if resolved_thinking_mode == VQA_THINKING_MODE_OFF else () - ), + suppress_token_ids=answer_suppress_token_ids, ) answer, metadata = self._decode_generation_result( diff --git a/cosmos_framework/model/tokenizer/models/utils.py b/cosmos_framework/model/tokenizer/models/utils.py index d105228f..7792e397 100644 --- a/cosmos_framework/model/tokenizer/models/utils.py +++ b/cosmos_framework/model/tokenizer/models/utils.py @@ -203,18 +203,23 @@ def batch_tensor_to_sparse(batch_tensor: torch.Tensor, patch_size: tuple[int, in """Convert a batch tensor to SparseTensor. Args: - batch_tensor: Input tensor of shape [B, C, H, W] or [B, T, H, W, C]. + batch_tensor: Input tensor of shape [B, C, H, W], [B, H, W, C], or [B, T, H, W, C]. patch_size: Patch size (Pt, Ph, Pw). Returns: SparseTensor with patches as features. """ + Pt, Ph, Pw = patch_size if batch_tensor.dim() == 4: if batch_tensor.shape[1] == 3: - batch_tensor = rearrange(batch_tensor, "b c h w -> b h w c") - batch_tensor = batch_tensor.unsqueeze(1) + batch_tensor = rearrange(batch_tensor, "b c h w -> b h w c") # [B,H,W,C] + batch_tensor = batch_tensor.unsqueeze(1) # [B,1,H,W,C] + if Pt > 1: + temporal_padding = batch_tensor.new_zeros( + (batch_tensor.shape[0], Pt - 1, *batch_tensor.shape[2:]) + ) # [B,Pt-1,H,W,C] + batch_tensor = torch.cat((batch_tensor, temporal_padding), dim=1) # [B,Pt,H,W,C] batch, T, H, W, _ = batch_tensor.shape - Pt, Ph, Pw = patch_size assert T % Pt == 0 and H % Ph == 0 and W % Pw == 0 num_temporal_patches = T // Pt num_height_patches = H // Ph diff --git a/cosmos_framework/model/tokenizer/utils/tensors.py b/cosmos_framework/model/tokenizer/utils/tensors.py new file mode 100644 index 00000000..62810a19 --- /dev/null +++ b/cosmos_framework/model/tokenizer/utils/tensors.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Tensor operation helpers shared by tokenizer training and inference.""" + +from collections.abc import Sequence + +import torch + +# PyTorch's CUDA multi-input cat path issued an illegal memory access on +# Blackwell with 130 and 194 views, while 98 passed. Keep a safety margin. +MAX_CAT_INPUTS_PER_CALL = 64 + + +def _normalize_tensor_sequence( + tensors: Sequence[torch.Tensor], +) -> list[torch.Tensor] | tuple[torch.Tensor, ...]: + """Return a PyTorch-compatible tensor sequence without copying common inputs.""" + if isinstance(tensors, (list, tuple)): + return tensors + return list(tensors) + + +def cat_with_bounded_inputs(tensors: Sequence[torch.Tensor], dim: int = 0) -> torch.Tensor: + """Concatenate tensors without exceeding the CUDA multi-input kernel's safe fan-in.""" + tensor_inputs = _normalize_tensor_sequence(tensors) + if len(tensor_inputs) <= MAX_CAT_INPUTS_PER_CALL: + return torch.cat(tensor_inputs, dim=dim) # [*cat_shape] + + current_level = list(tensor_inputs) + first_device = current_level[0].device + if any(tensor.device != first_device for tensor in current_level[1:]): + devices = ", ".join(sorted({str(tensor.device) for tensor in current_level})) + raise RuntimeError( + "cat_with_bounded_inputs requires tensors to share one device when the input count " + f"exceeds {MAX_CAT_INPUTS_PER_CALL}; got devices: {devices}" + ) + + target_dtype = current_level[0].dtype + for tensor in current_level[1:]: + target_dtype = torch.promote_types(target_dtype, tensor.dtype) + if any(tensor.dtype != target_dtype for tensor in current_level): + current_level = [tensor.to(dtype=target_dtype) for tensor in current_level] # [*input_shapes] + + while len(current_level) > MAX_CAT_INPUTS_PER_CALL: + # Each partial preserves all dimensions except for the concatenated axis. + current_level = [ + torch.cat(current_level[start : start + MAX_CAT_INPUTS_PER_CALL], dim=dim) + for start in range(0, len(current_level), MAX_CAT_INPUTS_PER_CALL) + ] # [*partial_cat_shapes] + return torch.cat(current_level, dim=dim) # [*cat_shape] + + +def stack_with_bounded_inputs(tensors: Sequence[torch.Tensor], dim: int = 0) -> torch.Tensor: + """Stack tensors without exposing the underlying cat kernel to unbounded fan-in.""" + tensor_inputs = _normalize_tensor_sequence(tensors) + if len(tensor_inputs) <= MAX_CAT_INPUTS_PER_CALL: + return torch.stack(tensor_inputs, dim=dim) # [*stack_shape] + expanded_inputs = [tensor.unsqueeze(dim) for tensor in tensor_inputs] # [*expanded_input_shapes] + return cat_with_bounded_inputs(expanded_inputs, dim=dim) # [*stack_shape] + + +__all__ = ["MAX_CAT_INPUTS_PER_CALL", "cat_with_bounded_inputs", "stack_with_bounded_inputs"] diff --git a/cosmos_framework/tools/flops/qwen3_vl.py b/cosmos_framework/tools/flops/qwen3_vl.py index c84a2444..7705795b 100644 --- a/cosmos_framework/tools/flops/qwen3_vl.py +++ b/cosmos_framework/tools/flops/qwen3_vl.py @@ -65,6 +65,7 @@ def compute_attention_flops( num_kv_heads: int, head_dim: int | None = None, is_causal: bool = False, + has_bias: bool = True, ) -> int: """Compute FLOPs for attention mechanism. @@ -81,6 +82,7 @@ def compute_attention_flops( count and the kernel work for the causal case. Defaults to False for bidirectional attention (e.g. the vision encoder); set True for causal decoders. + has_bias: Whether Q/K/V/O projections include bias terms. Returns: Total FLOPs for attention @@ -89,9 +91,9 @@ def compute_attention_flops( head_dim = hidden_size // num_heads # QKV projection: 3 linear layers (but KV uses num_kv_heads) - q_proj_flops = compute_linear_flops(hidden_size, num_heads * head_dim, seq_len, has_bias=True) - k_proj_flops = compute_linear_flops(hidden_size, num_kv_heads * head_dim, seq_len, has_bias=True) - v_proj_flops = compute_linear_flops(hidden_size, num_kv_heads * head_dim, seq_len, has_bias=True) + q_proj_flops = compute_linear_flops(hidden_size, num_heads * head_dim, seq_len, has_bias=has_bias) + k_proj_flops = compute_linear_flops(hidden_size, num_kv_heads * head_dim, seq_len, has_bias=has_bias) + v_proj_flops = compute_linear_flops(hidden_size, num_kv_heads * head_dim, seq_len, has_bias=has_bias) # Causal masking halves the work that scales with S^2 (QK^T, softmax, # attn @ V). The QKV / output projections are not affected. @@ -109,7 +111,7 @@ def compute_attention_flops( attn_v_matmul_flops = int(2 * num_heads * seq_len * seq_len * head_dim * causal_factor) # Output projection - o_proj_flops = compute_linear_flops(num_heads * head_dim, hidden_size, seq_len, has_bias=True) + o_proj_flops = compute_linear_flops(num_heads * head_dim, hidden_size, seq_len, has_bias=has_bias) total_flops = ( q_proj_flops @@ -335,6 +337,7 @@ def compute_text_decoder_flops( num_text_layers: int, head_dim: int = 128, is_causal: bool = True, + attention_bias: bool = False, # MoE parameters num_experts: int = 0, num_experts_per_tok: int = 0, @@ -357,6 +360,7 @@ def compute_text_decoder_flops( causal, so this defaults to True; set to False to recover the bidirectional upper bound (e.g. when comparing against a fitted runtime curve calibrated against the upper bound). + attention_bias: Whether Q/K/V/O attention projections include bias. num_experts: Total number of experts in MoE (0 means no MoE) num_experts_per_tok: Number of experts activated per token (top-k) moe_intermediate_size: Intermediate dimension for each MoE expert @@ -387,6 +391,7 @@ def compute_text_decoder_flops( num_key_value_heads, head_dim, is_causal=is_causal, + has_bias=attention_bias, ) # MLP or MoE (Qwen uses SwiGLU) @@ -409,7 +414,7 @@ def compute_text_decoder_flops( ln_flops = 2 * compute_layernorm_flops(total_tokens, hidden_size) # RMSNorm for Q and K (applied to each head dimension) - qk_norm_flops = 2 * compute_layernorm_flops(total_tokens * num_attention_heads, head_dim) + qk_norm_flops = compute_layernorm_flops(total_tokens * (num_attention_heads + num_key_value_heads), head_dim) decoder_flops += attn_flops + mlp_flops + ln_flops + qk_norm_flops @@ -435,6 +440,7 @@ def compute_qwen3vl_flops( include_lm_head: bool = True, spatial_merge_size: int = 2, is_causal: bool = True, + attention_bias: bool = False, # MoE parameters num_experts: int = 0, num_experts_per_tok: int = 0, @@ -466,6 +472,8 @@ def compute_qwen3vl_flops( is_causal: If True (default), apply the causal-attention factor of 0.5 to the text decoder's S^2 attention terms. The vision encoder is always treated as bidirectional. + attention_bias: Whether text-decoder Q/K/V/O attention projections + include bias. num_experts: Total number of experts in MoE (0 means no MoE) num_experts_per_tok: Number of experts activated per token (top-k) moe_intermediate_size: Intermediate dimension for each MoE expert @@ -540,6 +548,7 @@ def compute_qwen3vl_flops( num_text_layers=num_text_layers, head_dim=head_dim, is_causal=is_causal, + attention_bias=attention_bias, # MoE parameters num_experts=num_experts, num_experts_per_tok=num_experts_per_tok, @@ -612,6 +621,7 @@ def compute_qwen3vl_flops_from_config( """ # Extract MoE parameters if available (MoE models) text_config = config.text_config + attention_bias = getattr(text_config, "attention_bias", False) num_experts = getattr(text_config, "num_experts", 0) num_experts_per_tok = getattr(text_config, "num_experts_per_tok", 0) moe_intermediate_size = getattr(text_config, "moe_intermediate_size", None) @@ -635,6 +645,7 @@ def compute_qwen3vl_flops_from_config( head_dim=text_config.head_dim, spatial_merge_size=config.vision_config.spatial_merge_size, is_causal=is_causal, + attention_bias=attention_bias, # MoE parameters num_experts=num_experts, num_experts_per_tok=num_experts_per_tok, diff --git a/cosmos_framework/tools/flops/test_qwen3_vl.py b/cosmos_framework/tools/flops/test_qwen3_vl.py new file mode 100644 index 00000000..5413598c --- /dev/null +++ b/cosmos_framework/tools/flops/test_qwen3_vl.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from types import SimpleNamespace + +import pytest + +from cosmos_framework.tools.flops.qwen3_vl import ( + compute_attention_flops, + compute_layernorm_flops, + compute_mlp_flops, + compute_qwen3vl_flops, + compute_qwen3vl_flops_from_config, + compute_text_decoder_flops, +) + + +@pytest.mark.L0 +def test_attention_bias_controls_qkvo_projection_bias_terms() -> None: + seq_len = 3 + hidden_size = 8 + num_heads = 2 + num_kv_heads = 1 + head_dim = 4 + + without_bias = compute_attention_flops( + seq_len=seq_len, + hidden_size=hidden_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + has_bias=False, + ) + with_bias = compute_attention_flops( + seq_len=seq_len, + hidden_size=hidden_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + has_bias=True, + ) + + q_bias_flops = seq_len * num_heads * head_dim + k_bias_flops = seq_len * num_kv_heads * head_dim + v_bias_flops = seq_len * num_kv_heads * head_dim + o_bias_flops = seq_len * hidden_size + assert with_bias - without_bias == q_bias_flops + k_bias_flops + v_bias_flops + o_bias_flops + + +@pytest.mark.L0 +def test_text_decoder_qk_norm_counts_query_and_kv_heads_for_gqa() -> None: + total_tokens = 3 + hidden_size = 8 + intermediate_size = 16 + num_attention_heads = 2 + num_key_value_heads = 1 + head_dim = 4 + + actual = compute_text_decoder_flops( + total_tokens=total_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + num_text_layers=1, + head_dim=head_dim, + is_causal=False, + attention_bias=False, + ) + + attn_flops = compute_attention_flops( + seq_len=total_tokens, + hidden_size=hidden_size, + num_heads=num_attention_heads, + num_kv_heads=num_key_value_heads, + head_dim=head_dim, + is_causal=False, + has_bias=False, + ) + mlp_flops = compute_mlp_flops(total_tokens, hidden_size, intermediate_size, use_swiglu=True) + layernorm_flops = 2 * compute_layernorm_flops(total_tokens, hidden_size) + qk_norm_flops = compute_layernorm_flops( + total_tokens * (num_attention_heads + num_key_value_heads), + head_dim, + ) + + assert actual == attn_flops + mlp_flops + layernorm_flops + qk_norm_flops + + +@pytest.mark.L0 +def test_qwen3vl_flops_from_config_uses_attention_bias() -> None: + text_config = SimpleNamespace( + num_hidden_layers=1, + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=4, + vocab_size=32, + attention_bias=False, + ) + vision_config = SimpleNamespace( + depth=0, + hidden_size=4, + intermediate_size=8, + num_heads=1, + spatial_merge_size=2, + ) + config = SimpleNamespace(text_config=text_config, vision_config=vision_config) + + from_config = compute_qwen3vl_flops_from_config( + config=config, + total_tokens=5, + visual_tokens=1, + num_patches=None, + is_causal=False, + ) + direct_without_bias = compute_qwen3vl_flops( + num_text_layers=text_config.num_hidden_layers, + num_vision_layers=vision_config.depth, + hidden_size=text_config.hidden_size, + intermediate_size=text_config.intermediate_size, + num_attention_heads=text_config.num_attention_heads, + num_key_value_heads=text_config.num_key_value_heads, + vision_hidden_size=vision_config.hidden_size, + vision_intermediate_size=vision_config.intermediate_size, + vision_num_heads=vision_config.num_heads, + vocab_size=text_config.vocab_size, + total_tokens=5, + visual_tokens=1, + num_patches=None, + head_dim=text_config.head_dim, + spatial_merge_size=vision_config.spatial_merge_size, + is_causal=False, + attention_bias=False, + ) + direct_with_bias = compute_qwen3vl_flops( + num_text_layers=text_config.num_hidden_layers, + num_vision_layers=vision_config.depth, + hidden_size=text_config.hidden_size, + intermediate_size=text_config.intermediate_size, + num_attention_heads=text_config.num_attention_heads, + num_key_value_heads=text_config.num_key_value_heads, + vision_hidden_size=vision_config.hidden_size, + vision_intermediate_size=vision_config.intermediate_size, + vision_num_heads=vision_config.num_heads, + vocab_size=text_config.vocab_size, + total_tokens=5, + visual_tokens=1, + num_patches=None, + head_dim=text_config.head_dim, + spatial_merge_size=vision_config.spatial_merge_size, + is_causal=False, + attention_bias=True, + ) + + assert from_config == direct_without_bias + assert from_config["total_flops"] < direct_with_bias["total_flops"] diff --git a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py index fdc944f8..c417cf59 100644 --- a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py +++ b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py @@ -145,6 +145,18 @@ def dump_to_fileobj( obj = obj.cpu().numpy() h, w = obj.shape[1:-1] + # H.264 + yuv420p requires both dimensions to be even. Some eval videos + # keep the raw camera aspect ratio (for example 569x640), and ffmpeg + # fails with "width not divisible by 2" unless we pad before encoding. + # Pad minimally on the bottom/right edge so content coordinates remain + # unchanged and update ffmpeg's raw input size to match the padded array. + pad_h = h % 2 + pad_w = w % 2 + if pad_h or pad_w: + log.debug(f"Padding video from {w}x{h} to {w + pad_w}x{h + pad_h} for MP4 encoding") + obj = np.pad(obj, ((0, 0), (0, pad_h), (0, pad_w), (0, 0)), mode="edge") + h, w = obj.shape[1:-1] + # Default ffmpeg params that ensure width and height are set default_ffmpeg_params = ["-s", f"{w}x{h}"] @@ -160,13 +172,13 @@ def dump_to_fileobj( "output_params": ["-f", "mp4"], } else: - # Use provided ffmpeg_params if any, otherwise use defaults - final_ffmpeg_params = ffmpeg_params if ffmpeg_params is not None else default_ffmpeg_params + # Preserve caller-provided ffmpeg params, but always append the post-padding + # input size so custom params cannot leave ffmpeg seeing stale odd dimensions. mimsave_kwargs = { "fps": fps, "quality": quality, "macro_block_size": 1, - "ffmpeg_params": final_ffmpeg_params, + "ffmpeg_params": (ffmpeg_params or []) + default_ffmpeg_params, "output_params": ["-f", "mp4"], } # Update with any other kwargs diff --git a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler_test.py b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler_test.py new file mode 100644 index 00000000..121ea02e --- /dev/null +++ b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler_test.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from io import BytesIO + +import numpy as np +import pytest + +from cosmos_framework.utils.easy_io.handlers.imageio_video_handler import ImageioVideoHandler + + +@pytest.mark.L0 +def test_dump_to_fileobj_pads_odd_dimensions_for_mp4(monkeypatch): + calls = {} + + def fake_mimsave(file, obj, format, **kwargs): # pylint: disable=redefined-builtin + calls["file"] = file + calls["obj"] = obj + calls["format"] = format + calls["kwargs"] = kwargs + + monkeypatch.setattr( + "cosmos_framework.utils.easy_io.handlers.imageio_video_handler.imageio.mimsave", + fake_mimsave, + ) + + video = np.arange(2 * 3 * 5 * 3, dtype=np.uint8).reshape(2, 3, 5, 3) + ImageioVideoHandler().dump_to_fileobj(video, BytesIO(), crf=25) + + saved = calls["obj"] + assert saved.shape == (2, 4, 6, 3) + np.testing.assert_array_equal(saved[:, :3, :5, :], video) + np.testing.assert_array_equal(saved[:, 3:, :5, :], video[:, 2:3, :, :]) + np.testing.assert_array_equal(saved[:, :3, 5:, :], video[:, :, 4:5, :]) + assert calls["kwargs"]["ffmpeg_params"][-2:] == ["-s", "6x4"] + + +@pytest.mark.L0 +def test_dump_to_fileobj_appends_padded_size_to_custom_non_crf_ffmpeg_params(monkeypatch): + calls = {} + + def fake_mimsave(file, obj, format, **kwargs): # pylint: disable=redefined-builtin + calls["file"] = file + calls["obj"] = obj + calls["format"] = format + calls["kwargs"] = kwargs + + monkeypatch.setattr( + "cosmos_framework.utils.easy_io.handlers.imageio_video_handler.imageio.mimsave", + fake_mimsave, + ) + + video = np.zeros((2, 3, 5, 3), dtype=np.uint8) + ImageioVideoHandler().dump_to_fileobj(video, BytesIO(), ffmpeg_params=["-an"]) + + assert calls["obj"].shape == (2, 4, 6, 3) + assert calls["kwargs"]["ffmpeg_params"] == ["-an", "-s", "6x4"] diff --git a/cosmos_framework/utils/generator/aux_optimizer_utils.py b/cosmos_framework/utils/generator/aux_optimizer_utils.py new file mode 100644 index 00000000..c12356d4 --- /dev/null +++ b/cosmos_framework/utils/generator/aux_optimizer_utils.py @@ -0,0 +1,316 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared parameter-categorization for the orthogonalizing optimizers (Muon / Dion2). + +Both ``MuonWithAuxAdamW`` and ``Dion2WithAuxAdamW`` only apply their +orthogonalized (Newton-Schulz) update to *hidden* ``nn.Linear`` weight matrices. +Everything else -- token / positional embeddings, the output head (``lm_head``), +layer-norm scales, biases, and any non-Linear parameter -- must use the auxiliary +AdamW. + +The tricky part is reliably identifying the **embeddings** and the **output head** +across every model architecture in the repo (Qwen3, GPT-OSS, DeepSeekV3, Gemma4, +Qwen3-VL-MoE, unified MoT, ...). We do not rely on a single hard-coded name. +Instead :func:`split_orthogonalizable_params` combines four signals: + +1. **Module type** -- ``nn.Embedding`` weights always go to AdamW (and they are + never ``nn.Linear`` so they would never reach Muon anyway). +2. **Name keywords** -- a configurable set of substrings (default ``{"lm_head"}``, + the convention used by every LLM in this repo) marks output-head Linears. +3. **Tied weights** -- a Linear whose ``weight`` tensor is *the same object* as an + ``nn.Embedding`` weight (``tie_word_embeddings=True``) is the tied output head. +4. **Vocabulary shape** -- a Linear whose output dimension equals some + ``nn.Embedding.num_embeddings`` (the vocab size) is treated as an output + projection. + +Signals (2)-(4) are OR-ed, so an output head is excluded from Muon/Dion2 even if a +new architecture names it differently or ties it to the embedding. The failure +mode of the heuristics is conservative: a misclassified weight falls back to +AdamW (always safe) rather than being orthogonalized. +""" + +import torch +import torch.nn as nn + +# Substrings (matched against the dotted module name) that force an ``nn.Linear`` +# onto the AdamW side. Every LLM in the repo names its head ``lm_head``; the extra +# entries cover common alternative names used elsewhere. +DEFAULT_ADAMW_MODULE_KEYWORDS: tuple[str, ...] = ("lm_head", "embed_out", "output_layer") + + +def is_output_head_linear( + module_name: str, + *, + adamw_module_keywords: tuple[str, ...], +) -> bool: + """Whether an ``nn.Linear`` is an output head (and must use AdamW, not Muon). + + Name-based: every LLM backbone in the repo names its head with an + ``adamw_module_keywords`` substring (default ``lm_head``). Embeddings are + excluded separately by module type in :func:`split_orthogonalizable_params` + (``nn.Embedding`` never reaches this function), and a tied head shares its + weight object with an embedding, so it is caught there too -- hence the name + check alone is sufficient for every current model. + """ + return any(keyword in module_name for keyword in adamw_module_keywords) + + +def split_orthogonalizable_params( + model: nn.Module, + optimizer_param_ids: set[int], + adamw_module_keywords: tuple[str, ...] = DEFAULT_ADAMW_MODULE_KEYWORDS, + expert_param_keywords: tuple[str, ...] = (), +) -> tuple[list[nn.Parameter], list[nn.Parameter], dict[nn.Parameter, str]]: + """Split ``model``'s trainable params by whether they can be orthogonalized. + + Args: + model: The (sub)module to walk -- typically the trainable ``net``. + optimizer_param_ids: ``id()`` of every parameter actually owned by the + optimizer's param groups. Only these are categorized, so that + ``keys_to_select`` filtering is respected and ``state_dict`` stays + consistent. + adamw_module_keywords: Name substrings that force an ``nn.Linear`` onto the + unorthogonalizable (AdamW) side (output heads). + expert_param_keywords: Name substrings that mark **stacked MoE expert** + parameters -- raw ``nn.Parameter`` tensors of shape ``[num_experts, M, + N]`` (e.g. ``gate_up_proj`` / ``down_proj``). When a 3-D+ param matches, + it is routed to the orthogonalizable side so the optimizer can treat + each expert slice as its own matrix. Empty (default) keeps expert + params on AdamW, i.e. no behavior change. + + NOTE (sharding assumption): the optimizer orthogonalizes these per + expert slice assuming the tensor is sharded on dim 0 (the expert axis), + which holds for the FSDP2 ``fully_shard`` path used here (FSDP2 shards + every parameter on dim 0) and makes the update communication-free. It is + NOT guaranteed if tensor/expert parallelism shards *within* an expert + matrix (dim 1/2); that case is unsupported and rejected at step time by + ``MuonWithAuxAdamW._step_stacked_muon`` / ``Dion2WithAuxAdamW._step_stacked_dion2``. + This was not exhaustively audited across every parallelization config, + hence the runtime guard there. + + TODO(expert-parallelism): support 2-D sharding of stacked expert params. + With expert parallelism, dim 0 would be sharded along the expert axis and + dim 1 along the FSDP axis simultaneously, so each rank holds a shard of a + *slice* of each expert matrix rather than whole expert matrices. The + per-expert-slice orthogonalization here assumes dim-0-only sharding (whole + expert matrices are local), so it would need to reconstruct each expert + matrix across the FSDP axis (an all-gather along dim 1, like the dense + 2-D path) before running Newton-Schulz. Doable, but needs changes. + + Returns: + ``(orthogonalizable, unorthogonalizable, param_to_name)`` where + ``orthogonalizable`` holds the hidden ``nn.Linear`` weights (2-D) and, when + ``expert_param_keywords`` is set, the stacked expert tensors (3-D); the + optimizer splits those by rank. ``unorthogonalizable`` is everything else + (embeddings, output head, biases, norms, other non-Linear params), which + uses the auxiliary AdamW. + """ + orthogonalizable: list[nn.Parameter] = [] + unorthogonalizable: list[nn.Parameter] = [] + param_to_name: dict[nn.Parameter, str] = {} + categorized: set[int] = set() + + for name, param in model.named_parameters(): + if param.requires_grad: + param_to_name[param] = name + + def _eligible(p: nn.Parameter) -> bool: + return p.requires_grad and id(p) in optimizer_param_ids and id(p) not in categorized + + def _is_stacked_expert(param: nn.Parameter) -> bool: + if not expert_param_keywords or param.ndim < 3: + return False + name = param_to_name.get(param, "") + return any(keyword in name for keyword in expert_param_keywords) + + for module_name, module in model.named_modules(): + if isinstance(module, nn.Linear): + is_head = is_output_head_linear( + module_name, + adamw_module_keywords=adamw_module_keywords, + ) + if _eligible(module.weight): + (unorthogonalizable if is_head else orthogonalizable).append(module.weight) + categorized.add(id(module.weight)) + if module.bias is not None and _eligible(module.bias): + unorthogonalizable.append(module.bias) + categorized.add(id(module.bias)) + else: + # Embeddings (nn.Embedding), norms, conv, and any raw nn.Parameter -> + # unorthogonalizable (AdamW), except stacked MoE expert tensors when + # opted in. ``recurse=False`` so each owning module handles its own + # params exactly once. + for param in module.parameters(recurse=False): + if not _eligible(param): + continue + (orthogonalizable if _is_stacked_expert(param) else unorthogonalizable).append(param) + categorized.add(id(param)) + + return orthogonalizable, unorthogonalizable, param_to_name + + +# ----------------------------------------------------------------------------- +# Shared orthogonalization / momentum math (used by both MuonWithAuxAdamW and +# Dion2WithAuxAdamW). Kept here so the two optimizers do not duplicate them. +# ----------------------------------------------------------------------------- + + +@torch.compile +def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int = 5) -> torch.Tensor: + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. + + Uses a quintic iteration whose coefficients are selected to maximize the slope at zero. + This produces US'V^T where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which + empirically does not hurt model performance relative to exact UV^T. + + Args: + G: Input tensor of shape (m, n) where m, n >= 1. + steps: Number of Newton-Schulz iterations. + + Returns: + Orthogonalized tensor of same shape as G. + """ + assert G.ndim == 2, "Input must be 2D" + a, b, c = (3.4445, -4.7750, 2.0315) + X = G.bfloat16() + + # Transpose if tall matrix for numerical stability + if G.size(0) > G.size(1): + X = X.T + + # Ensure spectral norm is at most 1 (global norm, matching Moonlight) + X = X / (X.norm() + 1e-7) + + # Perform Newton-Schulz iterations + for _ in range(steps): + A = X @ X.T + B = b * A + c * A @ A + X = a * X + B @ X + + # Transpose back if we transposed earlier + if G.size(0) > G.size(1): + X = X.T + + return X + + +@torch.compile +def zeropower_via_newtonschulz5_batched(G: torch.Tensor, steps: int = 5) -> torch.Tensor: + """Batched Newton-Schulz over a stack of matrices. + + Same quintic iteration as :func:`zeropower_via_newtonschulz5`, but applied to a + batch ``G`` of shape ``[..., M, N]`` (e.g. stacked MoE experts ``[E, M, N]``) + using batched matmuls. All matrices in the batch share ``M, N``, so the + transpose-if-tall decision is uniform across the batch. + """ + assert G.ndim >= 3, "Batched input must be at least 3D ([..., M, N])" + a, b, c = (3.4445, -4.7750, 2.0315) + X = G.bfloat16() + + transposed = False + if X.size(-2) > X.size(-1): + X = X.mT + transposed = True + + # Per-matrix spectral-norm normalization (norm over the last two dims). + X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) + + for _ in range(steps): + A = X @ X.mT + B = b * A + c * A @ A + X = a * X + B @ X + + if transposed: + X = X.mT + + return X + + +def compute_pre_ns_update( + grad: torch.Tensor, + momentum_buffer: torch.Tensor, + momentum: float = 0.95, + nesterov: bool = True, +) -> torch.Tensor: + """ + Compute the pre-Newton-Schulz update (momentum + optional Nesterov). + + This is separated from NS so that momentum/Nesterov can be applied on shards + before all-gathering for distributed NS. + + Args: + grad: Gradient tensor. + momentum_buffer: Momentum buffer (modified in-place). + momentum: Momentum coefficient. + nesterov: Whether to use Nesterov momentum. + + Returns: + Pre-NS update tensor (same shape as grad). + """ + # SGD-style momentum: buf = momentum * buf + grad (matching Moonlight) + momentum_buffer.mul_(momentum).add_(grad) + + # Nesterov: g = g + momentum * buf, else just use buf + if nesterov: + return grad.add(momentum_buffer, alpha=momentum) + else: + return momentum_buffer.clone() + + +def compute_pre_ns_update_moe_expert( + grad: torch.Tensor, + momentum_buffer: torch.Tensor, + momentum: float = 0.95, + nesterov: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Per-expert masked momentum for stacked MoE experts. + + Identical to :func:`compute_pre_ns_update` for *active* experts (those whose + gradient slice has any nonzero element, i.e. tokens were routed to them this + step), but FREEZES the momentum of *inactive* experts -- no momentum decay + and no accumulation -- instead of decaying it. + + This distinction matters for Muon/Dion but not for Adam: Newton-Schulz + renormalizes whatever it is handed to unit spectral norm, so an inactive + expert's (decayed but nonzero) stale momentum would be blown back up into a + full-strength update in a stale direction. Freezing keeps the momentum in + reserve until the expert is routed again, matching the reference + (microsoft/dion), which sits out params/experts that received no gradient. + + This function only handles the momentum recurrence. The caller MUST also, for + inactive experts: (1) zero the Newton-Schulz output, and (2) skip the weight + update and weight decay. The returned ``active`` mask is provided for exactly + that. + + Args: + grad: Local expert gradients, shape ``[E, M, N]`` (E = local experts). + momentum_buffer: Momentum buffer ``[E, M, N]``, modified in place. + momentum: Momentum coefficient. + nesterov: Whether to use Nesterov momentum. + + Returns: + pre_ns: Pre-Newton-Schulz update, shape ``[E, M, N]``. + active: Per-expert bool mask ``[E]``; True where the expert got a gradient. + """ + # Per-expert active mask: True iff the expert's gradient slice has any nonzero + # element. An expert with no tokens routed produces an exactly-zero gradient + # slice, so this is an exact "was this expert updated" test. + active = (grad != 0).flatten(1).any(dim=1) # [E] bool + a = active.view(-1, 1, 1).to(momentum_buffer.dtype) # [E, 1, 1] in {0, 1} + + # Masked SGD-style momentum: + # active -> momentum * buf + grad (== compute_pre_ns_update) + # inactive -> buf unchanged (factor 1.0; grad is 0 there anyway) => frozen + momentum_buffer.mul_(1 - a * (1 - momentum)).add_(grad) + + # Nesterov: g = g + momentum * buf, else just use buf. Inactive experts get a + # bogus (renormalized) value here, but the caller zeros their NS output. + if nesterov: + pre_ns = grad.add(momentum_buffer, alpha=momentum) + else: + pre_ns = momentum_buffer.clone() + + return pre_ns, active diff --git a/cosmos_framework/utils/generator/dion2_with_aux_adamw.py b/cosmos_framework/utils/generator/dion2_with_aux_adamw.py new file mode 100644 index 00000000..61b56682 --- /dev/null +++ b/cosmos_framework/utils/generator/dion2_with_aux_adamw.py @@ -0,0 +1,1149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +""" +Dion2WithAuxAdamW optimizer implementation. + +DION2 (Distributed Orthogonalization) for nn.Linear weight matrices, +with auxiliary AdamW for embeddings, biases, norms, and output layers (lm_head). + +This implementation combines elements from: + +1. Microsoft DION2 (https://github.com/microsoft/dion): + - All-to-all communication pattern for efficient distributed orthogonalization + - Submatrix selection (top-k rows/columns by L1 norm) + - Error feedback for unselected parts + - Async operations for overlapping communication with computation + (TBD: not realized yet -- see the all-to-all sites in + ``_process_dion2_batch_distributed``, which currently ``wait()`` immediately) + +2. KellerJordan/Muon (https://github.com/KellerJordan/Muon): + - Newton-Schulz orthogonalization algorithm + - Quintic iteration coefficients (a=3.4445, b=-4.7750, c=2.0315) + +3. FusedAdam (cosmos_framework/utils/generator/fused_adam.py): + - DTensor handling for FSDP/TP compatibility + - Transformer Engine fused AdamW kernel + +Key differences from MuonWithAuxAdamW: +- Uses all-to-all instead of all-gather (no redundant NS computation) +- Batches params in groups of world_size for efficient distribution +- Supports submatrix selection (fraction parameter) for sparse orthogonalization +- Each rank computes NS for exactly one param per batch (truly parallel) + +Sharding -> orthogonalization (the core idea) +--------------------------------------------- +Under FSDP each weight matrix is a DTensor split across GPUs, but Newton-Schulz +(NS) is a *whole-matrix* op and needs the full matrix in one place. The two +optimizers assemble it differently: + +Weight A sharded over 4 GPUs (one row-slice each): + + GPU0:[A0] GPU1:[A1] GPU2:[A2] GPU3:[A3] + +Muon -- all-gather: every GPU rebuilds the *same* full A and runs NS on it, so +the NS work is duplicated world_size times: + + all_gather(A) -> each GPU holds full A -> every GPU runs NS(A) (N x redundant) + +DION2 -- all-to-all: process world_size matrices (A,B,C,D) together and give each +GPU one complete matrix, so the N matrices are orthogonalized in parallel with no +redundant compute: + + before (each GPU has one slice of every matrix): + + A B C D + GPU0 [A0] [B0] [C0] [D0] + GPU1 [A1] [B1] [C1] [D1] + GPU2 [A2] [B2] [C2] [D2] + GPU3 [A3] [B3] [C3] [D3] + + --- all_to_all #1 (transpose) ---> each GPU now owns one FULL matrix: + + GPU0: [A0 A1 A2 A3] = A -> NS(A) + GPU1: [B0 B1 B2 B3] = B -> NS(B) + GPU2: [C0 C1 C2 C3] = C -> NS(C) + GPU3: [D0 D1 D2 D3] = D -> NS(D) + + --- all_to_all #2 (transpose back) ---> each GPU gets its own + orthogonalized slice of every matrix, then applies the update. + +Matrices are grouped by local shard shape into batches of world_size so the +all-to-all tensors are uniform (see ``_create_dion2_batches``). With +``fraction < 1.0`` only the top-k rows/cols (by L1 norm) are sent through this +dance, and the unselected part is carried forward via error feedback. + +TODO(hsdp-replicate-redundancy): eliminate redundant orthogonalization work in +the replicate dimension under HSDP. ``world_size`` above is the *shard* mesh dim +only -- the all-to-all is confined to the FSDP shard group and the replicate +(data-parallel) dim keeps its ``Replicate`` placement. That is correct, but each +replica group independently reruns the full Newton-Schulz on identical (post +all-reduce) gradients, so the NS *compute* is duplicated ``R`` times (R = +replicate degree; e.g. 2x for the 30B-A3B run at shard_degree=64 on 128 GPUs). +This is the standard data-parallel optimizer redundancy (Adam has it too) but is +pricier here because NS is several matmuls per matrix rather than an element-wise +step. It could be removed by distributing the matrices across the *full* 2-D mesh +(shard x replicate) so every rank orthogonalizes a distinct matrix, then +broadcasting the results back across the replicate dim -- trading the duplicate +compute for an extra cross-replica collective. Worth doing only if R is large or +the optimizer step becomes a step-time bottleneck. +""" + +import math + +import torch +import torch.distributed as dist +import torch.nn as nn +import transformer_engine as te +import transformer_engine_torch as tex +from torch.distributed.tensor import DTensor, Shard + +from cosmos_framework.utils import log +from cosmos_framework.utils.misc import get_local_tensor_if_DTensor +from cosmos_framework.utils.generator.aux_optimizer_utils import ( + DEFAULT_ADAMW_MODULE_KEYWORDS, + compute_pre_ns_update, + compute_pre_ns_update_moe_expert, + split_orthogonalizable_params, + zeropower_via_newtonschulz5, + zeropower_via_newtonschulz5_batched, +) + + +class Dion2WithAuxAdamW(torch.optim.Optimizer): + """ + Dion2WithAuxAdamW optimizer. + + Uses DION2 (Distributed Orthogonalization) for nn.Linear hidden weight matrices, + and AdamW for embeddings, biases, layer norms, and output heads (lm_head). + + Key features: + - All-to-all communication for efficient distributed NS (no redundant compute) + - Submatrix selection: only orthogonalize top-k rows/columns + - Error feedback: maintains momentum for unselected parts + - Batched processing: handles world_size params per batch + + Args: + params: Iterable of parameters to optimize. + lr: Base learning rate. + muon_momentum: Momentum coefficient for Muon/DION2. + muon_lr_scale: Scale factor for Muon LR adjustment. + ns_steps: Number of Newton-Schulz iterations. + nesterov: Whether to use Nesterov momentum. + fraction: Fraction of rows/columns to orthogonalize (0 < fraction <= 1). + ef_decay: Error feedback decay factor for selected submatrix. + weight_decay: Weight decay for all parameters. + adam_betas: Beta coefficients for the auxiliary AdamW side. + eps: Epsilon for AdamW numerical stability. + use_distributed: Whether to use distributed operations. + """ + + def __init__( + self, + params, + lr: float = 1e-4, + muon_momentum: float = 0.95, + muon_lr_scale: float = 0.2, + ns_steps: int = 5, + nesterov: bool = True, + fraction: float = 1.0, # 1.0 = full matrix, <1.0 = submatrix selection + ef_decay: float = 0.95, + weight_decay: float = 0.1, + adam_betas: tuple[float, float] = (0.9, 0.95), + eps: float = 1e-8, + use_distributed: bool = True, + capturable: bool = False, + master_weights: bool = False, + adamw_module_keywords: tuple[str, ...] | None = None, + expert_param_keywords: tuple[str, ...] | None = None, + **kwargs, + ): + if kwargs: + ignored_keys = list(kwargs.keys()) + expected_ignored = {"fused", "keys_to_select", "adamw_betas", "adamw_eps"} + unexpected = set(ignored_keys) - expected_ignored + if unexpected: + import warnings + + warnings.warn(f"Dion2WithAuxAdamW ignoring unexpected kwargs: {unexpected}") + + if not (0.0 < fraction <= 1.0): + raise ValueError(f"fraction must be in (0, 1], got {fraction}") + + # Master weights requires capturable mode + if master_weights and not capturable: + raise RuntimeError("Master weights is currently only supported with capturable=True.") + + # Store hyperparameters + # Note: lr is accessed via property that reads from param_groups + # to support LR schedulers (which update param_groups[X]["lr"]) + self.wd = weight_decay + self.muon_momentum = muon_momentum + self.muon_lr_scale = muon_lr_scale + self.ns_steps = ns_steps + self.nesterov = nesterov + self.fraction = fraction + self.ef_decay = ef_decay + self.adam_betas = tuple(adam_betas) if isinstance(adam_betas, list) else adam_betas + self.eps = eps + + # Name substrings that route an nn.Linear to AdamW (output heads). Used by + # categorize_params alongside tied-weight / vocab-shape detection. + self.adamw_module_keywords = ( + tuple(adamw_module_keywords) if adamw_module_keywords else DEFAULT_ADAMW_MODULE_KEYWORDS + ) + # Name substrings that route stacked MoE expert params ([E, M, N]) to the + # DION2 side (each expert slice orthogonalized). Empty = experts stay on + # AdamW (no behavior change). + self.expert_param_keywords = tuple(expert_param_keywords) if expert_param_keywords else () + + # Distributed settings + self.use_distributed = use_distributed and dist.is_initialized() + self._world_size = 1 + self._device_rank = 0 + self._process_group = None + self._device_mesh = None + + # Master weights settings (for mixed-precision training stability) + self.capturable = capturable + self.master_weights = master_weights + + # Parameter lists + self.dion2_params: list[nn.Parameter] = [] + self.adamw_params: list[nn.Parameter] = [] + # Stacked MoE expert params ([E, M, N]); orthogonalized per expert slice. + self.stacked_dion2_params: list[nn.Parameter] = [] + self.param_to_name: dict[nn.Parameter, str] = {} + self._dion2_batches: list[list[nn.Parameter]] = [] + + # Master weight copies (populated by _create_master_weights after categorize_params) + self._dion2_masters: list[torch.Tensor] = [] + self._adamw_masters: list[torch.Tensor] = [] + + # Transformer Engine fused Adam. The zero buffer is the noop flag required + # as the second argument of TE's multi_tensor_applier; it is a fixed + # constant here (no AMP overflow handling). + self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device="cuda") + self._multi_tensor_adam = tex.multi_tensor_adam + self._multi_tensor_adam_capturable = tex.multi_tensor_adam_capturable + self._multi_tensor_adam_capturable_master = tex.multi_tensor_adam_capturable_master + + # Initialize base optimizer. betas / eps go in the defaults so each param + # group carries them (FusedAdam convention); the AdamW step reads them + # per-group, enabling per-group overrides and exact FusedAdam parity. + defaults = dict(lr=lr, weight_decay=weight_decay, betas=self.adam_betas, eps=eps) + super().__init__(params, defaults) + + # Convert LR to tensor for capturable mode + if capturable: + for idx, group in enumerate(self.param_groups): + if len(group["params"]) == 0: + continue + device = group["params"][0].device + if isinstance(group["lr"], float): + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32) + self.param_groups[idx]["lr"] = group["lr"].to(device=device) + + # id(param) -> owning param_group, so the DION2 and AdamW updates can read + # the *per-group* lr / weight_decay (honors lr_multipliers and + # disable_weight_decay_for_1d_params). With a single group it degenerates + # to a single global lr/wd, matching the original (reference) behavior. + self._param_group_map: dict[int, dict] = {} + self._adamw_param_ids: set[int] = set() + # Master weights are created lazily on the first step() (FusedAdam-style), + # so that after a checkpoint resume they are rebuilt from the *restored* + # params rather than the freshly-initialized ones. + self._masters_initialized = False + self._param_to_master: dict[int, torch.Tensor] = {} + + log.info(f"Dion2WithAuxAdamW master_weights: {master_weights} capturable: {capturable}") + + def categorize_params(self, model: nn.Module) -> None: + """ + Categorize parameters into DION2 and AdamW groups; also set up distributed + configuration from the DTensor DeviceMesh. + + DION2 is used for hidden ``nn.Linear`` weights only. Embeddings, the output + head (``lm_head`` / tied / vocab-shaped projection), biases, norms, and any + non-Linear parameter go to AdamW. See + :func:`split_orthogonalizable_params` for the architecture-agnostic + embedding / output-head detection. + """ + optimizer_param_ids = {id(p) for group in self.param_groups for p in group["params"]} + + orthogonalizable, self.adamw_params, self.param_to_name = split_orthogonalizable_params( + model, + optimizer_param_ids, + adamw_module_keywords=self.adamw_module_keywords, + expert_param_keywords=self.expert_param_keywords, + ) + + # Every optimizer param lands in exactly ONE of three disjoint buckets, + # each updated by a different function in step(): + # 1. self.dion2_params -> _step_dion2 (dense 2D Linear weights) + # 2. self.stacked_dion2_params -> _step_stacked_dion2 (3D MoE experts [E, M, N]) + # 3. self.adamw_params -> _step_adamw (embeddings/head/norms/biases/1D) + # split_orthogonalizable_params separates the orthogonalizable Linear + # weights (buckets 1+2) from everything else (bucket 3); here we further + # split the orthogonalizable set into 2D (DION2 all-to-all path) vs stacked + # 3D MoE experts (orthogonalized per expert slice in _step_stacked_dion2). + self.dion2_params = [p for p in orthogonalizable if p.ndim == 2] + self.stacked_dion2_params = [p for p in orthogonalizable if p.ndim >= 3] + + # Sort by size for load balancing + self.dion2_params = sorted(self.dion2_params, key=lambda x: x.numel(), reverse=True) + + # Setup distributed from first DTensor param + self._setup_distributed_from_params() + + # Create batches of world_size params + self._create_dion2_batches() + + dion2_numel = sum(p.numel() for p in self.dion2_params) + adamw_numel = sum(p.numel() for p in self.adamw_params) + stacked_dion2_numel = sum(p.numel() for p in self.stacked_dion2_params) + + log.info( + f"Dion2WithAuxAdamW: {len(self.dion2_params)} Muon params ({dion2_numel:,} elements), " + f"{len(self.stacked_dion2_params)} stacked-expert params ({stacked_dion2_numel:,} elements), " + f"{len(self.adamw_params)} AdamW params ({adamw_numel:,} elements), " + f"world_size={self._world_size}, {len(self._dion2_batches)} batches" + ) + + # Log Muon param details + log.info("DION2 parameters (layer name -> shape):") + for p in self.dion2_params: + name = self.param_to_name.get(p, "unknown") + log.info(f" {name}: {tuple(p.shape)}") + + # Build the param -> owning group map for per-group lr / weight_decay + # lookups during the DION2 and AdamW steps. + self._param_group_map = {} + for group in self.param_groups: + for p in group["params"]: + self._param_group_map[id(p)] = group + self._adamw_param_ids = {id(p) for p in self.adamw_params} + + # NOTE: master weights are intentionally NOT created here. They are + # created lazily on the first step() (see _maybe_init_master_weights) so + # that a checkpoint resume rebuilds them from the restored params. + + def _assert_homogeneous_sharding(self) -> None: + """Verify every DION2 param shares the first param's mesh and sharding. + + The all-to-all path caches a single distributed config (device_mesh, shard + mesh dim, shard tensor dim, world_size, process group) derived from + ``dion2_params[0]`` and applies it to *every* DION2 param. That is only + correct if all params are sharded identically. Heterogeneous sharding -- a + param on a different mesh, sharded on a different dim, or replicated while + others are sharded -- would be processed with the wrong layout and silently + corrupt the update, so it is rejected up front (once, at setup). VFM's FSDP2 + shards every 2-D weight on dim 0 with one mesh, so this is a no-op today; the + guard fails loudly if TP/EP or mixed replication is ever introduced. + """ + p0 = self.dion2_params[0] + for p in self.dion2_params[1:]: + mismatch = isinstance(p, DTensor) != isinstance(p0, DTensor) or ( + isinstance(p0, DTensor) and (p.device_mesh != p0.device_mesh or p.placements != p0.placements) + ) + if mismatch: + name = self.param_to_name.get(p, "unknown") + raise NotImplementedError( + f"DION2 requires all params to share the first param's sharding, but '{name}' " + f"differs (placements={getattr(p, 'placements', None)}). " + f"Heterogeneous sharding is unsupported." + ) + + def _setup_distributed_from_params(self) -> None: + """Extract distributed config from DTensor DeviceMesh.""" + if not self.dion2_params: + return + + # The config below is derived from dion2_params[0] and reused for every + # param, so first confirm they are all sharded identically. + self._assert_homogeneous_sharding() + + first_param = self.dion2_params[0] + if isinstance(first_param, DTensor): + device_mesh = first_param.device_mesh + placements = first_param.placements + + # Find the shard dimension in the mesh + for mesh_dim_idx, placement in enumerate(placements): + if placement.is_shard(): + self._device_mesh = device_mesh + self._world_size = device_mesh.size(mesh_dim=mesh_dim_idx) + self._device_rank = device_mesh.get_local_rank(mesh_dim=mesh_dim_idx) + self._process_group = device_mesh.get_group(mesh_dim=mesh_dim_idx) + self._shard_mesh_dim = mesh_dim_idx + self._shard_tensor_dim = placement.dim + log.info( + f"DION2 distributed setup: world_size={self._world_size}, " + f"rank={self._device_rank}, shard_dim={self._shard_tensor_dim}" + ) + return + + # Fallback: not distributed or not sharded + self._world_size = 1 + self._device_rank = 0 + + def _create_dion2_batches(self) -> None: + """ + Group Muon params by GLOBAL shape, then batch within each shape group. + + The distributed step (``_process_dion2_batch_distributed``) stacks a batch of + params into a single DTensor and redistributes it, so the batches must be + identical ACROSS ranks. Group by the *global* shape (same on every rank), NOT + the local shard shape -- under uneven sharding the local shard shape differs + per rank and would make ranks build inconsistent batches. ``batch_size = + world_size`` ensures each rank orthogonalizes exactly one param per batch. + """ + self._dion2_batches = [] + batch_size = self._world_size + + # Step 1: Group params by global shape (identical on all ranks). + shape_groups: dict[tuple, list[nn.Parameter]] = {} + for p in self.dion2_params: + shape = tuple(p.shape) + if shape not in shape_groups: + shape_groups[shape] = [] + shape_groups[shape].append(p) + + # Step 2: Create batches within each shape group + for shape, params in shape_groups.items(): + for i in range(0, len(params), batch_size): + batch = params[i : i + batch_size] + self._dion2_batches.append(batch) + + # Log batch info + num_shape_groups = len(shape_groups) + num_batches = len(self._dion2_batches) + if self._dion2_batches: + # Count batches that need padding + padded_batches = sum(1 for b in self._dion2_batches if len(b) < batch_size) + log.info( + f"DION2: {len(self.dion2_params)} params grouped into {num_shape_groups} shape groups, " + f"{num_batches} batches (world_size={batch_size}, {padded_batches} need padding)" + ) + # Log shape group details + for shape, params in shape_groups.items(): + log.info(f" Shape {shape}: {len(params)} params") + + def _base_lr_for(self, p: nn.Parameter) -> float | torch.Tensor: + """Per-group base learning rate for ``p`` (honors lr_multipliers).""" + return self._param_group_map[id(p)]["lr"] + + def _wd_for(self, p: nn.Parameter) -> float: + """Per-group weight decay for ``p`` (honors disable_weight_decay_for_1d_params).""" + return self._param_group_map[id(p)]["weight_decay"] + + def _get_adjusted_lr(self, param_shape: tuple[int, ...], base_lr: float | torch.Tensor) -> float | torch.Tensor: + """Compute adjusted learning rate based on parameter matrix size and the + owning param-group's base lr.""" + A, B = param_shape[:2] + adjusted_ratio = self.muon_lr_scale * math.sqrt(max(A, B)) + return base_lr * adjusted_ratio + + def _maybe_init_master_weights(self) -> None: + """Create FP32 master weights on first use (FusedAdam-style lazy init).""" + if self.master_weights and not self._masters_initialized: + self._create_master_weights() + + def _create_master_weights(self) -> None: + """ + Create FP32 master weight copies for mixed-precision training stability. + + Creates param_groups_master (for LowPrecisionCallback compatibility) and + indexed lists for efficient lookup during DION2/AdamW steps. + """ + # Create param_groups_master mirroring param_groups (like FusedAdam) + # This enables LowPrecisionCallback to copy masters -> params periodically + self.param_groups_master = [] + for pg in self.param_groups: + param_list = pg["params"] + self.param_groups_master.append( + { + "params": [p.clone().detach().float() if self.master_weights else None for p in param_list], + } + ) + + # Build param_id -> master mapping for efficient lookup + self._param_to_master: dict[int, torch.Tensor] = {} + for group, group_master in zip(self.param_groups, self.param_groups_master): + for p, p_master in zip(group["params"], group_master["params"]): + if p_master is not None: + self._param_to_master[id(p)] = p_master + + # Create indexed lists for DION2/AdamW step() methods + self._dion2_masters = [self._param_to_master[id(p)] for p in self.dion2_params] + self._adamw_masters = [self._param_to_master[id(p)] for p in self.adamw_params] + + muon_master_numel = sum(m.numel() for m in self._dion2_masters) + adamw_master_numel = sum(m.numel() for m in self._adamw_masters) + log.info( + f"Created FP32 master weights: {len(self._dion2_masters)} DION2 ({muon_master_numel:,} elements), " + f"{len(self._adamw_masters)} AdamW ({adamw_master_numel:,} elements)" + ) + self._masters_initialized = True + + @torch.no_grad() + def step(self, closure=None): + """Perform a single optimization step.""" + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + # Lazily build FP32 master weights from the current (possibly + # checkpoint-restored) params before the first update. + self._maybe_init_master_weights() + + # Params are split into three disjoint buckets at init (see + # categorize_params) and each is updated by a different function below: + # 1. self.dion2_params (dense 2D linears) -> _step_dion2 + # 2. self.stacked_dion2_params (3D MoE experts [E,M,N]) -> _step_stacked_dion2 + # 3. self.adamw_params (embeddings/head/norms/1D) -> _step_adamw + # Every parameter belongs to exactly one bucket, so exactly one of these + # updates it. Order does not matter (buckets are disjoint). + + # 1. Dense 2D linears: orthogonalized via all-to-all distributed Newton-Schulz. + self._step_dion2() + + # 2. MoE expert weights: each [E, M, N] param orthogonalized one expert + # slice at a time (local NS; per-expert masking for inactive experts). + self._step_stacked_dion2() + + # 3. Everything else (embeddings, lm_head, norms, biases, 1D): fused AdamW. + self._step_adamw() + + return loss + + def _step_stacked_dion2(self) -> None: + """Orthogonalize stacked MoE expert params, one expert slice at a time. + + Each param has shape ``[E, M, N]`` (E experts, each an M x N matrix). Under + FSDP2 these are sharded on the expert dim (dim 0), so every rank holds whole + expert matrices -- Newton-Schulz is therefore fully local (no all-to-all), + and is batched across the local experts via ``zeropower_via_newtonschulz5_batched``. + + NOTE (sharding assumption): the "no communication" property relies on the + expert tensor being sharded on dim 0 (the expert axis). This holds for the + FSDP2 ``fully_shard`` path used by LLM/VFM here, because FSDP2 shards every + parameter on dim 0. It is NOT guaranteed in general -- e.g. tensor/expert + parallelism could shard *within* an expert matrix (dim 1/2). That case is + unsupported and is rejected by the placement check below (fails loudly + rather than silently computing a wrong update); supporting it would require + a per-expert gather. The assumption was not exhaustively audited against + every parallelization config, which is exactly why it is enforced here. + """ + for p in self.stacked_dion2_params: + if p.grad is None: + continue + + # Validate sharding: only the expert axis (tensor dim 0) may be sharded. + if isinstance(p, DTensor): + for placement in p.placements: + if placement.is_shard() and placement.dim != 0: + raise NotImplementedError( + "Stacked-expert orthogonalization requires sharding on the expert " + f"dim (0); got placement {placement} for " + f"'{self.param_to_name.get(p, 'unknown')}'." + ) + + local_grad = get_local_tensor_if_DTensor(p.grad) + local_param = get_local_tensor_if_DTensor(p) + if local_grad.ndim != 3: + raise NotImplementedError( + f"Stacked-expert orthogonalization supports 3D params [E, M, N]; " + f"got shape {tuple(local_grad.shape)} for '{self.param_to_name.get(p, 'unknown')}'." + ) + + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + + # Per-expert masked momentum + Nesterov (element-wise over [E, M, N]). + # Active experts follow the standard mu*M + G recurrence; inactive + # experts (no gradient this step) keep their momentum frozen. ``active`` + # ([E] bool) is used below to zero the update and skip weight decay for + # inactive experts -- required because Newton-Schulz would otherwise + # renormalize their stale momentum into a full-strength spurious update. + pre_ns, active = compute_pre_ns_update_moe_expert( + local_grad, + get_local_tensor_if_DTensor(state["momentum_buffer"]), + momentum=self.muon_momentum, + nesterov=self.nesterov, + ) + + # Batched Newton-Schulz over the local experts, then zero the update for + # inactive experts (their NS result is a bogus unit-norm matrix). + ortho = zeropower_via_newtonschulz5_batched(pre_ns, steps=self.ns_steps) + ortho = ortho * active.view(-1, 1, 1).to(ortho.dtype) + + # LR scaling uses the per-expert matrix shape (M, N), shared across experts. + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(tuple(p.shape[-2:]), base_lr) + + # Per-expert weight-decay factor: (1 - base_lr*wd) for active experts, + # 1.0 (no decay) for inactive ones. Combined with the zeroed update + # above, inactive experts are left completely untouched. + if self.master_weights: + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) + a_wd = active.view(-1, 1, 1).to(master.dtype) + master.mul_(1 - a_wd * (base_lr * wd)) + master.add_(ortho.float() * (-adjusted_lr)) + local_param.copy_(master) + else: + a_wd = active.view(-1, 1, 1).to(local_param.dtype) + local_param.mul_(1 - a_wd * (base_lr * wd)) + local_param.add_(ortho.to(local_param.dtype) * (-adjusted_lr)) + + def _step_dion2(self) -> None: + """ + DION2 step with all-to-all distributed Newton-Schulz. + + For each batch of world_size params: + 1. Compute momentum + select submatrix on local shards + 2. All-to-all to redistribute shards (each rank gets full submatrix for its param) + 3. Newton-Schulz on full submatrix (each rank does different param) + 4. All-to-all to scatter results back + 5. Apply weight decay and update (to FP32 master if enabled) + """ + if not self.dion2_params: + return + + for batch in self._dion2_batches: + self._process_dion2_batch(batch) + + def _process_dion2_batch(self, batch: list[nn.Parameter]) -> None: + """Process a single batch of params using DION2 all-to-all pattern.""" + world_size = self._world_size + + # Pad batch if needed + actual_batch_size = len(batch) + if actual_batch_size < world_size: + # Pad with the last param (will be masked out) + padding = [batch[-1]] * (world_size - actual_batch_size) + batch = batch + padding + + # Check if using DTensor (FSDP) + is_dtensor = isinstance(batch[0], DTensor) + + if is_dtensor and world_size > 1: + self._process_dion2_batch_distributed(batch, actual_batch_size) + else: + self._process_dion2_batch_single(batch, actual_batch_size) + + def _process_dion2_batch_distributed(self, batch: list[nn.Parameter], actual_batch_size: int) -> None: + """Process a batch via DTensor collectives (the "each rank orthogonalizes one + whole param" transpose), correct for uneven / non-divisible shard dims. + + Rather than a hand-rolled ``all_to_all`` + ``cat``/``narrow`` (which assumed + FSDP2 padded every local shard uniformly -- it does not; ``_local_tensor`` is + unpadded and uneven), the gather/scatter is expressed through DTensor: + + 1. Momentum + Nesterov per param, kept as a DTensor so its shard metadata + (including uneven, unpadded local sizes) is preserved. + 2. ``torch.stack`` the world_size params -> a ``[W, ...]`` DTensor; the shard + tensor dim shifts to ``shard_dim + 1``. + 3. ``redistribute`` so the PARAM axis is sharded on the FSDP shard mesh dim + -> each rank owns one whole param (forward all-to-all). + 4. Newton-Schulz on that whole param (each rank a different one). + 5. ``redistribute`` back to shard the data axis (backward all-to-all), + unstack, and apply the update to each local shard. + + DTensor owns the (possibly uneven) per-rank size bookkeeping, so this is + correct regardless of divisibility. This transpose has been validated across + even/uneven/partial/bf16/shard-dim configs by the standalone multi-GPU script + ``unit_tests/dion2_uneven_shard_dtensor_check.py``. + + TODO(dion2-optionb-unittest): provide a standalone torchrun script into a + committed, CI-run multi-GPU unit test (it currently must be launched manually + via ``torchrun`` and is not part of the automated suite). + + Submatrix selection (``fraction < 1``) is not supported on this path -- it is + unused (every config runs fraction=1.0) -- and is rejected up front. + """ + # distributed path. Will need to wire it in is a + # well-scoped change: per-param DTensor select (norms via + # ``pre.abs().sum(sharded_dim).full_tensor()`` -> top-k -> ``index_select`` + # with a Replicate index), error-feedback decay on the momentum buffer, and a + # branched apply that reuses ``_apply_submatrix_update[_master]`` (index_add + # into the selected indices). Deferred for now: every config runs fraction=1.0, + # so this path is unused and not worth the added complexity yet. Single-device + # fraction<1 still works via ``_process_dion2_batch_single``. + if self.fraction != 1.0: + raise NotImplementedError( + "DION2 distributed path supports only fraction=1.0 (full-matrix " + f"orthogonalization); got fraction={self.fraction}. Submatrix selection " + "under FSDP is not implemented yet (see TODO(dion2-fsdp-fraction))." + ) + + world_size = self._world_size + mesh = self._device_mesh + shard_mesh_dim = self._shard_mesh_dim + shard_dim = self._shard_tensor_dim + stack_axis = shard_dim + 1 # torch.stack adds a leading param axis + + # Step 1: momentum + Nesterov, kept in DTensor space (metadata preserved). + # compute_pre_ns_update does not mutate the gradient, so p.grad is passed + # directly; it mutates the (sharded DTensor) momentum buffer in place. + # + # None-grad handling: a param with no gradient this step is sat out -- + # momentum frozen (compute_pre_ns_update NOT called, so no mu-decay) and no + # update applied (skipped in the apply loop via ``active``). We cannot just + # drop the slot: the stack + redistribute all-to-alls are a fixed-size + # collective every rank must enter identically, so an inactive slot instead + # contributes a zero placeholder (same DTensor sharding/dtype) to keep the + # collective shapes uniform. Newton-Schulz on zeros stays finite (norm+1e-7) + # and the result is discarded on apply. This relies on ``p.grad is None`` + # being identical across ranks -- true for dense params, where a missing grad + # is structural (an unused param is None on every rank), not data-dependent. + pre_ns_list = [] + active = [] + for i in range(actual_batch_size): + p = batch[i] + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + if p.grad is None: + active.append(False) + pre_ns_list.append(torch.zeros_like(state["momentum_buffer"]).to(torch.bfloat16)) + continue + active.append(True) + pre_ns = compute_pre_ns_update( + p.grad, state["momentum_buffer"], momentum=self.muon_momentum, nesterov=self.nesterov + ) + # Cast to bf16 BEFORE the transpose so the two all-to-alls move bf16 (not + # fp32) -- matching the previous path's communication volume. This is + # numerically identical: Newton-Schulz casts to bf16 anyway, and bf16 + # rounding commutes with the reconstruction (round-then-gather == + # gather-then-round elementwise). + pre_ns_list.append(pre_ns.to(torch.bfloat16)) + + # Pad the param axis up to world_size (mirrors batch padding); the dummy + # entries' Newton-Schulz results are discarded on apply. + padded = pre_ns_list + [pre_ns_list[-1]] * (world_size - actual_batch_size) + + # Step 2: stack -> [W, ...]; the shard tensor dim moves to stack_axis. + stacked = torch.stack(padded, dim=0) + expected = Shard(stack_axis) + if stacked.placements[shard_mesh_dim] != expected: + raise RuntimeError( + f"DION2 expected stacked placement {expected} on mesh dim {shard_mesh_dim}, " + f"got {stacked.placements} (torch.stack should shift Shard({shard_dim}) -> " + f"Shard({stack_axis}))." + ) + + # Step 3: forward all-to-all -- shard the PARAM axis on the FSDP shard mesh dim + # (keep any other mesh-dim placements, e.g. Replicate under HSDP). Each rank + # then owns one whole param. + fwd_placements = list(stacked.placements) + fwd_placements[shard_mesh_dim] = Shard(0) + per_matrix = stacked.redistribute(mesh, fwd_placements) + full_p = per_matrix.to_local()[0] # this rank's whole param + + # Step 4: Newton-Schulz on the whole param. + ortho_p = zeropower_via_newtonschulz5(full_p, steps=self.ns_steps) + + # Step 5: backward all-to-all -- re-shard the data axis, then unstack. + ortho_dt = DTensor.from_local(ortho_p.unsqueeze(0), mesh, fwd_placements, run_check=False) + back = ortho_dt.redistribute(mesh, list(stacked.placements)) + back_local = back.to_local() # [W, , ...] + + # Apply the orthogonalized update to each real param's local shard. + for i in range(actual_batch_size): + if not active[i]: + # No gradient this step: momentum was frozen above; skip weight + # decay and the update entirely (matches the reference, which + # leaves a None-grad param completely untouched). + continue + p = batch[i] + local_param = p._local_tensor + update_local = back_local[i] + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(tuple(p.shape), base_lr) + if self.master_weights: + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) + master.mul_(1 - base_lr * wd) + master.add_(update_local.float() * (-adjusted_lr)) + local_param.copy_(master) + else: + local_param.mul_(1 - base_lr * wd) + local_param.add_(update_local.to(local_param.dtype) * (-adjusted_lr)) + + def _process_dion2_batch_single(self, batch: list[nn.Parameter], actual_batch_size: int) -> None: + """Process batch on single device (no distribution).""" + for i, p in enumerate(batch): + if i >= actual_batch_size: + continue + + # No gradient this step -> sit the param out entirely: momentum buffer + # stays frozen (not created/decayed) and no update is applied. Newton- + # Schulz renormalizes any nonzero input back to unit norm, so feeding a + # stale/zero momentum would emit a full-strength spurious update; hence + # we skip rather than treat a missing grad as grad=0. Matches the + # reference (microsoft/dion), which filters None-grad params up front. + if p.grad is None: + continue + + grad = get_local_tensor_if_DTensor(p.grad) + param = get_local_tensor_if_DTensor(p) + + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + + mom = get_local_tensor_if_DTensor(state["momentum_buffer"]) + + # Momentum + submatrix selection. Single-device fallback is not + # sharded; pass shard_dim=1 so the selection dim is rows (-2), + # matching the apply call below. (Fixes the original ``select_dim=`` + # keyword-name bug, which raised TypeError on any single-GPU / + # non-DTensor run.) + if self.fraction == 1.0: + # Muon baseline: pre-decayed momentum (M <- mu*M + G) + Nesterov, + # full-matrix NS. compute_pre_ns_update does not mutate the + # gradient, so grad can be passed directly. + pre_ns = compute_pre_ns_update( + grad, + mom, + momentum=self.muon_momentum, + nesterov=self.nesterov, + ) + submatrix, indices = self._select_submatrix(pre_ns, state, shard_dim=1) + else: + # Dion2 Algorithm 1 (fractional): pure error-feedback accumulation + # with selective decay. No whole-buffer decay and no Nesterov -- NS + # runs on the accumulated M[K], and only the selected slice is + # decayed (by ef_decay) inside _select_submatrix. This keeps the + # unselected rows as a running residual, matching the reference. + mom.add_(grad) + submatrix, indices = self._select_submatrix(mom, state, shard_dim=1) + + # Newton-Schulz + ortho = zeropower_via_newtonschulz5(submatrix, steps=self.ns_steps) + + # Get adjusted LR / wd from the owning param-group. + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(p.shape, base_lr) + + if self.master_weights: + # Update FP32 master, then write the param directly inside + # _apply_submatrix_update_master. + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) + master.mul_(1 - base_lr * wd) + self._apply_submatrix_update_master(master, param, ortho, indices, adjusted_lr, select_dim=-2) + else: + # Apply weight decay + param.mul_(1 - base_lr * wd) + # Apply update to selected indices + self._apply_submatrix_update(param, ortho, indices, adjusted_lr, select_dim=-2) + + def _select_submatrix( + self, + tensor: torch.Tensor, + state: dict, + shard_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Select submatrix based on L1 norm (DION2 style). + + Args: + tensor: Input tensor (local shard or full matrix) + state: Optimizer state dict + shard_dim: Dimension along which tensor is sharded (for FSDP) + + Returns: + submatrix: Selected rows/columns + indices: Indices of selected rows/columns + """ + if self.fraction == 1.0: + # Full matrix, no selection + if tensor.ndim == 2: + indices = torch.arange(tensor.size(0), device=tensor.device) + else: + indices = None + # Convert to BF16 for all_to_all compatibility (tensor may be FP32 from momentum) + return tensor.to(torch.bfloat16), indices + + # Determine selection dimension (opposite of shard dim for efficiency) + # If sharded along rows, select columns; if sharded along cols, select rows + if shard_dim == 0: + select_dim = -1 # Select columns + norm_dim = -2 # Compute norm over rows + else: + select_dim = -2 # Select rows + norm_dim = -1 # Compute norm over columns + + num_select = tensor.size(select_dim) + k = max(1, int(math.ceil(self.fraction * num_select))) + + # Compute L1 norm along norm_dim + slice_norms = tensor.abs().sum(dim=norm_dim) + + # All-reduce norms across ranks so all ranks select the same indices + # This is critical for FSDP where each rank has different rows/cols + # The all-reduce sums the partial norms to get global norms + if self._process_group is not None: + dist.all_reduce(slice_norms, group=self._process_group) + + # Top-k selection (now deterministic across all ranks) + _, indices = torch.topk(slice_norms, k, dim=-1, sorted=False) + + # Extract selected submatrix + if select_dim == -2: + submatrix = tensor.index_select(dim=0, index=indices) + else: + submatrix = tensor.index_select(dim=1, index=indices) + + # Apply error feedback decay to the selected rows/cols of the momentum buffer. + # Operate on the LOCAL shard (not the DTensor wrapper): ``indices`` are + # computed from the local pre-NS slice, and the selection dimension is the + # non-sharded one (opposite of shard_dim), so local and global indices + # coincide on that axis. This keeps the buffer in the same (local) coordinate + # frame as the update applied later, and avoids an unsupported in-place + # index_copy_ on a DTensor. + if "momentum_buffer" in state and self.ef_decay < 1.0: + momentum = get_local_tensor_if_DTensor(state["momentum_buffer"]) + dim = 0 if select_dim == -2 else 1 + selected = momentum.index_select(dim=dim, index=indices) + momentum.index_copy_(dim=dim, index=indices, source=selected * self.ef_decay) + + return submatrix.to(torch.bfloat16), indices + + def _apply_submatrix_update( + self, + param: torch.Tensor, + ortho: torch.Tensor, + indices: torch.Tensor | None, + lr: float | torch.Tensor, + select_dim: int, + ) -> None: + """Apply orthogonalized update to selected indices.""" + ortho = ortho.to(param.dtype) + + if indices is None or self.fraction == 1.0: + # Full matrix update. lr may be a tensor (capturable), so scale via + # multiply rather than ``alpha=`` (Tensor.add_ only accepts a Number). + param.add_(ortho * (-lr)) + else: + # Submatrix update at selected indices + scaled_ortho = -lr * ortho + if select_dim == -2 or select_dim == 0: + param.index_add_(dim=0, index=indices, source=scaled_ortho) + else: + param.index_add_(dim=1, index=indices, source=scaled_ortho) + + def _apply_submatrix_update_master( + self, + master: torch.Tensor, + param: torch.Tensor, + ortho: torch.Tensor, + indices: torch.Tensor | None, + lr: float | torch.Tensor, + select_dim: int, + ) -> None: + """Apply orthogonalized update to FP32 master, then copy to BF16 param.""" + ortho = ortho.float() + + if indices is None or self.fraction == 1.0: + # Full matrix update. lr may be a tensor (capturable), so scale via + # multiply rather than ``alpha=`` (Tensor.add_ only accepts a Number). + master.add_(ortho * (-lr)) + else: + # Submatrix update at selected indices + scaled_ortho = -lr * ortho + if select_dim == -2 or select_dim == 0: + master.index_add_(dim=0, index=indices, source=scaled_ortho) + else: + master.index_add_(dim=1, index=indices, source=scaled_ortho) + # Write the BF16 param directly from the updated FP32 master so we do not + # depend on LowPrecisionCallback (the OptimizersContainer hides + # master_weights from it). Matches FusedAdam's in-kernel param write. + param.copy_(master) + + def _step_adamw(self) -> None: + """ + AdamW step using Transformer Engine's fused kernel. + + Iterates over param groups so each group's lr / betas / eps / weight_decay + (set by the factory's lr_multipliers and disable_weight_decay_for_1d_params) + is honored, then batches by dtype within the group. The per-group step + counter lives on ``group["step"]`` (FusedAdam-style) so it is + round-tripped by the distributed-checkpoint optimizer state dict. + + When master_weights=True and capturable=True, uses + multi_tensor_adam_capturable_master which maintains FP32 master weights. + """ + if not self.adamw_params: + return + + adam_w_mode = 1 + bias_correction = 1 + + for group in self.param_groups: + # Only the AdamW-categorized params of this group are handled here; + # the DION2-categorized params were already updated in _step_dion2. + group_params = [p for p in group["params"] if id(p) in self._adamw_param_ids] + if not group_params: + continue + + device = group_params[0].device + + # Per-group step counter stored on the group (FusedAdam convention) so + # DCP round-trips it on resume. + if group.get("step", None) is not None: + if self.capturable and not isinstance(group["step"], torch.Tensor): + group["step"] = torch.tensor(group["step"], dtype=torch.int32, device=device) + group["step"] = ( + group["step"].to(device=device) if isinstance(group["step"], torch.Tensor) else group["step"] + ) + group["step"] += 1 + else: + group["step"] = torch.tensor([1], dtype=torch.int32, device=device) if self.capturable else 1 + + if self.capturable and not isinstance(group["lr"], torch.Tensor): + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32, device=device) + + lr = group["lr"] + wd = group["weight_decay"] + step = group["step"] + beta1, beta2 = group["betas"] + eps = group["eps"] + + g_16, p_16, m_16, v_16 = [], [], [], [] + g_bf, p_bf, m_bf, v_bf = [], [], [], [] + g_32, p_32, m_32, v_32 = [], [], [], [] + p_16_master, p_bf_master, p_32_master = [], [], [] + + for p in group_params: + if p.grad is None: + continue + + grad = get_local_tensor_if_DTensor(p.grad) + param = get_local_tensor_if_DTensor(p) + + state = self.state[p] + if len(state) == 0: + state["exp_avg"] = torch.zeros_like(p).float() + state["exp_avg_sq"] = torch.zeros_like(p).float() + + exp_avg = get_local_tensor_if_DTensor(state["exp_avg"]) + exp_avg_sq = get_local_tensor_if_DTensor(state["exp_avg_sq"]) + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) if self.master_weights else None + + if p.dtype == torch.float16: + g_16.append(grad) + p_16.append(param) + m_16.append(exp_avg) + v_16.append(exp_avg_sq) + if self.master_weights: + p_16_master.append(master) + elif p.dtype == torch.bfloat16: + g_bf.append(grad) + p_bf.append(param) + m_bf.append(exp_avg) + v_bf.append(exp_avg_sq) + if self.master_weights: + p_bf_master.append(master) + elif p.dtype == torch.float32: + g_32.append(grad) + p_32.append(param) + m_32.append(exp_avg) + v_32.append(exp_avg_sq) + if self.master_weights: + p_32_master.append(master) + else: + raise RuntimeError(f"Unsupported dtype {p.dtype}") + + if self.capturable: + # The capturable-master kernel requires an inverse-scale argument; + # bf16-only training has no grad scaler, so it is a constant one. + kernel_inv_scale = torch.ones((1,), device=device, dtype=torch.float32) + dtype_batches = ( + (g_16, p_16, m_16, v_16, p_16_master), + (g_bf, p_bf, m_bf, v_bf, p_bf_master), + (g_32, p_32, m_32, v_32, p_32_master), + ) + kernel = ( + self._multi_tensor_adam_capturable_master + if self.master_weights + else self._multi_tensor_adam_capturable + ) + for g_, p_, m_, v_, pm_ in dtype_batches: + if len(g_) == 0: + continue + tensor_lists = [g_, p_, m_, v_, pm_] if self.master_weights else [g_, p_, m_, v_] + te.pytorch.optimizers.multi_tensor_applier( + kernel, + self._dummy_overflow_buf, + tensor_lists, + lr, + beta1, + beta2, + eps, + step, + adam_w_mode, + bias_correction, + wd, + kernel_inv_scale, + ) + else: + dtype_batches = ( + (g_16, p_16, m_16, v_16), + (g_bf, p_bf, m_bf, v_bf), + (g_32, p_32, m_32, v_32), + ) + for g_, p_, m_, v_ in dtype_batches: + if len(g_) == 0: + continue + te.pytorch.optimizers.multi_tensor_applier( + self._multi_tensor_adam, + self._dummy_overflow_buf, + [g_, p_, m_, v_], + lr, + beta1, + beta2, + eps, + step, + adam_w_mode, + bias_correction, + wd, + ) + + def load_state_dict(self, state_dict: dict) -> None: + """Load optimizer state. + + The optimizer state (per-param momentum / exp_avg / exp_avg_sq and the + per-group ``step``) round-trips through the base ``torch.optim.Optimizer`` + state dict, so the distributed-checkpoint container can save/restore it + with FSDP2 resharding just like FusedAdam. Master weights are *not* + checkpointed; they are rebuilt from the restored params on the next + ``step()`` (see ``_maybe_init_master_weights``). + """ + super().load_state_dict(state_dict) + + # Force master weights to be rebuilt from the (now restored) params. + self._masters_initialized = False + self.param_groups_master = None + + for group in self.param_groups: + device = group["params"][0].device if group["params"] else "cuda" + if self.capturable: + if isinstance(group["lr"], torch.Tensor): + group["lr"] = group["lr"].to(device=device) + else: + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32, device=device) + if group.get("step", None) is not None and not isinstance(group["step"], torch.Tensor): + group["step"] = torch.tensor(group["step"], dtype=torch.int32, device=device) + for p in group["params"]: + state = self.state[p] + if "exp_avg" in state: + state["exp_avg"] = state["exp_avg"].float() + state["exp_avg_sq"] = state["exp_avg_sq"].float() + if "momentum_buffer" in state: + state["momentum_buffer"] = state["momentum_buffer"].float() diff --git a/cosmos_framework/utils/generator/muon_with_aux_adamw.py b/cosmos_framework/utils/generator/muon_with_aux_adamw.py new file mode 100644 index 00000000..6a05e21d --- /dev/null +++ b/cosmos_framework/utils/generator/muon_with_aux_adamw.py @@ -0,0 +1,835 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +""" +MuonWithAuxAdamW optimizer implementation. + +Muon (MomentUm Orthogonalized by Newton-schulz) for nn.Linear weight matrices, +with auxiliary AdamW for embeddings, biases, norms, and output layers (lm_head). + +This implementation combines elements from: + +1. FusedAdam (cosmos_framework/utils/generator/fused_adam.py): + - DTensor handling via `get_local_tensor_if_DTensor` for FSDP/TP compatibility + - Distributed step synchronization in `load_state_dict` for checkpoint compatibility + +2. KellerJordan/Muon (https://github.com/KellerJordan/Muon): + - Newton-Schulz orthogonalization algorithm (`zeropower_via_newtonschulz5`) + - Quintic iteration coefficients (a=3.4445, b=-4.7750, c=2.0315) + - Nesterov momentum formulation + +3. MoonshotAI/Moonlight (https://arxiv.org/pdf/2502.16982): + - SGD-style momentum: `buf = momentum * buf + grad` (not EMA-style) + - Learning rate scaling by matrix size: `adjusted_lr = 0.2 * sqrt(max(A, B)) * lr` + - `@torch.compile` decoration for kernel fusion + - Parameter separation: Muon for nn.Linear weights only, AdamW for everything else + - Distributed Newton-Schulz: all-gather gradients, NS on full matrix, scatter back + (required because NS is a matrix-level operation, not element-wise) + +Sharding -> orthogonalization (the core idea) +--------------------------------------------- +Under FSDP each weight matrix is a DTensor split across GPUs, but Newton-Schulz +(NS) is a *whole-matrix* op and needs the full matrix in one place. Muon handles +each matrix one at a time with an all-gather: every GPU rebuilds the *same* full +matrix and runs NS on it. + +Weight A sharded over 4 GPUs (one row-slice each): + + GPU0:[A0] GPU1:[A1] GPU2:[A2] GPU3:[A3] + + --- all_gather(A) ---> every GPU holds the full matrix: + + GPU0:[A0 A1 A2 A3] GPU1:[A0 A1 A2 A3] GPU2:[...] GPU3:[...] + | | | | + NS(A) NS(A) NS(A) NS(A) (all identical) + | | | | + --- each GPU slices back its own rows and applies the update to its shard --- + +This is simple and correct, but the NS compute is duplicated world_size times +(every rank orthogonalizes the same matrix). The sibling ``Dion2WithAuxAdamW`` +removes that redundancy by trading shards with all-to-all so each GPU +orthogonalizes a *different* matrix in parallel (see its module docstring). +""" + +import math + +import torch +import torch.distributed as dist +import torch.nn as nn +import transformer_engine as te +import transformer_engine_torch as tex +from torch.distributed.tensor import DTensor, Replicate + +from cosmos_framework.utils import log +from cosmos_framework.utils.misc import get_local_tensor_if_DTensor +from cosmos_framework.utils.generator.aux_optimizer_utils import ( + DEFAULT_ADAMW_MODULE_KEYWORDS, + compute_pre_ns_update, + compute_pre_ns_update_moe_expert, + split_orthogonalizable_params, + zeropower_via_newtonschulz5, + zeropower_via_newtonschulz5_batched, +) + + +class MuonWithAuxAdamW(torch.optim.Optimizer): + """ + MuonWithAuxAdamW optimizer. + + Uses Muon (MomentUm Orthogonalized by Newton-schulz) for nn.Linear hidden weight matrices, + and AdamW for embeddings, biases, layer norms, and output heads (lm_head). + + See module docstring for full attribution of borrowed components. + + Args: + params: Iterable of parameters to optimize. + lr: Base learning rate. Muon scales this by muon_lr_scale*sqrt(max(A,B)), AdamW uses directly. + muon_momentum: Momentum coefficient for Muon. + muon_lr_scale: Scale factor for Muon LR adjustment. Final LR = muon_lr_scale * sqrt(max(A,B)) * lr. + ns_steps: Number of Newton-Schulz iterations. + nesterov: Whether to use Nesterov momentum for Muon. + weight_decay: Weight decay for all parameters. + adam_betas: Beta coefficients for the auxiliary AdamW side (matches VFM convention). + eps: Epsilon for AdamW numerical stability. + use_distributed: Whether to sync step counters across ranks when loading checkpoints. + """ + + def __init__( + self, + params, + lr: float = 1e-4, + muon_momentum: float = 0.95, + muon_lr_scale: float = 0.2, + ns_steps: int = 5, + nesterov: bool = True, + weight_decay: float = 0.1, + adam_betas: tuple[float, float] = (0.9, 0.95), + eps: float = 1e-8, + use_distributed: bool = True, + capturable: bool = False, + master_weights: bool = False, + adamw_module_keywords: tuple[str, ...] | None = None, + expert_param_keywords: tuple[str, ...] | None = None, + **kwargs, # Absorb VFM-specific args (fused, keys_to_select, etc.) + ): + # Log any ignored kwargs for debugging + if kwargs: + ignored_keys = list(kwargs.keys()) + # These are expected VFM args that we silently ignore + expected_ignored = {"fused", "keys_to_select", "adamw_betas", "adamw_eps"} + unexpected = set(ignored_keys) - expected_ignored + if unexpected: + import warnings + + warnings.warn(f"MuonWithAuxAdamW ignoring unexpected kwargs: {unexpected}") + + # Master weights requires capturable mode + if master_weights and not capturable: + raise RuntimeError("Master weights is currently only supported with capturable=True.") + + # Store shared hyperparameters + # Note: lr is accessed via property that reads from param_groups + # to support LR schedulers (which update param_groups[X]["lr"]) + self.wd = weight_decay + + # Store Muon-specific hyperparameters + self.muon_momentum = muon_momentum + self.muon_lr_scale = muon_lr_scale + self.ns_steps = ns_steps + self.nesterov = nesterov + + # Name substrings that route an nn.Linear to AdamW (output heads). Used by + # categorize_params alongside tied-weight / vocab-shape detection. + self.adamw_module_keywords = ( + tuple(adamw_module_keywords) if adamw_module_keywords else DEFAULT_ADAMW_MODULE_KEYWORDS + ) + # Name substrings that route stacked MoE expert params ([E, M, N]) to the + # Muon side (each expert slice orthogonalized). Empty = experts stay on + # AdamW (no behavior change). + self.expert_param_keywords = tuple(expert_param_keywords) if expert_param_keywords else () + + # Store AdamW-specific hyperparameters + self.adam_betas = tuple(adam_betas) if isinstance(adam_betas, list) else adam_betas + self.eps = eps + + # Distributed settings + self.use_distributed = use_distributed and dist.is_initialized() + + # Master weights settings (for mixed-precision training stability) + self.capturable = capturable + self.master_weights = master_weights + + # Parameter lists (populated by categorize_params) + self.muon_params: list[nn.Parameter] = [] + self.adamw_params: list[nn.Parameter] = [] + # Stacked MoE expert params ([E, M, N]); orthogonalized per expert slice. + self.stacked_muon_params: list[nn.Parameter] = [] + self.param_to_name: dict[nn.Parameter, str] = {} + + # Master weight copies (populated by _create_master_weights after categorize_params) + self._muon_masters: list[torch.Tensor] = [] + self._adamw_masters: list[torch.Tensor] = [] + + # Transformer Engine fused Adam for AdamW params. The zero buffer is the + # noop flag required as the second argument of TE's multi_tensor_applier; + # it is a fixed constant here (no AMP overflow handling). + self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device="cuda") + self._multi_tensor_adam = tex.multi_tensor_adam + self._multi_tensor_adam_capturable = tex.multi_tensor_adam_capturable + self._multi_tensor_adam_capturable_master = tex.multi_tensor_adam_capturable_master + + # Initialize base optimizer. betas / eps go in the defaults so each param + # group carries them (FusedAdam convention); the AdamW step reads them + # per-group, enabling per-group overrides and exact FusedAdam parity. + defaults = dict(lr=lr, weight_decay=weight_decay, betas=self.adam_betas, eps=eps) + super().__init__(params, defaults) + + # Convert LR to tensor for capturable mode + if capturable: + for idx, group in enumerate(self.param_groups): + if len(group["params"]) == 0: + continue + device = group["params"][0].device + if isinstance(group["lr"], float): + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32) + self.param_groups[idx]["lr"] = group["lr"].to(device=device) + + # id(param) -> owning param_group, so the Muon and AdamW updates can read + # the *per-group* lr / weight_decay. This is what lets the factory's + # lr_multipliers and disable_weight_decay_for_1d_params flow through. With + # a single param group it degenerates to a single global lr/wd, matching + # the original (reference) behavior. Populated by categorize_params. + self._param_group_map: dict[int, dict] = {} + self._adamw_param_ids: set[int] = set() + # Master weights are created lazily on the first step() (FusedAdam-style), + # so that after a checkpoint resume they are rebuilt from the *restored* + # params rather than the freshly-initialized ones. + self._masters_initialized = False + self._param_to_master: dict[int, torch.Tensor] = {} + + log.info(f"MuonWithAuxAdamW master_weights: {master_weights} capturable: {capturable}") + + def categorize_params(self, model: nn.Module) -> None: + """ + Categorize parameters into Muon and AdamW groups. + + Muon is used for hidden ``nn.Linear`` weights only. Embeddings, the output + head (``lm_head`` / tied / vocab-shaped projection), biases, norms, and any + non-Linear parameter go to AdamW. See + :func:`split_orthogonalizable_params` for the architecture-agnostic + embedding / output-head detection. + + Args: + model: The model (typically the trainable ``net``) to categorize. + """ + # Only categorize params that are in this optimizer's param_groups. This + # respects keys_to_select filtering - params not passed to __init__ should + # not be categorized, otherwise state_dict() would fail with KeyError when + # mapping state entries to param indices. + optimizer_param_ids = {id(p) for group in self.param_groups for p in group["params"]} + + orthogonalizable, self.adamw_params, self.param_to_name = split_orthogonalizable_params( + model, + optimizer_param_ids, + adamw_module_keywords=self.adamw_module_keywords, + expert_param_keywords=self.expert_param_keywords, + ) + + # Split orthogonalizable params into 2-D (regular Linear weights -> Muon) and + # stacked 3-D MoE experts ([E, M, N] -> orthogonalized per expert slice). + self.muon_params = [p for p in orthogonalizable if p.ndim == 2] + self.stacked_muon_params = [p for p in orthogonalizable if p.ndim >= 3] + + # Sort Muon params by size (largest first) for distributed load balancing + self.muon_params = sorted(self.muon_params, key=lambda x: x.numel(), reverse=True) + + # Check if using DTensor (FSDP) - determines whether we use distributed NS + self._params_are_dtensor = isinstance(self.muon_params[0], DTensor) if self.muon_params else False + + muon_numel = sum(p.numel() for p in self.muon_params) + adamw_numel = sum(p.numel() for p in self.adamw_params) + stacked_muon_numel = sum(p.numel() for p in self.stacked_muon_params) + + log.info( + f"MuonWithAuxAdamW: {len(self.muon_params)} Muon params ({muon_numel:,} elements), " + f"{len(self.stacked_muon_params)} stacked-expert params ({stacked_muon_numel:,} elements), " + f"{len(self.adamw_params)} AdamW params ({adamw_numel:,} elements)" + f"{', using distributed NS (DTensor/FSDP detected)' if self._params_are_dtensor else ''}" + ) + + # Log Muon param details (layer names and shapes) + log.info("Muon parameters (layer name -> shape):") + for p in self.muon_params: + name = self.param_to_name.get(p, "unknown") + log.info(f" {name}: {tuple(p.shape)}") + + # Build the param -> owning group map for per-group lr / weight_decay + # lookups during the Muon and AdamW steps. + self._param_group_map = {} + for group in self.param_groups: + for p in group["params"]: + self._param_group_map[id(p)] = group + self._adamw_param_ids = {id(p) for p in self.adamw_params} + + # NOTE: master weights are intentionally NOT created here. They are + # created lazily on the first step() (see _maybe_init_master_weights) so + # that a checkpoint resume rebuilds them from the restored params. + + def _base_lr_for(self, p: nn.Parameter) -> float | torch.Tensor: + """Per-group base learning rate for ``p`` (honors lr_multipliers).""" + return self._param_group_map[id(p)]["lr"] + + def _wd_for(self, p: nn.Parameter) -> float: + """Per-group weight decay for ``p`` (honors disable_weight_decay_for_1d_params).""" + return self._param_group_map[id(p)]["weight_decay"] + + def _get_adjusted_lr(self, param_shape: tuple[int, ...], base_lr: float | torch.Tensor) -> float | torch.Tensor: + """ + Compute adjusted learning rate based on parameter matrix size. + + Based on Moonlight: adjusted_lr = muon_lr_scale * sqrt(max(A, B)) * base_lr + + Args: + param_shape: Shape of the parameter tensor. + base_lr: The owning param-group's learning rate (after any + lr_multiplier). Muon's matrix-size scaling layers on top of it. + + Returns: + Adjusted learning rate for this parameter. + """ + A, B = param_shape[:2] + adjusted_ratio = self.muon_lr_scale * math.sqrt(max(A, B)) + return base_lr * adjusted_ratio + + def _maybe_init_master_weights(self) -> None: + """Create FP32 master weights on first use (FusedAdam-style lazy init).""" + if self.master_weights and not self._masters_initialized: + self._create_master_weights() + + def _create_master_weights(self) -> None: + """ + Create FP32 master weight copies for mixed-precision training stability. + + Creates param_groups_master (for LowPrecisionCallback compatibility) and + indexed lists for efficient lookup during Muon/AdamW steps. + """ + # Create param_groups_master mirroring param_groups (like FusedAdam) + # This enables LowPrecisionCallback to copy masters -> params periodically + self.param_groups_master = [] + for pg in self.param_groups: + param_list = pg["params"] + self.param_groups_master.append( + { + "params": [p.clone().detach().float() if self.master_weights else None for p in param_list], + } + ) + + # Build param_id -> master mapping for efficient lookup + self._param_to_master: dict[int, torch.Tensor] = {} + for group, group_master in zip(self.param_groups, self.param_groups_master): + for p, p_master in zip(group["params"], group_master["params"]): + if p_master is not None: + self._param_to_master[id(p)] = p_master + + # Create indexed lists for Muon/AdamW step() methods + self._muon_masters = [self._param_to_master[id(p)] for p in self.muon_params] + self._adamw_masters = [self._param_to_master[id(p)] for p in self.adamw_params] + + muon_master_numel = sum(m.numel() for m in self._muon_masters) + adamw_master_numel = sum(m.numel() for m in self._adamw_masters) + log.info( + f"Created FP32 master weights: {len(self._muon_masters)} Muon ({muon_master_numel:,} elements), " + f"{len(self._adamw_masters)} AdamW ({adamw_master_numel:,} elements)" + ) + self._masters_initialized = True + + @torch.no_grad() + def step(self, closure=None): + """ + Perform a single optimization step. + + Args: + closure: A closure that reevaluates the model and returns the loss. + + Returns: + Loss value if closure is provided, else None. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + # Lazily build FP32 master weights from the current (possibly + # checkpoint-restored) params before the first update. + self._maybe_init_master_weights() + + # Muon updates (with distributed NS for FSDP) + self._step_muon() + + # Stacked MoE expert updates (orthogonalize each expert slice) + self._step_stacked_muon() + + # AdamW updates + self._step_adamw() + + return loss + + def _step_stacked_muon(self) -> None: + """Orthogonalize stacked MoE expert params, one expert slice at a time. + + Each param has shape ``[E, M, N]`` (E experts, each an M x N matrix). Under + FSDP2 these are sharded on the expert dim (dim 0), so every rank holds whole + expert matrices -- Newton-Schulz is therefore fully local (no all-gather), + and is batched across the local experts via ``zeropower_via_newtonschulz5_batched``. + + NOTE (sharding assumption): the "no communication" property relies on the + expert tensor being sharded on dim 0 (the expert axis). This holds for the + FSDP2 ``fully_shard`` path used by LLM/VFM here, because FSDP2 shards every + parameter on dim 0. It is NOT guaranteed in general -- e.g. tensor/expert + parallelism could shard *within* an expert matrix (dim 1/2). That case is + unsupported and is rejected by the placement check below (fails loudly + rather than silently computing a wrong update); supporting it would require + a per-expert gather. The assumption was not exhaustively audited against + every parallelization config, which is exactly why it is enforced here. + """ + for p in self.stacked_muon_params: + if p.grad is None: + continue + + # Validate sharding: only the expert axis (tensor dim 0) may be sharded. + if isinstance(p, DTensor): + for placement in p.placements: + if placement.is_shard() and placement.dim != 0: + raise NotImplementedError( + "Stacked-expert orthogonalization requires sharding on the expert " + f"dim (0); got placement {placement} for " + f"'{self.param_to_name.get(p, 'unknown')}'." + ) + + local_grad = get_local_tensor_if_DTensor(p.grad) + local_param = get_local_tensor_if_DTensor(p) + if local_grad.ndim != 3: + raise NotImplementedError( + f"Stacked-expert orthogonalization supports 3D params [E, M, N]; " + f"got shape {tuple(local_grad.shape)} for '{self.param_to_name.get(p, 'unknown')}'." + ) + + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + + # Per-expert masked momentum + Nesterov (element-wise over [E, M, N]). + # Active experts follow the standard mu*M + G recurrence; inactive + # experts (no gradient this step) keep their momentum frozen. ``active`` + # ([E] bool) is used below to zero the update and skip weight decay for + # inactive experts -- required because Newton-Schulz would otherwise + # renormalize their stale momentum into a full-strength spurious update. + pre_ns, active = compute_pre_ns_update_moe_expert( + local_grad, + get_local_tensor_if_DTensor(state["momentum_buffer"]), + momentum=self.muon_momentum, + nesterov=self.nesterov, + ) + + # Batched Newton-Schulz over the local experts, then zero the update for + # inactive experts (their NS result is a bogus unit-norm matrix). + ortho = zeropower_via_newtonschulz5_batched(pre_ns, steps=self.ns_steps) + ortho = ortho * active.view(-1, 1, 1).to(ortho.dtype) + + # LR scaling uses the per-expert matrix shape (M, N), shared across experts. + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(tuple(p.shape[-2:]), base_lr) + + # Per-expert weight-decay factor: (1 - base_lr*wd) for active experts, + # 1.0 (no decay) for inactive ones. Combined with the zeroed update + # above, inactive experts are left completely untouched. + if self.master_weights: + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) + a_wd = active.view(-1, 1, 1).to(master.dtype) + master.mul_(1 - a_wd * (base_lr * wd)) + master.add_(ortho.float() * (-adjusted_lr)) + local_param.copy_(master) + else: + a_wd = active.view(-1, 1, 1).to(local_param.dtype) + local_param.mul_(1 - a_wd * (base_lr * wd)) + local_param.add_(ortho.to(local_param.dtype) * (-adjusted_lr)) + + def _step_muon(self) -> None: + """ + Muon step with distributed Newton-Schulz (following Moonlight paper). + + For DTensor/FSDP parameters: + 1. Apply momentum + Nesterov on local shards (element-wise, shard-safe) + 2. All-gather shards to reconstruct full pre-NS gradient + 3. Run Newton-Schulz on full matrix (must be done globally) + 4. Scatter back to get local orthogonalized update + 5. Apply weight decay and update to local shard (or FP32 master if enabled) + + For non-DTensor parameters, runs single-device Muon. + + When master_weights=True, updates are applied to FP32 masters and + then copied back to BF16 params for numerical stability. + """ + for idx, p in enumerate(self.muon_params): + if p.grad is None: + continue + + # Sharding info drives the routing below: a non-empty list means the + # param is FSDP-sharded (-> distributed NS); an empty list (a plain + # tensor, or a fully-replicated DTensor) takes the single-device path. + shard_info: list[tuple[int, int]] = [] + + if isinstance(p, DTensor): + # Local shard of the parameter (used to write the update in Step 5). + # The gradient / momentum stay as DTensors and are handled below. + local_param = p._local_tensor + + # Get sharding info from DTensor (FSDP typically shards dim 0 / rows). + # Each entry: (mesh_dim, tensor_shard_dim) + device_mesh = p.device_mesh + placements = p.placements + for mesh_dim_idx, placement in enumerate(placements): + if placement.is_shard(): + shard_info.append((mesh_dim_idx, placement.dim)) + + if not shard_info: + # Fully replicated DTensor (e.g. a DDP-like config where + # placements == (Replicate(),)). Legitimate -- fall back to + # single-device NS and log at info level (a warning here would + # needlessly alarm anyone reading the logs). + param_name = self.param_to_name.get(p, "unknown") + log.info( + f"DTensor '{param_name}' has no Shard placement (placements={placements}); " + "running single-device Newton-Schulz (replicated param)." + ) + + if shard_info: + # Distributed Muon: all-gather → NS → scatter + # Handles both 1D sharding (FSDP) and 2D sharding (FSDP + TP) + + # Initialize state with local shard shape + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + + # Step 1: Momentum + Nesterov, kept in DTensor space so the buffer + # stays sharded (memory-efficient) and the gather below is + # padding-aware. These ops are element-wise, so DTensor runs them on + # the local shards under the hood. compute_pre_ns_update does not + # mutate the gradient, so p.grad can be passed directly. + pre_ns_dtensor = compute_pre_ns_update( + p.grad, + state["momentum_buffer"], + momentum=self.muon_momentum, + nesterov=self.nesterov, + ) + + # Step 2: Reconstruct the full matrix from its shards. + # + # ``full_tensor()`` is DTensor's own all-gather (a Shard -> Replicate + # redistribute). Because the DTensor carries the logical shape, it + # rebuilds exactly ``p.shape`` and strips any FSDP2 padding of a shard + # dim that is not divisible by the mesh size. A hand-rolled + # ``all_gather`` + ``cat`` would keep those padding rows and feed a + # wrong-sized / misaligned matrix into Newton-Schulz (the uneven-shard + # bug this path used to have). + # + # PERF TODO: this is a *synchronous* all-gather (we block until the + # full matrix is reconstructed before running Newton-Schulz). If the + # optimizer step ever becomes a bottleneck, explore pipelining: overlap + # the gather for parameter (i) with the Newton-Schulz compute for + # parameter (i-1) across the muon_params loop. + # + # Cast to bf16 BEFORE the all-gather so the collective moves bf16 (not + # fp32), halving its communication volume (matches the dion2 path). This + # is numerically identical: Newton-Schulz casts to bf16 anyway, and bf16 + # rounding commutes with the reconstruction (round-then-gather == + # gather-then-round elementwise). Keep in sync with the input dtype of + # ``zeropower_via_newtonschulz5``. The fp32 momentum buffer is untouched + # (only this transient gathered copy is downcast). + full_pre_ns = pre_ns_dtensor.bfloat16().full_tensor() + + # Sanity check: full_tensor() must return the logical (unpadded) shape. + expected_shape = tuple(p.shape) + actual_shape = tuple(full_pre_ns.shape) + assert actual_shape == expected_shape, ( + f"Failed to reconstruct full matrix for '{self.param_to_name.get(p, 'unknown')}': " + f"got {actual_shape}, expected {expected_shape}" + ) + + # Log shapes on first step to confirm distributed NS is working + if "logged_shapes" not in state: + state["logged_shapes"] = True + log.info( + f"Muon distributed NS: '{self.param_to_name.get(p, 'unknown')}' " + f"local={tuple(local_param.shape)} → full={actual_shape} (verified)" + ) + + # Step 3: Newton-Schulz on full matrix + full_ortho = zeropower_via_newtonschulz5(full_pre_ns, steps=self.ns_steps) + + # Step 4: Scatter the orthogonalized full matrix back to p's sharding. + # + # full_ortho is identical on every rank, so we treat it as Replicate + # and redistribute to p's placements. Replicate -> Shard is pure local + # slicing (no communication) and re-applies the exact same (possibly + # padded) shard layout p has, so local_update lines up with the local + # shard / FP32 master. + ortho_dtensor = DTensor.from_local( + full_ortho, + device_mesh, + [Replicate()] * device_mesh.ndim, + run_check=False, + ) + local_update = ortho_dtensor.redistribute(device_mesh, placements).to_local() + + # Step 5: Apply weight decay and update + # Use full shape for LR scaling (not local shard shape) and the + # owning param-group's lr/wd (honors lr_multipliers / WD grouping). + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(full_pre_ns.shape, base_lr) + + if self.master_weights: + # Update FP32 master, then write the BF16 param directly so we + # do not depend on LowPrecisionCallback (the OptimizersContainer + # hides master_weights from it). Matches FusedAdam, whose kernel + # writes the param in-place. + # + # NOTE: adjusted_lr may be a tensor (capturable mode), so we + # scale via multiply rather than ``alpha=`` (Tensor.add_ only + # accepts a Python-number alpha). + master = get_local_tensor_if_DTensor(self._muon_masters[idx]) + master.mul_(1 - base_lr * wd) + master.add_(local_update.float() * (-adjusted_lr)) + local_param.copy_(master) + else: + local_param.mul_(1 - base_lr * wd) + local_param.add_(local_update * (-adjusted_lr)) + + else: + # Single-device Muon (non-FSDP or non-sharded) + grad = get_local_tensor_if_DTensor(p.grad) + param = get_local_tensor_if_DTensor(p) + + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p).float() + + # Compute pre-NS update (momentum + Nesterov). compute_pre_ns_update + # does not mutate the gradient, so grad can be passed directly. + pre_ns = compute_pre_ns_update( + grad, + get_local_tensor_if_DTensor(state["momentum_buffer"]), + momentum=self.muon_momentum, + nesterov=self.nesterov, + ) + + # Newton-Schulz orthogonalization + update = zeropower_via_newtonschulz5(pre_ns, steps=self.ns_steps) + + # Get adjusted LR based on matrix size (Moonlight scaling), using + # the owning param-group's lr/wd. + base_lr = self._base_lr_for(p) + wd = self._wd_for(p) + adjusted_lr = self._get_adjusted_lr(p.shape, base_lr) + + # Apply weight decay (using base LR) and update (using adjusted LR) + if self.master_weights: + # Update FP32 master, then write the BF16 param directly (see + # the distributed branch above for rationale). adjusted_lr may be + # a tensor (capturable), so scale via multiply, not ``alpha=``. + master = get_local_tensor_if_DTensor(self._muon_masters[idx]) + master.mul_(1 - base_lr * wd) + master.add_(update.float() * (-adjusted_lr)) + param.copy_(master) + else: + param.mul_(1 - base_lr * wd) + param.add_(update * (-adjusted_lr)) + + def _step_adamw(self) -> None: + """ + AdamW step using Transformer Engine's fused multi-tensor Adam kernel. + + Iterates over param groups so each group's lr / betas / eps / weight_decay + (set by the factory's lr_multipliers and disable_weight_decay_for_1d_params) + is honored, then batches by dtype (fp16, bf16, fp32) within the group. The + per-group step counter lives on ``group["step"]`` (FusedAdam-style) so + it is round-tripped by the distributed-checkpoint optimizer state dict. + + When master_weights=True and capturable=True, uses + multi_tensor_adam_capturable_master which maintains FP32 master weights. + """ + if not self.adamw_params: + return + + adam_w_mode = 1 # Decoupled weight decay + bias_correction = 1 + + for group in self.param_groups: + # Only the AdamW-categorized params of this group are handled here; + # the Muon-categorized params were already updated in _step_muon. + group_params = [p for p in group["params"] if id(p) in self._adamw_param_ids] + if not group_params: + continue + + device = group_params[0].device + + # Per-group step counter stored on the group (FusedAdam convention) so + # DCP round-trips it on resume. + if group.get("step", None) is not None: + if self.capturable and not isinstance(group["step"], torch.Tensor): + group["step"] = torch.tensor(group["step"], dtype=torch.int32, device=device) + group["step"] = ( + group["step"].to(device=device) if isinstance(group["step"], torch.Tensor) else group["step"] + ) + group["step"] += 1 + else: + group["step"] = torch.tensor([1], dtype=torch.int32, device=device) if self.capturable else 1 + + if self.capturable and not isinstance(group["lr"], torch.Tensor): + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32, device=device) + + lr = group["lr"] + wd = group["weight_decay"] + step = group["step"] + beta1, beta2 = group["betas"] + eps = group["eps"] + + # Batch parameters by dtype for multi-tensor apply. + g_16, p_16, m_16, v_16 = [], [], [], [] + g_bf, p_bf, m_bf, v_bf = [], [], [], [] + g_32, p_32, m_32, v_32 = [], [], [], [] + p_16_master, p_bf_master, p_32_master = [], [], [] + + for p in group_params: + if p.grad is None: + continue + + grad = get_local_tensor_if_DTensor(p.grad) + param = get_local_tensor_if_DTensor(p) + + state = self.state[p] + if len(state) == 0: + state["exp_avg"] = torch.zeros_like(p).float() + state["exp_avg_sq"] = torch.zeros_like(p).float() + + exp_avg = get_local_tensor_if_DTensor(state["exp_avg"]) + exp_avg_sq = get_local_tensor_if_DTensor(state["exp_avg_sq"]) + master = get_local_tensor_if_DTensor(self._param_to_master[id(p)]) if self.master_weights else None + + if p.dtype == torch.float16: + g_16.append(grad) + p_16.append(param) + m_16.append(exp_avg) + v_16.append(exp_avg_sq) + if self.master_weights: + p_16_master.append(master) + elif p.dtype == torch.bfloat16: + g_bf.append(grad) + p_bf.append(param) + m_bf.append(exp_avg) + v_bf.append(exp_avg_sq) + if self.master_weights: + p_bf_master.append(master) + elif p.dtype == torch.float32: + g_32.append(grad) + p_32.append(param) + m_32.append(exp_avg) + v_32.append(exp_avg_sq) + if self.master_weights: + p_32_master.append(master) + else: + raise RuntimeError(f"Unsupported dtype {p.dtype} for fused AdamW") + + if self.capturable: + # The capturable-master kernel requires an inverse-scale argument; + # bf16-only training has no grad scaler, so it is a constant one. + kernel_inv_scale = torch.ones((1,), device=device, dtype=torch.float32) + dtype_batches = ( + (g_16, p_16, m_16, v_16, p_16_master), + (g_bf, p_bf, m_bf, v_bf, p_bf_master), + (g_32, p_32, m_32, v_32, p_32_master), + ) + kernel = ( + self._multi_tensor_adam_capturable_master + if self.master_weights + else self._multi_tensor_adam_capturable + ) + for g_, p_, m_, v_, pm_ in dtype_batches: + if len(g_) == 0: + continue + tensor_lists = [g_, p_, m_, v_, pm_] if self.master_weights else [g_, p_, m_, v_] + te.pytorch.optimizers.multi_tensor_applier( + kernel, + self._dummy_overflow_buf, + tensor_lists, + lr, + beta1, + beta2, + eps, + step, + adam_w_mode, + bias_correction, + wd, + kernel_inv_scale, + ) + else: + dtype_batches = ( + (g_16, p_16, m_16, v_16), + (g_bf, p_bf, m_bf, v_bf), + (g_32, p_32, m_32, v_32), + ) + for g_, p_, m_, v_ in dtype_batches: + if len(g_) == 0: + continue + te.pytorch.optimizers.multi_tensor_applier( + self._multi_tensor_adam, + self._dummy_overflow_buf, + [g_, p_, m_, v_], + lr, + beta1, + beta2, + eps, + step, + adam_w_mode, + bias_correction, + wd, + ) + + def load_state_dict(self, state_dict: dict) -> None: + """Load optimizer state. + + The optimizer state (per-param momentum / exp_avg / exp_avg_sq and the + per-group ``step``) round-trips through the base ``torch.optim.Optimizer`` + state dict, so the distributed-checkpoint container can save/restore it + with FSDP2 resharding just like FusedAdam. Master weights are *not* + checkpointed; they are rebuilt from the restored params on the next + ``step()`` (see ``_maybe_init_master_weights``). + + This direct-call path (used outside the OptimizersContainer / DCP flow) + keeps the moments in FP32 and normalizes the capturable LR/step tensors. + """ + super().load_state_dict(state_dict) + + # Force master weights to be rebuilt from the (now restored) params. + self._masters_initialized = False + self.param_groups_master = None + + for group in self.param_groups: + device = group["params"][0].device if group["params"] else "cuda" + if self.capturable: + if isinstance(group["lr"], torch.Tensor): + group["lr"] = group["lr"].to(device=device) + else: + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32, device=device) + if group.get("step", None) is not None and not isinstance(group["step"], torch.Tensor): + group["step"] = torch.tensor(group["step"], dtype=torch.int32, device=device) + for p in group["params"]: + state = self.state[p] + if "exp_avg" in state: + state["exp_avg"] = state["exp_avg"].float() + state["exp_avg_sq"] = state["exp_avg_sq"].float() + if "momentum_buffer" in state: + state["momentum_buffer"] = state["momentum_buffer"].float() diff --git a/cosmos_framework/utils/generator/optimizer.py b/cosmos_framework/utils/generator/optimizer.py index 2c1c798b..fe446f37 100644 --- a/cosmos_framework/utils/generator/optimizer.py +++ b/cosmos_framework/utils/generator/optimizer.py @@ -15,6 +15,12 @@ from cosmos_framework.utils.functional.lr_scheduler import LambdaLinearScheduler, LambdaWarmUpCosineScheduler, WSDScheduler from cosmos_framework.utils import log +# Hybrid orthogonalizing optimizers (Muon / Dion2) that own their parameter +# categorization and run their own collective communication. They must be built +# as a single optimizer instance over all selected params (no per-device-mesh +# split) and need ``categorize_params`` called after construction. +_AUX_ADAMW_OPTIMIZERS = ("muonwithauxadamw", "dion2withauxadamw") + class ParamMetadata(NamedTuple): lr: float @@ -68,6 +74,23 @@ def _optimizer_cls( optimizer_kwargs["capturable"] = True optimizer_kwargs["master_weights"] = True optimizer = FusedAdam(params, **optimizer_kwargs) + elif optimizer_type.lower() == "muonwithauxadamw": + from cosmos_framework.utils.generator.muon_with_aux_adamw import MuonWithAuxAdamW + + # Muon's AdamW side is the TE-fused kernel; it is fused by construction and + # absorbs ``fused`` via **kwargs, but we pop it here to be explicit. We force + # capturable + master_weights to match FusedAdam's mixed-precision setup. + optimizer_kwargs.pop("fused", None) + optimizer_kwargs["capturable"] = True + optimizer_kwargs["master_weights"] = True + optimizer = MuonWithAuxAdamW(params, **optimizer_kwargs) + elif optimizer_type.lower() == "dion2withauxadamw": + from cosmos_framework.utils.generator.dion2_with_aux_adamw import Dion2WithAuxAdamW + + optimizer_kwargs.pop("fused", None) + optimizer_kwargs["capturable"] = True + optimizer_kwargs["master_weights"] = True + optimizer = Dion2WithAuxAdamW(params, **optimizer_kwargs) else: raise NotImplementedError(f"Optimizer {optimizer_type} not found.") return optimizer @@ -292,27 +315,49 @@ def __init__( disable_weight_decay_for_1d_params=disable_weight_decay_for_1d_params, ) - # Sub-group by device mesh so fused optimizers operate on same-mesh params. - mesh_groups: dict[str, list[tuple[nn.Parameter, ParamMetadata]]] = collections.defaultdict(list) - for param, metadata in params_with_metadata: - if hasattr(param, "device_mesh"): - # ``mesh_dim_names`` is ``tuple[str, ...] | None`` on DeviceMesh — - # fall back to ``default`` when names weren't assigned. - names = param.device_mesh.mesh_dim_names - mesh_key = "-".join(names) if names else "default" - else: - mesh_key = "default" - mesh_groups[mesh_key].append((param, metadata)) - - # Create one optimizer per mesh, each with per-LR,weight-decay param groups. - for mesh_key, mesh_params in mesh_groups.items(): - log.info(f"Building optimizer for mesh '{mesh_key}'") + if optimizer_type.lower() in _AUX_ADAMW_OPTIMIZERS: + # Muon / Dion2 categorize parameters into an orthogonalized-matrix group + # and an auxiliary AdamW group internally, and run their own collective + # communication across the FULL device mesh. They must therefore be a + # single optimizer instance over all selected params (no per-mesh split). + # + # The factory's param grouping is preserved: ``_build_optimizer_internal`` + # still produces the same (lr, weight_decay) param groups (from + # ``lr_multipliers`` / ``disable_weight_decay_for_1d_params``), and + # Muon/Dion2 read lr/weight_decay per group — degenerating to a single + # global lr/wd when there is only one group (the reference behavior). optimizer = _build_optimizer_internal( - mesh_params, + params_with_metadata, optimizer_type, **optimizer_kwargs, ) + # Categorize against the trainable network only, mirroring + # ``_build_params_with_metadata`` which optimizes the ``net`` subtree + # and skips ``net_ema``. + optimizer.categorize_params(model.net) self.optimizers.append(optimizer) + else: + # Sub-group by device mesh so fused optimizers operate on same-mesh params. + mesh_groups: dict[str, list[tuple[nn.Parameter, ParamMetadata]]] = collections.defaultdict(list) + for param, metadata in params_with_metadata: + if hasattr(param, "device_mesh"): + # ``mesh_dim_names`` is ``tuple[str, ...] | None`` on DeviceMesh — + # fall back to ``default`` when names weren't assigned. + names = param.device_mesh.mesh_dim_names + mesh_key = "-".join(names) if names else "default" + else: + mesh_key = "default" + mesh_groups[mesh_key].append((param, metadata)) + + # Create one optimizer per mesh, each with per-LR,weight-decay param groups. + for mesh_key, mesh_params in mesh_groups.items(): + log.info(f"Building optimizer for mesh '{mesh_key}'") + optimizer = _build_optimizer_internal( + mesh_params, + optimizer_type, + **optimizer_kwargs, + ) + self.optimizers.append(optimizer) log.info(f"Created {len(self.optimizers)} optimizers") @@ -323,6 +368,14 @@ def __len__(self) -> int: return len(self.optimizers) def step(self) -> None: + # NOTE: This container does not forward a ``torch.amp.GradScaler`` to the + # inner optimizers, and training here is bf16 with the scaler disabled + # (``grad_scaler_args={"enabled": False}``), so step() always runs unscaled. + # MuonWithAuxAdamW / Dion2WithAuxAdamW are intentionally bf16-only (their + # AMP grad-scaler support was removed to keep the Newton-Schulz path simple); + # only FusedAdam still carries a scaler hook. Enabling fp16 + dynamic loss + # scaling end-to-end would require re-adding scaler passthrough here and in + # the Muon/Dion2 step paths, which is deliberately not supported. for optimizer in self.optimizers: optimizer.step() diff --git a/cosmos_framework/utils/generator/torchcodec_video.py b/cosmos_framework/utils/generator/torchcodec_video.py index 6a27ac0d..3f8163d8 100644 --- a/cosmos_framework/utils/generator/torchcodec_video.py +++ b/cosmos_framework/utils/generator/torchcodec_video.py @@ -46,11 +46,16 @@ def _build_decoder( num_threads: int = 1, seek_mode: str = "exact", device: str = "cpu", + custom_frame_mappings: bytes | None = None, ) -> Any: normalized_source = _normalize_source(source) # Preserve FFmpeg/TorchCodec's 0 sentinel so callers can request automatic thread selection. num_ffmpeg_threads = 0 if num_threads == 0 else max(num_threads, 1) - kwargs: dict[str, Any] = {"seek_mode": seek_mode, "num_ffmpeg_threads": num_ffmpeg_threads} + kwargs: dict[str, Any] = {"num_ffmpeg_threads": num_ffmpeg_threads} + if custom_frame_mappings is None: + kwargs["seek_mode"] = seek_mode + else: + kwargs["custom_frame_mappings"] = custom_frame_mappings if device != "cpu": kwargs["device"] = device video_decoder_cls = _get_video_decoder_cls() @@ -107,8 +112,15 @@ def __init__( seek_mode: str = "exact", device: str = "cpu", include_dimensions: bool = False, + custom_frame_mappings: bytes | None = None, ) -> None: - self._decoder = _build_decoder(source, num_threads=num_threads, seek_mode=seek_mode, device=device) + self._decoder = _build_decoder( + source, + num_threads=num_threads, + seek_mode=seek_mode, + device=device, + custom_frame_mappings=custom_frame_mappings, + ) self.metadata = _metadata_from_frame(self._decoder, include_dimensions=include_dimensions) def __len__(self) -> int: From 755bff48dd66b467995366796e747ffc649dd83c Mon Sep 17 00:00:00 2001 From: yy-code-nv Date: Tue, 14 Jul 2026 13:17:17 +0800 Subject: [PATCH 23/26] Release 2026-07-13 (i4 c0d7b004) (#108) Automated release from i4. _source_commit: `c0d7b0044b6eb0239bfcd13a31fbb05dc789cb22-dirty` _dest_commit (base): `0fa3ba479a7c1c5be28f81fe084b9ea544d697c9` --- .file_mapping.json | 15 +- .../configs/base/defaults/quantization.py | 33 + .../data/generator/action/transforms.py | 48 +- .../augmentors/interleaved_image_transform.py | 18 +- .../interleaved_image_transform_test.py | 44 + .../generator/sequence_packing/__init__.py | 4 +- .../generator/sequence_packing/modality.py | 252 ++++ .../generator/sequence_packing/packers.py | 159 +-- .../generator/sequence_packing/sequence.py | 1097 +++++++++++++++++ .../sequence_packing/temporal_causal.py | 123 +- .../model/generator/omni_mot_model.py | 9 +- .../utils/generator/image_resize.py | 67 + .../utils/generator/image_resize_test.py | 42 + .../utils/generator/model_loader.py | 85 +- .../utils/generator/quantization.py | 135 ++ 15 files changed, 1924 insertions(+), 207 deletions(-) create mode 100644 cosmos_framework/configs/base/defaults/quantization.py create mode 100644 cosmos_framework/data/generator/augmentors/interleaved_image_transform_test.py create mode 100644 cosmos_framework/data/generator/sequence_packing/modality.py create mode 100644 cosmos_framework/data/generator/sequence_packing/sequence.py create mode 100644 cosmos_framework/utils/generator/image_resize.py create mode 100644 cosmos_framework/utils/generator/image_resize_test.py create mode 100644 cosmos_framework/utils/generator/quantization.py diff --git a/.file_mapping.json b/.file_mapping.json index 5b0e65ba..a084b446 100644 --- a/.file_mapping.json +++ b/.file_mapping.json @@ -1,7 +1,7 @@ { - "_source_commit": "7392c4b59ab7c1f0a94e59fbc4a75a0f3684b66f-dirty", - "_dest_commit": "3d9c0878fd0dde76eac98161aed0493d85a036fd", - "_generated_at": "2026-07-11T05:50:38Z", + "_source_commit": "c0d7b0044b6eb0239bfcd13a31fbb05dc789cb22-dirty", + "_dest_commit": "0fa3ba479a7c1c5be28f81fe084b9ea544d697c9", + "_generated_at": "2026-07-13T07:36:20Z", "files": { "imaginaire/__init__.py": "cosmos_framework/__init__.py", "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", @@ -222,6 +222,7 @@ "projects/cosmos3/cosmos3/configs/base/defaults/model_config.py": "cosmos_framework/configs/base/defaults/model_config.py", "projects/cosmos3/cosmos3/configs/base/defaults/optimizer.py": "cosmos_framework/configs/base/defaults/optimizer.py", "projects/cosmos3/cosmos3/configs/base/defaults/parallelism.py": "cosmos_framework/configs/base/defaults/parallelism.py", + "projects/cosmos3/cosmos3/configs/base/defaults/quantization.py": "cosmos_framework/configs/base/defaults/quantization.py", "projects/cosmos3/cosmos3/configs/base/defaults/reasoner.py": "cosmos_framework/configs/base/defaults/reasoner.py", "projects/cosmos3/cosmos3/configs/base/defaults/tokenizer.py": "cosmos_framework/configs/base/defaults/tokenizer.py", "projects/cosmos3/cosmos3/configs/base/experiment/action/__init__.py": "cosmos_framework/configs/base/experiment/action/__init__.py", @@ -263,6 +264,7 @@ "projects/cosmos3/cosmos3/datasets/augmentors/image_editing_transform_test.py": "cosmos_framework/data/generator/augmentors/image_editing_transform_test.py", "projects/cosmos3/cosmos3/datasets/augmentors/image_resolution_filter.py": "cosmos_framework/data/generator/augmentors/image_resolution_filter.py", "projects/cosmos3/cosmos3/datasets/augmentors/interleaved_image_transform.py": "cosmos_framework/data/generator/augmentors/interleaved_image_transform.py", + "projects/cosmos3/cosmos3/datasets/augmentors/interleaved_image_transform_test.py": "cosmos_framework/data/generator/augmentors/interleaved_image_transform_test.py", "projects/cosmos3/cosmos3/datasets/augmentors/interleaved_video_parsing.py": "cosmos_framework/data/generator/augmentors/interleaved_video_parsing.py", "projects/cosmos3/cosmos3/datasets/augmentors/merge_datadict.py": "cosmos_framework/data/generator/augmentors/merge_datadict.py", "projects/cosmos3/cosmos3/datasets/augmentors/multi_reference_transform.py": "cosmos_framework/data/generator/augmentors/multi_reference_transform.py", @@ -368,13 +370,13 @@ "projects/cosmos3/cosmos3/processors/nemotronvl_processor.py": "cosmos_framework/data/generator/processors/nemotronvl_processor.py", "projects/cosmos3/cosmos3/processors/qwen3vl_processor.py": "cosmos_framework/data/generator/processors/qwen3vl_processor.py", "projects/cosmos3/cosmos3/sequence_packing/__init__.py": "cosmos_framework/data/generator/sequence_packing/__init__.py", - "projects/cosmos3/cosmos3/sequence_packing/modalities.py": "cosmos_framework/data/generator/sequence_packing/modalities.py", + "projects/cosmos3/cosmos3/sequence_packing/modality.py": "cosmos_framework/data/generator/sequence_packing/modality.py", "projects/cosmos3/cosmos3/sequence_packing/mrope.py": "cosmos_framework/data/generator/sequence_packing/mrope.py", "projects/cosmos3/cosmos3/sequence_packing/natten.py": "cosmos_framework/data/generator/sequence_packing/natten.py", "projects/cosmos3/cosmos3/sequence_packing/packers.py": "cosmos_framework/data/generator/sequence_packing/packers.py", "projects/cosmos3/cosmos3/sequence_packing/runtime.py": "cosmos_framework/data/generator/sequence_packing/runtime.py", + "projects/cosmos3/cosmos3/sequence_packing/sequence.py": "cosmos_framework/data/generator/sequence_packing/sequence.py", "projects/cosmos3/cosmos3/sequence_packing/temporal_causal.py": "cosmos_framework/data/generator/sequence_packing/temporal_causal.py", - "projects/cosmos3/cosmos3/sequence_packing/types.py": "cosmos_framework/data/generator/sequence_packing/types.py", "projects/cosmos3/cosmos3/tokenizers/audio/__init__.py": "cosmos_framework/model/generator/tokenizers/audio/__init__.py", "projects/cosmos3/cosmos3/tokenizers/audio/avae.py": "cosmos_framework/model/generator/tokenizers/audio/avae.py", "projects/cosmos3/cosmos3/tokenizers/audio/avae_utils/__init__.py": "cosmos_framework/model/generator/tokenizers/audio/avae_utils/__init__.py", @@ -413,6 +415,8 @@ "projects/cosmos3/cosmos3/utils/flash_attn.py": "cosmos_framework/utils/generator/flash_attn.py", "projects/cosmos3/cosmos3/utils/fused_adam.py": "cosmos_framework/utils/generator/fused_adam.py", "projects/cosmos3/cosmos3/utils/hf_attention_cosmos.py": "cosmos_framework/utils/generator/hf_attention_cosmos.py", + "projects/cosmos3/cosmos3/utils/image_resize.py": "cosmos_framework/utils/generator/image_resize.py", + "projects/cosmos3/cosmos3/utils/image_resize_test.py": "cosmos_framework/utils/generator/image_resize_test.py", "projects/cosmos3/cosmos3/utils/lora.py": "cosmos_framework/utils/generator/lora.py", "projects/cosmos3/cosmos3/utils/model_loader.py": "cosmos_framework/utils/generator/model_loader.py", "projects/cosmos3/cosmos3/utils/model_weights_stats.py": "cosmos_framework/utils/generator/model_weights_stats.py", @@ -420,6 +424,7 @@ "projects/cosmos3/cosmos3/utils/muon_with_aux_adamw.py": "cosmos_framework/utils/generator/muon_with_aux_adamw.py", "projects/cosmos3/cosmos3/utils/optimizer.py": "cosmos_framework/utils/generator/optimizer.py", "projects/cosmos3/cosmos3/utils/parallelism.py": "cosmos_framework/utils/generator/parallelism.py", + "projects/cosmos3/cosmos3/utils/quantization.py": "cosmos_framework/utils/generator/quantization.py", "projects/cosmos3/cosmos3/utils/rand_state.py": "cosmos_framework/utils/generator/rand_state.py", "projects/cosmos3/cosmos3/utils/reasoner/__init__.py": "cosmos_framework/utils/generator/reasoner/__init__.py", "projects/cosmos3/cosmos3/utils/reasoner/constant.py": "cosmos_framework/utils/generator/reasoner/constant.py", diff --git a/cosmos_framework/configs/base/defaults/quantization.py b/cosmos_framework/configs/base/defaults/quantization.py new file mode 100644 index 00000000..87713dfc --- /dev/null +++ b/cosmos_framework/configs/base/defaults/quantization.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +import attrs + + +@attrs.define(slots=False) +class QuantizationConfig: + """Configuration for low-precision quantization of model parameters. + + Controls which quantization method is applied (mxfp8, nvfp4), and which + parameters are selected for quantization via include/exclude key filters. + When ``method`` is None, quantization is disabled and all other fields are + inert. + """ + + # Quantization method for the model. + method: str | None = attrs.field( + default=None, + validator=attrs.validators.optional(attrs.validators.in_({"mxfp8", "nvfp4"})), + ) + + # How to select parameters to select for the quantization. Each key is a + # regular expression matched against a module's fully-qualified name with + # `re.search` (a plain substring is still a valid pattern, so substring-style + # keys keep working, while anchors like `^`/`$`, alternation `a|b`, and + # character classes are also supported). A module is selected only if its FQN + # matches at least one pattern in `include_regex` and matches none in + # `exclude_regex`. If `include_regex` is empty, all parameters are + # considered as included. If `exclude_regex` is empty, no parameters are + # considered as excluded. + include_regex: list[str] = attrs.field(factory=list) + exclude_regex: list[str] = attrs.field(factory=list) diff --git a/cosmos_framework/data/generator/action/transforms.py b/cosmos_framework/data/generator/action/transforms.py index 6417a6e9..05696536 100644 --- a/cosmos_framework/data/generator/action/transforms.py +++ b/cosmos_framework/data/generator/action/transforms.py @@ -301,20 +301,58 @@ def build_sequence_plan_from_mode( # forward_dynamics: all action steps are clean (conditioning) # inverse_dynamics/policy: action is supervised (predicted) # History frames (prepended) are always conditioning. + # + # `base_action_length` is the number of *predicted* action steps, i.e. + # `action_length` minus the prepended history block. There are three + # supported alignments between the action stream and the video stream + # (`video_length` counts the leading conditioning frame at index 0): + # + # Case A (dense video, action = video_length - 1) + # Standard forward-dynamics / policy layout: one action per gap + # between consecutive video frames, no explicit initial-state + # action. Conditioning actions are only the prepended history. + # + # Case B (dense video, action = video_length) + # `use_state=True` layout: one extra action is prepended to + # encode the initial state, so history + one initial-state + # action are conditioning. + # + # Case C (subsampled video, general action_length > video_length) + # Video is subsampled by an integer factor while the action + # stream stays dense. `has_initial_state` mirrors Case B: it is + # True iff `(base_action_length - 1)` divides evenly by + # `(video_length - 1)`, i.e. the extra initial-state action is + # present. base_action_length = action_length - num_history_actions + has_initial_state = video_length > 1 and (base_action_length - 1) % (video_length - 1) == 0 if mode == "forward_dynamics": condition_frame_indexes_action = list(range(action_length)) - # This currently assumes that the action length is the same as the video length - 1 - # and if action length is the same as the video length, then the first action is the conditioning action elif base_action_length == video_length - 1: + # Case A: only the history block is conditioning. condition_frame_indexes_action = list(range(num_history_actions)) elif base_action_length == video_length: + # Case B: history block + one initial-state action. condition_frame_indexes_action = list(range(num_history_actions + 1)) + else: + # Case C: same rule as A/B based on whether an initial-state action exists. + condition_frame_indexes_action = list( + range(num_history_actions + 1 if has_initial_state else num_history_actions) + ) + # `action_start_frame_offset` shifts the action stream so that the first + # non-conditioning action aligns with video frame index 1 (the first + # frame the model must predict). The `-num_history_actions` term + # reserves room for the history block; the additional `+1` (i.e. + # `1 - num_history_actions`) is applied only when there is no + # initial-state action, so the offset skips past the first action + # slot that is instead consumed by the initial state. if base_action_length == video_length - 1: - action_start_frame_offset = 1 - num_history_actions - if base_action_length == video_length: - action_start_frame_offset = -num_history_actions + action_start_frame_offset = 1 - num_history_actions # Case A + elif base_action_length == video_length: + action_start_frame_offset = -num_history_actions # Case B + else: + # Case C: same rule as A/B based on whether an initial-state action exists. + action_start_frame_offset = -num_history_actions if has_initial_state else 1 - num_history_actions return SequencePlan( has_text=has_text, diff --git a/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py b/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py index 24f61a07..b303910f 100644 --- a/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py +++ b/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py @@ -14,6 +14,7 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.data.imaginaire.webdataset.augmentors.image.misc import obtain_image_size +from cosmos_framework.utils.generator.image_resize import DEFAULT_MAX_PIXELS, get_max_pixels_resized_size Image.MAX_IMAGE_PIXELS = 933120000 @@ -290,7 +291,7 @@ class InterleavedMediaResizeByMaxPixels(Augmentor): def __init__( self, input_keys: Optional[List] = None, - max_pixels: int = 1048576, + max_pixels: int = DEFAULT_MAX_PIXELS, padding_divisor: int = 16, args: Optional[dict] = None, ) -> None: @@ -301,15 +302,12 @@ def __init__( def _compute_target_size(self, width: int, height: int) -> tuple[int, int]: """Scale to fit within max_pixels, then align down to padding_divisor.""" - total_pixels = width * height - if total_pixels > self.max_pixels: - scale = math.sqrt(self.max_pixels / total_pixels) - width = int(width * scale) - height = int(height * scale) - - width = max(self.padding_divisor, (width // self.padding_divisor) * self.padding_divisor) - height = max(self.padding_divisor, (height // self.padding_divisor) * self.padding_divisor) - return width, height + return get_max_pixels_resized_size( + width=width, + height=height, + max_pixels=self.max_pixels, + padding_constant=self.padding_divisor, + ) def _resize_image(self, img: Image.Image) -> Image.Image: w, h = img.size diff --git a/cosmos_framework/data/generator/augmentors/interleaved_image_transform_test.py b/cosmos_framework/data/generator/augmentors/interleaved_image_transform_test.py new file mode 100644 index 00000000..a1588ce1 --- /dev/null +++ b/cosmos_framework/data/generator/augmentors/interleaved_image_transform_test.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Tests for interleaved image resize augmentors.""" + +from __future__ import annotations + +import pytest +from PIL import Image + +from cosmos_framework.data.generator.augmentors.interleaved_image_transform import ( + InterleavedMediaResizeByMaxPixels, +) + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +def test_interleaved_media_resize_by_max_pixels_resizes_images_and_videos() -> None: + transform = InterleavedMediaResizeByMaxPixels(max_pixels=1024 * 1024, padding_divisor=16) + wide_image = Image.new("RGB", (1920, 1080), color="blue") + small_image = Image.new("RGB", (640, 480), color="red") + + result = transform( + { + "media_list": { + "image_0": wide_image, + "image_1": small_image, + "video_0": [wide_image.copy()], + } + } + ) + + assert result is not None + resized_media = result["diffusion_media_list"] + assert resized_media["image_0"].size == (1360, 768) + assert resized_media["image_1"].size == (640, 480) + assert resized_media["video_0"][0].size == (1360, 768) + + +def test_interleaved_media_resize_by_max_pixels_rejects_impossible_budget() -> None: + transform = InterleavedMediaResizeByMaxPixels(max_pixels=255, padding_divisor=16) + + with pytest.raises(ValueError, match="too small"): + transform({"media_list": {"image_0": Image.new("RGB", (64, 64))}}) diff --git a/cosmos_framework/data/generator/sequence_packing/__init__.py b/cosmos_framework/data/generator/sequence_packing/__init__.py index 33349353..fbb87d0e 100644 --- a/cosmos_framework/data/generator/sequence_packing/__init__.py +++ b/cosmos_framework/data/generator/sequence_packing/__init__.py @@ -3,9 +3,9 @@ """High-level entry points for VFM sequence packing.""" +from cosmos_framework.data.generator.sequence_packing.modality import ModalityData from cosmos_framework.data.generator.sequence_packing.packers import pack_input_sequence -from cosmos_framework.data.generator.sequence_packing.types import ( - ModalityData, +from cosmos_framework.data.generator.sequence_packing.sequence import ( PackedSequence, SequencePlan, build_sequence_plans_from_data_batch, diff --git a/cosmos_framework/data/generator/sequence_packing/modality.py b/cosmos_framework/data/generator/sequence_packing/modality.py new file mode 100644 index 00000000..05c3a8e4 --- /dev/null +++ b/cosmos_framework/data/generator/sequence_packing/modality.py @@ -0,0 +1,252 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Modality utility helpers for VFM sequence packing.""" + +from dataclasses import dataclass, field + +import torch + + +def _empty_long_tensor() -> torch.Tensor: + return torch.empty(0, dtype=torch.long) # [0] + + +def _empty_float_tensor() -> torch.Tensor: + return torch.empty(0, dtype=torch.float32) # [0] + + +@dataclass +class ModalitySpan: + """One contiguous packed span paired with its logical modality payload slice. + + Attributes: + sequence_start: First global packed-sequence index owned by this span. + sequence_len: Number of contiguous tokens in the packed sequence. + payload_index: Index into ``ModalityData.tokens`` for the backing payload tensor. + payload_start: First flattened token offset within the backing payload tensor. + payload_len: Number of flattened payload tokens covered by this span. + payload_shape: Logical payload slice shape. For example, vision frame spans use + ``(1, patch_h, patch_w)``, action spans use ``(tcf,)``, and sound spans use + ``(1, 1, 1)``. + """ + + sequence_start: int + sequence_len: int + payload_index: int + payload_start: int + payload_len: int + payload_shape: tuple[int, ...] + + +@dataclass +class ModalityDataBuilder: + """Append-only construction state for a single generation modality. + + Attributes: + spans: Contiguous packed spans pointing back into grouped payload tensors. + sequence_indexes: Global packed-sequence indexes for all tokens of this modality. + timesteps: Diffusion timesteps for noised tokens only. + mse_loss_indexes: Global packed-sequence indexes where MSE loss should be computed. + token_shapes: Shape metadata for each payload tensor. Vision and sound use + ``(T, H, W)``-style tuples; action uses ``(T,)``. + tokens: Original modality payload tensors, kept grouped by sample/item. + condition_mask: Per-payload masks where 1 indicates clean/conditioning tokens + and 0 indicates noised/supervised tokens. + noisy_frame_indexes: Per-payload indexes of noised frames. These are constructed + during packing to avoid GPU-to-CPU synchronization later. + """ + + spans: list[ModalitySpan] = field(default_factory=list) + sequence_indexes: list[int] = field(default_factory=list) + timesteps: list[float] = field(default_factory=list) + mse_loss_indexes: list[int] = field(default_factory=list) + # list[tuple[int,int,int]] for vision, list[tuple[int]] for action, list[tuple[int,int,int]] for sound + token_shapes: list[tuple[int, ...]] = field(default_factory=list) + tokens: list[torch.Tensor] = field(default_factory=list) + condition_mask: list[torch.Tensor] = field(default_factory=list) + noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list) + + +@dataclass +class ModalityData: + """Finalized model-facing data for a single generation modality. + + Index-like fields are tensors after packing. Payload fields remain lists + because samples may have variable token shapes and because downstream code + mutates token payloads after noise injection and clean replay. + + Attributes: + sequence_indexes: Tensor of global packed-sequence indexes for all tokens of + this modality. + timesteps: Tensor of diffusion timesteps for noised tokens only. + mse_loss_indexes: Tensor of global packed-sequence indexes where MSE loss + should be computed. + spans: Contiguous packed spans pointing back into grouped payload tensors. + token_shapes: Shape metadata for each payload tensor. Vision and sound use + ``(T, H, W)``-style tuples; action uses ``(T,)``. + tokens: Original modality payload tensors, kept grouped by sample/item. + condition_mask: Per-payload masks where 1 indicates clean/conditioning tokens + and 0 indicates noised/supervised tokens. + noisy_frame_indexes: Per-payload indexes of noised frames. + domain_id: Domain IDs for multi-domain training. Only used for action. + raw_action_dim: Raw action dimensions. Only used for action-channel masking. + """ + + sequence_indexes: torch.Tensor = field(default_factory=_empty_long_tensor) + timesteps: torch.Tensor = field(default_factory=_empty_float_tensor) + mse_loss_indexes: torch.Tensor = field(default_factory=_empty_long_tensor) + spans: list[ModalitySpan] = field(default_factory=list) + # list[tuple[int,int,int]] for vision, list[tuple[int]] for action, list[tuple[int,int,int]] for sound + token_shapes: list[tuple[int, ...]] = field(default_factory=list) + tokens: list[torch.Tensor] = field(default_factory=list) + condition_mask: list[torch.Tensor] = field(default_factory=list) + noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list) + domain_id: list[torch.Tensor] = field(default_factory=list) + raw_action_dim: list[torch.Tensor | None] | None = field(default_factory=list) + + def __post_init__(self) -> None: + assert isinstance(self.sequence_indexes, torch.Tensor), "ModalityData.sequence_indexes must be finalized" + assert isinstance(self.timesteps, torch.Tensor), "ModalityData.timesteps must be finalized" + assert isinstance(self.mse_loss_indexes, torch.Tensor), "ModalityData.mse_loss_indexes must be finalized" + + def to_cuda(self) -> None: + """Move all tensor fields to CUDA in-place.""" + self.sequence_indexes = self.sequence_indexes.cuda() + self.timesteps = self.timesteps.cuda() + self.mse_loss_indexes = self.mse_loss_indexes.cuda() + self.tokens = [token.cuda() for token in self.tokens] + self.condition_mask = [cm.cuda() for cm in self.condition_mask] + self.noisy_frame_indexes = [ni.cuda() for ni in self.noisy_frame_indexes] + self.domain_id = [d.cuda() for d in self.domain_id] + # raw_action_dim is optional (e.g., when action-channel masking is disabled). + if self.raw_action_dim is not None: + self.raw_action_dim = [d.cuda() if d is not None else None for d in self.raw_action_dim] + + +def prepare_attention_mask_per_sample(split_lens, attn_modes, device="cpu"): + """Prepare dense attention mask for a single sample with multiple splits. + + Args: + split_lens: List of integers indicating length of each split within the sample + attn_modes: List of attention modes for each split ('causal', 'full', or 'noise') + device: Device to place the attention mask tensor on + + Returns: + Attention mask tensor of shape (sample_len, sample_len) with -inf for masked positions + """ + sample_len = sum(split_lens) + attention_mask = torch.zeros((sample_len, sample_len), dtype=torch.bool, device=device) # [sample_len,sample_len] + + # First pass: Set up basic attention patterns for each split + current_pos = 0 + for split_len, attn_mode in zip(split_lens, attn_modes): + assert attn_mode in ["causal", "full", "noise"], f"Invalid attention mode: {attn_mode}" + + split_start = current_pos + split_end = current_pos + split_len + + if attn_mode == "causal": + # Causal: lower triangular within split + full attention to previous splits + attention_mask[split_start:split_end, split_start:split_end] = torch.ones( + (split_len, split_len), device=device + ).tril() # [split_len,split_len] + attention_mask[split_start:split_end, :split_start] = 1 + else: # "full" or "noise" + # Full attention within split and to previous splits + attention_mask[split_start:split_end, split_start:split_end] = torch.ones( + (split_len, split_len), device=device + ) # [split_len,split_len] + attention_mask[split_start:split_end, :split_start] = 1 + + current_pos += split_len + + # Second pass: Handle noise mode - mask out noise columns except within same split + current_pos = 0 + for split_len, attn_mode in zip(split_lens, attn_modes): + if attn_mode == "noise": + split_start = current_pos + split_end = current_pos + split_len + + # Zero out the entire column for noise tokens + attention_mask[:, split_start:split_end] = 0 + # But allow self-attention within the noise split + attention_mask[split_start:split_end, split_start:split_end] = 1 + + current_pos += split_len + + # Convert boolean mask to float with -inf for masked positions + attention_mask = torch.zeros_like(attention_mask, dtype=torch.float).masked_fill_( + ~attention_mask, float("-inf") + ) # [sample_len,sample_len] + + return attention_mask + + +# ============================================================================ +# Tokenizer utilities +# ============================================================================ + + +def add_special_tokens(tokenizer): + """Add image-related special tokens to tokenizer if not already present. + + Args: + tokenizer: Tokenizer to add special tokens to + + Returns: + Tuple of (modified tokenizer, dict of new token IDs) + """ + # Collect existing special tokens + existing_special_tokens = [] + for key, value in tokenizer.special_tokens_map.items(): + if isinstance(value, str): + existing_special_tokens.append(value) + elif isinstance(value, list): + existing_special_tokens.extend(value) + + # Define image boundary tokens to add if missing + tokens_to_add = [] + if "<|vision_start|>" not in existing_special_tokens: + tokens_to_add.append("<|vision_start|>") + if "<|vision_end|>" not in existing_special_tokens: + tokens_to_add.append("<|vision_end|>") + + # Add new tokens to tokenizer vocabulary + if tokens_to_add: + tokenizer.add_tokens(tokens_to_add) + + # Get token IDs for image boundary tokens + new_token_ids = { + "start_of_generation": tokenizer.convert_tokens_to_ids("<|vision_start|>"), + "end_of_generation": tokenizer.convert_tokens_to_ids("<|vision_end|>"), + } + + return tokenizer, new_token_ids + + +def compute_text_split_length( + num_caption_tokens: int, + special_tokens: dict[str, int], + has_generation: bool = True, +) -> int: + """Compute the total text split length without mutating any state. + + This is the number of token positions occupied by the text split in a + packed sequence: caption tokens + optional BOS + EOS + optional BOV. + + Args: + num_caption_tokens: Number of raw caption token IDs (before special tokens). + special_tokens: Dictionary of special token IDs (checked for ``"bos_token_id"``). + has_generation: Whether a start-of-generation (BOV) token follows text. + + Returns: + Total text split length (positions consumed in the packed sequence). + """ + n = num_caption_tokens + if "bos_token_id" in special_tokens: + n += 1 + n += 1 # EOS + if has_generation: + n += 1 # start-of-generation / BOV + return n diff --git a/cosmos_framework/data/generator/sequence_packing/packers.py b/cosmos_framework/data/generator/sequence_packing/packers.py index b6a52234..68333931 100644 --- a/cosmos_framework/data/generator/sequence_packing/packers.py +++ b/cosmos_framework/data/generator/sequence_packing/packers.py @@ -9,19 +9,34 @@ import torch -from cosmos_framework.data.generator.sequence_packing.modalities import ( - pack_action_tokens, - pack_sound_tokens, - pack_text_tokens, - pack_vision_tokens, -) +from cosmos_framework.data.generator.sequence_packing.sequence import PackedSequence, PackedSequenceBuilder, SequencePlan from cosmos_framework.data.generator.sequence_packing.temporal_causal import pack_supertokens_temporal_causal -from cosmos_framework.data.generator.sequence_packing.types import PackedSequence, SequencePlan if TYPE_CHECKING: from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean +def _get_optional_fps( + fps_values: torch.Tensor | list[torch.Tensor | float] | None, + index: int, +) -> float | None: + """Return an optional FPS value as a float. + + Args: + fps_values: Optional tensor/list of per-sample or per-item FPS values. + index: Entry to read from ``fps_values``. + + Returns: + FPS value as ``float`` when present, otherwise ``None``. + """ + if fps_values is None or index >= len(fps_values): + return None + fps_value = fps_values[index] # [] + if isinstance(fps_value, torch.Tensor): + return float(fps_value.item()) + return float(fps_value) + + def pack_input_sequence( sequence_plans: list[SequencePlan], input_text_indexes: list[list[int]], @@ -67,6 +82,8 @@ def pack_input_sequence( unified_3d_mrope_reset_spatial_ids: If True (default), spatial (H, W) indices start from 0 for each vision segment. If False, spatial indices are offset by the temporal offset (Qwen2VL-style). + unified_3d_mrope_temporal_modality_margin: Extra temporal offset inserted between + text and generation modalities. enable_fps_modulation: If True, scale temporal position IDs based on video FPS to reflect real time. Requires fps_vision in gen_data_clean. Uses the same flag as diffusion_expert_config.enable_fps_modulation. @@ -79,6 +96,13 @@ def pack_input_sequence( vision_temporal_position_mode: Temporal coordinates used for unified_3d_mrope vision tokens. "latent_index" keeps legacy positions; "uniae_source_right_edge" uses per-latent positions from gen_data_clean.temporal_positions_vision. + video_temporal_causal: If True, pack vision and optional action as temporal-causal + supertokens instead of separate modality blocks. + action_dim: Action feature dimension used when temporal-causal packing creates + null action tokens. + initial_mrope_temporal_offset: Initial temporal cursor for each sample, used by + autoregressive inference to seed mRoPE positions. + Returns: PackedSequence containing all packed tensors and metadata. See PackedSequence for field details. """ @@ -125,11 +149,11 @@ def pack_input_sequence( ) use_float_mrope_positions = enable_fps_modulation or explicit_vision_temporal_positions_active - # Initialize packed sequence (acts as builder during packing) - packed_seq = PackedSequence() + # Initialize mutable builder state for sequence construction. + seq_builder = PackedSequenceBuilder() # Configure 3D mRoPE on the builder. - packed_seq._mrope_reset_spatial = unified_3d_mrope_reset_spatial_ids + seq_builder._mrope_reset_spatial = unified_3d_mrope_reset_spatial_ids # Maintain separate indices for each modality idx_text = 0 @@ -150,7 +174,7 @@ def pack_input_sequence( # mRoPE temporal offset resets per sample. # initial_mrope_temporal_offset is non-zero only for AR inference (frame N seeds at N*tcf). - packed_seq._mrope_temporal_offset = initial_mrope_temporal_offset + seq_builder.begin_sample(initial_mrope_temporal_offset) _ts = input_timesteps[sample_idx] input_timestep = _ts.item() if _ts.numel() == 1 else _ts # float (TF) or Tensor(T_max,) (DF) @@ -161,8 +185,7 @@ def pack_input_sequence( idx_text += 1 has_generation_for_sample = sequence_plan.has_vision or sequence_plan.has_action or sequence_plan.has_sound - text_sample_len = pack_text_tokens( - packed_seq, + text_sample_len = seq_builder.pack_text_tokens( text_ids, special_tokens, has_generation=has_generation_for_sample, @@ -171,10 +194,10 @@ def pack_input_sequence( sample_len += text_sample_len # End of text modality, add an offset as the boundary between text and vision. - packed_seq._mrope_temporal_offset += unified_3d_mrope_temporal_modality_margin + seq_builder.advance_mrope_temporal_offset(unified_3d_mrope_temporal_modality_margin) # Save temporal offset before vision for action tokens (action uses same offset as vision start) - vision_start_temporal_offset = packed_seq._mrope_temporal_offset + vision_start_temporal_offset = seq_builder.mrope_temporal_offset # Pack vision (and optionally action) tokens if video_temporal_causal and sequence_plan.has_vision: @@ -183,28 +206,17 @@ def pack_input_sequence( input_vision_tokens = gen_data_clean.x0_tokens_vision[idx_vision] idx_vision += 1 - vision_fps = None - if ( - enable_fps_modulation - and gen_data_clean.fps_vision is not None - and idx_vision - 1 < len(gen_data_clean.fps_vision) - ): - vision_fps = float(gen_data_clean.fps_vision[idx_vision - 1].item()) + vision_fps = _get_optional_fps(gen_data_clean.fps_vision, idx_vision - 1) input_action_tokens_tc: torch.Tensor | None = None - action_fps_tc: float | None = None + action_fps_tc = None if sequence_plan.has_action: input_action_tokens_tc = gen_data_clean.x0_tokens_action[idx_action] - if ( - enable_fps_modulation - and gen_data_clean.fps_action is not None - and idx_action < len(gen_data_clean.fps_action) - ): - action_fps_tc = float(gen_data_clean.fps_action[idx_action].item()) + action_fps_tc = _get_optional_fps(gen_data_clean.fps_action, idx_action) idx_action += 1 supertoken_split_len, null_flag = pack_supertokens_temporal_causal( - packed_seq=packed_seq, + seq_builder=seq_builder, input_vision_tokens=input_vision_tokens, input_action_tokens=input_action_tokens_tc, condition_frame_indexes_vision=sequence_plan.condition_frame_indexes_vision, @@ -223,7 +235,9 @@ def pack_input_sequence( # stamp the supertoken layout constant directly here. This is the # single source of truth read by downstream attention / KV-cache # code (no recomputation in the network). - packed_seq.num_action_tokens_per_supertoken = temporal_compression_factor if sequence_plan.has_action else 0 + seq_builder.num_action_tokens_per_supertoken = ( + temporal_compression_factor if sequence_plan.has_action else 0 + ) sample_len += supertoken_split_len vision_split_len = supertoken_split_len action_split_len = 0 # Already absorbed into supertoken_split_len @@ -254,7 +268,7 @@ def pack_input_sequence( # offset equals snapshot + latent_t (single-clip semantics for # downstream EOV / next-modality tokens). shared_grid = sequence_plan.share_vision_temporal_positions and num_vis > 1 - items_temporal_offset_snapshot = packed_seq._mrope_temporal_offset + items_temporal_offset_snapshot = seq_builder.mrope_temporal_offset shared_latent_t: int | None = None shared_patch_h: int | None = None shared_patch_w: int | None = None @@ -264,28 +278,21 @@ def pack_input_sequence( # the same conditioning FPS, so we read by sample_idx, not by the # flat idx_vision counter (which would alias to a neighbor sample's # fps and corrupt RoPE FPS modulation). - sample_vision_fps: float | None = None - if ( - enable_fps_modulation - and gen_data_clean.fps_vision is not None - and sample_idx < len(gen_data_clean.fps_vision) - ): - sample_vision_fps = float(gen_data_clean.fps_vision[sample_idx].item()) + sample_vision_fps = _get_optional_fps(gen_data_clean.fps_vision, sample_idx) for item_idx in range(num_vis): flat_vision_idx = idx_vision - input_vision_tokens = gen_data_clean.x0_tokens_vision[flat_vision_idx] + input_vision_tokens = gen_data_clean.x0_tokens_vision[flat_vision_idx] # [1,C,T,H,W] vision_temporal_positions: torch.Tensor | None = None if explicit_vision_temporal_positions_active: assert gen_data_clean.temporal_positions_vision is not None - vision_temporal_positions = gen_data_clean.temporal_positions_vision[flat_vision_idx] + vision_temporal_positions = gen_data_clean.temporal_positions_vision[flat_vision_idx] # [T] if vision_temporal_positions.shape[0] != input_vision_tokens.shape[2]: raise ValueError( "vision_temporal_positions must match latent_t for each vision item, " f"got {vision_temporal_positions.shape[0]} positions and " f"latent_t={input_vision_tokens.shape[2]} for item {flat_vision_idx}." ) - vision_fps = sample_vision_fps idx_vision += 1 # Determine conditioning for this vision item. @@ -332,15 +339,14 @@ def pack_input_sequence( f"{shared_temporal_positions.tolist()}." ) # Rewind so this item starts at the same temporal offset as item 0. - packed_seq._mrope_temporal_offset = items_temporal_offset_snapshot + seq_builder.set_mrope_temporal_offset(items_temporal_offset_snapshot) - item_split_len = pack_vision_tokens( - packed_seq=packed_seq, + item_split_len = seq_builder.pack_vision_tokens( input_vision_tokens=input_vision_tokens, condition_frame_indexes_vision=item_condition_frames, input_timestep=input_timestep, latent_patch_size=latent_patch_size, - vision_fps=vision_fps, + vision_fps=sample_vision_fps, enable_fps_modulation=enable_fps_modulation, base_fps=base_fps, temporal_compression_factor=temporal_compression_factor, @@ -349,8 +355,9 @@ def pack_input_sequence( vision_split_len += item_split_len if track_item_split_lens: sample_item_split_lens.append(item_split_len) + if track_item_split_lens: - packed_seq.vision_item_split_lens.append(sample_item_split_lens) + seq_builder.vision_item_split_lens.append(sample_item_split_lens) sample_len += vision_split_len else: @@ -359,20 +366,10 @@ def pack_input_sequence( # Pack action tokens if has_action=True if sequence_plan.has_action: input_action_tokens = gen_data_clean.x0_tokens_action[idx_action] - - # Get FPS for action (action may have its own FPS independent of vision) - action_fps: float | None = None - if ( - enable_fps_modulation - and gen_data_clean.fps_action is not None - and idx_action < len(gen_data_clean.fps_action) - ): - action_fps = float(gen_data_clean.fps_action[idx_action].item()) - + action_fps = _get_optional_fps(gen_data_clean.fps_action, idx_action) idx_action += 1 - action_split_len = pack_action_tokens( - packed_seq=packed_seq, + action_split_len = seq_builder.pack_action_tokens( input_action_tokens=input_action_tokens, condition_frame_indexes_action=sequence_plan.condition_frame_indexes_action, input_timestep=input_timestep, @@ -390,20 +387,10 @@ def pack_input_sequence( # Pack sound tokens if has_sound=True if sequence_plan.has_sound: input_sound_tokens = gen_data_clean.x0_tokens_sound[idx_sound] - - # Get FPS for sound (from gen_data_clean, like vision and action) - sound_fps: float | None = None - if ( - enable_fps_modulation - and gen_data_clean.fps_sound is not None - and idx_sound < len(gen_data_clean.fps_sound) - ): - sound_fps = float(gen_data_clean.fps_sound[idx_sound].item()) - + sound_fps = _get_optional_fps(gen_data_clean.fps_sound, idx_sound) idx_sound += 1 - sound_split_len = pack_sound_tokens( - packed_seq=packed_seq, + sound_split_len = seq_builder.pack_sound_tokens( input_sound_tokens=input_sound_tokens, condition_frame_indexes_sound=sequence_plan.condition_frame_indexes_sound, input_timestep=input_timestep, @@ -421,28 +408,14 @@ def pack_input_sequence( eov_len = 0 has_any_generation = sequence_plan.has_vision or sequence_plan.has_action or sequence_plan.has_sound if include_end_of_generation_token and has_any_generation: - # Type narrowing: we're in build mode, fields are lists - assert isinstance(packed_seq.text_ids, list) - assert isinstance(packed_seq.text_indexes, list) - assert isinstance(packed_seq.position_ids, list) - - packed_seq.text_ids.append(special_tokens["end_of_generation"]) - packed_seq.text_indexes.append(packed_seq.curr) - - # Use float dtype when any vision mRoPE positions are fractional. - eov_dtype = torch.float32 if use_float_mrope_positions else torch.long - eov_mrope_ids = torch.full((3, 1), packed_seq._mrope_temporal_offset, dtype=eov_dtype) # [3,1] - packed_seq.position_ids.append(eov_mrope_ids) # type: ignore[arg-type] - packed_seq._mrope_temporal_offset += 1 - - packed_seq.curr += 1 - eov_len = 1 - sample_len += 1 + eov_len = seq_builder.append_end_of_generation_token( + token_id=special_tokens["end_of_generation"], + use_float_mrope_positions=use_float_mrope_positions, + ) + sample_len += eov_len combined_split_len = vision_split_len + action_split_len + sound_split_len + eov_len - packed_seq.attn_modes.append("full") - packed_seq.split_lens.append(combined_split_len) - packed_seq.sample_lens.append(sample_len) + seq_builder.finish_sample(combined_split_len, sample_len) # Assert consistent null_action_supertokens across all TC samples, then set once if null_action_flags: @@ -450,9 +423,9 @@ def pack_input_sequence( f"Inconsistent null_action_supertokens across samples: {null_action_flags}. " "All samples in a batch must have the same structure (all training or all AR inference)." ) - packed_seq.null_action_supertokens = null_action_flags[0] + seq_builder.null_action_supertokens = null_action_flags[0] # Finalize and return packed data - return packed_seq.finalize( + return seq_builder.finalize( gen_data_clean=gen_data_clean, ) diff --git a/cosmos_framework/data/generator/sequence_packing/sequence.py b/cosmos_framework/data/generator/sequence_packing/sequence.py new file mode 100644 index 00000000..095e19dc --- /dev/null +++ b/cosmos_framework/data/generator/sequence_packing/sequence.py @@ -0,0 +1,1097 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Sequence builder/output and plan helpers for VFM sequence packing.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import torch + +from cosmos_framework.data.generator.sequence_packing.modality import ModalityData, ModalityDataBuilder, ModalitySpan +from cosmos_framework.data.generator.sequence_packing.mrope import ( + get_3d_mrope_ids_text_tokens, + get_3d_mrope_ids_vae_tokens, +) + +if TYPE_CHECKING: + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean + + +def _empty_long_tensor() -> torch.Tensor: + return torch.empty(0, dtype=torch.long) # [0] + + +@dataclass +class PackedSequenceBuilder: + """Mutable construction state for sequence packing. + + Attributes: + sample_lens: Length of each sample in the packed sequence. + split_lens: Length of each attention split. Each sample contributes a causal + text split and a full generation split. + attn_modes: Attention mode for each split, such as ``"causal"`` or ``"full"``. + is_image_batch: Whether this batch contains images rather than videos. + sequence_length: Total packed sequence length. Filled during finalization. + current_seq_index: Next global packed-sequence index to append. + text_ids: Text token IDs accumulated during packing, including special tokens. + text_indexes: Global packed-sequence indexes for text tokens. + position_ids: mRoPE position ID blocks accumulated as ``[3, N]`` tensors. + label_ids: Label IDs for text cross-entropy loss. + ce_loss_indexes: Global packed-sequence indexes for text cross-entropy loss. + ce_loss_weights: Per-token weights for text cross-entropy loss. + _mrope_temporal_offset: Running temporal cursor for per-sample mRoPE generation. + _mrope_reset_spatial: Whether spatial mRoPE IDs reset for each segment. + null_action_supertokens: Whether temporal-causal supertoken 0 contains null + action tokens. + num_action_tokens_per_supertoken: Number of action tokens prefixing each + temporal-causal vision supertoken. + vision: Vision modality construction state, or ``None`` if no vision was appended. + action: Action modality construction state, or ``None`` if no action was appended. + sound: Sound modality construction state, or ``None`` if no sound was appended. + vision_item_split_lens: Per-sample per-vision-item token counts for multi-control + transfer. + control_weights: Per-sample per-control weights for multi-control weighted V-scaling. + """ + + # Sequence structure + sample_lens: list[int] = field(default_factory=list) + split_lens: list[int] = field(default_factory=list) + attn_modes: list[str] = field(default_factory=list) + is_image_batch: bool = False + sequence_length: int = 0 + + # Build-time tracking (used during packing, not after finalize) + current_seq_index: int = 0 + + # Text modality append state + text_ids: list[int] = field(default_factory=list) + text_indexes: list[int] = field(default_factory=list) + position_ids: list[torch.Tensor] = field(default_factory=list) + + # Loss computation - Cross Entropy (text) + label_ids: list[int] = field(default_factory=list) + ce_loss_indexes: list[int] = field(default_factory=list) + ce_loss_weights: list[float] = field(default_factory=list) + + # Build-time mRoPE tracking (used during packing, not after finalize). + # position_ids accumulates (3, N) tensors and finalize() produces a + # (3, total_seq_len) tensor. + # Running temporal index for mRoPE position ID generation within a single sample. + # Reset to 0 at the start of each sample, then advanced by text and vision helpers + # as segments are packed. Action reuses the pre-vision snapshot (parallel temporal + # range) without advancing it. Float when FPS modulation is enabled. + # E.g. offset=0 -> text(4 tokens) -> offset=4 -> vision(3 frames) -> offset=7. + _mrope_temporal_offset: int | float = 0 + _mrope_reset_spatial: bool = True + + # Temporal causal: whether supertoken 0's action slot contains null tokens. + # True for all training calls and AR frame 0; False for AR frame N>0 (real actions). + # Used by three_way_attention to zero out V for null action tokens (inline when attention_meta.null_action_supertokens=True). + null_action_supertokens: bool = False + + # Temporal causal: number of action tokens prefixing each vision supertoken. + # Equals temporal_compression_factor when actions are packed inline; 0 when + # action_gen=False or for non-temporal-causal layouts. Single source of truth + # for downstream attention/KV-cache code (per-supertoken layout is + # num_action_tokens_per_supertoken + H_p * W_p). + num_action_tokens_per_supertoken: int = 0 + + # Generation modality construction state + vision: ModalityDataBuilder | None = None + action: ModalityDataBuilder | None = None + sound: ModalityDataBuilder | None = None + + def ensure_vision(self) -> ModalityDataBuilder: + """Return the vision builder, creating it on first use. + + Returns: + Vision ``ModalityDataBuilder`` for subsequent append operations. + """ + if self.vision is None: + self.vision = ModalityDataBuilder() + return self.vision + + def ensure_action(self) -> ModalityDataBuilder: + """Return the action builder, creating it on first use. + + Returns: + Action ``ModalityDataBuilder`` for subsequent append operations. + """ + if self.action is None: + self.action = ModalityDataBuilder() + return self.action + + def ensure_sound(self) -> ModalityDataBuilder: + """Return the sound builder, creating it on first use. + + Returns: + Sound ``ModalityDataBuilder`` for subsequent append operations. + """ + if self.sound is None: + self.sound = ModalityDataBuilder() + return self.sound + + # Multi-control transfer: per-sample list of per-vision-item token counts. + # For a multi-control transfer sample with N controls + 1 noisy target, + # vision_item_split_lens[i] = [L_ctrl0, L_ctrl1, ..., L_ctrlN-1, L_noisy]. + # Used by cosmos3_vfm_network.py to derive gen-relative control/noisy ranges + # for multi_control_two_way_attention. + vision_item_split_lens: list[list[int]] = field(default_factory=list) + + # Per-sample per-control weights for multi-control weighted V-scaling. + # Parallel to vision_item_split_lens[i][:-1] (excludes noisy item). + # None for non-transfer or standard single-control samples. + control_weights: list[list[float]] | None = None + + @property + def mrope_temporal_offset(self) -> int | float: + """Current per-sample temporal cursor used for mRoPE position generation.""" + return self._mrope_temporal_offset + + def set_mrope_temporal_offset(self, temporal_offset: int | float) -> None: + """Set the per-sample temporal cursor used for mRoPE position generation. + + Args: + temporal_offset: New mRoPE temporal cursor value. + """ + self._mrope_temporal_offset = temporal_offset + + def advance_mrope_temporal_offset(self, delta: int | float) -> None: + """Advance the per-sample temporal cursor by ``delta``. + + Args: + delta: Amount to add to the current mRoPE temporal cursor. + """ + self._mrope_temporal_offset += delta + + def begin_sample(self, initial_mrope_temporal_offset: int | float) -> None: + """Reset per-sample position cursors before appending sample tokens. + + Args: + initial_mrope_temporal_offset: Initial mRoPE temporal cursor for this sample. + """ + self._mrope_temporal_offset = initial_mrope_temporal_offset + + def pack_text_tokens( + self, + text_ids: list[int], + special_tokens: dict[str, int], + has_generation: bool, + use_float_positions: bool = False, + ) -> int: + """Pack text tokens into the sequence. + + Args: + text_ids: List of text token IDs (integers). + special_tokens: Dictionary of special token IDs. + has_generation: Whether there's media/action after text. + use_float_positions: If True, generate float position IDs for 3D mRoPE + (for consistency with FPS-modulated vision tokens). + + Returns: + Text sample length. + """ + # Prepend BOS token if available + if "bos_token_id" in special_tokens: + shifted_text_ids = [special_tokens["bos_token_id"]] + text_ids + else: + shifted_text_ids = text_ids + + span_token_ids = shifted_text_ids + [special_tokens["eos_token_id"]] + # Add start-of-generation token, but only if there's media/action present. + if has_generation: + span_token_ids.append(special_tokens["start_of_generation"]) + + split_len = len(span_token_ids) + expected_split_len = len(text_ids) + if "bos_token_id" in special_tokens: + expected_split_len += 1 + expected_split_len += 1 # EOS + if has_generation: + expected_split_len += 1 # start-of-generation / BOV + assert split_len == expected_split_len + + position_ids, self._mrope_temporal_offset = get_3d_mrope_ids_text_tokens( + num_tokens=split_len, + temporal_offset=self._mrope_temporal_offset, + use_float_positions=use_float_positions, + ) # position_ids: [3,split_len] + + span_start, _ = self.append_text_span(span_token_ids, position_ids) + + # Configure loss computation for text tokens + self.ce_loss_indexes.extend(range(span_start, span_start + len(shifted_text_ids))) + self.ce_loss_weights.extend([1.0] * len(shifted_text_ids)) + self.label_ids.extend(text_ids[1:] + [special_tokens["eos_token_id"]]) + + self.attn_modes.append("causal") + self.split_lens.append(split_len) + + return split_len + + def pack_vision_tokens( + self, + input_vision_tokens: torch.Tensor, + condition_frame_indexes_vision: list[int], + input_timestep: float | torch.Tensor, + latent_patch_size: int = 1, + vision_fps: float | None = None, + enable_fps_modulation: bool = False, + base_fps: float = 24.0, + temporal_compression_factor: int = 4, + vision_temporal_positions: torch.Tensor | None = None, + ) -> int: + """Pack vision tokens into the sequence. + + Args: + input_vision_tokens: Vision latent tokens (C, T, H, W). + condition_frame_indexes_vision: Indexes of conditioning frames. + input_timestep: Diffusion timestep. Either a float (teacher_forcing/none — all frames + share the same sigma) or a Tensor(T_max,) (diffusion_forcing — per-frame sigma; + indexed as input_timestep[frame_idx] for each noisy frame). + latent_patch_size: Patch size for latent patchification. + vision_fps: Frames per second of the video. Used when enable_fps_modulation=True. + enable_fps_modulation: If True, scale temporal position IDs based on video FPS. + base_fps: Base FPS for normalization (default 24.0). + temporal_compression_factor: VAE temporal compression factor (default 4). + vision_temporal_positions: Optional explicit temporal coordinate per latent + frame, shape ``(T,)``. Used by UniAE to account for kept boundary latents. + + Returns: + Vision split length. + """ + vision = self.ensure_vision() + + # Compute position IDs for image patches + _, _, latent_t, latent_h, latent_w = input_vision_tokens.shape + if latent_patch_size < 1: + raise ValueError(f"latent_patch_size must be >= 1, got {latent_patch_size}") + # Use ceil to support latent dims not divisible by patch size (padding handled in network) + patch_h = math.ceil(latent_h / latent_patch_size) + patch_w = math.ceil(latent_w / latent_patch_size) + vision.token_shapes.append((latent_t, patch_h, patch_w)) + vision.tokens.append(input_vision_tokens) + vision_payload_index = len(vision.tokens) - 1 + + # Supervise vision tokens based on conditioning frames + condition_set = {idx for idx in condition_frame_indexes_vision if 0 <= idx < latent_t} + + vision_condition_mask = torch.zeros( + (latent_t, 1, 1), device=input_vision_tokens.device, dtype=input_vision_tokens.dtype + ) # [T,1,1] + for frame_idx in condition_set: + vision_condition_mask[frame_idx, 0, 0] = 1.0 + vision.condition_mask.append(vision_condition_mask) + + vision_noisy_frame_indexes = torch.tensor( + [idx for idx in range(latent_t) if idx not in condition_set], + device=input_vision_tokens.device, + dtype=torch.long, + ) # [N_noisy_frames] + vision.noisy_frame_indexes.append(vision_noisy_frame_indexes) + + frame_token_stride = patch_h * patch_w + effective_fps = vision_fps if enable_fps_modulation else None + if vision_temporal_positions is not None: + vision_temporal_positions = vision_temporal_positions.to(device="cpu", dtype=torch.float32) # [T] + + vision_mrope_ids, self._mrope_temporal_offset = get_3d_mrope_ids_vae_tokens( + grid_t=latent_t, + grid_h=patch_h, + grid_w=patch_w, + temporal_offset=self._mrope_temporal_offset, + reset_spatial_indices=self._mrope_reset_spatial, + fps=effective_fps, + base_fps=base_fps, + temporal_compression_factor=temporal_compression_factor, + temporal_positions=vision_temporal_positions, + actual_temporal_compression_factor=temporal_compression_factor, + ) # vision_mrope_ids: [3,N_vision_tokens] + vision_mrope_ids = vision_mrope_ids.reshape(3, latent_t, frame_token_stride) # [3,T,H*W] + + vision_split_len = 0 + for frame_idx in range(latent_t): + position_ids = vision_mrope_ids[:, frame_idx, :] # [3,H*W] + frame_indexes = self.append_vision_span( + frame_token_stride, + position_ids, + payload_index=vision_payload_index, + payload_start=frame_idx * frame_token_stride, + payload_shape=(1, patch_h, patch_w), + ) + vision_split_len += frame_token_stride + + if frame_idx in condition_set: + continue + vision.mse_loss_indexes.extend(frame_indexes) + if isinstance(input_timestep, torch.Tensor): + frame_ts = input_timestep[frame_idx].item() + else: + frame_ts = input_timestep + vision.timesteps.extend([frame_ts] * frame_token_stride) + + return vision_split_len + + def pack_action_tokens( + self, + input_action_tokens: torch.Tensor, + condition_frame_indexes_action: list[int], + input_timestep: float, + action_temporal_offset: int | float = 0, + enable_fps_modulation: bool = False, + base_fps: float = 24.0, + action_fps: float | None = None, + base_temporal_compression_factor: int | None = None, + action_start_frame_offset: int = 1, + ) -> int: + """Pack action tokens into the sequence. + + Args: + input_action_tokens: Action latent tokens (T, D). + condition_frame_indexes_action: Indexes of conditioning action steps. + input_timestep: Diffusion timestep. + action_temporal_offset: Temporal offset for action mRoPE IDs (typically + the vision start offset so action aligns temporally with vision). + enable_fps_modulation: If True, scale temporal position IDs based on FPS. + base_fps: Base FPS for normalization (default 24.0). + action_fps: Frames per second of the action data. Used when enable_fps_modulation=True. + base_temporal_compression_factor: Base temporal compression factor for FPS scaling. + Should be set to the vision temporal compression factor (e.g. 4) so that action + tokens advance at frame rate (4x finer) relative to vision latent frames. + Only affects behavior when FPS modulation is enabled. + action_start_frame_offset: Frame offset for aligning action[0] with the + corresponding vision frame. Default 1 aligns action[0] with vision frame 1. + + Returns: + Number of action tokens added. + """ + action_split_len = input_action_tokens.shape[0] + + action = self.ensure_action() + + # Add token indexes and loss information + action.token_shapes.append((action_split_len,)) + action.tokens.append(input_action_tokens) + action_payload_index = len(action.tokens) - 1 + + condition_set = {idx for idx in condition_frame_indexes_action if 0 <= idx < action_split_len} + + action_condition_mask = torch.zeros( + (action_split_len, 1), device=input_action_tokens.device, dtype=input_action_tokens.dtype + ) # [T_action,1] + for frame_idx in condition_set: + action_condition_mask[frame_idx, 0] = 1.0 + action.condition_mask.append(action_condition_mask) + + action_noisy_frame_indexes = torch.tensor( + [idx for idx in range(action_split_len) if idx not in condition_set], + device=input_action_tokens.device, + dtype=torch.long, + ) # [N_noisy_action_frames] + action.noisy_frame_indexes.append(action_noisy_frame_indexes) + + # Action tokens use a 1x1 spatial grid with start_frame_offset=1 by default, + # so action[0] (null token) aligns with vision frame 1, not frame 0. + effective_fps = action_fps if enable_fps_modulation else None + + action_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( + grid_t=action_split_len, + grid_h=1, + grid_w=1, + temporal_offset=action_temporal_offset, + reset_spatial_indices=self._mrope_reset_spatial, + fps=effective_fps, + base_fps=base_fps, + temporal_compression_factor=1, # Action is at frame rate (no temporal compression) + base_temporal_compression_factor=base_temporal_compression_factor, + start_frame_offset=action_start_frame_offset, # Align action[0] with vision frame action_start_frame_offset + ) # action_mrope_ids: [3,N_action_tokens] + # Note: we don't update _mrope_temporal_offset here because action tokens + # share the temporal space with vision tokens (they run in parallel). + + for frame_idx in range(action_split_len): + position_ids = action_mrope_ids[:, frame_idx : frame_idx + 1] # [3,1] + frame_indexes = self.append_action_span( + 1, + position_ids, + payload_index=action_payload_index, + payload_start=frame_idx, + payload_shape=(1,), + ) + + if frame_idx in condition_set: + continue + action.mse_loss_indexes.extend(frame_indexes) + action.timesteps.extend([input_timestep]) + + return action_split_len + + def pack_sound_tokens( + self, + input_sound_tokens: torch.Tensor, + condition_frame_indexes_sound: list[int], + input_timestep: float, + sound_temporal_offset: int | float = 0, + enable_fps_modulation: bool = False, + base_fps: float = 24.0, + sound_fps: float | None = None, + sound_base_temporal_compression_factor: int | None = None, + ) -> int: + """Pack sound/audio tokens into the sequence. + + Sound latents have shape [C, T] where C is channels and T is temporal frames. + Sound tokens are added to the unified generation split to maintain SequencePack's + 2-split invariant (causal + full). + + Args: + input_sound_tokens: Sound latent tokens (C, T). + condition_frame_indexes_sound: Indexes of conditioning frames. + [] means all frames are noised/supervised. + All frames specified means all frames are clean (no MSE supervision). + input_timestep: Diffusion timestep. + sound_temporal_offset: Temporal offset for m-RoPE position IDs (aligned with vision start). + enable_fps_modulation: If True, scale temporal positions by FPS ratio. + base_fps: Base FPS for normalization (default 24.0). + sound_fps: Sound latent FPS (e.g., 25.0). Used for FPS-aware m-RoPE positions. + sound_base_temporal_compression_factor: Base temporal compression factor for sound FPS scaling. + ``None`` preserves the current behavior where sound advances at ``base_fps`` positions/sec. + + Returns: + Number of sound tokens added. + """ + # Sound latent shape: [C, T] → T tokens + _, sound_split_len = input_sound_tokens.shape + + sound = self.ensure_sound() + + # Add token indexes - sound uses (T, 1, 1) shape for compatibility with 3D RoPE + sound.token_shapes.append((sound_split_len, 1, 1)) + sound.tokens.append(input_sound_tokens) + sound_payload_index = len(sound.tokens) - 1 + + # Supervise sound tokens based on conditioning frames + condition_set = {idx for idx in condition_frame_indexes_sound if 0 <= idx < sound_split_len} + + # Condition mask: shape (T, 1) — 1 = clean/conditioning, 0 = noised/supervised + sound_condition_mask = torch.zeros( + (sound_split_len, 1), device=input_sound_tokens.device, dtype=input_sound_tokens.dtype + ) # [T_sound,1] + for frame_idx in condition_set: + sound_condition_mask[frame_idx, 0] = 1.0 + sound.condition_mask.append(sound_condition_mask) + + sound_noisy_frame_indexes = torch.tensor( + [idx for idx in range(sound_split_len) if idx not in condition_set], + device=input_sound_tokens.device, + dtype=torch.long, + ) # [N_noisy_sound_frames] + sound.noisy_frame_indexes.append(sound_noisy_frame_indexes) + + # Sound tokens use a 1x1 spatial grid, aligned with vision temporal positions. + # sound[0] aligns with vision frame 0 (start_frame_offset=0, unlike action which offsets by 1). + effective_fps = sound_fps if enable_fps_modulation else None + + sound_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( + grid_t=sound_split_len, + grid_h=1, + grid_w=1, + temporal_offset=sound_temporal_offset, + reset_spatial_indices=self._mrope_reset_spatial, + fps=effective_fps, + base_fps=base_fps, + temporal_compression_factor=1, # Sound latent is already at sound_latent_fps (no further compression) + base_temporal_compression_factor=sound_base_temporal_compression_factor, + start_frame_offset=0, # Sound[0] aligns with vision frame 0 + ) # sound_mrope_ids: [3,N_sound_tokens] + # Note: we don't update _mrope_temporal_offset here because sound tokens + # share the temporal space with vision tokens (they run in parallel). + + # Add to MSE loss indexes and timesteps for non-conditioning frames + for frame_idx in range(sound_split_len): + position_ids = sound_mrope_ids[:, frame_idx : frame_idx + 1] # [3,1] + frame_indexes = self.append_sound_span( + 1, + position_ids, + payload_index=sound_payload_index, + payload_start=frame_idx, + payload_shape=(1, 1, 1), + ) + + if frame_idx in condition_set: + continue + sound.mse_loss_indexes.extend(frame_indexes) + sound.timesteps.extend([input_timestep]) + + return sound_split_len + + def _append_position_ids(self, position_ids: torch.Tensor, span_len: int) -> None: + """Append one block of position IDs. + + Args: + position_ids: Position ID tensor with shape ``[3, span_len]``. + span_len: Number of token positions described by ``position_ids``. + """ + assert position_ids.shape[-1] == span_len, ( + f"position_ids last dimension must match span_len={span_len}; got {tuple(position_ids.shape)}" + ) + self.position_ids.append(position_ids) + + def _append_modality_span( + self, + modality: ModalityDataBuilder, + span_len: int, + position_ids: torch.Tensor, + *, + payload_index: int, + payload_start: int, + payload_shape: tuple[int, ...], + ) -> list[int]: + """Append one modality span and return its global sequence indexes. + + Args: + modality: Modality builder receiving sequence indexes and span metadata. + span_len: Number of contiguous tokens to append. + position_ids: Position ID tensor with shape ``[3, span_len]``. + payload_index: Index into ``modality.tokens`` for the backing payload tensor. + payload_start: First flattened token offset within the backing payload tensor. + payload_shape: Logical payload slice shape represented by this span. + + Returns: + Global packed-sequence indexes covered by the appended span. + """ + assert payload_index >= 0, f"payload_index must be non-negative, got {payload_index}" + assert payload_start >= 0, f"payload_start must be non-negative, got {payload_start}" + payload_len = 1 + for dim in payload_shape: + assert dim > 0, f"payload_shape dimensions must be positive, got {payload_shape}" + payload_len *= dim + assert payload_len == span_len, ( + f"payload_shape={payload_shape} describes {payload_len} tokens, but span_len={span_len}" + ) + + span_start = self.current_seq_index + span_indexes = list(range(span_start, span_start + span_len)) + modality.spans.append( + ModalitySpan( + sequence_start=span_start, + sequence_len=span_len, + payload_index=payload_index, + payload_start=payload_start, + payload_len=payload_len, + payload_shape=payload_shape, + ) + ) + modality.sequence_indexes.extend(span_indexes) + self._append_position_ids(position_ids, span_len) + self.current_seq_index += span_len + return span_indexes + + def append_text_span( + self, + token_ids: list[int], + position_ids: torch.Tensor, + ) -> tuple[int, int]: + """Append one contiguous text span. + + Args: + token_ids: Text token IDs to append. + position_ids: Position ID tensor with shape ``[3, len(token_ids)]``. + + Returns: + Tuple of ``(start_index, span_len)`` for the appended span. + """ + span_start = self.current_seq_index + span_len = len(token_ids) + self.text_ids.extend(token_ids) + self.text_indexes.extend(range(span_start, span_start + span_len)) + self._append_position_ids(position_ids, span_len) + self.current_seq_index += span_len + return span_start, span_len + + def append_vision_span( + self, + span_len: int, + position_ids: torch.Tensor, + *, + payload_index: int, + payload_start: int, + payload_shape: tuple[int, int, int], + ) -> list[int]: + """Append one contiguous vision span. + + Args: + span_len: Number of contiguous vision tokens to append. + position_ids: Position ID tensor with shape ``[3, span_len]``. + payload_index: Index into ``vision.tokens`` for the backing payload tensor. + payload_start: First flattened token offset within the backing vision payload. + payload_shape: Logical vision slice shape, usually ``(1, patch_h, patch_w)``. + + Returns: + Global packed-sequence indexes covered by the appended vision span. + """ + return self._append_modality_span( + self.ensure_vision(), + span_len, + position_ids, + payload_index=payload_index, + payload_start=payload_start, + payload_shape=payload_shape, + ) + + def append_action_span( + self, + span_len: int, + position_ids: torch.Tensor, + *, + payload_index: int, + payload_start: int, + payload_shape: tuple[int], + ) -> list[int]: + """Append one contiguous action span. + + Args: + span_len: Number of contiguous action tokens to append. + position_ids: Position ID tensor with shape ``[3, span_len]``. + payload_index: Index into ``action.tokens`` for the backing payload tensor. + payload_start: First flattened token offset within the backing action payload. + payload_shape: Logical action slice shape, usually ``(span_len,)``. + + Returns: + Global packed-sequence indexes covered by the appended action span. + """ + return self._append_modality_span( + self.ensure_action(), + span_len, + position_ids, + payload_index=payload_index, + payload_start=payload_start, + payload_shape=payload_shape, + ) + + def append_sound_span( + self, + span_len: int, + position_ids: torch.Tensor, + *, + payload_index: int, + payload_start: int, + payload_shape: tuple[int, int, int], + ) -> list[int]: + """Append one contiguous sound span. + + Args: + span_len: Number of contiguous sound tokens to append. + position_ids: Position ID tensor with shape ``[3, span_len]``. + payload_index: Index into ``sound.tokens`` for the backing payload tensor. + payload_start: First flattened token offset within the backing sound payload. + payload_shape: Logical sound slice shape, usually ``(1, 1, 1)``. + + Returns: + Global packed-sequence indexes covered by the appended sound span. + """ + return self._append_modality_span( + self.ensure_sound(), + span_len, + position_ids, + payload_index=payload_index, + payload_start=payload_start, + payload_shape=payload_shape, + ) + + def append_end_of_generation_token( + self, + token_id: int, + use_float_mrope_positions: bool, + ) -> int: + """Append an end-of-generation text token. + + Args: + token_id: Token ID for the end-of-generation marker. + use_float_mrope_positions: Whether to write float position IDs for this token. + + Returns: + Span length for the appended end-of-generation token. + """ + eov_dtype = torch.float32 if use_float_mrope_positions else torch.long + position_ids = torch.full((3, 1), self._mrope_temporal_offset, dtype=eov_dtype) # [3,1] + self._mrope_temporal_offset += 1 + + _, span_len = self.append_text_span([token_id], position_ids) + return span_len + + def finish_sample(self, generation_split_len: int, sample_len: int) -> None: + """Record the non-causal generation split and total sample length. + + Args: + generation_split_len: Length of the full-attention generation split. + sample_len: Total number of tokens appended for this sample. + """ + self.attn_modes.append("full") + self.split_lens.append(generation_split_len) + self.sample_lens.append(sample_len) + + def _finalize_modality( + self, + modality: ModalityDataBuilder | None, + *, + domain_id: list[torch.Tensor] | None = None, + raw_action_dim: list[torch.Tensor | None] | None = None, + include_raw_action_dim: bool = False, + ) -> ModalityData | None: + """Finalize one modality builder into model-facing modality data. + + Args: + modality: Modality builder to finalize, or ``None`` if no tokens were appended. + domain_id: Optional action domain IDs to attach to finalized action data. + raw_action_dim: Optional raw action dimensions for action-channel masking. + include_raw_action_dim: Whether to include ``raw_action_dim`` in the finalized data. + + Returns: + Finalized ``ModalityData`` or ``None`` when the modality has no sequence indexes. + """ + if modality is None or len(modality.sequence_indexes) == 0: + return None + + kwargs = { + "sequence_indexes": torch.tensor(modality.sequence_indexes, dtype=torch.long), # [N_modality_tokens] + "timesteps": torch.tensor(modality.timesteps, dtype=torch.float32), # [N_modality_noisy_tokens] + "mse_loss_indexes": torch.tensor(modality.mse_loss_indexes, dtype=torch.long), # [N_modality_noisy_tokens] + "spans": list(modality.spans), + "token_shapes": list(modality.token_shapes), + "tokens": modality.tokens, + "condition_mask": list(modality.condition_mask), + "noisy_frame_indexes": list(modality.noisy_frame_indexes), + } + if domain_id is not None: + kwargs["domain_id"] = domain_id + if include_raw_action_dim: + kwargs["raw_action_dim"] = raw_action_dim + return ModalityData(**kwargs) + + def finalize( + self, + gen_data_clean: GenerationDataClean, + ) -> "PackedSequence": + """Convert all lists to tensors and compute derived values. + + Args: + gen_data_clean: GenerationDataClean for metadata (e.g., action domain IDs). + + Returns: + New PackedSequence instance with tensors instead of lists. + """ + # Compute sequence length + sequence_length = sum(self.sample_lens) + sample_lens = self.sample_lens.copy() + split_lens = self.split_lens.copy() + attn_modes = self.attn_modes.copy() + + # Prepare loss-related tensors (cross-entropy) + label_ids: torch.Tensor | None = None + ce_loss_indexes: torch.Tensor | None = None + ce_loss_weights: torch.Tensor | None = None + if self.label_ids and len(self.label_ids) > 0: + label_ids = torch.tensor(self.label_ids) # [N_ce_tokens] + ce_loss_indexes = torch.tensor(self.ce_loss_indexes) # [N_ce_tokens] + ce_loss_weights = torch.tensor(self.ce_loss_weights) # [N_ce_tokens] + + # The condition_mask and noisy_frame_indexes are kept as lists to support variable shapes. + + vision = self._finalize_modality(self.vision) + action_domain_id = None + if self.action is not None: + if gen_data_clean.action_domain_id is not None: + action_domain_id = gen_data_clean.action_domain_id + else: + default_action_domain_id = torch.zeros(1, dtype=torch.long) # [1] + action_domain_id = [default_action_domain_id] * len(self.action.token_shapes) + action = self._finalize_modality( + self.action, + domain_id=action_domain_id, + raw_action_dim=gen_data_clean.raw_action_dim, + include_raw_action_dim=True, + ) + sound = self._finalize_modality(self.sound) + + # Finalize position IDs. + assert isinstance(self.position_ids, list) + assert all(isinstance(position_id, torch.Tensor) for position_id in self.position_ids) + if len(self.position_ids) > 0: + position_ids = torch.cat(self.position_ids, dim=1) # [3,actual_seq_len] + else: + position_ids = torch.empty((3, 0), dtype=torch.long) # [3,0] + + return PackedSequence( + # Sequence structure + sequence_length=sequence_length, + sample_lens=sample_lens, + split_lens=split_lens, + attn_modes=attn_modes, + is_image_batch=gen_data_clean.is_image_batch, + # Text modality (converted to tensors) + text_ids=torch.tensor(self.text_ids, dtype=torch.long), # [N_text_tokens] + text_indexes=torch.tensor(self.text_indexes, dtype=torch.long), # [N_text_tokens] + position_ids=position_ids, # [3,seq_len] + # Loss computation - Cross Entropy + label_ids=label_ids, + ce_loss_indexes=ce_loss_indexes, + ce_loss_weights=ce_loss_weights, + # Generation modalities + vision=vision, + action=action, + sound=sound, + # Temporal causal + null_action_supertokens=self.null_action_supertokens, + num_action_tokens_per_supertoken=self.num_action_tokens_per_supertoken, + # Multi-control transfer + vision_item_split_lens=list(self.vision_item_split_lens), + control_weights=gen_data_clean.control_weights, + ) + + +@dataclass +class PackedSequence: + """Finalized model-facing packed sequence. + + The object remains mutable for model-side payload replacement, clean replay, + and in-place device transfer. Build-only cursor and mRoPE state live on + ``PackedSequenceBuilder``. + + Attributes: + sample_lens: Length of each sample in the packed sequence. + split_lens: Length of each attention split. Each sample contributes a causal + text split and a full generation split. + attn_modes: Attention mode for each split, such as ``"causal"`` or ``"full"``. + is_image_batch: Whether this batch contains images rather than videos. + sequence_length: Total length of the packed sequence. + text_ids: Tensor of all text token IDs, including special tokens. + text_indexes: Tensor of global packed-sequence indexes for text tokens. + position_ids: Tensor of mRoPE position IDs for all packed tokens with shape + ``[3, sequence_length]``. + label_ids: Optional tensor of label IDs for text cross-entropy loss. + ce_loss_indexes: Optional tensor of global packed-sequence indexes for text + cross-entropy loss. + ce_loss_weights: Optional tensor of per-token weights for text cross-entropy loss. + null_action_supertokens: Whether temporal-causal supertoken 0 contains null + action tokens. + num_action_tokens_per_supertoken: Number of action tokens prefixing each + temporal-causal vision supertoken. + vision: Finalized vision modality data, or ``None`` if no vision is present. + action: Finalized action modality data, or ``None`` if no action is present. + sound: Finalized sound modality data, or ``None`` if no sound is present. + vision_item_split_lens: Per-sample per-vision-item token counts for multi-control + transfer. + control_weights: Per-sample per-control weights for multi-control weighted V-scaling. + """ + + # Sequence structure + sample_lens: list[int] = field(default_factory=list) + split_lens: list[int] = field(default_factory=list) + attn_modes: list[str] = field(default_factory=list) + is_image_batch: bool = False + sequence_length: int = 0 + + # Text modality + text_ids: torch.Tensor = field(default_factory=_empty_long_tensor) + text_indexes: torch.Tensor = field(default_factory=_empty_long_tensor) + position_ids: torch.Tensor = field(default_factory=_empty_long_tensor) + + # Loss computation - Cross Entropy (text) + label_ids: torch.Tensor | None = None + ce_loss_indexes: torch.Tensor | None = None + ce_loss_weights: torch.Tensor | None = None + + # Temporal causal: whether supertoken 0's action slot contains null tokens. + # True for all training calls and AR frame 0; False for AR frame N>0 (real actions). + # Used by three_way_attention to zero out V for null action tokens (inline when attention_meta.null_action_supertokens=True). + null_action_supertokens: bool = False + + # Temporal causal: number of action tokens prefixing each vision supertoken. + # Equals temporal_compression_factor when actions are packed inline; 0 when + # action_gen=False or for non-temporal-causal layouts. Single source of truth + # for downstream attention/KV-cache code (per-supertoken layout is + # num_action_tokens_per_supertoken + H_p * W_p). + num_action_tokens_per_supertoken: int = 0 + + # Generation modalities - NAMED FIELDS for type safety + vision: ModalityData | None = None + action: ModalityData | None = None + sound: ModalityData | None = None + + # Multi-control transfer: per-sample list of per-vision-item token counts. + # For a multi-control transfer sample with N controls + 1 noisy target, + # vision_item_split_lens[i] = [L_ctrl0, L_ctrl1, ..., L_ctrlN-1, L_noisy]. + # Used by cosmos3_vfm_network.py to derive gen-relative control/noisy ranges + # for multi_control_two_way_attention. + vision_item_split_lens: list[list[int]] = field(default_factory=list) + + # Per-sample per-control weights for multi-control weighted V-scaling. + # Parallel to vision_item_split_lens[i][:-1] (excludes noisy item). + # None for non-transfer or standard single-control samples. + control_weights: list[list[float]] | None = None + + def __post_init__(self) -> None: + assert isinstance(self.text_ids, torch.Tensor), "PackedSequence.text_ids must be finalized" + assert isinstance(self.text_indexes, torch.Tensor), "PackedSequence.text_indexes must be finalized" + assert isinstance(self.position_ids, torch.Tensor), "PackedSequence.position_ids must be finalized" + if self.label_ids is not None: + assert isinstance(self.label_ids, torch.Tensor), "PackedSequence.label_ids must be finalized" + if self.ce_loss_indexes is not None: + assert isinstance(self.ce_loss_indexes, torch.Tensor), "PackedSequence.ce_loss_indexes must be finalized" + if self.ce_loss_weights is not None: + assert isinstance(self.ce_loss_weights, torch.Tensor), "PackedSequence.ce_loss_weights must be finalized" + for modality in [self.vision, self.action, self.sound]: + assert modality is None or isinstance(modality, ModalityData), ( + "PackedSequence modality fields must be finalized ModalityData" + ) + + def to_cuda(self) -> None: + """Move all tensor fields to CUDA in-place.""" + self.text_ids = self.text_ids.cuda() + self.text_indexes = self.text_indexes.cuda() + self.position_ids = self.position_ids.cuda() + if isinstance(self.label_ids, torch.Tensor): + self.label_ids = self.label_ids.cuda() + if isinstance(self.ce_loss_indexes, torch.Tensor): + self.ce_loss_indexes = self.ce_loss_indexes.cuda() + if isinstance(self.ce_loss_weights, torch.Tensor): + self.ce_loss_weights = self.ce_loss_weights.cuda() + if self.vision is not None: + self.vision.to_cuda() + if self.action is not None: + self.action.to_cuda() + if self.sound is not None: + self.sound.to_cuda() + + +@dataclass +class SequencePlan: + """Plan describing which modalities are present in a sample. + + This dataclass tracks the presence of different modalities (text, vision, action) + and their conditioning configurations for a dataset sample. Unlike SequencePlan + which holds the actual tensor data, this class provides a lightweight summary + of what modalities exist and how they should be conditioned. + + Attributes: + has_text: Whether text/caption tokens are present for this sample. + Used for text-conditioned generation (e.g., text-to-image/video). + has_vision: Whether vision input (image or video latents) is present. + Defaults to False. + condition_frame_indexes_vision: Indexes of latent vision frames that are clean/conditioning. + [] means all frames are noised/supervised. + All frames specified means all frames are clean (no MSE supervision). + For multi-item samples (e.g. image editing where each sample has multiple + separately-encoded images), this applies to each vision item individually. + The number of items per sample is tracked by + ``GenerationDataClean.num_vision_items_per_sample``. + share_vision_temporal_positions: Whether all vision items in this sample share + the same temporal mRoPE grid. + has_action: Whether action input is present for robotics/embodied AI tasks. + Defaults to False. + condition_frame_indexes_action: Indexes of action steps that are clean/conditioning. + [] means all steps are noised/supervised. + All steps specified means all steps are clean (no MSE supervision). + action_start_frame_offset: Frame offset for aligning action[0] to vision frames. + has_sound: Whether sound/audio input is present. + condition_frame_indexes_sound: Indexes of sound frames that are clean/conditioning. + [] means all frames are noised/supervised. + All frames specified means all frames are clean (no MSE supervision). + """ + + # -- understanding (text conditioning) -- + has_text: bool + + # -- vision modality -- + has_vision: bool = False + condition_frame_indexes_vision: list[int] = field(default_factory=list) + # If True, all vision items in this sample share the same temporal mRoPE grid + # (controlnet-style transfer: target frame i is spatio-temporally aligned with + # control frame i). Each item gets the same temporal_offset; spatial reset + # behavior is unchanged. Requires num_vision_items_per_sample > 1, equal latent_t, + # and equal fps across items. Default False preserves single-clip and + # image-editing semantics where items represent distinct time states. + share_vision_temporal_positions: bool = False + + # -- action modality -- + has_action: bool = False + condition_frame_indexes_action: list[int] = field(default_factory=list) + action_start_frame_offset: int = 1 + + # -- sound modality -- + has_sound: bool = False + condition_frame_indexes_sound: list[int] = field(default_factory=list) + + def as_dict(self) -> dict: + return { + "has_text": self.has_text, + "has_vision": self.has_vision, + "has_action": self.has_action, + "has_sound": self.has_sound, + "condition_frame_indexes_vision": self.condition_frame_indexes_vision, + "condition_frame_indexes_action": self.condition_frame_indexes_action, + "condition_frame_indexes_sound": self.condition_frame_indexes_sound, + "share_vision_temporal_positions": self.share_vision_temporal_positions, + } + + +def build_sequence_plans_from_data_batch( + data_batch: dict, + input_video_key, + input_image_key: str, +) -> list[SequencePlan]: + """Build or retrieve sequence plans from a data batch dictionary. + + This function extracts sequence plans from the data batch if they exist, + otherwise creates default SequencePlan objects for each sample + in the batch. + + Args: + data_batch: Dictionary containing the data batch from the dataloader. + Expected keys include 'video' or other tensors to determine batch size. + If 'sequence_plan' key exists, those plans are returned directly. + input_video_key: Data-batch key used to find video tensors when inferring batch size. + input_image_key: Data-batch key used to find image tensors when inferring batch size. + + Returns: + List of SequencePlan objects, one per sample in the batch. + """ + # For new modalities, please generate the sequence_plan in the dataset class!!!! + + # If sequence_plan already exists in data_batch, return it + if "sequence_plan" in data_batch: + return data_batch["sequence_plan"] + + assert "action" not in data_batch or data_batch["action"] is None, "Action data SHOULD have sequence_plans!" + assert "sound" not in data_batch or data_batch["sound"] is None, "Sound data SHOULD have sequence_plans!" + + # Determine batch size from available tensors + batch_size = 0 + for key in [input_video_key, input_image_key]: + if key in data_batch: + val = data_batch[key] + if isinstance(val, torch.Tensor): + batch_size = val.shape[0] + break + elif isinstance(val, list): + batch_size = len(val) + break + + if batch_size == 0: + raise ValueError( + f"Cannot determine batch size from data_batch. Expected {input_video_key}, {input_image_key}, or similar key." + ) + + # Build default SequencePlan objects + return [ + SequencePlan( + has_text=True, # Has text prompt! + has_vision=True, + condition_frame_indexes_vision=[], # No conditioning frames! + ) + for _ in range(batch_size) + ] diff --git a/cosmos_framework/data/generator/sequence_packing/temporal_causal.py b/cosmos_framework/data/generator/sequence_packing/temporal_causal.py index bcbe26d2..5a4ee07e 100644 --- a/cosmos_framework/data/generator/sequence_packing/temporal_causal.py +++ b/cosmos_framework/data/generator/sequence_packing/temporal_causal.py @@ -8,11 +8,11 @@ import torch from cosmos_framework.data.generator.sequence_packing.mrope import get_3d_mrope_ids_vae_tokens -from cosmos_framework.data.generator.sequence_packing.types import ModalityData, PackedSequence +from cosmos_framework.data.generator.sequence_packing.sequence import PackedSequenceBuilder def pack_supertokens_temporal_causal( - packed_seq: "PackedSequence", + seq_builder: PackedSequenceBuilder, input_vision_tokens: torch.Tensor, input_action_tokens: torch.Tensor | None, condition_frame_indexes_vision: list[int], @@ -53,35 +53,37 @@ def pack_supertokens_temporal_causal( ``input_timestep`` is float (TF/none) or Tensor(T_max,) (DF, per-frame sigma). Conditioning frames are excluded from mse_loss_indexes either way. - Returns (total_split_len, null_action_flag); null_action_flag is False when - pack_action_tokens=False. + Args: + seq_builder: Mutable sequence builder receiving packed spans and metadata. + input_vision_tokens: Vision latent tokens with shape ``[1, C, T, H, W]``. + input_action_tokens: Optional action tokens. Whole-clip training uses + ``(T - 1) * temporal_compression_factor`` rows; AR chunks use + ``T * temporal_compression_factor`` rows. + condition_frame_indexes_vision: Vision frame indexes treated as clean conditioning. + input_timestep: Diffusion timestep as a scalar float or per-frame tensor. + latent_patch_size: Spatial patch size used to derive the vision patch grid. + temporal_compression_factor: Number of frame-rate action tokens per latent + vision frame. + action_dim: Action feature dimension used when null action tokens are created. + vision_fps: Optional video FPS for FPS-modulated vision mRoPE positions. + action_fps: Optional action FPS for FPS-modulated action mRoPE positions. + enable_fps_modulation: If True, scale temporal position IDs using FPS values. + base_fps: Base FPS for temporal position normalization. + pack_action_tokens: If True, append action tokens before each vision supertoken. + If False, append only vision tokens and report no null-action supertokens. + + Returns: + Tuple of ``(total_split_len, null_action_flag)``. ``null_action_flag`` is False + when ``pack_action_tokens=False``. """ - assert isinstance(packed_seq.position_ids, list), "PackedSequence must be in build mode" - _, _, latent_t, latent_h, latent_w = input_vision_tokens.shape patch_h = math.ceil(latent_h / latent_patch_size) patch_w = math.ceil(latent_w / latent_patch_size) tcf = temporal_compression_factor patches_per_frame = patch_h * patch_w - supertoken_len = tcf + patches_per_frame if pack_action_tokens else patches_per_frame # S - - # Initialize modalities if needed - if packed_seq.vision is None: - packed_seq.vision = ModalityData() - if pack_action_tokens and packed_seq.action is None: - packed_seq.action = ModalityData() - - assert isinstance(packed_seq.vision.sequence_indexes, list) - assert isinstance(packed_seq.vision.mse_loss_indexes, list) - assert isinstance(packed_seq.vision.timesteps, list) - assert isinstance(packed_seq.vision.tokens, list) - assert isinstance(packed_seq.vision.condition_mask, list) - if pack_action_tokens: - assert isinstance(packed_seq.action.sequence_indexes, list) - assert isinstance(packed_seq.action.mse_loss_indexes, list) - assert isinstance(packed_seq.action.timesteps, list) - assert isinstance(packed_seq.action.tokens, list) - assert isinstance(packed_seq.action.condition_mask, list) + + vision = seq_builder.ensure_vision() + action = seq_builder.ensure_action() if pack_action_tokens else None device = input_vision_tokens.device dtype = input_vision_tokens.dtype @@ -134,40 +136,42 @@ def pack_supertokens_temporal_causal( null_action_flag = False # Record vision token shapes and tokens - packed_seq.vision.token_shapes.append((latent_t, patch_h, patch_w)) - packed_seq.vision.tokens.append(input_vision_tokens) + vision.token_shapes.append((latent_t, patch_h, patch_w)) + vision.tokens.append(input_vision_tokens) + vision_payload_index = len(vision.tokens) - 1 # Vision conditioning mask: (T, 1, 1) condition_set_vision = {idx for idx in condition_frame_indexes_vision if 0 <= idx < latent_t} vision_condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=dtype) # [T,1,1] for fidx in condition_set_vision: vision_condition_mask[fidx, 0, 0] = 1.0 - packed_seq.vision.condition_mask.append(vision_condition_mask) + vision.condition_mask.append(vision_condition_mask) vision_noisy_frame_indexes = torch.tensor( [idx for idx in range(latent_t) if idx not in condition_set_vision], device=device, dtype=torch.long, ) # [N_noisy_frames] - packed_seq.vision.noisy_frame_indexes.append(vision_noisy_frame_indexes) + vision.noisy_frame_indexes.append(vision_noisy_frame_indexes) if pack_action_tokens: + assert action is not None # Action token shapes: latent_t * tcf total (including null tokens) - packed_seq.action.token_shapes.append((latent_t * tcf,)) - packed_seq.action.tokens.append(all_action_tokens) + action.token_shapes.append((latent_t * tcf,)) + action.tokens.append(all_action_tokens) + action_payload_index = len(action.tokens) - 1 # Action conditioning mask: all action tokens are conditioning (not supervised) # Null tokens are always conditioning; real actions are conditioning too (they are inputs) action_condition_mask = torch.ones((latent_t * tcf, 1), device=device, dtype=dtype) # [T*tcf,1] - packed_seq.action.condition_mask.append(action_condition_mask) + action.condition_mask.append(action_condition_mask) # Pack in interleaved supertoken order: [action_t, vision_t] for each frame t # (or just [vision_t] per frame when pack_action_tokens=False) - curr = packed_seq.curr total_split_len = 0 # Snapshot the offset before this sample and compute mRoPE IDs. - temporal_offset = packed_seq._mrope_temporal_offset + temporal_offset = seq_builder._mrope_temporal_offset effective_vision_fps = vision_fps if enable_fps_modulation else None # AR generation (single frame OR chunk) is detected by every frame carrying a @@ -187,13 +191,15 @@ def pack_supertokens_temporal_causal( grid_h=patch_h, grid_w=patch_w, temporal_offset=temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, + reset_spatial_indices=seq_builder._mrope_reset_spatial, fps=effective_vision_fps, base_fps=base_fps, temporal_compression_factor=tcf, start_frame_offset=vision_sfo, ) # vision_ids_flat: [3,T*patch_h*patch_w] + vision_ids_3d = vision_ids_flat.reshape(3, latent_t, patches_per_frame) # [3,T,patch_h*patch_w] + action_ids_3d: torch.Tensor | None = None if pack_action_tokens: effective_action_fps = action_fps if enable_fps_modulation else None @@ -214,7 +220,7 @@ def _real_action_ids(n_frames: int, start_frame_offset: int) -> torch.Tensor: grid_h=1, grid_w=1, temporal_offset=temporal_offset, - reset_spatial_indices=packed_seq._mrope_reset_spatial, + reset_spatial_indices=seq_builder._mrope_reset_spatial, fps=effective_action_fps, base_fps=base_fps, temporal_compression_factor=1, @@ -242,40 +248,39 @@ def _real_action_ids(n_frames: int, start_frame_offset: int) -> torch.Tensor: # AR frame 0 / image2video (latent_t == 1, no action): only null. action_ids_3d = null_ids.reshape(3, 1, tcf) # [3,1,tcf] - # (3, T*H*W) -> (3, T, H*W) - vision_ids_3d = vision_ids_flat.reshape(3, latent_t, patches_per_frame) # [3,T,patch_h*patch_w] - - # Interleave per frame: (3, T, tcf+H*W) -> (3, T*S) - interleaved_ids = torch.cat([action_ids_3d, vision_ids_3d], dim=2).reshape( - 3, latent_t * supertoken_len - ) # [3,T*S] - packed_seq.position_ids.append(interleaved_ids) - else: - # No action tokens: just vision IDs, already in (3, T*H*W) order. - packed_seq.position_ids.append(vision_ids_flat) - - packed_seq._mrope_temporal_offset = new_offset + seq_builder._mrope_temporal_offset = new_offset for frame_t in range(latent_t): if pack_action_tokens: - # Pack action tokens for this frame (indexes only; tokens already stored in packed_seq.action.tokens) - action_indexes = list(range(curr, curr + tcf)) - packed_seq.action.sequence_indexes.extend(action_indexes) + assert action is not None + assert action_ids_3d is not None + # Pack action tokens for this frame (indexes only; tokens already stored in seq_builder.action.tokens) + action_position_ids = action_ids_3d[:, frame_t, :] # [3,tcf] + seq_builder.append_action_span( + tcf, + action_position_ids, + payload_index=action_payload_index, + payload_start=frame_t * tcf, + payload_shape=(tcf,), + ) # Action tokens are never in MSE loss (always conditioning) - curr += tcf total_split_len += tcf # Pack vision tokens for this frame - frame_indexes = list(range(curr, curr + patches_per_frame)) - packed_seq.vision.sequence_indexes.extend(frame_indexes) - curr += patches_per_frame + vision_position_ids = vision_ids_3d[:, frame_t, :] # [3,patch_h*patch_w] + frame_indexes = seq_builder.append_vision_span( + patches_per_frame, + vision_position_ids, + payload_index=vision_payload_index, + payload_start=frame_t * patches_per_frame, + payload_shape=(1, patch_h, patch_w), + ) total_split_len += patches_per_frame # Vision MSE loss: supervise non-conditioning frames if frame_t not in condition_set_vision: - packed_seq.vision.mse_loss_indexes.extend(frame_indexes) + vision.mse_loss_indexes.extend(frame_indexes) frame_ts = input_timestep[frame_t].item() if isinstance(input_timestep, torch.Tensor) else input_timestep - packed_seq.vision.timesteps.extend([frame_ts] * patches_per_frame) + vision.timesteps.extend([frame_ts] * patches_per_frame) - packed_seq.curr = curr return total_split_len, null_action_flag diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index 6fed5dee..6f3a0b41 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -54,7 +54,7 @@ build_sequence_plans_from_data_batch, pack_input_sequence, ) -from cosmos_framework.data.generator.sequence_packing.modalities import add_special_tokens +from cosmos_framework.data.generator.sequence_packing.modality import add_special_tokens from cosmos_framework.model.generator.tokenizers.interface import VideoTokenizerInterface from cosmos_framework.model.generator.upsampler.prompts import build_messages, clean_response from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution @@ -3068,6 +3068,13 @@ def get_data_and_condition( fps_raw = torch.stack(fps_raw).flatten() # list of scalar tensors -> (B,) fps_vision = fps_raw.to(**self.tensor_kwargs) if fps_raw is not None else None fps_action = fps_raw.to(**self.tensor_kwargs) if fps_raw is not None else None + if "conditioning_fps_action" in data_batch: + fps_action_raw = data_batch["conditioning_fps_action"] + if isinstance(fps_action_raw, list): + # serve_policy.py wraps every tensor in a list; align with the conditioning_fps + # branch above so inference paths don't hit AttributeError on .to(). + fps_action_raw = torch.stack(fps_action_raw).flatten() + fps_action = fps_action_raw.to(**self.tensor_kwargs) # Sound FPS for RoPE alignment (constant, from config) if x0_tokens_sound is not None: diff --git a/cosmos_framework/utils/generator/image_resize.py b/cosmos_framework/utils/generator/image_resize.py new file mode 100644 index 00000000..4c1b38e0 --- /dev/null +++ b/cosmos_framework/utils/generator/image_resize.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Image resizing helpers shared by Cosmos3 image-edit generation paths.""" + +from __future__ import annotations + +import math + +from PIL import Image + +DEFAULT_MAX_PIXELS = 1024 * 1024 +DEFAULT_PADDING_CONSTANT = 32 + + +def get_max_pixels_resized_size( + width: int, + height: int, + max_pixels: int = DEFAULT_MAX_PIXELS, + padding_constant: int = DEFAULT_PADDING_CONSTANT, +) -> tuple[int, int]: + """Return an aspect-preserving size capped by max pixels and rounded down.""" + if width <= 0 or height <= 0: + raise ValueError(f"Image dimensions must be positive, got {width}x{height}.") + if padding_constant <= 0: + raise ValueError(f"padding_constant must be positive, got {padding_constant}.") + if max_pixels < padding_constant * padding_constant: + raise ValueError( + f"max_pixels={max_pixels} is too small for padding_constant={padding_constant}; " + f"minimum is {padding_constant * padding_constant}." + ) + + scale = min(1.0, math.sqrt(max_pixels / (width * height))) + resized_width = int(width * scale) + resized_height = int(height * scale) + + resized_width = (resized_width // padding_constant) * padding_constant + resized_height = (resized_height // padding_constant) * padding_constant + resized_width = max(resized_width, padding_constant) + resized_height = max(resized_height, padding_constant) + + while resized_width * resized_height > max_pixels: + if resized_width >= resized_height and resized_width > padding_constant: + resized_width -= padding_constant + elif resized_height > padding_constant: + resized_height -= padding_constant + else: + break + + return resized_width, resized_height + + +def resize_pil_image( + image: Image.Image, + max_pixels: int = DEFAULT_MAX_PIXELS, + padding_constant: int = DEFAULT_PADDING_CONSTANT, +) -> Image.Image: + """Resize a PIL image to a max-pixels budget while preserving aspect ratio.""" + resized_size = get_max_pixels_resized_size( + width=image.size[0], + height=image.size[1], + max_pixels=max_pixels, + padding_constant=padding_constant, + ) + if resized_size == image.size: + return image.copy() + return image.resize(resized_size, Image.LANCZOS) diff --git a/cosmos_framework/utils/generator/image_resize_test.py b/cosmos_framework/utils/generator/image_resize_test.py new file mode 100644 index 00000000..81be25a0 --- /dev/null +++ b/cosmos_framework/utils/generator/image_resize_test.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Tests for Cosmos3 max-pixels image resize helpers.""" + +from __future__ import annotations + +import pytest +from PIL import Image + +from cosmos_framework.utils.generator.image_resize import get_max_pixels_resized_size, resize_pil_image + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +def test_get_max_pixels_resized_size_downscales_wide_image() -> None: + assert get_max_pixels_resized_size(1920, 1080, max_pixels=1024 * 1024, padding_constant=32) == (1344, 768) + + +def test_get_max_pixels_resized_size_preserves_square_at_budget() -> None: + assert get_max_pixels_resized_size(1024, 1024, max_pixels=1024 * 1024, padding_constant=32) == (1024, 1024) + + +def test_get_max_pixels_resized_size_does_not_upscale_small_image() -> None: + assert get_max_pixels_resized_size(640, 480, max_pixels=1024 * 1024, padding_constant=32) == (640, 480) + + +def test_get_max_pixels_resized_size_rounds_down_under_budget() -> None: + assert get_max_pixels_resized_size(641, 481, max_pixels=1024 * 1024, padding_constant=32) == (640, 480) + + +def test_get_max_pixels_resized_size_rejects_impossible_budget() -> None: + with pytest.raises(ValueError, match="too small"): + get_max_pixels_resized_size(128, 128, max_pixels=1023, padding_constant=32) + + +def test_resize_pil_image_uses_computed_size() -> None: + image = Image.new("RGB", (1920, 1080), (10, 20, 30)) + + resized = resize_pil_image(image, max_pixels=1024 * 1024, padding_constant=32) + + assert resized.size == (1344, 768) diff --git a/cosmos_framework/utils/generator/model_loader.py b/cosmos_framework/utils/generator/model_loader.py index a53089de..2f68a9b9 100644 --- a/cosmos_framework/utils/generator/model_loader.py +++ b/cosmos_framework/utils/generator/model_loader.py @@ -41,7 +41,9 @@ def write_lock(self) -> FileLock: from cosmos_framework.utils.config_helper import get_config_module, override from cosmos_framework.utils.easy_io import easy_io from cosmos_framework.checkpoint.dcp import CustomLoadPlanner, CustomSavePlanner, ModelWrapper +from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig from cosmos_framework.model.generator.utils.safetensors_loader import load_vfm_model +from cosmos_framework.utils.generator.quantization import apply_quantization_inplace ################################################### # below are the load_model function for inference # @@ -303,6 +305,7 @@ def load_model_from_checkpoint( load_ema_to_reg: bool = False, parallelism_config: dict[str, Any] = {}, compile_config: dict[str, Any] = {}, + quantization_config: dict[str, Any] = {}, seed: int = 0, experiment_opts: list[str] = [], use_cache_checkpoint: bool = False, @@ -338,6 +341,12 @@ def load_model_from_checkpoint( ``config.model.config.parallelism``. compile_config: Dictionary of torch.compile configuration options. Keys are applied to ``config.model.config.compile``. + quantization_config: Post training quantization (PTQ) config for inference. Dictionary + of quantization options (e.g. ``method``, ``include_regex``, ``exclude_regex``). + Quantization is applied in-place on the model parameters. Only works when model + sharding (FSDP) is disabled. When ``method`` resolves to ``None``, + quantization is disabled. Only Blackwell architectures are supported. Valid + quantization methods are ``nvfp4`` and ``mxfp8``. seed: Random seed used for initialization (if applicable). experiment_opts: Extra experiment/config override options. use_cache_checkpoint: If True, locally save & read remote checkpoints to speed up repeated loads. @@ -389,6 +398,14 @@ def load_model_from_checkpoint( else: raise ValueError(f"Key {key} not found in config.model.config.compile") + quantization_config_obj = QuantizationConfig(**quantization_config) + quantization_method = quantization_config_obj.method + + if quantization_method is not None: + dp_shard_degree = config.model.config.parallelism.data_parallel_shard_degree + if dp_shard_degree > 1: + raise ValueError("Quantization is not supported for DP sharded models.") + # Disable activation checkpointing for inference. config.model.config.activation_checkpointing.mode = "none" @@ -463,43 +480,47 @@ def load_model_from_checkpoint( f"Successfully loaded safetensors VFM checkpoint from {checkpoint_path}; " f"time taken: {time.time() - start_time:.2f} seconds" ) - return model, config - - # DCP path with optional local cache reshard. - def load_model(checkpoint_load_path: str) -> None: - _load_model( - model, - checkpoint_path=checkpoint_load_path, - credential_path=credential_path, - enable_gcs_patch_in_boto3=enable_gcs_patch_in_boto3, - load_ema_to_reg=load_ema_to_reg, - keys_to_skip_loading=keys_to_skip_loading, - ) - checkpoint_cache_path = None - if use_cache_checkpoint: - checkpoint_cache_path = checkpoint_path_to_cached_path(checkpoint_path, cache_checkpoint_rootdir) + else: + # DCP path with optional local cache reshard. + def load_model(checkpoint_load_path: str) -> None: + _load_model( + model, + checkpoint_path=checkpoint_load_path, + credential_path=credential_path, + enable_gcs_patch_in_boto3=enable_gcs_patch_in_boto3, + load_ema_to_reg=load_ema_to_reg, + keys_to_skip_loading=keys_to_skip_loading, + ) - if checkpoint_cache_path is None: - load_model(checkpoint_path) - _reload_pretrained_reasoner_after_checkpoint_load(model) - return model, config + checkpoint_cache_path = None + if use_cache_checkpoint: + checkpoint_cache_path = checkpoint_path_to_cached_path(checkpoint_path, cache_checkpoint_rootdir) - cache_lock_path = f"{checkpoint_cache_path}.lock" - cache_action = _CheckpointCacheAction.ERROR - log.info(f"Acquiring checkpoint cache write lock: {cache_lock_path}") - with _checkpoint_cache_group_lock(checkpoint_cache_path, cache_lock_path) as cache_action: - if cache_action == _CheckpointCacheAction.POPULATE_CACHE: + if checkpoint_cache_path is None: load_model(checkpoint_path) - _save_model( - model, - checkpoint_path=checkpoint_cache_path, - save_reg_to_ema=load_ema_to_reg, - ) - if cache_action == _CheckpointCacheAction.LOAD_CACHE: - load_model(checkpoint_cache_path) + else: + cache_lock_path = f"{checkpoint_cache_path}.lock" + cache_action = _CheckpointCacheAction.ERROR + log.info(f"Acquiring checkpoint cache write lock: {cache_lock_path}") + with _checkpoint_cache_group_lock(checkpoint_cache_path, cache_lock_path) as cache_action: + if cache_action == _CheckpointCacheAction.POPULATE_CACHE: + load_model(checkpoint_path) + _save_model( + model, + checkpoint_path=checkpoint_cache_path, + save_reg_to_ema=load_ema_to_reg, + ) + + if cache_action == _CheckpointCacheAction.LOAD_CACHE: + load_model(checkpoint_cache_path) + + _reload_pretrained_reasoner_after_checkpoint_load(model) - _reload_pretrained_reasoner_after_checkpoint_load(model) + if quantization_method is not None: + # Replicated path: apply the (config-resident) recipe in place on the + # unsharded (plain-tensor) params. + apply_quantization_inplace(model, quantization_config_obj) return model, config diff --git a/cosmos_framework/utils/generator/quantization.py b/cosmos_framework/utils/generator/quantization.py new file mode 100644 index 00000000..436296be --- /dev/null +++ b/cosmos_framework/utils/generator/quantization.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Low-precision quantization helpers for the Cosmos3 VFM MoT network. + +Quantization is applied via torchao's :func:`apply_quantization_inplace`, which +uses the ``quantize_`` path to replace each selected weight with a quantized +tensor subclass in place. Because the live parameter becomes a tensor subclass, +this only works on unsharded (plain-tensor) params and is therefore restricted +to replicated inference (``data_parallel_shard_degree == 1``); it cannot be +applied to an FSDP-sharded model whose params are ``DTensor`` shards. + +This is an inference-only path: the ``quantize_`` PTQ configs have no backward +support. Module selection is delegated to the filter built by +:func:`_get_filter_fn`. +""" + +import re + +from torch import nn + +from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig + +# NOTE: ``torchao`` is imported lazily inside the functions below rather than at +# module top level. These two helpers are the only torchao consumers, but this +# module is imported transitively by the model package (e.g. during tests that +# never quantize). Keeping the imports lazy means importing this module does not +# require torchao to be installed; the imports only run — and only fail — when +# quantization is actually requested. + + +def _get_filter_fn(quantization_config: QuantizationConfig): + """Build a module-selection predicate from the quantization config. + + The returned closure captures ``include_regex`` / ``exclude_regex`` and + implements the selection policy documented on :class:`QuantizationConfig`. + Each key is treated as a regular expression and matched against a module's + FQN with :func:`re.search` (a plain substring remains a valid pattern, so + existing substring-style keys keep working). It is passed to torchao as + ``filter_fn`` (for ``quantize_``), which expects a + ``(module, fqn) -> bool`` signature. + + Args: + quantization_config: Config carrying the include/exclude key lists. + + Returns: + A predicate suitable for torchao's ``filter_fn`` / ``module_filter_fn``. + """ + + def _filter_fn(mod: nn.Module, name: str) -> bool: + """Decide whether a single module should be quantized. + + Called once per module as torchao walks the model recursively. A module + is selected only when ALL of the following hold: + + 1. It is an ``nn.Linear`` (the only layer type these recipes support). + 2. ``include_regex`` is empty (include everything) OR the module's FQN + matches at least one include pattern. + 3. The module's FQN matches none of the ``exclude_regex`` patterns. + + Each include/exclude key is treated as a regular expression and matched + against the FQN with :func:`re.search`, so the pattern can match anywhere + in the name (a plain substring is still a valid regex, preserving the + previous substring-match behavior, while enabling anchors like ``^``/``$``, + alternation, character classes, etc.). + + Note the parenthesization around the include check: ``and`` binds tighter + than ``or`` in Python, so without it the ``nn.Linear`` and exclude + checks would not apply across both include branches. + + Args: + mod (torch.nn.Module): The module that is being processed. + name (str): A fully qualified name of the module that is being processed. + + Return: + True if the module should be quantized, False otherwise. + """ + include_keys = quantization_config.include_regex + exclude_keys = quantization_config.exclude_regex + return ( + isinstance(mod, nn.Linear) + and (len(include_keys) == 0 or any(re.search(key, name) for key in include_keys)) + and not any(re.search(key, name) for key in exclude_keys) + ) + + return _filter_fn + + +def apply_quantization_inplace(model: nn.Module, quantization_config: QuantizationConfig): + """Apply quantization in place via ``quantize_`` (replaces weights with quantized tensors). + + This is the replication path. ``quantize_`` replaces each weight with a + quantized tensor subclass as the live parameter, which only works when the + parameters are plain tensors. It therefore cannot be applied to an already + FSDP-sharded model (the params are ``DTensor`` shards), so it is restricted + to replicated inference (``data_parallel_shard_degree == 1``). + + These configs (``MXDynamicActivationMXWeightConfig`` / + ``NVFP4DynamicActivationNVFP4WeightConfig``) are inference-only (PTQ) and + have no backward support. For the sharded case use ``apply_quantization`` + (the module-swap path) instead; both functions are currently inference + paths, selected by whether FSDP is sharding the model. + """ + # No-op when quantization is disabled. + if quantization_config.method is None: + return + + from torchao.prototype.mx_formats import ( + MXDynamicActivationMXWeightConfig, + NVFP4DynamicActivationNVFP4WeightConfig, + ) + from torchao.quantization import ( + quantize_, + ) + + if quantization_config.method == "mxfp8": + # mxfp8 / nvfp4 use fixed block scales. + quantize_( + model, + config=MXDynamicActivationMXWeightConfig(), + filter_fn=_get_filter_fn(quantization_config), + ) + elif quantization_config.method == "nvfp4": + # use_triton_kernel=False avoids torchao's fused NVFP4 Triton kernel, which + # requires the external `mslk` package. Prebuilt mslk wheels are linked + # against upstream torch and fail to load against NVIDIA's NGC custom torch + # builds (ABI mismatch on `torch::Library::_def`), so we use torchao's + # built-in NVFP4 path instead. + quantize_( + model, + config=NVFP4DynamicActivationNVFP4WeightConfig(use_triton_kernel=False), + filter_fn=_get_filter_fn(quantization_config), + ) + else: + raise ValueError(f"Unsupported quantization method: {quantization_config.method}") From ace614064b25b19c0b8452bd6f82fc5a7c040e89 Mon Sep 17 00:00:00 2001 From: lfengad Date: Wed, 15 Jul 2026 14:31:30 +0800 Subject: [PATCH 24/26] feat: wire MXFP8/NVFP4 quantization into inference CLI (#113) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Ports the inference-CLI plumbing from imaginaire4 **MR 10201** (`Inference optimizations [4/n]: MXFP8 / NVFP4 PTQ`) that was missing on this side. The quantization **backend** was already synced (via the 2026-07-10 release): `QuantizationConfig` (`configs/base/defaults/quantization.py`), `apply_quantization_inplace` (`utils/generator/quantization.py`), and the `model_loader.py` hook. But the user-facing `OmniInference` entrypoint was never wired, so `--quantization-method` had no effect. This PR connects them. The missing pieces mapped exactly to MR 10201's `packages/cosmos3/cosmos3/*` files (→ `cosmos_framework/inference/*`); the `projects/cosmos3/*` files were already in. ## Changes - **`inference/common/args.py`** — add `QuantizationMethod` / `QuantizationArgs` / `QuantizationOverrides`; mix them into `SetupArgs` / `SetupOverrides`. Fields flow through `build_setup`'s `model_dump → model_validate` (same as guardrails), so no extra build call is needed. - **`inference/model.py`** — add `Cosmos3OmniConfig.quantization` property + setter and thread `quantization_config` through `from_pretrained_dcp`. Uses the **local** direct-construction style `unstructure_config(QuantizationConfig(**value))` rather than upstream's `LazyCall` wrapper, matching the sibling `parallelism`/`compile` setters in this repo. - **`inference/inference.py`** — add `_get_quantization_config` and pass it into **both** `_create` branches (`load_model_from_checkpoint` and `from_pretrained_dcp`). ## Effect - **Standard experiment DCP path** (`load_model_from_checkpoint`): quantization applied end-to-end via `apply_quantization_inplace`. ✅ - **HF `from_pretrained_dcp` path**: stores `config.quantization` without an apply hook — **matches upstream MR behavior** (not a regression introduced here). - Only supported on Blackwell + replicated inference (no FSDP sharding); guarded by the existing `model_loader.py` DP-shard check. ## Not included - `websocket_policy_server/serve_policy.py` + README from the MR — that action-eval subtree is not present in this repo. - The MR's cosmetic blank-line change in `packages/cosmos3/cosmos3/args.py`. ## Verification - `py_compile` passes on all three files. - Load-bearing pydantic logic verified in-container (isolated replica): defaults, mutable-list-default isolation, `build_quantization`, `Literal` rejection of invalid methods, MRO/inheritance, and `SetupOverrides → SetupArgs` field flow — all assertions pass. - Full module import not exercised: cosmos-framework requires Python ≥3.12 while the available i4 container env is 3.10; needs the cosmos-framework 3.13 venv to import the real module. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/common/args.py | 36 +++++++++++++++++-- .../inference/common/public_model_config.py | 1 + cosmos_framework/inference/inference.py | 12 +++++++ cosmos_framework/inference/model.py | 15 ++++++++ 4 files changed, 62 insertions(+), 2 deletions(-) diff --git a/cosmos_framework/inference/common/args.py b/cosmos_framework/inference/common/args.py index 4fa87c97..13ee789e 100644 --- a/cosmos_framework/inference/common/args.py +++ b/cosmos_framework/inference/common/args.py @@ -603,6 +603,38 @@ def build_checkpoint(self, *, checkpoints: dict[str, CheckpointConfig]) -> Check CfgpSize = Annotated[int, pydantic.Field(ge=1, le=2)] CompiledRegion = Literal["all", "language"] +# Low-precision quantization method to apply to the model at load time. +# One of ``mxfp8`` / ``nvfp4``, or ``None`` (default) to disable. +# Routed to the VFM model loader, which selects an FSDP-compatible +# (module-swap) path when sharded (``dp_shard_size > 1``) and an in-place +# path when replicated (``dp_shard_size == 1``). Note ``mxfp8`` / ``nvfp4`` +# are only supported on the replicated path. +QuantizationMethod = Literal["mxfp8", "nvfp4"] + + +class QuantizationArgs(ArgsBase): + """Low-precision quantization arguments applied to the model at load time.""" + + quantization_method: QuantizationMethod | None + quantization_include_regex: list[str] + quantization_exclude_regex: list[str] + + +class QuantizationOverrides(OverridesBase): + quantization_method: QuantizationMethod | None = None + """Quantization method (``mxfp8`` / ``nvfp4``), or ``None`` to disable. + + Post-training quantization (PTQ) is applied in-place to the model at load + time. Only supported on Blackwell architectures and when FSDP sharding is disabled. + """ + quantization_include_regex: list[str] = ["language_model.model.layers"] + """Regexes matched against module FQNs; a Linear is quantized only if it matches one (empty = all).""" + quantization_exclude_regex: list[str] = pydantic.Field(default_factory=list) + """Regexes matched against module FQNs; a Linear is skipped if it matches any.""" + + def build_quantization(self) -> QuantizationArgs: + return self._build(QuantizationArgs) + class ParallelismArgs(ArgsBase): """Parallelism arguments.""" @@ -702,7 +734,7 @@ class GuardrailOverrides(OverridesBase): """Offload guardrail models to CPU.""" -class SetupArgs(ABC, CheckpointArgs, ParallelismArgs, GuardrailArgs): +class SetupArgs(ABC, CheckpointArgs, ParallelismArgs, QuantizationArgs, GuardrailArgs): output_dir: ResolvedPath keep_going: bool skip_invalid_samples: bool @@ -737,7 +769,7 @@ def get_variant(cls) -> str: return cls.model_fields["variant"].default -class SetupOverrides(ABC, CheckpointOverrides, ParallelismOverrides, GuardrailOverrides): +class SetupOverrides(ABC, CheckpointOverrides, ParallelismOverrides, QuantizationOverrides, GuardrailOverrides): """Inference setup arguments.""" output_dir: Annotated[ResolvedPath | None, tyro.conf.arg(aliases=("-o",))] = None diff --git a/cosmos_framework/inference/common/public_model_config.py b/cosmos_framework/inference/common/public_model_config.py index 584a65c9..aff1a057 100644 --- a/cosmos_framework/inference/common/public_model_config.py +++ b/cosmos_framework/inference/common/public_model_config.py @@ -37,6 +37,7 @@ "projects.cosmos3.vfm.configs.base.defaults.model_config.RectifiedFlowInferenceConfig": "rectified_flow_inference_config", "projects.cosmos3.vfm.configs.base.defaults.model_config.RectifiedFlowTrainingConfig": "rectified_flow_training_config", "projects.cosmos3.vfm.configs.base.defaults.parallelism.ParallelismConfig": "parallelism_config", + "projects.cosmos3.vfm.configs.base.defaults.quantization.QuantizationConfig": "quantization_config", "projects.cosmos3.vfm.configs.base.defaults.vlm.PretrainedWeightsConfig": "pretrained_weights_config", "projects.cosmos3.vfm.configs.base.defaults.vlm.VLMConfig": "vlm_config", } diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index a7f338ca..30bd04b5 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -23,6 +23,7 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig +from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig from cosmos_framework.inference.args import ( ModelMode, NegativeMetadataMode, @@ -1055,6 +1056,14 @@ def _get_compile_config(cls, setup_args: ParallelismArgs) -> CompileConfig: compile_dynamic=setup_args.compile_dynamic, ) + @classmethod + def _get_quantization_config(cls, setup_args: SetupArgs) -> QuantizationConfig: + return QuantizationConfig( + method=setup_args.quantization_method, + include_regex=list(setup_args.quantization_include_regex), + exclude_regex=list(setup_args.quantization_exclude_regex), + ) + @override @classmethod def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: @@ -1064,6 +1073,7 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: sampler_override = setup_args.sampler parallelism_config = cls._get_parallelism_config(setup_args) compile_config = cls._get_compile_config(setup_args) + quantization_config = cls._get_quantization_config(setup_args) if setup_args.checkpoint_type == CheckpointType.DCP and setup_args.config_file_type == ConfigFileType.MODULE: from cosmos_framework.inference.common.config import save_config from cosmos_framework.utils.generator.model_loader import load_model_from_checkpoint @@ -1081,6 +1091,7 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: credential_path=setup_args.credential_path or None, parallelism_config=attrs.asdict(parallelism_config), compile_config=attrs.asdict(compile_config), + quantization_config=attrs.asdict(quantization_config), load_ema_to_reg=setup_args.use_ema_weights, experiment_opts=[ *setup_args.experiment_overrides, @@ -1130,6 +1141,7 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: config=config, parallelism_config=parallelism_config, compile_config=compile_config, + quantization_config=quantization_config, ).model if model.config.rectified_flow_inference_config.scheduler_type != sampler_override: model.config.rectified_flow_inference_config.scheduler_type = sampler_override diff --git a/cosmos_framework/inference/model.py b/cosmos_framework/inference/model.py index e9c38b48..042bab5b 100644 --- a/cosmos_framework/inference/model.py +++ b/cosmos_framework/inference/model.py @@ -34,6 +34,7 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig +from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig from cosmos_framework.inference.common.args import CheckpointType from cosmos_framework.inference.common.checkpoints import register_checkpoints from cosmos_framework.inference.common.config import structure_config, undo_config_dict_replacements, unstructure_config @@ -416,6 +417,16 @@ def compile(self, value: dict | None): return self.model.setdefault("config", {})["compile"] = unstructure_config(CompileConfig(**value)) + @property + def quantization(self) -> dict: + return self.model.get("config", {}).get("quantization", {}) + + @quantization.setter + def quantization(self, value: dict | None): + if value is None: + return + self.model.setdefault("config", {})["quantization"] = unstructure_config(QuantizationConfig(**value)) + class Cosmos3OmniModel(transformers.PreTrainedModel): config_class = Cosmos3OmniConfig # type: ignore @@ -448,6 +459,7 @@ def from_pretrained_dcp( config: Cosmos3OmniConfig | None = None, parallelism_config: ParallelismConfig | None = None, compile_config: CompileConfig | None = None, + quantization_config: QuantizationConfig | None = None, ): if config is None: config = Cosmos3OmniConfig.from_pretrained(checkpoint_path) @@ -455,8 +467,11 @@ def from_pretrained_dcp( parallelism_config = ParallelismConfig() if compile_config is None: compile_config = CompileConfig() + if quantization_config is None: + quantization_config = QuantizationConfig() config.parallelism = attrs.asdict(parallelism_config) config.compile = attrs.asdict(compile_config) + config.quantization = attrs.asdict(quantization_config) model = cls(config) checkpoint_type = CheckpointType.from_path(checkpoint_path) match checkpoint_type: From bc44b1e5fbe76402aa60f8503acf15b153bb3b51 Mon Sep 17 00:00:00 2001 From: lfengad Date: Wed, 15 Jul 2026 14:39:38 +0800 Subject: [PATCH 25/26] Sync Cosmos3 Diffusers converter support (#110) ## Summary - support Cosmos3 checkpoints with optional action and sound generation - support current and legacy AVAE/tokenizer layouts and framework-native VLM config layouts - make language-model remapping reject collisions and exclude visual weights from the text transformer - support both legacy and current action projection module names - add `--skip-vision-encoder` and make `--config-only` work with an empty output directory - generate a top-level safetensors index that matches the files actually exported - keep the converter implementation aligned with the corresponding Cosmos3 package implementation, apart from repository-specific imports and licensing - add a GPU CI regression that converts the staged Nano DCP, validates the complete Diffusers output against the public Nano reference indexes, and always removes the generated output ## Validation - converted a Cosmos3-Nano DCP checkpoint to a complete Diffusers checkpoint - verified all 10 exported safetensors files byte-for-byte against `nvidia/Cosmos3-Nano` - verified the unified index contains 1,165 tensors with no dangling file references - loaded the output with `Cosmos3OmniPipeline` and ran one-step video-plus-sound inference with finite outputs - verified pytest collection selects the new Nano Diffusers conversion test in the generator GPU workflow - Ruff check, `py_compile`, YAML parse, and `git diff --check` pass --- .github/workflows/gpu-tests.yml | 28 +++- .../scripts/_convert_model_to_diffusers.py | 46 ++++-- .../scripts/convert_model_to_diffusers.py | 121 ++++++++++----- tests/launch_regression_test.py | 140 +++++++++++++++++- 4 files changed, 280 insertions(+), 55 deletions(-) diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index 5142f40a..31119267 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -10,7 +10,7 @@ # * training-smoke — Nano SFT pipeline (convert -> train 5 -> export -> t2i) # * generator-training-regression — vision_sft_nano loss vs goldens (4-GPU subset) # * generator-inference-smoke — Nano multi-modality inference (t2vs + policy + forward_dynamics) -# * reasoner-inference-smoke — Nano reasoner inference first-token logits vs golden (image-conditioned, 4-GPU) +# * reasoner-inference-smoke — Nano reasoner inference golden + Qwen conversion coverage (4-GPU) # * reasoner-training-regression — llava_ov loss vs goldens (4-GPU subset) # # Requires: @@ -76,7 +76,7 @@ jobs: generator-training-regression: needs: pre-commit runs-on: [self-hosted, gpu, h200] - timeout-minutes: 60 + timeout-minutes: 75 env: HF_TOKEN: ${{ secrets.HF_TOKEN }} HF_HUB_DISABLE_XET: "1" @@ -90,12 +90,14 @@ jobs: - name: Sync environment (cu128-train) run: uv sync --all-extras --group=cu128-train - # Generator (vision_sft_nano) loss vs the h100 goldens. -s streams the live log. - - name: Generator regression (vision_sft_nano, 4-GPU subset) + # Generator (vision_sft_nano) loss vs the h100 goldens, plus a Nano DCP -> + # Diffusers conversion that validates the exported component files and indexes. + # Both tests share the module-scoped staged Nano DCP. -s streams the live log. + - name: Generator regression and Nano Diffusers export (4-GPU subset) run: | export LD_LIBRARY_PATH= uv run --all-extras --group=cu128-train python -m pytest -v -s \ - tests/launch_regression_test.py -k vision_sft_nano \ + tests/launch_regression_test.py -k 'vision_sft_nano or convert_model_to_diffusers' \ --num-gpus=4 --levels=2 -o addopts= # The h100_inputs fixture removes its DCP stage on teardown; clear the @@ -169,8 +171,20 @@ jobs: tests/nano_reasoner_inference_smoke_test.py::test_nano_reasoner_image_first_token_logits \ --num-gpus=4 --levels=2 -o addopts= - # Reasoner inference writes only the pytest tmp dir (generated text + logs); - # the checkpoint download stays in the HF cache (kept). + # Merge Cosmos3-Nano into the public Qwen3-VL shell and verify that every + # Qwen tensor is present and both the vision and language towers come from + # Cosmos3-Nano. The exact node ID avoids collecting the training regressions + # in the same file. + - name: Nano reasoner Qwen tensor conversion coverage (4 GPU collection) + run: | + export LD_LIBRARY_PATH= + uv run --all-extras --group=cu128-train python -m pytest -v -s \ + tests/launch_regression_test.py::test_convert_reasoner_converts_all_qwen_tensors \ + --num-gpus=4 --levels=2 -o addopts= + + # Reasoner inference and conversion write to pytest tmp dirs (generated + # text, logs, and the merged checkpoint); downloaded checkpoints stay in + # the HF cache (kept). - name: Clean up run outputs if: always() run: | diff --git a/cosmos_framework/scripts/_convert_model_to_diffusers.py b/cosmos_framework/scripts/_convert_model_to_diffusers.py index da5549ab..ceab0ef2 100644 --- a/cosmos_framework/scripts/_convert_model_to_diffusers.py +++ b/cosmos_framework/scripts/_convert_model_to_diffusers.py @@ -99,6 +99,8 @@ (".k_norm.", ".norm_k."), ) +_LANGUAGE_MODEL_VISION_PREFIXES = ("model.visual.", "visual.") + # Legacy TimestepEmbedder stored its MLP as `nn.Sequential([Linear, SiLU, Linear])`, # so state-dict keys were `mlp.0.*` / `mlp.2.*`. The diffusers `TimestepEmbedding` # stores them as named attributes `linear_1` / `linear_2`. Index 1 (SiLU) has no @@ -141,6 +143,22 @@ def _remap_language_model_key(key: str) -> str: return key +def _remap_language_model_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Remap language-model keys and reject ambiguous source layouts.""" + remapped_state_dict: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if key.startswith(_LANGUAGE_MODEL_VISION_PREFIXES): + continue + remapped_key = _remap_language_model_key(key) + if remapped_key in remapped_state_dict: + raise RuntimeError( + "Language-model key remap collision while applying diffusers key remap: " + f"{key!r} maps to existing key {remapped_key!r}." + ) + remapped_state_dict[remapped_key] = value + return remapped_state_dict + + def _load_sound_tokenizer_state_dict(checkpoint_path: pathlib.Path) -> dict[str, torch.Tensor]: if checkpoint_path.suffix == ".safetensors": try: @@ -615,18 +633,22 @@ def convert_model_to_diffusers(args: Args) -> None: unified_3d_mrope_temporal_modality_margin = ( _tmp.config.diffusion_expert_config.unified_3d_mrope_temporal_modality_margin ) - action2llm = getattr(_tmp.net, "action2llm", None) - llm2action = getattr(_tmp.net, "llm2action", None) + action_proj_in = getattr(_tmp.net, "action2llm", None) + if action_proj_in is None: + action_proj_in = getattr(_tmp.net, "action_proj_in", None) + action_proj_out = getattr(_tmp.net, "llm2action", None) + if action_proj_out is None: + action_proj_out = getattr(_tmp.net, "action_proj_out", None) action_modality_embed = getattr(_tmp.net, "action_modality_embed", None) has_action_projection_weights = any( - module is not None for module in (action2llm, llm2action, action_modality_embed) + module is not None for module in (action_proj_in, action_proj_out, action_modality_embed) ) action_gen = bool( _get_config_value(net_cfg, model_cfg, name="action_gen", default=False) or has_action_projection_weights ) action_dim = _get_config_value(net_cfg, model_cfg, name="action_dim", default=None) - if action_dim is None and action2llm is not None: - action_dim = getattr(action2llm, "input_size", None) + if action_dim is None and action_proj_in is not None: + action_dim = getattr(action_proj_in, "input_size", None) if action_dim is None: action_dim = max_action_dim num_embodiment_domains = int(_get_config_value(net_cfg, model_cfg, name="num_embodiment_domains", default=32)) @@ -662,8 +684,8 @@ def convert_model_to_diffusers(args: Args) -> None: missing_action_modules = [ name for name, module in ( - ("action2llm", action2llm), - ("llm2action", llm2action), + ("action_proj_in/action2llm", action_proj_in), + ("action_proj_out/llm2action", action_proj_out), ("action_modality_embed", action_modality_embed), ) if module is None @@ -718,7 +740,7 @@ def convert_model_to_diffusers(args: Args) -> None: unified_3d_mrope_temporal_modality_margin=unified_3d_mrope_temporal_modality_margin, vocab_size=lm_cfg.vocab_size, ) - state_dict = {_remap_language_model_key(k): v for k, v in language_model.state_dict().items()} + state_dict = _remap_language_model_state_dict(language_model.state_dict()) for k, v in vae2llm.state_dict().items(): state_dict[f"proj_in.{k}"] = v for k, v in llm2vae.state_dict().items(): @@ -726,9 +748,9 @@ def convert_model_to_diffusers(args: Args) -> None: for k, v in time_embedder.state_dict().items(): state_dict[f"time_embedder.{_TIME_EMBEDDER_KEY_REMAP[k]}"] = v if action_gen: - for k, v in action2llm.state_dict().items(): + for k, v in action_proj_in.state_dict().items(): state_dict[f"action_proj_in.{k}"] = v - for k, v in llm2action.state_dict().items(): + for k, v in action_proj_out.state_dict().items(): state_dict[f"action_proj_out.{k}"] = v state_dict["action_modality_embed"] = action_modality_embed if sound_gen: @@ -743,8 +765,8 @@ def convert_model_to_diffusers(args: Args) -> None: vae2llm, llm2vae, time_embedder, - action2llm, - llm2action, + action_proj_in, + action_proj_out, action_modality_embed, sound2llm, llm2sound, diff --git a/cosmos_framework/scripts/convert_model_to_diffusers.py b/cosmos_framework/scripts/convert_model_to_diffusers.py index 98aa328a..63531771 100644 --- a/cosmos_framework/scripts/convert_model_to_diffusers.py +++ b/cosmos_framework/scripts/convert_model_to_diffusers.py @@ -11,11 +11,12 @@ } ) +import copy import json import shutil import struct from pathlib import Path -from typing import Annotated +from typing import Annotated, Any import pydantic import tyro @@ -28,6 +29,7 @@ build_public_model_config, load_model_config_from_hf_config, ) +from cosmos_framework.utils import log from cosmos_framework.utils.checkpoint_db import CheckpointConfig, CheckpointDirHf @@ -40,6 +42,9 @@ class Args(pydantic.BaseModel): config_only: bool = False """If True, only save config.""" + skip_vision_encoder: bool = False + """Do not save the vision encoder sidecar in the Diffusers checkpoint.""" + class SafetensorsIndexMetadata(pydantic.BaseModel): total_size: int = 0 @@ -67,7 +72,26 @@ def update_dir(self, safetensors_dir: Path, rel_path: str): self.update(safetensors_path, f"{rel_path}/{safetensors_path.name}") -def convert_model_to_diffusers(args: Args): +def _build_public_export_model_config(model_dict: dict[str, Any]) -> dict[str, Any]: + """Remove unsupported internal-only settings before building the public config.""" + public_model_dict = copy.deepcopy(model_dict) + quantization = public_model_dict.get("config", {}).pop("quantization", None) + if quantization is not None: + quantization_values = {key: value for key, value in quantization.items() if key not in {"_type", "_target_"}} + disabled_quantization = { + "exclude_regex": [], + "include_regex": [], + "method": None, + } + if quantization_values != disabled_quantization: + raise ValueError( + "Cannot export an enabled or non-default internal quantization config to the public Cosmos3 schema: " + f"{quantization_values}" + ) + return build_public_model_config(public_model_dict) + + +def convert_model_to_diffusers(args: Args) -> None: args.output_path.mkdir(parents=True, exist_ok=True) register_checkpoints() @@ -75,26 +99,45 @@ def convert_model_to_diffusers(args: Args): checkpoint_path = checkpoint_config.download_checkpoint() model_dict = load_model_config_from_hf_config(deserialize_config_dict(checkpoint_path / "config.json")) - assert model_dict["config"]["action_gen"] - assert model_dict["config"]["sound_gen"] - - sound_tokenizer_dir, sound_tokenizer_name = model_dict["config"]["sound_tokenizer"]["avae_path"].rsplit("/", 1) - sound_tokenizer_checkpoint = CheckpointConfig.maybe_from_uri(f"s3://bucket/{sound_tokenizer_dir}") - assert sound_tokenizer_checkpoint is not None - sound_tokenizer_local = Path(sound_tokenizer_checkpoint.hf.download()) - # HF-published checkpoints ship sound_tokenizer/ in the diffusers layout - # (config.json + diffusion_pytorch_model.safetensors), which the converter - # consumes directly; fall back to the legacy AVAE file pair named by the - # model config's avae_path. - sound_tokenizer_path = sound_tokenizer_local / "diffusion_pytorch_model.safetensors" - sound_tokenizer_config_path = sound_tokenizer_local / "config.json" - if not sound_tokenizer_path.is_file(): - sound_tokenizer_path = sound_tokenizer_local / sound_tokenizer_name - sound_tokenizer_config_path = sound_tokenizer_path.with_suffix(".json") - assert sound_tokenizer_path.is_file(), f"Sound tokenizer checkpoint not found: {sound_tokenizer_path}" - assert sound_tokenizer_config_path.is_file(), f"Sound tokenizer config not found: {sound_tokenizer_config_path}" - - vision_encoder_model = model_dict["config"]["vlm_config"]["tokenizer"]["pretrained_model_name"] + supports_action = model_dict["config"]["action_gen"] + supports_sound = model_dict["config"]["sound_gen"] + if not supports_action: + log.warning( + "The checkpoint does not support action generation. For some checkpoints, like " + "'Cosmos3-Super-Image2Video' it's fine, but make sure it's expected." + ) + if not supports_sound: + log.warning( + "The checkpoint does not support sound generation. For some checkpoints, like " + "'Cosmos3-Nano-Policy-DROID' it's fine, but make sure it's expected." + ) + + sound_tokenizer_path: Path | None = None + sound_tokenizer_config_path: Path | None = None + if supports_sound and not args.config_only: + sound_tokenizer_dir, sound_tokenizer_name = model_dict["config"]["sound_tokenizer"]["avae_path"].rsplit("/", 1) + sound_tokenizer_checkpoint = CheckpointConfig.maybe_from_uri(f"s3://bucket/{sound_tokenizer_dir}") + assert sound_tokenizer_checkpoint is not None + sound_tokenizer_local = Path(sound_tokenizer_checkpoint.hf.download()) + # HF-published checkpoints ship sound_tokenizer/ in the diffusers layout + # (config.json + diffusion_pytorch_model.safetensors), which the converter + # consumes directly; fall back to the legacy AVAE file pair named by the + # model config's avae_path. + sound_tokenizer_path = sound_tokenizer_local / "diffusion_pytorch_model.safetensors" + sound_tokenizer_config_path = sound_tokenizer_local / "config.json" + if not sound_tokenizer_path.is_file(): + sound_tokenizer_path = sound_tokenizer_local / sound_tokenizer_name + sound_tokenizer_config_path = sound_tokenizer_path.with_suffix(".json") + assert sound_tokenizer_path.is_file(), f"Sound tokenizer checkpoint not found: {sound_tokenizer_path}" + assert sound_tokenizer_config_path.is_file(), f"Sound tokenizer config not found: {sound_tokenizer_config_path}" + + vlm_config = model_dict["config"]["vlm_config"] + tokenizer_config = vlm_config["tokenizer"] + vision_encoder_model = ( + tokenizer_config.get("pretrained_model_name") + or tokenizer_config.get("tokenizer_type") + or vlm_config["model_name"] + ) if not args.config_only: from cosmos_framework.scripts._convert_model_to_diffusers import ( @@ -109,11 +152,11 @@ def convert_model_to_diffusers(args: Args): output=str(args.output_path), save_pipeline=True, dtype="bf16", - sound_tokenizer_path=str(sound_tokenizer_path), - sound_tokenizer_config_path=str(sound_tokenizer_config_path), - include_sound_tokenizer=True, + sound_tokenizer_path=str(sound_tokenizer_path) if sound_tokenizer_path else None, + sound_tokenizer_config_path=str(sound_tokenizer_config_path) if sound_tokenizer_config_path else None, + include_sound_tokenizer=supports_sound, vision_encoder_model=vision_encoder_model, - skip_vision_encoder=False, + skip_vision_encoder=args.skip_vision_encoder, ) _convert_model_to_diffusers(_args) @@ -127,26 +170,34 @@ def convert_model_to_diffusers(args: Args): vlm_checkpoint_path = vlm_checkpoint.download() for pattern in vlm_checkpoint.include: for p in Path(vlm_checkpoint_path).glob(pattern): + if p.name == "model.safetensors.index.json": + continue shutil.copy(p, args.output_path / p.name) # Add top-level config config_dict = deserialize_config_dict(args.output_path / "config.json") config_dict["architectures"] = ["Cosmos3ForConditionalGeneration"] + config_dict["model_type"] = "cosmos3_omni" # vLLM's `_prepare_weights` breaks after the first pattern with any match, so # collapse to a single glob spanning both component subdirs. The unified # `model.safetensors.index.json` written below dedupes the consolidated shard. config_dict["allow_patterns_overrides"] = ["*/*.safetensors"] - config_dict["model"] = build_public_model_config(model_dict) + config_dict["model"] = _build_public_export_model_config(model_dict) serialize_config_dict(config_dict, args.output_path / "config.json") - # Add top-level index - index = SafetensorsIndex() - index.update_dir(args.output_path / "transformer", "transformer") - vision_encoder_rel = "vision_encoder/model.safetensors" - index.update(args.output_path / vision_encoder_rel, vision_encoder_rel) - (args.output_path / "model.safetensors.index.json").write_text(index.model_dump_json(indent=2)) - - shutil.copy(checkpoint_path / "checkpoint.json", args.output_path / "checkpoint.json") + if not args.config_only: + # Add top-level index + index = SafetensorsIndex() + index.update_dir(args.output_path / "transformer", "transformer") + vision_encoder_rel = "vision_encoder/model.safetensors" + if (args.output_path / vision_encoder_rel).is_file(): + index.update(args.output_path / vision_encoder_rel, vision_encoder_rel) + (args.output_path / "model.safetensors.index.json").write_text(index.model_dump_json(indent=2)) + + checkpoint_metadata_path = checkpoint_path / "checkpoint.json" + if not checkpoint_metadata_path.is_file(): + checkpoint_metadata_path = checkpoint_path.parent / "checkpoint.json" + shutil.copy(checkpoint_metadata_path, args.output_path / "checkpoint.json") print(f"Saved diffusers checkpoint to {args.output_path}") diff --git a/tests/launch_regression_test.py b/tests/launch_regression_test.py index 24aec9d6..7d1055ab 100644 --- a/tests/launch_regression_test.py +++ b/tests/launch_regression_test.py @@ -73,6 +73,7 @@ from __future__ import annotations +import json import os import re import shutil @@ -506,6 +507,101 @@ def _load_st_tensor(directory: Path, key: str, weight_map: dict[str, str]): return f.get_tensor(key) +def _safetensors_keys(path: Path) -> set[str]: + """Read tensor names from a safetensors header without materializing weights.""" + from safetensors import safe_open + + assert path.is_file() and path.stat().st_size > 8, f"safetensors file missing or empty: {path}" + with safe_open(str(path), framework="pt") as f: + return set(f.keys()) + + +def _assert_diffusers_export_complete(model_dir: Path, reference_dir: Path) -> None: + """Validate the component files and safetensors indexes of a Nano Diffusers export.""" + required_files = ( + "checkpoint.json", + "config.json", + "model_index.json", + "model.safetensors.index.json", + "scheduler/scheduler_config.json", + "sound_tokenizer/config.json", + "sound_tokenizer/diffusion_pytorch_model.safetensors", + "tokenizer_config.json", + "transformer/config.json", + "transformer/diffusion_pytorch_model.safetensors.index.json", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", + "vision_encoder/config.json", + "vision_encoder/model.safetensors", + ) + for rel_path in required_files: + path = model_dir / rel_path + assert path.is_file() and path.stat().st_size > 0, f"Diffusers export missing or empty: {path}" + + pipeline_config = json.loads((model_dir / "model_index.json").read_text()) + assert pipeline_config["_class_name"] == "Cosmos3OmniPipeline" + # The public Diffusers pipeline does not register the Qwen vision encoder as + # a pipeline component; it is saved as a sidecar for transformers/vLLM and + # validated below through its required files, weight index, and tensor header. + for component in ("scheduler", "sound_tokenizer", "transformer", "vae"): + assert component in pipeline_config, f"Diffusers model_index.json is missing component {component!r}" + + model_config = json.loads((model_dir / "config.json").read_text()) + assert model_config["architectures"] == ["Cosmos3ForConditionalGeneration"] + assert model_config["model_type"] == "cosmos3_omni" + assert isinstance(model_config.get("model"), dict) and model_config["model"] + + top_index = json.loads((model_dir / "model.safetensors.index.json").read_text()) + top_weight_map: dict[str, str] = top_index["weight_map"] + assert top_weight_map, "top-level Diffusers safetensors index has no tensors" + assert top_index.get("metadata", {}).get("total_size", 0) > 0 + + reference_index = json.loads((reference_dir / "model.safetensors.index.json").read_text()) + assert top_weight_map == reference_index["weight_map"] + assert {str(path.relative_to(model_dir)) for path in model_dir.rglob("*.safetensors")} == { + str(path.relative_to(reference_dir)) for path in reference_dir.rglob("*.safetensors") + } + + indexed_files = set(top_weight_map.values()) + assert indexed_files, "top-level Diffusers safetensors index references no files" + expected_indexed_files = { + f"transformer/{path.name}" for path in (model_dir / "transformer").glob("*.safetensors") + } + expected_indexed_files.add("vision_encoder/model.safetensors") + assert indexed_files == expected_indexed_files + indexed_tensor_names: set[str] = set() + for rel_path in sorted(indexed_files): + path = model_dir / rel_path + expected_names = {name for name, filename in top_weight_map.items() if filename == rel_path} + actual_names = _safetensors_keys(path) + assert actual_names == expected_names, ( + f"top-level index/header mismatch for {rel_path}: " + f"missing={sorted(expected_names - actual_names)[:10]}, " + f"extra={sorted(actual_names - expected_names)[:10]}" + ) + indexed_tensor_names.update(actual_names) + assert indexed_tensor_names == set(top_weight_map) + + transformer_index = json.loads( + (model_dir / "transformer/diffusion_pytorch_model.safetensors.index.json").read_text() + ) + transformer_weight_map: dict[str, str] = transformer_index["weight_map"] + top_transformer_weight_map = { + name: filename.removeprefix("transformer/") + for name, filename in top_weight_map.items() + if filename.startswith("transformer/") + } + assert transformer_weight_map == top_transformer_weight_map + assert set(transformer_weight_map.values()) == { + path.name for path in (model_dir / "transformer").glob("*.safetensors") + } + + assert _safetensors_keys(model_dir / "vision_encoder/model.safetensors") + assert _safetensors_keys(model_dir / "vae/diffusion_pytorch_model.safetensors") + assert _safetensors_keys(model_dir / "sound_tokenizer/diffusion_pytorch_model.safetensors") + assert not list(model_dir.rglob("*.distcp")), "Diffusers export unexpectedly contains DCP shards" + + # --- tests ------------------------------------------------------------------- @@ -667,9 +763,10 @@ def test_convert_reasoner_converts_all_qwen_tensors( # ``vision_encoder/`` source bit-for-bit, and a non-trivial subset differs # from stock Qwen3-VL (so (2) actually distinguishes a converted tower # from a stock one — the vision-drop regression). + from safetensors import safe_open + from cosmos_framework.inference.args import OmniSetupOverrides from cosmos_framework.inference.common.args import CheckpointOverrides - from safetensors import safe_open nano_dir = Path( CheckpointOverrides(checkpoint_path="Cosmos3-Nano") @@ -720,6 +817,47 @@ def test_convert_reasoner_converts_all_qwen_tensors( f"LM tensor {k} still equals stock Qwen3-VL (not converted from Cosmos3-Nano)" ) + @pytest.mark.level(2) + @pytest.mark.gpus(4) + def test_convert_model_to_diffusers_exports_complete_nano( + tmp_path: Path, h100_inputs: dict[str, str], request: pytest.FixtureRequest + ) -> None: + """Convert the staged Nano DCP and validate the complete Diffusers output layout and indexes.""" + del h100_inputs # The fixture stages and exports BASE_CHECKPOINT_PATH for the subprocess. + checkpoint_path = Path(os.environ["BASE_CHECKPOINT_PATH"]) + config_path = checkpoint_path / "model/config.json" + assert config_path.is_file(), f"Nano DCP config is missing: {config_path}" + reference_dir = Path(_hf_download(["nvidia/Cosmos3-Nano"])) + + out_dir = tmp_path / "Cosmos3-Nano-Diffusers" + + def cleanup_output() -> None: + shutil.rmtree(out_dir, ignore_errors=True) + assert not out_dir.exists(), f"failed to clean Diffusers test output: {out_dir}" + + request.addfinalizer(cleanup_output) + + env = os.environ.copy() + env["PYTHONPATH"] = f".:{env.get('PYTHONPATH', '')}" + result = subprocess.run( + [ + sys.executable, + "-m", + "cosmos_framework.scripts.convert_model_to_diffusers", + "--checkpoint-path", + str(checkpoint_path), + "--config-file", + str(config_path), + "--no-use-ema-weights", + "-o", + str(out_dir), + ], + cwd=str(REPO_ROOT), + env=env, + ) + assert result.returncode == 0, f"convert_model_to_diffusers failed with exit code {result.returncode}" + _assert_diffusers_export_complete(out_dir, reference_dir) + if MAX_GPUS == 8: From 692c410583a2e1271cdc355ffe2ffe574c2ae625 Mon Sep 17 00:00:00 2001 From: yy-code-nv Date: Wed, 15 Jul 2026 15:37:22 +0800 Subject: [PATCH 26/26] Release 2026-07-14 (i4 b66cfd81) (#109) Automated release from i4. _source_commit: `b66cfd8155158744e16d598b7515fbd0b5081a24-dirty` _dest_commit (base): `755bff48dd66b467995366796e747ffc649dd83c` Co-authored-by: lfengad --- .file_mapping.json | 9 +- .../callbacks/sampled_media_recorder.py | 295 ++++++++++++++++++ .../callbacks/sampled_media_recorder_test.py | 94 ++++++ .../configs/base/defaults/callbacks.py | 10 + .../generator/action/action_processing.py | 15 +- .../data/generator/action/domain_utils.py | 2 + .../data/generator/action/transforms.py | 10 +- .../data/generator/action/transforms_test.py | 1 + .../data/generator/augmentor_provider.py | 2 +- .../augmentors/caption_embedding_keys.py | 23 ++ .../data/generator/joint_dataloader.py | 7 +- .../model/tokenizer/evaluation/metric.py | 54 ++-- .../model/tokenizer/models/dense_runtime.py | 123 ++++++-- .../models/modules/quantizers/residual_vq.py | 97 ++++-- .../tokenizer/models/sparse_autoencoder.py | 6 +- .../model/tokenizer/models/text_decoder.py | 108 +++++-- 16 files changed, 728 insertions(+), 128 deletions(-) create mode 100644 cosmos_framework/callbacks/sampled_media_recorder.py create mode 100644 cosmos_framework/callbacks/sampled_media_recorder_test.py create mode 100644 cosmos_framework/data/generator/augmentors/caption_embedding_keys.py diff --git a/.file_mapping.json b/.file_mapping.json index a084b446..9d9b2eee 100644 --- a/.file_mapping.json +++ b/.file_mapping.json @@ -1,7 +1,7 @@ { - "_source_commit": "c0d7b0044b6eb0239bfcd13a31fbb05dc789cb22-dirty", - "_dest_commit": "0fa3ba479a7c1c5be28f81fe084b9ea544d697c9", - "_generated_at": "2026-07-13T07:36:20Z", + "_source_commit": "b66cfd8155158744e16d598b7515fbd0b5081a24-dirty", + "_dest_commit": "755bff48dd66b467995366796e747ffc649dd83c", + "_generated_at": "2026-07-14T08:42:29Z", "files": { "imaginaire/__init__.py": "cosmos_framework/__init__.py", "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", @@ -201,6 +201,8 @@ "projects/cosmos3/cosmos3/callbacks/norm_monitor.py": "cosmos_framework/callbacks/norm_monitor.py", "projects/cosmos3/cosmos3/callbacks/ofu.py": "cosmos_framework/callbacks/ofu.py", "projects/cosmos3/cosmos3/callbacks/param_count.py": "cosmos_framework/callbacks/param_count.py", + "projects/cosmos3/cosmos3/callbacks/sampled_media_recorder.py": "cosmos_framework/callbacks/sampled_media_recorder.py", + "projects/cosmos3/cosmos3/callbacks/sampled_media_recorder_test.py": "cosmos_framework/callbacks/sampled_media_recorder_test.py", "projects/cosmos3/cosmos3/callbacks/sequence_packing_padding.py": "cosmos_framework/callbacks/sequence_packing_padding.py", "projects/cosmos3/cosmos3/callbacks/sigma_loss_analysis.py": "cosmos_framework/callbacks/sigma_loss_analysis.py", "projects/cosmos3/cosmos3/callbacks/skip_nan_step.py": "cosmos_framework/callbacks/skip_nan_step.py", @@ -256,6 +258,7 @@ "projects/cosmos3/cosmos3/datasets/augmentors/append_fps_frames_for_image.py": "cosmos_framework/data/generator/augmentors/append_fps_frames_for_image.py", "projects/cosmos3/cosmos3/datasets/augmentors/audio_caption.py": "cosmos_framework/data/generator/augmentors/audio_caption.py", "projects/cosmos3/cosmos3/datasets/augmentors/audio_parsing.py": "cosmos_framework/data/generator/augmentors/audio_parsing.py", + "projects/cosmos3/cosmos3/datasets/augmentors/caption_embedding_keys.py": "cosmos_framework/data/generator/augmentors/caption_embedding_keys.py", "projects/cosmos3/cosmos3/datasets/augmentors/caption_filter.py": "cosmos_framework/data/generator/augmentors/caption_filter.py", "projects/cosmos3/cosmos3/datasets/augmentors/cropping.py": "cosmos_framework/data/generator/augmentors/cropping.py", "projects/cosmos3/cosmos3/datasets/augmentors/duration_fps_text_timestamps.py": "cosmos_framework/data/generator/augmentors/duration_fps_text_timestamps.py", diff --git a/cosmos_framework/callbacks/sampled_media_recorder.py b/cosmos_framework/callbacks/sampled_media_recorder.py new file mode 100644 index 00000000..e6e6b195 --- /dev/null +++ b/cosmos_framework/callbacks/sampled_media_recorder.py @@ -0,0 +1,295 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Record the image and video IDs consumed by training. + +The callback is disabled by default. When enabled, every rank buffers the +``__key__`` and ``__url__`` values from batches that reached the training step. +At a configurable interval the buffers are gathered to rank 0, which appends +every consumed sample occurrence to one Lance table. + +The resulting table can be inspected with the Streamlit viewer in +``tools/lance_sample_viewer``. +""" + +from __future__ import annotations + +import json +import os +from datetime import datetime, timezone +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import torch +import torch.distributed as dist + +from cosmos_framework.model._base import ImaginaireModel +from cosmos_framework.utils import log +from cosmos_framework.utils.callback import Callback + +_REMOTE_URI_SCHEMES = {"gs", "s3"} +# Keep the training callback independent of the optional cosmos-sila package. +_LANCE_DATA_STORAGE_VERSION = "2.1" + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _as_list(value: Any, count: int, default: str = "") -> list[str]: + """Normalize scalar or batched metadata to exactly ``count`` strings.""" + if isinstance(value, str): + values = [value] + elif isinstance(value, (list, tuple)): + values = [str(item) for item in value] + elif value is None: + values = [] + else: + values = [str(value)] + + if len(values) == 1 and count > 1: + values *= count + if len(values) < count: + values.extend([default] * (count - len(values))) + return values[:count] + + +class SampledMediaRecorder(Callback): + """Append consumed sample metadata to a Lance table from rank 0. + + Args: + enabled: Enable recording. Disabled by default to avoid training overhead. + output_uri: Destination ``.lance`` table. Local, ``s3://``, and ``gs://`` + URIs are supported. + creds_path: Optional S3-compatible credential JSON used for remote Lance + writes. + flush_every_n_batches: Number of consumed microbatches buffered per rank + between distributed gathers and table appends. + """ + + def __init__( + self, + enabled: bool = False, + output_uri: str = "", + creds_path: str | None = None, + flush_every_n_batches: int = 100, + ) -> None: + super().__init__() + if flush_every_n_batches < 1: + raise ValueError(f"flush_every_n_batches must be >= 1, got {flush_every_n_batches}.") + if enabled and not output_uri.endswith(".lance"): + raise ValueError(f"output_uri must end with '.lance', got {output_uri!r}.") + + self.enabled: bool = enabled + self.output_uri: str = output_uri + self.creds_path: str | None = creds_path + self.flush_every_n_batches: int = flush_every_n_batches + + self._pending: list[dict[str, Any]] = [] + self._batch_index: int = 0 + self._batches_since_flush: int = 0 + + @staticmethod + def _rank() -> int: + if dist.is_available() and dist.is_initialized(): + return dist.get_rank() + return 0 + + @staticmethod + def _media_type(data_batch: dict[str, Any]) -> str: + has_images = data_batch.get("images") is not None + has_video = data_batch.get("video") is not None + if has_images and not has_video: + return "image" + if has_video and not has_images: + return "video" + return "image_video" if has_images and has_video else "unknown" + + def _job_name(self) -> str: + job_config = getattr(getattr(self, "config", None), "job", None) + return str(getattr(job_config, "name", "")) + + def _extract_records(self, data_batch: dict[str, Any], iteration: int, rank: int) -> list[dict[str, Any]]: + sample_ids_raw = data_batch.get("__key__") + if sample_ids_raw is None: + return [] + if isinstance(sample_ids_raw, str): + sample_ids = [sample_ids_raw] + elif isinstance(sample_ids_raw, (list, tuple)): + sample_ids = [str(sample_id) for sample_id in sample_ids_raw] + else: + sample_ids = [str(sample_ids_raw)] + if not sample_ids: + return [] + + count = len(sample_ids) + media_urls = _as_list(data_batch.get("__url__"), count) + dataset_names = _as_list(data_batch.get("dataset_name"), count) + source_names = _as_list(data_batch.get("source_dataset_name"), count) + recorded_at = _utc_now() + media_type = self._media_type(data_batch) + run_id = os.environ.get("SLURM_JOB_ID") or os.environ.get("WANDB_RUN_ID", "local") + + return [ + { + "recorded_at": recorded_at, + "run_id": run_id, + "job_name": self._job_name(), + "iteration": int(iteration), + "batch_index": self._batch_index, + "sample_index": sample_index, + "rank": rank, + "media_type": media_type, + "dataset_name": dataset_names[sample_index], + "source_dataset_name": source_names[sample_index], + "sample_id": sample_id, + "media_url": media_urls[sample_index], + } + for sample_index, sample_id in enumerate(sample_ids) + ] + + @staticmethod + def _table_schema() -> Any: + import pyarrow as pa + + return pa.schema( + [ + pa.field("recorded_at", pa.string(), nullable=False), + pa.field("run_id", pa.string(), nullable=False), + pa.field("job_name", pa.string(), nullable=False), + pa.field("iteration", pa.int64(), nullable=False), + pa.field("batch_index", pa.int64(), nullable=False), + pa.field("sample_index", pa.int32(), nullable=False), + pa.field("rank", pa.int32(), nullable=False), + pa.field("media_type", pa.string(), nullable=False), + pa.field("dataset_name", pa.string(), nullable=False), + pa.field("source_dataset_name", pa.string(), nullable=False), + pa.field("sample_id", pa.string(), nullable=False), + pa.field("media_url", pa.string(), nullable=False), + ] + ) + + def _load_credentials(self) -> dict[str, Any]: + if self.creds_path is None: + raise ValueError("creds_path is required for remote sampled-media output.") + with Path(self.creds_path).open(encoding="utf-8") as handle: + return json.load(handle) + + def _lance_storage_options(self) -> dict[str, str]: + scheme = urlparse(self.output_uri).scheme + if scheme not in _REMOTE_URI_SCHEMES: + return {} + + raw = self._load_credentials() + options: dict[str, str] = { + "aws_access_key_id": str(raw["aws_access_key_id"]), + "aws_secret_access_key": str(raw["aws_secret_access_key"]), + } + endpoint = raw.get("endpoint_url") + if endpoint: + options["aws_endpoint"] = str(endpoint) + if "storage.googleapis.com" in str(endpoint): + options["aws_virtual_hosted_style_request"] = "false" + if raw.get("region_name"): + options["aws_region"] = str(raw["region_name"]) + return options + + def _write_lance_records(self, records: list[dict[str, Any]]) -> None: + if not records: + return + + import lance + import pyarrow as pa + + schema = self._table_schema() + table = pa.Table.from_pylist(records, schema=schema) + storage_options = self._lance_storage_options() + parsed = urlparse(self.output_uri) + is_remote = parsed.scheme in _REMOTE_URI_SCHEMES + if not is_remote: + local_path = Path(parsed.path if parsed.scheme == "file" else self.output_uri) + local_path.parent.mkdir(parents=True, exist_ok=True) + exists = local_path.exists() + else: + try: + lance.dataset(self.output_uri, storage_options=storage_options) + except (FileNotFoundError, OSError): + exists = False + else: + exists = True + + if not exists: + lance.write_dataset( + table, + self.output_uri, + mode="create", + data_storage_version=_LANCE_DATA_STORAGE_VERSION, + enable_v2_manifest_paths=True, + storage_options=storage_options, + ) + return + + existing = lance.dataset(self.output_uri, storage_options=storage_options) + if not existing.schema.equals(schema, check_metadata=False): + raise ValueError( + f"Sample recorder schema mismatch for {self.output_uri!r}: expected {schema}, found {existing.schema}." + ) + lance.write_dataset(table, self.output_uri, mode="append", storage_options=storage_options) + + def _flush(self) -> None: + if not self.enabled: + return + + distributed = dist.is_available() and dist.is_initialized() + rank = self._rank() + if distributed: + gathered: list[list[dict[str, Any]] | None] | None = ( + [None for _ in range(dist.get_world_size())] if rank == 0 else None + ) + dist.gather_object(self._pending, object_gather_list=gathered, dst=0) + else: + gathered = [self._pending] + + error_message = "" + if rank == 0: + try: + records = [record for rank_records in gathered or [] if rank_records for record in rank_records] + self._write_lance_records(records) + if records: + log.info(f"[SampledMediaRecorder] Appended {len(records):,} records to {self.output_uri!r}.") + except Exception as error: + error_message = f"SampledMediaRecorder failed to flush: {error}" + + if distributed: + messages = [error_message] + dist.broadcast_object_list(messages, src=0) + error_message = messages[0] + if error_message: + raise RuntimeError(error_message) + + self._pending.clear() + self._batches_since_flush = 0 + + def on_training_step_batch_end( + self, + model: ImaginaireModel, + data_batch: dict[str, Any], + output_batch: dict[str, Any], + loss: torch.Tensor, # [] + iteration: int = 0, + ) -> None: + del model, output_batch, loss + if not self.enabled: + return + + self._pending.extend(self._extract_records(data_batch, iteration, self._rank())) + self._batch_index += 1 + self._batches_since_flush += 1 + if self._batches_since_flush >= self.flush_every_n_batches: + self._flush() + + def on_train_end(self, model: ImaginaireModel, iteration: int = 0) -> None: + del model, iteration + if self.enabled: + self._flush() diff --git a/cosmos_framework/callbacks/sampled_media_recorder_test.py b/cosmos_framework/callbacks/sampled_media_recorder_test.py new file mode 100644 index 00000000..41a286c5 --- /dev/null +++ b/cosmos_framework/callbacks/sampled_media_recorder_test.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from cosmos_framework.callbacks.sampled_media_recorder import SampledMediaRecorder + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +def test_extract_records_from_consumed_image_batch(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SLURM_JOB_ID", "12345") + callback = SampledMediaRecorder(enabled=True, output_uri="/tmp/samples.lance") + callback.config = SimpleNamespace(job=SimpleNamespace(name="test_experiment")) + batch = { + "images": [object(), object()], + "__key__": ["image-a", "image-b"], + "__url__": ["s3r:profile//bucket/a:0-10", "s3r:profile//bucket/b:10-20"], + "dataset_name": ["images", "images"], + "source_dataset_name": ["source-a", "source-b"], + } + + records = callback._extract_records(batch, iteration=7, rank=3) + + assert [record["sample_id"] for record in records] == ["image-a", "image-b"] + assert [record["media_type"] for record in records] == ["image", "image"] + assert [record["source_dataset_name"] for record in records] == ["source-a", "source-b"] + assert all(record["run_id"] == "12345" for record in records) + assert all(record["iteration"] == 7 for record in records) + assert all(record["rank"] == 3 for record in records) + + +def test_media_type_accepts_batched_tensors() -> None: + images = torch.empty(2, 3, 16, 16) # [B,C,H,W] + video = torch.empty(2, 3, 4, 16, 16) # [B,C,T,H,W] + + assert SampledMediaRecorder._media_type({"images": images}) == "image" + assert SampledMediaRecorder._media_type({"video": video}) == "video" + assert SampledMediaRecorder._media_type({"images": images, "video": video}) == "image_video" + assert SampledMediaRecorder._media_type({}) == "unknown" + + +def test_extract_records_preserves_repeated_sample_occurrences() -> None: + callback = SampledMediaRecorder(enabled=True, output_uri="/tmp/samples.lance") + callback.config = SimpleNamespace(job=SimpleNamespace(name="test_experiment")) + batch = { + "video": [object(), object()], + "__key__": ["video-a", "video-a"], + "__url__": ["s3://bucket/video-a.mp4", "s3://bucket/video-a.mp4"], + "dataset_name": ["video_480", "video_480"], + "source_dataset_name": ["source", "source"], + } + + records = callback._extract_records(batch, iteration=9, rank=0) + + assert [record["sample_id"] for record in records] == ["video-a", "video-a"] + assert [record["sample_index"] for record in records] == [0, 1] + + +def test_local_lance_append_without_cosmos_sila(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + lance = pytest.importorskip("lance") + monkeypatch.setitem(sys.modules, "cosmos_sila", None) + output_uri = str(tmp_path / "samples.lance") + callback = SampledMediaRecorder( + enabled=True, + output_uri=output_uri, + ) + records = [ + { + "recorded_at": "2026-07-03T00:00:00+00:00", + "run_id": "run", + "job_name": "job", + "iteration": iteration, + "batch_index": iteration, + "sample_index": 0, + "rank": 0, + "media_type": "video", + "dataset_name": "video_256", + "source_dataset_name": "source", + "sample_id": f"video-{iteration}", + "media_url": f"s3://bucket/video-{iteration}.mp4", + } + for iteration in (1, 2) + ] + + callback._write_lance_records(records[:1]) + callback._write_lance_records(records[1:]) + + assert lance.dataset(output_uri).count_rows() == 2 diff --git a/cosmos_framework/configs/base/defaults/callbacks.py b/cosmos_framework/configs/base/defaults/callbacks.py index dfe03317..2521e748 100644 --- a/cosmos_framework/configs/base/defaults/callbacks.py +++ b/cosmos_framework/configs/base/defaults/callbacks.py @@ -25,6 +25,7 @@ from cosmos_framework.callbacks.norm_monitor import NormMonitor from cosmos_framework.callbacks.ofu import OFUCallback from cosmos_framework.callbacks.param_count import ParamCount +from cosmos_framework.callbacks.sampled_media_recorder import SampledMediaRecorder from cosmos_framework.callbacks.sequence_packing_padding import SequencePackingPadding from cosmos_framework.callbacks.sigma_loss_analysis import SigmaLossAnalysis from cosmos_framework.callbacks.skip_nan_step import SkipNaNStep @@ -70,6 +71,15 @@ save_s3="${upload_reproducible_setup}", ), sequence_packing_padding=L(SequencePackingPadding)(every_n="${trainer.logging_iter}"), + sampled_media=L(SampledMediaRecorder)( + enabled=False, + output_uri=( + "${oc.env:IMAGINAIRE_OUTPUT_ROOT,/tmp/imaginaire4-output}/" + "${job.project}/${job.group}/${job.name}/sampled_media.lance" + ), + creds_path=None, + flush_every_n_batches=100, + ), mfu=L(MFUCallback)(every_n="${trainer.logging_iter}", grad_accum_iter="${trainer.grad_accum_iter}"), ofu=L(OFUCallback)(every_n="${trainer.logging_iter}"), ) diff --git a/cosmos_framework/data/generator/action/action_processing.py b/cosmos_framework/data/generator/action/action_processing.py index 2d445da6..6008c883 100644 --- a/cosmos_framework/data/generator/action/action_processing.py +++ b/cosmos_framework/data/generator/action/action_processing.py @@ -196,7 +196,7 @@ def make_batched_action_processing_fields( class ActionProcessor: - """Forward and inverse Action tensor processing for a single sample.""" + """Build model-space actions and decode model outputs for a single sample.""" def __init__(self, max_action_dim: int, action_channel_masking: bool = True) -> None: self.max_action_dim = int(max_action_dim) @@ -209,9 +209,17 @@ def preprocess_action( *, action_normalizer: ActionNormalizer | None, ) -> dict[str, Any]: - """Return a sample with normalized, padded action fields and the inverse record.""" + """Return a sample with canonical raw and model-space action fields. + + ``action_raw`` is the exact unnormalized, unpadded action accepted by + this method. It is the canonical external-space target for evaluation. + ``action`` is the normalized and padded model-space representation. + """ raw_action_dim = int(action.shape[-1]) + if "action_raw" in data_dict: + raise ValueError("action_raw is reserved for the canonical ActionProcessor input") + action_raw = action.clone() # [T,D] if action_normalizer is not None: action = action_normalizer.normalize_action(action) # [T,D] if int(action.shape[-1]) != raw_action_dim: @@ -220,6 +228,7 @@ def preprocess_action( ) processed_data_dict = dict(data_dict) + processed_data_dict["action_raw"] = action_raw processed_data_dict["action"] = pad_action_to_max_dim(action, self.max_action_dim) # [T,D_model] record = ActionProcessingRecord( raw_action_dim=raw_action_dim, @@ -243,7 +252,7 @@ def postprocess_action( action: torch.Tensor, record: ActionProcessingRecord, ) -> torch.Tensor: - """Unpad and denormalize a model-space action tensor.""" + """Unpad and denormalize a generated model-space action tensor.""" action = ActionProcessor._unpad_action(action, record.raw_action_dim) # [...,D_raw] if record.action_normalizer is not None: action = record.action_normalizer.denormalize_action(action) # [...,D_raw] diff --git a/cosmos_framework/data/generator/action/domain_utils.py b/cosmos_framework/data/generator/action/domain_utils.py index c75f7f38..66cc51cb 100644 --- a/cosmos_framework/data/generator/action/domain_utils.py +++ b/cosmos_framework/data/generator/action/domain_utils.py @@ -21,6 +21,7 @@ "embodiment_c_gripper": 15, "embodiment_c_gripper_ext": 15, "xdof_yam": 16, + "molmoact2_yam": 16, # MolmoAct2 uses the same YAM 20D FK action contract "fractal": 20, } @@ -40,6 +41,7 @@ "embodiment_c_gripper": 29, "embodiment_c_gripper_ext": 29, "xdof_yam": 20, + "molmoact2_yam": 20, "fractal": 10, # NOTE: ``libero`` (7/10/13 depending on ``rotation_space``) and ``hand_pose`` # (variable with ``keypoint_option`` and ``rotation_format``) are absent diff --git a/cosmos_framework/data/generator/action/transforms.py b/cosmos_framework/data/generator/action/transforms.py index 05696536..e451ee2b 100644 --- a/cosmos_framework/data/generator/action/transforms.py +++ b/cosmos_framework/data/generator/action/transforms.py @@ -617,11 +617,12 @@ def __call__( sample is in inverse dynamics mode (if enabled). 7. Tokenize caption text (if enabled). 8. Build a ``SequencePlan`` from the ``"mode"`` key (if present). - 9. If action is needed by the plan, normalize real channels, pad - ``"action"`` to ``max_action_dim``, and attach + 9. If action is needed by the plan, preserve the canonical unnormalized + and unpadded action as ``"action_raw"``, normalize real channels, + pad ``"action"`` to ``max_action_dim``, and attach ``"action_processing_record"``. - 10. Otherwise, nullify ``"action"`` and ``"domain_id"`` (e.g. in - ``"image2video"`` mode). + 10. Otherwise, nullify ``"action_raw"``, ``"action"``, and + ``"domain_id"`` (e.g. in ``"image2video"`` mode). Args: data_dict: A sample dictionary as returned by a Action dataset. @@ -709,6 +710,7 @@ def __call__( # Nullify action-related fields when action is not needed so the # collate function can simply stack all non-None actions. data_dict["raw_action_dim"] = None + data_dict["action_raw"] = None data_dict["action"] = None data_dict["domain_id"] = None data_dict["action_processing_record"] = None diff --git a/cosmos_framework/data/generator/action/transforms_test.py b/cosmos_framework/data/generator/action/transforms_test.py index 93b62f05..ad182708 100644 --- a/cosmos_framework/data/generator/action/transforms_test.py +++ b/cosmos_framework/data/generator/action/transforms_test.py @@ -169,6 +169,7 @@ def test_action_transform_pipeline_json_prompt_toggle() -> None: assert prompt["resolution"] == {"H": 192, "W": 320} assert prompt["aspect_ratio"] == "16,9" assert result["action"].shape == (16, 4) + torch.testing.assert_close(result["action_raw"], action) @pytest.mark.L0 diff --git a/cosmos_framework/data/generator/augmentor_provider.py b/cosmos_framework/data/generator/augmentor_provider.py index 8840d160..fa2a1c01 100644 --- a/cosmos_framework/data/generator/augmentor_provider.py +++ b/cosmos_framework/data/generator/augmentor_provider.py @@ -582,7 +582,7 @@ def get_video_augmentor_v3( "use_dynamic_fps": use_dynamic_fps, "max_stride": max_stride, "min_stride": min_stride, - "seek_mode": "exact", # Change to "approximate"? + "seek_mode": "exact", "dataset_resolution_type": dataset_resolution_type, "resolution": resolution, "causal_vae": causal_vae, diff --git a/cosmos_framework/data/generator/augmentors/caption_embedding_keys.py b/cosmos_framework/data/generator/augmentors/caption_embedding_keys.py new file mode 100644 index 00000000..ddf579f5 --- /dev/null +++ b/cosmos_framework/data/generator/augmentors/caption_embedding_keys.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Caption-to-embedding key mappings shared by dataset configuration and workers. + +Keep these mappings in a dependency-light leaf module rather than in +``data_sources.data_registration``. Spawned Lance DataLoader workers import the +image augmentor modules; importing a shared constant from ``data_registration`` +would also pull WebDataset registration code and its transitive dependencies +into every worker. Registration and augmentor code can instead import this +module without coupling the Lance runtime import path to WebDataset setup. + +The current Lance JSON-caption pipeline does not dereference this mapping, but +it still imports the image augmentor module that also serves embedding-backed +WebDataset pipelines. +""" + +# embeddings are packed together. Need to clean data to reduce entropy. +_CAPTION_EMBEDDING_KEY_MAPPING_IMAGES: dict[str, str] = { + "ai_v3p1": "ai_v3p1", + "qwen2p5_7b_v4": "qwen2p5_7b_v4", + "prompts": "qwen2p5_7b_v4", +} diff --git a/cosmos_framework/data/generator/joint_dataloader.py b/cosmos_framework/data/generator/joint_dataloader.py index 64f430eb..ddf7c6d4 100644 --- a/cosmos_framework/data/generator/joint_dataloader.py +++ b/cosmos_framework/data/generator/joint_dataloader.py @@ -40,6 +40,7 @@ def custom_collate_fn(batch): "images", "video", "action", + "action_raw", "domain_id", "sequence_plan", "sound", @@ -473,7 +474,7 @@ def __len__(self) -> int: # Keys where each sample may hold multiple tensors (e.g. multiple video # clips in a packed sequence). Kept as single-element lists per sample # via v[i:i+1] so that _update_output_batch yields list[list[Tensor]]. - _MULTI_ITEM_KEYS = {"text_token_ids", "images", "video", "action", "sound"} + _MULTI_ITEM_KEYS = {"text_token_ids", "images", "video", "action", "action_raw", "sound"} def _get_next_sample(self, index_id: int) -> dict: """Pop the next single-sample dict from the buffer for the given dataloader. @@ -494,8 +495,8 @@ def _get_next_sample(self, index_id: int) -> dict: After ``_update_output_batch`` accumulates samples, the packed output batch has the following shapes: - Multi-item keys (``text_token_ids``, ``video``, ``images``, - ``action``): ``list[list[Tensor]]`` — each inner list has one - element from one sub-sample. + ``action_raw``, ``action``): ``list[list[Tensor]]`` — each inner + list has one element from one sub-sample. - Per-sequence metadata keys (``sequence_plan``, ``domain_id``, ``dataset_name``, etc.): ``list[element]`` — flat list. - Tensor-origin keys: ``list[Tensor(1, ...)]``. diff --git a/cosmos_framework/model/tokenizer/evaluation/metric.py b/cosmos_framework/model/tokenizer/evaluation/metric.py index 6659dde6..88e13836 100644 --- a/cosmos_framework/model/tokenizer/evaluation/metric.py +++ b/cosmos_framework/model/tokenizer/evaluation/metric.py @@ -18,8 +18,6 @@ import torch from loguru import logger as logging -from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs - def compute_psnr( original: torch.Tensor, @@ -127,42 +125,26 @@ def compute_codebook_usage( Returns: Dictionary with usage statistics. """ - # Flatten indices - flat_indices = indices.flatten().long() + if num_codes <= 0: + raise ValueError(f"num_codes must be positive, got {num_codes}.") - # Handle empty indices - if flat_indices.numel() == 0: - return { - "perplexity": 0.0, - "active_codes": 0, - "active_ratio": 0.0, - } - - # Gather indices across all GPUs for accurate codebook usage + flat_indices = indices.flatten().long() # [N] + invalid_mask = (flat_indices < 0) | (flat_indices >= num_codes) # [N] + safe_indices = flat_indices.masked_fill(invalid_mask, num_codes) # [N] + histogram_with_invalid = torch.bincount(safe_indices, minlength=num_codes + 1) # [K+1] if torch.distributed.is_initialized(): - world_size = torch.distributed.get_world_size() - # Gather sizes first (indices may have different lengths per GPU) - local_size = torch.tensor([flat_indices.numel()], device=flat_indices.device) - sizes = [torch.zeros(1, dtype=torch.long, device=flat_indices.device) for _ in range(world_size)] - torch.distributed.all_gather(sizes, local_size) - max_size = max(s.item() for s in sizes) - - # Pad indices to max_size for gathering - padded = torch.zeros(max_size, dtype=flat_indices.dtype, device=flat_indices.device) - padded[: flat_indices.numel()] = flat_indices - gathered = [ - torch.zeros(max_size, dtype=flat_indices.dtype, device=flat_indices.device) for _ in range(world_size) - ] - torch.distributed.all_gather(gathered, padded) - - # Concatenate only valid indices from each GPU - all_indices = [] - for i, g in enumerate(gathered): - all_indices.append(g[: sizes[i].item()]) - flat_indices = cat_with_bounded_inputs(all_indices, dim=0) - - # Compute code histogram - histogram = torch.bincount(flat_indices, minlength=num_codes).float() + # Reducing a fixed-size histogram avoids gathering every token index and + # keeps collective participation valid for empty or invalid local ranks. + torch.distributed.all_reduce(histogram_with_invalid, op=torch.distributed.ReduceOp.SUM) + + invalid_count = int(histogram_with_invalid[-1].item()) + if invalid_count > 0: + raise ValueError( + f"Code indices must be in [0, {num_codes}); found {invalid_count} out-of-range value(s) globally." + ) + + histogram_counts = histogram_with_invalid[:-1] # [K] + histogram = histogram_counts.to(dtype=torch.float32) # [K] total = histogram.sum() if total == 0: return { diff --git a/cosmos_framework/model/tokenizer/models/dense_runtime.py b/cosmos_framework/model/tokenizer/models/dense_runtime.py index 57c08b2d..386c852f 100644 --- a/cosmos_framework/model/tokenizer/models/dense_runtime.py +++ b/cosmos_framework/model/tokenizer/models/dense_runtime.py @@ -6,7 +6,9 @@ from __future__ import annotations import math +from collections import OrderedDict from dataclasses import dataclass +from typing import Literal import torch import torch.nn as nn @@ -48,7 +50,8 @@ class DenseGridMetadata: max_seq_len: int -DenseGridMetadataKey = tuple[str, int, int, int, int, str, str] +DenseGridMetadataKey = tuple[str, int, int, int, int, str, str, int] +DenseImageTemporalPadding = Literal["repeat", "zero"] class DenseAutoencoderRuntime(nn.Module): @@ -61,7 +64,9 @@ class DenseAutoencoderRuntime(nn.Module): autoencoder: AutoencoderKL backend: DenseRuntimeBackend - _metadata_cache: dict[DenseGridMetadataKey, DenseGridMetadata] + image_temporal_padding: DenseImageTemporalPadding + metadata_cache_max_entries: int + _metadata_cache: OrderedDict[DenseGridMetadataKey, DenseGridMetadata] def __init__( self, @@ -70,6 +75,8 @@ def __init__( pad_frames: int = 0, pixel_trim: bool = True, chunk_size: int = 16, + image_temporal_padding: DenseImageTemporalPadding = "zero", + metadata_cache_max_entries: int = 32, ) -> None: """Initialize the dense runtime wrapper. @@ -92,8 +99,17 @@ def __init__( ``autoencoder.num_sample_frames_batch_size`` and used to slice the input video into encode batches. Must satisfy ``2 * pad_frames < chunk_size``. Default ``16``. + image_temporal_padding: How a one-frame image fills a temporal + patch when ``patch_size[0] > 1``. ``"zero"`` matches the + canonical sparse path and keeps the first decoded frame. + ``"repeat"`` preserves the legacy deployed dense-runtime + contract and keeps the last decoded frame. + metadata_cache_max_entries: Maximum number of device-resident + dense-grid metadata entries retained by the runtime. Set to + ``0`` to disable caching. """ super().__init__() + self._validate_autoencoder(autoencoder) self.autoencoder = autoencoder self.backend = backend autoencoder.num_sample_frames_batch_size = chunk_size @@ -101,9 +117,17 @@ def __init__( raise ValueError(f"pad_frames must be non-negative, got {pad_frames}.") if 2 * pad_frames >= chunk_size: raise ValueError(f"pad_frames must be less than chunk_size / 2, got {pad_frames=}, {chunk_size=}.") + if image_temporal_padding not in {"repeat", "zero"}: + raise ValueError( + f"image_temporal_padding must be either 'repeat' or 'zero', got {image_temporal_padding!r}." + ) + if metadata_cache_max_entries < 0: + raise ValueError(f"metadata_cache_max_entries must be non-negative, got {metadata_cache_max_entries}.") self.pad_frames = pad_frames self.pixel_trim = pixel_trim - self._metadata_cache: dict[DenseGridMetadataKey, DenseGridMetadata] = {} + self.image_temporal_padding = image_temporal_padding + self.metadata_cache_max_entries = metadata_cache_max_entries + self._metadata_cache: OrderedDict[DenseGridMetadataKey, DenseGridMetadata] = OrderedDict() self.cg_compiled = False @classmethod @@ -114,15 +138,18 @@ def from_autoencoder( pad_frames: int = 0, pixel_trim: bool = True, chunk_size: int = 16, + image_temporal_padding: DenseImageTemporalPadding = "zero", + metadata_cache_max_entries: int = 32, ) -> "DenseAutoencoderRuntime": """Build a dense runtime from a supported sparse autoencoder.""" - cls._validate_autoencoder(autoencoder) return cls( autoencoder=autoencoder, backend=backend, pad_frames=pad_frames, pixel_trim=pixel_trim, chunk_size=chunk_size, + image_temporal_padding=image_temporal_padding, + metadata_cache_max_entries=metadata_cache_max_entries, ) @staticmethod @@ -138,6 +165,18 @@ def _validate_autoencoder(autoencoder: AutoencoderKL) -> None: raise ValueError("Dense runtime V1 does not support concat_latent.") if autoencoder.use_dual_latent: raise ValueError("Dense runtime V1 does not support dual latent.") + if autoencoder.use_quantizer: + raise ValueError("Dense runtime V1 does not support quantized latent paths.") + if autoencoder.decoder_temporal_mode != "bidirectional": + raise ValueError( + "Dense runtime V1 only supports decoder_temporal_mode='bidirectional', " + f"got {autoencoder.decoder_temporal_mode!r}." + ) + if int(autoencoder.inference_kv_cache_size) != 0: + raise ValueError( + "Dense runtime V1 does not support decoder KV cache; " + f"got inference_kv_cache_size={autoencoder.inference_kv_cache_size}." + ) if decoder.multiscale is not None or decoder.multiscale_outputs is not None: raise ValueError("Dense runtime V1 does not support decoder multiscale outputs.") if any(getattr(block, "multiscale", None) is not None for block in encoder.blocks): @@ -193,7 +232,7 @@ def resolve_backend(self, use_compile: bool = False) -> DenseResolvedBackend: return resolve_dense_backend(self.backend, use_compile=use_compile) def clear_metadata_cache(self) -> None: - """Drop cached dense-grid metadata.""" + """Drop cached dense-grid metadata after an eval-time state mutation.""" self._metadata_cache.clear() def encode( @@ -257,9 +296,10 @@ def encode_moments( ``[1, 28, 480, 832, 3]``; the latent fed to a downstream DiT is ``[1, 32, 30, 52, 2C]``. - For images (``T = 1``) the input is repeated to one temporal patch - (``T = patch_time``) and ``latents_per_boundary = 0``, so the - DiT-facing shape is ``[B, 1, H_p, W_p, 2C]``. + For images (``T = 1``) the input is padded to one temporal patch + (``T = patch_time``) according to ``image_temporal_padding`` and + ``latents_per_boundary = 0``, so the DiT-facing shape is + ``[B, 1, H_p, W_p, 2C]``. """ if video.ndim != 5: raise ValueError(f"Dense runtime expects 5D video tensor, got {video.ndim}D") @@ -290,7 +330,11 @@ def encode_moments( # if input is an image, we pad to form single temporal patch if raw_frames == 1: is_image = True - video = video.repeat(1, patch_time, 1, 1, 1) + if self.image_temporal_padding == "repeat": + video = video.repeat(1, patch_time, 1, 1, 1) # [B,Pt,H,W,3] + else: + temporal_padding = video.new_zeros((video.shape[0], patch_time - 1, *video.shape[2:])) # [B,Pt-1,H,W,3] + video = torch.cat((video, temporal_padding), dim=1) # [B,Pt,H,W,3] raw_frames = patch_time else: is_image = False @@ -417,11 +461,11 @@ def decode( **Output shape contract**: - Video (``temporal_patches > 1``): ``[B, T, H, W, C]`` where T is the total number of decoded pixel frames across all chunks (after trim). - - Image (``temporal_patches == 1``): ``[B, 1, H, W, C]``. The image - latent is decoded into ``patch_time`` identical frames (it was encoded - from ``patch_time`` copies of the same frame); only the last frame is - kept. This differs from pre-``dense_runtime`` behaviour where the - full ``[B, patch_time, H, W, C]`` was returned. + - Image (``temporal_patches == 1``): ``[B, 1, H, W, C]``. The image + latent decodes to ``patch_time`` frames. Sparse-compatible zero + padding keeps the first frame; legacy repeat padding keeps the last. + This differs from pre-``dense_runtime`` behavior where the full + ``[B, patch_time, H, W, C]`` was returned. """ if self.decoder_cache_spec.patch_frames != 0: raise NotImplementedError("Dense runtime decoder V1 does not support KV cache.") @@ -447,13 +491,17 @@ def decode( is_image = temporal_patches == 1 # Patch 0 is always a single-latent chunk — either the noncausal first - # frame (video) or the sole image latent. Both were encoded from - # [frame × patch_time] copies, so all decoded frames are equivalent; - # keep the last one. For images temporal_patches == 1, so the loop - # below is empty and this is the only chunk. + # frame (video) or the sole image latent. The first video frame was + # repeated across one temporal patch, so keep the last decoded frame. + # Images keep either the last frame for the deployed repeat contract or + # the first frame for sparse-compatible zero padding. For images + # temporal_patches == 1, so the loop below is empty. decoded_chunks: list[torch.Tensor] = [] decoded_first = self._decode_latent_chunk(latent[:, 0:1]) # [B, patch_time, H, W, C] - decoded_chunks.append(decoded_first[:, -1:]) + if is_image and self.image_temporal_padding == "zero": + decoded_chunks.append(decoded_first[:, :1]) # [B,1,H,W,C] + else: + decoded_chunks.append(decoded_first[:, -1:]) # [B,1,H,W,C] for start_patch in range(1, temporal_patches, chunk_patch_frames): end_patch = min(start_patch + chunk_patch_frames, temporal_patches) @@ -468,6 +516,7 @@ def decode( def _metadata_cache_key( self, module_name: str, + module: SparseTransformerBase, batch_size: int, temporal_patches: int, height_patches: int, @@ -476,6 +525,10 @@ def _metadata_cache_key( dtype: torch.dtype, ) -> DenseGridMetadataKey: """Build a stable metadata-cache key for one dense grid shape.""" + if isinstance(module.pos_embedder, LearnedPositionEmbedder): + position_embedding_pointer = module.pos_embedder.position_embedding.weight.data_ptr() + else: + position_embedding_pointer = -1 return ( module_name, int(batch_size), @@ -484,8 +537,15 @@ def _metadata_cache_key( int(width_patches), str(device), str(dtype), + position_embedding_pointer, ) + def train(self, mode: bool = True) -> "DenseAutoencoderRuntime": + """Set module mode and discard metadata tied to prior parameter state.""" + super().train(mode) + self._metadata_cache.clear() + return self + def _raw_frames_to_patch_frames(self, raw_frames: int) -> int: """Convert raw video frames into temporal patch steps.""" patch_time = self.patch_size[0] @@ -819,6 +879,7 @@ def _get_or_build_grid_metadata( """Fetch or create dense-grid metadata for one uniform chunk shape.""" key = self._metadata_cache_key( module_name, + module, batch_size, temporal_patches, height_patches, @@ -826,8 +887,30 @@ def _get_or_build_grid_metadata( device, dtype, ) + learned_position_requires_grad = bool( + module.pe_mode in {"joint", "learned"} + and isinstance(module.pos_embedder, LearnedPositionEmbedder) + and module.pos_embedder.position_embedding.weight.requires_grad + ) + # Trainable learned positions are cacheable only during eval without + # gradients. Entering either train or eval mode clears prior metadata, + # while the storage pointer above detects parameter replacement. + cache_enabled = self.metadata_cache_max_entries > 0 and not ( + learned_position_requires_grad and (self.training or torch.is_grad_enabled()) + ) + if not cache_enabled: + return self._build_grid_metadata( + module=module, + batch_size=batch_size, + temporal_patches=temporal_patches, + height_patches=height_patches, + width_patches=width_patches, + device=device, + ) + cached = self._metadata_cache.get(key) if cached is not None: + self._metadata_cache.move_to_end(key) return cached metadata = self._build_grid_metadata( @@ -839,6 +922,8 @@ def _get_or_build_grid_metadata( device=device, ) self._metadata_cache[key] = metadata + if len(self._metadata_cache) > self.metadata_cache_max_entries: + self._metadata_cache.popitem(last=False) return metadata def _build_grid_metadata( diff --git a/cosmos_framework/model/tokenizer/models/modules/quantizers/residual_vq.py b/cosmos_framework/model/tokenizer/models/modules/quantizers/residual_vq.py index 806ac123..ccfd9e19 100644 --- a/cosmos_framework/model/tokenizer/models/modules/quantizers/residual_vq.py +++ b/cosmos_framework/model/tokenizer/models/modules/quantizers/residual_vq.py @@ -41,6 +41,8 @@ class VQEmbedding(nn.Embedding): decay: EMA decay rate. restart_unused_codes: Whether to restart unused codebook entries. eps: Small epsilon for numerical stability. + distance_chunk_size: Maximum number of input and codebook rows in one + distance tile. ``None`` computes the full distance matrix. """ def __init__( @@ -51,15 +53,20 @@ def __init__( decay: float = 0.99, restart_unused_codes: bool = True, eps: float = 1e-5, - ): + distance_chunk_size: int | None = 4096, + ) -> None: """Initialize VQEmbedding.""" super().__init__(n_embed + 1, embed_dim, padding_idx=n_embed) + if distance_chunk_size is not None and distance_chunk_size <= 0: + raise ValueError(f"distance_chunk_size must be positive when set, got {distance_chunk_size}.") + self.ema = ema self.decay = decay self.eps = eps self.restart_unused_codes = restart_unused_codes self.n_embed = n_embed + self.distance_chunk_size: int | None = distance_chunk_size if self.ema: _ = [p.requires_grad_(False) for p in self.parameters()] @@ -78,22 +85,22 @@ def compute_distances(self, inputs: torch.Tensor) -> torch.Tensor: Returns: Distance tensor of shape (..., n_embed). """ - codebook_t = self.weight[:-1, :].t() + codebook_t = self.weight[:-1, :].t() # [E,K] (embed_dim, _) = codebook_t.shape inputs_shape = inputs.shape assert inputs_shape[-1] == embed_dim - inputs_flat = inputs.reshape(-1, embed_dim) + inputs_flat = inputs.reshape(-1, embed_dim).to(dtype=codebook_t.dtype) # [N,E] - inputs_norm_sq = inputs_flat.pow(2.0).sum(dim=1, keepdim=True) - codebook_t_norm_sq = codebook_t.pow(2.0).sum(dim=0, keepdim=True) + inputs_norm_sq = inputs_flat.pow(2.0).sum(dim=1, keepdim=True) # [N,1] + codebook_t_norm_sq = codebook_t.pow(2.0).sum(dim=0, keepdim=True) # [1,K] distances = torch.addmm( inputs_norm_sq + codebook_t_norm_sq, inputs_flat, codebook_t, alpha=-2.0, - ) - distances = distances.reshape(*inputs_shape[:-1], -1) + ) # [N,K] + distances = distances.reshape(*inputs_shape[:-1], -1) # [...,K] return distances @torch.no_grad() @@ -106,9 +113,47 @@ def find_nearest_embedding(self, inputs: torch.Tensor) -> torch.Tensor: Returns: Index tensor of shape (...,). """ - distances = self.compute_distances(inputs) - embed_idxs = distances.argmin(dim=-1) - return embed_idxs + input_shape = inputs.shape + embed_dim = input_shape[-1] + inputs_flat = inputs.reshape(-1, embed_dim).to(dtype=self.weight.dtype) # [N,E] + codebook = self.weight[:-1, :] # [K,E] + chunk_size = self.distance_chunk_size + if chunk_size is None: + distances = self.compute_distances(inputs) # [...,K] + return distances.argmin(dim=-1) # [...] + + best_indices = torch.zeros(inputs_flat.shape[0], dtype=torch.long, device=inputs.device) # [N] + for input_start in range(0, inputs_flat.shape[0], chunk_size): + input_end = min(input_start + chunk_size, inputs_flat.shape[0]) + input_chunk = inputs_flat[input_start:input_end] # [Nc,E] + input_norm_sq = input_chunk.pow(2.0).sum(dim=1, keepdim=True) # [Nc,1] + input_best_distances = input_chunk.new_full((input_chunk.shape[0],), float("inf")) # [Nc] + input_best_indices = torch.zeros(input_chunk.shape[0], dtype=torch.long, device=inputs.device) # [Nc] + for codebook_start in range(0, codebook.shape[0], chunk_size): + codebook_end = min(codebook_start + chunk_size, codebook.shape[0]) + codebook_chunk = codebook[codebook_start:codebook_end] # [Kc,E] + codebook_norm_sq = codebook_chunk.pow(2.0).sum(dim=1).unsqueeze(0) # [1,Kc] + chunk_distances = torch.addmm( + input_norm_sq + codebook_norm_sq, + input_chunk, + codebook_chunk.t(), + alpha=-2.0, + ) # [Nc,Kc] + chunk_best_distances, chunk_best_indices = chunk_distances.min(dim=1) # [Nc], [Nc] + use_chunk = chunk_best_distances < input_best_distances # [Nc] + input_best_distances = torch.where( # [Nc] + use_chunk, + chunk_best_distances, + input_best_distances, + ) + input_best_indices = torch.where( # [Nc] + use_chunk, + chunk_best_indices + codebook_start, + input_best_indices, + ) + best_indices[input_start:input_end] = input_best_indices # [Nc] + + return best_indices.reshape(input_shape[:-1]) # [...] @torch.no_grad() def _tile_with_noise(self, x: torch.Tensor, target_n: int) -> torch.Tensor: @@ -138,19 +183,19 @@ def _update_buffers(self, vectors: torch.Tensor, idxs: torch.Tensor) -> None: """ n_embed, embed_dim = self.weight.shape[0] - 1, self.weight.shape[-1] - vectors = vectors.reshape(-1, embed_dim) - idxs = idxs.reshape(-1) + vectors = vectors.reshape(-1, embed_dim).to(dtype=self.embed_ema.dtype) # [N,E] + idxs = idxs.reshape(-1).long() # [N] n_vectors = vectors.shape[0] - n_total_embed = n_embed - - one_hot_idxs = vectors.new_zeros(n_total_embed, n_vectors) - one_hot_idxs.scatter_(dim=0, index=idxs.unsqueeze(0), src=vectors.new_ones(1, n_vectors)) - - cluster_size = one_hot_idxs.sum(dim=1) - vectors_sum_per_cluster = one_hot_idxs @ vectors - - if dist.is_initialized(): + cluster_size = torch.bincount(idxs, minlength=n_embed).to(dtype=self.cluster_size_ema.dtype) # [K] + vectors_sum_per_cluster = torch.zeros_like(self.embed_ema) # [K,E] + # PyTorch selects its deterministic CUDA implementation here when deterministic algorithms are enabled. + vectors_sum_per_cluster.index_add_(0, idxs, vectors) # [K,E] + + # Collectives are no-ops at world size one. Skipping them also avoids + # backend/device mismatches such as CPU tensors with an NCCL group. + is_multi_rank: bool = dist.is_initialized() and dist.get_world_size() > 1 + if is_multi_rank: dist.all_reduce(vectors_sum_per_cluster, op=dist.ReduceOp.SUM) dist.all_reduce(cluster_size, op=dist.ReduceOp.SUM) @@ -163,7 +208,7 @@ def _update_buffers(self, vectors: torch.Tensor, idxs: torch.Tensor) -> None: n_vectors = vectors.shape[0] _vectors_random = vectors[torch.randperm(n_vectors, device=vectors.device)][:n_embed] - if dist.is_initialized(): + if is_multi_rank: dist.broadcast(_vectors_random, 0) usage = (self.cluster_size_ema.view(-1, 1) >= 1).float() @@ -194,7 +239,7 @@ def forward(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: if self.ema: self._update_buffers(inputs, embed_idxs) - embeds = self.embed(embed_idxs) + embeds = self.embed(embed_idxs).to(dtype=inputs.dtype) # [...,E] if self.ema and self.training: self._update_embedding() @@ -229,6 +274,7 @@ class RQBottleneck(nn.Module): shared_codebook: Whether to share codebook across depth. restart_unused_codes: Whether to restart unused codes. commitment_loss: Type of commitment loss ("cumsum"). + distance_chunk_size: Maximum input/codebook rows in each distance tile. """ def __init__( @@ -240,7 +286,8 @@ def __init__( shared_codebook: bool = False, restart_unused_codes: bool = True, commitment_loss: str = "cumsum", - ): + distance_chunk_size: int | None = 4096, + ) -> None: """Initialize RQBottleneck.""" super().__init__() @@ -275,6 +322,7 @@ def __init__( embed_dim, decay=self.decay[0], restart_unused_codes=restart_unused_codes, + distance_chunk_size=distance_chunk_size, ) self.codebooks = nn.ModuleList([codebook0 for _ in range(self.code_shape[-1])]) else: @@ -284,6 +332,7 @@ def __init__( embed_dim, decay=self.decay[idx], restart_unused_codes=restart_unused_codes, + distance_chunk_size=distance_chunk_size, ) for idx in range(self.code_shape[-1]) ] diff --git a/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py b/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py index 40b79fd0..5de038d0 100644 --- a/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py +++ b/cosmos_framework/model/tokenizer/models/sparse_autoencoder.py @@ -1555,6 +1555,7 @@ def decode_text( image_feats: "SparseTensor", image_patch_indices: torch.Tensor, segment_ids: torch.Tensor | None = None, + return_hidden_states: bool = False, ) -> tuple[Any, int]: """Decode text from image features using the configured text decoder. @@ -1566,9 +1567,11 @@ def decode_text( image_patch_indices: [N_pooled] flat indices into [B*S]. segment_ids: [B, S] segment IDs for packed sequences. Enables segment-isolated attention with per-segment position reset. + return_hidden_states: Return decoder hidden states instead of full + vocabulary logits for chunked training loss projection. Returns: - Tuple of (lm_logits [B, S, vocab_size], num_pooled_tokens int). + Tuple of decoder output [B,S,Vocab] or [B,S,D] and pooled-token count. """ if self.text_decoder_wrapper is None: raise RuntimeError("Text decoder not initialized. Set use_text_decoder=True and text_decoder_model_name.") @@ -1579,6 +1582,7 @@ def decode_text( image_patch_indices=image_patch_indices, image_layout=image_feats.layout if hasattr(image_feats, "layout") else None, segment_ids=segment_ids, + return_hidden_states=return_hidden_states, ) def encode_text(self, text: torch.Tensor, normalize: bool = False) -> torch.Tensor: diff --git a/cosmos_framework/model/tokenizer/models/text_decoder.py b/cosmos_framework/model/tokenizer/models/text_decoder.py index 3cddc135..0eb922df 100644 --- a/cosmos_framework/model/tokenizer/models/text_decoder.py +++ b/cosmos_framework/model/tokenizer/models/text_decoder.py @@ -951,6 +951,11 @@ def _lm_head(self, hidden_states: torch.Tensor) -> torch.Tensor: """Project decoder hidden states to token logits.""" return self.text_decoder.lm_head(hidden_states) + @property + def output_projection(self) -> nn.Module: + """Return the LM head used by the chunked training loss.""" + return self.text_decoder.lm_head + def _sample_next_token( self, logits: torch.Tensor, @@ -1093,6 +1098,7 @@ def forward( image_patch_indices: torch.Tensor, image_layout: list[slice] | None = None, segment_ids: torch.Tensor | None = None, + return_hidden_states: bool = False, ) -> tuple[torch.Tensor, int]: """Forward pass: inject image features into text and run causal LM. @@ -1112,9 +1118,13 @@ def forward( segment_ids: [B, S] segment IDs for packed sequences. Values >= 0 indicate valid segments, -1 indicates padding. When provided, enables segment-isolated attention and per-segment position IDs. + return_hidden_states: Return decoder hidden states instead of projecting + the complete sequence to vocabulary logits. This is used by the + chunked training loss to bound LM-head memory. Returns: - Dense [B, S, vocab_size] logits. + Dense [B, S, vocab_size] logits, or [B, S, hidden_size] hidden + states when ``return_hidden_states=True``. num_pooled_tokens: Number of image tokens after spatial merging. """ image_features = image_feats_tensor @@ -1175,23 +1185,39 @@ def forward( # Segment-aware forward pass for packed sequences if segment_ids is not None: - lm_logits = self._forward_packed( - text_embeds, - input_ids, - segment_ids, - ) + if return_hidden_states: + decoder_output = self._forward_packed( # [B,S,D] + text_embeds, + input_ids, + segment_ids, + return_hidden_states=True, + ) + else: + decoder_output = self._forward_packed( # [B,S,Vocab] + text_embeds, + input_ids, + segment_ids, + ) else: - lm_logits = self._forward_standard( - text_embeds, - input_ids=input_ids, - ) + if return_hidden_states: + decoder_output = self._forward_standard( # [B,S,D] + text_embeds, + input_ids=input_ids, + return_hidden_states=True, + ) + else: + decoder_output = self._forward_standard( # [B,S,Vocab] + text_embeds, + input_ids=input_ids, + ) - return lm_logits, num_pooled_tokens + return decoder_output, num_pooled_tokens def _forward_standard( self, text_embeds: torch.Tensor, input_ids: torch.Tensor | None = None, + return_hidden_states: bool = False, ) -> torch.Tensor: """Standard forward pass without segment isolation.""" effective_text_embeds = text_embeds @@ -1223,19 +1249,27 @@ def _forward_standard( cache_position=cache_position, use_cache=False, ) - lm_logits = self._lm_head(outputs.last_hidden_state) + decoder_hidden_states = outputs.last_hidden_state # [B,T,D] + decoder_output = ( + decoder_hidden_states if return_hidden_states else self._lm_head(decoder_hidden_states) + ) # [B,T,D|Vocab] if effective_seq_len == original_seq_len: - return lm_logits + return decoder_output pad_len = original_seq_len - effective_seq_len - padded_logits = lm_logits.new_zeros(lm_logits.shape[0], pad_len, lm_logits.shape[-1]) - return torch.cat([lm_logits, padded_logits], dim=1) + padded_output = decoder_output.new_zeros( # [B,S-T,D|Vocab] + decoder_output.shape[0], + pad_len, + decoder_output.shape[-1], + ) + return torch.cat([decoder_output, padded_output], dim=1) # [B,S,D|Vocab] def _forward_packed( self, text_embeds: torch.Tensor, input_ids: torch.Tensor, segment_ids: torch.Tensor, + return_hidden_states: bool = False, ) -> torch.Tensor: """Segment-isolated forward pass for packed sequences. @@ -1249,14 +1283,15 @@ def _forward_packed( segment_ids: [B, S] segment IDs (>=0 for valid, -1 for padding). Returns: - Dense [B, S, vocab_size] logits. + Dense [B, S, vocab_size] logits, or [B, S, hidden_size] hidden + states when ``return_hidden_states=True``. """ B, S, d = text_embeds.shape device = text_embeds.device dtype = text_embeds.dtype # Process each batch element (typically B=1 for packed sequences) - all_logits: list[torch.Tensor] = [] + all_outputs: list[torch.Tensor] = [] for b in range(B): seg_ids = segment_ids[b] # [S] embeds = text_embeds[b] # [S, d] @@ -1265,8 +1300,8 @@ def _forward_packed( valid_indices = (seg_ids >= 0).nonzero(as_tuple=True)[0] if valid_indices.numel() == 0: # All padding — return zeros - V = self.lm_config.vocab_size - all_logits.append(torch.zeros(S, V, device=device, dtype=dtype)) + output_size = self.lm_config.hidden_size if return_hidden_states else self.lm_config.vocab_size + all_outputs.append(torch.zeros(S, output_size, device=device, dtype=dtype)) # [S,D|Vocab] continue valid_seg_ids = seg_ids.index_select(0, valid_indices) @@ -1303,10 +1338,13 @@ def _forward_packed( packed_cumulative_seqlen=cumulative_seqlen, packed_max_seqlen=max_seqlen, ) - valid_logits = self._lm_head(outputs.last_hidden_state) # [1,V,Vocab] - packed_logits = valid_logits.new_zeros((S, valid_logits.shape[-1])) # [S,Vocab] - packed_logits[valid_indices] = valid_logits[0] # [V,Vocab] - all_logits.append(packed_logits) + valid_hidden_states = outputs.last_hidden_state # [1,V,D] + valid_output = ( + valid_hidden_states if return_hidden_states else self._lm_head(valid_hidden_states) + ) # [1,V,D|Vocab] + packed_output = valid_output.new_zeros((S, valid_output.shape[-1])) # [S,D|Vocab] + packed_output[valid_indices] = valid_output[0] # [V,D|Vocab] + all_outputs.append(packed_output) continue segment_rows = torch.repeat_interleave( @@ -1327,17 +1365,19 @@ def _forward_packed( cache_position=position_template, use_cache=False, ) - seg_logits = self._lm_head(outputs.last_hidden_state) # [R,T,Vocab] - packed_logits = seg_logits.new_zeros((S, seg_logits.shape[-1])) # [S,Vocab] - packed_logits[valid_indices] = seg_logits[segment_rows, segment_positions] # [V,Vocab] - all_logits.append(packed_logits) - - # Guard: empty batch (all samples skipped by collate) -> empty logits tensor. - if len(all_logits) == 0: - V = self.lm_config.vocab_size - return torch.zeros(0, S, V, device=device, dtype=dtype) - # Stack batch: [B, S, vocab_size] - return stack_with_bounded_inputs(all_logits, dim=0) # [B,S,Vocab] + segment_hidden_states = outputs.last_hidden_state # [R,T,D] + segment_output = ( + segment_hidden_states if return_hidden_states else self._lm_head(segment_hidden_states) + ) # [R,T,D|Vocab] + packed_output = segment_output.new_zeros((S, segment_output.shape[-1])) # [S,D|Vocab] + packed_output[valid_indices] = segment_output[segment_rows, segment_positions] # [V,D|Vocab] + all_outputs.append(packed_output) + + # Guard: empty batch (all samples skipped by collate) -> empty output tensor. + if len(all_outputs) == 0: + output_size = self.lm_config.hidden_size if return_hidden_states else self.lm_config.vocab_size + return torch.zeros(0, S, output_size, device=device, dtype=dtype) # [0,S,D|Vocab] + return stack_with_bounded_inputs(all_outputs, dim=0) # [B,S,D|Vocab] def _decode_generation_result( self,