|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import json
|
| import zlib
|
| from pathlib import Path
|
|
|
| import numpy as np
|
| import pandas as pd
|
| import torch
|
| from safetensors import safe_open
|
| from torch.utils.data import Dataset
|
|
|
| LABELED = "latents_256_manifest_labeled.csv"
|
| MERGED = "latents_256_manifest.csv"
|
| PARTS_DIR = "latents_256_manifests"
|
| CAPTIONS = "captions.json"
|
|
|
|
|
| def source_key(df):
|
| """A string per row that identifies the source video (the split must keep every clip of a video on one side)."""
|
| if "source_file" in df.columns:
|
| return df["source_file"]
|
|
|
| return df["clip_path"].str.rsplit("/", n=1).str[-1].str.split("__").str[1]
|
|
|
|
|
| def is_val(source, val_mod):
|
| return np.fromiter((zlib.crc32(s.encode()) % val_mod == 0 for s in source), dtype=bool, count=len(source))
|
|
|
|
|
| class BoraLatents(Dataset):
|
| def __init__(self, root, split="train", mode="block", block_len=4, overlap=1, random_block=True,
|
| dtype=torch.float32, val_mod=100, manifest=None, captions_file=None, text_embeds=None,
|
| random_caption=None, primary_p=0.5):
|
| assert mode in ("clip", "block") and split in ("train", "val") and 0 <= overlap < block_len
|
| root = Path(root)
|
| self.root, self.mode, self.dtype = root, mode, dtype
|
| self.block_len, self.random_block = block_len, random_block
|
| self.random_caption = random_block if random_caption is None else random_caption
|
| self.primary_p = primary_p
|
|
|
|
|
| if manifest is not None:
|
| manifest = Path(manifest)
|
| elif (root / LABELED).exists():
|
| manifest = root / LABELED
|
| elif (root / MERGED).exists():
|
| manifest = root / MERGED
|
| if manifest is not None:
|
| df = pd.read_csv(manifest, dtype=str, keep_default_na=False)
|
| self.manifest_source = manifest.name
|
| else:
|
| parts = sorted((root / PARTS_DIR).glob("part_*.csv"))
|
| assert parts, f"no manifest found under {root}"
|
| df = pd.concat([pd.read_csv(p, dtype=str, keep_default_na=False) for p in parts], ignore_index=True)
|
| self.manifest_source = "part_*.csv"
|
| df["_g"] = df["global_idx"].astype(int)
|
| df["latent_row"] = df["latent_row"].astype(int)
|
| df = df.sort_values("_g").reset_index(drop=True)
|
| self.manifest = df
|
| self.labeled = "caption_ids" in df.columns
|
|
|
|
|
| self.captions = None
|
| if self.labeled:
|
| cf = Path(captions_file) if captions_file else (manifest.parent / CAPTIONS)
|
| assert cf.exists(), f"{cf} not found (it is written next to the labeled manifest)"
|
| self.captions = json.load(open(cf, encoding="utf-8"))
|
| assert self.captions[0] == "", "captions.json must start with the empty caption"
|
|
|
|
|
| val = is_val(source_key(df), val_mod)
|
| sub = (df[val] if split == "val" else df[~val]).reset_index(drop=True)
|
| assert len(sub), f"empty {split} split"
|
| self.n_videos = int(source_key(sub).nunique())
|
|
|
|
|
| self._handles = {}
|
| first = sub["latent_file"].iloc[0]
|
| self.latents_mean, self.latents_std = self._read_stats(first)
|
| _, self.C, self.T, self.H, self.W = self._open(first).get_slice("latents").get_shape()
|
| self.starts = list(range(0, self.T - block_len + 1, block_len - overlap))
|
|
|
| if mode == "block" and not random_block:
|
| sub = sub.loc[sub.index.repeat(len(self.starts))].reset_index(drop=True)
|
| self._bidx = [i % len(self.starts) for i in range(len(sub))]
|
| self._file = sub["latent_file"].tolist()
|
| self._row = sub["latent_row"].tolist()
|
| self._gidx = sub["_g"].tolist()
|
| self._cids = [json.loads(s) for s in sub["caption_ids"]] if self.labeled else None
|
|
|
|
|
| self.text_embeds = Path(text_embeds) if text_embeds else None
|
| self._emb = None
|
| if self.text_embeds is not None:
|
| assert self.labeled, "text_embeds needs the labeled manifest"
|
| with safe_open(str(self.text_embeds), framework="pt") as f:
|
| n, self.L, self.D = f.get_slice("emb").get_shape()
|
| assert n == len(self.captions), f"{self.text_embeds} has {n} rows, captions.json has {len(self.captions)}"
|
|
|
| def _open(self, rel):
|
| h = self._handles.get(rel)
|
| if h is None:
|
| h = self._handles[rel] = safe_open(str(self.root / rel), framework="pt")
|
| return h
|
|
|
| def _read_stats(self, rel):
|
| meta = self._open(rel).metadata()
|
| shape = (1, -1, 1, 1, 1)
|
| return (torch.tensor(json.loads(meta["latents_mean"])).view(shape),
|
| torch.tensor(json.loads(meta["latents_std"])).view(shape))
|
|
|
| def __len__(self):
|
| return len(self._file)
|
|
|
| def __getitem__(self, i):
|
| row = self._row[i]
|
| z = self._open(self._file[i]).get_slice("latents")[row:row + 1][0]
|
| item = {"global_idx": self._gidx[i]}
|
| if self.mode == "block":
|
| b = int(torch.randint(len(self.starts), (1,))) if self.random_block else self._bidx[i]
|
| s = self.starts[b]
|
| z = z[:, s:s + self.block_len]
|
| item["block_index"], item["start"] = b, s
|
| item["latents"] = z.to(self.dtype)
|
| if self._cids is not None:
|
| cids = self._cids[i]
|
| if self.random_caption and len(cids) > 1 and torch.rand(()) >= self.primary_p:
|
| cid = cids[int(torch.randint(len(cids), (1,)))]
|
| else:
|
| cid = cids[0]
|
| item["caption_id"] = cid
|
| item["text"] = self.captions[cid]
|
| if self.text_embeds is not None:
|
| if self._emb is None:
|
| self._emb = safe_open(str(self.text_embeds), framework="pt")
|
| item["text_emb"] = self._emb.get_slice("emb")[cid:cid + 1][0]
|
| item["text_mask"] = self._emb.get_slice("mask")[cid:cid + 1][0]
|
| return item
|
|
|
