From 01088f4f1915e36a7d0d8d8756f62f8207a48911 Mon Sep 17 00:00:00 2001
From: soryu <soryu@soryu.co>
Date: Sat, 20 Dec 2025 15:36:04 +0000
Subject: Implement makima listen websockets server

---
 makima/src/server/state.rs | 50 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)
 create mode 100644 makima/src/server/state.rs

(limited to 'makima/src/server/state.rs')

diff --git a/makima/src/server/state.rs b/makima/src/server/state.rs
new file mode 100644
index 0000000..8eaf788
--- /dev/null
+++ b/makima/src/server/state.rs
@@ -0,0 +1,50 @@
+//! Application state holding shared ML models.
+
+use std::sync::Arc;
+use tokio::sync::Mutex;
+
+use crate::listen::{DiarizationConfig, ParakeetTDT, Sortformer};
+use crate::tts::ChatterboxTTS;
+
+/// Shared application state containing ML models.
+///
+/// Models are wrapped in `Mutex` for thread-safe mutable access during inference.
+pub struct AppState {
+    /// Speech-to-text model (Parakeet)
+    pub parakeet: Mutex<ParakeetTDT>,
+    /// Speaker diarization model (Sortformer)
+    pub sortformer: Mutex<Sortformer>,
+    /// Text-to-speech model (ChatterboxTTS)
+    pub chatterbox: Mutex<ChatterboxTTS>,
+}
+
+impl AppState {
+    /// Load all ML models from the specified directories.
+    ///
+    /// # Arguments
+    /// * `parakeet_model_dir` - Path to the Parakeet STT model directory
+    /// * `sortformer_model_path` - Path to the Sortformer diarization model file
+    /// * `tts_model_dir` - Optional path to the ChatterboxTTS model directory
+    pub fn new(
+        parakeet_model_dir: &str,
+        sortformer_model_path: &str,
+        tts_model_dir: Option<&str>,
+    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
+        let parakeet = ParakeetTDT::from_pretrained(parakeet_model_dir, None)?;
+        let sortformer = Sortformer::with_config(
+            sortformer_model_path,
+            None,
+            DiarizationConfig::callhome(),
+        )?;
+        let chatterbox = ChatterboxTTS::from_pretrained(tts_model_dir)?;
+
+        Ok(Self {
+            parakeet: Mutex::new(parakeet),
+            sortformer: Mutex::new(sortformer),
+            chatterbox: Mutex::new(chatterbox),
+        })
+    }
+}
+
+/// Type alias for the shared application state.
+pub type SharedState = Arc<AppState>;
-- 
cgit v1.2.3