SoundGame/tools/stable-audio-3/optimized/tflite/scripts/spec.py

"""Mel-spectrogram renderer — numpy-only port of the TensorRT release's spec.py.

Same algorithm and constants as optimized/tensorRT/scripts/spec.py (the 3-band
tinted stereo mel spectrogram from our run dashboards), with the torch pieces
(STFT, resample, filterbank matmul) rewritten in numpy so the MLX release stays
torch-free. Numerics match: periodic Hann window, center=True reflect padding,
Slaney mel scale, librosa-compatible filterbank, linear resample with
align_corners=False sampling positions.

Exposes render_spectrogram_png(pcm_int16, sample_rate, width, height).
"""
from __future__ import annotations
import io
import math
from functools import lru_cache

import numpy as np
from PIL import Image


# 3-band tint
SPEC_BANDS = [
    (0, 200, (1.0, 0.0, 0.0)),       # Bass -> Red
    (200, 1500, (0.0, 1.0, 0.0)),    # Mid  -> Green
    (1500, 16000, (0.0, 0.0, 1.0)),  # High -> Blue
]
_BAND_COLORS = np.array([c for _, _, c in SPEC_BANDS], dtype=np.float32)

# Slaney mel-scale constants (linear below 1 kHz, log above), matches librosa default.
_F_SP = 200.0 / 3
_MIN_LOG_HZ = 1000.0
_MIN_LOG_MEL = _MIN_LOG_HZ / _F_SP
_LOGSTEP = math.log(6.4) / 27.0


def _hz_to_mel(hz):
    hz = np.asarray(hz, dtype=np.float64)
    log_term = np.log(np.maximum(hz, _MIN_LOG_HZ) / _MIN_LOG_HZ) / _LOGSTEP
    return np.where(hz >= _MIN_LOG_HZ, _MIN_LOG_MEL + log_term, hz / _F_SP)


def _mel_to_hz(mels):
    mels = np.asarray(mels, dtype=np.float64)
    return np.where(mels >= _MIN_LOG_MEL,
                    _MIN_LOG_HZ * np.exp(_LOGSTEP * (mels - _MIN_LOG_MEL)),
                    _F_SP * mels)


def _mel_frequencies(n_mels, fmax, fmin=0.0):
    return _mel_to_hz(np.linspace(_hz_to_mel(fmin), _hz_to_mel(fmax), n_mels))


