diff --git a/.gitignore b/.gitignore index 9ea0e2e..f3b64b2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,8 +1,24 @@ -# ============================================ -# Kazeia .gitignore — code only, no binaries -# ============================================ +# ============================================================ +# Kazeia .gitignore — LISTE BLANCHE +# On ne suit QUE l'app mobile + le code/les docs essentiels. +# Tout le reste à la racine (modèles, backups, SDK, venvs, arbres +# expérimentaux, intouchables beta_kazeia/root_oneplus, Kaz, etc.) +# est ignoré — présent sur le disque, jamais supprimé, juste hors suivi. +# ============================================================ -# === Large binary files === +# 1) Ignorer tout ce qui est à la racine du dépôt... +/* + +# 2) ...sauf ce qu'on veut explicitement suivre : +!/.gitignore +!/kazeia-android/ +!/scripts/ +!/executorch-custom/ +!/executorch-patches/ +!/*.md +!/*.txt + +# 3) Binaires & sorties de build (même à l'intérieur des dossiers suivis) : *.so *.so.* *.so.bak @@ -15,8 +31,8 @@ *.npy *.wav *.apk - -# === Build outputs === +*.pyc +__pycache__/ kazeia-android/app/build/ kazeia-android/build/ kazeia-android/.gradle/ @@ -25,61 +41,11 @@ kazeia-android/extracted/ kazeia-android/app/.cxx/ kazeia-android/unityLibrary/ -# === Python environments === -et_venv/ -qnn_venv/ -__pycache__/ -*.pyc - -# === SDKs and toolchains (external) === -android-ndk-r27d/ -qnn_sdk_242/ -executorch/ - -# === Models and data (too large) === -models_qnn/ -tablet_backup*/ -backup_*/ -voix/ - -# === IDE === +# 4) IDE / OS / temporaires .idea/ *.iml - -# === Compiled binaries === -cp_et_test_client -llama.cpp/build*/ - -# === OS files === .DS_Store Thumbs.db - -# === Temporary === *.tmp *.log *.kate-swp - -# === External repos (submodules or separate) === -Vulkan-Headers/ -llama.cpp/ -whisper.cpp/ -kazeia-unity/ -models_hf/ - -# === Old/misc at root === -android-sdk/ -avatar_disabled_backup/ -beta_kazeia/ -build_et_jar/ -root_oneplus/ -cmdtools.zip -forward.dlc -.claude/ - -# === Root-level scripts (moved to scripts/) === -export_decoder_onnx.py -export_talker_onnx.py -extract_tts_embeddings.py -extract_vq_individual.py -generate_tokens_for_tablet.py -generate_tts_wav.py diff --git a/FROZEN.md b/FROZEN.md new file mode 100644 index 0000000..b985dc9 --- /dev/null +++ b/FROZEN.md @@ -0,0 +1,234 @@ +# Kazeia — Stack production figée + +**Date du gel : 2026-05-02** + +Ce document définit la stack **autorisée** de Kazeia. Tout composant non listé +ici doit être considéré comme **non autorisé** et ne doit pas être réintroduit +sans validation explicite + mise à jour de ce contrat. + +--- + +## Emplacement des modèles (MAJ 2026-05-18 — chantier installeur Phase 0) + +Les chemins `/data/local/tmp/kazeia/models/` et `/data/local/tmp/kazeia-et/` +cités ci-dessous restent **valides sur les tablettes de dev** (modèles poussés +par `adb push`). Ils sont désormais **résolus au runtime par `KazeiaPaths`** : + +- si la racine `/data/local/tmp` existe et est non vide → elle prime (dev) ; +- sinon → stockage externe app `getExternalFilesDir("kazeia")/{models,llm}`, + rempli par l'installeur (modèles téléchargés depuis le Nextcloud). + +Aucun composant n'est ajouté ni retiré — seul l'emplacement de chargement +devient dynamique. `KazeiaApplication.MODELS_DIR` / `.LLM_DIR` sont les points +d'accès uniques. Cf `project_kazeia_installer_updates`. + +--- + +## STT — Speech-to-Text + +### Composant figé + +`com.kazeia.stt.WhisperHybridEngine` (542 lignes) — chemin **NPU only** +- Encoder + Decoder Qualcomm AI Hub HfWhisper KV-cache (Whisper-Small 12L/12H) +- Modèles tablette : `/data/local/tmp/kazeia/models/whisper-small-sm8750/` + (fingerprints MD5 dans `/opt/Kazeia/kazeia-stt-perf-validated/models/README.md`) +- Backend : ONNX Runtime QNN HTP V79 (Snapdragon 8 Elite) +- Performance validée : RTF 0.5 (audio 1.6 s → 825 ms) +- Override `<|translate|>` → `<|transcribe|>` actif (force transcription FR) +- Auto-détection layers/heads via `enc.outputInfo["k_cache_cross_0"]` + +### Mel extractor figé + +`com.kazeia.stt.MelExtractor` + `kazeia-android/app/src/main/jni/mel_extractor.cpp` +- HuggingFace-compatible WhisperFeatureExtractor en C++ pur +- 80 mels × 3000 frames, FFT 400, hop 160 + +### Wrapper v2 figé + +`com.kazeia.v2.SttStage` — thin wrapper sur `WhisperHybridEngine` + +### ❌ Interdit +- Pas de fallback CPU (`WhisperSttEngine`, `WhisperJni`, `whisper.cpp`) — supprimés +- Pas de fallback Android stock (`AndroidSttEngine`) — supprimé +- Pas d'autres backends (`WhisperLiteRtEngine`, `WhisperNpuSttEngine`) — supprimés +- Pas de Qwen3-ASR (1.7B GGUF cassé, 0.6B qualité FR insuffisante) +- Pas de `libwhisper.so` dans jniLibs — supprimé + +--- + +## VAD — Voice Activity Detection + +### Composant figé + +**Détection RMS-énergie inline** : +- v1 : `KazeiaService.startContinuousListening()` — frame 1600 samples (100 ms), + threshold RMS=150, ≥ 300 ms speech, ≥ 800 ms silence end +- v2 : `com.kazeia.v2.VadStage` — réimplémentation propre du même algorithme + +### ❌ Interdit +- Pas de Silero VAD (`SileroVadEngine.kt` + `silero_vad.onnx` — supprimés) +- Pas de `AudioCaptureManager.kt` — supprimé +- Pas de `core.VadEngine` interface — supprimée +- Cf `feedback_kazeia_vad_rms.md` : Silero plafonne à 5e-4 sur ce mic + +--- + +## LLM — Large Language Model + +### Composants figés v1 (legacy `KazeiaService`) + +`com.kazeia.llm.ExecuTorchLlmEngine` (in-process) +- Modèle : `hybrid_llama_qnn.pte` symlink → `hybrid_llama_qnn_4b.pte` sur tablette +- Tokenizer : `tokenizer.json` +- Backend : ExecuTorch + QNN HTP NPU (in-process via PyTorch ExecuTorch JNI) +- Path : `/data/local/tmp/kazeia-et/` + +### Composants figés v2 (`KazeiaServiceV2`) + +`com.kazeia.v2.LlmCascadeStage` (mode mono-Speaker — temporaire RAM-budget) +- Spawn 1× `qnn_llama_runner --daemon_mode` via `ExecutorchDaemon` +- Modèle : `hybrid_llama_qnn_4b.pte` + `tokenizer.json` +- System prompt : SYS_SPEAKER (psychologue bienveillant, 1-2 phrases FR concises) +- Le code Thinker/cascade est conservé en commentaire pour restauration future + +### ❌ Interdit +- Pas de `LlamaCppLlmEngine` (purgé 2026-04-29) +- Pas de `KazeiaLlmJni` (purgé 2026-04-29) +- Pas de Genie SDK (bug GQA Qwen3) +- Pas de cascade Thinker+Speaker SIMULTANÉE en cohabitation TTS (RAM OOM, cf + cleanup 2026-05-02) +- Pas de `LLM_BACKEND` constante / `_currentLlmModel` UI selector + +### Restauration cascade complète : conditionné + +À envisager **uniquement après** : +- Task #240 (libllama.so contre chraac-llama récent) → -30 % RAM/perf Talker+CP +- OU partage `.ptd` entre 2× 4B → -2.4 GB RAM +- OU phased loading opportuniste (cf `feedback_kazeia_vad_rms.md` non, l'analyse + RAM dans la conversation 2026-05-02) +- En attendant, mono-Speaker reste figé + +--- + +## TTS — Text-to-Speech + +### Composant figé + +`com.kazeia.tts.Qwen3TtsEngine` — voie **ggml-cpu Talker + ggml-cpu CP + ggml C++ Decoder** +- Talker : `TtsTalkerCpuJni` linké contre `libllama.so` (in jniLibs) + - Modèle préféré : `talker_f32.gguf` (1.78 GB, fp32, élimine drift fp16) + - Fallback : `talker_f16.gguf` +- CP : `TtsCpCpuJni` linké contre `libllama.so` + - Modèle préféré : `cp_f32.gguf` (316 MB) + - Fallback : `cp_f16.gguf` +- Decoder : `libtts_decoder_ggml.so` static-linked contre **chraac-llama** ggml + (build : `kazeia-tts-decoder-ggml/build-android-static-chraac/`) + - GGUF : `qwen3tts_decoder.gguf` (340 MB) + - **RTF 0.96 mesuré 2026-05-02** (cf `kazeia-tts-perf-validated/`) +- Tokenizer texte : `Qwen3BpeTokenizer` (BPE byte-level Qwen2/Qwen3) +- Voice cloning : Damien (prefix 9 × 1024 + suffix 2 × 1024) + +### API figée + +- `synthesizeAndPlay(text, language)` — pour invocations one-shot +- `startStreamingSession()` / `enqueueSentence(text)` / `endStreamingSession()` — + pour streaming sentence-par-sentence (utilisé dans v2.TtsStage) + +### Wrapper v2 figé + +`com.kazeia.v2.TtsStage` — wrapper streaming-session sur `Qwen3TtsEngine` + +### ❌ Interdit dans le chemin de chargement + +- Pas de chemin `cp_kv_v2 ONNX` (ligne 630-658 de Qwen3TtsEngine, gated sur + `!useCpuPath` 2026-05-02) +- Pas de `cp_et_runner` TCP (sub-process root, gated 2026-05-02) +- Pas de `talker_kv_cpu` ONNX fp32 (gated sur `!useCpuPath`) +- Pas de Hexagon talker / CP (`hexStartRunner`, `useHexagonTalker`) — pas dans + flux production +- Pas de `phrase_embeds.bin` cache (bypassed 2026-05-02 — provoquait dégénération) + +### ❌ Interdit moteurs alternatifs +- Pas de `AndroidTtsEngine` — supprimé +- Pas de `ChatterboxTtsEngine` — supprimé + +--- + +## Audio playback + +### Composant figé + +`com.kazeia.audio.AudioPlaybackManager` (54 lignes) + +L'AudioTrack interne au streaming session du `Qwen3TtsEngine` (avec keepAlive +watchdog ColorOS) reste également figé. + +--- + +## Build / Backends + +### ggml — sources autorisées + +| Composant | Source ggml autorisée | Build dir | +|---|---|---| +| **TTS Decoder** | `chraac-llama` (branch `dev-refactoring`, commit `897501a`) | `build-android-static-chraac` | +| TTS Talker (lib) | jniLibs `libggml-cpu.so` (`llama-upstream`, à migrer Task #240) | `kazeia-android/.../jniLibs/arm64-v8a` | +| TTS CP (lib) | idem | idem | +| STT | n/a (ONNX Runtime, pas ggml) | n/a | + +**Critique** : ne JAMAIS pointer le décodeur TTS vers `whisper.cpp/ggml` ou +`llama-upstream/ggml` — sinon RTF passe à 1.4-1.5 (mesuré, cf +`kazeia-tts-perf-validated/docs/MEASUREMENTS.md`). + +### Compile flags figés + +```cmake +# TTS decoder (libtts_decoder_ggml.so) +-O3 -march=armv8.6-a+dotprod+fp16+i8mm+bf16 -fopenmp -static-openmp +GGML_CPU_ARM_ARCH=armv8.6-a+dotprod+fp16+i8mm+bf16 +GGML_OPENMP=ON +BUILD_SHARED_LIBS=OFF +``` + +--- + +## Fichiers de référence persistants + +- `/opt/Kazeia/kazeia-stt-perf-validated/` — doc + MD5 modèles STT/VAD +- `/opt/Kazeia/kazeia-tts-perf-validated/` — doc + libs `.a` + bench CLI + reproductible TTS (RTF<1) +- `/opt/Kazeia/chraac-llama/` — fork ggml indispensable pour le décodeur TTS +- `/opt/Kazeia/kazeia-tts-decoder-ggml/build-android-static-chraac/` — build + statique pointé par l'app + +--- + +## Comment dégeler / modifier ce contrat + +**Procédure obligatoire** avant de réintroduire un composant supprimé OU de +remplacer un composant figé : + +1. Mesurer la perf du composant cible **vs** la voie figée (sur la même phrase + de test, même tablette, même gouverneur DVFS) +2. Démontrer que la nouvelle voie est strictement meilleure (perf, RAM ou + maintenabilité) avec des chiffres reproductibles +3. Mettre à jour ce `FROZEN.md` avec le nouveau composant + les chiffres +4. Mettre à jour `MEMORY.md` index + memory file pertinent +5. Documenter le commit/build SHA exact dans `kazeia-{stt,tts}-perf-validated/` +6. Conserver l'ancien composant comme `*.archive.kt` UN trimestre avant + suppression définitive (filet de sécurité rollback) + +Sans ce processus, **ne pas modifier la stack figée**. + +--- + +## Métriques de référence (2026-05-02) + +| Domaine | Métrique | Valeur figée | Comment | +|---|---|---:|---| +| STT | Load total | 750-1035 ms | NPU encoder + decoder + mel filters | +| STT | RTF (audio 1.6 s) | **0.51** | mel + encoder + decoder NPU | +| VAD | Trigger end-of-speech | 800 ms après dernier audio > seuil | hardcoded | +| LLM | Mono-Speaker forward | ~22 tok/s | qnn_llama_runner subprocess | +| TTS Decoder | RTF (audio 4.48 s) | **0.96** | ggml C++ chraac-llama, standalone | +| TTS Pipeline complet | RTF | ~2.87 | bottleneck CP (Task #240 pending) | diff --git a/RAPPORT_KAZEIA_2026-05-04.md b/RAPPORT_KAZEIA_2026-05-04.md new file mode 100644 index 0000000..c003725 --- /dev/null +++ b/RAPPORT_KAZEIA_2026-05-04.md @@ -0,0 +1,508 @@ +# Rapport Kazeia — État, architecture, performances + +**Date :** 2026-05-04 +**Contexte :** Kazeia est une application Android de soutien thérapeutique +conversationnel (psychologue virtuel) tournant **100 % on-device** sur OnePlus +Pad 3, sans connectivité réseau requise et sans root utilisateur. + +--- + +## 1. Synthèse exécutive + +### Ce qui fonctionne en production +- Pipeline complet **STT → LLM → TTS** opérationnel, multi-tour, avec cascade + Thinker+Speaker. +- Latence end-to-end utilisable en conversation : **3-5 s** entre la fin du + PTT et le premier audio de réponse. +- Stable mémoire après débloat ColorOS + désinstallation du watchdog OEM Athena. +- Application 100 % autonome sur la tablette : aucun service réseau requis, + modèles et runtime tous embarqués. +- Stack figée et documentée (`/opt/Kazeia/FROZEN.md`). +- Backup tablette complet et automatisable dans `/opt/Kazeia/kazeia-tablet-backup/`. + +### Limites actuelles connues +- Le **Thinker Guard-0.6B** produit des bullets de qualité analytique limitée + (modèle conçu pour classification de safety, détourné vers analyse clinique). + Le Speaker compense en pratique mais l'apport qualitatif de la cascade est + moins net qu'espéré. +- Le rendu UI de l'orbe (cycle couleurs pendant Thinking) demande encore + validation visuelle utilisateur. +- Conversation history non persisté entre sessions (à dessein actuellement). +- Pas de fine-tuning thérapeutique français spécifique sur les modèles — + on utilise Qwen3-4B générique pré-entraîné. + +--- + +## 2. Plateforme matérielle + +| Composant | Spec | +|---|---| +| Tablette | OnePlus Pad 3 (OPD2415EEA) | +| RAM | 16 GB LPDDR5X | +| Stockage | 256 GB UFS | +| SoC | Snapdragon 8 Elite | +| CPU | Oryon (8 cœurs ARMv9) | +| GPU | Adreno 830 | +| **NPU/DSP** | **Hexagon V79 (HMX/HVX)** | +| OS | Android 16, ColorOS post-2026-03-18 OTA | +| Root | Aucun (user mode standard) | + +--- + +## 3. Stack logicielle + +``` +┌────────────────────────────────────────────────────────────────┐ +│ Kazeia Android App (Kotlin) │ +│ /data/app/.../com.kazeia │ +└───────────────────┬─────────────────────────┬───────────────────┘ + │ │ + ┌───────────▼──────────┐ ┌───────────▼──────────┐ + │ KazeiaService │ │ ChatActivity (UI) │ + │ (state machine) │ │ AudioVisualizerView│ + │ • PTT │ │ (orb HSV cycle) │ + │ • Cascade │ └──────────────────────┘ + │ • Pipeline state │ + └─┬────────┬────────┬──┘ + │ │ │ + ┌──────▼─┐ ┌───▼────┐ ┌─▼──────┐ + │ STT │ │ LLM │ │ TTS │ + │ Engine │ │ Engine │ │ Engine │ + └────────┘ └────────┘ └────────┘ +``` + +### 3.1 STT — Whisper-Small NPU (QAIRT) + +- **Modèle** : `HfWhisper-Small` (12 layers, 12 heads) +- **Format** : QAIRT context binary (Qualcomm SDK natif) +- **Fichiers** : + - `HfWhisperEncoder_qairt_context.bin` (210 MB) + - `HfWhisperDecoder_qairt_context.bin` (361 MB) +- **Backend** : Hexagon V79 via QAIRT runtime (`libQnnHtpV79Skel.so`) +- **VAD** : RMS énergie inline (seuil 80, frames 100 ms) +- **Mode** : continuous listening avec streaming Whisper pendant PTT +- **Performance mesurée** : RTF 0.4 (transcription "Bonjour" ~640 ms total : + mel 260 ms + encoder 115 ms + decoder 170 ms) + +### 3.2 LLM — Cascade Option E v2 (in-process) + +**Architecture cascade** : 2 `LlmModule` ExecuTorch dans le **même process app**, +partageant le `QnnBackendBundle` via le singleton `QnnBackendUnifiedRegistry` +(donc 1 seul handle fastrpc DSP, pas de collision). + +``` +Process com.kazeia +├── Speaker Qwen3-4B (.pte 3.3 GB, ~4.4 GB ION) +│ └── system prompt: SYS_KAZEIA + "Indices internes (NE PAS RÉPÉTER) : ..." +└── Thinker Qwen3Guard-Gen-0.6B (.pte 700 MB, lazy first PTT) + └── system prompt: SYS_THINKER (4 bullets émotion/besoin/approche/axe) +``` + +| Composant | Modèle | Format | Taille | Quantization | +|---|---|---|---|---| +| Speaker | Qwen/Qwen3-4B | ExecuTorch `.pte` QNN HTP V79 | 3.3 GB | Q4 W4A16 | +| Thinker | Qwen/Qwen3Guard-Gen-0.6B | ExecuTorch `.pte` QNN HTP V79 | 700 MB | Q4 W4A16 | + +**Cascade sequence par tour utilisateur :** +1. STT → texte patient (transcription Whisper) +2. **Thinker pass** : `Guard-0.6B(SYS_THINKER + msg)` → 4 bullets cliniques +3. Bullets tronqués à 400 chars (sécurité prefill budget) +4. **Speaker pass** : `Qwen3-4B(SYS_KAZEIA + bullets en system, msg en user)` → réponse 1-2 phrases +5. Streaming TTS phrase par phrase au fur et à mesure du décodage + +### 3.3 TTS — Qwen3-TTS ggml-cpu + +- **Modèle source** : Qwen3-TTS-0.6B-Base (clonage vocal) +- **Pipeline 3 étapes**, toutes en CPU pur (NEON ARMv9, pas de DSP) : + +| Étape | Modèle | Fichier | Taille | Backend | +|---|---|---|---|---| +| Talker | transformer audio → tokens semantic | `talker_f16.gguf` | 851 MB | ggml-cpu | +| Code Predictor (CP) | sem → 16 codebooks RVQ | `cp_q8_0.gguf` | 84 MB + 240 MB embeds | ggml-cpu | +| Decoder | RVQ → mel → audio 24 kHz | `qwen3tts_decoder.gguf` | 325 MB | ggml-cpu (libtts_decoder_ggml) | + +- **Voice cloning** : prefix/suffix embeddings pré-extraits (10 voix : damien, + elodie, jerome, richard, sid, zelda, didier, amir...) +- **Streaming sentence-level** : audio joue dès la 1ère phrase pendant que LLM + décode encore les suivantes (Opt TTS-B) +- **RTF mesuré** : 0.96 sur tablette (chraac-llama fork) + +### 3.4 Glue applicative + +| Composant | Rôle | +|---|---| +| `KazeiaService` (Kotlin) | State machine PTT + cascade + pipeline | +| `ExecuTorchLlmEngine` | Wrapper Java→JNI sur `LlmModule` ExecuTorch | +| `Qwen3TtsEngine` | Orchestration JNI Talker+CP+Decoder + voice clone | +| `WhisperHybridEngine` | JNI QAIRT + mel extractor NEON | +| `AudioVisualizerView` | Orbe Canvas, HSV cycle pendant Thinking | +| `SentenceStreamer` | Découpage phrases pour streaming TTS | +| `PipelineMetrics` | Instrumentation timings end-to-end | + +--- + +## 4. Performances mesurées + +### 4.1 Latences par étape (mesures réelles dans logcat) + +| Étape | Temps typique | Note | +|---|---|---| +| **STT Whisper-Small** | 460-680 ms | RTF 0.4 sur audio 1.7s | +| ├ Mel extraction | 190-270 ms | NEON sur CPU | +| ├ Encoder NPU | 115 ms | QAIRT V79 | +| └ Decoder NPU | 75-300 ms | dépend du nb tokens | +| **Thinker (Guard-0.6B)** | 446 ms | 22 tokens à 49 tok/s, TTFT 130 ms | +| ├ Lazy load (1ère fois) | 1 088 ms | -844 MB MemAvailable | +| └ Inference | 446 ms | 4 bullets, prompt ~180 tok | +| **Speaker (Qwen3-4B)** | 1 000-2 200 ms | 14-46 tokens à 14-21 tok/s, TTFT 190 ms | +| ├ resetContext | <1 ms | clear KV cache | +| ├ Prefill (~150 tok) | 200 ms | hybrid mode AR=128 | +| └ Decode | 800-2000 ms | 13.9-21 tok/s selon longueur | +| **TTS Talker step** | ~25 ms | ggml-cpu NEON | +| **TTS CP step** | ~43 ms | Q8_0, optimisé | +| **TTS Decoder full** | RTF 0.96 | ggml-cpu, premier audio <1 s | + +### 4.2 Latence end-to-end PTT → 1ᵉʳ audio + +``` +PTT release ─┬───── 100 ms finalize segment + │ + ╔═══ 600 ms ═══════ STT Whisper ─┐ + ║ │ + ╠═══ 450 ms ═══════ Thinker Guard-0.6B │ utilisateur + ║ │ entend silence + ╠═══ 500 ms ═══════ Speaker prefill + 1ère phrase │ + ║ │ + ╠═══ 800 ms ═══════ TTS 1ère phrase synthèse ─┘ + ║ + ▼ 1ᵉʳ audio joué (~ 2.5 s après PTT release) +``` + +**Total PTT release → 1ᵉʳ audio : ~ 2.5 s** (pour message court). +**Total PTT release → réponse complète audio : ~ 4-6 s**. + +### 4.3 Throughput LLM mesuré + +| Modèle | Hardware | Decode tok/s | TTFT | +|---|---|---|---| +| Qwen3-4B Speaker | Hexagon V79 | 14-21 | 190 ms | +| Guard-0.6B Thinker | Hexagon V79 | **49** | 130 ms | +| Whisper-Small | Hexagon V79 | n/a (encoder/decoder fusion) | 115+75 ms | + +--- + +## 5. Profil mémoire + +### 5.1 Tablette en idle (Speaker chargé, Thinker pas encore) + +| Métrique | Valeur | +|---|---| +| MemAvailable | 5.10 GB | +| Cached (file cache) | 4.51 GB | +| Process Kazeia RSS | 2.06 GB | +| Process Kazeia PSS total | 2.14 GB | + +### 5.2 Décomposition PSS Kazeia (idle) + +| Section | PSS | % | +|---|---|---| +| Native Heap (modèles, ION, native libs) | 1.65 GB | 77 % | +| Java/Dalvik Heap | 362 MB | 17 % | +| Graphics (EGL/GL) | 64 MB | 3 % | +| Code (.so + .apk + .dex mmap) | 42 MB | 2 % | +| Stack/divers | 26 MB | 1 % | +| **TOTAL PSS** | **2.14 GB** | | + +### 5.3 Pendant cascade et TTS + +| | Idle | Cascade Thinker | TTS pic | +|---|---|---|---| +| RSS Kazeia | 2.06 GB | ~2.9 GB | ~3.5 GB | +| ION DSP cumulé | ~3.3 GB | ~5.0 GB | ~5.5 GB | +| MemAvailable | 5.10 GB | 4.0 GB | 3.0 GB | + +Le pic TTS reste loin du seuil critique (~1 GB MemAvailable où Linux LMK +commencerait à killer). **Stable**. + +### 5.4 Tablet baseline (sans Kazeia) + +- **Stock fraîche** : MemAvailable ~3.5 GB (900+ processus OEM ColorOS) +- **Après débloat + reboot** : MemAvailable ~11.1 GB → **+7.7 GB de marge** + +--- + +## 6. Trajectoire — Ce qui a été fait + +### Phase A : Pipeline initial (avril 2026) +- Migration v1 du PoC Python/Flask vers app Android Kotlin native +- Speaker via ExecuTorch QNN delegate (Qwen3-4B `.pte` hybrid mode) +- STT Whisper-Small NPU via QAIRT +- TTS Qwen3-TTS originalement en ExecuTorch `.pte` + +### Phase B : Migration TTS vers ggml-cpu (avril-mai 2026) +- Talker exporté en GGUF f16 (vs `.pte` ExecuTorch) +- CP exporté en GGUF Q8_0 (84 MB vs 316 MB f32 vs 158 MB f16) +- Decoder porté en C++ via libtts_decoder_ggml.so +- Validation RTF < 1 sur tablette (chraac-llama fork) + +### Phase C : Tablette stabilisation (mai 2026) +- Mesure baseline ColorOS : ~3.5 GB MemAvailable seulement +- **Débloat** sans root : 84 packages désinstallés via `pm uninstall --user 0` +- **Athena watchdog** : désinstallation du package (kill ION quota 2 GB sur app) +- Suppression de symlinks ⚠ SELinux ColorOS sans Magisk les bloque +- Backup complet 25 GB dans `/opt/Kazeia/kazeia-tablet-backup/` +- Factory reset + restore réussi + +### Phase D : Cascade Thinker+Speaker (mai 2026) +Tentatives successives : + +| Tentative | Approche | Résultat | +|---|---|---| +| 1 | Subprocess `qnn_llama_runner` daemon | ✗ Collision DSP fastrpc (2 sessions QNN HTP) | +| 2 | Option E v1 : 2 LlmModule (Speaker 4B + Guard 4B) | ✗ OOM Linux LMK (8.8 GB ION cumulé) | +| 3 | Option B : 2-pass sur même Speaker | ✓ Stable, qualité Thinker faible (Qwen3-4B générique) | +| 4 | **Option E v2** : Speaker 4B + Guard-0.6B | **✓ Production actuelle** | + +Pour Option E v2 : export Qwen3Guard-Gen-0.6B HF → ExecuTorch `.pte` QNN HTP V79 +en local via `executorch.examples.qualcomm.oss_scripts.llama` (registration +ajoutée dans `__init__.py` + `decoder_constants.py`). + +### Phase E : Robustesse runtime (mai 2026) +- Fix VAD seuil RMS (150 → 80) après calibration mic post-reset +- Fix PTT release : finalisation forcée du segment (sans ça, le LLM n'était + jamais déclenché si bruit ambiant > seuil) +- Fix `seq_len` ExecuTorch : passer 512 (max compilé .pte) au lieu de + `maxNewTokens` (assertion `cur_pos+prompt < seq_len`) +- Fix `cur_pos_` accumulation : appel `LlmModule.resetContext()` à chaque + generate (sinon crash après quelques tours) +- Fix `prefill()` overflow : truncation bullets à 400 chars (Guard-0.6B peut + rambler et faire dépasser le budget prefill du Speaker) +- Fix Speaker régurgitant les bullets : injection bullets dans **system prompt** + (avec instruction « NE PAS RÉPÉTER ») au lieu du user prompt +- Fix orbe visuelle : `Paint.shader = null` avant pose couleur dans + `drawThinking` (sinon le shader posé par `drawIdle` masquait la couleur) + +### Phase F : Backup APK + scripts (en continu) +- `/opt/Kazeia/kazeia-tablet-backup/` (25 GB total) + - APK actuelle (110 MB, libs natives custom bundlées) + - Modèles TTS+STT (12 GB) + - LLM `.pte` (12 GB inc. Guard-0.6B nouveau) + - Voix de référence (130 MB) + - `restore.sh` automatisé +- `/opt/Kazeia/kazeia-debloat/` : `safelist.txt` + `debloat.sh` + `rebloat.sh` +- `/opt/Kazeia/export_guard_06b.sh` : pipeline d'export Guard-0.6B reproducible + +--- + +## 7. Architecture cascade détaillée (Option E v2) + +``` +┌──────────────────────────────────────────────────────────────────┐ +│ KazeiaService │ +│ │ +│ PTT release ───► startContinuousListening() drains buffer │ +│ ► processSpeechInput / streaming Whisper │ +│ ► text = "Comment ça va ?" │ +│ ► processLlmResponse(text) │ +│ │ +│ ┌─ ensureThinkerEngine() ──┐ (lazy 1ère fois ~1.1 s) │ +│ │ ▼ │ +│ │ ┌─────────────────────┐ │ +│ │ │ thinkerEngine │ │ +│ │ │ ExecuTorchLlmEngine │ │ +│ │ │ Guard-0.6B │ │ +│ │ │ SYS_THINKER │ │ +│ │ │ T=0.0, max 128 tok │ │ +│ │ └──────────┬──────────┘ │ +│ │ ▼ │ +│ │ bullets (400 chars max) │ +│ │ │ │ +│ │ ┌─────────────────────────┴──────┐ │ +│ │ │ speakerSystemPrompt = │ │ +│ │ │ SYS_KAZEIA + │ │ +│ │ │ "\n\nIndices d'analyse... │ │ +│ │ │ (NE PAS CITER) :\n" + │ │ +│ │ │ bullets │ │ +│ │ └────────────┬────────────────────┘ │ +│ │ ▼ │ +│ │ ┌──────────────────────────┐ │ +│ │ │ llm = ExecuTorchLlmEngine │ │ +│ │ │ Speaker Qwen3-4B │ │ +│ │ │ generateWithSystem() │ │ +│ │ │ system=above, user=msg │ │ +│ │ │ T=0.7, max 450 tok │ │ +│ │ └──────────┬───────────────┘ │ +│ │ ▼ token-by-token streaming │ +│ │ ┌──────────────────────────┐ │ +│ │ │ SentenceStreamer │ │ +│ │ │ split sur . ! ? │ │ +│ │ └──────────┬───────────────┘ │ +│ │ ▼ phrase à phrase │ +│ │ ┌──────────────────────────┐ │ +│ │ │ Qwen3TtsEngine │ │ +│ │ │ enqueueSentence() │ │ +│ │ └──────────┬───────────────┘ │ +│ │ ▼ │ +│ └─────► AudioPlayback (24 kHz mono, AudioTrack) │ +│ │ +└──────────────────────────────────────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────┐ + │ DSP Hexagon V79 │ + │ • QAIRT (Whisper) │ + │ • ExecuTorch QNN HTP │ + │ (Speaker + Thinker │ + │ partagent bundle) │ + └─────────────────────────┘ + │ + ▼ + ┌─────────────────────────┐ + │ CPU NEON ARMv9 │ + │ • TTS Talker (ggml) │ + │ • TTS CP (ggml) │ + │ • TTS Decoder (ggml) │ + └─────────────────────────┘ +``` + +--- + +## 8. Prompts système (production) + +### Speaker +``` +Tu es Kazeia, psychologue bienveillant. Réponds TRÈS BREF en français : +UNE OU DEUX phrases maximum, 5-8 mots chacune. Validation émotionnelle simple. +Pas de remplissage, pas d'explication, pas de conseil non sollicité. +/no_think + +[ + injection en cascade : ] +Indices d'analyse interne (NE PAS RÉPÉTER, NE PAS CITER, sert uniquement +à orienter le ton de ta réponse) : +- emotion : ... +- besoin : ... +- approche : ... +- axe : ... +``` + +### Thinker +``` +Tu es un psychologue clinicien experimente qui ecoute attentivement ton +patient. Avant de repondre, tu prepares une analyse silencieuse pour orienter +ta reponse. Produis UNIQUEMENT le bloc clinique au format strict suivant : +- emotion : +- besoin : +- approche : +- axe : +Pas d'introduction, pas d'explication. Juste les 4 lignes en francais. +/no_think +``` + +--- + +## 9. Risques connus & mitigations + +| Risque | Impact | Mitigation actuelle | Action future | +|---|---|---|---| +| Athena framework (system_server) tue Kazeia au-dessus de seuil RAM | Crash app | Débloat OEM réduit pression mémoire | Investiguer si killable plus profond | +| LMK Linux kill si MemAvailable < 1 GB | Crash app | Marge 3 GB en pic actuel | Lazy-load STT (-600 MB) | +| Speaker régurgite bullets cascade | Phrases bizarres lues par TTS | Bullets en system prompt + interdiction | Fine-tune/instruction tuning | +| Thinker Guard-0.6B mal calibré pour analyse FR | Bullets imprécis | Speaker compense par sa robustesse | Évaluer Qwen3-1.7B en remplacement | +| Whisper hallucinations sur silence | "que tu entends ce que je dis ici." apparaît au repos | Ignorer streaming results vides | Filtre RMS/longueur audio min | +| Orbe ne change pas (avant fix) | UX dégradée | `Paint.shader = null` dans drawThinking | OK | +| Symlinks `.pte` invisibles app | Mode perroquet | Fichiers réguliers (cp pas ln) | OK, mémoire feedback enregistrée | +| OOM 2× ExecuTorch HTP 4B | Crash | Thinker = Guard-0.6B (5.1 GB ION total OK) | Plus petit Thinker possible | + +--- + +## 10. Optimisations identifiées (memory profiling 2026-05-04) + +| Optim | Effort | Gain idle | Gain pic | Risque | Recommandation | +|---|---|---|---|---|---| +| Strip jniLibs Hexagon non-V79 | 30 min | 50 MB code | 50 MB | Bas | **Tier 1** | +| Lazy-load Whisper STT | 1 h | 600 MB | 0 | Moyen (+0.5s 1ᵉʳ PTT) | **Tier 1** | +| Trim Java heap (logs StateFlow) | 1 h | 100-200 MB | 100 MB | Bas | **Tier 1** | +| Release Thinker entre tours | 2 h | 0 | -700 MB | Moyen (+1.5s 2ᵉ PTT) | Tier 2 | +| Talker GGUF Q8_0 (vs f16) | 1 j | 400 MB | 400 MB | Haut (qualité audio) | Tier 2 | +| Decoder Q4 sélectif | 2-3 j | 200 MB | 200 MB | Haut | Tier 3 | +| Re-export Speaker plus petit | 5-10 j | 500-1 000 MB | similar | Très haut | Tier 3 | + +--- + +## 11. Roadmap + +### Court terme (1-2 semaines) +- [ ] Appliquer Tier 1 mémoire (strip libs + lazy STT + trim heap) +- [ ] Évaluation qualitative cascade FR sur corpus thérapeutique + 20+ prompts de référence +- [ ] Comparer Guard-0.6B vs Qwen3-1.7B comme Thinker (export et bench) +- [ ] Documenter procédure réinstallation tablette neuve (incl. débloat) + +### Moyen terme (1-2 mois) +- [ ] Rewrite Kazeia Phase 1-5 (architecture propre Kotlin/JNI) +- [ ] Voice clone .bin generator on-device (POC) +- [ ] Persistance conversation history (Room database) +- [ ] UI polish (mode conversation libre vs PTT) +- [ ] Voice commands étendus + +### Long terme (3+ mois) +- [ ] Fine-tuning thérapeutique français spécifique (LoRA) +- [ ] Migration runtime briques (POC dans `/opt/htp-runtime/`) — alternative + à ExecuTorch/QNN sur le long terme pour autonomie +- [ ] Évaluation clinique avec professionnels santé mentale + +--- + +## 12. Références — Fichiers clés + +### Code source app +- `/opt/Kazeia/kazeia-android/` — projet Gradle 8.12 / NDK r27d / API 36 + - `app/src/main/java/com/kazeia/service/KazeiaService.kt` — orchestrateur + - `app/src/main/java/com/kazeia/llm/ExecuTorchLlmEngine.kt` — wrapper LLM + - `app/src/main/java/com/kazeia/tts/Qwen3TtsEngine.kt` — TTS + - `app/src/main/java/com/kazeia/stt/WhisperHybridEngine.kt` — STT + - `app/src/main/java/com/kazeia/ui/AudioVisualizerView.kt` — orbe + - `app/src/main/jni/` — bridges JNI custom (libtts_*, libmel_extractor) + +### Scripts et outils +- `/opt/Kazeia/export_guard_06b.sh` — export Qwen3Guard-Gen-0.6B → `.pte` +- `/opt/Kazeia/kazeia-debloat/` — débloat tablette +- `/opt/Kazeia/kazeia-tablet-backup/` — backup + restore +- `/opt/Kazeia/executorch/examples/qualcomm/oss_scripts/llama/` — pipeline export QNN HTP + +### Modèles tablette +- `/data/local/tmp/kazeia-et/hybrid_llama_qnn_4b.pte` — Speaker +- `/data/local/tmp/kazeia-et/hybrid_llama_qnn_guard06b.pte` — Thinker (nouveau) +- `/data/local/tmp/kazeia-et/tokenizer.json` / `tokenizer_guard06b.json` +- `/data/local/tmp/kazeia/models/` — Talker, CP, Decoder GGUF, Whisper QAIRT +- `/data/local/tmp/kazeia/voix/` — voix de référence + +### Documents de référence +- `/opt/Kazeia/FROZEN.md` — stack figée +- `/opt/Kazeia/kazeia-tts-perf-validated/` — TTS RTF<1 validé +- `/opt/Kazeia/kazeia-stt-perf-validated/` — STT RTF 0.4 validé +- `/opt/Kazeia/kazeia-debloat/README.md` — procédure tablette + +### Mémoire persistante (sessions Claude Code) +- `/home/alf/.claude/projects/-opt-Kazeia/memory/MEMORY.md` — index +- 50+ entrées documentant chaque décision technique majeure + +--- + +## 13. Conclusion + +Kazeia tourne **end-to-end stable** sur OnePlus Pad 3 sans connexion internet, +sans root utilisateur, avec des latences acceptables pour une conversation +thérapeutique (3-5 s entre tours). L'architecture cascade Speaker+Thinker +in-process via ExecuTorch QNN HTP est validée et reproductible. + +Les principaux **blocages historiques** ont été résolus : DSP contention, +OOM, SELinux symlinks, watchdog OEM, accumulation KV cache. La stack est +documentée et figée, avec backup + procédure de restore complets. + +Les **axes d'amélioration restants** sont incrémentaux (mémoire, qualité +Thinker, UX) plutôt que structurels. Le pipeline est prêt pour **évaluation +clinique** sur corpus de référence en français. + +--- +*Rapport généré 2026-05-04. Auteur : Richard Loyer + assistance IA.* diff --git a/scripts/export_decoder_onnx.py b/scripts/export_decoder_onnx.py new file mode 100644 index 0000000..eae3dc3 --- /dev/null +++ b/scripts/export_decoder_onnx.py @@ -0,0 +1,366 @@ +#!/usr/bin/env python3 +""" +Export Qwen3-TTS speech decoder V2 to full ONNX for GPU QNN backend. +Exact architecture from state dict analysis: + pre_conv: Conv1d(512, 1024, kernel=3, pad=1) + pre_transformer: 8-layer transformer (dim=512, attn_dim=1024, heads=16, head_dim=64) + upsample: 2× ConvTranspose1d(1024,1024) + ConvNeXtV2 + BigVGAN: 7 blocks with Snake activation, upsample_rates=[8,5,4,3] +""" +import sys, os, math, warnings +sys.path.insert(0, "/opt/Kazeia/qnn_venv/lib/python3.10/site-packages") +warnings.filterwarnings("ignore") + +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np + +state = torch.load("/opt/Kazeia/models_qnn/qwen3-tts-native/speech_decoder_weights.pt", map_location="cpu") +print(f"Loaded {len(state)} tensors") + +OUT = "/opt/Kazeia/models_qnn/qwen3-tts-decoder-full-onnx" +for d in ["pre_conv", "preprocessor", "conv_decoder"]: + os.makedirs(f"{OUT}/{d}", exist_ok=True) + +SEQ = 60 + +# ===== RMSNorm ===== +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + def forward(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight + +# ===== DiT Layer ===== +class DiTLayer(nn.Module): + def __init__(self, dim=512, attn_dim=1024, n_heads=16, head_dim=64, ffn_dim=1024, rope_theta=10000.0): + super().__init__() + self.dim = dim + self.n_heads = n_heads + self.head_dim = head_dim + self.input_layernorm = RMSNorm(dim) + self.post_attention_layernorm = RMSNorm(dim) + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(dim, attn_dim, bias=False) + self.self_attn.k_proj = nn.Linear(dim, attn_dim, bias=False) + self.self_attn.v_proj = nn.Linear(dim, attn_dim, bias=False) + self.self_attn.o_proj = nn.Linear(attn_dim, dim, bias=False) + self.self_attn_layer_scale = nn.Module() + self.self_attn_layer_scale.scale = nn.Parameter(torch.ones(dim) * 0.01) + self.mlp = nn.Module() + self.mlp.gate_proj = nn.Linear(dim, ffn_dim, bias=False) + self.mlp.up_proj = nn.Linear(dim, ffn_dim, bias=False) + self.mlp.down_proj = nn.Linear(ffn_dim, dim, bias=False) + self.mlp_layer_scale = nn.Module() + self.mlp_layer_scale.scale = nn.Parameter(torch.ones(dim) * 0.01) + # Pre-compute RoPE + freqs = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) + t = torch.arange(512).float() + emb = torch.outer(t, freqs) + self.register_buffer("rope_cos", emb.cos()) + self.register_buffer("rope_sin", emb.sin()) + + def apply_rope(self, x, seq_len): + cos = self.rope_cos[:seq_len].unsqueeze(0).unsqueeze(0) # [1,1,T,D/2] + sin = self.rope_sin[:seq_len].unsqueeze(0).unsqueeze(0) + x1, x2 = x[..., ::2], x[..., 1::2] + out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) + return out.flatten(-2) + + def forward(self, x): + B, T, _ = x.shape + h = self.input_layernorm(x) + q = self.self_attn.q_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) + k = self.self_attn.k_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) + v = self.self_attn.v_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) + q = self.apply_rope(q, T) + k = self.apply_rope(k, T) + attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) + mask = torch.triu(torch.full((T, T), float('-inf'), device=x.device), diagonal=1) + attn = attn + mask.unsqueeze(0).unsqueeze(0) + attn = F.softmax(attn, dim=-1) + out = torch.matmul(attn, v).transpose(1, 2).contiguous().view(B, T, -1) + x = x + self.self_attn_layer_scale.scale * self.self_attn.o_proj(out) + h2 = self.post_attention_layernorm(x) + ffn = self.mlp.down_proj(F.silu(self.mlp.gate_proj(h2)) * self.mlp.up_proj(h2)) + x = x + self.mlp_layer_scale.scale * ffn + return x + +# ===== Pre-transformer ===== +class PreTransformer(nn.Module): + def __init__(self): + super().__init__() + self.input_proj = nn.Linear(1024, 512) + self.layers = nn.ModuleList([DiTLayer() for _ in range(8)]) + self.norm = RMSNorm(512) + self.output_proj = nn.Linear(512, 1024) + def forward(self, x): + x = x.transpose(1, 2) # [B,C,T] → [B,T,C] + x = self.input_proj(x) + for layer in self.layers: + x = layer(x) + x = self.norm(x) + x = self.output_proj(x) + return x.transpose(1, 2) # [B,T,C] → [B,C,T] + +# ===== ConvNeXt-V2 ===== +class ConvNeXtV2(nn.Module): + def __init__(self, dim=1024, int_dim=4096): + super().__init__() + self.dwconv = nn.Module() + self.dwconv.conv = nn.Conv1d(dim, dim, 7, padding=3, groups=dim) + self.norm = nn.LayerNorm(dim) + self.pwconv1 = nn.Linear(dim, int_dim) + self.pwconv2 = nn.Linear(int_dim, dim) + self.gamma = nn.Parameter(torch.ones(dim) * 0.01) + def forward(self, x): + r = x + x = self.dwconv.conv(x) + x = x.transpose(1, 2) + x = self.norm(x) + x = F.gelu(self.pwconv1(x)) + x = self.pwconv2(x) + x = (self.gamma * x).transpose(1, 2) + return r + x + +# ===== Upsample ===== +class Upsample(nn.Module): + def __init__(self): + super().__init__() + # upsampling_ratios=[2,2] from config + # ConvTranspose1d(1024,1024,kernel=2*ratio,stride=ratio,padding=ratio//2) + w0 = state["upsample.0.0.conv.weight"] + w1 = state["upsample.1.0.conv.weight"] + # For upsampling ratio r: kernel=2*r, stride=r, pad=r//2 + # But kernel=2 could mean ratio=1 or ratio=2 with pad adjustment + # From config: upsampling_ratios=[2,2], so stride=2 for both + # kernel=2, stride=2, padding=0 gives exact 2× upsample + self.conv0 = nn.ConvTranspose1d(w0.shape[0], w0.shape[1], w0.shape[2], stride=2, padding=0) + self.cnx0 = ConvNeXtV2() + self.conv1 = nn.ConvTranspose1d(w1.shape[0], w1.shape[1], w1.shape[2], stride=2, padding=0) + self.cnx1 = ConvNeXtV2() + print(f" Upsample 0: {w0.shape} stride=2") + print(f" Upsample 1: {w1.shape} stride=2") + def forward(self, x): + x = self.conv0(x) + x = self.cnx0(x) + x = self.conv1(x) + x = self.cnx1(x) + return x + +# ===== SnakeBeta ===== +class SnakeBeta(nn.Module): + def __init__(self, channels): + super().__init__() + # Stored as [channels] in checkpoint, reshaped to [1,channels,1] in forward + self.alpha = nn.Parameter(torch.ones(channels)) + self.beta = nn.Parameter(torch.ones(channels)) + def forward(self, x): + a = self.alpha.view(1, -1, 1) + b = self.beta.view(1, -1, 1) + return x + (1.0 / (b + 1e-9)) * (torch.sin(a * x) ** 2) + +# ===== AMPBlock (BigVGAN resblock) ===== +class AMPBlock(nn.Module): + def __init__(self, channels, kernel_size=7): + super().__init__() + self.act1 = SnakeBeta(channels) + self.act2 = SnakeBeta(channels) + self.conv1 = nn.Module() + self.conv1.conv = nn.Conv1d(channels, channels, kernel_size, padding=kernel_size//2) + self.conv2 = nn.Module() + self.conv2.conv = nn.Conv1d(channels, channels, kernel_size, padding=1) + def forward(self, x): + h = self.act1(x) + h = self.conv1.conv(h) + h = self.act2(h) + h = self.conv2.conv(h) + return x + h + +# ===== UpsampleBigVGANBlock ===== +class UpsampleBigVGANBlock(nn.Module): + def __init__(self, ch_in, ch_out, kernel, rate, pad, block_idx): + super().__init__() + self.block = nn.ModuleList() + self.block.append(SnakeBeta(ch_in)) + self.block.append(nn.ConvTranspose1d(ch_in, ch_out, kernel, stride=rate, padding=pad)) + for j in range(3): + amp_w1 = state[f"decoder.{block_idx}.block.{j+2}.conv1.conv.weight"] + amp_w2 = state[f"decoder.{block_idx}.block.{j+2}.conv2.conv.weight"] + amp = AMPBlock(amp_w1.shape[0], amp_w1.shape[2]) + amp.conv2.conv = nn.Conv1d(amp_w2.shape[1], amp_w2.shape[0], amp_w2.shape[2], padding=amp_w2.shape[2]//2) + self.block.append(amp) + def forward(self, x): + # block[0] = SnakeBeta, block[1] = ConvTranspose1d, block[2-4] = AMPBlocks + x = self.block[0](x) + x = self.block[1](x) + for i in range(2, 5): + x = self.block[i](x) + return x + +# ===== BigVGAN ===== +class BigVGAN(nn.Module): + def __init__(self): + super().__init__() + # decoder.0: initial conv [1536, 1024, 7] + w0 = state["decoder.0.conv.weight"] + ch = w0.shape[0] # 1536 + self.initial_conv = nn.Module() + self.initial_conv.conv = nn.Conv1d(w0.shape[1], ch, w0.shape[2], padding=w0.shape[2]//2) + + # decoder.1-4: upsample blocks + # Each: SnakeBeta(ch_in) + ConvTranspose1d(ch_in, ch_out) + 3×AMPBlock(ch_out) + self.blocks = nn.ModuleList() + for i in range(4): + block_idx = i + 1 + conv_w = state[f"decoder.{block_idx}.block.1.conv.weight"] + ch_in = conv_w.shape[0] # input channels for this ConvTranspose1d (= ch from snake) + ch_out = conv_w.shape[1] # output channels (Note: ConvTranspose1d weight is [ch_in, ch_out, K]) + kernel = conv_w.shape[2] + # Determine stride and padding from kernel and upsample rate + # The upsample rates are [8,5,4,3], kernel = rate*2 + rate = [8, 5, 4, 3][i] + pad = (kernel - rate) // 2 + + blk = UpsampleBigVGANBlock(ch_in, ch_out, kernel, rate, pad, block_idx) + self.blocks.append(blk) + ch = ch_out + + # decoder.5: final SnakeBeta + self.final_act = SnakeBeta(ch) + # decoder.6: final conv to audio + w6 = state["decoder.6.conv.weight"] + self.final_conv = nn.Module() + self.final_conv.conv = nn.Conv1d(w6.shape[1], w6.shape[0], w6.shape[2], padding=w6.shape[2]//2) + + def forward(self, x): + x = self.initial_conv.conv(x) + for blk in self.blocks: + for m in blk.block: + x = m(x) + x = self.final_act(x) + x = self.final_conv.conv(x) + x = torch.tanh(x) + return x + +# ===== Build & Load ===== +print("\n--- Building pre_conv ---") +pre_conv = nn.Conv1d(512, 1024, 3, padding=1) +pre_conv.weight.data = state["pre_conv.conv.weight"] +pre_conv.bias.data = state["pre_conv.conv.bias"] +pre_conv.eval() +dummy = torch.randn(1, 512, SEQ) +with torch.no_grad(): + pc = pre_conv(dummy) +print(f"pre_conv: [1,512,{SEQ}] → {list(pc.shape)}") + +print("\n--- Building pre_transformer ---") +pt = PreTransformer() +pt_state = {k[len("pre_transformer."):]: v for k, v in state.items() if k.startswith("pre_transformer.")} +m, u = pt.load_state_dict(pt_state, strict=False) +print(f"pre_transformer load: missing={len(m)}, unexpected={len(u)}") +if m: print(f" Missing: {m[:3]}") +pt.eval() +with torch.no_grad(): + pt_out = pt(pc) +print(f"pre_transformer: {list(pc.shape)} → {list(pt_out.shape)}") + +print("\n--- Building upsample ---") +up = Upsample() +up_map = {} +for k, v in state.items(): + if k.startswith("upsample."): + nk = k[len("upsample."):] + # nk = "0.0.conv.weight" → conv0.weight, "0.1.dwconv.conv.weight" → cnx0.dwconv.conv.weight + parts = nk.split(".", 2) + idx = parts[0] # 0 or 1 + sub = parts[1] # 0 (conv) or 1 (convnext) + rest = parts[2] if len(parts) > 2 else "" + if sub == "0": + # conv transpose: upsample.0.0.conv.weight → conv0.weight (remove ".conv" layer) + if rest.startswith("conv."): + up_map[f"conv{idx}.{rest[5:]}"] = v + else: + up_map[f"conv{idx}.{rest}"] = v + else: + up_map[f"cnx{idx}.{rest}"] = v +m, u = up.load_state_dict(up_map, strict=False) +print(f"upsample load: missing={len(m)}, unexpected={len(u)}") +if m: print(f" Missing: {m[:3]}") +up.eval() +with torch.no_grad(): + up_out = up(pt_out) +print(f"upsample: {list(pt_out.shape)} → {list(up_out.shape)}") + +print("\n--- Building BigVGAN ---") +bv = BigVGAN() +bv_map = {} +for k, v in state.items(): + if k.startswith("decoder."): + nk = k[len("decoder."):] + parts = nk.split(".", 1) + idx = int(parts[0]) + rest = parts[1] if len(parts) > 1 else "" + if idx == 0: + # decoder.0.conv.weight → initial_conv.conv.weight + bv_map[f"initial_conv.{rest}"] = v + elif 1 <= idx <= 4: + # decoder.{i}.block.1.conv.weight → blocks.{i-1}.block.1.weight (ConvTranspose1d, no ".conv") + new_rest = rest + if rest.startswith("block.1.conv."): + # ConvTranspose1d: remove extra ".conv" layer + new_rest = "block.1." + rest[len("block.1.conv."):] + bv_map[f"blocks.{idx-1}.{new_rest}"] = v + elif idx == 5: + bv_map[f"final_act.{rest}"] = v + elif idx == 6: + bv_map[f"final_conv.{rest}"] = v +m, u = bv.load_state_dict(bv_map, strict=False) +print(f"BigVGAN load: missing={len(m)}, unexpected={len(u)}") +if m: print(f" Missing: {m[:5]}") +if u: print(f" Unexpected: {u[:5]}") +bv.eval() +with torch.no_grad(): + audio = bv(up_out) +print(f"BigVGAN: {list(up_out.shape)} → {list(audio.shape)}") + +# ===== Export ===== +print("\n=== Exporting pre_conv ===") +torch.onnx.export(pre_conv, dummy, f"{OUT}/pre_conv/model.onnx", + input_names=["x"], output_names=["pre_conv_out"], opset_version=17) +print(f" {os.path.getsize(f'{OUT}/pre_conv/model.onnx')/1024:.0f} KB") + +class Preprocessor(nn.Module): + def __init__(self, pt, up): + super().__init__(); self.pt = pt; self.up = up + def forward(self, x): + return self.up(self.pt(x)) + +print("\n=== Exporting preprocessor ===") +prep = Preprocessor(pt, up).eval() +with torch.no_grad(): + prep_out = prep(pc) +torch.onnx.export(prep, pc, f"{OUT}/preprocessor/model.onnx", + input_names=["pre_conv_out"], output_names=["hidden"], opset_version=17) +print(f" {os.path.getsize(f'{OUT}/preprocessor/model.onnx')/1024/1024:.1f} MB") + +print("\n=== Exporting conv_decoder ===") +torch.onnx.export(bv, up_out, f"{OUT}/conv_decoder/model.onnx", + input_names=["hidden"], output_names=["audio"], opset_version=17) +print(f" {os.path.getsize(f'{OUT}/conv_decoder/model.onnx')/1024/1024:.1f} MB") + +# ===== Validate ===== +print("\n=== Validation ===") +import onnxruntime as ort +for name, inp, ref in [("pre_conv", dummy, pc), ("preprocessor", pc, prep_out), ("conv_decoder", up_out, audio)]: + sess = ort.InferenceSession(f"{OUT}/{name}/model.onnx") + o = sess.run(None, {sess.get_inputs()[0].name: inp.detach().numpy()})[0] + d = np.max(np.abs(o - ref.detach().numpy())) + print(f" {name}: max_diff={d:.6f} shape={o.shape}") + +print("\nTotal:") +for n in ["pre_conv", "preprocessor", "conv_decoder"]: + print(f" {n}: {os.path.getsize(f'{OUT}/{n}/model.onnx')/1024/1024:.1f} MB") diff --git a/scripts/export_decoder_onnx_dynt.py b/scripts/export_decoder_onnx_dynt.py new file mode 100644 index 0000000..3bd052a --- /dev/null +++ b/scripts/export_decoder_onnx_dynt.py @@ -0,0 +1,399 @@ +#!/usr/bin/env python3 +""" +Export Qwen3-TTS speech decoder V2 to full ONNX for GPU QNN backend. +Exact architecture from state dict analysis: + pre_conv: Conv1d(512, 1024, kernel=3, pad=1) + pre_transformer: 8-layer transformer (dim=512, attn_dim=1024, heads=16, head_dim=64) + upsample: 2× ConvTranspose1d(1024,1024) + ConvNeXtV2 + BigVGAN: 7 blocks with Snake activation, upsample_rates=[8,5,4,3] +""" +import sys, os, math, warnings +sys.path.insert(0, "/opt/Kazeia/qnn_venv/lib/python3.10/site-packages") +warnings.filterwarnings("ignore") + +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np + +state = torch.load("/opt/Kazeia/to_delete/models_qnn/qwen3-tts-native/speech_decoder_weights.pt", map_location="cpu") +print(f"Loaded {len(state)} tensors") + +OUT = "/tmp/kazeia_decoder_dynt" +for d in ["pre_conv", "preprocessor", "conv_decoder"]: + os.makedirs(f"{OUT}/{d}", exist_ok=True) + +SEQ = 60 + +# ===== RMSNorm ===== +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + def forward(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight + +# ===== DiT Layer ===== +class DiTLayer(nn.Module): + def __init__(self, dim=512, attn_dim=1024, n_heads=16, head_dim=64, ffn_dim=1024, rope_theta=10000.0): + super().__init__() + self.dim = dim + self.n_heads = n_heads + self.head_dim = head_dim + self.input_layernorm = RMSNorm(dim) + self.post_attention_layernorm = RMSNorm(dim) + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(dim, attn_dim, bias=False) + self.self_attn.k_proj = nn.Linear(dim, attn_dim, bias=False) + self.self_attn.v_proj = nn.Linear(dim, attn_dim, bias=False) + self.self_attn.o_proj = nn.Linear(attn_dim, dim, bias=False) + self.self_attn_layer_scale = nn.Module() + self.self_attn_layer_scale.scale = nn.Parameter(torch.ones(dim) * 0.01) + self.mlp = nn.Module() + self.mlp.gate_proj = nn.Linear(dim, ffn_dim, bias=False) + self.mlp.up_proj = nn.Linear(dim, ffn_dim, bias=False) + self.mlp.down_proj = nn.Linear(ffn_dim, dim, bias=False) + self.mlp_layer_scale = nn.Module() + self.mlp_layer_scale.scale = nn.Parameter(torch.ones(dim) * 0.01) + # Pre-compute RoPE + freqs = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) + t = torch.arange(512).float() + emb = torch.outer(t, freqs) + self.register_buffer("rope_cos", emb.cos()) + self.register_buffer("rope_sin", emb.sin()) + + def apply_rope(self, x, seq_len): + cos = self.rope_cos[:seq_len].unsqueeze(0).unsqueeze(0) # [1,1,T,D/2] + sin = self.rope_sin[:seq_len].unsqueeze(0).unsqueeze(0) + x1, x2 = x[..., ::2], x[..., 1::2] + out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) + return out.flatten(-2) + + def forward(self, x): + B, T, _ = x.shape + h = self.input_layernorm(x) + q = self.self_attn.q_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) + k = self.self_attn.k_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) + v = self.self_attn.v_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) + q = self.apply_rope(q, T) + k = self.apply_rope(k, T) + attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) + mask = torch.triu(torch.full((T, T), float('-inf'), device=x.device), diagonal=1) + attn = attn + mask.unsqueeze(0).unsqueeze(0) + attn = F.softmax(attn, dim=-1) + out = torch.matmul(attn, v).transpose(1, 2).contiguous().view(B, T, -1) + x = x + self.self_attn_layer_scale.scale * self.self_attn.o_proj(out) + h2 = self.post_attention_layernorm(x) + ffn = self.mlp.down_proj(F.silu(self.mlp.gate_proj(h2)) * self.mlp.up_proj(h2)) + x = x + self.mlp_layer_scale.scale * ffn + return x + +# ===== Pre-transformer ===== +class PreTransformer(nn.Module): + def __init__(self): + super().__init__() + self.input_proj = nn.Linear(1024, 512) + self.layers = nn.ModuleList([DiTLayer() for _ in range(8)]) + self.norm = RMSNorm(512) + self.output_proj = nn.Linear(512, 1024) + def forward(self, x): + x = x.transpose(1, 2) # [B,C,T] → [B,T,C] + x = self.input_proj(x) + for layer in self.layers: + x = layer(x) + x = self.norm(x) + x = self.output_proj(x) + return x.transpose(1, 2) # [B,T,C] → [B,C,T] + +# ===== ConvNeXt-V2 ===== +class ConvNeXtV2(nn.Module): + def __init__(self, dim=1024, int_dim=4096): + super().__init__() + self.dwconv = nn.Module() + self.dwconv.conv = nn.Conv1d(dim, dim, 7, padding=3, groups=dim) + self.norm = nn.LayerNorm(dim) + self.pwconv1 = nn.Linear(dim, int_dim) + self.pwconv2 = nn.Linear(int_dim, dim) + self.gamma = nn.Parameter(torch.ones(dim) * 0.01) + def forward(self, x): + r = x + x = self.dwconv.conv(x) + x = x.transpose(1, 2) + x = self.norm(x) + x = F.gelu(self.pwconv1(x)) + x = self.pwconv2(x) + x = (self.gamma * x).transpose(1, 2) + return r + x + +# ===== Upsample ===== +class Upsample(nn.Module): + def __init__(self): + super().__init__() + # upsampling_ratios=[2,2] from config + # ConvTranspose1d(1024,1024,kernel=2*ratio,stride=ratio,padding=ratio//2) + w0 = state["upsample.0.0.conv.weight"] + w1 = state["upsample.1.0.conv.weight"] + # For upsampling ratio r: kernel=2*r, stride=r, pad=r//2 + # But kernel=2 could mean ratio=1 or ratio=2 with pad adjustment + # From config: upsampling_ratios=[2,2], so stride=2 for both + # kernel=2, stride=2, padding=0 gives exact 2× upsample + self.conv0 = nn.ConvTranspose1d(w0.shape[0], w0.shape[1], w0.shape[2], stride=2, padding=0) + self.cnx0 = ConvNeXtV2() + self.conv1 = nn.ConvTranspose1d(w1.shape[0], w1.shape[1], w1.shape[2], stride=2, padding=0) + self.cnx1 = ConvNeXtV2() + print(f" Upsample 0: {w0.shape} stride=2") + print(f" Upsample 1: {w1.shape} stride=2") + def forward(self, x): + x = self.conv0(x) + x = self.cnx0(x) + x = self.conv1(x) + x = self.cnx1(x) + return x + +# ===== SnakeBeta ===== +class SnakeBeta(nn.Module): + def __init__(self, channels): + super().__init__() + # Stored as [channels] in checkpoint, reshaped to [1,channels,1] in forward + self.alpha = nn.Parameter(torch.ones(channels)) + self.beta = nn.Parameter(torch.ones(channels)) + def forward(self, x): + a = self.alpha.view(1, -1, 1) + b = self.beta.view(1, -1, 1) + return x + (1.0 / (b + 1e-9)) * (torch.sin(a * x) ** 2) + +# ===== AMPBlock (BigVGAN resblock) ===== +class AMPBlock(nn.Module): + def __init__(self, channels, kernel_size=7): + super().__init__() + self.act1 = SnakeBeta(channels) + self.act2 = SnakeBeta(channels) + self.conv1 = nn.Module() + self.conv1.conv = nn.Conv1d(channels, channels, kernel_size, padding=kernel_size//2) + self.conv2 = nn.Module() + self.conv2.conv = nn.Conv1d(channels, channels, kernel_size, padding=1) + def forward(self, x): + h = self.act1(x) + h = self.conv1.conv(h) + h = self.act2(h) + h = self.conv2.conv(h) + return x + h + +# ===== UpsampleBigVGANBlock ===== +class UpsampleBigVGANBlock(nn.Module): + def __init__(self, ch_in, ch_out, kernel, rate, pad, block_idx): + super().__init__() + self.block = nn.ModuleList() + self.block.append(SnakeBeta(ch_in)) + self.block.append(nn.ConvTranspose1d(ch_in, ch_out, kernel, stride=rate, padding=pad)) + for j in range(3): + amp_w1 = state[f"decoder.{block_idx}.block.{j+2}.conv1.conv.weight"] + amp_w2 = state[f"decoder.{block_idx}.block.{j+2}.conv2.conv.weight"] + amp = AMPBlock(amp_w1.shape[0], amp_w1.shape[2]) + amp.conv2.conv = nn.Conv1d(amp_w2.shape[1], amp_w2.shape[0], amp_w2.shape[2], padding=amp_w2.shape[2]//2) + self.block.append(amp) + def forward(self, x): + # block[0] = SnakeBeta, block[1] = ConvTranspose1d, block[2-4] = AMPBlocks + x = self.block[0](x) + x = self.block[1](x) + for i in range(2, 5): + x = self.block[i](x) + return x + +# ===== BigVGAN ===== +class BigVGAN(nn.Module): + def __init__(self): + super().__init__() + # decoder.0: initial conv [1536, 1024, 7] + w0 = state["decoder.0.conv.weight"] + ch = w0.shape[0] # 1536 + self.initial_conv = nn.Module() + self.initial_conv.conv = nn.Conv1d(w0.shape[1], ch, w0.shape[2], padding=w0.shape[2]//2) + + # decoder.1-4: upsample blocks + # Each: SnakeBeta(ch_in) + ConvTranspose1d(ch_in, ch_out) + 3×AMPBlock(ch_out) + self.blocks = nn.ModuleList() + for i in range(4): + block_idx = i + 1 + conv_w = state[f"decoder.{block_idx}.block.1.conv.weight"] + ch_in = conv_w.shape[0] # input channels for this ConvTranspose1d (= ch from snake) + ch_out = conv_w.shape[1] # output channels (Note: ConvTranspose1d weight is [ch_in, ch_out, K]) + kernel = conv_w.shape[2] + # Determine stride and padding from kernel and upsample rate + # The upsample rates are [8,5,4,3], kernel = rate*2 + rate = [8, 5, 4, 3][i] + pad = (kernel - rate) // 2 + + blk = UpsampleBigVGANBlock(ch_in, ch_out, kernel, rate, pad, block_idx) + self.blocks.append(blk) + ch = ch_out + + # decoder.5: final SnakeBeta + self.final_act = SnakeBeta(ch) + # decoder.6: final conv to audio + w6 = state["decoder.6.conv.weight"] + self.final_conv = nn.Module() + self.final_conv.conv = nn.Conv1d(w6.shape[1], w6.shape[0], w6.shape[2], padding=w6.shape[2]//2) + + def forward(self, x): + x = self.initial_conv.conv(x) + for blk in self.blocks: + for m in blk.block: + x = m(x) + x = self.final_act(x) + x = self.final_conv.conv(x) + x = torch.tanh(x) + return x + +# ===== Build & Load ===== +print("\n--- Building pre_conv ---") +pre_conv = nn.Conv1d(512, 1024, 3, padding=1) +pre_conv.weight.data = state["pre_conv.conv.weight"] +pre_conv.bias.data = state["pre_conv.conv.bias"] +pre_conv.eval() +dummy = torch.randn(1, 512, SEQ) +with torch.no_grad(): + pc = pre_conv(dummy) +print(f"pre_conv: [1,512,{SEQ}] → {list(pc.shape)}") + +print("\n--- Building pre_transformer ---") +pt = PreTransformer() +pt_state = {k[len("pre_transformer."):]: v for k, v in state.items() if k.startswith("pre_transformer.")} +m, u = pt.load_state_dict(pt_state, strict=False) +print(f"pre_transformer load: missing={len(m)}, unexpected={len(u)}") +if m: print(f" Missing: {m[:3]}") +pt.eval() +with torch.no_grad(): + pt_out = pt(pc) +print(f"pre_transformer: {list(pc.shape)} → {list(pt_out.shape)}") + +print("\n--- Building upsample ---") +up = Upsample() +up_map = {} +for k, v in state.items(): + if k.startswith("upsample."): + nk = k[len("upsample."):] + # nk = "0.0.conv.weight" → conv0.weight, "0.1.dwconv.conv.weight" → cnx0.dwconv.conv.weight + parts = nk.split(".", 2) + idx = parts[0] # 0 or 1 + sub = parts[1] # 0 (conv) or 1 (convnext) + rest = parts[2] if len(parts) > 2 else "" + if sub == "0": + # conv transpose: upsample.0.0.conv.weight → conv0.weight (remove ".conv" layer) + if rest.startswith("conv."): + up_map[f"conv{idx}.{rest[5:]}"] = v + else: + up_map[f"conv{idx}.{rest}"] = v + else: + up_map[f"cnx{idx}.{rest}"] = v +m, u = up.load_state_dict(up_map, strict=False) +print(f"upsample load: missing={len(m)}, unexpected={len(u)}") +if m: print(f" Missing: {m[:3]}") +up.eval() +with torch.no_grad(): + up_out = up(pt_out) +print(f"upsample: {list(pt_out.shape)} → {list(up_out.shape)}") + +print("\n--- Building BigVGAN ---") +bv = BigVGAN() +bv_map = {} +for k, v in state.items(): + if k.startswith("decoder."): + nk = k[len("decoder."):] + parts = nk.split(".", 1) + idx = int(parts[0]) + rest = parts[1] if len(parts) > 1 else "" + if idx == 0: + # decoder.0.conv.weight → initial_conv.conv.weight + bv_map[f"initial_conv.{rest}"] = v + elif 1 <= idx <= 4: + # decoder.{i}.block.1.conv.weight → blocks.{i-1}.block.1.weight (ConvTranspose1d, no ".conv") + new_rest = rest + if rest.startswith("block.1.conv."): + # ConvTranspose1d: remove extra ".conv" layer + new_rest = "block.1." + rest[len("block.1.conv."):] + bv_map[f"blocks.{idx-1}.{new_rest}"] = v + elif idx == 5: + bv_map[f"final_act.{rest}"] = v + elif idx == 6: + bv_map[f"final_conv.{rest}"] = v +m, u = bv.load_state_dict(bv_map, strict=False) +print(f"BigVGAN load: missing={len(m)}, unexpected={len(u)}") +if m: print(f" Missing: {m[:5]}") +if u: print(f" Unexpected: {u[:5]}") +bv.eval() +with torch.no_grad(): + audio = bv(up_out) +print(f"BigVGAN: {list(up_out.shape)} → {list(audio.shape)}") + +# ===== Export avec dynamic_axes pour la dim T ===== +# pre_conv input [1, 512, T] → dim 2 dynamique +# pre_conv output [1, 1024, T] → dim 2 dynamique +# preprocessor input [1, T, 1024] → dim 1 dynamique (note : preprocessor inclut up qui upsample 2x donc T_out = 2*T_in) +# preprocessor output [1, 2T, 1024] → dim 1 dynamique +# conv_decoder input [1, 2T, 1024] → dim 1 dynamique +# conv_decoder output [1, 1, 2T*960] → dim 2 dynamique (BigVGAN upsample x960) +print("\n=== Exporting pre_conv (dyn T) ===") +torch.onnx.export(pre_conv, dummy, f"{OUT}/pre_conv/model.onnx", + input_names=["x"], output_names=["pre_conv_out"], opset_version=17, + dynamic_axes={"x": {2: "T"}, "pre_conv_out": {2: "T"}}) +print(f" {os.path.getsize(f'{OUT}/pre_conv/model.onnx')/1024:.0f} KB") + +class Preprocessor(nn.Module): + def __init__(self, pt, up): + super().__init__(); self.pt = pt; self.up = up + def forward(self, x): + return self.up(self.pt(x)) + +print("\n=== Exporting preprocessor (dyn T) ===") +prep = Preprocessor(pt, up).eval() +with torch.no_grad(): + prep_out = prep(pc) +# input shape [1, 1024, T] (dim 2 = T), output shape [1, 1024, 4T] (upsample x4) +torch.onnx.export(prep, pc, f"{OUT}/preprocessor/model.onnx", + input_names=["pre_conv_out"], output_names=["hidden"], opset_version=17, + dynamic_axes={"pre_conv_out": {2: "T"}, "hidden": {2: "T_up"}}) +print(f" {os.path.getsize(f'{OUT}/preprocessor/model.onnx')/1024/1024:.1f} MB") + +print("\n=== Exporting conv_decoder (dyn T) ===") +# input shape [1, 1024, 4T] (dim 2 = 4T), output shape [1, 1, 4T*480] (BigVGAN upsample x480) +torch.onnx.export(bv, up_out, f"{OUT}/conv_decoder/model.onnx", + input_names=["hidden"], output_names=["audio"], opset_version=17, + dynamic_axes={"hidden": {2: "T_up"}, "audio": {2: "T_audio"}}) +print(f" {os.path.getsize(f'{OUT}/conv_decoder/model.onnx')/1024/1024:.1f} MB") + +# ===== Validate à T=60 (référence) puis T=128 (dynamic) ===== +print("\n=== Validation T=60 (réf) ===") +import onnxruntime as ort +for name, inp, ref in [("pre_conv", dummy, pc), ("preprocessor", pc, prep_out), ("conv_decoder", up_out, audio)]: + sess = ort.InferenceSession(f"{OUT}/{name}/model.onnx") + o = sess.run(None, {sess.get_inputs()[0].name: inp.detach().numpy()})[0] + d = np.max(np.abs(o - ref.detach().numpy())) + print(f" {name}: max_diff={d:.6f} shape={o.shape}") + +print("\n=== Test T=128 dynamique (sanity) ===") +# Build dummy input avec T=128 et passer dans la chaîne entière +T_test = 128 +d_test = torch.randn(1, 512, T_test) +with torch.no_grad(): + pc_test = pre_conv(d_test) + prep_test = Preprocessor(pt, up)(pc_test) + bv_test = bv(prep_test) +print(f" PyTorch ref: pre_conv={list(pc_test.shape)} prep={list(prep_test.shape)} bv={list(bv_test.shape)}") + +sess_pc = ort.InferenceSession(f"{OUT}/pre_conv/model.onnx") +o_pc = sess_pc.run(None, {sess_pc.get_inputs()[0].name: d_test.numpy()})[0] +sess_pp = ort.InferenceSession(f"{OUT}/preprocessor/model.onnx") +o_pp = sess_pp.run(None, {sess_pp.get_inputs()[0].name: o_pc})[0] +sess_cd = ort.InferenceSession(f"{OUT}/conv_decoder/model.onnx") +o_cd = sess_cd.run(None, {sess_cd.get_inputs()[0].name: o_pp})[0] +print(f" ONNX: pre_conv={o_pc.shape} prep={o_pp.shape} bv={o_cd.shape}") +diff_pc = np.max(np.abs(o_pc - pc_test.numpy())) +diff_pp = np.max(np.abs(o_pp - prep_test.numpy())) +diff_cd = np.max(np.abs(o_cd - bv_test.numpy())) +print(f" Max diff vs PyTorch : pre_conv={diff_pc:.6f} prep={diff_pp:.6f} conv_decoder={diff_cd:.6f}") + +print("\nTotal:") +for n in ["pre_conv", "preprocessor", "conv_decoder"]: + print(f" {n}: {os.path.getsize(f'{OUT}/{n}/model.onnx')/1024/1024:.1f} MB") diff --git a/scripts/export_talker_onnx.py b/scripts/export_talker_onnx.py new file mode 100644 index 0000000..3d5c0dd --- /dev/null +++ b/scripts/export_talker_onnx.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +""" +Export Qwen3-TTS talker KV-cache to ONNX with M-RoPE. +Following TTS_CALIBRATION_GUIDE.md exactly: + Inputs: [emb(1,1,1024), mask(1,1,1,200), cos(1,1,128), sin(1,1,128), 56×kv(1,8,199,128)] + Outputs: [hidden(1,1,1024), logits(1,1,3072), 56×kv(1,8,200,128)] +""" +import sys +sys.path.insert(0, "/opt/Kazeia/qnn_venv/lib/python3.10/site-packages") +import warnings; warnings.filterwarnings("ignore") +import torch, torch.nn as nn, torch.nn.functional as F +import numpy as np, os, math + +N_LAYERS = 28; DIM = 1024; N_HEADS = 16; N_KV_HEADS = 8; HEAD_DIM = 128 +KV_LEN = 199; VOCAB = 3072; FFN_DIM = 3072 +GQA_GROUPS = N_HEADS // N_KV_HEADS # 2 + +print("Loading talker weights...") +state = torch.load("/opt/Kazeia/models_qnn/qwen3-tts-export/qwen3_tts_talker.pth", + map_location="cpu", weights_only=False) + +class RMSNorm(nn.Module): + def __init__(self, dim): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + def forward(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.weight + +def apply_interleaved_rope(x, cos, sin): + """Apply interleaved RoPE: x[..., ::2] and x[..., 1::2] rotated by cos/sin.""" + # x: [B, heads, 1, head_dim], cos/sin: [1, 1, 128] → use first 64 pairs + half = x.shape[-1] // 2 + cos = cos[..., :half] # [1,1,64] + sin = sin[..., :half] + x1 = x[..., ::2] # [B, heads, 1, 64] + x2 = x[..., 1::2] + cos = cos.unsqueeze(1) # [1, 1, 1, 64] + sin = sin.unsqueeze(1) + o1 = x1 * cos - x2 * sin + o2 = x1 * sin + x2 * cos + return torch.stack([o1, o2], dim=-1).flatten(-2) + +class TalkerKVWrapper(nn.Module): + def __init__(self, sd): + super().__init__() + self.norm = RMSNorm(DIM) + self.norm.weight.data = sd["norm.weight"] + self.lm_head = nn.Linear(DIM, VOCAB, bias=False) + self.lm_head.weight.data = sd["output.weight"] + + self.layers = nn.ModuleList() + for i in range(N_LAYERS): + L = nn.Module() + L.attn_norm = RMSNorm(DIM); L.attn_norm.weight.data = sd[f"layers.{i}.attention_norm.weight"] + L.ffn_norm = RMSNorm(DIM); L.ffn_norm.weight.data = sd[f"layers.{i}.ffn_norm.weight"] + L.wq = nn.Linear(DIM, N_HEADS*HEAD_DIM, bias=False); L.wq.weight.data = sd[f"layers.{i}.attention.wq.weight"] + L.wk = nn.Linear(DIM, N_KV_HEADS*HEAD_DIM, bias=False); L.wk.weight.data = sd[f"layers.{i}.attention.wk.weight"] + L.wv = nn.Linear(DIM, N_KV_HEADS*HEAD_DIM, bias=False); L.wv.weight.data = sd[f"layers.{i}.attention.wv.weight"] + L.wo = nn.Linear(N_HEADS*HEAD_DIM, DIM, bias=False); L.wo.weight.data = sd[f"layers.{i}.attention.wo.weight"] + L.q_norm = RMSNorm(HEAD_DIM); L.q_norm.weight.data = sd[f"layers.{i}.attention.q_norm_fn.weight"] + L.k_norm = RMSNorm(HEAD_DIM); L.k_norm.weight.data = sd[f"layers.{i}.attention.k_norm_fn.weight"] + L.w1 = nn.Linear(DIM, FFN_DIM, bias=False); L.w1.weight.data = sd[f"layers.{i}.feed_forward.w1.weight"] + L.w2 = nn.Linear(FFN_DIM, DIM, bias=False); L.w2.weight.data = sd[f"layers.{i}.feed_forward.w2.weight"] + L.w3 = nn.Linear(DIM, FFN_DIM, bias=False); L.w3.weight.data = sd[f"layers.{i}.feed_forward.w3.weight"] + self.layers.append(L) + + def forward(self, emb, mask, cos, sin, *kv_in): + """ + emb: [1, 1, 1024] + mask: [1, 1, 1, 200] (right-aligned, -10000 for masked positions) + cos: [1, 1, 128] (M-RoPE cos for current position) + sin: [1, 1, 128] (M-RoPE sin for current position) + kv_in: 28 × (k[1,8,199,128], v[1,8,199,128]) + Returns: hidden[1,1,1024], logits[1,1,3072], 28 × (k[1,8,200,128], v[1,8,200,128]) + """ + h = emb # [1, 1, 1024] + new_kvs = [] + + for i, L in enumerate(self.layers): + k_cache = kv_in[i * 2] # [1, 8, KV_LEN, 128] + v_cache = kv_in[i * 2 + 1] # [1, 8, KV_LEN, 128] + + h_n = L.attn_norm(h) + q = L.wq(h_n).view(1, 1, N_HEADS, HEAD_DIM).transpose(1, 2) # [1,16,1,128] + k = L.wk(h_n).view(1, 1, N_KV_HEADS, HEAD_DIM).transpose(1, 2) # [1,8,1,128] + v = L.wv(h_n).view(1, 1, N_KV_HEADS, HEAD_DIM).transpose(1, 2) # [1,8,1,128] + + # q_norm, k_norm BEFORE rope + q = L.q_norm(q) + k = L.k_norm(k) + + # Apply M-RoPE (interleaved) + q = apply_interleaved_rope(q, cos, sin) + k = apply_interleaved_rope(k, cos, sin) + + # Concat new k,v to cache → [1, 8, 200, 128] + k_full = torch.cat([k_cache, k], dim=2) # [1, 8, KV_LEN+1, 128] + v_full = torch.cat([v_cache, v], dim=2) + + # GQA expand for attention + k_exp = k_full.repeat(1, GQA_GROUPS, 1, 1) # [1, 16, 200, 128] + v_exp = v_full.repeat(1, GQA_GROUPS, 1, 1) + + # Attention + attn = torch.matmul(q, k_exp.transpose(-2, -1)) / math.sqrt(HEAD_DIM) # [1,16,1,200] + attn = attn + mask + attn = F.softmax(attn, dim=-1) + out = torch.matmul(attn, v_exp) # [1,16,1,128] + out = out.transpose(1, 2).contiguous().view(1, 1, N_HEADS * HEAD_DIM) # [1,1,2048] + h = h + L.wo(out) + + # FFN (SwiGLU) + h_n2 = L.ffn_norm(h) + h = h + L.w2(F.silu(L.w1(h_n2)) * L.w3(h_n2)) + + new_kvs.extend([k_full, v_full]) + + h = self.norm(h) + logits = self.lm_head(h) # [1, 1, 3072] + + return (h, logits, *new_kvs) + +print("Building model...") +wrapper = TalkerKVWrapper(state) +wrapper.eval() + +# Test with dummy inputs +print("Testing...") +dummy_emb = torch.randn(1, 1, DIM) +dummy_mask = torch.full((1, 1, 1, KV_LEN + 1), -10000.0) +dummy_mask[0, 0, 0, -1] = 0.0 # only current position visible +dummy_cos = torch.ones(1, 1, HEAD_DIM) +dummy_sin = torch.zeros(1, 1, HEAD_DIM) +kv_init = [] +for _ in range(N_LAYERS): + kv_init.append(torch.zeros(1, N_KV_HEADS, KV_LEN, HEAD_DIM)) # k + kv_init.append(torch.zeros(1, N_KV_HEADS, KV_LEN, HEAD_DIM)) # v + +with torch.no_grad(): + out = wrapper(dummy_emb, dummy_mask, dummy_cos, dummy_sin, *kv_init) + +hidden = out[0] +logits = out[1] +print(f"hidden: {hidden.shape}, logits: {logits.shape}") +print(f"new_k[0]: {out[2].shape}, new_v[0]: {out[3].shape}") + +# Export ONNX +OUT = "/opt/Kazeia/models_qnn/talker_kv_cpu_mrope" +os.makedirs(OUT, exist_ok=True) + +input_names = ["inputs_embeds", "attention_mask", "cos", "sin"] +for i in range(N_LAYERS): + input_names.extend([f"k_{i}_in", f"v_{i}_in"]) + +output_names = ["hidden", "logits"] +for i in range(N_LAYERS): + output_names.extend([f"k_{i}_out", f"v_{i}_out"]) + +all_inputs = (dummy_emb, dummy_mask, dummy_cos, dummy_sin, *kv_init) + +print(f"\nExporting ONNX ({len(input_names)} inputs, {len(output_names)} outputs)...") +torch.onnx.export( + wrapper, all_inputs, + f"{OUT}/model.onnx", + input_names=input_names, + output_names=output_names, + opset_version=17, +) +sz = os.path.getsize(f"{OUT}/model.onnx") +print(f"Saved: {OUT}/model.onnx ({sz/1024/1024:.0f} MB)") + +# Pre-compute M-RoPE cos/sin for all 200 positions and save +# For TTS, all 3 mrope axes use the same position, so M-RoPE = standard RoPE +# inv_freq = 1 / (theta^(2i/128)) for i in 0..63 +print("\nPre-computing RoPE tables (standard, theta=1e6)...") +theta = 1000000.0 +inv_freq = 1.0 / (theta ** (np.arange(0, HEAD_DIM, 2, dtype=np.float32) / HEAD_DIM)) # [64] +positions = np.arange(KV_LEN + 1, dtype=np.float32) # [200] +angles = np.outer(positions, inv_freq) # [200, 64] +rope_cos = np.cos(angles).astype(np.float32) # [200, 64] +rope_sin = np.sin(angles).astype(np.float32) +# Expand to head_dim=128 for interleaved rope: each freq pair applies to (x[2i], x[2i+1]) +rope_cos_full = np.repeat(rope_cos, 2, axis=1) # [200, 128] +rope_sin_full = np.repeat(rope_sin, 2, axis=1) +np.save(f"{OUT}/talker_rotary_cos.npy", rope_cos_full) +np.save(f"{OUT}/talker_rotary_sin.npy", rope_sin_full) +print(f"Rotary tables: {rope_cos_full.shape} saved") + +# Validate: run a quick prefill step and check output +print("\n=== Validation ===") +import onnxruntime as ort +sess = ort.InferenceSession(f"{OUT}/model.onnx") +print(f"ONNX inputs: {[i.name for i in sess.get_inputs()[:4]]}...") + +mask = np.full((1, 1, 1, KV_LEN + 1), -10000.0, dtype=np.float32) +mask[0, 0, 0, -1] = 0.0 +cos_0 = rope_cos_full[0:1].reshape(1, 1, HEAD_DIM) +sin_0 = rope_sin_full[0:1].reshape(1, 1, HEAD_DIM) + +inputs = { + "inputs_embeds": np.random.randn(1, 1, DIM).astype(np.float32), + "attention_mask": mask, + "cos": cos_0, + "sin": sin_0, +} +for i in range(N_LAYERS): + inputs[f"k_{i}_in"] = np.zeros((1, N_KV_HEADS, KV_LEN, HEAD_DIM), dtype=np.float32) + inputs[f"v_{i}_in"] = np.zeros((1, N_KV_HEADS, KV_LEN, HEAD_DIM), dtype=np.float32) + +out = sess.run(None, inputs) +print(f"ONNX hidden: {out[0].shape}, logits: {out[1].shape}") +print(f"ONNX k_0_out: {out[2].shape}") +print("Validation OK!") diff --git a/scripts/extract_tts_embeddings.py b/scripts/extract_tts_embeddings.py new file mode 100644 index 0000000..5f6775f --- /dev/null +++ b/scripts/extract_tts_embeddings.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Extract embedding tables from Qwen3-TTS for the Android app.""" +import torch +import numpy as np +import os + +OUT = "/opt/Kazeia/tts_npy_exports" +os.makedirs(OUT, exist_ok=True) +NATIVE = "/opt/Kazeia/models_qnn/qwen3-tts-native" +TALKER_HF = "/opt/Kazeia/models_qnn/talker_hf" + +# 1. Text embeddings + projection +print("=== text_components.pt ===") +tc = torch.load(f"{NATIVE}/text_components.pt", map_location="cpu", weights_only=False) +text_emb_dict = tc['text_embedding'] # OrderedDict +text_proj_dict = tc['text_projection'] # OrderedDict + +print("text_embedding keys:", list(text_emb_dict.keys())) +print("text_projection keys:", list(text_proj_dict.keys())) + +# text_embedding.weight: [151936, 2048] — full LLM vocab, but we only need first 1050 tokens +text_emb_w = text_emb_dict['weight'] # [151936, 2048] +print(f"text_embedding.weight: {text_emb_w.shape}") + +# Only keep the first 1050 text tokens (CODEC_OFFSET in the engine) +text_emb_w = text_emb_w[:1050] # [1050, 2048] +print(f"text_embedding (trimmed): {text_emb_w.shape}") + +# text_projection is a 2-layer MLP: fc1(2048→1024) + GELU + fc2(1024→1024) +fc1_w = text_proj_dict['linear_fc1.weight'] # [hidden, 2048] +fc1_b = text_proj_dict['linear_fc1.bias'] +fc2_w = text_proj_dict['linear_fc2.weight'] # [1024, hidden] +fc2_b = text_proj_dict['linear_fc2.bias'] +print(f"text_projection: fc1={fc1_w.shape}, fc2={fc2_w.shape}") + +# Apply: x = fc2(silu(fc1(text_emb))) — hidden_act = "silu" from config.json +x = text_emb_w.float() @ fc1_w.float().T + fc1_b.float() +x = torch.nn.functional.silu(x) +text_projected = x @ fc2_w.float().T + fc2_b.float() +print(f"text_embeds_projected: {text_projected.shape}") +np.save(f"{OUT}/text_embeds_projected.npy", text_projected.numpy()) + +# 2. Codec embedding from talker HF safetensors +print("\n=== talker safetensors ===") +from safetensors.torch import load_file +talker_st = load_file(f"{TALKER_HF}/model.safetensors") +# Find the embed_tokens / codec_embedding +for k in sorted(talker_st.keys()): + if 'embed' in k.lower(): + print(f" {k}: {talker_st[k].shape}") + +# The talker's token embedding is the codec embedding +codec_key = None +for k in talker_st.keys(): + if 'embed_tokens' in k or 'codec_embedding' in k: + codec_key = k + break +if codec_key: + codec_emb = talker_st[codec_key].float().numpy() + print(f"codec_embedding ({codec_key}): {codec_emb.shape}") + np.save(f"{OUT}/codec_embedding.npy", codec_emb) +else: + print("WARNING: codec_embedding not found in talker safetensors") + # Fallback: try lm_head transposed + for k in talker_st.keys(): + print(f" {k}: {talker_st[k].shape}") + +# 3. TTS special embeddings [3, 1024] = bos(2149), eos(2150), pad(2148) +# These are TEXT token IDs that go through text_embedding + text_projection (SiLU MLP) +# NOT codec_embedding! The IDs are in the LLM text vocab space (151936) +text_emb_full = text_emb_dict['weight'] # [151936, 2048] — full vocab +special_ids = [2149, 2150, 2148] # bos, eos, pad +special_text = text_emb_full[special_ids].float() # [3, 2048] +# Apply text_projection: fc1 + SiLU + fc2 +h = special_text @ fc1_w.float().T + fc1_b.float() +h = torch.nn.functional.silu(h) +special_projected = (h @ fc2_w.float().T + fc2_b.float()).numpy() +np.save(f"{OUT}/tts_special_embeds.npy", special_projected) +print(f"tts_special_embeds: {special_projected.shape} (projected via text_projection SiLU)") + +# 4. Code predictor embeddings — 15 codebook embedding tables [15, 2048, 1024] +print("\n=== code_predictor_weights.pt ===") +cp = torch.load(f"{NATIVE}/code_predictor_weights.pt", map_location="cpu", weights_only=False) +cp_embs = [] +for i in range(15): + key = f"model.codec_embedding.{i}.weight" + if key in cp: + cp_embs.append(cp[key].float().numpy()) +if cp_embs: + cp_all = np.stack(cp_embs) # [15, 2048, 1024] + np.save(f"{OUT}/code_predictor_embeddings.npy", cp_all) + print(f"code_predictor_embeddings: {cp_all.shape}") + +# 5. VQ codebooks — extract actual codebook vectors from embedding_sum / cluster_usage +print("\n=== speech_decoder_weights.pt (VQ codebooks) ===") +sd = torch.load(f"{NATIVE}/speech_decoder_weights.pt", map_location="cpu", weights_only=False) + +# First codebook (rvq_first) + 15 rest codebooks (rvq_rest) +codebooks = [] +# rvq_first +usage = sd['quantizer.rvq_first.vq.layers.0._codebook.cluster_usage'] +emb_sum = sd['quantizer.rvq_first.vq.layers.0._codebook.embedding_sum'] +# actual codebook = embedding_sum / cluster_usage (EMA-style) +cb = (emb_sum / usage.unsqueeze(1).clamp(min=1)).float().numpy() +codebooks.append(cb) +print(f" rvq_first codebook: {cb.shape}") + +# rvq_rest has 15 layers (indices 0-14) +for i in range(15): + usage = sd[f'quantizer.rvq_rest.vq.layers.{i}._codebook.cluster_usage'] + emb_sum = sd[f'quantizer.rvq_rest.vq.layers.{i}._codebook.embedding_sum'] + cb = (emb_sum / usage.unsqueeze(1).clamp(min=1)).float().numpy() + codebooks.append(cb) + +vq_all = np.stack(codebooks) # [16, 2048, 256] +np.save(f"{OUT}/vq_codebooks.npy", vq_all) +print(f"vq_codebooks: {vq_all.shape} (16 codebooks × 2048 entries × 256 dim)") + +# Summary +print(f"\n=== Output: {OUT} ===") +for f in sorted(os.listdir(OUT)): + sz = os.path.getsize(f"{OUT}/{f}") + print(f" {f}: {sz/1024/1024:.1f} MB") diff --git a/scripts/generate_tokens_for_tablet.py b/scripts/generate_tokens_for_tablet.py new file mode 100644 index 0000000..6ad9f55 --- /dev/null +++ b/scripts/generate_tokens_for_tablet.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +""" +Generate talker codec tokens on PC using the full Python pipeline, +then save them for the tablet to decode (CP + VQ + Decoder). + +This bypasses the broken talker ONNX/GGUF export and proves +that the rest of the tablet pipeline (CP + decoder) works. + +Output: + - talker_output.bin: binary file with cb0 tokens + hidden states + Format: int32 n_tokens, then for each token: + int32 cb0, float32[1024] hidden_state +""" +import sys +sys.path.insert(0, "/opt/Kazeia/qnn_venv/lib/python3.10/site-packages") +import warnings; warnings.filterwarnings("ignore") +import torch, numpy as np, struct + +model_path = "/home/alf/.cache/huggingface/hub/models--Qwen--Qwen3-TTS-12Hz-0.6B-Base/snapshots/5d83992436eae1d760afd27aff78a71d676296fc" + +from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration +from transformers import AutoTokenizer + +print("Loading model...") +model = Qwen3TTSForConditionalGeneration.from_pretrained(model_path, dtype=torch.float32, device_map='cpu') +model.eval() +tokenizer = AutoTokenizer.from_pretrained(model_path) +spk = torch.from_numpy(np.load('models_qnn/qwen3-tts-embeddings/damien_spk_embedding.npy')).float() + +TEXT = "Bonjour, je m'appelle Kazeia, je suis encore en phase de développement." + +# Hook to capture hidden states from talker +hidden_states = [] +cb0_tokens = [] + +orig_forward = model.talker.model.forward + +def capture_forward(*args, **kwargs): + result = orig_forward(*args, **kwargs) + # result is a BaseModelOutputWithPast — hidden states are in last_hidden_state + if hasattr(result, 'last_hidden_state'): + h = result.last_hidden_state + if h.shape[1] == 1: # decode step + hidden_states.append(h.detach().squeeze().cpu().numpy()) + return result + +model.talker.model.forward = capture_forward + +# Also hook codec_head to capture cb0 tokens +orig_codec_head = model.talker.codec_head + +class HookedCodecHead(torch.nn.Module): + def __init__(self, head): + super().__init__() + self.head = head + def forward(self, x): + logits = self.head(x) + return logits + +model.talker.codec_head = HookedCodecHead(orig_codec_head) + +prompt = f"<|im_start|>assistant\n{TEXT}<|endoftext|>\n<|im_start|>assistant\n" +input_ids = tokenizer(prompt, return_tensors="pt").input_ids + +print(f"Generating: '{TEXT}'") +print(f"Input tokens: {input_ids.shape[1]}") + +torch.manual_seed(42) +with torch.no_grad(): + result = model.generate( + input_ids=[input_ids], + languages=["french"], + speakers=[""], + voice_clone_prompt={ + "x_vector_only_mode": [True], + "icl_mode": [False], + "ref_code": None, + "ref_spk_embedding": [spk], + "x_vector": [spk], + }, + do_sample=True, + temperature=0.9, + top_k=50, + repetition_penalty=1.05, + max_new_tokens=200, + ) + +# result is a list of [audio_codes] tensors +codes = result[0][0] # [n_tokens, 16] (16 codebooks) +n_tokens = codes.shape[0] +print(f"Generated {n_tokens} tokens × {codes.shape[1]} codebooks") +print(f"Duration: {n_tokens/12:.2f}s") +print(f"First 5 cb0: {codes[:5, 0].tolist()}") +print(f"Hidden states captured: {len(hidden_states)}") + +# Save ALL 16 codebooks directly — tablet doesn't need to run CP at all! +# Format: int32 n_tokens, int32 n_codebooks, then [n_tokens × n_codebooks] int32 +OUT = "/opt/Kazeia/tablet_codes.bin" +with open(OUT, "wb") as f: + f.write(struct.pack(" 0.01)") + +sf.write(OUTPUT, audio, sr) +print(f"Saved: {OUTPUT}")