97 lines
2.7 KiB
Python
97 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Chunk boundary and overlap arithmetic for MiniMax Music 3.
|
|
|
|
The AR stage emits fixed 200-frame windows that overlap by half. The acoustic
|
|
stage decodes each window with the previous window's latent as a prompt, then
|
|
crops the overlap back off so the concatenated waveform is seamless. All of
|
|
that is pure integer arithmetic and lives here so it can be unit tested
|
|
without a GPU.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from .constants import (
|
|
AR_CHUNK_FRAMES,
|
|
AR_CHUNK_HOP_FRAMES,
|
|
BLEND_MEL_FRAMES,
|
|
DAV_HOP_SAMPLES,
|
|
HOP_MEL_FRAMES,
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ChunkWindow:
|
|
"""One acoustic window, in AR frame coordinates."""
|
|
|
|
index: int
|
|
start: int
|
|
end: int
|
|
is_first: bool
|
|
is_last: bool
|
|
|
|
@property
|
|
def length(self) -> int:
|
|
return self.end - self.start
|
|
|
|
|
|
def chunk_windows(frames: int) -> list[ChunkWindow]:
|
|
"""Return the fixed-size windows covering ``frames``, overlapping by half.
|
|
|
|
A song shorter than one window produces a single short window; otherwise
|
|
windows advance by ``AR_CHUNK_HOP_FRAMES`` and the final one is truncated
|
|
at the end of the song.
|
|
"""
|
|
frames = int(frames)
|
|
if frames <= 0:
|
|
return []
|
|
if frames <= AR_CHUNK_FRAMES:
|
|
return [ChunkWindow(0, 0, frames, True, True)]
|
|
|
|
windows: list[ChunkWindow] = []
|
|
index = 0
|
|
start = 0
|
|
while start < frames:
|
|
end = min(start + AR_CHUNK_FRAMES, frames)
|
|
windows.append(
|
|
ChunkWindow(
|
|
index=index,
|
|
start=start,
|
|
end=end,
|
|
is_first=index == 0,
|
|
is_last=end >= frames,
|
|
)
|
|
)
|
|
if end >= frames:
|
|
break
|
|
index += 1
|
|
start += AR_CHUNK_HOP_FRAMES
|
|
return windows
|
|
|
|
|
|
def overlap_mel_length() -> int:
|
|
"""Latent frames carried forward as the next window's flow-matching prompt."""
|
|
return HOP_MEL_FRAMES // 2
|
|
|
|
|
|
def crop_sample_bounds(window: ChunkWindow) -> tuple[int, int]:
|
|
"""Return ``(left, right)`` waveform samples to trim from a decoded window.
|
|
|
|
The first window keeps its head and the last keeps its tail; every other
|
|
window is cropped on both sides so that concatenating the survivors
|
|
reconstructs the song exactly once.
|
|
"""
|
|
left = 0 if window.is_first else BLEND_MEL_FRAMES * DAV_HOP_SAMPLES
|
|
right = 0 if window.is_last else (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES
|
|
return left, right
|
|
|
|
|
|
__all__ = [
|
|
"ChunkWindow",
|
|
"chunk_windows",
|
|
"crop_sample_bounds",
|
|
"overlap_mel_length",
|
|
]
|