Files
Aime/app/inference/rknpu_embed.py
T
Claude f9e1eb82c4 feat: komplette Codebase vom Board ins Repo gespiegelt
app/ = /opt/callrec (Node-Backend, Vite-Frontend, Python-Inferenz,
Audio-Skripte, rkllm_server.cpp), system/ = systemd-Units,
WirePlumber/PipeWire-Drop-ins, sudoers. Modelle, Binaries, Aufnahmen
und node_modules bleiben draussen (.gitignore).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-16 11:18:29 +02:00

117 lines
4.6 KiB
Python

"""ctypes-Binding auf librknnrt.so für das wespeaker-Embedding auf der NPU (specs/007).
Warum kein rknn-toolkit-lite2: dessen PyPI-Wheels enden bei cp312, der R1 läuft auf
Python 3.13 → direktes C-API-Binding nach dem bewährten loop_native.py-Muster (DEC-708).
Struct-Layouts verifiziert gegen rknn_api.h v2.3.2 (airockchip). Läuft NUR auf dem R1
(braucht /usr/lib/librknnrt.so 2.3.2 + Kernel-Treiber 0.9.8) — auf Windows nicht
importier-testbar, wie alles unter device/inference/.
Modell: /opt/aime/models/diar/wespeaker-pyannote-fp16.rknn (Fixed-Shape,
feats [1,998,80] float32 + weights [1,125] float32 → embs [1,256] float32;
Konvertierung siehe specs/007 T702/T703, Toolchain aime-rknn:2.3.2 auf dem PC).
Nutzung:
e = RknnEmbedding("/opt/aime/models/diar/wespeaker-pyannote-fp16.rknn")
emb = e.run(feats, weights) # → np.ndarray [1, 256]
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 (T704b: 1,84x vs. 1 Kern, 39 ms/Chunk)
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 RknnEmbedding:
"""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(self, feats: np.ndarray, weights: np.ndarray) -> np.ndarray:
feats = np.ascontiguousarray(feats, dtype=np.float32)
weights = np.ascontiguousarray(weights, dtype=np.float32)
ins = (_RknnInput * 2)()
for i, arr in enumerate((feats, weights)):
ins[i].index = i
ins[i].buf = arr.ctypes.data_as(C.c_void_p)
ins[i].size = arr.nbytes
ins[i].pass_through = 0
ins[i].type = RKNN_TENSOR_FLOAT32
ins[i].fmt = RKNN_TENSOR_UNDEFINED_FMT
rc = self._lib.rknn_inputs_set(self._ctx, 2, 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
emb = 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 emb.reshape(1, -1)
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