File size: 6,984 Bytes
8fd1a55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | """Load-time grid-rate decimation for the grid-resolution ablation.
The cache is rendered once at 250 Hz. Coarser conditioning rates are obtained by
decimating that rendering here, at __getitem__ time, so the ablation needs no
extra caches and leaves the audio/target side byte-identical across arms.
Frames are partitioned into contiguous groups by target frame count, and each
channel is reduced with the rule that preserves its meaning:
* state velocity -> max (a sustained state stays present)
* onset velocity -> max (a hit keeps its velocity, never averaged away)
* onset count -> sum (two hits merging into one frame is the effect
under study, so counts must accumulate)
* onset flags -> any
* class ids -> value at the loudest onset in the group, else the state
Upsampling is refused: timing precision the 250 Hz raster discarded cannot be
recovered by interpolation, so a finer rate must be re-rendered from MIDI.
"""
from __future__ import annotations
import numpy as np
import torch
# grid_ft rows are [state_vel, onset_vel, onset_count] interleaved per family,
# matching stack_family_state_grid in diffusion_cache_utils.
ROWS_PER_FAMILY = 3
STATE_VEL_OFFSET = 0
ONSET_VEL_OFFSET = 1
ONSET_COUNT_OFFSET = 2
def target_frame_count(*, grid_num_frames: int, duration_sec: float, target_rate_hz: float) -> int:
"""Frame count for target_rate_hz, never exceeding what the cache holds."""
frames = int(max(0, int(grid_num_frames)))
if frames <= 0:
return 0
duration = float(max(0.0, float(duration_sec)))
rate = float(target_rate_hz)
if not rate > 0.0 or not duration > 0.0:
return frames
return int(min(int(frames), max(1, int(round(duration * rate)))))
def _group_bounds(num_frames: int, num_groups: int) -> np.ndarray:
"""Contiguous group edges spreading num_frames over num_groups."""
return np.linspace(0, int(num_frames), num=int(num_groups) + 1).round().astype(np.int64)
def downsample_grid_payload(
*,
grid_ft: torch.Tensor,
grid_ids_ft: torch.Tensor,
family_onsets_ft: torch.Tensor,
family_onset_count_ft: torch.Tensor,
grid_num_frames: int,
duration_sec: float,
target_rate_hz: float,
) -> dict[str, object]:
"""Decimate one cached example to target_rate_hz.
Returns the replacement tensors plus the new frame count / effective rate.
A target rate matching the cached frame count is a no-op returning the
inputs unchanged. A genuinely higher rate raises instead of fabricating
timing precision by interpolation.
"""
frames = int(grid_num_frames)
requested_frames = int(max(1, round(float(duration_sec) * float(target_rate_hz))))
if int(frames) > 0 and float(duration_sec) > 0.0 and int(requested_frames) > int(frames):
cached_rate = float(frames) / float(duration_sec)
raise ValueError(
"grid-rate upsampling is invalid: "
f"requested {float(target_rate_hz):g} Hz ({requested_frames} frames) from a "
f"{cached_rate:.3f} Hz cache ({frames} frames). Re-render the grid from MIDI."
)
out_frames = target_frame_count(
grid_num_frames=frames,
duration_sec=duration_sec,
target_rate_hz=target_rate_hz,
)
duration = float(max(1.0e-6, float(duration_sec)))
if int(out_frames) >= int(frames) or int(frames) <= 0:
# No-op: identical tensors, so the source-rate arm is bit-identical to
# bypassing this module entirely.
return {
"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(frames),
"grid_frame_rate": float(frames) / duration,
}
bounds = _group_bounds(int(frames), int(out_frames))
num_families = int(family_onsets_ft.shape[0])
grid_out = torch.zeros((int(grid_ft.shape[0]), int(out_frames)), dtype=grid_ft.dtype)
ids_out = torch.zeros((int(grid_ids_ft.shape[0]), int(out_frames)), dtype=grid_ids_ft.dtype)
onsets_out = torch.zeros((num_families, int(out_frames)), dtype=family_onsets_ft.dtype)
counts_out = torch.zeros((num_families, int(out_frames)), dtype=torch.int64)
onset_count_i64 = family_onset_count_ft.to(dtype=torch.int64)
for out_idx in range(int(out_frames)):
start = int(bounds[int(out_idx)])
stop = int(max(int(start) + 1, int(bounds[int(out_idx) + 1])))
window_grid = grid_ft[:, int(start) : int(stop)]
# Velocities take the max so a hit is never diluted by neighbouring
# silence; counts sum so merged hits stay visible to the loss.
grid_out[:, int(out_idx)] = window_grid.max(dim=1).values
for family_idx in range(int(num_families)):
count_row = int(ONSET_COUNT_OFFSET + ROWS_PER_FAMILY * int(family_idx))
if int(count_row) < int(grid_ft.shape[0]):
grid_out[int(count_row), int(out_idx)] = window_grid[int(count_row)].sum()
onsets_out[:, int(out_idx)] = family_onsets_ft[:, int(start) : int(stop)].any(dim=1)
counts_out[:, int(out_idx)] = onset_count_i64[:, int(start) : int(stop)].sum(dim=1)
# Ids follow the loudest onset in the group; with no onset the group is
# a sustained state, so carry the first frame's id.
window_onset_vel = grid_ft[ONSET_VEL_OFFSET::ROWS_PER_FAMILY, int(start) : int(stop)]
for family_idx in range(int(ids_out.shape[0])):
if int(family_idx) < int(window_onset_vel.shape[0]):
vel_row = window_onset_vel[int(family_idx)]
if float(vel_row.max()) > 0.0:
pick = int(start) + int(torch.argmax(vel_row).item())
else:
pick = int(start)
else:
pick = int(start)
ids_out[int(family_idx), int(out_idx)] = grid_ids_ft[int(family_idx), int(pick)]
count_dtype = family_onset_count_ft.dtype
if count_dtype == torch.uint8:
counts_out = counts_out.clamp_(0, 255)
return {
"grid_ft": grid_out.contiguous(),
"grid_ids_ft": ids_out.contiguous(),
"family_onsets_ft": onsets_out.contiguous(),
"family_onset_count_ft": counts_out.to(dtype=count_dtype).contiguous(),
"grid_num_frames": int(out_frames),
"grid_frame_rate": float(out_frames) / duration,
}
def uniform_grid_times(num_frames: int, duration_sec: float) -> torch.Tensor:
"""Frame-centre times, matching _uniform_frame_times_from_duration (bpm geometry)."""
frames = int(max(0, int(num_frames)))
duration = float(max(0.0, float(duration_sec)))
if frames <= 0 or not duration > 0.0:
return torch.zeros((0,), dtype=torch.float32)
idx = torch.arange(int(frames), dtype=torch.float32)
return ((idx + 0.5) * (float(duration) / float(frames))).contiguous()
|