# App de Gradio para despliegue en Hugging Face Spaces
# Super-resolución de audio con Attention Res-UNet 2D
import os
import torch
import torchaudio
import gradio as gr
import tempfile
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from model import UNetAudio2D
# ── Constantes ──────────────────────────────────────────────
MODEL_PATH = "unet2D_superres.pt"
TARGET_SR = 44100
POOL_FACTOR = 16
N_FFT = 2048
HOP_LENGTH = 512
FRAGMENT_LENGTH = 65536
# ── Carga del modelo ────────────────────────────────────────
device = torch.device("cpu")
model = UNetAudio2D().to(device)
model.load_state_dict(torch.load(MODEL_PATH, map_location=device, weights_only=True))
model.eval()
print("✅ Modelo cargado correctamente.")
# ═══════════════════════════════════════════════════════════
# Funciones de procesamiento (pipeline idéntico a inference.py)
# ═══════════════════════════════════════════════════════════
def waveform_to_stft(waveform, n_fft=N_FFT, hop_length=HOP_LENGTH):
"""Convierte una forma de onda a un STFT con canales real e imaginario."""
if waveform.ndim == 2:
waveform = waveform.squeeze(0)
window = torch.hann_window(n_fft, device='cpu')
stft = torch.stft(
waveform, n_fft=n_fft, hop_length=hop_length,
win_length=n_fft, window=window, return_complex=True,
)
return torch.stack([stft.real, stft.imag], dim=0)
def stft_to_waveform(stft, n_fft=N_FFT, hop_length=HOP_LENGTH, length=None):
"""Convierte un STFT con canales real e imaginario a forma de onda."""
stft_complex = torch.complex(stft[0], stft[1])
window = torch.hann_window(n_fft)
waveform = torch.istft(
stft_complex, n_fft=n_fft, hop_length=hop_length,
win_length=n_fft, window=window, length=length,
)
return waveform.unsqueeze(0)
def normalize_stft(stft):
"""Log-compresión de la magnitud STFT preservando la fase."""
real, imag = stft[0], stft[1]
magnitude = torch.sqrt(real**2 + imag**2 + 1e-8)
phase_cos = real / magnitude
phase_sin = imag / magnitude
mag_compressed = torch.log1p(magnitude)
return torch.stack([mag_compressed * phase_cos, mag_compressed * phase_sin], dim=0)
def denormalize_stft(stft):
"""Inversa de la log-compresión de la STFT (expm1 = exp(x)-1)."""
real, imag = stft[0], stft[1]
mag_compressed = torch.sqrt(real**2 + imag**2 + 1e-8)
phase_cos = real / mag_compressed
phase_sin = imag / mag_compressed
magnitude = torch.expm1(mag_compressed)
return torch.stack([magnitude * phase_cos, magnitude * phase_sin], dim=0)
def pad_stft(stft, pool_factor=POOL_FACTOR):
"""Padding para que las dimensiones sean divisibles por pool_factor."""
_, freq_bins, time_frames = stft.shape
pad_f = (pool_factor - (freq_bins % pool_factor)) % pool_factor
pad_t = (pool_factor - (time_frames % pool_factor)) % pool_factor
if pad_f > 0 or pad_t > 0:
stft = torch.nn.functional.pad(stft, (0, pad_t, 0, pad_f), mode='reflect')
return stft, freq_bins, time_frames
def process_audio_in_chunks(mdl, stft, orig_f, orig_t, chunk_frames, overlap=64):
"""Procesa el STFT por chunks con overlap para evitar artefactos."""
_, F, T = stft.shape
hop = chunk_frames - overlap
output = torch.zeros_like(stft)
weights = torch.zeros(T, device=stft.device)
window = torch.ones(chunk_frames, device=stft.device)
if overlap > 0:
window[:overlap] = torch.linspace(0, 1, overlap, device=stft.device)
window[-overlap:] = torch.linspace(1, 0, overlap, device=stft.device)
start = 0
while start < T:
end = min(start + chunk_frames, T)
chunk = stft[:, :, start:end]
pad_t = chunk_frames - chunk.shape[-1]
if pad_t > 0:
pad_mode = 'reflect' if pad_t < chunk.shape[-1] else 'replicate'
chunk = torch.nn.functional.pad(chunk, (0, pad_t), mode=pad_mode)
with torch.no_grad():
pred_chunk = mdl(chunk.unsqueeze(0)).squeeze(0)
actual_len = end - start
w = window[:actual_len].clone()
if start == 0 and overlap > 0:
w[:overlap] = 1.0
if end == T and overlap > 0:
w[-overlap:] = 1.0
output[:, :, start:end] += pred_chunk[:, :, :actual_len] * w
weights[start:end] += w
start += hop
output = output / weights.unsqueeze(0).unsqueeze(0)
return output
# ── Generación de espectrogramas ────────────────────────────
def generate_spectrogram(waveform, sample_rate, n_fft=N_FFT, hop_length=HOP_LENGTH):
"""Genera una imagen de espectrograma con estilo premium oscuro."""
spec_transform = torchaudio.transforms.Spectrogram(
n_fft=n_fft, hop_length=hop_length, power=2
)
spec = spec_transform(waveform.cpu()).squeeze().numpy()
spec_db = 10 * np.log10(spec + 1e-10)
nyquist = sample_rate / 2
duration = (spec.shape[1] * hop_length) / sample_rate
extent = [0, duration, 0, nyquist / 1000] # kHz
fig, ax = plt.subplots(figsize=(10, 3.5))
fig.patch.set_facecolor('#0c0f1a')
ax.set_facecolor('#0c0f1a')
im = ax.imshow(
spec_db, origin='lower', aspect='auto', cmap='magma',
extent=extent, vmin=-80, vmax=spec_db.max(),
)
ax.set_ylabel("Frecuencia (kHz)", color='#94a3b8', fontsize=10)
ax.set_xlabel("Tiempo (s)", color='#94a3b8', fontsize=10)
ax.tick_params(colors='#64748b', labelsize=8)
for spine in ax.spines.values():
spine.set_color('#1e293b')
cbar = fig.colorbar(im, ax=ax, pad=0.02)
cbar.set_label("dB", color='#94a3b8', fontsize=9)
cbar.ax.tick_params(colors='#64748b', labelsize=8)
plt.tight_layout(pad=0.5)
tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
plt.savefig(tmp.name, dpi=150, facecolor='#0c0f1a', edgecolor='none')
plt.close(fig)
return tmp.name
# ── Pipeline de inferencia ──────────────────────────────────
def inference(audio_path):
"""Ejecuta la super-resolución sobre el audio de entrada."""
if audio_path is None:
raise gr.Error("Por favor, sube un archivo de audio.")
waveform, original_sr = torchaudio.load(audio_path)
# Resampleo a TARGET_SR
if original_sr != TARGET_SR:
resampler = torchaudio.transforms.Resample(original_sr, TARGET_SR)
waveform = resampler(waveform)
# Mono
if waveform.size(0) > 1:
waveform = waveform.mean(dim=0, keepdim=True)
# Normalización por amplitud máxima
max_val = waveform.abs().max().item() + 1e-8
waveform_norm = waveform / max_val
original_length = waveform_norm.size(1)
input_for_plot = waveform_norm.clone()
# STFT → Log-compresión → Padding
stft_input = waveform_to_stft(waveform_norm)
stft_input = normalize_stft(stft_input)
stft_padded, orig_f, orig_t = pad_stft(stft_input)
chunk_frames = FRAGMENT_LENGTH // HOP_LENGTH
chunk_frames = chunk_frames + (POOL_FACTOR - (chunk_frames % POOL_FACTOR)) % POOL_FACTOR
# Inferencia
predicted_stft = process_audio_in_chunks(
model, stft_padded, orig_f, orig_t, chunk_frames=chunk_frames
)
predicted_stft = predicted_stft[:, :orig_f, :orig_t]
predicted_stft = denormalize_stft(predicted_stft)
# ISTFT
predicted_waveform = stft_to_waveform(predicted_stft, length=original_length)
predicted_waveform = torch.nan_to_num(predicted_waveform, nan=0.0, posinf=1.0, neginf=-1.0)
predicted_waveform = torch.clamp(predicted_waveform, -1.0, 1.0)
# Guardar WAV
tmp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
torchaudio.save(tmp_wav.name, predicted_waveform, TARGET_SR, bits_per_sample=16, encoding="PCM_S")
# Espectrogramas
spec_in = generate_spectrogram(input_for_plot, TARGET_SR)
spec_out = generate_spectrogram(predicted_waveform, TARGET_SR)
return tmp_wav.name, spec_in, spec_out
# ═══════════════════════════════════════════════════════════
# DISEÑO GRADIO
# ═══════════════════════════════════════════════════════════
custom_theme = gr.themes.Base(
primary_hue=gr.themes.colors.indigo,
secondary_hue=gr.themes.colors.violet,
neutral_hue=gr.themes.colors.slate,
font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "monospace"],
).set(
body_background_fill="#0b0f1a",
body_text_color="#e2e8f0",
block_background_fill="rgba(15, 23, 42, 0.6)",
block_border_color="rgba(99, 102, 241, 0.15)",
block_border_width="1px",
block_radius="16px",
block_label_text_color="#a5b4fc",
block_title_text_color="#c7d2fe",
input_background_fill="rgba(15, 23, 42, 0.8)",
input_border_color="rgba(99, 102, 241, 0.2)",
input_border_width="1px",
button_primary_background_fill="linear-gradient(135deg, #6366f1 0%, #8b5cf6 50%, #a78bfa 100%)",
button_primary_background_fill_hover="linear-gradient(135deg, #4f46e5 0%, #7c3aed 50%, #8b5cf6 100%)",
button_primary_text_color="#ffffff",
button_secondary_background_fill="rgba(99, 102, 241, 0.08)",
button_secondary_border_color="rgba(99, 102, 241, 0.25)",
button_secondary_text_color="#c7d2fe",
shadow_drop="0 4px 24px rgba(0,0,0,0.3)",
)
custom_css = """
/* ── Base ─────────────────────────────────────────── */
body, .gradio-container {
background: #0b0f1a !important;
background-image:
radial-gradient(ellipse 80% 60% at 50% -10%, rgba(99,102,241,0.12) 0%, transparent 60%),
radial-gradient(ellipse 60% 50% at 80% 50%, rgba(139,92,246,0.06) 0%, transparent 50%) !important;
min-height: 100vh;
}
/* ── Header ───────────────────────────────────────── */
.app-header {
text-align: center;
padding: 2.5rem 1rem 1rem;
animation: fadeDown 0.7s ease-out;
}
.app-header h1 {
font-size: 2.8rem !important;
font-weight: 800 !important;
background: linear-gradient(135deg, #818cf8 0%, #c084fc 50%, #f0abfc 100%);
-webkit-background-clip: text;
-webkit-text-fill-color: transparent;
background-clip: text;
letter-spacing: -0.03em;
margin-bottom: 0.25rem !important;
line-height: 1.2;
}
.app-subtitle {
color: #94a3b8 !important;
font-size: 1.05rem;
max-width: 640px;
margin: 0.5rem auto 0;
line-height: 1.65;
text-align: center;
animation: fadeIn 0.9s ease-out 0.25s both;
}
/* ── Badges ───────────────────────────────────────── */
.tech-badges {
display: flex;
justify-content: center;
gap: 0.5rem;
margin-top: 1rem;
flex-wrap: wrap;
animation: fadeIn 1s ease-out 0.4s both;
}
.tech-badge {
background: rgba(99,102,241,0.1);
border: 1px solid rgba(99,102,241,0.2);
color: #a5b4fc;
padding: 0.3rem 0.85rem;
border-radius: 999px;
font-size: 0.78rem;
font-weight: 500;
letter-spacing: 0.02em;
}
/* ── Glass panels ─────────────────────────────────── */
.panel-glass {
background: rgba(15, 23, 42, 0.55) !important;
backdrop-filter: blur(16px) saturate(1.2) !important;
-webkit-backdrop-filter: blur(16px) saturate(1.2) !important;
border: 1px solid rgba(99,102,241,0.12) !important;
border-radius: 20px !important;
padding: 1.5rem !important;
transition: border-color 0.35s ease, box-shadow 0.35s ease;
}
.panel-glass:hover {
border-color: rgba(99,102,241,0.25) !important;
box-shadow: 0 8px 32px rgba(99,102,241,0.08) !important;
}
.section-label {
font-size: 0.8rem !important;
font-weight: 600 !important;
text-transform: uppercase !important;
letter-spacing: 0.08em !important;
color: #818cf8 !important;
margin-bottom: 0.75rem !important;
}
/* ── Primary button ───────────────────────────────── */
.run-btn {
margin-top: 0.75rem !important;
font-weight: 700 !important;
letter-spacing: 0.04em !important;
text-transform: uppercase !important;
border-radius: 12px !important;
transition: all 0.3s cubic-bezier(0.4,0,0.2,1) !important;
box-shadow: 0 4px 20px rgba(99,102,241,0.25) !important;
}
.run-btn:hover {
transform: translateY(-2px) scale(1.01) !important;
box-shadow: 0 8px 30px rgba(99,102,241,0.4) !important;
}
/* ── Spectrogram section ──────────────────────────── */
.spec-section {
background: rgba(15, 23, 42, 0.45) !important;
backdrop-filter: blur(12px) !important;
border: 1px solid rgba(99,102,241,0.1) !important;
border-radius: 20px !important;
padding: 1.25rem !important;
margin-top: 1rem !important;
}
.spec-label {
text-align: center;
font-weight: 600;
color: #c7d2fe !important;
font-size: 0.85rem;
text-transform: uppercase;
letter-spacing: 0.06em;
margin-bottom: 0.5rem;
}
/* ── Footer ───────────────────────────────────────── */
.app-footer {
text-align: center;
color: #475569 !important;
font-size: 0.82rem;
margin-top: 2.5rem;
padding: 1.5rem 1rem;
border-top: 1px solid rgba(99,102,241,0.08);
line-height: 1.7;
}
.app-footer strong { color: #64748b; }
.app-footer a { color: #818cf8; text-decoration: none; }
.app-footer a:hover { text-decoration: underline; }
/* ── Animations ───────────────────────────────────── */
@keyframes fadeDown {
from { opacity: 0; transform: translateY(-18px); }
to { opacity: 1; transform: translateY(0); }
}
@keyframes fadeIn {
from { opacity: 0; }
to { opacity: 1; }
}
/* ── Download button ──────────────────────────────── */
.dl-btn {
border-radius: 10px !important;
margin-top: 0.5rem !important;
}
/* ── Panel alignment: audio players at same height ── */
#left-panel, #right-panel {
display: flex !important;
flex-direction: column !important;
justify-content: flex-start !important;
align-items: stretch !important;
}
#left-panel > div, #right-panel > div {
flex: 0 0 auto !important;
}
#audio-input, #audio-output {
margin-top: 0 !important;
}
/* ── Responsive ───────────────────────────────────── */
@media (max-width: 768px) {
.app-header h1 { font-size: 2rem !important; }
.app-subtitle { font-size: 0.95rem; }
}
"""
# ── Construcción de la interfaz ─────────────────────────────
with gr.Blocks(
title="Audio Super-Resolution · Attention Res-UNet 2D",
fill_width=True,
) as demo:
# Header
gr.HTML("""
""")
# ── Sección principal: entrada / salida ──
with gr.Row(equal_height=True):
# Panel izquierdo: entrada
with gr.Column(scale=1, elem_classes="panel-glass", elem_id="left-panel"):
gr.Markdown("#### ① Entrada de Audio")
audio_input = gr.Audio(
label="Sube o graba tu audio",
type="filepath",
sources=["upload", "microphone"],
elem_id="audio-input",
)
btn = gr.Button(
"⚡ Procesar Super-Resolución",
variant="primary",
size="lg",
elem_classes="run-btn",
elem_id="run-btn",
)
# Panel derecho: salida
with gr.Column(scale=1, elem_classes="panel-glass", elem_id="right-panel"):
gr.Markdown("#### ② Audio Restaurado")
audio_output = gr.Audio(
label="Resultado — Alta Resolución",
type="filepath",
interactive=False,
elem_id="audio-output",
)
download_btn = gr.DownloadButton(
label="⬇ Descargar WAV",
visible=False,
variant="secondary",
elem_classes="dl-btn",
elem_id="download-btn",
)
# ── Espectrogramas lado a lado ──
with gr.Row(equal_height=True, elem_classes="spec-section"):
with gr.Accordion("📉 Espectrograma · Entrada", open=False):
spec_input = gr.Image(
type="filepath",
interactive=False,
show_label=False,
elem_id="spec-input",
)
with gr.Accordion("📈 Espectrograma · Super-Resolución", open=False):
spec_output = gr.Image(
type="filepath",
interactive=False,
show_label=False,
elem_id="spec-output",
)
# ── Lógica de eventos ──
def run_pipeline(audio_path):
wav_path, spec_in_path, spec_out_path = inference(audio_path)
return (
wav_path,
spec_in_path,
spec_out_path,
gr.DownloadButton(visible=True, value=wav_path),
)
btn.click(
fn=run_pipeline,
inputs=[audio_input],
outputs=[audio_output, spec_input, spec_output, download_btn],
)
# Footer
gr.HTML("""
""")
if __name__ == "__main__":
demo.launch(theme=custom_theme, css=custom_css)