From a859bd427bf4d5eda205afd2e38b07305f8377e5 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 6 Mar 2025 08:42:44 -0500 Subject: [PATCH] download clip models (not using them yet) --- .../src/ml/downloads/download_all_models.rs | 6 +++-- .../src-tauri/src/ml/model_registry.rs | 22 ++++++++++++++----- .../infer_clip_text_embeddings.rs | 2 +- .../stable_diffusion_pipeline.rs | 15 ++++--------- 4 files changed, 26 insertions(+), 19 deletions(-) diff --git a/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs b/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs index 6699daf7a5..953e084345 100644 --- a/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs +++ b/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs @@ -2,8 +2,10 @@ use crate::ml::downloads::download_model_registry_file::download_model_registry_ use crate::ml::model_registry::ModelRegistry; pub async fn download_all_models() -> anyhow::Result<()> { - download_model_registry_file(ModelRegistry::Clip).await?; + download_model_registry_file(ModelRegistry::ClipJson).await?; download_model_registry_file(ModelRegistry::SdxlTurboUnet).await?; - + download_model_registry_file(ModelRegistry::SdxlTurboVae).await?; + download_model_registry_file(ModelRegistry::SdxlTurboClipEncoder).await?; + download_model_registry_file(ModelRegistry::SdxlTurboClipEncoder2).await?; Ok(()) } diff --git a/crates/desktop/src-tauri/src/ml/model_registry.rs b/crates/desktop/src-tauri/src/ml/model_registry.rs index 144431f4fd..e3a4e933b7 100644 --- a/crates/desktop/src-tauri/src/ml/model_registry.rs +++ b/crates/desktop/src-tauri/src/ml/model_registry.rs @@ -1,29 +1,41 @@ #[derive(Copy, Clone, Debug)] pub enum ModelRegistry { - Clip, + ClipJson, SdxlTurboUnet, + SdxlTurboVae, + SdxlTurboClipEncoder, + SdxlTurboClipEncoder2, } impl ModelRegistry { pub fn get_hf_id(&self) -> Option<&'static str> { match self { - Self::Clip => Some("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"), - Self::SdxlTurboUnet => Some("stabilityai/sdxl-turbo"), + Self::ClipJson => Some("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"), + Self::SdxlTurboUnet + | Self::SdxlTurboVae + | Self::SdxlTurboClipEncoder + | Self::SdxlTurboClipEncoder2 => Some("stabilityai/sdxl-turbo"), } } pub fn get_download_url(&self) -> &'static str { match self { - Self::Clip => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/tokenizer.json", + Self::ClipJson => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/tokenizer.json", Self::SdxlTurboUnet => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.unet.safetensors", + Self::SdxlTurboVae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/diffusion_pytorch_model.vae.safetensors", + Self::SdxlTurboClipEncoder => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder.safetensors", + Self::SdxlTurboClipEncoder2 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/clip_text_encoder_2.safetensors", } } pub fn get_filename(&self) -> &'static str { match self { - Self::Clip => "clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json", + Self::ClipJson => "clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json", Self::SdxlTurboUnet => "diffusion_pytorch_model.unet.safetensors", + Self::SdxlTurboVae => "diffusion_pytorch_model.vae.safetensors", + Self::SdxlTurboClipEncoder => "clip_text_encoder.safetensors", + Self::SdxlTurboClipEncoder2 => "clip_text_encoder_2.safetensors", } } } diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs b/crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs index 80cade0ff3..11a56137e1 100644 --- a/crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs +++ b/crates/desktop/src-tauri/src/ml/stable_diffusion/infer_clip_text_embeddings.rs @@ -87,7 +87,7 @@ fn do_infer_clip_text_embeddings( info!("Loading clip from download"); //let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?; - let tokenizer = PathBuf::from(ModelRegistry::Clip.get_filename()); + let tokenizer = PathBuf::from(ModelRegistry::ClipJson.get_filename()); info!("Tokenizer path: {:?}", tokenizer); diff --git a/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs b/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs index d5bb1955d5..60e3361078 100644 --- a/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs +++ b/crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs @@ -123,25 +123,18 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { None => { info!("No vae found in cache; loading..."); - let repo = configs.sd_version.repo(); - - println!("Building VAE model from : {:?} ... (3)", repo); - - let vae_file = configs.hf_api.model(repo.to_string()) - .get("vae/diffusion_pytorch_model.safetensors")?; - - println!("Building VAE model from file {:?}...", &vae_file); - let mut notify_download_complete = false; //if !vae_file.exists() { if true { notify_download_complete = true; app.emit("notification", NotificationEvent::ModelDownloadStarted { - model_name: repo, + model_name: "sdxl-turbo-vae", model_type: ModelType::Vae, })?; } + let vae_file = ModelRegistry::SdxlTurboVae.get_filename(); + let vae = configs .sd_config .build_vae(vae_file, &configs.device, configs.dtype)?; @@ -152,7 +145,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result { if notify_download_complete { app.emit("notification", NotificationEvent::ModelDownloadComplete { - model_name: repo, + model_name: "sdxl-turbo-vae", model_type: ModelType::Vae, })?; }