Aime/app/inference/rknpu_whisper_encoder.py

"""ctypes-Binding auf librknnrt.so für den Whisper-turbo-Encoder auf der NPU (specs/009 Option C).

Muster: rknpu_embed.py (specs/007, DEC-708 — kein rknn-toolkit-lite2, cp313-Lücke).
Läuft NUR auf dem R1 (librknnrt.so 2.3.2 + Kernel-Treiber 0.9.8); auf Windows nicht
importier-testbar, wie alles unter device/inference/.

Modell: /opt/aime/models/asr/turbo-encoder-fp16.rknn (Fixed-Shape,
mel [1, 128, 3000] float32 → enc_out [1, 1500, 1280] float32 = ein 30-s-Fenster;
Konvertierung: aime-rknn:2.3.2-Docker auf dem PC, specs/009 T903).

Integrations-Naht: faster-whisper `WhisperModel.encode()` ist die EINZIGE
Encoder-Stelle (auch für BatchedInferencePipeline); der Transcriber ersetzt sie
bei `transcription.compute: npu` durch run() + get_ctranslate2_storage.
Batch-Eingaben (B, 128, 3000) laufen als Schleife über das statische 1er-Modell.

Nutzung:
    e = RknnWhisperEncoder("/opt/aime/models/asr/turbo-encoder-fp16.rknn")
    out = e.run(mel)   # (B, 128, 3000) float32 → (B, 1500, 1280) float32
    e.close()
"""
import ctypes as C

import numpy as np

RKNN_SUCC = 0
RKNN_TENSOR_FLOAT32 = 0
RKNN_TENSOR_UNDEFINED_FMT = 3  # rknn_api.h: NCHW=0, NHWC=1, NC1HWC2=2, UNDEFINED=3
RKNN_NPU_CORE_0_1_2 = 7        # alle 3 NPU-Kerne

MEL_SHAPE = (128, 3000)
OUT_SHAPE = (1500, 1280)


class _RknnInput(C.Structure):
    _fields_ = [
        ("index", C.c_uint32),
        ("buf", C.c_void_p),
        ("size", C.c_uint32),
        ("pass_through", C.c_uint8),
        ("type", C.c_int),
        ("fmt", C.c_int),
    ]


class _RknnOutput(C.Structure):
    _fields_ = [
        ("want_float", C.c_uint8),
        ("is_prealloc", C.c_uint8),
        ("index", C.c_uint32),
        ("buf", C.c_void_p),
        ("size", C.c_uint32),
    ]


class RknnWhisperEncoder:
    """Eine geladene .rknn-Instanz; run() ist NICHT thread-safe (ein Kontext)."""

    def __init__(self, model_path: str, lib_path: str = "librknnrt.so",
                 core_mask: int = RKNN_NPU_CORE_0_1_2):
        self._lib = C.CDLL(lib_path)
        self._lib.rknn_init.argtypes = [C.POINTER(C.c_uint64), C.c_void_p,
                                        C.c_uint32, C.c_uint32, C.c_void_p]
        self._lib.rknn_set_core_mask.argtypes = [C.c_uint64, C.c_int]
        self._lib.rknn_inputs_set.argtypes = [C.c_uint64, C.c_uint32,
                                              C.POINTER(_RknnInput)]
        self._lib.rknn_run.argtypes = [C.c_uint64, C.c_void_p]
        self._lib.rknn_outputs_get.argtypes = [C.c_uint64, C.c_uint32,
                                               C.POINTER(_RknnOutput), C.c_void_p]
        self._lib.rknn_outputs_release.argtypes = [C.c_uint64, C.c_uint32,
                                                   C.POINTER(_RknnOutput)]
        self._lib.rknn_destroy.argtypes = [C.c_uint64]

        with open(model_path, "rb") as f:
            blob = f.read()
        self._ctx = C.c_uint64(0)
        buf = C.create_string_buffer(blob, len(blob))
        rc = self._lib.rknn_init(C.byref(self._ctx), buf, len(blob), 0, None)
        if rc != RKNN_SUCC:
            raise RuntimeError(f"rknn_init rc={rc} (librknnrt/Treiber prüfen)")
        if core_mask:  # best effort — Speed-Sache, kein Korrektheits-Thema
            self._lib.rknn_set_core_mask(self._ctx, core_mask)

    def _run_one(self, mel: np.ndarray) -> np.ndarray:
        ins = (_RknnInput * 1)()
        ins[0].index = 0
        ins[0].buf = mel.ctypes.data_as(C.c_void_p)
        ins[0].size = mel.nbytes
        ins[0].pass_through = 0
        ins[0].type = RKNN_TENSOR_FLOAT32
        ins[0].fmt = RKNN_TENSOR_UNDEFINED_FMT
        rc = self._lib.rknn_inputs_set(self._ctx, 1, ins)
        if rc != RKNN_SUCC:
            raise RuntimeError(f"rknn_inputs_set rc={rc}")
        rc = self._lib.rknn_run(self._ctx, None)
        if rc != RKNN_SUCC:
            raise RuntimeError(f"rknn_run rc={rc}")
        outs = (_RknnOutput * 1)()
        outs[0].want_float = 1
        outs[0].is_prealloc = 0
        outs[0].index = 0
        rc = self._lib.rknn_outputs_get(self._ctx, 1, outs, None)
        if rc != RKNN_SUCC:
            raise RuntimeError(f"rknn_outputs_get rc={rc}")
        n = outs[0].size // 4
        flat = np.ctypeslib.as_array(
            C.cast(outs[0].buf, C.POINTER(C.c_float)), shape=(n,)).copy()
        self._lib.rknn_outputs_release(self._ctx, 1, outs)
        return flat.reshape(OUT_SHAPE)

    def run(self, mel: np.ndarray) -> np.ndarray:
        """(B, 128, 3000) oder (128, 3000) float32 → (B, 1500, 1280) float32."""
        mel = np.ascontiguousarray(mel, dtype=np.float32)
        if mel.ndim == 2:
            mel = mel[np.newaxis]
        if mel.shape[1:] != MEL_SHAPE:
            raise ValueError(f"mel shape {mel.shape} != (B,)+{MEL_SHAPE}")
        out = np.empty((mel.shape[0],) + OUT_SHAPE, dtype=np.float32)
        for i in range(mel.shape[0]):
            out[i] = self._run_one(np.ascontiguousarray(mel[i]))
        return out

    def close(self):
        if getattr(self, "_ctx", None) and self._ctx.value:
            self._lib.rknn_destroy(self._ctx)
            self._ctx = C.c_uint64(0)

    def __del__(self):  # best effort
        try:
            self.close()
        except Exception:
            pass