download clip models (not using them yet)

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(())
 }

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",
     }
   }
 }

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);
   

crates/desktop/src-tauri/src/ml/stable_diffusion/stable_diffusion_pipeline.rs

@@ -123,25 +123,18 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
         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<RgbImage> {
             
             if notify_download_complete {
                 app.emit("notification", NotificationEvent::ModelDownloadComplete {
-                    model_name: repo,
+                    model_name: "sdxl-turbo-vae",
                     model_type: ModelType::Vae,
                 })?;
             }