import os, re, random, warnings import torch import numpy as np import pandas as pd from pathlib import Path _REPO_DIR = str(Path(__file__).resolve().parents[2] / "asset") _BLANK_PROMPT_EMB_FILENAME = "blank_string_prompt_emb.pth" def _default_blank_prompt_emb_path() -> str: return os.path.join(_REPO_DIR, _BLANK_PROMPT_EMB_FILENAME) def _parse_int_csv(value, *, default=None, name="value"): if value is None or value == "": return default if isinstance(value, (list, tuple)): return [int(v) for v in value] return [int(v.strip()) for v in str(value).split(",") if v.strip()] def _prompt_emb_to_cpu(prompt_emb: dict) -> dict: if not isinstance(prompt_emb, dict): raise TypeError(f"prompt emb should be dict, got {type(prompt_emb)}") return { k: v.detach().cpu() if torch.is_tensor(v) else v for k, v in prompt_emb.items() } def _save_prompt_emb_atomic(prompt_emb: dict, path: str) -> None: os.makedirs(os.path.dirname(path) or ".", exist_ok=True) rank = os.environ.get("RANK", "0") tmp_path = f"{path}.tmp.rank{rank}.pid{os.getpid()}" try: torch.save(prompt_emb, tmp_path) os.replace(tmp_path, path) finally: if os.path.exists(tmp_path): os.remove(tmp_path) def _resolve_path(base_path, p): """CSV 路径可为绝对路径；否则相对 dataset_path。""" p = str(p) return p if os.path.isabs(p) else os.path.join(base_path, p) def _camera12_from_absolute_c2w( c2w_abs: torch.Tensor, pos_scale: float ) -> torch.Tensor: """Convert absolute cameras to the model's shared-relative representation. Args: c2w_abs: Float tensor ``[P,T,4,4]`` in the shared CS world frame. pos_scale: Positive translation normalization in metres. Returns: Float32 tensor ``[P,T,12]``, anchored at player 0/frame 0. """ if c2w_abs.ndim != 4 or tuple(c2w_abs.shape[-2:]) != (4, 4): raise ValueError(f"absolute c2w 应为 [P,T,4,4]，got {tuple(c2w_abs.shape)}") if c2w_abs.shape[0] < 1 or c2w_abs.shape[1] < 1: raise ValueError(f"absolute c2w 的 P/T 必须为正，got {tuple(c2w_abs.shape)}") if not torch.isfinite(c2w_abs).all(): raise ValueError("absolute c2w 包含 NaN/Inf") if not np.isfinite(pos_scale) or float(pos_scale) <= 0: raise ValueError(f"camera pos_scale 必须为正，got {pos_scale}") c2w = c2w_abs.to(dtype=torch.float64) expected_bottom = torch.tensor( [0.0, 0.0, 0.0, 1.0], dtype=c2w.dtype, device=c2w.device ) if not torch.allclose(c2w[..., 3, :], expected_bottom, atol=1e-09, rtol=0.0): raise ValueError("absolute c2w 齐次矩阵最后一行不是 [0,0,0,1]") rotations = c2w[..., :3, :3] eye3 = torch.eye(3, dtype=c2w.dtype, device=c2w.device) if not torch.allclose( rotations.transpose(-1, -2) @ rotations, eye3, atol=2e-07, rtol=0.0 ): raise ValueError("absolute c2w rotation 不是正交矩阵") det = torch.linalg.det(rotations) if not torch.allclose(det, torch.ones_like(det), atol=2e-07, rtol=0.0): raise ValueError("absolute c2w rotation determinant 不是 +1") anchor = c2w[0, 0] anchor_w2c = torch.eye(4, dtype=c2w.dtype, device=c2w.device) anchor_w2c[:3, :3] = anchor[:3, :3].transpose(0, 1) anchor_w2c[:3, 3] = -(anchor[:3, :3].transpose(0, 1) @ anchor[:3, 3]) c2w_rel = anchor_w2c.view(1, 1, 4, 4) @ c2w c2w_rel[..., :3, 3] /= float(pos_scale) return c2w_rel[..., :3, :].reshape(c2w.shape[0], c2w.shape[1], 12).float() def _ensure_prompt_batch_dim(t: torch.Tensor, *, path: str, key: str) -> torch.Tensor: if not torch.is_tensor(t): raise TypeError(f"prompt emb {key!r} 不是 tensor: {path}") if t.ndim == 2: return t.unsqueeze(0) if t.ndim >= 3: return t raise ValueError(f"prompt emb {key!r} 形状异常 {tuple(t.shape)}: {path}") def _load_prompt_emb_many(paths: list[str]) -> dict: """Load per-player prompt embeddings and stack on P. Each file is expected to contain one player's context, e.g. [1,S,4096]. Return shape per key: [P,S,4096] before DataLoader adds B -> [B,P,S,4096]. """ embs = [torch.load(p, weights_only=True, map_location="cpu") for p in paths] if any((not isinstance(e, dict) for e in embs)): raise TypeError(f"prompt emb 文件应为 dict: {paths}") keys = list(embs[0].keys()) out = {} for k in keys: vals = [] for path, emb in zip(paths, embs): if k not in emb: raise KeyError(f"prompt emb key {k!r} 不存在于 {path}") v = _ensure_prompt_batch_dim(emb[k], path=path, key=k) if v.shape[0] != 1: raise ValueError( f"prompt_emb_path* 应各自只含 1 条 player emb，但 key {k!r} 在 {path} 的形状为 {tuple(v.shape)}" ) if vals and tuple(v.shape[1:]) != tuple(vals[0].shape[1:]): raise ValueError( f"prompt emb key {k!r} 除 player 维外形状不一致: {tuple(v.shape)} vs {tuple(vals[0].shape)} ({path})" ) vals.append(v) out[k] = torch.cat(vals, dim=0) return out def consecutive_precomputed_start_rows( metadata, consecutive_rows: int, num_players: int ) -> list[int]: """Return CSV row starts whose following rows are the same continuous clip.""" consecutive_rows = int(consecutive_rows) if consecutive_rows < 1: raise ValueError(f"consecutive_rows must be >= 1, got {consecutive_rows}") n_rows = len(metadata) if consecutive_rows == 1: return list(range(n_rows)) required = ["start_frame"] for pidx in range(1, int(num_players) + 1): required.extend((f"video_path{pidx}", f"video{pidx}_frame")) missing = [column for column in required if column not in metadata.columns] if missing: raise ValueError( f"consecutive precomputed rows require CSV columns {missing} to verify continuity" ) valid = [] for start in range(0, n_rows - consecutive_rows + 1): ok = True for row in range(start, start + consecutive_rows - 1): lengths = [ int(metadata.iloc[row][f"video{pidx}_frame"]) for pidx in range(1, num_players + 1) ] if min(lengths) <= 0 or len(set(lengths)) != 1: ok = False break if ( int(metadata.iloc[row + 1]["start_frame"]) != int(metadata.iloc[row]["start_frame"]) + lengths[0] ): ok = False break for pidx in range(1, num_players + 1): current_dir = os.path.dirname( str(metadata.iloc[row][f"video_path{pidx}"]).replace("\\", "/") ) next_dir = os.path.dirname( str(metadata.iloc[row + 1][f"video_path{pidx}"]).replace("\\", "/") ) if current_dir != next_dir: ok = False break if not ok: break if ok: valid.append(start) return valid class PrecomputedLatentDataset(torch.utils.data.Dataset): """直接从预处理好的 CSV 读取压好的 VAE latent + action（跳过在线读视频 + VAE 编码）。 CSV 列（路径相对 base_path=preprocess root；绝对路径则原样用）： latent_path1..latent_pathP（或 vae_path1..P）, action_path1..action_pathP camera pose（仅 --use_camera_pose 时读取）： - 只读预计算好的 camera_pose_path/camera12_path，或 per-player camera_pose_path1..P； - per-player 文件可为旧 [T,12]，也可为 canonical absolute c2w [T,4,4]；后者在 P 路全部加载后，以 compact player 0 / frame 0 统一计算 shared-relative camera12； - 训练时不再从 json_file1/2 现场生成 camera12，避免 dataloader 反复读 action JSON。 每段（CSV 一行）对应的物理文件： *.latent.pt : torch.save 的张量 [C=48, T_lat, h, w]（bf16） *.action.pt : torch.save({"keyboard": float[T,23], "mouse": float[T,2]})（mouse 已 compress_mouse_linear） 返回的 batch key 与 TensorDataset 默认路径下游一致： latents[P,48,T_lat,h,w]、keyboard[P,T,23]、mouse[P,T,2]、prompt_emb[P,S,4096]。 为兼容旧 2-player 路径，当 P=2 时也会额外返回 latents1/latents2。 latent 与 action 由同一 VAE 和同一切片离线算出。 """ def __init__( self, base_path, metadata_path, steps_per_epoch=2000, num_players=2, target_rgb_frames=165, blank_prompt_emb_path=None, vae_latent_channels=48, debug_no_retry=False, num_frame_per_block=3, local_attn_size=24, use_camera_pose=False, use_group_labels=False, num_groups=1, group_path_column="group_path", group_label_player_ids=None, hier_group_mode="csv", hier_group_seed=0, hier_random_preserve_sizes=True, hier_group_format="labels", consecutive_rows=1, grouping_sample_indices=None, grouping_cache_dir="outputs/grouping_cache", grouping_model_profile=(5, 6, 1.0, 3), allow_first_frame_latents=False, ): metadata = pd.read_csv(metadata_path) self.num_players = int(num_players) if self.num_players < 2: raise ValueError(f"num_players must be >= 2, got {self.num_players}") latent_cols = [] action_cols = [] missing = [] for pidx in range(1, self.num_players + 1): latent_col = None for candidate in (f"latent_path{pidx}", f"vae_path{pidx}"): if candidate in metadata.columns: latent_col = candidate break if latent_col is None: missing.append(f"latent_path{pidx}/vae_path{pidx}") latent_cols.append(latent_col) action_col = f"action_path{pidx}" if action_col not in metadata.columns: missing.append(action_col) action_cols.append(action_col) if missing: raise ValueError( f"[PrecomputedLatentDataset] CSV 缺少预处理列 {missing}；该 CSV 似乎不是 *.preprocessed.csv。现有列：{list(metadata.columns)}" ) self.metadata = metadata self.base_path = base_path self.latent_column_names = latent_cols self.action_column_names = action_cols self.latent_paths = [ [_resolve_path(base_path, x) for x in metadata[col]] for col in self.latent_column_names ] self.action_paths = [ [_resolve_path(base_path, x) for x in metadata[col]] for col in self.action_column_names ] self.latent_path1 = self.latent_paths[0] self.latent_path2 = self.latent_paths[1] self.action_path1 = self.action_paths[0] self.action_path2 = self.action_paths[1] self.use_camera_pose = bool(use_camera_pose) self.camera_pose_path = None self.camera_pose_paths = None combined_camera_col = next( (c for c in ("camera_pose_path", "camera12_path") if c in metadata.columns), None, ) pair_camera_cols = next( ( prefix for prefix in ("camera_pose_path", "camera12_path") if all( ( f"{prefix}{pidx}" in metadata.columns for pidx in range(1, self.num_players + 1) ) ) ), None, ) if combined_camera_col is not None and pair_camera_cols is not None: raise ValueError( "CSV camera pose 列不能同时使用 combined 和 per-player 两种格式" ) for prefix in ("camera_pose_path", "camera12_path"): present = [ f"{prefix}{pidx}" in metadata.columns for pidx in range(1, self.num_players + 1) ] if any(present) and (not all(present)): raise ValueError( f"CSV 若使用 per-player camera pose，必须同时包含 {prefix}1..{prefix}{self.num_players}" ) if combined_camera_col is not None: self.camera_pose_path = [ _resolve_path(base_path, x) for x in metadata[combined_camera_col] ] elif pair_camera_cols is not None: self.camera_pose_paths = [ [ _resolve_path(base_path, x) for x in metadata[f"{pair_camera_cols}{pidx}"] ] for pidx in range(1, self.num_players + 1) ] if ( self.use_camera_pose and self.camera_pose_path is None and (self.camera_pose_paths is None) ): raise ValueError( f"--use_camera_pose requires camera_pose_path/camera12_path, or per-player camera_pose_path1..{self.num_players} / camera12_path1..{self.num_players}. Camera poses must already be precomputed." ) self.prompt_emb_paths = None prompt_cols = [ f"prompt_emb_path{pidx}" for pidx in range(1, self.num_players + 1) ] prompt_present = [c in metadata.columns for c in prompt_cols] if any(prompt_present) and (not all(prompt_present)): raise ValueError( f"CSV 若使用 prompt emb，必须同时包含 prompt_emb_path1..prompt_emb_path{self.num_players}" ) if all(prompt_present): self.prompt_emb_paths = [ [_resolve_path(base_path, x) for x in metadata[col]] for col in prompt_cols ] self.prompt_emb_path1 = ( self.prompt_emb_paths[0] if self.prompt_emb_paths else None ) self.prompt_emb_path2 = ( self.prompt_emb_paths[1] if self.prompt_emb_paths else None ) self.captions = None caption_cols = [f"caption{pidx}" for pidx in range(1, self.num_players + 1)] caption_present = [c in metadata.columns for c in caption_cols] if any(caption_present) and (not all(caption_present)): raise ValueError( f"CSV 若使用 caption，必须同时包含 caption1..caption{self.num_players}" ) if all(caption_present): self.captions = [ [str(v).strip() if pd.notna(v) else "" for v in metadata[col]] for col in caption_cols ] print( f"[PrecomputedLatentDataset] 使用 caption1..caption{self.num_players} 作为 validation prompt。" ) self.caption1 = self.captions[0] if self.captions else None self.caption2 = self.captions[1] if self.captions else None self.use_group_labels = bool(use_group_labels) self.num_groups = int(num_groups) self.group_path_column = str(group_path_column) self.hier_group_mode = str(hier_group_mode) self.hier_group_seed = int(hier_group_seed) self.hier_random_preserve_sizes = bool(hier_random_preserve_sizes) self.hier_group_format = str(hier_group_format) if self.hier_group_mode not in ("csv", "random_static", "random_block"): raise ValueError( f"hier_group_mode must be csv/random_static/random_block, got {self.hier_group_mode!r}" ) if self.hier_group_format not in ("labels", "membership"): raise ValueError( f"hier_group_format must be labels or membership, got {self.hier_group_format!r}" ) if self.hier_group_format == "membership" and self.hier_group_mode != "csv": raise ValueError( f"hier_group_format=membership does not support hier_group_mode={self.hier_group_mode!r}; use csv. Random modes only support the labels format." ) self.group_label_player_ids = _parse_int_csv( group_label_player_ids, default=list(range(1, self.num_players + 1)), name="group_label_player_ids", ) if len(self.group_label_player_ids) != self.num_players: raise ValueError( f"group_label_player_ids length {len(self.group_label_player_ids)} must equal num_players={self.num_players}: {self.group_label_player_ids}" ) if any((pid < 1 or pid > 10 for pid in self.group_label_player_ids)): raise ValueError( f"group_label_player_ids must be 1-based ids in [1,10], got {self.group_label_player_ids}" ) self.group_paths = None self.source_n_rows = len(metadata) self.consecutive_rows = int(consecutive_rows) if self.consecutive_rows < 1: raise ValueError( f"consecutive_rows must be >= 1, got {self.consecutive_rows}" ) self.sample_start_rows = consecutive_precomputed_start_rows( metadata, self.consecutive_rows, self.num_players ) self.n_rows = len(self.sample_start_rows) if self.n_rows == 0: raise ValueError( f"[PrecomputedLatentDataset] no valid {self.consecutive_rows}-row continuous windows in {metadata_path}" ) self.allow_first_frame_latents = bool(allow_first_frame_latents) if self.allow_first_frame_latents and self.consecutive_rows != 1: raise ValueError("First-frame latents require consecutive_rows=1") self.steps_per_epoch = steps_per_epoch self.target_rgb_frames = int(target_rgb_frames) self.vae_latent_channels = int(vae_latent_channels) self.num_frame_per_block = int(num_frame_per_block) self.local_attn_size = int(local_attn_size) self.blank_prompt_emb_path = ( blank_prompt_emb_path or _default_blank_prompt_emb_path() ) self.blank_prompt_emb = None self.debug_no_retry = debug_no_retry if self.use_group_labels: from multics.grouping.prepare import prepare_dataset_groups prepare_dataset_groups( self, grouping_sample_indices, grouping_cache_dir, (*grouping_model_profile[:3], self.num_frame_per_block), ) print( f"[PrecomputedLatentDataset] rows={self.n_rows} target_rgb_frames={self.target_rgb_frames} latent_channels={self.vae_latent_channels} use_camera_pose={self.use_camera_pose} camera_source={('path' if self.camera_pose_path or self.camera_pose_paths else 'none')} prompt_emb={self.prompt_emb_paths is not None} num_players={self.num_players} use_group_labels={self.use_group_labels} group_label_player_ids={(self.group_label_player_ids if self.use_group_labels else 'none')} hier_group_mode={(self.hier_group_mode if self.use_group_labels else 'none')} hier_group_format={(self.hier_group_format if self.use_group_labels else 'none')} hier_group_seed={self.hier_group_seed} hier_random_preserve_sizes={self.hier_random_preserve_sizes} consecutive_rows={self.consecutive_rows} valid_windows={self.n_rows}" ) def __len__(self): return max(self.n_rows, self.steps_per_epoch) def _load_action(self, path): d = torch.load(path, weights_only=True, map_location="cpu") return (d["keyboard"].float(), d["mouse"].float()) def _load_camera_pose_file(self, path: str): ext = os.path.splitext(path)[1].lower() if ext == ".npy": obj = np.load(path) elif ext == ".npz": with np.load(path) as npz: if "c2w_abs" in npz and ( not any( ( key in npz for key in ("camera_pose", "camera12", "pose", "camera") ) ) ): obj = {key: np.array(npz[key], copy=True) for key in npz.files} else: key = next( ( k for k in ("camera_pose", "camera12", "pose", "camera") if k in npz ), npz.files[0], ) obj = np.array(npz[key], copy=True) else: obj = torch.load(path, weights_only=True, map_location="cpu") if isinstance(obj, dict): if "c2w_abs" in obj and ( not any( ( key in obj for key in ("camera_pose", "camera12", "pose", "camera") ) ) ): return obj key = next( (k for k in ("camera_pose", "camera12", "pose", "camera") if k in obj), None, ) if key is None: raise KeyError( f"camera pose 文件缺少 camera_pose/camera12/pose/camera key: {path}" ) obj = obj[key] if not torch.is_tensor(obj): obj = torch.as_tensor(obj) return obj.float() @staticmethod def _camera_payload_tensor(payload: dict, key: str, *, path: str) -> torch.Tensor: """Read one required sidecar field as a tensor. Args: payload: Loaded numeric per-video camera dictionary. key: Required dictionary field name. path: Source path included in validation errors. Returns: The field as a torch tensor without changing its dtype. """ if key not in payload: raise KeyError(f"per-video camera 文件缺少 {key!r}: {path}") value = payload[key] return value if torch.is_tensor(value) else torch.as_tensor(value) def _camera12_from_per_video_absolute( self, row: int, payloads: list[dict] ) -> torch.Tensor: """Stack one row's physical-video sidecars and choose a shared anchor. Args: row: Metadata row index used for CSV/path consistency checks. payloads: Per-video absolute-camera dictionaries in compact CSV order. Returns: Float32 ``camera12[P,T_lat,12]`` anchored at the first compact player's first latent frame. Physical ``player_key`` is checked against each corresponding ``video_pathN``. """ paths = [player_paths[row] for player_paths in self.camera_pose_paths] c2w_players = [] reference_ticks = None reference_indices = None reference_start = None reference_stride = None reference_scale = None reference_unit = None csv_start = None if "start_frame" in self.metadata.columns: csv_start = int(self.metadata.iloc[row]["start_frame"]) for player_index, (payload, path) in enumerate(zip(payloads, paths), start=1): version = int( self._camera_payload_tensor(payload, "format_version", path=path) .reshape(-1)[0] .item() ) if version not in (2, 3): raise ValueError( f"per-video absolute camera format_version 应为 2/3，got {version}: {path}" ) c2w = self._camera_payload_tensor(payload, "c2w_abs", path=path) if c2w.ndim != 3 or tuple(c2w.shape[-2:]) != (4, 4): raise ValueError( f"per-video c2w_abs 应为 [T,4,4]，got {tuple(c2w.shape)}: {path}" ) c2w_players.append(c2w) ticks = self._camera_payload_tensor( payload, "latent_ticks", path=path ).long() indices = self._camera_payload_tensor( payload, "latent_rgb_indices", path=path ).long() start = int( self._camera_payload_tensor(payload, "start_frame", path=path) .reshape(-1)[0] .item() ) stride = int( self._camera_payload_tensor(payload, "latent_stride", path=path) .reshape(-1)[0] .item() ) scale = float( self._camera_payload_tensor(payload, "pos_scale", path=path) .reshape(-1)[0] .item() ) unit = float( self._camera_payload_tensor(payload, "world_unit_to_metre", path=path) .reshape(-1)[0] .item() ) player_key = int( self._camera_payload_tensor(payload, "player_key", path=path) .reshape(-1)[0] .item() ) video_col = f"video_path{player_index}" if video_col in self.metadata.columns: video_path = str(self.metadata.iloc[row][video_col]).replace("\\", "/") match = _CAMERA_PLAYER_PART.search(video_path) if match is not None and int(match.group(1)) != player_key: raise ValueError( f"camera player_key={player_key:02d} 与 {video_col} 中的 player={int(match.group(1)):02d} 不一致: {path}" ) if csv_start is not None and start != csv_start: raise ValueError( f"camera start_frame={start} 与 CSV start_frame={csv_start} 不一致: {path}" ) if reference_ticks is None: reference_ticks = ticks reference_indices = indices reference_start = start reference_stride = stride reference_scale = scale reference_unit = unit else: if not torch.equal(ticks, reference_ticks): raise ValueError( f"不同 player 的 camera latent_ticks 不一致: {path}" ) if not torch.equal(indices, reference_indices): raise ValueError( f"不同 player 的 camera latent_rgb_indices 不一致: {path}" ) if start != reference_start or stride != reference_stride: raise ValueError( f"不同 player 的 camera start/stride 不一致: {path}" ) if scale != reference_scale or unit != reference_unit: raise ValueError( f"不同 player 的 camera unit/pos_scale 不一致: {path}" ) if len({tuple(c2w.shape) for c2w in c2w_players}) != 1: raise ValueError( f"不同 player 的 per-video c2w_abs shape 不一致: {[tuple(c2w.shape) for c2w in c2w_players]}" ) c2w_abs = torch.stack(c2w_players, dim=0) return _camera12_from_absolute_c2w(c2w_abs, float(reference_scale)) def _load_camera_pose(self, row: int, t_lat: int) -> torch.Tensor | None: if self.camera_pose_path is not None: cam = self._load_camera_pose_file(self.camera_pose_path[row]) if isinstance(cam, dict): raise ValueError( "combined camera_pose_path 不能指向 per-video absolute c2w；请使用 camera_pose_path1..N" ) elif self.camera_pose_paths is not None: cams = [ self._load_camera_pose_file(paths[row]) for paths in self.camera_pose_paths ] absolute = [isinstance(cam_i, dict) for cam_i in cams] if any(absolute) and (not all(absolute)): raise ValueError( "per-player camera 文件不能混用 absolute c2w 与旧 camera12 格式" ) if all(absolute): cam = self._camera12_from_per_video_absolute(row, cams) else: if any((cam_i.dim() != 2 for cam_i in cams)): raise ValueError("旧 per-player camera pose 文件应为 [T,12]") cam = torch.stack(cams, dim=0) else: return None if cam.dim() != 3: raise ValueError( f"camera_pose 应为 [P,T,12] 或 [T,P,12]，got {tuple(cam.shape)}" ) if cam.shape[0] != self.num_players and cam.shape[1] == self.num_players: cam = cam.permute(1, 0, 2).contiguous() if cam.shape[0] != self.num_players or cam.shape[-1] != 12: raise ValueError( f"camera_pose 形状应为 [P={self.num_players},T,12]，got {tuple(cam.shape)}" ) if cam.shape[1] < t_lat: raise ValueError( f"camera_pose latent T={cam.shape[1]} 短于 latent T={t_lat}" ) return cam[:, :t_lat, :].contiguous() def _randomize_group_labels( self, group_labels: torch.Tensor, row: int ) -> torch.Tensor: """Apply a deterministic random grouping while preserving [T,P].""" if self.hier_group_mode == "csv": return group_labels t_lat, num_players = group_labels.shape n_block = max(1, int(self.num_frame_per_block)) generator = torch.Generator() seed = (self.hier_group_seed + int(row) * 1000003) % ((1 << 63) - 1) generator.manual_seed(seed) if self.hier_random_preserve_sizes: if self.hier_group_mode == "random_static": perm = torch.randperm(num_players, generator=generator) return group_labels.index_select(1, perm).contiguous() randomized = torch.empty_like(group_labels) for start in range(0, t_lat, n_block): end = min(start + n_block, t_lat) perm = torch.randperm(num_players, generator=generator) assignment = group_labels[start].index_select(0, perm) randomized[start:end] = assignment return randomized randomized = torch.empty_like(group_labels) if self.hier_group_mode == "random_static": assignment = torch.randint( self.num_groups, (num_players,), generator=generator ) randomized[:] = assignment else: for start in range(0, t_lat, n_block): end = min(start + n_block, t_lat) assignment = torch.randint( self.num_groups, (num_players,), generator=generator ) randomized[start:end] = assignment return randomized def _select_group_player_columns( self, tensor: torch.Tensor, player_dim: int, path: str, name: str ) -> torch.Tensor: """Select compact player columns using 1-based group_label_player_ids. tensor player axis may be the full 10-player roster or already-compact P. """ cols = [pid - 1 for pid in self.group_label_player_ids] n_src = int(tensor.shape[player_dim]) max_id = max(self.group_label_player_ids) if n_src < max_id: already_compact = ( n_src == self.num_players and self.group_label_player_ids == list(range(1, self.num_players + 1)) ) if not already_compact: raise ValueError( f"{name} player dim {n_src} cannot cover ids {self.group_label_player_ids} at {path}" ) return tensor return tensor.index_select(player_dim, torch.tensor(cols, dtype=torch.long)) @staticmethod def _assert_block_constant( tensor: torch.Tensor, n_block: int, path: str, name: str ) -> None: """Grouping must be constant inside each latent temporal block (n=3).""" if n_block <= 1: return t_lat = int(tensor.shape[0]) for start in range(0, t_lat, n_block): block = tensor[start : start + n_block] if block.numel() and (not torch.equal(block, block[:1].expand_as(block))): raise ValueError( f"{name} must be constant inside each latent temporal block; path={path} block=[{start},{min(start + n_block, t_lat)})" ) def _load_group_labels(self, row: int, t_lat: int) -> dict: """Load grouping tensors for one CSV row. labels mode returns: group_labels [T,P] long in [0,N), active_group [T,N] bool membership mode additionally returns: membership [T,N,P] bool, global_only [T,P] bool group_labels is debug-only (argmax; -1 = not in any local group) eligible stays in the .pt and is only asserted, never returned. """ if self.group_paths is None: raise RuntimeError( "use_group_labels=True but group_paths is not initialized" ) path = self.group_paths[row] d = torch.load(path, weights_only=True, map_location="cpu") n_block = max(1, int(self.num_frame_per_block)) if self.hier_group_format == "membership": if not isinstance(d, dict) or "membership" not in d: raise ValueError( f"hier_group_format=membership requires key 'membership' in {path}; refusing to silently fall back to labels (anchors / orphans would be wrong)" ) membership = d["membership"] if not torch.is_tensor(membership): membership = torch.as_tensor(membership) if membership.dim() != 3: raise ValueError( f"membership must be [T,N,P] bool, got {tuple(membership.shape)} at {path}" ) if membership.shape[0] < t_lat: raise ValueError( f"membership T={membership.shape[0]} shorter than latent T={t_lat}: {path}" ) if int(membership.shape[1]) != self.num_groups: raise ValueError( f"membership N={int(membership.shape[1])} must equal num_groups={self.num_groups} at {path}" ) membership = membership[:t_lat].to(dtype=torch.bool) membership = self._select_group_player_columns( membership, player_dim=2, path=path, name="membership" ) if membership.shape[2] != self.num_players: raise ValueError( f"membership P={membership.shape[2]} must equal num_players={self.num_players} at {path}" ) if "global_only" not in d: raise ValueError( f"hier_group_format=membership requires key 'global_only' in {path}" ) global_only = d["global_only"] if not torch.is_tensor(global_only): global_only = torch.as_tensor(global_only) if global_only.dim() != 2: raise ValueError( f"global_only must be [T,P] bool, got {tuple(global_only.shape)} at {path}" ) if global_only.shape[0] < t_lat: raise ValueError( f"global_only T={global_only.shape[0]} shorter than latent T={t_lat}: {path}" ) global_only = global_only[:t_lat].to(dtype=torch.bool) global_only = self._select_group_player_columns( global_only, player_dim=1, path=path, name="global_only" ) if global_only.shape[1] != self.num_players: raise ValueError( f"global_only P={global_only.shape[1]} must equal num_players={self.num_players} at {path}" ) if "eligible" not in d: raise ValueError( f"hier_group_format=membership requires key 'eligible' [P] in {path}" ) eligible = d["eligible"] if not torch.is_tensor(eligible): eligible = torch.as_tensor(eligible) eligible = eligible.to(dtype=torch.bool).reshape(-1) eligible = self._select_group_player_columns( eligible.unsqueeze(0), player_dim=1, path=path, name="eligible" ).squeeze(0) if eligible.numel() != self.num_players: raise ValueError( f"eligible P={eligible.numel()} must equal num_players={self.num_players} at {path}" ) if not bool(eligible.all().item()): raise ValueError( f"eligible must be all True for this training CSV (not used in the mask); got {eligible.tolist()} at {path}" ) self._assert_block_constant(membership, n_block, path, "membership") self._assert_block_constant(global_only, n_block, path, "global_only") active_group = membership.any(dim=-1).contiguous() if "labels" in d: labels = d["labels"] if not torch.is_tensor(labels): labels = torch.as_tensor(labels) group_labels = labels[:t_lat].to(dtype=torch.long) group_labels = self._select_group_player_columns( group_labels, player_dim=1, path=path, name="labels" ).contiguous() else: any_g = membership.any(dim=1) group_labels = membership.to(dtype=torch.long).argmax(dim=1) group_labels = torch.where( any_g, group_labels, torch.full_like(group_labels, -1) ) return { "group_labels": group_labels, "active_group": active_group, "membership": membership.contiguous(), "global_only": global_only.contiguous(), } if isinstance(d, dict) and "labels" not in d: raise ValueError( f"hier_group_format=labels requires key 'labels' in {path}; this looks like a membership-only file. Pass --hier_group_format membership." ) labels = d["labels"] if isinstance(d, dict) else d if not torch.is_tensor(labels): labels = torch.as_tensor(labels) if labels.dim() != 2 or labels.shape[1] < max(self.group_label_player_ids): already_compact = ( labels.dim() == 2 and labels.shape[1] == self.num_players and ( self.group_label_player_ids == list(range(1, self.num_players + 1)) ) ) if not already_compact: raise ValueError( f"group labels should be [T,10] and cover ids {self.group_label_player_ids}, got {tuple(labels.shape)} at {path}" ) if labels.shape[0] < t_lat: raise ValueError( f"group labels T={labels.shape[0]} shorter than latent T={t_lat}: {path}" ) group_labels = labels[:t_lat].to(dtype=torch.long) group_labels = self._select_group_player_columns( group_labels, player_dim=1, path=path, name="labels" ).contiguous() if group_labels.numel() and ( int(group_labels.min().item()) < 0 or int(group_labels.max().item()) >= self.num_groups ): raise ValueError( f"group_labels values out of range [0,{self.num_groups}) at {path}: min={int(group_labels.min().item())} max={int(group_labels.max().item())}" ) self._assert_block_constant(group_labels, n_block, path, "group_labels") group_labels = self._randomize_group_labels(group_labels, row) active_group = torch.zeros(t_lat, self.num_groups, dtype=torch.bool) active_group.scatter_(1, group_labels, True) return {"group_labels": group_labels, "active_group": active_group} def _load_camera_pose_rows( self, rows: list[int], latent_lengths: list[int], target_t_lat: int ) -> torch.Tensor | None: """Load consecutive camera segments in one shared relative frame. Multi-row inference requires canonical per-player absolute c2w files. Concatenating legacy camera12 tensors would silently reset the relative anchor at every CSV boundary, so that ambiguous case is rejected. """ if len(rows) == 1: return self._load_camera_pose(rows[0], target_t_lat) if self.camera_pose_paths is None: raise ValueError( "consecutive_rows>1 with camera pose requires per-player canonical camera_pose_path1..P containing absolute c2w" ) segments = [] reference_scale = None for row, t_lat in zip(rows, latent_lengths): payloads = [ self._load_camera_pose_file(paths[row]) for paths in self.camera_pose_paths ] if not all((isinstance(payload, dict) for payload in payloads)): raise ValueError( "consecutive_rows>1 cannot concatenate legacy relative camera12; use canonical per-player absolute c2w sidecars" ) self._camera12_from_per_video_absolute(row, payloads) scales = [ float( self._camera_payload_tensor(payload, "pos_scale", path=paths[row]) .reshape(-1)[0] .item() ) for payload, paths in zip(payloads, self.camera_pose_paths) ] if len(set(scales)) != 1: raise ValueError( f"camera pos_scale differs between players at CSV row {row}: {scales}" ) if reference_scale is None: reference_scale = scales[0] elif scales[0] != reference_scale: raise ValueError( f"camera pos_scale differs across consecutive rows: {reference_scale} vs {scales[0]}" ) c2w = torch.stack( [ self._camera_payload_tensor(payload, "c2w_abs", path=paths[row]) for payload, paths in zip(payloads, self.camera_pose_paths) ], dim=0, ) if c2w.shape[1] < t_lat: raise ValueError( f"camera latent T={c2w.shape[1]} shorter than row latent T={t_lat} at CSV row {row}" ) segments.append(c2w[:, :t_lat]) absolute = torch.cat(segments, dim=1) if absolute.shape[1] < target_t_lat: raise ValueError( f"combined camera latent T={absolute.shape[1]} shorter than target {target_t_lat}" ) return _camera12_from_absolute_c2w( absolute[:, :target_t_lat], float(reference_scale) ).contiguous() def _load_group_rows( self, rows: list[int], latent_lengths: list[int], target_t_lat: int ) -> dict: parts = [ self._load_group_labels(row, t_lat) for row, t_lat in zip(rows, latent_lengths) ] keys = set(parts[0]) if any((set(part) != keys for part in parts[1:])): raise ValueError("group tensor keys differ across consecutive CSV rows") result = {key: torch.cat([part[key] for part in parts], dim=0) for key in keys} for key, tensor in result.items(): if tensor.shape[0] < target_t_lat: raise ValueError( f"combined {key} T={tensor.shape[0]} shorter than target latent T={target_t_lat}" ) result[key] = tensor[:target_t_lat].contiguous() return result def __getitem__(self, index): for _ in range(20): sample_index = index % self.n_rows row = self.sample_start_rows[sample_index] if self.use_group_labels and any( r not in self.prepared_group_rows for r in range(row, row + self.consecutive_rows) ): raise ValueError(f"Grouping was not prepared for sample {sample_index}") try: rows = list(range(row, row + self.consecutive_rows)) row_latents = [ [ torch.load(paths[r], weights_only=True, map_location="cpu") for paths in self.latent_paths ] for r in rows ] row_actions = [ [self._load_action(paths[r]) for paths in self.action_paths] for r in rows ] latent_lengths = [] for segment_idx, ( source_row, segment_latents, segment_actions, ) in enumerate(zip(rows, row_latents, row_actions)): for pidx, lat in enumerate(segment_latents, start=1): if lat.dim() != 4 or lat.shape[0] != self.vae_latent_channels: raise RuntimeError( f"latent{pidx} shape {tuple(lat.shape)}, expected [{self.vae_latent_channels}, T_lat, h, w]: {self.latent_paths[pidx - 1][source_row]}" ) if tuple(lat.shape) != tuple(segment_latents[0].shape): raise RuntimeError( f"latent{pidx} shape {tuple(lat.shape)} differs from latent1 {tuple(segment_latents[0].shape)} at CSV row {source_row}" ) segment_rgb_t = int(segment_actions[0][0].shape[0]) expected_latent_t = (segment_rgb_t - 1) // 4 + 1 first_frame_only = ( self.allow_first_frame_latents and int(segment_latents[0].shape[1]) == 1 ) if ( not first_frame_only and int(segment_latents[0].shape[1]) != expected_latent_t ): raise RuntimeError( f"CSV row {source_row} latent T={segment_latents[0].shape[1]} does not match action T={segment_rgb_t} (expected {expected_latent_t})" ) for pidx, (keyboard_i, mouse_i) in enumerate( segment_actions, start=1 ): if keyboard_i.shape != (segment_rgb_t, 23) or mouse_i.shape != ( segment_rgb_t, 2, ): raise RuntimeError( f"action{pidx} shape invalid at CSV row {source_row}: keyboard={tuple(keyboard_i.shape)} mouse={tuple(mouse_i.shape)}" ) latent_lengths.append(expected_latent_t) target_t_lat = (self.target_rgb_frames - 1) // 4 + 1 latents_per_player = [ torch.cat([segment[pidx] for segment in row_latents], dim=1)[ :, :target_t_lat ] for pidx in range(self.num_players) ] keyboards = [ torch.cat([segment[pidx][0] for segment in row_actions], dim=0)[ : self.target_rgb_frames ] for pidx in range(self.num_players) ] mice = [ torch.cat([segment[pidx][1] for segment in row_actions], dim=0)[ : self.target_rgb_frames ] for pidx in range(self.num_players) ] got_T = int(keyboards[0].shape[0]) if got_T < self.target_rgb_frames: warnings.warn( f"[PrecomputedLatentDataset] action 比 target 短：got={got_T} want={self.target_rgb_frames} (row={row}, {self.action_paths[0][row]})。" ) for pidx, (kb, ms) in enumerate(zip(keyboards, mice), start=1): if kb.dim() != 2 or ms.dim() != 2: raise RuntimeError( f"action{pidx} 形状异常 keyboard={tuple(kb.shape)} mouse={tuple(ms.shape)}" ) if kb.shape[0] != got_T or ms.shape[0] != got_T: raise RuntimeError( f"action{pidx} T 与 action1 不一致: keyboard={tuple(kb.shape)} mouse={tuple(ms.shape)}, action1_T={got_T}" ) if kb.shape[-1] != 23 or ms.shape[-1] != 2: raise RuntimeError( f"action{pidx} feature dim 异常，期望 keyboard[...,23] / mouse[...,2]，got {tuple(kb.shape)} / {tuple(ms.shape)}" ) exp_T_lat = (got_T - 1) // 4 + 1 if latents_per_player[0].shape[1] != exp_T_lat and not ( self.allow_first_frame_latents and latents_per_player[0].shape[1] == 1 ): raise RuntimeError( f"latent T({latents_per_player[0].shape[1]}) 与 action T({got_T}) 不一致：期望 latent_T=(T-1)//4+1={exp_T_lat}（{self.latent_paths[0][row]}）" ) latents = torch.stack(latents_per_player, dim=0) keyboard = torch.stack(keyboards, dim=0) mouse = torch.stack(mice, dim=0) data = { "latents": latents, "keyboard": keyboard, "mouse": mouse, "source_rows": torch.tensor(rows, dtype=torch.long), } if self.num_players == 2: data["latents1"] = latents[0] data["latents2"] = latents[1] if self.use_camera_pose: camera_pose = self._load_camera_pose_rows( rows, latent_lengths, exp_T_lat ) if camera_pose is None: raise RuntimeError( "use_camera_pose=True 但 CSV 没有可读取的 camera_pose_path/camera12_path" ) data["camera_pose"] = camera_pose if self.use_group_labels: data.update( self._load_group_rows(rows, latent_lengths, exp_T_lat) ) if self.prompt_emb_paths is not None: data["prompt_emb"] = _load_prompt_emb_many( [paths[row] for paths in self.prompt_emb_paths] ) else: data["prompt_emb"] = {} if self.captions is not None: data["prompt"] = [caps[row] for caps in self.captions] return data except Exception as e: print( f"ERROR WHEN LOADING (precomputed) sample={sample_index} source_row={row}: {e}" ) if self.debug_no_retry: raise if self.use_group_labels: index = self.grouping_sample_indices[ random.randrange(len(self.grouping_sample_indices)) ] else: index = random.randrange(self.n_rows) raise RuntimeError("PrecomputedLatentDataset: 连续加载失败") _CAMERA_PLAYER_PART = re.compile("(?:^|/)player=(\\d+)(?:/|$)")