app.py

"""AudioSR — wzbogacanie krótkich próbek głosu (HF Space, CPU/DARMOWY). UI Lektor AI.
Wejście: gr.Audio (upload) → BEZPOŚREDNIO do enhance (bez State — pewne okablowanie).
Zamiast pustej nuty: WYKRES WIDMA (wierzchołki/góra pasma) jako <img>. Pasek postępu + ETA.
CPU: 1 klip ≈ ~2 min. Model „basic", domyślnie 24 kroki.

Tryb NO-FIDDLE: backend jest w pełni ustawiony (model='basic', guidance=3.5, seed=42,
latent_t_per_second=12.8). UŻYTKOWNIK steruje TYLKO dwoma rzeczami:
  • suwak DDIM steps (jakość vs czas),
  • przełącznik „Normalizuj" (włącza dwustopniową normalizację LUFS → sufit szczytu).
Całe czyszczenie + normalizacja (force-mono, DC, high-pass 75 Hz, LUFS -18, sufit -1 dB)
żyje w postproc.py (czysty DSP, testowalny lokalnie). app.py tylko orkiestruje."""
import os
import time
import threading
import tempfile
import numpy as np
import soundfile as sf
import gradio as gr
import torch
from audiosr import build_model, super_resolution
import viz
import postproc

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
OUTPUT_SR = 48000

# ── BACKEND HARDCODE (user tego NIE rusza) ───────────────────────────────────
MODEL_NAME = "basic"          # „basic" — sprawdzony do tego głosu (vs „speech")
GUIDANCE_SCALE = 3.5          # classifier-free guidance — stały punkt pracy
SEED = 42                     # stały seed → powtarzalny wynik dla tej samej próbki
LATENT_T_PER_SECOND = 12.8    # rozdzielczość czasowa latentu AudioSR (zalecana)

CYAN, ORANGE, MUTED = "#00ffff", "#ff9900", "#8a8aa0"
HINT = ("<div style='color:#8a8aa0;font-size:12px;text-align:center;padding:10px'>"
        "Wrzuć próbkę, aby zobaczyć widmo (wierzchołki / góra pasma)</div>")
_CACHE = {}

# ── Hook REALNEGO postępu AudioSR ─────────────────────────────────────────────
# AudioSR liczy dyfuzję w wątku; jedyny wiarygodny sygnał postępu to jego własny
# pasek tqdm „DDIM Sampler". Podpinamy się pod tqdm.update() i czytamy n/total
# TYLKO dla tego paska → UI pokazuje FAKTYCZNY krok zamiast zgadywanej ETA.
# (Na CPU krok ≈ 25 s, więc estymata czasowa była bezużyteczna i pasek „stał" na 95%.)
_SR_PROGRESS = {"frac": 0.0, "step": 0, "total": 0}
try:
    from tqdm import std as _tqdm_std
    _orig_tqdm_update = _tqdm_std.tqdm.update

    def _tqdm_update_hook(self, n=1):
        r = _orig_tqdm_update(self, n)
        try:
            if "DDIM" in (getattr(self, "desc", "") or "") and self.total:
                _SR_PROGRESS["step"] = int(self.n)
                _SR_PROGRESS["total"] = int(self.total)
                _SR_PROGRESS["frac"] = max(0.0, min(0.999, self.n / self.total))
        except Exception:
            pass
        return r

    _tqdm_std.tqdm.update = _tqdm_update_hook
except Exception:  # pragma: no cover - brak tqdm = zostaje fallback czasowy
    pass


def _model(model_name: str):
    if model_name not in _CACHE:
        _CACHE[model_name] = build_model(model_name=model_name, device=DEVICE)
    return _CACHE[model_name]


def on_audio(audio_path):
    """Po wrzuceniu/zmianie pliku → widmo wejścia (lub podpowiedź gdy puste)."""
    if not audio_path:
        return HINT
    try:
        return viz.spectrum_html([(audio_path, "wejście", CYAN)], "Widmo wejścia — wierzchołki (góra pasma)")
    except Exception as e:
        return f"<div style='color:#ff5a5a;font-size:12px;padding:8px'>Nie udało się narysować widma: {e}</div>"


