diff --git a/Cargo.lock b/Cargo.lock index 98f7affe53..c650215771 100644 --- a/Cargo.lock +++ b/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" diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 52db39a4bf..defd056341 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/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" diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index c38d90ab3b..d8522b802d 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/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, + image: &str, + model_config: State, model_cache: State, prompt_cache: State, ) -> Result { @@ -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, ) -> Result { @@ -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, }; diff --git a/crates/desktop/src-tauri/src/endpoints/sd.rs b/crates/desktop/src-tauri/src/endpoints/sd.rs index b884e8fe5c..e5c6d0c18c 100644 --- a/crates/desktop/src-tauri/src/endpoints/sd.rs +++ b/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, - pub width: Option, - pub n_steps: Option, - pub num_samples: usize, - pub seed: Option, - pub sd_version: StableDiffusionVersion, pub guidance_scale: Option, 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, }; diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 14ed9d6ff3..1d6fe53f1e 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/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) diff --git a/crates/desktop/src-tauri/src/model_config.rs b/crates/desktop/src-tauri/src/model_config.rs deleted file mode 100644 index 7dd709fbc6..0000000000 --- a/crates/desktop/src-tauri/src/model_config.rs +++ /dev/null @@ -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 { - - 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, - }) - } -} diff --git a/crates/desktop/src-tauri/src/state/app_config.rs b/crates/desktop/src-tauri/src/state/app_config.rs new file mode 100644 index 0000000000..e9064e1876 --- /dev/null +++ b/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, + + /// Probably shouldn't be here. + pub hf_api: Api, +} + +impl AppConfig { + + pub fn init() -> anyhow::Result { + + 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() +} diff --git a/crates/desktop/src-tauri/src/state/mod.rs b/crates/desktop/src-tauri/src/state/mod.rs new file mode 100644 index 0000000000..196a0ce480 --- /dev/null +++ b/crates/desktop/src-tauri/src/state/mod.rs @@ -0,0 +1,2 @@ +pub mod yaml_config; +pub mod app_config; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/state/yaml_config.rs b/crates/desktop/src-tauri/src/state/yaml_config.rs new file mode 100644 index 0000000000..cd0626338c --- /dev/null +++ b/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, + pub image_width: Option, + pub scheduler_steps: Option, + pub scheduler_samples: Option, + pub seed: Option, +} + +impl YamlConfig { + pub fn read_from_file>(path: P) -> anyhow::Result { + let file = File::open(path)?; + let reader = BufReader::new(file); + let config = serde_yaml::from_reader(reader)?; + Ok(config) + } +} diff --git a/crates/lib/files/filesys/src/lib.rs b/crates/lib/files/filesys/src/lib.rs index f2781858d1..604dff5349 100644 --- a/crates/lib/files/filesys/src/lib.rs +++ b/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; diff --git a/crates/lib/files/filesys/src/recursively_find_file_by_name.rs b/crates/lib/files/filesys/src/recursively_find_file_by_name.rs new file mode 100644 index 0000000000..a454238611 --- /dev/null +++ b/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, Q: AsRef>( + search_filename: P, + search_root: Q, + max_depth: usize, +) -> AnyhowResult> { + + 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()); + } +} \ No newline at end of file