[WIP takeover] Preserve three-commit local runtime patch history

cosmos_predict2/_src/predict2/callbacks/every_n_draw_sample.py

@@ -328,6 +328,15 @@ def sample(self, trainer, model, data_batch, output_batch, loss, iteration):
         return None
 
     def run_save(self, to_show, batch_size, base_fp_wo_ext) -> Optional[str]:
+        frame_counts = [item.shape[2] for item in to_show]
+        min_frame_count = min(frame_counts)
+        if any(frame_count != min_frame_count for frame_count in frame_counts):
+            log.warning(
+                f"EveryNDrawSample received mixed temporal lengths {frame_counts}; "
+                f"cropping visualization tensors to {min_frame_count} frames.",
+                rank0_only=False,
+            )
+            to_show = [item[:, :, :min_frame_count] for item in to_show]
         to_show = (1.0 + torch.stack(to_show, dim=0).clamp(-1, 1)) / 2.0  # [n, b, c, t, h, w]
         is_single_frame = to_show.shape[3] == 1
         n_viz_sample = min(self.n_viz_sample, batch_size)

cosmos_predict2/_src/predict2/callbacks/validation_draw_sample.py

@@ -358,6 +358,15 @@ def sample(self, model, data_batch, iteration, should_save):
         return None
 
     def run_save(self, to_show, batch_size, iteration, save_name) -> Optional[str]:
+        frame_counts = [item.shape[2] for item in to_show]
+        min_frame_count = min(frame_counts)
+        if any(frame_count != min_frame_count for frame_count in frame_counts):
+            log.warning(
+                f"ValidationDrawSample received mixed temporal lengths {frame_counts}; "
+                f"cropping visualization tensors to {min_frame_count} frames.",
+                rank0_only=False,
+            )
+            to_show = [item[:, :, :min_frame_count] for item in to_show]
         to_show = (1.0 + torch.stack(to_show, dim=0).clamp(-1, 1)) / 2.0  # [n, b, c, t, h, w]
         is_single_frame = to_show.shape[3] == 1
         n_viz_sample = min(self.n_viz_sample, batch_size)

cosmos_predict2/_src/predict2/datasets/local_datasets/dataset_video.py

@@ -15,8 +15,10 @@
 
 """Generic video dataset loader for Cosmos Predict2."""
 
+import hashlib
 import json
 import os
+import pickle
 import random
 import traceback
 from pathlib import Path
@@ -43,6 +45,7 @@ def __init__(
         prompt_type: str | None = None,  # "long", "short", "medium", or None for auto
         caption_format: str = "auto",  # "text", "json", or "auto"
         video_paths: Optional[list[str]] = None,
+        text_embeddings_path: Optional[str] = None,
     ) -> None:
         """Dataset class for loading image-text-to-video generation data.
 
@@ -65,6 +68,8 @@ def __init__(
         self.sequence_length = num_frames
         self.prompt_type = prompt_type
         self.caption_format = caption_format
+        self.text_embeddings_path = text_embeddings_path
+        self.text_embeddings = self._load_text_embeddings(text_embeddings_path)
 
         # Determine caption format and directory
         self._setup_caption_format()
@@ -81,6 +86,39 @@ def __init__(
         self.num_failed_loads = 0
         self.preprocess = T.Compose([ToTensorVideo(), ResizePreprocess((video_size[0], video_size[1]))])
 
+    def _load_text_embeddings(self, text_embeddings_path: Optional[str]) -> Optional[dict[str, torch.Tensor]]:
+        if not text_embeddings_path:
+            return None
+        with open(text_embeddings_path, "rb") as handle:
+            cache = pickle.load(handle)
+        if isinstance(cache, dict) and "embeddings" in cache:
+            cache = cache["embeddings"]
+        if not isinstance(cache, dict):
+            raise ValueError(f"text embedding cache must be a dict, got {type(cache)} from {text_embeddings_path}")
+        log.info(f"Loaded {len(cache)} cached text embeddings from {text_embeddings_path}")
+        return cache
+
+    @staticmethod
+    def _caption_hash(caption: str) -> str:
+        return hashlib.sha256(caption.encode("utf-8")).hexdigest()
+
+    def _get_text_embedding(self, caption: str) -> torch.Tensor:
+        if self.text_embeddings is None:
+            raise RuntimeError("text embedding cache was not initialized")
+        caption_hash = self._caption_hash(caption)
+        embedding = self.text_embeddings.get(caption_hash)
+        if embedding is None:
+            embedding = self.text_embeddings.get(caption)
+        if embedding is None:
+            raise KeyError(f"caption hash {caption_hash} missing from {self.text_embeddings_path}: {caption!r}")
+        if not isinstance(embedding, torch.Tensor):
+            embedding = torch.as_tensor(embedding)
+        if embedding.dim() == 3 and embedding.shape[0] == 1:
+            embedding = embedding.squeeze(0)
+        if embedding.dim() != 2:
+            raise ValueError(f"cached text embedding for hash {caption_hash} must be 2D, got {tuple(embedding.shape)}")
+        return embedding.to(dtype=torch.bfloat16).contiguous()
+
     def __str__(self) -> str:
         return f"{len(self.video_paths)} samples from {self.dataset_dir}"
 
@@ -98,7 +136,7 @@ def _load_video(self, video_path: str) -> tuple[np.ndarray, float]:
 
         # randomly sample a sequence of frames
         max_start_idx = total_frames - self.sequence_length
-        start_frame = np.random.randint(0, max_start_idx)
+        start_frame = np.random.randint(0, max_start_idx + 1)
         end_frame = start_frame + self.sequence_length
         frame_ids = np.arange(start_frame, end_frame).tolist()
 
@@ -209,6 +247,8 @@ def __getitem__(self, index: int) -> dict | Any:
 
             data["video"] = video
             data["ai_caption"] = caption
+            if self.text_embeddings is not None:
+                data["t5_text_embeddings"] = self._get_text_embedding(caption)
 
             _, _, h, w = video.shape
 

cosmos_predict2/_src/reason1/tokenizer/processor.py

@@ -66,10 +66,11 @@ def __init__(self, name="Qwen/Qwen2.5-VL-3B-Instruct", cache_dir=None):
         else:
             self.is_vision_tokenizer = True
 
-        s3_uri = f"s3://bucket/cosmos_reasoning1/pretrained/Qwen_tokenizer/{name}/"
-        from cosmos_predict2._src.imaginaire.utils.checkpoint_db import get_checkpoint_path
+        if cache_dir is None:
+            s3_uri = f"s3://bucket/cosmos_reasoning1/pretrained/Qwen_tokenizer/{name}/"
+            from cosmos_predict2._src.imaginaire.utils.checkpoint_db import get_checkpoint_path
 
-        cache_dir = get_checkpoint_path(s3_uri)
+            cache_dir = get_checkpoint_path(s3_uri)
 
         self.processor = AutoProcessor.from_pretrained(cache_dir)
         log.info("Successfully loaded processor from local cache")