Release 2026-07-27 (i4 a78ca94d)

.file_mapping.json

@@ -1,7 +1,7 @@
 {
-  "_source_commit": "97901d395925c227163a7370ec2763b4057783d5",
-  "_dest_commit": "058c8c046877f087e07b3472a76cac214f3ad39c",
-  "_generated_at": "2026-07-23T11:16:20Z",
+  "_source_commit": "a78ca94df154c804f1c32c0dff6be38a3d22abeb-dirty",
+  "_dest_commit": "2190a72a5799ab1bc8e4ddc19ae09cd1f14d7cbc",
+  "_generated_at": "2026-07-27T08:25:15Z",
   "files": {
     "imaginaire/__init__.py": "cosmos_framework/__init__.py",
     "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py",
@@ -334,6 +334,8 @@
     "projects/cosmos3/cosmos3/models/mot/cosmos3_vfm_network.py": "cosmos_framework/model/generator/mot/cosmos3_vfm_network.py",
     "projects/cosmos3/cosmos3/models/mot/domain_aware_linear.py": "cosmos_framework/model/generator/mot/domain_aware_linear.py",
     "projects/cosmos3/cosmos3/models/mot/dot_product_attention.py": "cosmos_framework/model/generator/mot/dot_product_attention.py",
+    "projects/cosmos3/cosmos3/models/mot/flex_attention.py": "cosmos_framework/model/generator/mot/flex_attention.py",
+    "projects/cosmos3/cosmos3/models/mot/flex_attention_test.py": "cosmos_framework/model/generator/mot/flex_attention_test.py",
     "projects/cosmos3/cosmos3/models/mot/inference_text_kv_memory.py": "cosmos_framework/model/generator/mot/inference_text_kv_memory.py",
     "projects/cosmos3/cosmos3/models/mot/modeling_utils.py": "cosmos_framework/model/generator/mot/modeling_utils.py",
     "projects/cosmos3/cosmos3/models/mot/parallelize_unified_mot.py": "cosmos_framework/model/generator/mot/parallelize_unified_mot.py",

cosmos_framework/checkpoint/dcp.py

@@ -898,7 +898,6 @@ def load(
 
                     # Use rank-specific key for RNG state to support correct per-rank restoration
                     rng_key = f"rng_state_{dist.get_rank()}"
-                    current_rng_state = get_rand_state_dict()
                     _state_dict = {
                         "grad_scaler": grad_scaler.state_dict(),
                         "iteration": iteration,
@@ -909,7 +908,7 @@ def load(
                         k.startswith(f"{rng_key}.") or k == rng_key for k in metadata.state_dict_metadata.keys()
                     )
                     if rng_key_exists:
-                        _state_dict[rng_key] = current_rng_state
+                        _state_dict[rng_key] = get_rand_state_dict()
 
                     dcp.load(
                         _state_dict,
@@ -919,7 +918,8 @@ def load(
                     )
                     grad_scaler.load_state_dict(_state_dict["grad_scaler"])
                     iteration = _state_dict["iteration"]
-                    set_rand_state_dict(_state_dict.get(rng_key, current_rng_state))
+                    if rng_key_exists:
+                        set_rand_state_dict(_state_dict[rng_key])
 
                 elif key == "dataloader":
                     if not easy_io.exists(cur_key_ckpt_full_path, backend_key=self.load_s3_backend_key):

cosmos_framework/data/generator/action/domain_utils.py

@@ -23,6 +23,7 @@
     "xdof_yam": 16,
     "molmoact2_yam": 16,  # MolmoAct2 uses the same YAM 20D FK action contract
     "fractal": 20,
+    "drawanything": 21,
 }
 
 
@@ -43,6 +44,7 @@
     "xdof_yam": 20,
     "molmoact2_yam": 20,
     "fractal": 10,
+    "drawanything": 3,
     # NOTE: ``libero`` (7/10/13 depending on ``rotation_space``) and ``hand_pose``
     # (variable with ``keypoint_option`` and ``rotation_format``) are absent
     # because their raw width is set per-dataset at construction time. Inference

cosmos_framework/data/generator/action/viewpoint_utils.py

@@ -15,13 +15,14 @@
 from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor
 from cosmos_framework.utils import log
 
-Viewpoint = Literal["ego_view", "third_person_view", "wrist_view", "concat_view"]
+Viewpoint = Literal["ego_view", "third_person_view", "wrist_view", "concat_view", "top_down_2d_view"]
 
 DEFAULT_VIEWPOINT_TEMPLATES: dict[str, str] = {
     "ego_view": "This video is captured from a first-person perspective looking at the scene.",
     "third_person_view": "This video is captured from a third-person perspective looking towards the agent from the front.",
     "wrist_view": "This video is captured from a wrist-mounted camera.",
     "concat_view": "This video contains concatenated views from multiple camera perspectives.",
+    "top_down_2d_view": "This video is captured from a top-down view of a flat 2D scene.",
 }
 
 

cosmos_framework/inference/common/distillation_export.py

@@ -3,9 +3,45 @@
 
 """Portable student-only checkpoint export helpers."""
 
+import warnings
 from collections.abc import Callable
+from pathlib import Path, PurePath
 from typing import Any
 
+_PUBLIC_WAN_VAE_PATHS = {
+    "Wan2.2_VAE.pth": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth",
+}
+_PUBLIC_QWEN3_VL_CONFIG_PATHS = {
+    filename: f"cosmos_framework/model/generator/reasoner/qwen3_vl/configs/{filename}"
+    for filename in (
+        "Qwen3-VL-2B-Instruct.json",
+        "Qwen3-VL-4B-Instruct.json",
+        "Qwen3-VL-8B-Instruct.json",
+        "Qwen3-VL-32B-Instruct.json",
+    )
+}
+
+
+def _normalize_public_dependency_path(
+    value: Any,
+    *,
+    field_name: str,
+    public_paths: dict[str, str],
+) -> str:
+    if not isinstance(value, str):
+        raise TypeError(f"Expected {field_name} to be a string.")
+    filename = PurePath(value).name
+    if filename in public_paths:
+        return public_paths[filename]
+    if Path(value).is_absolute():
+        warnings.warn(
+            f"{field_name} contains an unrecognized absolute path and the exported artifact may not be portable: "
+            f"{value}",
+            UserWarning,
+            stacklevel=3,
+        )
+    return value
+
 
 def build_student_checkpoint_metadata(*, use_ema_weights: bool) -> dict[str, str | bool]:
     """Build portable metadata without source checkpoint or credential paths."""
@@ -59,11 +95,29 @@ def sanitize_student_public_model_config(
             tokenizer_config["bucket_name"] = "bucket"
         if "object_store_credential_path_pretrained" in tokenizer_config:
             tokenizer_config["object_store_credential_path_pretrained"] = ""
+        if tokenizer_key == "tokenizer" and "vae_path" in tokenizer_config:
+            tokenizer_config["vae_path"] = _normalize_public_dependency_path(
+                tokenizer_config["vae_path"],
+                field_name="tokenizer.vae_path",
+                public_paths=_PUBLIC_WAN_VAE_PATHS,
+            )
 
     vlm_config = config.get("vlm_config")
     if not isinstance(vlm_config, dict):
         return
 
+    model_instance = vlm_config.get("model_instance")
+    if isinstance(model_instance, dict):
+        model_instance_config = model_instance.get("config")
+        if isinstance(model_instance_config, dict):
+            base_config = model_instance_config.get("base_config")
+            if isinstance(base_config, dict) and "json_file" in base_config:
+                base_config["json_file"] = _normalize_public_dependency_path(
+                    base_config["json_file"],
+                    field_name="vlm_config.model_instance.config.base_config.json_file",
+                    public_paths=_PUBLIC_QWEN3_VL_CONFIG_PATHS,
+                )
+
     pretrained_weights = vlm_config.get("pretrained_weights")
     if isinstance(pretrained_weights, dict):
         pretrained_weights["enabled"] = False

cosmos_framework/inference/common/distillation_export_test.py

@@ -1,6 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: OpenMDW-1.1
 
+import copy
+
+import pytest
 
 from cosmos_framework.inference.common import distillation_export
 from cosmos_framework.inference.common.distillation_export import (
@@ -75,7 +78,7 @@ def test_sanitize_student_public_model_config_removes_internal_loaders() -> None
             "tokenizer": {
                 "bucket_name": "internal-checkpoint-bucket",
                 "object_store_credential_path_pretrained": "/path/to/source.secret",
-                "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth",
+                "vae_path": "/lustre/training/checkpoints/wan22_vae/Wan2.2_VAE.pth",
             },
             "sound_tokenizer": {
                 "bucket_name": "internal-checkpoint-bucket",
@@ -97,6 +100,16 @@ def test_sanitize_student_public_model_config_removes_internal_loaders() -> None
                     "config_variant": "gcp",
                     "pretrained_model_name": "Qwen/Qwen3-VL-32B-Instruct",
                 },
+                "model_instance": {
+                    "config": {
+                        "base_config": {
+                            "json_file": (
+                                "/checkout/cosmos-framework/cosmos_framework/model/generator/"
+                                "reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json"
+                            )
+                        }
+                    }
+                },
             },
         }
     }
@@ -134,10 +147,141 @@ def test_sanitize_student_public_model_config_removes_internal_loaders() -> None
                     "config_variant": "hf",
                     "pretrained_model_name": "Qwen/Qwen3-VL-32B-Instruct",
                 },
+                "model_instance": {
+                    "config": {
+                        "base_config": {
+                            "json_file": (
+                                "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json"
+                            )
+                        }
+                    }
+                },
+            },
+        }
+    }
+
+
+@pytest.mark.parametrize("model_size", ("2B", "4B", "8B", "32B"))
+def test_sanitize_student_public_model_config_preserves_qwen_variant(model_size: str) -> None:
+    filename = f"Qwen3-VL-{model_size}-Instruct.json"
+    model_dict = {
+        "config": {
+            "vlm_config": {
+                "model_instance": {
+                    "config": {"base_config": {"json_file": f"/checkout/private/qwen3_vl/configs/{filename}"}}
+                }
+            }
+        }
+    }
+
+    distillation_export.sanitize_student_public_model_config(model_dict)
+
+    assert (
+        model_dict["config"]["vlm_config"]["model_instance"]["config"]["base_config"]["json_file"]
+        == f"cosmos_framework/model/generator/reasoner/qwen3_vl/configs/{filename}"
+    )
+
+
+def test_sanitize_student_public_model_config_is_idempotent() -> None:
+    model_dict = {
+        "config": {
+            "tokenizer": {
+                "bucket_name": "bucket",
+                "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth",
+            },
+            "vlm_config": {
+                "model_instance": {
+                    "config": {
+                        "base_config": {
+                            "json_file": (
+                                "cosmos_framework/model/generator/reasoner/qwen3_vl/configs/Qwen3-VL-32B-Instruct.json"
+                            )
+                        }
+                    }
+                }
             },
         }
     }
 
