From acd2d7d21d5381bd1915dc6237e2dd7d5350e458 Mon Sep 17 00:00:00 2001 From: yangyangt Date: Thu, 2 Jul 2026 06:50:25 -0700 Subject: [PATCH] =?UTF-8?q?Release:=20vfm=20=E2=86=92=20generator,=20vlm?= =?UTF-8?q?=20=E2=86=92=20reasoner=20dest=20rename=20(moves=20+=20import?= =?UTF-8?q?=20paths=20only)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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.