| """ |
| LatentDataset: Load precomputed ctx and target latents for Context-as-Memory dataset. |
| |
| Use after running precompute_ctx_target_latents.py. Returns samples with: |
| - context_latents: (1, C, K, H//8, W//8) - 1 latent per context frame |
| - target_latents: (1, C, T, H//8, W//8) - 1 latent per 4 target frames |
| - prompt, video_name, start_frame, end_frame, actions |
| |
| Compatible with training that uses precomputed latents instead of encoding on the fly. |
| """ |
|
|
| import json |
| import os |
| import warnings |
|
|
| import torch |
|
|
|
|
| class LatentDataset(torch.utils.data.Dataset): |
| """ |
| Dataset that loads precomputed ctx and target latents. |
| """ |
|
|
| def __init__( |
| self, |
| latent_dir, |
| metadata_path=None, |
| action_base_path=None, |
| repeat=1, |
| num_frames=81, |
| context_frames=5, |
| target_frames_per_latent=4, |
| ): |
| """ |
| Args: |
| latent_dir: Directory containing ctx_latents/ and target_latents/ subdirs. |
| metadata_path: Optional. If provided, used to get total_samples and validate. |
| action_base_path: Base path for action JSON files (for loading actions if not in .pt). |
| repeat: Dataset repeat factor. |
| num_frames: Expected num_frames per segment. |
| context_frames: Number of context frames (K). |
| target_frames_per_latent: Target: 1 latent per N frames. |
| """ |
| self.latent_dir = latent_dir |
| self.ctx_dir = os.path.join(latent_dir, "ctx_latents") |
| self.target_dir = os.path.join(latent_dir, "target_latents") |
| self.action_base_path = action_base_path or latent_dir |
| self.repeat = repeat |
| self.num_frames = num_frames |
| self.context_frames = context_frames |
| self.target_frames_per_latent = target_frames_per_latent |
|
|
| |
| self._indices = [] |
| if os.path.isdir(self.ctx_dir) and os.path.isdir(self.target_dir): |
| ctx_files = {f.replace(".pt", "") for f in os.listdir(self.ctx_dir) if f.endswith(".pt")} |
| target_files = {f.replace(".pt", "") for f in os.listdir(self.target_dir) if f.endswith(".pt")} |
| common = sorted([int(x) for x in ctx_files & target_files]) |
| self._indices = common |
| if not self._indices: |
| meta_path = os.path.join(latent_dir, "metadata_precompute.json") |
| if os.path.isfile(meta_path): |
| with open(meta_path) as f: |
| meta = json.load(f) |
| self._total = meta.get("total_samples", 0) |
| self._indices = list(range(self._total)) |
| else: |
| self._total = 0 |
| else: |
| self._total = len(self._indices) |
|
|
| def __len__(self): |
| return self._total * self.repeat |
|
|
| def __getitem__(self, idx): |
| real_idx = idx % self._total |
| if self._indices is not None: |
| real_idx = self._indices[real_idx] |
|
|
| ctx_path = os.path.join(self.ctx_dir, f"{real_idx:08d}.pt") |
| target_path = os.path.join(self.target_dir, f"{real_idx:08d}.pt") |
|
|
| if not os.path.isfile(ctx_path) or not os.path.isfile(target_path): |
| warnings.warn(f"Latent files not found for idx {real_idx}. Returning None.") |
| return None |
|
|
| ctx_data = torch.load(ctx_path, map_location="cpu", weights_only=True) |
| target_data = torch.load(target_path, map_location="cpu", weights_only=True) |
|
|
| ctx_latent = ctx_data["latent"] |
| target_latent = target_data["latent"] |
|
|
| |
| if ctx_latent.dim() == 4: |
| ctx_latent = ctx_latent.unsqueeze(0) |
| if target_latent.dim() == 4: |
| target_latent = target_latent.unsqueeze(0) |
|
|
| out = { |
| "context_latents": ctx_latent, |
| "target_latents": target_latent, |
| "prompt": ctx_data.get("prompt", ""), |
| "video_name": ctx_data.get("video_name"), |
| "start_frame": ctx_data.get("start_frame"), |
| "end_frame": ctx_data.get("end_frame"), |
| } |
| if "actions" in ctx_data and ctx_data["actions"] is not None: |
| out["actions"] = ctx_data["actions"] |
| elif "actions" in target_data and target_data["actions"] is not None: |
| out["actions"] = target_data["actions"] |
|
|
| return out |
|
|
|
|
| def get_latent_dataset_args(latent_dir, action_base_path=None, **kwargs): |
| """Build argparse.Namespace for LatentDataset from precompute metadata.""" |
| meta_path = os.path.join(latent_dir, "metadata_precompute.json") |
| if not os.path.isfile(meta_path): |
| return None |
| with open(meta_path) as f: |
| meta = json.load(f) |
| from argparse import Namespace |
| return Namespace( |
| latent_dir=latent_dir, |
| action_base_path=action_base_path or meta.get("dataset_base_path", latent_dir), |
| num_frames=meta.get("num_frames", 81), |
| context_frames=meta.get("context_frames", 5), |
| target_frames_per_latent=meta.get("target_frames_per_latent", 4), |
| **kwargs, |
| ) |
|
|
