weights directory, assets directory

2026-10-20 17:02:53 -04:00 · 2025-03-06 21:06:58 -05:00
parent 1a887d7414
commit ef3527dd85
9 changed files with 289 additions and 34 deletions
@@ -2,6 +2,7 @@
  <dictionary name="bt">
    <words>
      <w>arpabet</w>
+      <w>artcraft</w>
      <w>backfill</w>
      <w>backfilling</w>
      <w>backfills</w>
@@ -664,6 +664,7 @@ dependencies = [
 "candle-nn",
 "candle-transformers",
 "cudarc",
+ "expanduser",
 "filesys",
 "hf-hub 0.3.2",
 "image",
@@ -732,6 +733,12 @@ version = "0.3.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb"

+[[package]]
+name = "arrayvec"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23b62fc65de8e4e7f52534fb52b0f3ed04746ae267519eef2a83941e8085068b"
+
 [[package]]
 name = "arrayvec"
 version = "0.7.6"
@@ -1015,7 +1022,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6678909d8c5d46a42abcf571271e15fdbc0a225e3646cf23762cd415046c78bf"
 dependencies = [
 "anyhow",
- "arrayvec",
+ "arrayvec 0.7.6",
 "log",
 "nom",
 "num-rational",
@@ -1028,7 +1035,7 @@ version = "0.8.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e335041290c43101ca215eed6f43ec437eb5a42125573f600fc3fa42b9bddd62"
 dependencies = [
- "arrayvec",
+ "arrayvec 0.7.6",
 ]

 [[package]]
@@ -1298,6 +1305,17 @@ dependencies = [
 "wyz",
 ]

+[[package]]
+name = "blake2b_simd"
+version = "0.5.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "afa748e348ad3be8263be728124b24a24f268266f6f5d58af9d75f6a40b5c587"
+dependencies = [
+ "arrayref",
+ "arrayvec 0.5.2",
+ "constant_time_eq 0.1.5",
+]
+
 [[package]]
 name = "blake2b_simd"
 version = "1.0.2"
@@ -1305,7 +1323,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "23285ad32269793932e830392f2fe2f83e26488fd3ec778883a93c8323735780"
 dependencies = [
 "arrayref",
- "arrayvec",
+ "arrayvec 0.7.6",
 "constant_time_eq 0.3.0",
 ]

@@ -3188,6 +3206,17 @@ dependencies = [
 "dirs-sys 0.4.1",
 ]

+[[package]]
+name = "dirs"
+version = "1.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3fd78930633bd1c6e35c4b42b1df7b0cbc6bc191146e512bb3bedf243fcc3901"
+dependencies = [
+ "libc",
+ "redox_users 0.3.5",
+ "winapi 0.3.9",
+]
+
 [[package]]
 name = "dirs"
 version = "5.0.1"
@@ -3838,6 +3867,17 @@ version = "2.5.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0"

+[[package]]
+name = "expanduser"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "14e0b79235da57db6b6c2beed9af6e5de867d63a973ae3e91910ddc33ba40bc0"
+dependencies = [
+ "dirs 1.0.5",
+ "lazy_static",
+ "pwd",
+]
+
 [[package]]
 name = "exr"
 version = "1.73.0"
@@ -5860,7 +5900,7 @@ checksum = "50ae7e0018905a795d6f2a60ac32a547490abdd8df509906a8c6171e6d861711"
 dependencies = [
 "anyhow",
 "binstring",
- "blake2b_simd",
+ "blake2b_simd 1.0.2",
 "coarsetime",
 "ct-codecs",
 "ed25519-compact",
@@ -6819,7 +6859,7 @@ dependencies = [
 "actix-rt",
 "ahash 0.8.11",
 "anstream",
- "arrayvec",
+ "arrayvec 0.7.6",
 "base64 0.21.7",
 "base64 0.22.1",
 "base64ct",
@@ -8596,6 +8636,16 @@ dependencies = [
 "version_check",
 ]

+[[package]]
+name = "pwd"
+version = "1.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72c71c0c79b9701efe4e1e4b563b2016dd4ee789eb99badcb09d61ac4b92e4a2"
+dependencies = [
+ "libc",
+ "thiserror 1.0.63",
+]
+
 [[package]]
 name = "pyo3"
 version = "0.22.6"
@@ -8921,7 +8971,7 @@ checksum = "cd87ce80a7665b1cce111f8a16c1f3929f6547ce91ade6addf4ec86a8dda5ce9"
 dependencies = [
 "arbitrary",
 "arg_enum_proc_macro",
- "arrayvec",
+ "arrayvec 0.7.6",
 "av1-grain",
 "bitstream-io",
 "built",
@@ -9150,6 +9200,17 @@ dependencies = [
 "bitflags 2.5.0",
 ]

+[[package]]
+name = "redox_users"
+version = "0.3.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "de0737333e7a9502c789a36d7c7fa6092a49895d4faa31ca5df163857ded2e9d"
+dependencies = [
+ "getrandom 0.1.16",
+ "redox_syscall 0.1.57",
+ "rust-argon2",
+]
+
 [[package]]
 name = "redox_users"
 version = "0.4.6"
@@ -9436,6 +9497,18 @@ dependencies = [
 "realfft",
 ]

+[[package]]
+name = "rust-argon2"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4b18820d944b33caa75a71378964ac46f58517c92b6ae5f762636247c09e78fb"
+dependencies = [
+ "base64 0.13.1",
+ "blake2b_simd 0.5.11",
+ "constant_time_eq 0.1.5",
+ "crossbeam-utils 0.8.20",
+]
+
 [[package]]
 name = "rust-embed"
 version = "8.0.0"
@@ -9521,7 +9594,7 @@ version = "1.36.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b082d80e3e3cc52b2ed634388d436fe1f4de6af5786cc2de9ba9737527bdf555"
 dependencies = [
- "arrayvec",
+ "arrayvec 0.7.6",
 "borsh",
 "bytes 1.10.0",
 "num-traits",
@@ -11199,7 +11272,7 @@ name = "symphonia-core"
 version = "0.5.0"
 source = "git+https://github.com/pdeljanov/Symphonia.git?rev=e1a70095c4409478c1ded499fb8dd4afa540fb90#e1a70095c4409478c1ded499fb8dd4afa540fb90"
 dependencies = [
- "arrayvec",
+ "arrayvec 0.7.6",
 "bitflags 1.3.2",
 "bytemuck",
 "lazy_static",
@@ -11212,7 +11285,7 @@ version = "0.5.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "798306779e3dc7d5231bd5691f5a813496dc79d3f56bf82e25789f2094e022c3"
 dependencies = [
- "arrayvec",
+ "arrayvec 0.7.6",
 "bitflags 1.3.2",
 "bytemuck",
 "lazy_static",
@@ -38,6 +38,7 @@ image = "0.25.5"
 serde_yaml = "0.9"
 log = "0.4"
 once_cell = "1.20.3"
+expanduser = "1.2.2"
 rand = "0.8"
 reqwest = "0.12.12"
 serde = { version = "1.0", features = ["derive"] }
@@ -1,16 +1,17 @@
+use crate::ml::model_registry::ModelRegistry;
+use log::info;
 use std::fs::File;
 use std::path::{Path, PathBuf};
-use log::info;
-use crate::ml::model_registry::ModelRegistry;

 pub async fn download_model_registry_file(model_type: ModelRegistry) -> anyhow::Result<PathBuf> {
-  info!("downloading model: {:?}", model_type);
  let path = PathBuf::from(model_type.get_filename());
+
  if path.exists() {
-    info!("model already exists: {:?}", path);
    return Ok(path);
  }

+  info!("downloading model: {:?}", model_type);
+
  let body = reqwest::get(model_type.get_download_url())
    .await?
    .bytes()
@@ -21,4 +22,4 @@ pub async fn download_model_registry_file(model_type: ModelRegistry) -> anyhow::
  std::fs::write(&path, body.as_ref())?;
  
  Ok(path)
-}
+}
@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+use expanduser::expanduser;
 use crate::ml::model_file::StableDiffusionVersion;
 use crate::state::yaml_config::YamlConfig;
 use candle_core::{DType, Device};
@@ -5,8 +7,8 @@ use candle_transformers::models::stable_diffusion::StableDiffusionConfig;
 use filesys::recursively_find_file_by_name::recursively_find_file_by_name;
 use hf_hub::api::sync::Api;
 use log::info;
-
-const CONFIG_FILENAME : &str = "app_config.yaml";
+use crate::state::app_dir::AppDataRoot;
+use crate::state::os_platform::OsPlatform;

 const DEFAULT_SD_IMAGE_WIDTH: usize = 512;
 const DEFAULT_SD_IMAGE_HEIGHT: usize = 512;
@@ -14,6 +16,8 @@ const DEFAULT_SD_IMAGE_HEIGHT: usize = 512;
 const DEFAULT_SD_NUMERIC_DATATYPE: DType = DType::F32;
 const SD_VERSION: StableDiffusionVersion = StableDiffusionVersion::Turbo;

+const DEFAULT_DATA_DIR_UNIX : &str = "~/artcraft";
+
 /// A central place to configure app-wide details.
 pub struct AppConfig {
  /// We only support one device type currently, and we can't multiplex 
@@ -37,13 +41,15 @@ pub struct AppConfig {
  
  /// Probably shouldn't be here.
  pub hf_api: Api,
+  
+  /// Root location for application data.
+  pub app_data_root: AppDataRoot,
 }

 impl AppConfig {
  
  pub fn init() -> anyhow::Result<Self> {
-    println!("Configuration filename: {}", CONFIG_FILENAME);
-    let yaml_configs = load_yaml_configs();
+    let yaml_configs = YamlConfig::load_from_config_file_recursive();
    
    println!("Configs: {:?}", yaml_configs);
    
@@ -64,6 +70,26 @@ impl AppConfig {
    
    let hf_api = Api::new()?;
    
+    let os_platform = OsPlatform::get();
+
+    let default_data_path = match os_platform {
+      OsPlatform::Linux => yaml_configs.linux_default_data_path.as_deref(),
+      OsPlatform::MacOs => yaml_configs.mac_default_data_path.as_deref(),
+      OsPlatform::Windows => yaml_configs.windows_default_data_path.as_deref(),
+    };
+    
+    let data_root_path = match default_data_path {
+      Some(default_data_path) => default_data_path.to_string(),
+      None => match os_platform { 
+        OsPlatform::Linux | OsPlatform::MacOs => DEFAULT_DATA_DIR_UNIX.to_string(),
+        OsPlatform::Windows => ".".to_string(),
+      }
+    };
+      
+    info!("Using data root: {}", data_root_path);
+
+    let app_data_root = AppDataRoot::create_existing(data_root_path)?;
+    
    Ok(Self {
      device,
      dtype: DEFAULT_SD_NUMERIC_DATATYPE,
@@ -76,21 +102,7 @@ impl AppConfig {
      seed: yaml_configs.seed,
      cfg_scale: yaml_configs.cfg_scale,
      hf_api,
+      app_data_root,
    })
  }
 }
-
-fn load_yaml_configs() -> YamlConfig {
-  let maybe_path = recursively_find_file_by_name(CONFIG_FILENAME, "../../../", 5)
-    .ok()
-    .flatten();
-  if let Some(path) = maybe_path {
-    let path = path.canonicalize().unwrap_or(path);
-    println!("Attempting to load configs from: {:?}", &path);
-    if let Ok(config) = YamlConfig::read_from_file(path) {
-      return config;
-    }
-  }
-  println!("Loading default configs.");
-  YamlConfig::default()
-}
@@ -0,0 +1,118 @@
+use crate::state::os_platform::OsPlatform;
+use expanduser::expanduser;
+use std::path::{Path, PathBuf};
+
+const ASSETS_SUBDIRECTORY : &str = "assets";
+const WEIGHTS_SUBDIRECTORY : &str = "weights";
+
+/// The path to the application data directory, which includes "asset" and "weights" data.
+#[derive(Clone)]
+pub struct AppDataRoot {
+  path: PathBuf,
+  assets_dir: AppAssetsDir,
+  weights_dir: AppWeightsDir,
+}
+
+#[derive(Clone)]
+pub struct AppAssetsDir {
+  path: PathBuf,
+}
+
+#[derive(Clone)]
+pub struct AppWeightsDir {
+  path: PathBuf,
+}
+
+impl AppDataRoot {
+  pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
+    let mut dir = dir.as_ref().to_path_buf();
+    
+    match OsPlatform::get() {
+      OsPlatform::Linux | OsPlatform::MacOs => {
+        if let Some(d) = dir.as_os_str().to_str() {
+          dir = expanduser(d)?;
+        }
+      },
+      OsPlatform::Windows => {}
+    }
+    
+    if !dir.is_dir() {
+      println!("Creating directory {:?}", dir);
+      std::fs::create_dir_all(&dir)?;
+    }
+
+    match dir.canonicalize() {
+      Ok(d) => dir = d,
+      Err(err) => {
+        println!("Error canonicalizing {:?}: {}", dir, err);
+      }
+    }
+    
+    let assets_dir = AppAssetsDir::create_existing(dir.join(ASSETS_SUBDIRECTORY))?;
+    let weights_dir = AppWeightsDir::create_existing(dir.join(WEIGHTS_SUBDIRECTORY))?;
+    
+    Ok(Self {
+      path: dir,
+      assets_dir,
+      weights_dir,
+    })
+  }
+  
+  pub fn assets_dir(&self) -> &AppAssetsDir {
+    &self.assets_dir
+  }
+  
+  pub fn weights_dir(&self) -> &AppWeightsDir {
+    &self.weights_dir
+  }
+  
+  pub fn path(&self) -> &Path {
+    &self.path
+  }
+}
+
+impl AppAssetsDir {
+  pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
+    let mut dir = dir.as_ref().to_path_buf();
+    match dir.canonicalize() {
+      Ok(d) => dir = d,
+      Err(err) => {
+        println!("Error canonicalizing {:?}: {}", dir, err);
+      }
+    }
+    if !dir.exists() {
+      println!("Creating directory {:?}", dir);
+      std::fs::create_dir(&dir)?;
+    }
+    Ok(Self {
+      path: dir,
+    })
+  }
+
+  pub fn path(&self) -> &Path {
+    &self.path
+  }
+}
+
+impl AppWeightsDir {
+  pub fn create_existing<P: AsRef<Path>>(dir: P) -> anyhow::Result<Self> {
+    let mut dir = dir.as_ref().to_path_buf();
+    match dir.canonicalize() {
+      Ok(d) => dir = d,
+      Err(err) => {
+        println!("Error canonicalizing {:?}: {}", dir, err);
+      }
+    }
+    if !dir.exists() {
+      println!("Creating directory {:?}", dir);
+      std::fs::create_dir(&dir)?;
+    }
+    Ok(Self {
+      path: dir,
+    })
+  }
+
+  pub fn path(&self) -> &Path {
+    &self.path
+  }
+}
@@ -1,2 +1,4 @@
 pub mod yaml_config;
-pub mod app_config;
+pub mod app_config;
+mod app_dir;
+mod os_platform;
@@ -0,0 +1,20 @@
+
+#[derive(Copy,Clone,Eq,PartialEq,Debug)]
+pub enum OsPlatform {
+  Linux,
+  MacOs,
+  Windows,
+}
+
+impl OsPlatform {
+  pub fn get() -> OsPlatform {
+    // Totally scientific.
+    // https://doc.rust-lang.org/std/env/consts/constant.OS.html
+    match std::env::consts::OS {
+      "linux" => OsPlatform::Linux,
+      "windows" => OsPlatform::Windows,
+      "macos" | "apple" => OsPlatform::MacOs,
+      _ => OsPlatform::Linux,
+    }
+  }
+}
@@ -1,7 +1,10 @@
+use filesys::recursively_find_file_by_name::recursively_find_file_by_name;
 use std::fs::File;
 use std::io::BufReader;
 use std::path::Path;

+const CONFIG_FILENAME : &str = "app_config.yaml";
+
 /// 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.
@@ -13,6 +16,15 @@ pub struct YamlConfig {
  pub scheduler_samples: Option<usize>,
  pub seed: Option<u64>,
  pub cfg_scale: Option<f64>,
+
+  /// Override the default location to store data on Linux
+  pub linux_default_data_path: Option<String>,
+
+  /// Override the default location to store data on Windows
+  pub windows_default_data_path: Option<String>,
+  
+  /// Override the default location to store data on Mac
+  pub mac_default_data_path: Option<String>,
 }

 impl YamlConfig {
@@ -22,4 +34,19 @@ impl YamlConfig {
    let config = serde_yaml::from_reader(reader)?;
    Ok(config)
  }
+
+  pub fn load_from_config_file_recursive() -> YamlConfig {
+    let maybe_path = recursively_find_file_by_name(CONFIG_FILENAME, "../../../", 5)
+      .ok()
+      .flatten();
+    if let Some(path) = maybe_path {
+      let path = path.canonicalize().unwrap_or(path);
+      println!("Attempting to load configs from: {:?}", &path);
+      if let Ok(config) = YamlConfig::read_from_file(path) {
+        return config;
+      }
+    }
+    println!("Couldn't load config file. Using default configs.");
+    YamlConfig::default()
+  }
 }