def enhance(audio_path, ddim_steps, norm, progress=gr.Progress()):
    if not audio_path:
        raise gr.Error("Najpierw wrzuć plik audio.")

    progress(0.03, desc="Ładowanie modelu (pierwszy raz: pobieranie ~GB)…")
    mdl = _model(MODEL_NAME)

    res = {}
    def work():
        try:
            res["w"] = super_resolution(mdl, audio_path, seed=SEED,
                                        guidance_scale=GUIDANCE_SCALE,
                                        ddim_steps=int(ddim_steps),
                                        latent_t_per_second=LATENT_T_PER_SECOND)
        except Exception as e:
            res["err"] = e

    # reset stanu postępu przed nowym przebiegiem
    _SR_PROGRESS["frac"] = 0.0; _SR_PROGRESS["step"] = 0; _SR_PROGRESS["total"] = 0
    t = threading.Thread(target=work, daemon=True)
    t.start()
    # fallback czasowy gdyby hook tqdm zawiódł: CPU-basic ≈ 25-28 s/krok
    est = max(60.0, int(ddim_steps) * 28.0)
    t0 = time.time()
    while t.is_alive():
        el = time.time() - t0
        real = _SR_PROGRESS["frac"]
        if real > 0:  # mamy REALNY postęp dyfuzji
            frac = real
            eta = int(el / real * (1.0 - real))
            desc = f"Wzbogacam… krok {_SR_PROGRESS['step']}/{_SR_PROGRESS['total']} · ETA ~{eta} s"
        else:  # przed startem samplera (ładowanie/inicjalizacja) — pełzanie czasowe
            frac = min(0.90, el / est)
            desc = f"Inicjalizacja dyfuzji… ETA ~{max(0, int(est - el))} s"
        progress(min(0.97, 0.05 + 0.92 * frac), desc=desc)
        time.sleep(0.4)
    t.join()
    if "err" in res:
        raise gr.Error(f"Błąd AudioSR: {res['err']}")

    progress(0.97, desc="Składanie wyniku…")
    wav = np.asarray(res["w"]).squeeze()
    if wav.ndim > 1:
        wav = wav.T
    # Przytnij padding AudioSR do długości wejścia (frames * OUTPUT_SR / src_sr).
    # UWAGA: dla mp3 (główne wejście, „spłaszczone mp3") sf.info().frames bywa
    # PRZYBLIŻONE lub 0 na niektórych buildach libsndfile → tnij TYLKO gdy wynik jest
    # WYRAŹNIE dłuższy (>1%), żeby nie ściąć kilku ważnych próbek ogona; przy frames=0
    # warunek `exp` i tak pomija trim.
    info = sf.info(audio_path)
    exp = int(round(info.frames / info.samplerate * OUTPUT_SR))
    if exp and wav.shape[0] > int(exp * 1.01):
        wav = wav[:exp]

    # Czyszczenie + normalizacja → JEDNO wejście postproc (czysty DSP, testowalny).
    # norm=True  : mono → DC → high-pass 75 Hz → LUFS -18 → sufit szczytu -1 dB.
    # norm=False : mono → dopasowanie RMS do oryginału + bezpiecznik 0.99.
    wav, _pp = postproc.process(wav, OUTPUT_SR, audio_path, normalize=bool(norm))

    base = os.path.splitext(os.path.basename(audio_path))[0]
    out_path = os.path.join(tempfile.gettempdir(), f"{base}_48k.wav")
    # Jawny subtype: WAV 16-bit PCM (w spec agy: „16-bit PCM lub 32-bit float").
    # Sufit szczytu -1 dB istnieje WŁAŚNIE dlatego, że kontener obcina na ±1.0 —
    # nie rozluźniaj go zakładając, że plik jest float.
    sf.write(out_path, wav.astype(np.float32), OUTPUT_SR, subtype="PCM_16")

    # Anty-silent-skip: jeśli przy normalizacji cel LUFS -18 NIE został osiągnięty
    # (pojedynczy głośny transjent — plozja/klik — wymusił zejście sufitem szczytu),
    # POWIEDZ to wprost zamiast cicho udawać sukces.
    warn = ""
    short_db = _pp.get("lufs_shortfall_db", 0.0) if isinstance(_pp, dict) else 0.0
    if bool(norm) and short_db > 3.0:
        warn = ("<div style='color:#ff9900;font-size:12px;padding:8px;border:1px solid #ff9900;"
                "border-radius:8px;margin-bottom:8px'>Uwaga: docelowy poziom LUFS -18 nie został "
                f"w pełni osiągnięty (brak ~{short_db:.1f} dB) — pojedynczy głośny transjent "
                "(plozja/klik) wymusił zejście sufitem szczytu -1 dB. Rozważ ręczne wyciszenie "
                "tego transjentu w Audacity i ponów.</div>")

    html = warn + viz.spectrum_html([(audio_path, "wejście", CYAN), (out_path, "wyjście", ORANGE)],
                                    "Widmo: wejście vs wyjście — odbudowane wierzchołki")
    progress(1.0, desc="Gotowe")
    return (gr.Audio(value=out_path, visible=True),
            gr.DownloadButton(value=out_path, visible=True, interactive=True),
            html)


