Release 2026-10-26 (i4 09e71fb6)

.file_mapping.json

@@ -1,7 +1,7 @@
 {
-  "_source_commit": "b3866eb30721bf68dce73380563235d39fa5a622-dirty",
-  "_dest_commit": "bf492fe2fda2108aae9b9561dc9fb9eec3fe2027",
-  "_generated_at": "2026-10-24T15:53:46Z",
+  "_source_commit": "09e71fb667f046f9f81f3935f13cd3b2143fd576-dirty",
+  "_dest_commit": "4f6e539f98d1e7951e1c0907e61e2ff3ccdbdc92",
+  "_generated_at": "2026-10-26T05:53:07Z",
   "files": {
     "imaginaire/__init__.py": "cosmos_framework/__init__.py",
     "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py",

cosmos_framework/configs/base/defaults/model_config.py

@@ -264,9 +264,9 @@ class OmniMoTModelConfig:
     # whose noised tokens share one timestep, take the unchunked path. None disables chunking.
     timestep_embedding_chunk_size: int | None = 128_000
 
-    # Keep the raw CP data window on the host and release raw video CUDA inputs
-    # immediately after slot-0 tokenization. H100 multiview experiments opt in;
-    # the model default remains disabled.
+    # Keep the raw CP data window on the host and release raw vision CUDA inputs
+    # immediately after slot-0 tokenization. Memory-constrained CP experiments
+    # opt in; the model default remains disabled.
     release_cp_raw_batch_after_preprocessing: bool = False
 
     # torch.compile knobs (enabled, compiled_region, dynamic, ...).

cosmos_framework/model/generator/mot/context_parallel_test.py

@@ -141,12 +141,13 @@ def test_context_parallel_fetch_is_release_unaware(
 
 @pytest.mark.L0
 @pytest.mark.CPU
-def test_next_step_batch_routes_release_mode_batches(monkeypatch: pytest.MonkeyPatch) -> None:
+@pytest.mark.parametrize("media_key", ["video", "images"])
+def test_next_step_batch_routes_release_mode_batches(monkeypatch: pytest.MonkeyPatch, media_key: str) -> None:
     """StepBatch encapsulates host retention, CUDA transfer, and callback merging."""
     trainer, model = _make_window_trainer(cp_size=2)
     model.config = SimpleNamespace(release_cp_raw_batch_after_preprocessing=True)
-    raw_video = torch.zeros(1, 3, 2, 4, 4, dtype=torch.uint8)
-    host_batch = {"video": raw_video, "dataset_name": "host"}
+    raw_vision = torch.zeros(1, 3, 2, 4, 4, dtype=torch.uint8)  # [B,C,T,H,W]
+    host_batch = {media_key: raw_vision, "dataset_name": "host"}
     dataloader_iter = _CountingIterator([host_batch])
     monkeypatch.setattr(dist, "get_backend", lambda group: dist.Backend.GLOO)
     monkeypatch.setattr(dist, "all_reduce", lambda tensor, op, group: None)
@@ -174,7 +175,7 @@ def test_next_step_batch_routes_release_mode_batches(monkeypatch: pytest.MonkeyP
     assert not step1._is_fetch_slot
     assert step1.training is prepared
     cb = step1.callback_post_step
-    assert cb["video"] is raw_video
+    assert cb[media_key] is raw_vision
     assert cb["packed_tokens"] is prepared["packed_tokens"]
     assert cb["dataset_name"] == "prepared"
     assert dataloader_iter.fetch_count == 1

cosmos_framework/model/generator/omni_mot_model.py

@@ -1409,29 +1409,27 @@ def _unpack_training_payload(
         vae_pixel_shapes = payload["vae_pixel_shapes"]
         return input_text_indexes, sequence_plans, gen_data_clean, memory_info, data_resolutions, vae_pixel_shapes
 
-    def _release_raw_video_after_tokenization(
+    def _release_raw_vision_after_tokenization(
         self,
         gen_data_clean: GenerationDataClean,
         data_batch: dict[str, torch.Tensor],
     ) -> None:
-        """Free raw video tensors from GPU after CP slot-0 tokenization.
+        """Free raw image or video tensors from GPU after CP slot-0 tokenization.
 
-        Only video (the largest modality) is dropped, and only for non-image
-        batches.  The trainer retains the original host batch for callbacks,
-        which re-encode from its uint8 video.
+        The trainer retains the original host batch for callbacks, which
+        re-encode from its image or video tensors.
 
         ``num_vision_items_per_sample`` and ``image_size`` are kept: both
         describe the flattened vision-item layout and must stay mutually
-        consistent.  :meth:`build_callback_batch` restores host video and
-        flattens it to match the retained counts so callbacks see a
-        self-consistent batch without recomputation.
+        consistent.  :meth:`build_callback_batch` restores host vision and
+        flattens it to match the retained counts so callbacks see a self-consistent
+        batch without recomputation.
         """
         if not self.config.release_cp_raw_batch_after_preprocessing:
             return
-        if gen_data_clean.is_image_batch:
-            return
         gen_data_clean.raw_state_vision = None
         data_batch.pop(self.input_video_key, None)
+        data_batch.pop(self.input_image_key, None)
 
     def _get_training_inputs(
         self, data_batch: dict[str, torch.Tensor], iteration: int
@@ -1476,7 +1474,7 @@ def _get_training_inputs(
         if cp_window_slot == 0:
             local_training_data = self._prepare_training_data(data_batch, iteration)
             self._cp_local_training_payload = self._pack_training_payload(*local_training_data)
-            self._release_raw_video_after_tokenization(local_training_data[2], data_batch)
+            self._release_raw_vision_after_tokenization(local_training_data[2], data_batch)
         if self._cp_local_training_payload is None:
             raise RuntimeError("CP training payload cache is empty before its data window is complete.")
         payload = broadcast_context_parallel_object(