flux2distill/data.py

"""Cached-latent dataset for the dev distillation run.

Loads packed VAE latents (N,1024,128) + captions produced by scripts/06_cache_data.py.
Text embeddings are computed online per batch (captions are cheap to re-encode and
caching the 7680-dim Qwen3 embeddings would be huge — plan.md S4)."""
from __future__ import annotations

import json

import torch
from torch.utils.data import Dataset


class LatentCaptionDataset(Dataset):
    def __init__(self, cache_dir="data/monet_cache"):
        self.latents = torch.load(f"{cache_dir}/latents.pt", map_location="cpu")  # (N,1024,128) bf16
        with open(f"{cache_dir}/captions.json") as f:
            self.captions = json.load(f)
        assert len(self.latents) == len(self.captions), "latent/caption count mismatch"

    def __len__(self):
        return len(self.latents)

    def __getitem__(self, i):
        return self.latents[i], self.captions[i]


def collate(batch):
    lats = torch.stack([b[0] for b in batch])
    caps = [b[1] for b in batch]
    return lats, caps