from __future__ import annotations import json from dataclasses import dataclass from pathlib import Path from typing import Any, Mapping, Sequence import torch from torch.utils.data import DataLoader, Dataset from data.encodec_utils import ( PRESNAP_TARGET_LAYOUT, normalize_target_payload, resolve_target_dim_from_cache_config, resolve_target_layout_from_cache_config, ) @dataclass(frozen=True) class DiffusionExample: example_path: Path source_id: str source_manifest_index: int source_pt_rel: str source_row_in_shard: int beat_index: int beat_index_end: int split: str kit_name: str source_row_id: str conditioning_mode: str class_names: tuple[str, ...] class_id_vocab_sizes: tuple[int, ...] feature_row_names: tuple[str, ...] grid_ft: torch.Tensor grid_ids_ft: torch.Tensor family_onsets_ft: torch.Tensor family_onset_count_ft: torch.Tensor grid_times_sec_t: torch.Tensor token_times_sec_t: torch.Tensor beat_boundaries_sec_rel: torch.Tensor grid_num_frames: int grid_frame_rate: float bpm: float duration_sec: float target_sum_td: torch.Tensor target_full_sum_td: torch.Tensor target_dim: int target_full_dim: int target_layout: str pca_basis_path: str target_num_frames: int source_codes_ct: torch.Tensor target_sum_pool_d: torch.Tensor @property def target_sum_t128(self) -> torch.Tensor: return self.target_sum_td @property def target_sum_pool_128(self) -> torch.Tensor: return self.target_sum_pool_d from data.grid_rate_downsample import downsample_grid_payload, uniform_grid_times def _load_jsonl(path: Path) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] with path.open("r", encoding="utf-8") as handle: for line in handle: text = str(line).strip() if not text: continue rows.append(dict(json.loads(text))) return rows def _require_payload_keys(payload: dict[str, Any], keys: Sequence[str], *, example_path: Path) -> None: missing = [str(key) for key in list(keys) if str(key) not in payload] if missing: raise RuntimeError( f"cache example {example_path} is missing required seconds-conditioning fields {missing}. " "Rebuild the cache with scripts/build_diffusion_cache.py." ) class DiffusionConditioningDataset(Dataset[DiffusionExample]): def __init__( self, cache_root: str | Path, *, split: str = "train", max_items: int = 0, conditioning_mode: str = "seconds", grid_rate_hz: float = 0.0, grid_overlay_root: str | Path | None = None, ) -> None: super().__init__() self.cache_root = Path(cache_root).resolve() self.split = str(split).strip().lower() self.conditioning_mode = str(conditioning_mode).strip().lower() if self.conditioning_mode not in {"seconds", "seconds_frontend"}: raise ValueError( f"unsupported conditioning_mode={conditioning_mode!r}; expected 'seconds' or 'seconds_frontend'" ) self.target_dim = int(resolve_target_dim_from_cache_config(self.cache_root)) self.target_layout = str(resolve_target_layout_from_cache_config(self.cache_root)) self.grid_overlay_root = ( None if grid_overlay_root is None or not str(grid_overlay_root).strip() else Path(grid_overlay_root).expanduser().resolve() ) if self.grid_overlay_root is not None: overlay_config_path = self.grid_overlay_root / "config.json" if not overlay_config_path.is_file(): raise FileNotFoundError(f"missing grid-overlay config: {overlay_config_path}") overlay_config = json.loads(overlay_config_path.read_text(encoding="utf-8")) if str(overlay_config.get("format")) != "grid2drum-grid-overlay-v1": raise ValueError(f"unsupported grid-overlay format in {overlay_config_path}") # 0 disables decimation entirely (cached rate is used as-is). self.grid_rate_hz = float(max(0.0, float(grid_rate_hz))) manifest_path = self.cache_root / "manifests" / f"{self.split}.jsonl" self.rows = _load_jsonl(manifest_path) if int(max_items) > 0: self.rows = self.rows[: int(max_items)] if not self.rows: raise RuntimeError(f"no rows found for split={self.split!r} under {self.cache_root}") def __len__(self) -> int: return int(len(self.rows)) def __getitem__(self, index: int) -> DiffusionExample: return self.example_from_row(self.rows[int(index)]) def example_from_row(self, row: Mapping[str, Any]) -> DiffusionExample: row = dict(row) example_path = (self.cache_root / str(row["out_pt"])).resolve() payload = dict(torch.load(example_path, map_location="cpu", weights_only=False)) if self.grid_overlay_root is not None: overlay_path = (self.grid_overlay_root / str(row["out_pt"])).resolve() if not overlay_path.is_file(): raise FileNotFoundError(f"missing grid overlay for cache example: {overlay_path}") overlay = dict(torch.load(overlay_path, map_location="cpu", weights_only=False)) if str(overlay.get("source_id") or "") != str(payload.get("source_id") or ""): raise RuntimeError(f"grid-overlay source mismatch for {example_path}") for key in ( "grid_ft", "grid_ids_ft", "family_onsets_ft", "family_onset_count_ft", "grid_times_sec_t", "grid_num_frames", "grid_frame_rate", ): payload[key] = overlay[key] _require_payload_keys( payload, ( "grid_ft", "grid_ids_ft", "family_onsets_ft", "family_onset_count_ft", "grid_times_sec_t", "token_times_sec_t", "beat_boundaries_sec_rel", "grid_num_frames", "grid_frame_rate", "bpm", "duration_sec", "target_num_frames", "source_codes_ct", ), example_path=example_path, ) target_payload = payload if self.target_layout == PRESNAP_TARGET_LAYOUT: target_payload = {**payload, "target_layout": PRESNAP_TARGET_LAYOUT, "pca_basis_path": ""} target_sum_td, target_sum_pool_d, target_dim = normalize_target_payload( target_payload, fallback_target_dim=int(self.target_dim), ) raw_target_sum_td = torch.as_tensor( payload.get("target_sum_td", payload.get("target_sum_t128")), dtype=torch.float32, ).contiguous() raw_target_sum_pool_d = torch.as_tensor( payload.get("target_sum_pool_d", payload.get("target_sum_pool_128")), dtype=torch.float32, ).contiguous() target_num_frames = int(payload["target_num_frames"]) if int(target_sum_td.shape[0]) != int(target_num_frames): raise RuntimeError( f"target_num_frames mismatch: target_sum_td has {target_sum_td.shape[0]} frames vs {target_num_frames}" ) if int(raw_target_sum_td.shape[0]) != int(target_num_frames): raise RuntimeError( f"raw target_num_frames mismatch: target_sum_td has {raw_target_sum_td.shape[0]} frames vs {target_num_frames}" ) grid_ft = torch.as_tensor(payload["grid_ft"], dtype=torch.float32).contiguous() grid_ids_ft = torch.as_tensor(payload["grid_ids_ft"], dtype=torch.long).contiguous() family_onsets_ft = torch.as_tensor(payload["family_onsets_ft"], dtype=torch.bool).contiguous() family_onset_count_ft = torch.as_tensor(payload["family_onset_count_ft"], dtype=torch.uint8).contiguous() grid_times_sec_t = torch.as_tensor(payload["grid_times_sec_t"], dtype=torch.float32).contiguous() token_times_sec_t = torch.as_tensor(payload["token_times_sec_t"], dtype=torch.float32).contiguous() beat_boundaries_sec_rel = torch.as_tensor(payload["beat_boundaries_sec_rel"], dtype=torch.float32).contiguous() source_codes_ct = torch.as_tensor(payload["source_codes_ct"], dtype=torch.int16).contiguous() grid_num_frames = int(payload["grid_num_frames"]) grid_frame_rate = float(payload["grid_frame_rate"]) bpm = float(payload["bpm"]) duration_sec = float(payload["duration_sec"]) if int(grid_ft.dim()) != 2: raise RuntimeError(f"expected grid_ft [F,Tg], got {tuple(grid_ft.shape)}") if int(grid_ids_ft.dim()) != 2: raise RuntimeError(f"expected grid_ids_ft [C,Tg], got {tuple(grid_ids_ft.shape)}") if int(family_onsets_ft.dim()) != 2 or int(family_onset_count_ft.dim()) != 2: raise RuntimeError( f"expected family onset tensors [C,Tg], got {tuple(family_onsets_ft.shape)} / {tuple(family_onset_count_ft.shape)}" ) if float(self.grid_rate_hz) > 0.0: decimated = downsample_grid_payload( grid_ft=grid_ft, grid_ids_ft=grid_ids_ft, family_onsets_ft=family_onsets_ft, family_onset_count_ft=family_onset_count_ft, grid_num_frames=int(grid_num_frames), duration_sec=float(payload["duration_sec"]), target_rate_hz=float(self.grid_rate_hz), ) grid_ft = decimated["grid_ft"] grid_ids_ft = decimated["grid_ids_ft"] family_onsets_ft = decimated["family_onsets_ft"] family_onset_count_ft = decimated["family_onset_count_ft"] grid_num_frames = int(decimated["grid_num_frames"]) grid_frame_rate = float(decimated["grid_frame_rate"]) grid_times_sec_t = uniform_grid_times( int(grid_num_frames), float(payload["duration_sec"]), ) if int(grid_ft.shape[-1]) != int(grid_num_frames): raise RuntimeError( f"grid_num_frames mismatch: grid_ft has {grid_ft.shape[-1]} frames vs {grid_num_frames}" ) if int(grid_ids_ft.shape[-1]) != int(grid_num_frames): raise RuntimeError( f"grid_ids_ft mismatch: grid_ids_ft has {grid_ids_ft.shape[-1]} frames vs {grid_num_frames}" ) if int(grid_times_sec_t.shape[0]) != int(grid_num_frames): raise RuntimeError( f"grid_times_sec_t mismatch: grid_times_sec_t has {grid_times_sec_t.shape[0]} frames vs {grid_num_frames}" ) if int(token_times_sec_t.shape[0]) != int(target_num_frames): raise RuntimeError( f"token_times_sec_t mismatch: token_times_sec_t has {token_times_sec_t.shape[0]} frames vs {target_num_frames}" ) class_names = tuple(str(x) for x in list(payload.get("class_names") or [])) class_id_vocab_sizes = tuple(int(x) for x in list(payload.get("class_id_vocab_sizes") or [])) feature_row_names = tuple(str(x) for x in list(payload.get("feature_row_names") or [])) return DiffusionExample( example_path=example_path, source_id=str(payload.get("source_id") or row.get("source_id") or ""), source_manifest_index=int(payload.get("source_manifest_index", row.get("source_manifest_index", -1))), source_pt_rel=str(payload.get("source_pt_rel") or row.get("source_pt_rel") or ""), source_row_in_shard=int(payload.get("source_row_in_shard", row.get("source_row_in_shard", -1))), beat_index=int(payload.get("beat_index", row.get("beat_index", 0))), beat_index_end=int(payload.get("beat_index_end", row.get("beat_index_end", 0))), split=str(payload.get("split") or row.get("split") or self.split), kit_name=str(payload.get("kit_name") or row.get("kit_name") or ""), source_row_id=str(payload.get("source_row_id") or row.get("source_row_id") or ""), conditioning_mode=str(payload.get("conditioning_mode") or "midi_family_state_onset_ids"), class_names=class_names, class_id_vocab_sizes=class_id_vocab_sizes, feature_row_names=feature_row_names, grid_ft=grid_ft, grid_ids_ft=grid_ids_ft, family_onsets_ft=family_onsets_ft, family_onset_count_ft=family_onset_count_ft, grid_times_sec_t=grid_times_sec_t, token_times_sec_t=token_times_sec_t, beat_boundaries_sec_rel=beat_boundaries_sec_rel, grid_num_frames=grid_num_frames, grid_frame_rate=grid_frame_rate, bpm=bpm, duration_sec=duration_sec, target_sum_td=target_sum_td, target_full_sum_td=raw_target_sum_td, target_dim=int(target_dim), target_full_dim=int(raw_target_sum_td.shape[-1]), target_layout=( PRESNAP_TARGET_LAYOUT if self.target_layout == PRESNAP_TARGET_LAYOUT else str(payload.get("target_layout") or "framewise_sum") ), pca_basis_path=( "" if self.target_layout == PRESNAP_TARGET_LAYOUT else str(payload.get("pca_basis_path") or "") ), target_num_frames=target_num_frames, source_codes_ct=source_codes_ct, target_sum_pool_d=raw_target_sum_pool_d, ) def collate_diffusion_examples(items: Sequence[DiffusionExample]) -> dict[str, Any]: if not items: raise ValueError("expected at least one DiffusionExample") batch_size = int(len(items)) grid_len = int(max(item.grid_num_frames for item in items)) grid_dim = int(items[0].grid_ft.shape[0]) class_dim = int(items[0].grid_ids_ft.shape[0]) target_len = int(max(item.target_num_frames for item in items)) codebooks = int(items[0].source_codes_ct.shape[0]) beat_boundary_len = int(max(item.beat_boundaries_sec_rel.shape[0] for item in items)) target_dim = int(items[0].target_dim) target_full_dim = int(items[0].target_full_dim) if any(int(item.target_dim) != int(target_dim) for item in items): raise ValueError("all items in a batch must share target_dim") if any(int(item.target_full_dim) != int(target_full_dim) for item in items): raise ValueError("all items in a batch must share target_full_dim") target_layout = str(items[0].target_layout) if any(str(item.target_layout) != str(target_layout) for item in items): raise ValueError("all items in a batch must share target_layout") pca_basis_path = str(items[0].pca_basis_path) if any(str(item.pca_basis_path) != str(pca_basis_path) for item in items): raise ValueError("all items in a batch must share pca_basis_path") grid_bft = torch.zeros((batch_size, grid_dim, grid_len), dtype=torch.float32) grid_ids_bct = torch.full((batch_size, class_dim, grid_len), -1, dtype=torch.long) family_onsets_bft = torch.zeros((batch_size, class_dim, grid_len), dtype=torch.bool) family_onset_count_bft = torch.zeros((batch_size, class_dim, grid_len), dtype=torch.uint8) grid_valid_mask_bt = torch.zeros((batch_size, grid_len), dtype=torch.bool) grid_times_sec_bt = torch.zeros((batch_size, grid_len), dtype=torch.float32) token_times_sec_bt = torch.zeros((batch_size, target_len), dtype=torch.float32) beat_boundaries_sec_bk = torch.zeros((batch_size, beat_boundary_len), dtype=torch.float32) beat_boundaries_valid_mask_bk = torch.zeros((batch_size, beat_boundary_len), dtype=torch.bool) target_btd = torch.zeros((batch_size, target_len, target_dim), dtype=torch.float32) target_sum_btd = torch.zeros((batch_size, target_len, target_full_dim), dtype=torch.float32) target_valid_mask_bt = torch.zeros((batch_size, target_len), dtype=torch.bool) source_codes_bct = torch.full((batch_size, codebooks, target_len), -1, dtype=torch.int16) target_sum_pool_bd = torch.zeros((batch_size, target_full_dim), dtype=torch.float32) bpm_b = torch.zeros((batch_size,), dtype=torch.float32) duration_sec_b = torch.zeros((batch_size,), dtype=torch.float32) grid_frame_rate_b = torch.zeros((batch_size,), dtype=torch.float32) for row_idx, item in enumerate(items): gf = int(item.grid_num_frames) tf = int(item.target_num_frames) bk = int(item.beat_boundaries_sec_rel.shape[0]) grid_bft[int(row_idx), :, : int(gf)] = item.grid_ft[:, : int(gf)] grid_ids_bct[int(row_idx), :, : int(gf)] = item.grid_ids_ft[:, : int(gf)] family_onsets_bft[int(row_idx), :, : int(gf)] = item.family_onsets_ft[:, : int(gf)] family_onset_count_bft[int(row_idx), :, : int(gf)] = item.family_onset_count_ft[:, : int(gf)] grid_valid_mask_bt[int(row_idx), : int(gf)] = True grid_times_sec_bt[int(row_idx), : int(gf)] = item.grid_times_sec_t[: int(gf)] token_times_sec_bt[int(row_idx), : int(tf)] = item.token_times_sec_t[: int(tf)] beat_boundaries_sec_bk[int(row_idx), : int(bk)] = item.beat_boundaries_sec_rel[: int(bk)] beat_boundaries_valid_mask_bk[int(row_idx), : int(bk)] = True target_btd[int(row_idx), : int(tf), :] = item.target_sum_td[: int(tf)] target_sum_btd[int(row_idx), : int(tf), :] = item.target_full_sum_td[: int(tf)] target_valid_mask_bt[int(row_idx), : int(tf)] = True source_codes_bct[int(row_idx), :, : int(tf)] = item.source_codes_ct[:, : int(tf)] target_sum_pool_bd[int(row_idx), :] = item.target_sum_pool_d bpm_b[int(row_idx)] = float(item.bpm) duration_sec_b[int(row_idx)] = float(item.duration_sec) grid_frame_rate_b[int(row_idx)] = float(item.grid_frame_rate) batch = { "conditioning_mode": str(items[0].conditioning_mode), "class_names": list(items[0].class_names), "class_id_vocab_sizes": [int(x) for x in list(items[0].class_id_vocab_sizes)], "feature_row_names": list(items[0].feature_row_names), "grid": grid_bft.contiguous(), "grid_ids": grid_ids_bct.contiguous(), "family_onsets_bft": family_onsets_bft.contiguous(), "family_onset_count_bft": family_onset_count_bft.contiguous(), "grid_valid_mask": grid_valid_mask_bt.contiguous(), "grid_times_sec": grid_times_sec_bt.contiguous(), "token_times_sec": token_times_sec_bt.contiguous(), "beat_boundaries_sec": beat_boundaries_sec_bk.contiguous(), "beat_boundaries_valid_mask": beat_boundaries_valid_mask_bk.contiguous(), "bpm": bpm_b.contiguous(), "duration_sec": duration_sec_b.contiguous(), "grid_frame_rate_b": grid_frame_rate_b.contiguous(), "grid_num_frames_b": torch.tensor([int(item.grid_num_frames) for item in items], dtype=torch.long), "target_layout": str(target_layout), "target_dim": int(target_dim), "target_full_dim": int(target_full_dim), "target_btd": target_btd.contiguous(), "target_sum_btd": target_sum_btd.contiguous(), "target_valid_mask_bt": target_valid_mask_bt.contiguous(), "source_codes_bct": source_codes_bct.contiguous(), "target_sum_pool_bd": target_sum_pool_bd.contiguous(), "target_pca_basis_path": str(pca_basis_path), "target_num_frames_b": torch.tensor([int(item.target_num_frames) for item in items], dtype=torch.long), "source_manifest_index_b": torch.tensor([int(item.source_manifest_index) for item in items], dtype=torch.long), "source_row_in_shard_b": torch.tensor([int(item.source_row_in_shard) for item in items], dtype=torch.long), "beat_index_b": torch.tensor([int(item.beat_index) for item in items], dtype=torch.long), "beat_index_end_b": torch.tensor([int(item.beat_index_end) for item in items], dtype=torch.long), "source_id": [str(item.source_id) for item in items], "source_pt_rel": [str(item.source_pt_rel) for item in items], "split": [str(item.split) for item in items], "kit_name": [str(item.kit_name) for item in items], "source_row_id": [str(item.source_row_id) for item in items], "example_path": [str(item.example_path) for item in items], } return batch @torch.no_grad() def estimate_target_normalization( dataloader: DataLoader, *, device: str | torch.device = "cpu", eps: float = 1.0e-6, ) -> tuple[torch.Tensor, torch.Tensor]: resolved_device = torch.device(device) target_sum_d: torch.Tensor | None = None target_sq_sum_d: torch.Tensor | None = None total_frames = 0 for batch in dataloader: target_btd = torch.as_tensor(batch["target_btd"], dtype=torch.float32, device=resolved_device) target_valid_mask_bt = torch.as_tensor( batch["target_valid_mask_bt"], dtype=torch.bool, device=resolved_device, ) if int(target_btd.dim()) != 3: raise ValueError(f"expected target_btd [B,T,D], got {tuple(target_btd.shape)}") if tuple(target_valid_mask_bt.shape) != tuple(target_btd.shape[:2]): raise ValueError( "target_valid_mask_bt must match target_btd batch/time dimensions, got " f"{tuple(target_valid_mask_bt.shape)} vs {tuple(target_btd.shape[:2])}" ) if target_sum_d is None: target_sum_d = torch.zeros((int(target_btd.shape[-1]),), dtype=torch.float32, device=resolved_device) target_sq_sum_d = torch.zeros_like(target_sum_d) valid_btd = target_btd * target_valid_mask_bt.unsqueeze(-1).to(dtype=target_btd.dtype) target_sum_d += valid_btd.sum(dim=(0, 1)) target_sq_sum_d += valid_btd.square().sum(dim=(0, 1)) total_frames += int(target_valid_mask_bt.sum().item()) if target_sum_d is None or target_sq_sum_d is None or int(total_frames) <= 0: raise RuntimeError("could not estimate target normalization from an empty dataloader") mean_d = target_sum_d / float(total_frames) var_d = (target_sq_sum_d / float(total_frames)) - mean_d.square() std_d = var_d.clamp_min(float(eps)).sqrt() return mean_d.contiguous(), std_d.contiguous() @torch.no_grad() def estimate_target_range( dataloader: DataLoader, *, device: str | torch.device = "cpu", ) -> tuple[torch.Tensor, torch.Tensor]: """Per-dimension minimum and maximum over the valid target frames.""" resolved_device = torch.device(device) target_min_d: torch.Tensor | None = None target_max_d: torch.Tensor | None = None for batch in dataloader: target_btd = torch.as_tensor(batch["target_btd"], dtype=torch.float32, device=resolved_device) valid_td = target_btd[torch.as_tensor(batch["target_valid_mask_bt"], dtype=torch.bool, device=resolved_device)] if int(valid_td.shape[0]) == 0: continue batch_min_d, batch_max_d = valid_td.amin(dim=0), valid_td.amax(dim=0) target_min_d = batch_min_d if target_min_d is None else torch.minimum(target_min_d, batch_min_d) target_max_d = batch_max_d if target_max_d is None else torch.maximum(target_max_d, batch_max_d) if target_min_d is None or target_max_d is None: raise RuntimeError("could not measure the target range from an empty dataloader") return target_min_d.contiguous(), target_max_d.contiguous() def build_diffusion_dataloader( cache_root: str | Path, *, split: str = "train", batch_size: int = 8, shuffle: bool = False, num_workers: int = 0, max_items: int = 0, conditioning_mode: str = "seconds", pin_memory: bool = False, persistent_workers: bool = False, multiprocessing_context: str | None = None, grid_rate_hz: float = 0.0, grid_overlay_root: str | Path | None = None, ) -> DataLoader: dataset = DiffusionConditioningDataset( cache_root, split=split, max_items=max_items, conditioning_mode=conditioning_mode, grid_rate_hz=float(grid_rate_hz), grid_overlay_root=grid_overlay_root, ) loader_kwargs: dict[str, Any] = {} if int(num_workers) > 0: loader_kwargs["persistent_workers"] = bool(persistent_workers) if multiprocessing_context is not None and str(multiprocessing_context).strip(): loader_kwargs["multiprocessing_context"] = str(multiprocessing_context).strip() return DataLoader( dataset, batch_size=int(batch_size), shuffle=bool(shuffle), num_workers=int(num_workers), pin_memory=bool(pin_memory), collate_fn=collate_diffusion_examples, **loader_kwargs, )