diff options
Diffstat (limited to 'makima/src/bin')
| -rw-r--r-- | makima/src/bin/server.rs | 32 |
1 files changed, 24 insertions, 8 deletions
diff --git a/makima/src/bin/server.rs b/makima/src/bin/server.rs index 06b6585..470e295 100644 --- a/makima/src/bin/server.rs +++ b/makima/src/bin/server.rs @@ -8,10 +8,10 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; use makima::server::{run_server, state::AppState}; -/// Default model paths relative to the working directory. -const PARAKEET_MODEL_DIR: &str = "models/parakeet-tdt-0.6b-v3"; -const PARAKEET_EOU_DIR: &str = "models/realtime_eou_120m-v1-onnx"; -const SORTFORMER_MODEL_PATH: &str = "models/diarization/diar_streaming_sortformer_4spk-v2.onnx"; +/// Default model paths (can be overridden via environment variables). +const DEFAULT_PARAKEET_MODEL_DIR: &str = "models/parakeet-tdt-0.6b-v3"; +const DEFAULT_PARAKEET_EOU_DIR: &str = "models/realtime_eou_120m-v1-onnx"; +const DEFAULT_SORTFORMER_MODEL_PATH: &str = "models/diarization/diar_streaming_sortformer_4spk-v2.onnx"; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -19,22 +19,38 @@ async fn main() -> anyhow::Result<()> { tracing_subscriber::registry() .with( tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| "makima=debug,tower_http=debug".into()), + .unwrap_or_else(|_| "makima=info,tower_http=info".into()), ) .with(tracing_subscriber::fmt::layer()) .init(); tracing::info!("Starting Makima Listening API Server"); - tracing::info!("Loading ML models..."); + + // Read configuration from environment + let port = std::env::var("PORT").unwrap_or_else(|_| "8080".to_string()); + let parakeet_dir = std::env::var("PARAKEET_MODEL_DIR") + .unwrap_or_else(|_| DEFAULT_PARAKEET_MODEL_DIR.to_string()); + let parakeet_eou_dir = std::env::var("PARAKEET_EOU_DIR") + .unwrap_or_else(|_| DEFAULT_PARAKEET_EOU_DIR.to_string()); + let sortformer_path = std::env::var("SORTFORMER_MODEL_PATH") + .unwrap_or_else(|_| DEFAULT_SORTFORMER_MODEL_PATH.to_string()); + + tracing::info!( + parakeet = %parakeet_dir, + eou = %parakeet_eou_dir, + sortformer = %sortformer_path, + "Loading ML models..." + ); // Load ML models let state = Arc::new( - AppState::new(PARAKEET_MODEL_DIR, PARAKEET_EOU_DIR, SORTFORMER_MODEL_PATH) + AppState::new(¶keet_dir, ¶keet_eou_dir, &sortformer_path) .map_err(|e| anyhow::anyhow!("Failed to load models: {}", e))?, ); tracing::info!("Models loaded successfully"); // Run the server - run_server(state, "0.0.0.0:8080").await + let addr = format!("0.0.0.0:{}", port); + run_server(state, &addr).await } |