app.py

"""
UNlossifier -- Lossy Audio Reconstructor & Sound Signature Simulator (Space demo)

Original project: aston89/UNlossifier-lossy-audio-reconstructor-and-sound-signature-simulator
https://github.com/aston89/UNlossifier-lossy-audio-reconstructor-and-sound-signature-simulator
(GPL-3.0)

This app is self-contained: the model (StereoUNet) and the v5 context-aware
overlap-add inference routine are reproduced directly from the author's
UNlossifier_(v5).py, since the network is small enough to embed rather than
depend on cloning the whole training repo. Only inference is exposed here --
training (the project's "core feature" for building custom sound-signature
models) is CLI-only in the original repo; see README_SETUP.md for why that's
left out of this Space and how to add it if you want it.

Five pretrained checkpoints ship in the source repo, each for a specific
codec / bitrate / sample-rate combo (the sample rate is NOT optional -- you
must match the checkpoint that was trained at that rate):

    model_mp3_64k_44100_epoch997.safetensors   -> mp3 64kbps,  sr=44100
    model_mp3_96k_32000_epoch393.safetensors   -> mp3 96kbps,  sr=32000
    model_mp3_128k_44100_epoch397.safetensors  -> mp3 128kbps, sr=44100 (v2)
    model_aac_128k_44100_epoch998.safetensors  -> aac 128kbps, sr=44100
    model_mp3_128k_44100_epoch500.safetensors  -> mp3 128kbps, sr=44100 (v5, noise-dataset trained)
"""

import os
from pathlib import Path

import gradio as gr
import numpy as np
import librosa
import soundfile as sf
import torch
import torch.nn as nn
import safetensors.torch as sf_torch
import requests

try:
    import spaces  # ZeroGPU support
    GPU_DECORATOR = spaces.GPU
except Exception:
    def GPU_DECORATOR(fn):
        return fn

torch.set_float32_matmul_precision("high")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

RAW_BASE = (
    "https://raw.githubusercontent.com/aston89/"
    "UNlossifier-lossy-audio-reconstructor-and-sound-signature-simulator/main"
)
CKPT_DIR = Path("checkpoints")
CKPT_DIR.mkdir(exist_ok=True)

# filename -> (display label, sample rate the model was trained/must run at)
MODELS = {
    "model_mp3_64k_44100_epoch997.safetensors": ("MP3 64 kbps (44.1 kHz) -- heaviest restoration", 44100),
    "model_mp3_96k_32000_epoch393.safetensors": ("MP3 96 kbps (32 kHz)", 32000),
    "model_mp3_128k_44100_epoch397.safetensors": ("MP3 128 kbps (44.1 kHz, v2)", 44100),
    "model_aac_128k_44100_epoch998.safetensors": ("AAC 128 kbps (44.1 kHz, e.g. YouTube rips)", 44100),
    "model_mp3_128k_44100_epoch500.safetensors": ("MP3 128 kbps (44.1 kHz, v5 noise-trained)", 44100),
}

SEG_LEN_SEC = 4
CTX_RATIO = 0.25


# =========================================================
# MODEL (verbatim from UNlossifier_(v5).py)
# =========================================================
class ConvBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv1d(in_ch, out_ch, 7, padding=3),
            nn.GroupNorm(8, out_ch),
            nn.GELU(),
            nn.Conv1d(out_ch, out_ch, 5, padding=2),
            nn.GroupNorm(8, out_ch),
            nn.GELU(),
        )

    def forward(self, x):
        return self.net(x)


class StereoUNet(nn.Module):
    def __init__(self, base=128):
        super().__init__()
        self.enc1 = ConvBlock(4, base)
        self.enc2 = ConvBlock(base, base)
        self.mid = ConvBlock(base, base)
        self.dec2 = ConvBlock(base * 2, base)
        self.dec1 = ConvBlock(base * 2, base)
        self.out = nn.Conv1d(base, 4, 7, padding=3)

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(e1)
        x = self.mid(e2)
        x = torch.cat([x, e2], dim=1)
        x = self.dec2(x)
        x = torch.cat([x, e1], dim=1)
        x = self.dec1(x)
        return self.out(x)


def to_ms(x):
    if isinstance(x, torch.Tensor):
        x = x.detach().cpu().numpy()
    L, R = x[0], x[1]
    M = 0.5 * (L + R)
    S = 0.5 * (L - R)
    return np.stack([L, R, M, S], axis=0).astype(np.float32)


def to_torch(x, device):
    return torch.from_numpy(x).float().to(device)


# =========================================================
# CHECKPOINT MANAGEMENT
# =========================================================
_MODEL_CACHE = {}


def ensure_checkpoint(filename):
    local_path = CKPT_DIR / filename
    if local_path.exists():
        return local_path
    url = f"{RAW_BASE}/{filename}"
    resp = requests.get(url, timeout=120)
    resp.raise_for_status()
    local_path.write_bytes(resp.content)
    return local_path


