cleanup - use tauri-managed state instead of hand-rolled mutexes

crates/desktop/src-tauri/src/endpoints/image_endpoint.rs

@@ -3,8 +3,6 @@ use crate::ml::model_cache::ModelCache;
 use crate::model_config::ModelConfig;
 use anyhow::Result;
 use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
-use once_cell::sync::Lazy;
-use std::sync::{Arc, Mutex, RwLock};
 use tauri::State;
 
 const PROMPT: &str = "A beautiful landscape with mountains and a lake";
@@ -12,79 +10,41 @@ const PROMPT: &str = "A beautiful landscape with mountains and a lake";
 const HEIGHT: usize = 512;
 const WIDTH: usize = 512;
 
-static MODEL_CACHE: Lazy<Arc<RwLock<Option<ModelCache>>>> = Lazy::new(|| Arc::new(RwLock::new(None)));
-
 #[tauri::command]
-pub fn infer_image(image: &str, model_config: State<ModelConfig>) -> Result<String, String> {
-    do_infer_image(image, &model_config)
+pub fn infer_image(image: &str, model_config: State<ModelConfig>, model_cache: State<ModelCache>) -> Result<String, String> {
+  do_infer_image(image, &model_config, &model_cache)
 }
 
-fn do_infer_image(image: &str, config: &ModelConfig) -> Result<String, String> {
-    println!("infer_image called; generating image with SDXL Turbo...");
-    
-    // TODO(bt,2025-02-14): Clean this up.
-    match MODEL_CACHE.write() {
-        Err(err) => return Err(err.to_string()),
-        Ok(mut maybe_cache) => {
-            match &*maybe_cache {
-                None => {
-                    let cache = ModelCache::new(
-                        config.device.clone(), 
-                        config.dtype, 
-                        config.sd_version.clone(), 
-                        config.sd_config.clone()
-                    ).map_err(|err| err.to_string())?;
+fn do_infer_image(image: &str, config: &ModelConfig, model_cache: &ModelCache) -> Result<String, String> {
+  println!("infer_image called; generating image with SDXL Turbo...");
+  
+  let args = Args {
+    api: config.hf_api.clone(),
+    prompt: PROMPT.to_string(),
+    uncond_prompt: "".to_string(),
+    cpu: config.device.is_cpu(),
+    height: Some(HEIGHT),
+    width: Some(WIDTH),
+    n_steps: Some(1),
+    num_samples: 1,
+    seed: None,
+    sd_version: config.sd_version.clone(),
+    guidance_scale: Some(0.0),
+    model_cache,
+  };
 
-                    *maybe_cache = Some(cache);
-                }
-                Some(_model) => {
-                    println!("Model cache already exists");
-                }
-            }
+  match run(args) {
+    Ok(_) => {
+      let img_data = std::fs::read("temp.png")
+        .map_err(|e| format!("Failed to read generated image: {}", e))?;
+        
+      let encoded = BASE64_STANDARD.encode(&img_data);
+      println!("Generated image encoded successfully");
 
-        }
+      let _ = std::fs::remove_file("temp.png");
+
+      Ok(encoded)
     }
-    
-    match MODEL_CACHE.read() {
-        Err(err) => return Err(err.to_string()),
-        Ok(maybe_model) => {
-            match &*maybe_model {
-                None => {
-                    Err("model does not exist".to_string())
-                }
-                Some(model_cache) => {
-                    let args = Args {
-                        api: config.hf_api.clone(),
-                        prompt: PROMPT.to_string(),
-                        uncond_prompt: "".to_string(),
-                        cpu: config.device.is_cpu(),
-                        height: Some(HEIGHT),
-                        width: Some(WIDTH),
-                        n_steps: Some(1),
-                        num_samples: 1,
-                        seed: None,
-                        sd_version: config.sd_version.clone(),
-                        guidance_scale: Some(0.0),
-                        model_cache,
-                    };
-
-                    match run(args) {
-                        Ok(_) => {
-                            let img_data = std::fs::read("temp.png")
-                              .map_err(|e| format!("Failed to read generated image: {}", e))?;
-                            let encoded = BASE64_STANDARD.encode(&img_data);
-                            println!("Generated image encoded successfully");
-
-                            let _ = std::fs::remove_file("temp.png");
-
-                            Ok(encoded)
-                        }
-                        Err(e) => Err(format!("Failed to generate image: {}", e))
-                    }
-
-                }
-            }
-        }
-    }
-    
+    Err(e) => Err(format!("Failed to generate image: {}", e))
+  }
 }

crates/desktop/src-tauri/src/lib.rs

@@ -2,20 +2,31 @@ pub mod ml;
 pub mod endpoints;
 pub mod model_config;
 
+use crate::ml::model_cache::ModelCache;
+use crate::model_config::ModelConfig;
 use endpoints::image_endpoint::infer_image;
+use log::info;
 use once_cell::sync::Lazy;
 use std::sync::{Arc, RwLock};
-use crate::model_config::ModelConfig;
 
 #[cfg_attr(mobile, tauri::mobile_entry_point)]
 pub fn run() {
   info!("Loading model config...");
-  
+
   let model_config = ModelConfig::init()
     .expect("config should load");
 
-  info!("Initializing backend runtime...");
+  info!("Creating model cache...");
   
+  let model_cache = ModelCache::new(
+    model_config.device.clone(),
+    model_config.dtype,
+    model_config.sd_version.clone(),
+    model_config.sd_config.clone()
+  ).expect("Model cache should create");
+
+  info!("Initializing backend runtime...");
+
   tauri::Builder::default()
     .setup(|app| {
       if cfg!(debug_assertions) {
@@ -28,6 +39,7 @@ pub fn run() {
       Ok(())
     })
     .manage(model_config)
+    .manage(model_cache)
     .invoke_handler(tauri::generate_handler![
       my_custom_command, 
       test_round_trip,