+    distillation_export.sanitize_student_public_model_config(model_dict)
+    first_result = copy.deepcopy(model_dict)
+    distillation_export.sanitize_student_public_model_config(model_dict)
+
+    assert model_dict == first_result
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [
+        ("vae_path", "/custom/tokenizers/CustomVAE.pth"),
+        ("json_file", "/custom/reasoners/CustomReasoner.json"),
+    ],
+)
+def test_sanitize_student_public_model_config_warns_for_unknown_absolute_paths(
+    field: str,
+    value: str,
+) -> None:
+    model_dict = {
+        "config": {
+            "tokenizer": {"vae_path": "relative/custom-vae.pth"},
+            "vlm_config": {
+                "model_instance": {"config": {"base_config": {"json_file": "relative/custom-reasoner.json"}}}
+            },
+        }
+    }
+    if field == "vae_path":
+        model_dict["config"]["tokenizer"]["vae_path"] = value
+    else:
+        model_dict["config"]["vlm_config"]["model_instance"]["config"]["base_config"]["json_file"] = value
+
+    with pytest.warns(UserWarning, match="may not be portable") as warning_records:
+        distillation_export.sanitize_student_public_model_config(model_dict)
+
+    assert warning_records[0].filename == __file__
+    if field == "vae_path":
+        assert model_dict["config"]["tokenizer"]["vae_path"] == value
+    else:
+        assert model_dict["config"]["vlm_config"]["model_instance"]["config"]["base_config"]["json_file"] == value
+
+
+def test_sanitize_student_public_model_config_preserves_unknown_relative_paths() -> None:
+    model_dict = {
+        "config": {
+            "tokenizer": {"vae_path": "relative/custom-vae.pth"},
+            "vlm_config": {
+                "model_instance": {"config": {"base_config": {"json_file": "relative/custom-reasoner.json"}}}
+            },
+        }
+    }
+
+    distillation_export.sanitize_student_public_model_config(model_dict)
+
+    assert model_dict["config"]["tokenizer"]["vae_path"] == "relative/custom-vae.pth"
+    assert (
+        model_dict["config"]["vlm_config"]["model_instance"]["config"]["base_config"]["json_file"]
+        == "relative/custom-reasoner.json"
+    )
+
+
+@pytest.mark.parametrize(
+    ("model_dict", "field_name"),
+    [
+        ({"config": {"tokenizer": {"vae_path": None}}}, "tokenizer.vae_path"),
+        (
+            {"config": {"vlm_config": {"model_instance": {"config": {"base_config": {"json_file": None}}}}}},
+            "vlm_config.model_instance.config.base_config.json_file",
+        ),
+    ],
+)
+def test_sanitize_student_public_model_config_rejects_non_string_dependency_path(
+    model_dict: dict,
+    field_name: str,
+) -> None:
+    with pytest.raises(TypeError) as error:
+        distillation_export.sanitize_student_public_model_config(model_dict)
+
+    assert str(error.value) == f"Expected {field_name} to be a string."
+
 
 def test_resolve_vision_checkpoint_path_prefers_local_override() -> None:
     fallback_called = False