# app.py
import os, time, torch, gradio as gr, spaces
from transformers import pipeline

# ── Аптымізацыя CPU ──────────────────────────────────────────────
num_cpu_cores = os.cpu_count() or 1
torch.set_num_threads(num_cpu_cores)
print(f"✅ PyTorch настроены на {num_cpu_cores} ядраў CPU.")

# ── Пайплайн ─────────────────────────────────────────────────────
pipe = pipeline(
    "audio-classification",
    model="MIT/ast-finetuned-audioset-10-10-0.448"
)

_model_on_gpu = False   # каб перанесці мадэль на GPU толькі адзін раз

@spaces.GPU(duration=90)          # ZeroGPU выдае GPU на час працы функцыі
def classify_audio(audio_path: str):
    """
    Вяртае:
      1) dict label→score (топ-3),
      2) радок з часам інферэнсу ў секундах.
    """
    global _model_on_gpu

    if audio_path is None:
        return {"⚠️": "Загрузіце файл"}, "—"

    # адзін раз пераносім мадэль на GPU
    if torch.cuda.is_available() and not _model_on_gpu:
        pipe.model.to("cuda")
        _model_on_gpu = True

    start = time.perf_counter()
    preds = pipe(audio_path)                # інферэнс
    elapsed = time.perf_counter() - start

    top3 = {p["label"]: p["score"] for p in preds[:3]}
    return top3, f"{elapsed:.2f} сек"

# ── Gradio-інтэрфейс ─────────────────────────────────────────────
app = gr.Interface(
    fn=classify_audio,
    inputs=gr.Audio(type="filepath", label="Upload Audio File"),
    outputs=[
        gr.Label(num_top_classes=3, label="Predictions"),
        gr.Textbox(label="⏱️ Inference time")
    ],
    title="Audio Classification (MIT/AST) · ZeroGPU",
    description="Загрузіце аўдыя-файл – атрымаеце 3 лепшыя катэгорыі гуку і час інферэнсу.",
    cache_examples=False,
)

if __name__ == "__main__":
    app.launch()