@lru_cache(maxsize=8)
def _mel_filterbank(n_mels, n_fft, sr, fmax):
    pts = _mel_to_hz(np.linspace(_hz_to_mel(0.0), _hz_to_mel(fmax), n_mels + 2))
    fft_f = np.linspace(0, sr / 2, n_fft // 2 + 1)
    filt = np.zeros((n_mels, n_fft // 2 + 1), dtype=np.float32)
    for i in range(n_mels):
        lo, ce, hi = pts[i], pts[i + 1], pts[i + 2]
        left = (fft_f - lo) / max(ce - lo, 1e-10)
        right = (hi - fft_f) / max(hi - ce, 1e-10)
        filt[i] = np.maximum(0, np.minimum(left, right))
    enorm = (2.0 / (pts[2:n_mels + 2] - pts[0:n_mels])).astype(np.float32)
    filt *= enorm[:, None]
    return filt


def _stft_power(y, n_fft, hop):
    """|STFT|^2 with periodic Hann + center=True reflect pad (torch.stft-compatible).
    Returns (n_fft//2+1, n_frames) float32."""
    y = np.asarray(y, dtype=np.float32)
    pad = n_fft // 2
    y = np.pad(y, pad, mode="reflect")
    n_frames = 1 + (len(y) - n_fft) // hop
    idx = np.arange(n_fft)[None, :] + hop * np.arange(n_frames)[:, None]
    # periodic Hann (torch.hann_window default), not numpy's symmetric hanning
    win = (0.5 - 0.5 * np.cos(2.0 * np.pi * np.arange(n_fft) / n_fft)).astype(np.float32)
    frames = y[idx] * win[None, :]
    spec = np.fft.rfft(frames, n=n_fft, axis=1)
    return (spec.real ** 2 + spec.imag ** 2).T.astype(np.float32)


def _melspectrogram(y_ch, sr, n_mels=30, fmax=16000, hop_length=2048, n_fft=2048):
    power = _stft_power(y_ch, n_fft, hop_length)
    return _mel_filterbank(n_mels, n_fft, sr, fmax) @ power


def _power_to_db(S, top_db=80.0):
    log_spec = 10.0 * np.log10(np.maximum(S, 1e-10))
    return np.maximum(log_spec - log_spec.max(), -top_db)


def _mel_channel(y_ch, sr, n_mels=30):
    """dB-scaled mel + band-tinted RGB for one mono channel."""
    S = _melspectrogram(y_ch, sr, n_mels=n_mels, fmax=16000, hop_length=2048)

    S_db = _power_to_db(S)
    np.clip(S_db, -60, 0, out=S_db)
    S_db += 60.0
    S_db /= 60.0
    np.power(S_db, 0.6, out=S_db)

    mel_f = _mel_frequencies(n_mels, fmax=16000)
    n_frames = S.shape[1]
    band_norms = np.empty((3, n_frames), dtype=np.float32)
    for i, (flo, fhi, _) in enumerate(SPEC_BANDS):
        mask = (mel_f >= flo) & (mel_f < fhi)
        if mask.any():
            power = np.sum(S[mask], axis=0)
            db = 10.0 * np.log10(power + 1e-10)
            np.clip(db, -20, None, out=db)
            db -= -20
            mx = db.max()
            if mx > 0:
                db /= mx
            band_norms[i] = db
        else:
            band_norms[i] = 0.0

    rgb = band_norms.T @ _BAND_COLORS
    for c in range(3):
        mx = rgb[:, c].max()
        if mx > 0:
            rgb[:, c] /= mx

    return S_db, rgb


def _resample_to_32k(pcm_int16, sr_in):
    """Resample (T, 2) int16 PCM to (2, T') float32 at 32 kHz.
    Linear interpolation at align_corners=False sample positions (matches the
    torch.nn.functional.interpolate call in the TRT original)."""
    target_sr = 32000
    y = pcm_int16.astype(np.float32).T / 32768.0      # (2, T)
    if sr_in == target_sr:
        return y, target_sr
    in_len = y.shape[-1]
    new_len = int(round(in_len * target_sr / sr_in))
    scale = in_len / new_len
    pos = np.clip((np.arange(new_len, dtype=np.float64) + 0.5) * scale - 0.5, 0, in_len - 1)
    out = np.stack([np.interp(pos, np.arange(in_len), ch) for ch in y]).astype(np.float32)
    return out, target_sr


def render_spectrogram_png(pcm_int16, sample_rate=44100, width=1200, height=240,
                           n_mels=30) -> bytes:
    """Render the 3-band tinted stereo mel spectrogram to PNG bytes.

    Args:
        pcm_int16: (T, 2) int16 numpy array — stereo PCM.
        sample_rate: input PCM sample rate (default 44100).
        width, height: final image size.
        n_mels: mel filterbank size (30 matches the dashboards).

    Returns: PNG bytes ready to base64 / save / etc.
    """
    if pcm_int16.ndim != 2 or pcm_int16.shape[1] != 2:
        raise ValueError(f"expected (T, 2) PCM, got {pcm_int16.shape}")
    y, sr = _resample_to_32k(pcm_int16, sample_rate)
    S_L, rgb_L = _mel_channel(y[0], sr, n_mels=n_mels)
    S_R, rgb_R = _mel_channel(y[1], sr, n_mels=n_mels)

    nf = min(S_L.shape[1], S_R.shape[1])
    S_L, S_R = S_L[:, :nf], S_R[:, :nf]
    rgb_L, rgb_R = rgb_L[:nf], rgb_R[:nf]
    nm = S_L.shape[0]

    # L channel: flip vertically so bass-band is at bottom.
    S_L = S_L[::-1]

    img = np.empty((nm * 2, nf, 3), dtype=np.float32)
    img[:nm] = S_L[:, :, np.newaxis] * rgb_L[np.newaxis, :, :]
    img[nm:] = S_R[:, :, np.newaxis] * rgb_R[np.newaxis, :, :]

    np.clip(img, 0, 1, out=img)
    img *= 255
    pil = Image.fromarray(img.astype(np.uint8)).resize((width, height), Image.LANCZOS)
    buf = io.BytesIO()
    pil.save(buf, format="PNG", optimize=True)
    return buf.getvalue()