enhance.py

"""input enhancement: coda hears the song, not the noise.

people's unfinished songs live on phone recordings, voice memos, old mp3
rips — low level, hissy, rumbly. feeding that straight to musicgen makes
the model imitate the *recording* (artifacts included) instead of the
*music*. this module cleans the copy of the input that feeds analysis and
the audio prompt: high-pass the rumble, spectrally gate the steady noise
floor, normalize the level. the user's untouched original still goes into
the final track (unless they choose remaster mode, which runs the same
cleanup on the original before the splice).

deliberately gentle, deliberately DSP-only: a denoiser that invents
content would be another model to fight. nothing here changes the length
or timing of the audio — the sample-aligned crossfade math downstream
depends on that.
"""
import os
import tempfile

import librosa
import numpy as np
import soundfile as sf
from scipy.signal import butter, filtfilt

# below this the input gets flagged as lo-fi in the ui
LOFI_BANDWIDTH_HZ = 13000
LOFI_NOISE_DB = -38.0

N_FFT = 2048
HOP = 512


def _highpass(y, sr, cutoff=35.0):
    """rumble filter. filtfilt = zero phase shift, length preserved."""
    b, a = butter(2, cutoff / (sr / 2), btype="high")
    return filtfilt(b, a, y, axis=-1).astype(np.float32)


def _gate_channel(y, sr, reduction_db=12.0, percentile=10):
    """
    spectral gating on one channel: estimate the per-band noise floor from
    the quietest frames, then softly attenuate anything near that floor.
    musical content sits well above the floor and passes untouched.
    """
    spec = librosa.stft(y, n_fft=N_FFT, hop_length=HOP)
    mag = np.abs(spec)

    frame_energy = mag.mean(axis=0)
    quiet = mag[:, frame_energy <= np.percentile(frame_energy, percentile)]
    if quiet.shape[1] < 2:  # uniformly loud clip — nothing to learn from
        return y
    noise_profile = np.median(quiet, axis=1, keepdims=True)

    # soft mask: 1 well above the floor, floor_gain at/below it
    floor_gain = 10 ** (-reduction_db / 20)
    ratio = mag / (noise_profile * 2.0 + 1e-10)
    mask = np.clip((ratio - 1.0) / 2.0, 0.0, 1.0)
    mask = floor_gain + (1.0 - floor_gain) * mask

    out = librosa.istft(spec * mask, n_fft=N_FFT, hop_length=HOP,
                        length=len(y))
    return out.astype(np.float32)


def enhance_audio(y, sr, normalize=True):
    """
    high-pass + spectral gate + peak normalize. accepts 1-D (mono) or
    2-D (channels, samples) float32; returns the same shape and length.
    gentle on clean input — a quiet clip just gets lifted to level.
    """
    orig_ndim = np.asarray(y).ndim
    y = np.atleast_2d(np.asarray(y, dtype=np.float32))
    out = _highpass(y, sr)
    out = np.stack([_gate_channel(ch, sr) for ch in out])
    if normalize:
        peak = float(np.abs(out).max())
        if peak > 1e-6:
            out = out * (0.891 / peak)  # -1 dBFS
    out = out.astype(np.float32)
    return out if orig_ndim == 2 else out[0]


def enhance_to_tempfile(path):
    """
    enhanced copy of `path` written as wav, for the analysis + musicgen
    prompt feed. returns the new path; falls back to the original path if
    anything goes sideways — enhancement must never block the pipeline.
    """
    try:
        y, sr = librosa.load(path, sr=None, mono=False)
        cleaned = enhance_audio(y, sr)
        out_path = os.path.join(tempfile.mkdtemp(), "coda_enhanced.wav")
        sf.write(out_path, cleaned.T if cleaned.ndim == 2 else cleaned, sr)
        return out_path
    except Exception as e:
        print(f"[coda] enhancement failed ({e}); using raw input", flush=True)
        return path


def input_quality(path):
    """
    judge the recording (not the song): effective bandwidth and noise
    floor. returns {"bandwidth_hz", "noise_db", "lofi"} or None on failure.
    """
    try:
        y, sr = librosa.load(path, sr=None, mono=True)
        if len(y) < sr:
            return None
        mag = np.abs(librosa.stft(y, n_fft=4096))
        med = np.median(mag, axis=1)
        ref = float(med.max())
        if ref <= 0:
            return None
        above = np.where(med > ref * 10 ** (-55 / 20))[0]
        freqs = librosa.fft_frequencies(sr=sr, n_fft=4096)
        bandwidth = float(freqs[above[-1]]) if len(above) else 0.0

        frames = librosa.util.frame(y, frame_length=2048, hop_length=1024)
        frame_rms = np.sqrt(np.mean(frames ** 2, axis=0))
        noise_db = float(20 * np.log10(np.percentile(frame_rms, 10) + 1e-12))

        return {
            "bandwidth_hz": round(bandwidth),
            "noise_db": round(noise_db, 1),
            "lofi": bandwidth < LOFI_BANDWIDTH_HZ or noise_db > LOFI_NOISE_DB,
        }
    except Exception as e:
        print(f"[coda] quality probe failed ({e})", flush=True)
        return None