don't download sdxl weights

Cargo.lock

@@ -663,6 +663,7 @@ dependencies = [
  "candle-core",
  "candle-nn",
  "candle-transformers",
+ "const_format",
  "cudarc",
  "directories 6.0.0",
  "dirs 6.0.0",
@@ -2300,6 +2301,26 @@ dependencies = [
  "tiny-keccak",
 ]
 
+[[package]]
+name = "const_format"
+version = "0.2.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "126f97965c8ad46d6d9163268ff28432e8f6a1196a55578867832e3049df63dd"
+dependencies = [
+ "const_format_proc_macros",
+]
+
+[[package]]
+name = "const_format_proc_macros"
+version = "0.2.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d57c2eccfb16dbac1f4e61e206105db5820c9d26c3c472bc17c774259ef7744"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-xid",
+]
+
 [[package]]
 name = "constant_time_eq"
 version = "0.1.5"
@@ -12884,6 +12905,12 @@ version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd"
 
+[[package]]
+name = "unicode-xid"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
+
 [[package]]
 name = "unicode_categories"
 version = "0.1.1"

crates/desktop/src-tauri/Cargo.toml

@@ -34,8 +34,9 @@ bytes = "1.10"
 candle-core = { path = "../../vendor/candle/candle-core", features = ["cuda"] }
 candle-nn = { path = "../../vendor/candle/candle-nn", features = ["cuda"] }
 candle-transformers = { path = "../../vendor/candle/candle-transformers", features = ["cuda"] }
-dirs = "6.0.0"
+const_format = "0.2.34"
 directories = "6.0.0"
+dirs = "6.0.0"
 futures = "0.3"
 hf-hub = "0.3.2"
 image = "0.25.5"

crates/desktop/src-tauri/src/endpoints/download_models.rs

@@ -1,16 +1,12 @@
-use crate::ml::downloads::download_all_models::download_all_models;
 use crate::state::app_dir::AppDataRoot;
 use tauri::State;
 
 #[tauri::command]
 pub async fn download_models(
-  app_data_root: State<'_, AppDataRoot>,
+  _app_data_root: State<'_, AppDataRoot>,
 ) -> Result<String, String> {
-  
-  download_all_models(&app_data_root).await
-    .map_err(|err| {
-      "error downloading models".to_string()
-    })?;
-  
+
+  // NB: This endpoint has been removed for now as we now download models in the app lifecycle.
+
   Ok(("downloaded".to_string()))
 }

crates/desktop/src-tauri/src/endpoints/image_endpoint.rs

@@ -1,8 +1,8 @@
-use crate::ml::downloads::download_all_models::download_all_models;
 use crate::ml::model_cache::ModelCache;
 use crate::ml::prompt_cache::PromptCache;
 use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
 use crate::state::app_config::AppConfig;
+use crate::state::app_dir::AppDataRoot;
 use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
 use bytes::BytesMut;
 use image::imageops::FilterType;
@@ -11,7 +11,6 @@ use log::{error, info};
 use std::io::Cursor;
 use std::path::PathBuf;
 use tauri::{AppHandle, State};
-use crate::state::app_dir::AppDataRoot;
 
 const PROMPT_FILENAME : &str = "prompt.txt";
 const PROMPT: &str = "green goblin, blood dripping from lips, detailed, photorealistic, 8k";
