app.py

# 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("""
    <div class="app-header">
        <h1>🎧 Audio Super-Resolution</h1>
        <p style="text-align: center;">
            Restaura las altas frecuencias de grabaciones de baja calidad
            utilizando un modelo <strong>Attention Res-UNet 2D</strong> entrenado
            mediante representaciones STFT. Sube tu audio y obtén calidad
            <strong>44.1 kHz</strong> en segundos.
        </p>
        <div class="tech-badges">
            <span class="tech-badge">PyTorch</span>
            <span class="tech-badge">STFT / iSTFT</span>
            <span class="tech-badge">Attention Gates</span>
            <span class="tech-badge">Residual UNet</span>
            <span class="tech-badge">44.1 kHz</span>
        </div>
    </div>
    """)

    # ── 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("#### <span class='section-label'>① Entrada de Audio</span>")
            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("#### <span class='section-label'>② Audio Restaurado</span>")
            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("""
    <div class="app-footer">
        <strong>Proyecto:</strong> Sistema de Deep Learning para la extensión del ancho de banda basado en representaciones STFT<br>
        <strong>Tecnología:</strong> STFT · Attention Res-UNet 2D · PyTorch<br>
        Trabajo de Fin de Grado — Álvaro Roca Nacarino ·
        <a href="https://huggingface.co/devAlvaro26/Unet2D_SuperRes" target="_blank">Hugging Face Spaces</a>
    </div>
    """)

if __name__ == "__main__":
    demo.launch(theme=custom_theme, css=custom_css)