def load_model(filename):
    if filename in _MODEL_CACHE:
        return _MODEL_CACHE[filename]
    path = ensure_checkpoint(filename)
    model = StereoUNet().to(DEVICE)
    sf_torch.load_model(model, str(path))
    model.eval()
    _MODEL_CACHE[filename] = model
    return model


# =========================================================
# INFERENCE (verbatim logic from UNlossifier_(v5).py's inference())
# =========================================================
@torch.no_grad()
def restore(model, audio, sr):
    total = audio.shape[1]
    seg_len = SEG_LEN_SEC * sr
    ctx = int(seg_len * CTX_RATIO)
    chunk = seg_len
    expected_in = chunk + 2 * ctx
    step = max(1, chunk - 2 * ctx)

    padded = np.pad(audio, ((0, 0), (ctx, ctx)), mode="edge")
    out = np.zeros((2, total), dtype=np.float32)
    w = np.zeros((2, total), dtype=np.float32)
    window = np.hanning(chunk).astype(np.float32)
    eps = 1e-8
    w_lr = 0.50
    w_ms = 0.50

    for i in range(0, total, step):
        x = padded[:, i:i + expected_in]
        if x.shape[1] < expected_in:
            pad = expected_in - x.shape[1]
            x = np.pad(x, ((0, 0), (0, pad)), mode="edge")
        x_t = to_torch(to_ms(x), DEVICE).unsqueeze(0)
        y = model(x_t).squeeze(0).cpu().numpy().astype(np.float32)
        L1, R1 = y[0], y[1]
        M, S = y[2], y[3]
        L2 = M + S
        R2 = M - S
        L = w_lr * L1 + w_ms * L2
        R = w_lr * R1 + w_ms * R2
        stereo = np.stack([L, R], axis=0)
        stereo = stereo[:, ctx:ctx + chunk]
        valid = min(chunk, total - i)
        win = window[:valid]
        out[:, i:i + valid] += stereo[:, :valid] * win
        w[:, i:i + valid] += win

    out = out / np.clip(w, eps, None)
    out = np.nan_to_num(out)
    out = np.clip(out, -1.0, 1.0)
    return out


@GPU_DECORATOR
def run_restoration(input_path, model_choice, progress=gr.Progress()):
    if input_path is None:
        raise gr.Error("Please upload an audio file first.")

    filename = [k for k, v in MODELS.items() if v[0] == model_choice][0]
    _, sr = MODELS[filename]

    progress(0.1, desc="Loading model...")
    model = load_model(filename)

    progress(0.3, desc=f"Decoding audio at {sr} Hz...")
    y, _ = librosa.load(input_path, sr=sr, mono=False)
    if y.ndim == 1:
        y = np.stack([y, y], axis=0)
    elif y.shape[0] > 2:
        y = y[:2]

    progress(0.5, desc="Running UNlossifier restoration...")
    restored = restore(model, y.astype(np.float32), sr)

    out_path = "restored_output.wav"
    sf.write(out_path, restored.T, sr, subtype="FLOAT")
    progress(1.0, desc="Done.")
    return out_path


with gr.Blocks(title="UNlossifier -- Lossy Audio Reconstructor") as demo:
    gr.Markdown(
        "# UNlossifier\n"
        "**Lossy audio reconstructor & sound signature simulator** by "
        "[aston89](https://github.com/aston89/UNlossifier-lossy-audio-reconstructor-and-sound-signature-simulator) "
        "(GPL-3.0). Upload a compressed track and pick the model matching how "
        "it was likely encoded -- it doesn't recover the exact original signal "
        "(that information is permanently gone), it generates a plausible, "
        "spectrally-richer reconstruction.\n\n"
        "*Tip: pick the checkpoint that matches your file's original codec/bitrate "
        "as closely as possible for best results.*"
    )
    with gr.Row():
        with gr.Column():
            audio_in = gr.Audio(type="filepath", label="Upload lossy audio")
            model_choice = gr.Dropdown(
                choices=[v[0] for v in MODELS.values()],
                value=list(MODELS.values())[0][0],
                label="Restoration model",
            )
            run_btn = gr.Button("Restore audio", variant="primary")
        with gr.Column():
            audio_out = gr.Audio(label="Restored output", type="filepath")

    run_btn.click(fn=run_restoration, inputs=[audio_in, model_choice], outputs=audio_out)

    gr.Markdown(
        "---\n"
        "Note: this Space only exposes **inference**. The project's other core "
        "feature -- training your own custom restoration or sound-signature "
        "model from a handful of paired WAV files -- is CLI-only in the "
        "original repo and isn't included here. See `README_SETUP.md`."
    )

if __name__ == "__main__":
    demo.queue().launch()