wip

Cargo.lock

@@ -609,6 +609,7 @@ dependencies = [
  "candle-nn",
  "candle-transformers",
  "env_logger",
+ "filesys",
  "hf-hub",
  "image",
  "log",
@@ -617,6 +618,7 @@ dependencies = [
  "rand 0.8.5",
  "serde",
  "serde_json",
+ "serde_yaml",
  "tauri",
  "tauri-build",
  "tauri-plugin-log",
@@ -8715,6 +8717,19 @@ dependencies = [
  "syn 2.0.87",
 ]
 
+[[package]]
+name = "serde_yaml"
+version = "0.9.34+deprecated"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47"
+dependencies = [
+ "indexmap 2.2.6",
+ "itoa 1.0.11",
+ "ryu",
+ "serde",
+ "unsafe-libyaml",
+]
+
 [[package]]
 name = "serial_test"
 version = "3.0.0"
@@ -10969,6 +10984,12 @@ dependencies = [
  "void",
 ]
 
+[[package]]
+name = "unsafe-libyaml"
+version = "0.2.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861"
+
 [[package]]
 name = "untrusted"
 version = "0.9.0"

crates/desktop/src-tauri/Cargo.toml

@@ -20,6 +20,7 @@ tauri-build = { version = "2.0.4", features = [] }
 [dependencies]
 
 # Internal
+filesys = { path = "../../lib/files/filesys" }
 memory_caching = { path = "../../lib/caching/memory_caching" }
 
 # External
@@ -31,6 +32,7 @@ candle-transformers = { version = "0.8.2", features = ["cuda"] }
 env_logger = "0.11.6"
 hf-hub = "0.3.2"
 image = "0.25.5"
+serde_yaml = "0.9"
 log = "0.4"
 once_cell = "1.20.3"
 rand = "0.8"

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

@@ -1,25 +1,23 @@
-use std::io::Cursor;
-use std::path::PathBuf;
 use crate::endpoints::sd::{run, Args};
+use crate::state::yaml_config::YamlConfig;
 use crate::ml::model_cache::ModelCache;
-use crate::model_config::ModelConfig;
+use crate::ml::prompt_cache::PromptCache;
+use crate::state::app_config::AppConfig;
 use anyhow::Result;
 use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
 use image::imageops::FilterType;
 use image::ImageReader;
 use log::error;
+use std::io::Cursor;
+use std::path::PathBuf;
 use tauri::State;
-use crate::ml::prompt_cache::PromptCache;
 
 const PROMPT: &str = "A beautiful landscape with mountains and a lake";
 
-const HEIGHT: usize = 512;
-const WIDTH: usize = 512;
-
 #[tauri::command]
 pub fn infer_image(
-  image: &str, 
-  model_config: State<ModelConfig>, 
+  image: &str,
+  model_config: State<AppConfig>,
   model_cache: State<ModelCache>,
   prompt_cache: State<PromptCache>,
 ) -> Result<String, String> {
@@ -64,9 +62,9 @@ pub fn infer_image(
 }
 
 fn do_infer_image(
-  prompt: &str, 
-  image_path: &PathBuf, 
-  config: &ModelConfig, 
+  prompt: &str,
+  image_path: &PathBuf,
+  config: &AppConfig,
   model_cache: &ModelCache,
   prompt_cache: State<PromptCache>,
 ) -> Result<String, String> {
@@ -74,20 +72,12 @@ fn do_infer_image(
   
   let args = Args {
     image_path,
-    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(15),
-    num_samples: 1,
-    seed: None,
-    sd_version: config.sd_version.clone(),
     //guidance_scale: Some(0.0),
     guidance_scale: None,
     model_cache,
-    model_configs: config,
+    configs: config,
     prompt_cache: &prompt_cache,
   };
 

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

@@ -24,7 +24,7 @@ use crate::ml::load_image_file_to_tensor_2::load_image_file_to_tensor_2;
 use crate::ml::model_cache::ModelCache;
 use crate::ml::models::unet_model::UNetModel;
 use crate::ml::prompt_cache::PromptCache;
-use crate::model_config::ModelConfig;
+use crate::state::app_config::AppConfig;
 // TODO: Clean up
 
 // Note for Kasisnu: I'm going to start using lifetimes as long as that doesn't slow your velocity
@@ -32,43 +32,36 @@ use crate::model_config::ModelConfig;
 // itself with the 'a lifetime.
 pub struct Args<'a> {
     pub image_path: &'a Path,
-    pub api: Api,
     pub prompt: String,
     pub uncond_prompt: String,
-    pub cpu: bool,
-    pub height: Option<usize>,
-    pub width: Option<usize>,
-    pub n_steps: Option<usize>,
-    pub num_samples: usize,
-    pub seed: Option<u64>,
-    pub sd_version: StableDiffusionVersion,
     pub guidance_scale: Option<f64>,
     pub model_cache: &'a ModelCache,
-    pub model_configs: &'a ModelConfig,
+    pub configs: &'a AppConfig,
     pub prompt_cache: &'a PromptCache,
 }
 
 pub fn run(args: Args<'_>) -> Result<()> {
     println!("Starting image generation with the following configuration:");
-    println!("  Model: {:?}", args.sd_version);
+    println!("  Model: {:?}", args.configs.sd_version);
     println!("  Prompt: {}", args.prompt);
-    println!("  Steps: {}", args.n_steps.unwrap_or(1));
-    println!("  Using CPU: {}", args.cpu);
+    println!("  Steps: {}", args.configs.scheduler_steps);
+    println!("  Device: {:?}", args.configs.device);
 
 
-    println!("Model dimensions: {}x{}", args.model_configs.sd_config.width, args.model_configs.sd_config.height);
+    println!("Model dimensions: {}x{}", args.configs.sd_config.width, args.configs.sd_config.height);
 
     println!("Building scheduler...");
-    let mut scheduler = args.model_configs.sd_config.build_scheduler(args.n_steps.unwrap_or(1))?;
+    let mut scheduler = args.configs.sd_config.build_scheduler(
+        args.configs.scheduler_steps)?;
 
-    let seed = args.seed.unwrap_or_else(|| rand::thread_rng().gen());
+    let seed = args.configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
     println!("Using seed: {}", seed);
 
-    args.model_configs.device.set_seed(seed)?;
+    args.configs.device.set_seed(seed)?;
 
     println!("Initializing Hugging Face API...");
 
-    let repo = args.sd_version.repo();
+    let repo = args.configs.sd_version.repo();
     println!("Downloading model files from: {}", repo);
     
     //println!("Downloading VAE ...");
@@ -109,11 +102,11 @@ pub fn run(args: Args<'_>) -> Result<()> {
             None, // tokenizer
             None, // clip_weights
             None, // clip2_weights
-            args.sd_version,
-            &args.model_configs.sd_config,
+            args.configs.sd_version,
+            &args.configs.sd_config,
             false, // use_f16
-            &args.model_configs.device,
-            args.model_configs.dtype,
+            &args.configs.device,
+            args.configs.dtype,
             false, // use_guide_scale
         )?;
         
@@ -128,8 +121,8 @@ pub fn run(args: Args<'_>) -> Result<()> {
 
     let input_image = load_image_file_to_tensor_2(args.image_path)?;
 
-    let input_image = input_image.to_device(&args.model_configs.device)?
-      .to_dtype(args.model_configs.dtype)?;
+    let input_image = input_image.to_device(&args.configs.device)?
+      .to_dtype(args.configs.dtype)?;
 
     // REFERENCE IMAGE shape: [1, 3, 1024, 1024]
     println!("REFERENCE IMAGE shape: {:?}", input_image.shape());
@@ -141,15 +134,15 @@ pub fn run(args: Args<'_>) -> Result<()> {
     let img2img_strength = 0.75f64;
     
     let t_start = {
-        let start = args.n_steps.unwrap_or(1) - (args.n_steps.unwrap_or(1) as f64 * img2img_strength) as usize;
+        let start = args.configs.scheduler_steps - (args.configs.scheduler_steps as f64 * img2img_strength) as usize;
         
-        println!("Starting from step {} of {} (strength: {})", start, args.n_steps.unwrap_or(1), img2img_strength);
+        println!("Starting from step {} of {} (strength: {})", start, args.configs.scheduler_steps, img2img_strength);
         start
     };
 
     let use_guide_scale = false; // TODO
 
-    let vae_scale = match args.model_configs.sd_version {
+    let vae_scale = match args.configs.sd_version {
         StableDiffusionVersion::V1_5
         | StableDiffusionVersion::V1_5Inpaint
         | StableDiffusionVersion::V2_1
@@ -165,7 +158,7 @@ pub fn run(args: Args<'_>) -> Result<()> {
     let init_latent_dist = encoded_image;
 
     println!("Generating latents from input image...");
-    let latents = (init_latent_dist.sample()? * vae_scale)?.to_device(&args.model_configs.device)?;
+    let latents = (init_latent_dist.sample()? * vae_scale)?.to_device(&args.configs.device)?;
     
     println!("Initial latents shape: {:?}", latents.shape());
 
@@ -182,7 +175,7 @@ pub fn run(args: Args<'_>) -> Result<()> {
 
 
 
-    let vae_scale = match args.sd_version {
+    let vae_scale = match args.configs.sd_version {
         StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint | StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Xl => 0.18215,
         StableDiffusionVersion::Turbo => 0.13025,
     };

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

@@ -1,14 +1,15 @@
-pub mod ml;
 pub mod endpoints;
-pub mod model_config;
+pub mod ml;
+pub mod state;
 
 use crate::ml::model_cache::ModelCache;
 use crate::ml::prompt_cache::PromptCache;
-use crate::model_config::ModelConfig;
+use crate::state::app_config::AppConfig;
 use endpoints::image_endpoint::infer_image;
 use env_logger;
 use log::info;
 use once_cell::sync::Lazy;
+use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
 
 #[cfg_attr(mobile, tauri::mobile_entry_point)]
@@ -17,7 +18,7 @@ pub fn run() {
 
   info!("Loading model config...");
 
-  let model_config = ModelConfig::init()
+  let model_config = AppConfig::init()
     .expect("config should load");
 
   let prompt_cache = PromptCache::with_capacity(8)

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

@@ -1,59 +0,0 @@
-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,
-    })
-  }
-}

crates/desktop/src-tauri/src/state/app_config.rs

@@ -0,0 +1,88 @@
+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;
+use filesys::recursively_find_file_by_name::recursively_find_file_by_name;
+use crate::state::yaml_config::YamlConfig;
+
+const DEFAULT_SD_IMAGE_WIDTH: usize = 512;
+const DEFAULT_SD_IMAGE_HEIGHT: usize = 512;
+
+const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
+const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::Turbo;
+
+/// A central place to configure app-wide details.
+pub struct AppConfig {
+  /// 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,
+
+  pub image_height: usize,
+  pub image_width: usize,
+  pub scheduler_steps: usize,
+  pub scheduler_samples: usize,
+  
+  pub seed: Option<u64>,
+  
+  /// Probably shouldn't be here.
+  pub hf_api: Api,
+}
+
+impl AppConfig {
+  
+  pub fn init() -> anyhow::Result<Self> {
+    
+    let yaml_configs = load_yaml_configs();
+    
+    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(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
+      _ => {
+        StableDiffusionConfig::v2_1(None, Some(DEFAULT_SD_IMAGE_HEIGHT), Some(DEFAULT_SD_IMAGE_WIDTH))
+      }
+    };
+    
+    let hf_api = Api::new()?;
+    
+    Ok(Self {
+      device,
+      dtype: DEFAULT_SD_NUMERIC_DATATYPE,
+      sd_version: SD_VERSION,
+      sd_config,
+      image_height: yaml_configs.image_height.unwrap_or(DEFAULT_SD_IMAGE_HEIGHT),
+      image_width: yaml_configs.image_width.unwrap_or(DEFAULT_SD_IMAGE_WIDTH),
+      scheduler_steps: yaml_configs.scheduler_steps.unwrap_or(1),
+      scheduler_samples: yaml_configs.scheduler_samples.unwrap_or(15),
+      seed: yaml_configs.seed,
+      hf_api,
+    })
+  }
+}
+
+fn load_yaml_configs() -> YamlConfig {
+  let maybe_path = recursively_find_file_by_name("app_config.yaml", "../../../", 5)
+    .ok()
+    .flatten();
+  if let Some(path) = maybe_path {
+    info!("Attempting to load configs from: {:?}", &path);
+    if let Ok(config) = YamlConfig::read_from_file(path) {
+      return config;
+    }
+  }
+  info!("Loading default configs.");
+  YamlConfig::default()
+}

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

@@ -0,0 +1,2 @@
+pub mod yaml_config;
+pub mod app_config;

crates/desktop/src-tauri/src/state/yaml_config.rs

@@ -0,0 +1,24 @@
+use std::fs::File;
+use std::io::BufReader;
+use std::path::Path;
+
+/// Configuration read or derived from a YAML config file at startup
+/// This is not meant to contain any runtime config values.
+/// This may not be exhaustive.
+#[derive(serde::Deserialize, serde::Serialize, Clone, Debug, Default)]
+pub struct YamlConfig {
+  pub image_height: Option<usize>,
+  pub image_width: Option<usize>,
+  pub scheduler_steps: Option<usize>,
+  pub scheduler_samples: Option<usize>,
+  pub seed: Option<u64>,
+}
+
+impl YamlConfig {
+  pub fn read_from_file<P: AsRef<Path>>(path: P) -> anyhow::Result<Self> {
+    let file = File::open(path)?;
+    let reader = BufReader::new(file);
+    let config = serde_yaml::from_reader(reader)?;
+    Ok(config)
+  }
+}

crates/lib/files/filesys/src/lib.rs

@@ -31,4 +31,5 @@ pub mod file_size;
 pub mod filename_concat;
 pub mod is_filesystem_full_error;
 pub mod path_to_string;
+pub mod recursively_find_file_by_name;
 pub mod rename_across_devices;

crates/lib/files/filesys/src/recursively_find_file_by_name.rs

@@ -0,0 +1,56 @@
+use std::path::{Path, PathBuf};
+use walkdir::WalkDir;
+use errors::AnyhowResult;
+
+pub fn recursively_find_file_by_name<P: AsRef<Path>, Q: AsRef<Path>>(
+  search_filename: P,
+  search_root: Q,
+  max_depth: usize,
+) -> AnyhowResult<Option<PathBuf>> {
+
+  let walker = WalkDir::new(search_root)
+    .max_depth(max_depth);
+
+  for entry in walker {
+    let entry = entry?;
+    if entry.path().is_dir() {
+      let check_path = entry.path().join(search_filename.as_ref());
+      if check_path.exists() && check_path.is_file() {
+        return Ok(Some(check_path));
+      }
+    }
+  }
+
+  Ok(None)
+}
+
+#[cfg(test)]
+mod tests {
+  use testing::test_file_path::test_file_path;
+  use crate::recursively_find_file_by_name::recursively_find_file_by_name;
+
+  #[test]
+  fn test_path_exists() {
+    let test_data_path = test_file_path("test_data/").expect("should work");
+    let result = recursively_find_file_by_name(
+      "super_mario_rpg_beware_the_forests_mushrooms.mp3",
+      test_data_path,
+      5
+    );
+    let result = result.expect("should be ok");
+    let result = result.expect("should exist");
+    assert!(result.ends_with("test_data/audio/mp3/super_mario_rpg_beware_the_forests_mushrooms.mp3"));
+  }
+
+  #[test]
+  fn test_path_does_not_exist() {
+    let test_data_path = test_file_path("test_data/").expect("should work");
+    let result = recursively_find_file_by_name(
+      "fake_name.foo",
+      test_data_path,
+      5
+    );
+    let result = result.expect("should be ok");
+    assert!(result.is_none());
+  }
+}