# ── Motyw Lektor AI + kompakt ────────────────────────────────────────────────
PANEL, BG, BG2, TXT, BORDER = "#0a0a0f", "#0d0d14", "#15151e", "#cfd4e0", "#222230"
CSS = f"""
.gradio-container, gradio-app, body {{ background:{PANEL} !important; color:{TXT} !important;
  font-family:'Segoe UI', system-ui, sans-serif !important; }}
.gradio-container {{ max-width:560px !important; margin:0 auto !important; padding:10px !important; }}
.gap, .form {{ gap:8px !important; }}
.block {{ background:{BG} !important; border:1px solid {BORDER} !important; border-radius:10px !important; padding:8px 10px !important; }}
#lk-menu {{ display:flex; align-items:center; gap:10px; padding:10px 12px; margin:0 0 10px 0; border-radius:10px;
  background:linear-gradient(180deg,{BG2},{BG}); border:1px solid {BORDER}; border-bottom:2px solid {CYAN}; box-shadow:0 6px 18px -12px {CYAN}; }}
#lk-menu .lk-title {{ font-size:17px; font-weight:700; color:{TXT}; }}
#lk-menu .lk-title b {{ color:{CYAN}; text-shadow:0 0 8px rgba(0,255,255,.4); }}
#lk-menu .lk-sub {{ font-size:11px; color:{MUTED}; margin-top:2px; }}
label, label span, span[data-testid="block-info"] {{ color:{MUTED} !important; font-size:12px !important; }}
/* UWAGA: NIE klobruj checkboxa tłem/borderem — gradio rysuje „ptaszek" przez tło
   przy :checked, a wymuszone background+border czyni go NIEWIDOCZNYM (kwadrat pusty
   mimo zaznaczenia). Dlatego wykluczamy [type=checkbox] z reguły poniżej i renderujemy
   go natywnie z cyan akcentem osobno. */
input:not([type=checkbox]), textarea, select {{ background:{BG2} !important; color:{TXT} !important; border:1px solid {BORDER} !important; }}
input[type=range] {{ accent-color:{CYAN} !important; }}
input[type=checkbox] {{ -webkit-appearance:auto !important; appearance:auto !important;
  accent-color:{CYAN} !important; background:initial !important; border:initial !important;
  width:18px !important; height:18px !important; min-width:18px !important; cursor:pointer; flex:none !important; }}
input[type=checkbox] + label, label:has(> input[type=checkbox]) {{ cursor:pointer; }}
/* wejście audio: niższe, bez wielkiej ikony placeholdera */
#ain .empty, #ain .icon {{ min-height:unset !important; }}
#ain .empty svg, #ain .icon svg {{ width:26px !important; height:26px !important; }}
#ain .wrap {{ min-height:90px !important; }}
#ain audio {{ height:40px !important; }}
#go, #dl {{ width:auto !important; min-width:180px; max-width:280px; margin:6px auto !important; display:block !important; padding:9px 18px !important; }}
#go {{ background:{CYAN} !important; color:#06141c !important; border:none !important; font-weight:700 !important; letter-spacing:.5px !important; }}
#go:hover {{ box-shadow:0 0 14px rgba(0,255,255,.5) !important; }}
#dl {{ background:transparent !important; color:{CYAN} !important; border:1px solid {CYAN} !important; }}
#aout audio {{ height:40px !important; }}
footer {{ display:none !important; }}
"""

ICON = (f"<svg width='26' height='26' viewBox='0 0 24 24' fill='none' stroke='{CYAN}' "
        f"stroke-width='2.2' stroke-linecap='round'>"
        f"<line x1='4' y1='15' x2='4' y2='9'/><line x1='9' y1='19' x2='9' y2='5'/>"
        f"<line x1='14' y1='16' x2='14' y2='8'/><line x1='19' y1='13' x2='19' y2='11'/></svg>")
MENU = (f"<div id='lk-menu'>{ICON}<div>"
        f"<div class='lk-title'>Audio<b>SR</b> · Lektor AI</div>"
        f"<div class='lk-sub'>Wzbogacanie próbek głosu → 48 kHz</div></div></div>")

with gr.Blocks(title="AudioSR · Lektor AI", css=CSS, theme=gr.themes.Base()) as demo:
    gr.HTML(MENU)
    audio = gr.Audio(sources=["upload"], type="filepath", elem_id="ain", label="Próbka — wrzuć plik (krótki klip)")
    plot = gr.HTML(HINT, elem_id="plot")
    # TYLKO dwie kontrolki: DDIM steps + przełącznik normalizacji. Reszta = backend.
    steps = gr.Slider(10, 100, value=24, step=1, label="DDIM steps (więcej = lepiej / wolniej)")
    norm = gr.Checkbox(value=True, label="Normalizuj (Zalecane)")
    btn = gr.Button("WZBOGAĆ", variant="primary", elem_id="go")
    out = gr.Audio(type="filepath", label="Wynik — 48 kHz", interactive=False, elem_id="aout", visible=False)
    dl = gr.DownloadButton("Pobierz WAV", elem_id="dl", interactive=False, visible=False)

    audio.change(on_audio, audio, plot)
    btn.click(enhance, [audio, steps, norm], [out, dl, plot])

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