diff --git a/cosmos_predict2/_src/predict2/callbacks/every_n_draw_sample.py b/cosmos_predict2/_src/predict2/callbacks/every_n_draw_sample.py index e124d34f7..764420385 100644 --- a/cosmos_predict2/_src/predict2/callbacks/every_n_draw_sample.py +++ b/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) diff --git a/cosmos_predict2/_src/predict2/callbacks/validation_draw_sample.py b/cosmos_predict2/_src/predict2/callbacks/validation_draw_sample.py index 71fd18fa6..7560b9fff 100644 --- a/cosmos_predict2/_src/predict2/callbacks/validation_draw_sample.py +++ b/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) diff --git a/cosmos_predict2/_src/predict2/datasets/local_datasets/dataset_video.py b/cosmos_predict2/_src/predict2/datasets/local_datasets/dataset_video.py index bf0d6d9f7..f0c6e71dc 100644 --- a/cosmos_predict2/_src/predict2/datasets/local_datasets/dataset_video.py +++ b/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 diff --git a/cosmos_predict2/_src/reason1/tokenizer/processor.py b/cosmos_predict2/_src/reason1/tokenizer/processor.py index 459cb2a5f..ad69266c5 100644 --- a/cosmos_predict2/_src/reason1/tokenizer/processor.py +++ b/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")