From 8b35b14126e6b6a5031af575a35cd4e567112b02 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 14 Feb 2025 19:32:15 -0500 Subject: [PATCH] cleanup - use tauri-managed state instead of hand-rolled mutexes --- .../src-tauri/src/endpoints/image_endpoint.rs | 102 ++++++------------ crates/desktop/src-tauri/src/lib.rs | 18 +++- 2 files changed, 46 insertions(+), 74 deletions(-) diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index 9aa261368b..06879fffb4 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/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>>> = Lazy::new(|| Arc::new(RwLock::new(None))); - #[tauri::command] -pub fn infer_image(image: &str, model_config: State) -> Result { - do_infer_image(image, &model_config) +pub fn infer_image(image: &str, model_config: State, model_cache: State) -> Result { + do_infer_image(image, &model_config, &model_cache) } -fn do_infer_image(image: &str, config: &ModelConfig) -> Result { - 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 { + 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)) + } } diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index ddd7f3c984..0a649cf8d5 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/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,