diff --git a/crates/desktop/ml_models/src/ml/stable_diffusion/app_dir.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/app_dir.rs deleted file mode 100644 index f59b9abf83..0000000000 --- a/crates/desktop/ml_models/src/ml/stable_diffusion/app_dir.rs +++ /dev/null @@ -1,220 +0,0 @@ -use crate::ml::weights_registry::weight_descriptor::WeightDescriptor; -use anyhow::anyhow; -use directories::UserDirs; -use once_cell::sync::Lazy; -use std::io; -use std::path::{Path, PathBuf, MAIN_SEPARATOR}; -use tempdir::TempDir; -use tempfile::{Builder, NamedTempFile}; - -static PREFIX: Lazy = Lazy::new(|| format!("~{}", MAIN_SEPARATOR)); - -const DEFAULT_DATA_DIR : &str = "artcraft"; -const ASSETS_SUBDIRECTORY : &str = "assets"; -const WEIGHTS_SUBDIRECTORY : &str = "weights"; - -const TEMPORARY_SUBDIRECTORY : &str = "temp"; - -/// 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, - temp_dir: TemporaryDir, -} - -#[derive(Clone)] -pub struct AppAssetsDir { - path: PathBuf, -} - -#[derive(Clone)] -pub struct AppWeightsDir { - path: PathBuf, -} - -#[derive(Clone)] -pub struct TemporaryDir { - path: PathBuf, -} - -impl AppDataRoot { - pub fn create_default() -> anyhow::Result { - let directory = get_default_data_dir()?; - Self::create_existing(directory) - } - - pub fn create_existing>(dir: P) -> anyhow::Result { - let mut dir = dir.as_ref().to_path_buf(); - - dir = expanduser(dir.to_str().map(|s| s.to_string()).unwrap())?; - - 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))?; - let temp_dir = TemporaryDir::create_existing(dir.join(TEMPORARY_SUBDIRECTORY))?; - - Ok(Self { - path: dir, - assets_dir, - weights_dir, - temp_dir, - }) - } - - pub fn assets_dir(&self) -> &AppAssetsDir { - &self.assets_dir - } - - pub fn weights_dir(&self) -> &AppWeightsDir { - &self.weights_dir - } - - pub fn temp_dir(&self) -> &TemporaryDir { - &self.temp_dir - } - - pub fn path(&self) -> &Path { - &self.path - } -} - -impl AppAssetsDir { - pub fn create_existing>(dir: P) -> anyhow::Result { - 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>(dir: P) -> anyhow::Result { - 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 - } - - pub fn weight_path(&self, descriptor: &WeightDescriptor) -> PathBuf { - self.path.join(descriptor.filename) - } -} - -impl TemporaryDir { - pub fn create_existing>(dir: P) -> anyhow::Result { - 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 with_prefix(&self, prefix: &str) -> anyhow::Result { - let tempdir = TempDir::new_in(&self.path, prefix)?; - Ok(tempdir) - } - - pub fn new_named_temp_file(&self) -> anyhow::Result { - let tempfile = Builder::new() - .prefix("temp_") - .suffix(".bin") - .tempfile_in(&self.path)?; - Ok(tempfile) - } -} - -// eg. /home/bob/artcraft, /Users/bob/artcraft, or C:\Users\Alice\artcraft -fn get_default_data_dir() -> anyhow::Result { - Ok(UserDirs::new() - .ok_or_else(|| anyhow!("could not determine user home directory"))? - .home_dir() - .join(DEFAULT_DATA_DIR)) -} - - -/// The "expanduser" crate doesn't compile on Windows, so we replace its functionality slightly -pub fn expanduser>(path: P) -> io::Result { - - Ok(match path.as_ref() { - // matches an exact "~" - s if s == "~" => { - home_dir()? - }, - // matches paths that start with `~/` - s if s.starts_with(&*PREFIX) => { - let home = home_dir()?; - home.join(&s[2..]) - }, - // // matches paths that start with `~` but not `~/`, might be a `~username/` path - // s if s.starts_with("~") => { - // let mut parts = s[1..].splitn(2, MAIN_SEPARATOR); - // let user = parts.next() - // .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "malformed path"))?; - // let user = Passwd::from_name(&user) - // .map_err(|_| io::Error::new(io::ErrorKind::Other, "error searching for user"))? - // .ok_or_else(|| io::Error::new(io::ErrorKind::Other, format!("user '{}', does not exist", &user)))?; - // if let Some(ref path) = parts.next() { - // PathBuf::from(user.dir).join(&path) - // } else { - // PathBuf::from(user.dir) - // } - // }, - // nothing to expand, just make a PathBuf - s => PathBuf::from(s) - }) -} - -pub fn home_dir() -> io::Result { - dirs::home_dir() - .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no home directory is set")) -} diff --git a/crates/desktop/ml_models/src/ml/stable_diffusion/clip.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/clip.rs deleted file mode 100644 index 007f32f6e7..0000000000 --- a/crates/desktop/ml_models/src/ml/stable_diffusion/clip.rs +++ /dev/null @@ -1,134 +0,0 @@ -#[cfg(feature = "accelerate")] -extern crate accelerate_src; -#[cfg(feature = "mkl")] -extern crate intel_mkl_src; - -use candle_transformers::models::stable_diffusion; -use std::path::PathBuf; - -use crate::ml::model_file::{ModelFile, StableDiffusionVersion}; -use crate::ml::stable_diffusion::app_dir::AppWeightsDir; -use crate::ml::weights_registry::weights::{CLIP_JSON, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2}; -use anyhow::Error as E; -use candle_core::{DType, Device, Module, Tensor, D}; -use log::info; -use tokenizers::Tokenizer; - -#[allow(clippy::too_many_arguments)] -pub fn old_clip( - prompt: &str, - uncond_prompt: &str, - tokenizer: Option, - clip_weights: Option, - clip2_weights: Option, - sd_version: StableDiffusionVersion, - sd_config: &stable_diffusion::StableDiffusionConfig, - use_f16: bool, - device: &Device, - dtype: DType, - use_guide_scale: bool, - weights_dir: &AppWeightsDir -) -> anyhow::Result { - info!("Preparing text embeddings... for {:?}", sd_version); - - let which = match sd_version { - StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Turbo => vec![true, false], - _ => vec![true], - }; - - let text_embeddings = which - .iter() - .map(|first| { - do_infer_clip_text_embeddings( - &prompt, - &uncond_prompt, - None, // tokenizer - clip_weights.clone(), - None, // clip2_weights - sd_version, - &sd_config, - false, // use_f16 - &device, - dtype, - use_guide_scale, - *first, - weights_dir, - ) - }) - .collect::>>()?; - - let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?; - - Ok(text_embeddings) -} - -#[allow(clippy::too_many_arguments)] -fn do_infer_clip_text_embeddings( - prompt: &str, - uncond_prompt: &str, - tokenizer: Option, - clip_weights: Option, - clip2_weights: Option, - sd_version: StableDiffusionVersion, - sd_config: &stable_diffusion::StableDiffusionConfig, - use_f16: bool, - device: &Device, - dtype: DType, - use_guide_scale: bool, - first: bool, - weights_dir: &AppWeightsDir -) -> anyhow::Result { - info!("do_infer_clip_text_embeddings called with args {:?}", (prompt, uncond_prompt, tokenizer.clone().unwrap_or_else(|| "None".to_string()), clip_weights.clone().unwrap_or_else(|| "None".to_string()), clip2_weights.clone().unwrap_or_else(|| "None".to_string()), sd_version, sd_config, use_f16, device, dtype, use_guide_scale, first)); - - let tokenizer = weights_dir.weight_path(&CLIP_JSON); - - info!("Loading Clip Tokenizer path: {:?}", tokenizer); - - let tokenizer = Tokenizer::from_file(tokenizer).map_err(E::msg)?; - - let pad_id = match &sd_config.clip.pad_with { - Some(padding) => *tokenizer.get_vocab(true).get(padding.as_str()).unwrap(), - None => *tokenizer.get_vocab(true).get("<|endoftext|>").unwrap(), - }; - - println!("Running with prompt \"{prompt}\"."); - let mut tokens = tokenizer.encode(prompt, true).map_err(E::msg)?.get_ids().to_vec(); - if tokens.len() > sd_config.clip.max_position_embeddings { - anyhow::bail!("the prompt is too long, {} > max-tokens ({})", tokens.len(), sd_config.clip.max_position_embeddings) - } - while tokens.len() < sd_config.clip.max_position_embeddings { - tokens.push(pad_id) - } - let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?; - - println!("Building the Clip transformer."); - let clip_weights = if first { - println!(">>>>>>>>> CLIP 1"); - weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER) - } else { - println!(">>>>>>>>> CLIP 2"); - weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER_2) - }; - let clip_config = if first { &sd_config.clip } else { sd_config.clip2.as_ref().unwrap() }; - let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?; - - let text_embeddings = text_model.forward(&tokens)?; - - let text_embeddings = if use_guide_scale { - let mut uncond_tokens = tokenizer.encode(uncond_prompt, true).map_err(E::msg)?.get_ids().to_vec(); - if uncond_tokens.len() > sd_config.clip.max_position_embeddings { - anyhow::bail!("the negative prompt is too long, {} > max-tokens ({})", uncond_tokens.len(), sd_config.clip.max_position_embeddings) - } - while uncond_tokens.len() < sd_config.clip.max_position_embeddings { - uncond_tokens.push(pad_id) - } - - let uncond_tokens = Tensor::new(uncond_tokens.as_slice(), device)?.unsqueeze(0)?; - let uncond_embeddings = text_model.forward(&uncond_tokens)?; - - Tensor::cat(&[uncond_embeddings, text_embeddings], 0)?.to_dtype(dtype)? - } else { - text_embeddings.to_dtype(dtype)? - }; - Ok(text_embeddings) -} diff --git a/crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs index 3d0a889cf9..778c43e1e7 100644 --- a/crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs +++ b/crates/desktop/ml_models/src/ml/stable_diffusion/lcm_pipeline.rs @@ -82,6 +82,8 @@ pub fn lcm_pipeline(args: Args<'_, P1, P2, P3, P4>) -> Result(args: Args<'_, P1, P2, P3, P4>) -> Result