mold_server/
lib.rs

1pub mod logging;
2pub mod model_cache;
3pub mod model_manager;
4pub mod queue;
5pub mod routes;
6pub mod state;
7
8#[cfg(test)]
9mod routes_test;
10
11use anyhow::Result;
12use mold_core::{Config, ModelPaths};
13use std::net::SocketAddr;
14use std::path::PathBuf;
15use tokio::net::TcpListener;
16use tower_http::cors::CorsLayer;
17use tower_http::trace::TraceLayer;
18use tracing::info;
19
20use state::QueueHandle;
21
22pub async fn run_server(bind: &str, port: u16, models_dir: PathBuf) -> Result<()> {
23    Config::install_runtime_models_dir_override(models_dir.clone());
24
25    let mut config = Config::load_or_default();
26    config.models_dir = models_dir.to_string_lossy().into_owned();
27    let model_name = config.resolved_default_model();
28
29    // Create the generation queue channel (bounded, 16 slots).
30    let (job_tx, job_rx) = tokio::sync::mpsc::channel(16);
31    let queue_handle = QueueHandle::new(job_tx);
32
33    let state = match ModelPaths::resolve(&model_name, &config) {
34        Some(paths) => {
35            info!(model = %model_name, "configured model");
36            info!(transformer = %paths.transformer.display());
37            info!(vae = %paths.vae.display());
38            if let Some(t5) = &paths.t5_encoder {
39                info!(t5 = %t5.display());
40            }
41            if let Some(clip) = &paths.clip_encoder {
42                info!(clip = %clip.display());
43            }
44            if let Some(t5_tok) = &paths.t5_tokenizer {
45                info!(t5_tok = %t5_tok.display());
46            }
47            if let Some(clip_tok) = &paths.clip_tokenizer {
48                info!(clip_tok = %clip_tok.display());
49            }
50            if let Some(clip2) = &paths.clip_encoder_2 {
51                info!(clip2 = %clip2.display());
52            }
53            if let Some(clip2_tok) = &paths.clip_tokenizer_2 {
54                info!(clip2_tok = %clip2_tok.display());
55            }
56            for (i, te) in paths.text_encoder_files.iter().enumerate() {
57                info!(text_encoder_shard = i, path = %te.display());
58            }
59            if let Some(text_tok) = &paths.text_tokenizer {
60                info!(text_tok = %text_tok.display());
61            }
62
63            let offload = std::env::var("MOLD_OFFLOAD").is_ok_and(|v| v == "1");
64            // Create shared pool before engine so the first engine can populate it.
65            let shared_pool = std::sync::Arc::new(std::sync::Mutex::new(
66                mold_inference::shared_pool::SharedPool::new(),
67            ));
68            let engine = mold_inference::create_engine_with_pool(
69                model_name,
70                paths,
71                &config,
72                mold_inference::LoadStrategy::Eager,
73                offload,
74                Some(shared_pool.clone()),
75            )?;
76            let mut state = state::AppState::new(engine, config, queue_handle);
77            state.shared_pool = shared_pool;
78            state
79        }
80        None => {
81            info!("no default model configured — models will be pulled on first request");
82            state::AppState::empty(config, queue_handle)
83        }
84    };
85
86    // Spawn the generation queue worker — processes jobs sequentially (single GPU).
87    let worker_state = state.clone();
88    tokio::spawn(queue::run_queue_worker(job_rx, worker_state));
89
90    // Ensure output directory exists and pre-generate thumbnails.
91    {
92        let config = state.config.read().await;
93        if config.is_output_disabled() {
94            tracing::warn!(
95                "image output is disabled (output_dir is empty) — \
96                 generated images will not be saved and the TUI gallery will be empty"
97            );
98        } else {
99            let output_dir = config.effective_output_dir();
100            let _ = std::fs::create_dir_all(&output_dir);
101            info!(output_dir = %output_dir.display(), "gallery output directory");
102            routes::spawn_thumbnail_warmup(&config);
103        }
104    }
105
106    let cors = build_cors_layer()?;
107
108    let app = routes::create_router(state)
109        .layer(TraceLayer::new_for_http())
110        .layer(cors);
111
112    let addr: SocketAddr = format!("{bind}:{port}").parse()?;
113    info!(%addr, "starting mold server");
114
115    let listener = TcpListener::bind(addr).await?;
116    axum::serve(listener, app).await?;
117
118    Ok(())
119}
120
121fn build_cors_layer() -> Result<CorsLayer> {
122    let cors = match std::env::var("MOLD_CORS_ORIGIN") {
123        Ok(origin) if !origin.is_empty() => {
124            let origin = origin
125                .parse::<axum::http::HeaderValue>()
126                .map_err(|_| anyhow::anyhow!("invalid MOLD_CORS_ORIGIN value: {origin}"))?;
127            CorsLayer::new()
128                .allow_origin(origin)
129                .allow_methods([
130                    axum::http::Method::GET,
131                    axum::http::Method::POST,
132                    axum::http::Method::DELETE,
133                ])
134                .allow_headers(tower_http::cors::Any)
135                .expose_headers([axum::http::header::HeaderName::from_static(
136                    "x-mold-seed-used",
137                )])
138        }
139        _ => CorsLayer::permissive(),
140    };
141    Ok(cors)
142}
143
144#[cfg(test)]
145mod tests {
146    use super::build_cors_layer;
147    use std::sync::Mutex;
148
149    static ENV_LOCK: Mutex<()> = Mutex::new(());
150
151    #[test]
152    fn invalid_cors_origin_returns_error() {
153        let _lock = ENV_LOCK.lock().unwrap();
154        std::env::set_var("MOLD_CORS_ORIGIN", "\nnot-a-header");
155        let result = build_cors_layer();
156        std::env::remove_var("MOLD_CORS_ORIGIN");
157        assert!(result.is_err());
158    }
159
160    #[test]
161    fn valid_cors_origin_builds_layer() {
162        let _lock = ENV_LOCK.lock().unwrap();
163        std::env::set_var("MOLD_CORS_ORIGIN", "https://example.com");
164        let result = build_cors_layer();
165        std::env::remove_var("MOLD_CORS_ORIGIN");
166        assert!(result.is_ok());
167    }
168}
mold_server/lib.rs

mold_server/
lib.rs
