From b032385ed83fa7c1114cb0527a273041caa33a4e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 14 Feb 2025 19:17:35 -0500 Subject: [PATCH] config object --- .../src/{ => endpoints}/flip_image.rs | 0 .../src/{ => endpoints}/image_endpoint.rs | 107 +++++++----------- crates/desktop/src-tauri/src/endpoints/mod.rs | 3 + .../src-tauri/src/{ => endpoints}/sd.rs | 0 crates/desktop/src-tauri/src/lib.rs | 7 +- crates/desktop/src-tauri/src/model_config.rs | 59 ++++++++++ 6 files changed, 108 insertions(+), 68 deletions(-) rename crates/desktop/src-tauri/src/{ => endpoints}/flip_image.rs (100%) rename crates/desktop/src-tauri/src/{ => endpoints}/image_endpoint.rs (53%) create mode 100644 crates/desktop/src-tauri/src/endpoints/mod.rs rename crates/desktop/src-tauri/src/{ => endpoints}/sd.rs (100%) create mode 100644 crates/desktop/src-tauri/src/model_config.rs diff --git a/crates/desktop/src-tauri/src/flip_image.rs b/crates/desktop/src-tauri/src/endpoints/flip_image.rs similarity index 100% rename from crates/desktop/src-tauri/src/flip_image.rs rename to crates/desktop/src-tauri/src/endpoints/flip_image.rs diff --git a/crates/desktop/src-tauri/src/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs similarity index 53% rename from crates/desktop/src-tauri/src/image_endpoint.rs rename to crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index 862aaa3eaf..b3a43b05eb 100644 --- a/crates/desktop/src-tauri/src/image_endpoint.rs +++ b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs @@ -1,80 +1,62 @@ -use std::env::args; -use std::sync::{Arc, RwLock}; -use crate::ml::model_file::StableDiffusionVersion; -use crate::sd::{self, Args}; +use crate::endpoints::sd::{run, Args}; +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 candle_core::{DType, Device}; -use candle_transformers::models::stable_diffusion; use once_cell::sync::Lazy; -use crate::ml::model_cache::ModelCache; -use crate::ml::models::unet_model::UNetModel; +use std::sync::{Arc, Mutex, RwLock}; const PROMPT: &str = "A beautiful landscape with mountains and a lake"; const HEIGHT: usize = 512; const WIDTH: usize = 512; +static MODEL_CONFIG: Lazy>>> = Lazy::new(|| Arc::new(Mutex::new(None))); + static MODEL_CACHE: Lazy>>> = Lazy::new(|| Arc::new(RwLock::new(None))); #[tauri::command] pub fn infer_image(image: &str) -> Result { + // TODO(bt,2025-02-14): There has to be a better way to inject dependencies. + + match MODEL_CONFIG.lock() { + Err(err) => Err(err.to_string()), + Ok(mut model_config) => { + match &*model_config { + Some(config) => { + do_infer_image(image, config) + } + None => { + // TODO(bt,2025-02-14): This is awful. + let config = ModelConfig::init() + .map_err(|err| err.to_string())?; + + let result = do_infer_image(image, &config)?; + + *model_config = Some(config); + + Ok(result) + } + } + } + } +} + +fn do_infer_image(image: &str, config: &ModelConfig) -> Result { println!("infer_image called; generating image with SDXL Turbo..."); - // Check CUDA availability - let use_cpu = match Device::new_cuda(0) { - Ok(_) => { - println!("CUDA device available"); - false - } - Err(e) => { - println!("CUDA not available ({}), falling back to CPU", e); - true - } - }; - - let sd_version = StableDiffusionVersion::Turbo; - - let sd_config = match sd_version { - StableDiffusionVersion::Turbo => { - println!("Configuring SDXL Turbo"); - stable_diffusion::StableDiffusionConfig::sdxl_turbo(None, Some(HEIGHT), Some(WIDTH)) - } - _ => { - println!("Configuring SD v2.1"); - stable_diffusion::StableDiffusionConfig::v2_1(None, Some(HEIGHT), Some(WIDTH)) - } - }; - - let api = hf_hub::api::sync::Api::new() - .map_err(|err| err.to_string())?; - - let dtype = DType::F32; - + // 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 device = if use_cpu { - println!("Using CPU for computation"); - Device::Cpu - } else { - println!("Attempting to use CUDA device"); - match Device::new_cuda(0) { - Ok(cuda_device) => { - println!("Successfully initialized CUDA device"); - cuda_device - } - Err(e) => { - println!("Failed to initialize CUDA device: {}. Falling back to CPU", e); - Device::Cpu - } - } - }; - - let cache = ModelCache::new(device, dtype, sd_version.clone(), sd_config.clone()) - .map_err(|err| err.to_string())?; + let cache = ModelCache::new( + config.device.clone(), + config.dtype, + config.sd_version.clone(), + config.sd_config.clone() + ).map_err(|err| err.to_string())?; *maybe_cache = Some(cache); } @@ -86,8 +68,6 @@ pub fn infer_image(image: &str) -> Result { } } - - match MODEL_CACHE.read() { Err(err) => return Err(err.to_string()), Ok(maybe_model) => { @@ -96,23 +76,22 @@ pub fn infer_image(image: &str) -> Result { Err("model does not exist".to_string()) } Some(model_cache) => { - let args = Args { - api, + api: config.hf_api.clone(), prompt: PROMPT.to_string(), uncond_prompt: "".to_string(), - cpu: use_cpu, + cpu: config.device.is_cpu(), height: Some(HEIGHT), width: Some(WIDTH), n_steps: Some(1), num_samples: 1, seed: None, - sd_version, + sd_version: config.sd_version.clone(), guidance_scale: Some(0.0), model_cache, }; - match sd::run(args) { + match run(args) { Ok(_) => { let img_data = std::fs::read("temp.png") .map_err(|e| format!("Failed to read generated image: {}", e))?; diff --git a/crates/desktop/src-tauri/src/endpoints/mod.rs b/crates/desktop/src-tauri/src/endpoints/mod.rs new file mode 100644 index 0000000000..89c5cf3414 --- /dev/null +++ b/crates/desktop/src-tauri/src/endpoints/mod.rs @@ -0,0 +1,3 @@ +pub mod flip_image; +pub mod image_endpoint; +pub mod sd; diff --git a/crates/desktop/src-tauri/src/sd.rs b/crates/desktop/src-tauri/src/endpoints/sd.rs similarity index 100% rename from crates/desktop/src-tauri/src/sd.rs rename to crates/desktop/src-tauri/src/endpoints/sd.rs diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 98f5d13927..285c576e07 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -1,9 +1,8 @@ -mod flip_image; -mod image_endpoint; -mod sd; pub mod ml; +pub mod endpoints; +pub mod model_config; -use image_endpoint::infer_image; +use endpoints::image_endpoint::infer_image; use once_cell::sync::Lazy; use std::sync::{Arc, RwLock}; diff --git a/crates/desktop/src-tauri/src/model_config.rs b/crates/desktop/src-tauri/src/model_config.rs new file mode 100644 index 0000000000..7dd709fbc6 --- /dev/null +++ b/crates/desktop/src-tauri/src/model_config.rs @@ -0,0 +1,59 @@ +use crate::ml::model_file::StableDiffusionVersion; +use candle_core::{DType, Device}; +use candle_transformers::models::stable_diffusion::StableDiffusionConfig; +use hf_hub::api::sync::Api; +use log::info; + +const SD_IMAGE_WIDTH: usize = 512; +const SD_IMAGE_HEIGHT: usize = 512; + +const SD_NUMERIC_DATATYPE: DType = DType::F32; +const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::Turbo; + +/// A central place to configure app-wide model details at startup. +/// We'll probably need much more fine-grained control later. +pub struct ModelConfig { + /// We only support one device type currently, and we can't multiplex + /// over several devices yet. + pub device: Device, + + /// The core numeric type to use for models. + pub dtype: DType, + + pub sd_version: StableDiffusionVersion, + pub sd_config: StableDiffusionConfig, + + /// Probably shouldn't be here. + pub hf_api: Api, +} + +impl ModelConfig { + + pub fn init() -> anyhow::Result { + + let device = Device::new_cuda(0) + .unwrap_or_else(|e| { + info!("CUDA not available ({}), falling back to CPU", e); + Device::Cpu + }); + + let sd_config = match SD_VERSION { + StableDiffusionVersion::Turbo => { + StableDiffusionConfig::sdxl_turbo(None, Some(SD_IMAGE_HEIGHT), Some(SD_IMAGE_WIDTH)) + } + _ => { + StableDiffusionConfig::v2_1(None, Some(SD_IMAGE_HEIGHT), Some(SD_IMAGE_WIDTH)) + } + }; + + let hf_api = Api::new()?; + + Ok(Self { + device, + dtype: SD_NUMERIC_DATATYPE, + sd_version: SD_VERSION, + sd_config, + hf_api, + }) + } +}