@@ -45,13 +44,6 @@ pub async fn infer_image(
   let image = hydrate_base64_image(image)
     .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
 
-  // TODO(bt): Move this to another endpoint
-  download_all_models(&app_data_root).await
-    .map_err(|err| {
-      error!("couldn't download models: {:?}", err);
-      "Couldn't download models".to_string()
-    })?;
-
   let result = do_infer_image(
     &prompt, 
     image, 

crates/desktop/src-tauri/src/events/notification_event.rs

@@ -6,8 +6,8 @@ use serde::Serialize;
 pub enum NotificationModelType {
   Unet,
   Vae,
-  Json,
-  ClipEncoder,
+  TextEncoder,
+  TextTokenizer,
 }
 
 #[derive(Clone, Serialize)]

crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs

@@ -1,15 +0,0 @@
-use crate::ml::downloads::download_model_registry_file::download_model_registry_file;
-use crate::ml::model_type::ModelType;
-use crate::state::app_dir::AppDataRoot;
-
-pub async fn download_all_models(app_data_root: &AppDataRoot) -> anyhow::Result<()> {
-  let weights_dir = app_data_root.weights_dir();
-  
-  //download_model_registry_file(ModelType::ClipJson, weights_dir).await?;
-  //download_model_registry_file(ModelType::SdxlTurboUnet, weights_dir).await?;
-  //download_model_registry_file(ModelType::SdxlTurboVae, weights_dir).await?;
-  //download_model_registry_file(ModelType::SdxlTurboClipEncoder, weights_dir).await?;
-  //download_model_registry_file(ModelType::SdxlTurboClipEncoder2, weights_dir).await?;
-  
-  Ok(())
-}

crates/desktop/src-tauri/src/ml/downloads/download_model_registry_file.rs

@@ -1,30 +0,0 @@
-use crate::ml::model_type::ModelType;
-use crate::state::app_dir::AppWeightsDir;
-use log::info;
-use std::fs::File;
-use std::path::{Path, PathBuf};
-
-pub async fn download_model_registry_file(
-  model_type: ModelType,
-  weights_dir: &AppWeightsDir,
-) -> anyhow::Result<PathBuf> {
-  
-  let path = weights_dir.model_path(&model_type);
-
-  if path.exists() {
-    return Ok(path);
-  }
-
-  info!("downloading model: {:?} to {:?}", model_type, path);
-
-  let body = reqwest::get(model_type.get_download_url())
-    .await?
-    .bytes()
-    .await?;
-  
-  let _file = File::create(&path)?;
-  
-  std::fs::write(&path, body.as_ref())?;
-  
-  Ok(path)
-}

crates/desktop/src-tauri/src/ml/downloads/mod.rs

@@ -1,2 +0,0 @@
-pub mod download_model_registry_file;
-pub mod download_all_models;

crates/desktop/src-tauri/src/ml/mod.rs

@@ -4,5 +4,4 @@ pub mod model_file;
 pub mod models;
 pub mod prompt_cache;
 pub mod stable_diffusion;
-pub mod model_type;
-pub mod downloads;
+pub mod weights_registry;

crates/desktop/src-tauri/src/ml/model_type.rs

@@ -1,86 +0,0 @@
-use crate::events::notification_event::NotificationModelType;
-use crate::state::app_dir::AppWeightsDir;
-use std::path::PathBuf;
-
-#[derive(Copy, Clone, Debug)]
-pub enum ModelType {
-  ClipJson,
-  SdxlTurboUnet,
-  SdxlTurboVae,
-  SdxlTurboClipEncoder,
-  SdxlTurboClipEncoder2,
-  SimianLuoLcmDreamshaperV7Unet,
-  LykonDreamshaper7Vae,
-  LykonDreamshaper7TextEncoderFp16,
-}
-
-impl ModelType {
-  pub fn get_name(&self) -> &'static str {
-    match self {
-      Self::ClipJson => "CLIP (JSON)",
-      Self::SdxlTurboUnet => "SDXL Turbo UNET",
-      Self::SdxlTurboVae => "SDXL Turbo VAE",
-      Self::SdxlTurboClipEncoder => "Clip Encoder #1",
-      Self::SdxlTurboClipEncoder2 => "Clip Encoder #2",
-      Self::SimianLuoLcmDreamshaperV7Unet => "SimianLuo LCM Dreamshaper v7 UNET",
-      Self::LykonDreamshaper7Vae => "Lykon Dreamshaper 7 VAE",
-      Self::LykonDreamshaper7TextEncoderFp16 => "Lykon Dreamshaper Text Encoder (fp16)",
-    }
-  }
-
-  pub fn get_notification_type(&self) -> NotificationModelType {
-    match self {
-      Self::ClipJson => NotificationModelType::Json,
-      Self::SdxlTurboUnet => NotificationModelType::Unet,
-      Self::SdxlTurboVae => NotificationModelType::Vae,
-      Self::SdxlTurboClipEncoder => NotificationModelType::ClipEncoder,
-      Self::SdxlTurboClipEncoder2 => NotificationModelType::ClipEncoder,
-      Self::SimianLuoLcmDreamshaperV7Unet => NotificationModelType::Unet,
-      Self::LykonDreamshaper7Vae => NotificationModelType::Vae,
-      Self::LykonDreamshaper7TextEncoderFp16 => NotificationModelType::ClipEncoder,
-    }
-  }
-
-  pub fn get_hf_id(&self) -> Option<&'static str> {
-    match self {
-      Self::ClipJson => Some("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"),
-      Self::SdxlTurboUnet 
-      | Self::SdxlTurboVae
-      | Self::SdxlTurboClipEncoder
-      | Self::SdxlTurboClipEncoder2 => Some("stabilityai/sdxl-turbo"),
-      Self::SimianLuoLcmDreamshaperV7Unet => Some("SimianLuo/LCM_Dreamshaper_v7"),
-      Self::LykonDreamshaper7Vae
-      | Self::LykonDreamshaper7TextEncoderFp16 => Some("Lykon/dreamshaper-7"),
-    }
-  }
-
-  pub fn get_download_url(&self) -> &'static str {
-    match self {
-      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",
-      Self::SimianLuoLcmDreamshaperV7Unet => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/simianluo_lcm_dreamshaper_v7_unet.safetensors",
-      Self::LykonDreamshaper7Vae => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_vae.safetensors",
-      Self::LykonDreamshaper7TextEncoderFp16 => "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_text_encoder.fp16.safetensors",
-    }
-  }
-
-  pub fn get_filename(&self) -> &'static str {
-    match self {
-      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",
-      Self::SimianLuoLcmDreamshaperV7Unet => "simianluo_lcm_dreamshaper_v7_unet.safetensors",
-      Self::LykonDreamshaper7Vae => "lykon_dreamshaper_7_vae.safetensors",
-      Self::LykonDreamshaper7TextEncoderFp16 => "lykon_dreamshaper_7_text_encoder.fp16.safetensors",
-    }
-  }
-  
-  pub fn get_path(&self, weights_dir: &AppWeightsDir) -> PathBuf {
-    weights_dir.model_path(self)
-  }
-}

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

