config object

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<Arc<Mutex<Option<ModelConfig>>>> = Lazy::new(|| Arc::new(Mutex::new(None)));
+
 static MODEL_CACHE: Lazy<Arc<RwLock<Option<ModelCache>>>> = Lazy::new(|| Arc::new(RwLock::new(None)));
 
 #[tauri::command]
 pub fn infer_image(image: &str) -> Result<String, String> {
+    // 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<String, String> {
     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<String, String> {
         }
     }
     
-    
-
     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<String, String> {
                     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))?;

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

@@ -0,0 +1,3 @@
+pub mod flip_image;
+pub mod image_endpoint;
+pub mod sd;

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};
 

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<Self> {
+    
+    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,
+    })
+  }
+}