# utils.py from __future__ import annotations import io, base64, math from math import gcd import numpy as np import soundfile as sf from scipy.signal import resample_poly # Magenta RT audio types from magenta_rt import audio as au # Optional loudness try: import pyloudnorm as pyln _HAS_LOUDNORM = True except Exception: _HAS_LOUDNORM = False # ---------- Loudness ---------- def _measure_lufs(wav: au.Waveform) -> float: meter = pyln.Meter(wav.sample_rate) # BS.1770-4 return float(meter.integrated_loudness(wav.samples)) def _rms(x: np.ndarray) -> float: if x.size == 0: return 0.0 return float(np.sqrt(np.mean(x**2))) def match_loudness_to_reference( ref: au.Waveform, target: au.Waveform, method: str = "auto", # "auto"|"lufs"|"rms"|"none" headroom_db: float = 1.0 ) -> tuple[au.Waveform, dict]: stats = {"method": method, "applied_gain_db": 0.0} if method == "none": return target, stats if method == "auto": method = "lufs" if _HAS_LOUDNORM else "rms" if method == "lufs" and _HAS_LOUDNORM: L_ref = _measure_lufs(ref) L_tgt = _measure_lufs(target) delta_db = L_ref - L_tgt gain = 10.0 ** (delta_db / 20.0) y = target.samples.astype(np.float32) * gain stats.update({"ref_lufs": L_ref, "tgt_lufs_before": L_tgt, "applied_gain_db": delta_db}) else: ra = _rms(ref.samples) rb = _rms(target.samples) if rb <= 1e-12: return target, stats gain = ra / rb y = target.samples.astype(np.float32) * gain stats.update({"ref_rms": ra, "tgt_rms_before": rb, "applied_gain_db": 20*np.log10(max(gain,1e-12))}) # simple peak “limiter” to keep headroom limit = 10 ** (-headroom_db / 20.0) # e.g., -1 dBFS peak = float(np.max(np.abs(y))) if y.size else 0.0 if peak > limit: y *= (limit / peak) stats["post_peak_limited"] = True else: stats["post_peak_limited"] = False target.samples = y.astype(np.float32) return target, stats # ---------- Stitch / fades / trims ---------- def stitch_generated(chunks, sr: int, xfade_s: float, drop_first_pre_roll: bool = True): if not chunks: raise ValueError("no chunks") xfade_n = int(round(xfade_s * sr)) if xfade_n <= 0: return au.Waveform(np.concatenate([c.samples for c in chunks], axis=0), sr) t = np.linspace(0, np.pi/2, xfade_n, endpoint=False, dtype=np.float32) eq_in, eq_out = np.sin(t)[:, None], np.cos(t)[:, None] first = chunks[0].samples if first.shape[0] < xfade_n: raise ValueError("chunk shorter than crossfade prefix") # 🔧 key change: out = first[xfade_n:].copy() if drop_first_pre_roll else first.copy() for i in range(1, len(chunks)): cur = chunks[i].samples if cur.shape[0] < xfade_n: continue head, tail = cur[:xfade_n], cur[xfade_n:] mixed = out[-xfade_n:] * eq_out + head * eq_in out = np.concatenate([out[:-xfade_n], mixed, tail], axis=0) return au.Waveform(out, sr) def hard_trim_seconds(wav: au.Waveform, seconds: float) -> au.Waveform: n = int(round(seconds * wav.sample_rate)) return au.Waveform(wav.samples[:n], wav.sample_rate) def apply_micro_fades(wav: au.Waveform, ms: int = 5) -> None: n = int(wav.sample_rate * ms / 1000.0) if n > 0 and wav.samples.shape[0] > 2*n: env = np.linspace(0.0, 1.0, n, dtype=np.float32)[:, None] wav.samples[:n] *= env wav.samples[-n:] *= env[::-1] # ---------- Token context helpers ---------- def make_bar_aligned_context(tokens, bpm, fps=25.0, ctx_frames=250, beats_per_bar=4): """ Return a ctx_frames-long slice of `tokens` whose **end** lands on an integer bar boundary in codec-frame space (model runs at `fps`, typically 25). """ if tokens is None: raise ValueError("tokens is None") tokens = np.asarray(tokens) if tokens.ndim == 1: tokens = tokens[:, None] T = tokens.shape[0] if T == 0: return tokens fps = float(fps) # float frames per bar (e.g., ~65.934 at 91 BPM for 4/4 @ 25fps) frames_per_bar_f = (beats_per_bar * 60.0 / float(bpm)) * fps # >>> KEY FIX: quantize bar length to an integer number of codec frames frames_per_bar_i = max(1, int(round(frames_per_bar_f))) # Tile so we can always snap the *end* to a bar boundary and still have ctx_frames reps = int(np.ceil((ctx_frames + T) / float(T))) + 1 tiled = np.tile(tokens, (reps, 1)) total = tiled.shape[0] # How many whole integer bars fit in the tiled sequence? k_bars = total // frames_per_bar_i if k_bars <= 0: return tiled[-ctx_frames:] # Snap END to an exact integer multiple of frames_per_bar_i end_idx = int(k_bars * frames_per_bar_i) end_idx = min(max(end_idx, ctx_frames), total) start_idx = end_idx - ctx_frames if start_idx < 0: start_idx = 0 end_idx = ctx_frames window = tiled[start_idx:end_idx] # Guard off-by-one if window.shape[0] < ctx_frames: pad = np.tile(tokens, (int(np.ceil((ctx_frames - window.shape[0]) / T)), 1)) window = np.vstack([window, pad])[:ctx_frames] elif window.shape[0] > ctx_frames: window = window[-ctx_frames:] return window def take_bar_aligned_tail( wav: au.Waveform, bpm: float, beats_per_bar: int, ctx_seconds: float, max_bars=None ) -> au.Waveform: """ Take a tail whose length is an integer number of bars, with the END aligned to a bar boundary. Uses ceil for bars_needed so we never under-fill the context. IMPORTANT: If the input loop is shorter than needed, we TILE it to create enough audio before taking the tail. This ensures the context window is always filled with actual looping audio, not zeros or padding. """ import math # seconds per bar spb = (60.0 / float(bpm)) * float(beats_per_bar) # Pick enough whole bars to cover ctx_seconds (avoid underfilling on round-down). # The small epsilon avoids an extra bar due to FP jitter when ctx_seconds ~= k * spb. eps = 1e-9 bars_needed = max(1, int(math.ceil((float(ctx_seconds) - eps) / spb))) if max_bars is not None: bars_needed = min(bars_needed, int(max_bars)) # Convert bars -> samples (do rounding once at the end for stability) samples_per_bar_f = spb * float(wav.sample_rate) n = int(round(bars_needed * samples_per_bar_f)) total = int(wav.samples.shape[0]) # ======================== KEY FIX ======================== # If the loop is too short, TILE it to fill the context window. # This ensures we always have continuous looping audio as context, # rather than padding with zeros or repeated tokens at the front. if n > total and total > 0: # Calculate how many times we need to tile reps = int(math.ceil(n / total)) + 1 # +1 for safety margin samples = wav.samples if samples.ndim == 1: samples = samples[:, None] tiled_samples = np.tile(samples, (reps, 1)) wav = au.Waveform(tiled_samples, wav.sample_rate) total = int(wav.samples.shape[0]) # ========================================================= if n >= total: # Still not enough after tiling (edge case: empty input) return wav start = total - n return au.Waveform(wav.samples[start:], wav.sample_rate) # ---------- SR normalize + snap ---------- def resample_and_snap(x: np.ndarray, cur_sr: int, target_sr: int, seconds: float) -> np.ndarray: """ x: np.ndarray shape (S, C), float32 Returns: exact-length array (round(seconds*target_sr), C) """ if x.ndim == 1: x = x[:, None] if cur_sr != target_sr: g = gcd(cur_sr, target_sr) up, down = target_sr // g, cur_sr // g x = resample_poly(x, up, down, axis=0) expected_len = int(round(seconds * target_sr)) if x.shape[0] < expected_len: pad = np.zeros((expected_len - x.shape[0], x.shape[1]), dtype=x.dtype) x = np.vstack([x, pad]) elif x.shape[0] > expected_len: x = x[:expected_len, :] return x.astype(np.float32, copy=False) # ---------- WAV encode ---------- def wav_bytes_base64(x: np.ndarray, sr: int) -> tuple[str, int, int]: """ x: np.ndarray shape (S, C) returns: (base64_wav, total_samples, channels) """ buf = io.BytesIO() sf.write(buf, x, sr, subtype="FLOAT", format="WAV") buf.seek(0) b64 = base64.b64encode(buf.read()).decode("utf-8") return b64, int(x.shape[0]), int(x.shape[1]) def _ratio(out_sr: int, in_sr: int) -> tuple[int, int]: g = gcd(int(out_sr), int(in_sr)) return int(out_sr) // g, int(in_sr) // g class StreamingResampler: """ Stateful streaming resampler. Prefers soxr (best), then libsamplerate; final fallback is block resample_poly. Always pass float32 arrays shaped (S, C). """ def __init__(self, in_sr: int, out_sr: int, channels: int = 2, quality: str = "VHQ"): self.in_sr = int(in_sr) self.out_sr = int(out_sr) self.channels = int(channels) self.quality = quality self._backend = None # Try soxr first try: import soxr # pip install soxr self._backend = "soxr" # dtype float32 keeps things consistent with the rest of your code self._rs = soxr.Resampler( self.in_sr, self.out_sr, channels=self.channels, dtype="float32", quality=self.quality, # "Q", "HQ", "VHQ" ) except Exception: # Try libsamplerate try: import samplerate # pip install samplerate self._backend = "samplerate" # sinc_best == highest quality; you can choose 'sinc_medium' for speed self._rs = samplerate.Resampler(converter_type="sinc_best", channels=self.channels) except Exception: # Last resort: block resample (not truly streaming) from scipy.signal import resample_poly self._backend = "scipy" self._resample_poly = resample_poly self._L, self._M = _ratio(self.out_sr, self.in_sr) # Keep a tiny tail to help transitions (still not perfect vs true streaming) self._hist = np.zeros((0, self.channels), dtype=np.float32) def process(self, x: np.ndarray, final: bool = False) -> np.ndarray: """Feed a chunk (S, C) and get resampled chunk (S', C). Keep calling in order.""" if x.size == 0 and not final: # nothing to do return np.zeros((0, self.channels), dtype=np.float32) if self._backend == "soxr": return self._rs.process(x, final=final) elif self._backend == "samplerate": import samplerate ratio = float(self.out_sr) / float(self.in_sr) # end_of_input=True flushes tail on the last call y = self._rs.process(x, ratio, end_of_input=final) # libsamplerate returns (S', C) return y.astype(np.float32, copy=False) # --- scipy fallback (block, not truly streaming) --- # We concatenate a short history to reduce block edge artifacts x_ext = x if self._hist.size == 0 else np.vstack([self._hist, x]) y = self._resample_poly(x_ext, up=self._L, down=self._M, axis=0).astype(np.float32, copy=False) # Heuristic: drop the portion corresponding roughly to the history to avoid duplicate content # (Not perfect, but helps a lot when chunks are reasonably sized.) drop = int(round(self._hist.shape[0] * self.out_sr / self.in_sr)) y = y[drop:] if drop < y.shape[0] else np.zeros((0, self.channels), dtype=np.float32) # Keep a small input tail for the next call (say ~ 4 ms at in_sr) tail_samples = max(int(0.004 * self.in_sr), 1) self._hist = x[-tail_samples:] if x.shape[0] >= tail_samples else x.copy() if final: self._hist = np.zeros((0, self.channels), dtype=np.float32) return y def flush(self) -> np.ndarray: """Drain converter tail (call at stop).""" if self._backend == "soxr": return self._rs.process(np.zeros((0, self.channels), dtype=np.float32), final=True) elif self._backend == "samplerate": ratio = float(self.out_sr) / float(self.in_sr) return self._rs.process(np.zeros((0, self.channels), dtype=np.float32), ratio, end_of_input=True) else: # nothing meaningful to flush in scipy fallback return np.zeros((0, self.channels), dtype=np.float32)