@@ -3,16 +3,16 @@ extern crate accelerate_src;
 #[cfg(feature = "mkl")]
 extern crate intel_mkl_src;
 
-use std::path::PathBuf;
 use candle_transformers::models::stable_diffusion;
+use std::path::PathBuf;
 
 use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
+use crate::ml::weights_registry::weights::{CLIP_JSON, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2};
+use crate::state::app_dir::AppWeightsDir;
 use anyhow::Error as E;
 use candle_core::{DType, Device, Module, Tensor, D};
 use log::info;
 use tokenizers::Tokenizer;
-use crate::ml::model_type::ModelType;
-use crate::state::app_dir::AppWeightsDir;
 
 #[allow(clippy::too_many_arguments)]
 pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option<String>, clip_weights: Option<String>, clip2_weights: Option<String>, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, weights_dir: &AppWeightsDir) -> anyhow::Result<Tensor> {
@@ -53,17 +53,7 @@ pub fn infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer:
 fn do_infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option<String>, clip_weights: Option<String>, clip2_weights: Option<String>, 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<Tensor> {
   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_file = if first {
-  //  info!("ModelFile::Tokenizer");
-  //  ModelFile::Tokenizer
-  //} else {
-  //  info!("ModelFile::Tokenizer2");
-  //  ModelFile::Tokenizer2
-  //};
-
-  //let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
-
-  let tokenizer = weights_dir.model_path(&ModelType::ClipJson);
+  let tokenizer = weights_dir.weight_path(&CLIP_JSON);
 
   info!("Loading Clip Tokenizer path: {:?}", tokenizer);
 
@@ -85,13 +75,10 @@ fn do_infer_clip_text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: O
   let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?;
 
   println!("Building the Clip transformer.");
-  let clip_weights_file = if first { ModelFile::Clip } else { ModelFile::Clip2 };
   let clip_weights = if first {
-    //clip_weights_file.get(clip_weights, sd_version, use_f16)?
-    weights_dir.model_path(&ModelType::SdxlTurboClipEncoder)
+    weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER)
   } else {
-    //clip_weights_file.get(clip2_weights, sd_version, use_f16)?
-    weights_dir.model_path(&ModelType::SdxlTurboClipEncoder2)
+    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)?;

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

@@ -6,11 +6,11 @@ use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
 use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
 use crate::ml::model_cache::ModelCache;
 use crate::ml::model_file::{ModelFile, StableDiffusionVersion};
-use crate::ml::model_type::ModelType;
 use crate::ml::models::unet_model::UNetModel;
 use crate::ml::prompt_cache::PromptCache;
 use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
 use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings;
+use crate::ml::weights_registry::weights::{LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
 use crate::state::app_config::AppConfig;
 use crate::state::app_dir::AppDataRoot;
 use anyhow::{anyhow, Error as E, Result};
@@ -48,22 +48,12 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
   let img2img_strength = args.i2i_strength.unwrap_or(75) as f64 / 100.0;
   let use_f16 = true;
 
-  //println!("Starting image generation with the following configuration:");
-  //println!("  Model: {:?}", configs.sd_version);
-  //println!("  Prompt: {}", prompt);
-  //println!("  Steps: {}", configs.scheduler_steps);
-  //println!("  Device: {:?}", configs.device);
-
-  //println!("Model dimensions: {}x{}", configs.sd_config.width, configs.sd_config.height);
-
   // Use LCM Scheduler instead of Euler Ancestral for better speed and quality
   let mut scheduler = LCMScheduler::new(configs.scheduler_steps, img2img_strength, candle_transformers::models::stable_diffusion::lcm::LCMSchedulerConfig::default())?;
 
   let seed = configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
   configs.device.set_seed(seed)?;
 
-  //info!("Using seed: {}", seed);
-
   let guidance_scale = match cfg_scale {
     Some(guidance_scale) => guidance_scale,
     None => match configs.sd_version {
@@ -76,26 +66,14 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
   // let use_guide_scale = guidance_scale > 1.0;
   let use_guide_scale = false;
 
-  //info!("Using guide scale: {}", use_guide_scale);
-
-  //info!("Checking if prompt is cached");
   let maybe_cached = prompt_cache.get_copy(&prompt)?;
 
   let mut text_embeddings = if let Some(tensor) = maybe_cached {
     tensor
   } else {
     info!("Loading clip weights...");
-    //println!("Loading clip weights for {:?} with model {:?}", configs.sd_version, configs.sd_version.repo());
-    //let api = configs.hf_api.clone();
 
-    //let clip_path = api.model(configs.sd_version.repo().to_string()).get(if use_f16 { "text_encoder/model.fp16.safetensors" } else { "text_encoder/model.safetensors" })?;
-
-    //println!(">>>> CLIP REPO = {:?}", configs.sd_version.repo());
-    //println!(">>>> CLIP PATH = {:?}", clip_path);
-    //println!(">>>> USE_FP16 = {:?}", use_f16);
-    //let clip_weights = Some(clip_path.to_string_lossy().to_string());
-
-    let clip_weights = weights_dir.model_path(&ModelType::LykonDreamshaper7TextEncoderFp16);
+    let clip_weights = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16);
     let clip_weights = Some(clip_weights.to_string_lossy().to_string());
 
     info!("Prompt is NOT cached! Calculating embedding...");
@@ -132,12 +110,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
   let vae = match maybe_vae {
     Some(vae) => vae,
     None => {
-      //let repo = configs.sd_version.repo();
-      //println!(">>>>> VAE REPO = {:?}", repo);
-      //println!("Building VAE model from : {:?} ... (3)", repo);
-      //let vae_file = configs.hf_api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
-
-      let vae_file = weights_dir.model_path(&ModelType::LykonDreamshaper7Vae);
+      let vae_file = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE);
 
       println!("Building VAE model from file {:?}...", &vae_file);
 
@@ -158,7 +131,7 @@ pub fn lcm_pipeline(args: Args<'_>) -> Result<RgbImage> {
     None => {
       info!("No unet found in cache; loading...");
 
-      let unet_weights = weights_dir.model_path(&ModelType::SimianLuoLcmDreamshaperV7Unet);
+      let unet_weights = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
 
       let in_channels = match configs.sd_version {
         StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => 9,

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

@@ -1,16 +1,17 @@
 use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
 
-use crate::events::notification_event::{NotificationModelType, NotificationEvent};
+use crate::events::notification_event::{NotificationEvent, NotificationModelType};
 use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
 use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
 use crate::ml::model_cache::ModelCache;
 use crate::ml::model_file::StableDiffusionVersion;
-use crate::ml::model_type::ModelType;
 use crate::ml::models::unet_model::UNetModel;
 use crate::ml::prompt_cache::PromptCache;
 use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
 use crate::ml::stable_diffusion::infer_clip_text_embeddings::infer_clip_text_embeddings;
+use crate::ml::weights_registry::weights::{SDXL_TURBO_UNET, SDXL_TURBO_VAE};
 use crate::state::app_config::AppConfig;
+use crate::state::app_dir::AppDataRoot;
 use anyhow::{anyhow, Error as E, Result};
 use candle_core::{DType, IndexOp, Tensor, D};
 use candle_transformers::models::stable_diffusion::vae::{AutoEncoderKL, DiagonalGaussianDistribution};
@@ -18,7 +19,6 @@ use image::DynamicImage;
 use log::info;
 use rand::Rng;
 use tauri::{AppHandle, Emitter};
-use crate::state::app_dir::AppDataRoot;
 
 pub struct Args<'a> {
     pub image: &'a DynamicImage,
@@ -129,17 +129,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
         None => {
             info!("No vae found in cache; loading...");
             
-            let mut notify_download_complete = false;
-            //if !vae_file.exists() {
-            if true {
-                notify_download_complete = true;
-                app.emit("notification", NotificationEvent::ModelDownloadStarted {
-                    model_name: "sdxl-turbo-vae",
-                    model_type: NotificationModelType::Vae,
-                })?;
-            }
-            
-            let vae_file = weights_dir.model_path(&ModelType::SdxlTurboVae);
+            let vae_file = weights_dir.weight_path(&SDXL_TURBO_VAE);
             
             let vae = configs
               .sd_config
@@ -148,14 +138,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
             let vae = Arc::new(vae);
             
             model_cache.set_vae(vae.clone())?;
-            
-            if notify_download_complete {
-                app.emit("notification", NotificationEvent::ModelDownloadComplete {
-                    model_name: "sdxl-turbo-vae",
-                    model_type: NotificationModelType::Vae,
-                })?;
-            }
-            
+
             vae
         }
     };
@@ -167,17 +150,7 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
         None => {
             info!("No unet found in cache; loading...");
 
-            let mut notify_download_complete = false;
-            //if !unet_file.exists() {
-            if true {
-                notify_download_complete = true;
-                app.emit("notification", NotificationEvent::ModelDownloadStarted {
-                    model_name: "sdxl-turbo-unet",
-                    model_type: NotificationModelType::Unet,
-                })?;
-            }
-            
-            let unet_file = weights_dir.model_path(&ModelType::SdxlTurboUnet);
+            let unet_file = weights_dir.weight_path(&SDXL_TURBO_UNET);
 
             let unet = UNetModel::new(&configs.sd_config, unet_file, &configs.device, configs.dtype)
               .map_err(|err| anyhow!("error initializing unet model: {:?}", err))?;
@@ -186,13 +159,6 @@ pub fn stable_diffusion_pipeline(args: Args<'_>) -> Result<RgbImage> {
 
             model_cache.set_unet(unet.clone())?;
 
-            if notify_download_complete {
-                app.emit("notification", NotificationEvent::ModelDownloadComplete {
-                    model_name: "sdxl-turbo-unet",
-                    model_type: NotificationModelType::Unet,
-                })?;
-            }
-
             unet
         }
     };

crates/desktop/src-tauri/src/ml/weights_registry/mod.rs

@@ -0,0 +1,3 @@
+mod weight_descriptor_builder;
+pub mod weight_descriptor;
+pub mod weights;

crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs

@@ -0,0 +1,43 @@
+
+// TODO(bt,2025-03-13): When it's time for dynamic, non-const allocation, refer to
+//  http::header::name::HeaderName: https://github.com/hyperium/http/blob/master/src/header/name.rs
+//  as a pattern that will support both compile-time constants and dynamically allocated String-based variants.
+
+#[derive(Clone, Copy)]
+pub enum WeightFunction {
+  TextEncoder,
+  TextTokenizer,
+  Unet,
+  Vae,
+}
+
+pub struct WeightDescriptor {
+  /// Name of the weight, eg. "Lykon Dreamshaper 7"
+  pub name: &'static str,
+
+  /// Name of the file on the filesystem, eg. lykon_dreamshaper_7_text_encoder.fp16.safetensors
+  pub filename: &'static str,
+
+  /// URL of the file in the R2 bucket,
+  /// eg. https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_text_encoder.fp16.safetensors
+  pub r2_download_url: &'static str,
+
+  pub function: WeightFunction,
+}
+
+impl WeightFunction {
+  pub fn as_descriptive_str(&self) -> &'static str {
+    match self {
+      WeightFunction::TextEncoder => "text encoder",
+      WeightFunction::TextTokenizer => "text tokenizer",
+      WeightFunction::Unet => "UNET",
+      WeightFunction::Vae => "VAE",
+    }
+  }
+}
+
+impl WeightDescriptor {
+  pub fn to_descriptive_name(&self) -> String {
+    format!("{} ({})", self.name, self.function.as_descriptive_str())
+  }
+}

crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor_builder.rs

@@ -0,0 +1,17 @@
+pub (super) const R2_BUCKET_URL : &str = "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/";
+
+// TODO(bt,2025-03-13): Try to use keyword args from the call site.
+// NB: We're using a macro because we can't use a builder pattern and then concat string literals
+//  and variables (even if they're known to be &'static) and keep allocations static.
+macro_rules! weight {
+    ($name: literal, $file_name:literal, $weight_function:expr) => {
+      $crate::ml::weights_registry::weight_descriptor::WeightDescriptor {
+        name: $name,
+        filename: $file_name,
+        r2_download_url: const_format::concatcp!(crate::ml::weights_registry::weight_descriptor_builder::R2_BUCKET_URL, $file_name),
+        function: $weight_function,
+      }
+    };
+}
+
+pub(crate) use weight;

crates/desktop/src-tauri/src/ml/weights_registry/weights.rs

@@ -0,0 +1,57 @@
+use crate::ml::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction};
+use crate::ml::weights_registry::weight_descriptor_builder::weight;
+
+pub const CLIP_JSON : WeightDescriptor = weight!(
+  "CLIP",
+  "clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json",
+  WeightFunction::TextTokenizer
+);
+
+/// For LCM-based image generation
+pub const LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16 : WeightDescriptor = weight!(
+  "Lykon Dreamshaper 7",
+  "lykon_dreamshaper_7_text_encoder.fp16.safetensors",
+  WeightFunction::TextEncoder
+);
+
+/// For LCM-based image generation
+pub const LYKON_DEAMSHAPER_7_VAE : WeightDescriptor = weight!(
+  "Lykon Dreamshaper 7",
+  "lykon_dreamshaper_7_vae.safetensors",
+  WeightFunction::Vae
+);
+
+/// For LCM-based image generation
+pub const SIMIANLUO_LCM_DREAMSHAPER_V7_UNET : WeightDescriptor = weight!(
+  "SimianLuo LCM Dreamshaper v7",
+  "simianluo_lcm_dreamshaper_v7_unet.safetensors",
+  WeightFunction::Unet
+);
+
+#[deprecated]
+pub const SDXL_TURBO_CLIP_TEXT_ENCODER : WeightDescriptor = weight!(
+  "SDXL Turbo CLIP",
+  "clip_text_encoder.safetensors",
+  WeightFunction::TextEncoder
+);
+
+#[deprecated]
+pub const SDXL_TURBO_CLIP_TEXT_ENCODER_2 : WeightDescriptor = weight!(
+  "SDXL Turbo CLIP #2",
+  "clip_text_encoder_2.safetensors",
+  WeightFunction::TextEncoder
+);
+
+#[deprecated]
+pub const SDXL_TURBO_VAE : WeightDescriptor = weight!(
+  "SDXL Turbo",
+  "diffusion_pytorch_model.vae.safetensors",
+  WeightFunction::Vae
+);
+
+#[deprecated]
+pub const SDXL_TURBO_UNET : WeightDescriptor = weight!(
+  "SDXL Turbo",
+  "diffusion_pytorch_model.unet.safetensors",
+  WeightFunction::Unet
+);

crates/desktop/src-tauri/src/state/app_dir.rs

@@ -1,4 +1,4 @@
-use crate::ml::model_type::ModelType;
+use crate::ml::weights_registry::weight_descriptor::WeightDescriptor;
 use crate::state::expanduser::expanduser;
 use crate::state::os_platform::OsPlatform;
 use anyhow::anyhow;
@@ -6,6 +6,7 @@ use directories::UserDirs;
 use std::path::{Path, PathBuf};
 use tempdir::TempDir;
 use tempfile::{Builder, NamedTempFile};
+
 const DEFAULT_DATA_DIR : &str = "artcraft";
 const ASSETS_SUBDIRECTORY : &str = "assets";
 const WEIGHTS_SUBDIRECTORY : &str = "weights";
@@ -139,9 +140,9 @@ impl AppWeightsDir {
   pub fn path(&self) -> &Path {
     &self.path
   }
-  
-  pub fn model_path(&self, model_registry: &ModelType) -> PathBuf {
-    self.path.join(model_registry.get_filename())
+
+  pub fn weight_path(&self, descriptor: &WeightDescriptor) -> PathBuf {
+    self.path.join(descriptor.filename)
   }
 }
 

crates/desktop/src-tauri/src/threads/downloader_thread.rs

@@ -1,5 +1,4 @@
 use crate::events::notification_event::{NotificationEvent, NotificationModelType};
-use crate::ml::model_type::ModelType;
 use crate::state::app_dir::AppDataRoot;
 use crate::transfer::download::{download_async, ProgressUpdate};
 use futures::stream::FuturesUnordered;
@@ -12,6 +11,8 @@ use std::sync::{Arc, Mutex};
 use tauri::{AppHandle, Emitter};
 use tempfile::NamedTempFile;
 use tokio::task::JoinHandle;
+use crate::ml::weights_registry::weight_descriptor::WeightFunction;
+use crate::ml::weights_registry::weights::{CLIP_JSON, LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16, LYKON_DEAMSHAPER_7_VAE, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2, SDXL_TURBO_UNET, SDXL_TURBO_VAE, SIMIANLUO_LCM_DREAMSHAPER_V7_UNET};
 
 const NOTIFICATION_CHANNEL_NAME : &str = "notification";
 const MAX_FILES : usize = 8;
@@ -39,22 +40,21 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !
 
   let mut download_queue = VecDeque::new();
 
-  // TODO: Automatic enqueue of known important models + enqueue new models on-demand
-  download_queue.push_back(ModelType::ClipJson);
-  //download_queue.push_back(ModelType::SdxlTurboUnet);
-  //download_queue.push_back(ModelType::SdxlTurboVae);
-  download_queue.push_back(ModelType::SdxlTurboClipEncoder); // TODO(bt): Why is this still needed?
-  //download_queue.push_back(ModelType::SdxlTurboClipEncoder2);
-  download_queue.push_back(ModelType::SimianLuoLcmDreamshaperV7Unet);
-  download_queue.push_back(ModelType::LykonDreamshaper7Vae);
-  download_queue.push_back(ModelType::LykonDreamshaper7TextEncoderFp16);
+  download_queue.push_back(CLIP_JSON);
+  download_queue.push_back(LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16);
+  download_queue.push_back(LYKON_DEAMSHAPER_7_VAE);
+  download_queue.push_back(SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
+  download_queue.push_back(SDXL_TURBO_CLIP_TEXT_ENCODER); // TODO(bt): Why is this still needed?
+  //download_queue.push_back(SDXL_TURBO_CLIP_TEXT_ENCODER_2); // TODO(bt): REMOVE
+  //download_queue.push_back(SDXL_TURBO_VAE); // TODO(bt): REMOVE
+  //download_queue.push_back(SDXL_TURBO_UNET); // TODO(bt): REMOVE
 
   let mut handles = FuturesUnordered::new();
 
   while let Some(model) = download_queue.pop_front() {
-    let url = model.get_download_url().to_string();
+    let url = model.r2_download_url.to_string();
 
-    let filename = app_data_root.weights_dir().model_path(&model);
+    let filename = app_data_root.weights_dir().weight_path(&model);
 
     if filename.exists() {
       info!("File already exists: {:?}", &filename);
@@ -75,9 +75,16 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !
       }
     };
 
+    let model_type = match model.function {
+      WeightFunction::TextEncoder => NotificationModelType::TextEncoder,
+      WeightFunction::TextTokenizer => NotificationModelType::TextTokenizer,
+      WeightFunction::Unet => NotificationModelType::Unet,
+      WeightFunction::Vae => NotificationModelType::Vae,
+    };
+
     let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadStarted {
-      model_name: model.get_name(),
-      model_type: model.get_notification_type(),
+      model_name: model.name,
+      model_type,
     });
 
     if let Err(err) = result {
@@ -102,8 +109,8 @@ pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> !
         lock.push(ModelDownloadStatusChannel {
           receiver: rx,
           final_file_path: filename.clone(),
-          model_name: model.get_name().to_string(),
-          model_type: model.get_notification_type(),
+          model_name: model.to_descriptive_name(),
+          model_type,
         })
       }
     }