"""Caption-embedding store for the GAIA `adverse_weather` subset. The `cv_unified_adverse_weather` videos carry only the caption *text* in each scene's `metadata.json` (`language.prompt`). The matching precomputed UMT5-XXL embedding lives under `captions_t5_all//.pt`. The text -> hash linkage is recovered by joining, on `(scene_id, frame_idx)`: - per-platform `captions.json` : scene_id -> frames[{frame_idx, caption}] (the text) - `captions_t5_all/index.pt` : {scene_to_hash: scene_id -> {frame_idx: hash}, negative: {embeddings, masks}, max_length} so: caption text -> hash -> `/.pt` = {emb (1,512,4096) bf16, mask (1,512) i64}. This module builds (and caches) the `text -> hash` map, then lazily loads the per-hash embedding files at training time. Returns embeddings in the same keys the Wan pipeline's `a4d_vae_wan` expects for the `lat_language` modality (`t5_text_embeddings`, `t5_text_mask`, plus the negative/uncond pair for CFG). Default data layout (see `download` to /home/ubuntu/data): /captions_t5_all/index.pt /captions_t5_all//.pt /adverse_weather_meta/captions_json/.json (the text side) /adverse_weather_meta/index.pt (copy of the above index) /adverse_weather_meta/caption_text2hash.json (built cache) """ import glob import json import os import torch DEFAULT_DATA_ROOT = os.environ.get('ADVERSE_WEATHER_DATA_ROOT', '/home/ubuntu/data') def _meta_dir(data_root): return os.path.join(data_root, 'adverse_weather_meta') def build_text2hash(data_root=DEFAULT_DATA_ROOT, captions_json_dir=None, index_path=None, out_path=None, force=False): """Build (and cache) the {caption_text: hash} map. :param captions_json_dir: dir with per-platform `.json` caption files. :param index_path: path to `index.pt` ({scene_to_hash, negative, max_length}). :param out_path: where to write the cached JSON (defaults under adverse_weather_meta/). :return: dict {caption_text: hash}. """ meta = _meta_dir(data_root) captions_json_dir = captions_json_dir or os.path.join(meta, 'captions_json') index_path = index_path or os.path.join(meta, 'index.pt') out_path = out_path or os.path.join(meta, 'caption_text2hash.json') if os.path.exists(out_path) and not force: with open(out_path) as f: return json.load(f) if not os.path.exists(index_path): raise FileNotFoundError(f'index.pt not found at {index_path}') scene_to_hash = torch.load(index_path, map_location='cpu', weights_only=False)['scene_to_hash'] text2hash = {} json_files = sorted(glob.glob(os.path.join(captions_json_dir, '*.json'))) if not json_files: raise FileNotFoundError(f'No per-platform captions.json under {captions_json_dir}') for jf in json_files: with open(jf) as f: scenes = json.load(f) for scene_id, scene in scenes.items(): frame_map = scene_to_hash.get(scene_id, {}) for fr in scene.get('frames', []): # JSON object keys are strings; scene_to_hash keys are ints — try both. h = frame_map.get(fr['frame_idx'], frame_map.get(str(fr['frame_idx']))) if h is not None: text2hash[fr['caption']] = h os.makedirs(os.path.dirname(out_path), exist_ok=True) with open(out_path, 'w') as f: json.dump(text2hash, f) return text2hash class AdverseWeatherCaptions: """Lazily maps a caption string to its precomputed UMT5-XXL embedding. Usage: store = AdverseWeatherCaptions() # uses /home/ubuntu/data emb = store.get("A sunny, clear day ...") # -> {'t5_text_embeddings', 't5_text_mask', ...} or None neg = store.negative() # cached uncond embedding for CFG """ def __init__(self, data_root=DEFAULT_DATA_ROOT, t5_subdir='captions_t5_all', device='cpu', dtype=torch.bfloat16, cache_embeddings=True): self.data_root = data_root self.t5_dir = os.path.join(data_root, t5_subdir) self.device = device self.dtype = dtype self.cache_embeddings = cache_embeddings self.text2hash = build_text2hash(data_root) self._emb_cache = {} self._negative = None # ------------------------------------------------------------------ paths def _hash_path(self, h): return os.path.join(self.t5_dir, h[:2], f'{h}.pt') def has(self, caption_text): return caption_text in self.text2hash # ------------------------------------------------------------------ load def _load_pt(self, h): path = self._hash_path(h) if not os.path.exists(path): return None d = torch.load(path, map_location='cpu', weights_only=False) emb = d['emb'].to(device=self.device, dtype=self.dtype) # (1, 512, 4096) mask = d['mask'].to(device=self.device) # (1, 512) i64 return emb, mask def get(self, caption_text): """Return {'t5_text_embeddings','t5_text_mask','t5_text_embeddings_neg','t5_text_mask_neg'} for a caption string, or None if the text/embedding is unavailable.""" h = self.text2hash.get(caption_text) if h is None: return None if self.cache_embeddings and h in self._emb_cache: emb, mask = self._emb_cache[h] else: loaded = self._load_pt(h) if loaded is None: return None emb, mask = loaded if self.cache_embeddings: self._emb_cache[h] = (emb, mask) neg = self.negative() return { 't5_text_embeddings': emb, 't5_text_mask': mask, 't5_text_embeddings_neg': neg['t5_text_embeddings'], 't5_text_mask_neg': neg['t5_text_mask'], } def negative(self): """Cached negative (uncond) embedding from index.pt, for CFG / text dropout.""" if self._negative is None: idx = torch.load(os.path.join(_meta_dir(self.data_root), 'index.pt'), map_location='cpu', weights_only=False) neg = idx['negative'] emb = neg['embeddings'] mask = neg['masks'] if emb.ndim == 2: emb = emb.unsqueeze(0) if mask.ndim == 1: mask = mask.unsqueeze(0) self._negative = { 't5_text_embeddings': emb.to(device=self.device, dtype=self.dtype), 't5_text_mask': mask.to(device=self.device), } return self._negative