core/model_loader.py

import threading
import time
import tempfile
import wave
import os
import numpy as np
import librosa
from openwakeword.model import Model

from config import (
    MODEL_PATH, SAMPLE_RATE, CHUNK_SIZE,
    WARMUP_ITERS, STABLE_COUNT, STABLE_MS, MODEL_TIMEOUT_S
)

# --- Singleton model + trạng thái ---
_model: Model | None = None
_model_ready = threading.Event()
_model_error: str | None = None


def get_model() -> Model | None:
    return _model


def is_ready() -> bool:
    return _model_ready.is_set() and _model is not None


def wait_until_ready(timeout: float = MODEL_TIMEOUT_S) -> bool:
    """Block cho đến khi model sẵn sàng. Trả về False nếu timeout hoặc lỗi."""
    return _model_ready.wait(timeout=timeout) and _model is not None


def get_error() -> str | None:
    return _model_error


# -------------------------------------------------------
# BƯỚC 1 — Warm-up librosa (lần đầu librosa.load rất chậm
#           do lazy import bên trong, phải kích hoạt trước)
# -------------------------------------------------------
def _warmup_librosa():
    t = time.perf_counter()

    dummy = np.zeros(SAMPLE_RATE, dtype=np.float32)
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        tmp_path = f.name
    with wave.open(tmp_path, "w") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(SAMPLE_RATE)
        wf.writeframes((dummy * 32767).astype(np.int16).tobytes())

    librosa.load(tmp_path, sr=SAMPLE_RATE)
    os.unlink(tmp_path)

    print(f"  [loader] librosa warm-up: {(time.perf_counter()-t)*1000:.0f}ms")


# -------------------------------------------------------
# BƯỚC 2 — Load ONNX model
# -------------------------------------------------------
def _load_oww_model() -> Model:
    t = time.perf_counter()
    m = Model(wakeword_model_paths=[MODEL_PATH])
    print(f"  [loader] ONNX load: {(time.perf_counter()-t)*1000:.0f}ms")
    return m


# -------------------------------------------------------
# BƯỚC 3 — JIT warm-up: chạy predict nhiều lần cho đến khi
#           ONNX runtime compile xong và latency ổn định
#           (tránh lần predict đầu của user bị chậm bất thường)
# -------------------------------------------------------
def _warmup_jit(model: Model):
    stable_count = 0
    silence = np.zeros(CHUNK_SIZE, dtype=np.int16)

    for i in range(WARMUP_ITERS):
        t_iter = time.perf_counter()
        model.predict(silence)
        elapsed_ms = (time.perf_counter() - t_iter) * 1000

        if i < 3:
            # Vài iter đầu luôn bất thường, bỏ qua
            continue

        if elapsed_ms < STABLE_MS:
            stable_count += 1
            if stable_count >= STABLE_COUNT:
                print(f"  [loader] JIT stable sau {i+1} iters ({elapsed_ms:.1f}ms/iter)")
                return
        else:
            stable_count = 0

    print(f"  [loader] JIT warm-up xong {WARMUP_ITERS} iters (chưa stable hoàn toàn)")


# -------------------------------------------------------
# MAIN — chạy trong background thread ngay khi app start
# -------------------------------------------------------
def _boot():
    global _model, _model_error

    t_total = time.perf_counter()
    print("[loader] Bắt đầu khởi động model...")

    try:
        _warmup_librosa()
        model = _load_oww_model()
        _warmup_jit(model)
        _model = model
        print(f"[loader] ✅ Sẵn sàng — tổng boot: {time.perf_counter()-t_total:.1f}s")
    except Exception as e:
        _model_error = str(e)
        print(f"[loader] ❌ Lỗi: {e}")
    finally:
        # Dù thành công hay lỗi đều set event
        # để wait_until_ready() không bị block mãi
        _model_ready.set()


def start_loading():
    """Gọi hàm này 1 lần duy nhất trong app.py khi khởi động."""
    thread = threading.Thread(target=_boot, daemon=True)
    thread.start()