""" 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()