cache SD model weights in a cache directory

crates/service/job/inference_job/config/inference-job.development.env

@@ -9,6 +9,9 @@ HTTP_BIND_ADDRESS=0.0.0.0:11223
 # NB: Disable if not testing Redis
 REDIS_FOR_KEEPALIVE_URL="redis://localhost/0"
 
+# General model weights cache directory
+MODEL_WEIGHTS_CACHE_DIR="/tmp/model_weights_cache"
+
 # ================================= #
 #                                   #
 #            Vall-E-X               #

crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job_inference.rs

@@ -97,14 +97,13 @@ pub async fn process_job_inference(
       .new_tempdir(&work_temp_dir)
       .map_err(|e| ProcessSingleJobError::from_io_error(e))?;
 
-  let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint.safetensors");
+  //let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint.safetensors");
   let mut lora_path = work_temp_dir.path().join("lora.safetensors");
   //let vae_path = work_temp_dir.path().join("vae.safetensors");
   let vae_path = work_temp_dir.path().join("vae.pt"); // TODO: Should this be `.safetensors` or `.pt`?
   let output_path = work_temp_dir.path().join("output");
 
   info!("Paths to download to:");
-  info!("sd_checkpoint_path: {:?}", sd_checkpoint_path);
   info!("lora_path: {:?}", lora_path);
   info!("vae_path: {:?}", vae_path);
   info!("output_path: {:?}", output_path);
@@ -114,91 +113,107 @@ pub async fn process_job_inference(
   // // Unpack loRA and Checkpoint
   // // run inference by downloading from google drive.
   let lora_token = sd_args.maybe_lora_model_token;
-  let weight_token = sd_args.maybe_sd_model_token.clone();
+///  let weight_token = sd_args.maybe_sd_model_token.clone();
 
-  let retrieved_sd_record = match weight_token {
-    Some(ref token) => {
-      let retrieved_sd_record = get_weight_by_token(
-        &token,
-        false,
-        &deps.db.mysql_pool
-      ).await?;
-      match retrieved_sd_record {
-        Some(record) => record,
-        None => {
-          return Err(
-            ProcessSingleJobError::from_anyhow_error(anyhow!("no record of model!"))
-          );
-        }
-      }
-    }
-    None => {
-      return Err(
-        ProcessSingleJobError::from_anyhow_error(anyhow!("no sd model token for job!"))
-      );
-    }
+  let sd_model_weight_token = match sd_args.maybe_sd_model_token {
+    None => return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no sd model token for job!"))),
+    Some(ref token) => token,
   };
+
+  let sd_model_weight = get_weight_by_token(
+    sd_model_weight_token,
+    false,
+    &deps.db.mysql_pool
+  ).await?;
+
+  let sd_model_weight = match sd_model_weight {
+    None => return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no record of sd model!"))),
+    Some(record) => record,
+  };
+
+  let sd_checkpoint_path = args.job_dependencies
+      .fs
+      .model_weights_cache_directory
+      .get_model_weight_from_cache_or_bucket(&sd_model_weight)
+      .await?;
+
+  info!("sd_checkpoint_path: {:?}", sd_checkpoint_path);
+
+//  let retrieved_sd_record = match sd_args.maybe_sd_model_token {
+//    None => return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no sd model token for job!"))),
+//    Some(ref token) => {
+//      let retrieved_sd_record = get_weight_by_token(
+//        &token,
+//        false,
+//        &deps.db.mysql_pool
+//      ).await?;
+//      match retrieved_sd_record {
+//        Some(record) => record,
+//        None => return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no record of model!"))),
+//      }
+//    }
+//  };
+
   // origin file name needs to be just the file name  /tmp/downloads_long_lived/temp_stable_diffusion_inference_32.8qJJljxWWZeD/output_0.png
   // ignore if no lora token
-  let retrieved_loRA_record = match lora_token {
-    Some(token) => {
-      let retrieved_loRA_record = get_weight_by_token(
-        &token,
-        false,
-        &deps.db.mysql_pool
-      ).await?;
 
-      Some(retrieved_loRA_record)
-    }
-    None => None,
-  };
+  let mut maybe_lora_record = None;
+
+  if let Some(token) = lora_token {
+    maybe_lora_record = get_weight_by_token(
+      &token,
+      false,
+      &deps.db.mysql_pool
+    ).await?;
+  }
 
   let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?;
 
-  // Details for SD checkpoint
-  let details = RemoteCloudBucketDetails::new(
-    retrieved_sd_record.public_bucket_hash.clone(),
-    retrieved_sd_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
-    retrieved_sd_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
-  );
-
-  remote_cloud_file_client.download_file(
-    details,
-    path_to_string(sd_checkpoint_path.clone())
-  ).await?;
+///  // Details for SD checkpoint
+///  let details = RemoteCloudBucketDetails::new(
+///    retrieved_sd_record.public_bucket_hash.clone(),
+///    retrieved_sd_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
+///    retrieved_sd_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
+///  );
+///
+///  remote_cloud_file_client.download_file(
+///    details,
+///    path_to_string(sd_checkpoint_path.clone())
+///  ).await?;
 
   let mut lora_name = String::from("");
   let mut lora_token = String::from("");
 
-  match retrieved_loRA_record {
-    Some(record) => {
-      match record {
-        Some(model_weight_record) => {
-          let lora_details = RemoteCloudBucketDetails::new(
-            model_weight_record.public_bucket_hash.clone(),
-            model_weight_record.maybe_public_bucket_prefix
-                .clone()
-                .unwrap_or_else(|| "".to_string()),
-            model_weight_record.maybe_public_bucket_extension
-                .clone()
-                .unwrap_or_else(|| "".to_string())
-          );
-
-          lora_name = model_weight_record.title;
-          lora_token = model_weight_record.token.to_string();
-          remote_cloud_file_client.download_file(
-            lora_details,
-            path_to_string(lora_path.clone())
-          ).await?;
-        }
-        None => {
-          lora_path.clear();
-        }
-      }
-    }
+  match maybe_lora_record {
     None => {
       lora_path.clear();
     }
+    Some(lora_record) => {
+/////      let lora_details = RemoteCloudBucketDetails::new(
+/////        lora_record.public_bucket_hash.clone(),
+/////        lora_record.maybe_public_bucket_prefix
+/////            .clone()
+/////            .unwrap_or_else(|| "".to_string()),
+/////        lora_record.maybe_public_bucket_extension
+/////            .clone()
+/////            .unwrap_or_else(|| "".to_string())
+/////      );
+/////
+/////
+/////      remote_cloud_file_client.download_file(
+/////        lora_details,
+/////        path_to_string(lora_path.clone())
+/////      ).await?;
+
+      lora_name = lora_record.title.to_string();
+      lora_token = lora_record.token.to_string();
+
+      lora_path = args.job_dependencies
+          .fs
+          .model_weights_cache_directory
+          .get_model_weight_from_cache_or_bucket(&lora_record)
+          .await?;
+    }
   }
 
   args.job_dependencies
@@ -211,33 +226,6 @@ pub async fn process_job_inference(
         ProcessSingleJobError::from_anyhow_error(anyhow!("could not download VAE: {:?}", err))
       })?;
 
-//    // VAE token for now
-//    let vae_token = String::from("REPLACE_ME");
-//    let model_weight_vae = ModelWeightToken(vae_token);
-//
-//    let vae_weight_record = get_weight_by_token(
-//        &model_weight_vae,
-//        false,
-//        &deps.db.mysql_pool
-//    ).await?;
-//
-//    let vae_weight_record = match vae_weight_record {
-//        Some(val) => val,
-//        None => {
-//            return Err(
-//                ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
-//            );
-//        }
-//    };
-//
-//    let vae_details = RemoteCloudBucketDetails::new(
-//        vae_weight_record.public_bucket_hash.clone(),
-//        vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
-//        vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
-//    );
-//
-//    remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?;
-
   let prompt = match sd_args.maybe_prompt {
     Some(val) => val,
     None => {
@@ -359,7 +347,7 @@ pub async fn process_job_inference(
       origin_category: MediaFileOriginCategory::Upload,
       origin_product_category: MediaFileOriginProductCategory::ImageGeneration,
       maybe_origin_model_type: Some(MediaFileOriginModelType::StableDiffusion15),
-      maybe_origin_model_token: weight_token.clone(),
+      maybe_origin_model_token: Some(sd_model_weight_token.clone()),
       maybe_origin_filename: Some(file_path),
       is_batch_generated: true,
       maybe_mime_type: Some(metadata.mimetype.as_ref()),

crates/service/job/inference_job/src/job_dependencies.rs

@@ -25,6 +25,7 @@ use mysql_queries::mediators::firehose_publisher::FirehosePublisher;
 
 use crate::job_specific_dependencies::JobSpecificDependencies;
 use crate::util::instrumentation::JobInstruments;
+use crate::util::model_weights_cache::model_weights_cache_directory::ModelWeightsCacheDirectory;
 use crate::util::scoped_execution::ScopedExecution;
 use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator;
 
@@ -164,6 +165,10 @@ pub struct FileSystemDetails {
   /// Directory to store long-term downloads (models)
   /// (In prod, typically model files from GCS / NFS PVC mount)
   pub semi_persistent_cache: SemiPersistentCacheDir,
+
+  /// The new cache scheme for model weights.
+  /// This should replace `semi_persistent_cache` for new users.
+  pub model_weights_cache_directory: ModelWeightsCacheDirectory,
 }
 
 // TODO: Move into the appropriate job-specific dependencies object.

crates/service/job/inference_job/src/main.rs

@@ -57,6 +57,7 @@ use crate::job::job_loop::main_loop::main_loop;
 use crate::job_dependencies::{BucketDependencies, ClientDependencies, DatabaseDependencies, FileSystemDetails, JobCaches, JobDependencies, JobInstanceInfo, JobSystemControls, JobSystemDependencies};
 use crate::job_specific_dependencies::JobSpecificDependencies;
 use crate::util::instrumentation::JobInstruments;
+use crate::util::model_weights_cache::model_weights_cache_directory::ModelWeightsCacheDirectory;
 use crate::util::scoped_execution::ScopedExecution;
 use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator;
 
@@ -234,6 +235,10 @@ async fn main() -> AnyhowResult<()> {
 
   let job_specific_dependencies = JobSpecificDependencies::setup_for_jobs(&scoped_execution)?;
 
+  let scoped_tempdir_for_downloads = ScopedTempDirCreator::for_directory(
+    easyenv::get_env_pathbuf_or_default(
+      "SCOPED_TEMP_DIR_LONG_LIVED_DOWNLOADS", PathBuf::from("/tmp/downloads_long_lived")));
+
   let job_dependencies = JobDependencies {
     db: DatabaseDependencies {
       mysql_pool,
@@ -245,13 +250,15 @@ async fn main() -> AnyhowResult<()> {
       scoped_temp_dir_creator_for_short_lived_downloads: ScopedTempDirCreator::for_directory(
         easyenv::get_env_pathbuf_or_default(
           "SCOPED_TEMP_DIR_SHORT_LIVED_DOWNLOADS", PathBuf::from("/tmp/downloads_short_lived"))),
-      scoped_temp_dir_creator_for_long_lived_downloads: ScopedTempDirCreator::for_directory(
-        easyenv::get_env_pathbuf_or_default(
-          "SCOPED_TEMP_DIR_LONG_LIVED_DOWNLOADS", PathBuf::from("/tmp/downloads_long_lived"))),
+      scoped_temp_dir_creator_for_long_lived_downloads: scoped_tempdir_for_downloads.clone(),
       scoped_temp_dir_creator_for_work: ScopedTempDirCreator::for_directory(
         easyenv::get_env_pathbuf_or_default(
           "SCOPED_TEMP_DIR_WORK", PathBuf::from("/tmp/downloads_long_lived"))),
       semi_persistent_cache,
+      model_weights_cache_directory: ModelWeightsCacheDirectory::setup_from_env_and_deps(
+        &scoped_tempdir_for_downloads,
+        &public_bucket_client,
+      )?,
     },
     buckets: BucketDependencies {
       public_bucket_client,

crates/service/job/inference_job/src/util/mod.rs

@@ -1,8 +1,9 @@
 pub mod common_commands;
 pub mod file_cache_dir;
 pub mod from_env;
-pub mod maybe_download_file_from_bucket;
 pub mod instrumentation;
+pub mod maybe_download_file_from_bucket;
+pub mod model_weights_cache;
 
 #[macro_use]
 pub mod model_downloader;

crates/service/job/inference_job/src/util/model_weights_cache/mod.rs

@@ -0,0 +1,2 @@
+pub mod model_weights_cache_directory;
+pub mod model_weights_cache_filename;

crates/service/job/inference_job/src/util/model_weights_cache/model_weights_cache_directory.rs

@@ -0,0 +1,200 @@
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use log::{error, info, warn};
+
+use buckets::public::media_files::bucket_file_path::MediaFileBucketPath;
+use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath;
+use cloud_storage::bucket_client::BucketClient;
+use errors::AnyhowResult;
+use filesys::file_size::file_size;
+use filesys::rename_across_devices::{rename_across_devices, RenameError};
+use filesys::safe_delete_temp_directory::safe_delete_temp_directory;
+use filesys::safe_delete_temp_file::safe_delete_temp_file;
+use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporter;
+use mysql_queries::queries::model_weights::get::get_weight::RetrievedModelWeight;
+
+use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
+use crate::util::model_weights_cache::model_weights_cache_filename::ModelWeightsCacheMapping;
+use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator;
+
+#[derive(Clone)]
+pub struct ModelWeightsCacheDirectory {
+  /// The path to the cache directory
+  pub local_cache_directory: PathBuf,
+
+  /// Create tempdirs during the download process.
+  pub scoped_tempdir_creator: ScopedTempDirCreator,
+
+  /// A client to download from the public bucket.
+  pub public_bucket_client: BucketClient,
+
+  /// NB: Sometimes the downloader incompletely downloads files. Typically, this is a zero file
+  /// size, but the Rust system may report a non-zero (but small) number of bytes. This should
+  /// later be investigated and the heuristic simplified. The intent here is merely to make sure
+  /// we don't consider these okay:
+  ///
+  ///   -rw-r--r-- 1 root root 55824433 Nov 30 00:12 VM:012rkwsv91zb
+  ///   -rw-r--r-- 1 root root 55823149 Nov 30 00:13 VM:79etbx4fdksv
+  ///   -rw-r--r-- 1 root root        0 Nov 30 00:52 VM:1dzepsnwzbkc (bad file)
+  ///   -rw-r--r-- 1 root root        0 Nov 30 00:29 VM:7c2df5a36qjb (bad file)
+  ///
+  pub maybe_cached_file_minimum_size_required: Option<u64>,
+
+  /// Update the status of jobs when we manage the cache.
+  pub maybe_job_progress_reporter: Option<Arc<dyn JobProgressReporter>>,
+}
+
+
+impl ModelWeightsCacheDirectory {
+  /// Configure and try to create the directory if it doesn't exist.
+  pub fn setup_from_env_and_deps(
+    scoped_temp_dir_creator: &ScopedTempDirCreator,
+    public_bucket_client: &BucketClient,
+  ) -> AnyhowResult<Self> {
+    let directory = easyenv::get_env_pathbuf_or_default(
+      "MODEL_WEIGHTS_CACHE_DIR",
+      "/model_weights_cache");
+
+    info!("Using model weights cache directory: {:?}", directory);
+
+    let directory = ModelWeightsCacheDirectory {
+      local_cache_directory: directory,
+      maybe_cached_file_minimum_size_required: Some(1000), // TODO(bt,2024-02-11): Make configurable
+      maybe_job_progress_reporter: None, // TODO(bt,2024-02-11): Make Sync/Send to hold a copy
+      scoped_tempdir_creator: scoped_temp_dir_creator.clone(),
+      public_bucket_client: public_bucket_client.clone(),
+    };
+
+    directory.try_create_directory()?;
+
+    Ok(directory)
+  }
+
+
+  /// Try to create the directory if it doesn't exist.
+  fn try_create_directory(&self) -> std::io::Result<()> {
+    if !self.local_cache_directory.exists() {
+      std::fs::create_dir_all(&self.local_cache_directory)?;
+    }
+    return Ok(())
+  }
+
+
+  pub fn model_weight_cache_path(&self, model_weight: &RetrievedModelWeight) -> PathBuf {
+    let filename = ModelWeightsCacheMapping::new_from_model(model_weight);
+    let full_path = self.local_cache_directory.join(filename.to_path_buf());
+    full_path
+  }
+
+
+  /// Download the model if we don't already have a copy.
+  /// Return the path from the cache.
+  pub async fn get_model_weight_from_cache_or_bucket(
+    &self,
+    model_weight: &RetrievedModelWeight
+  ) -> Result<PathBuf, ProcessSingleJobError> {
+    let cache_filesystem_path = self.model_weight_cache_path(model_weight);
+
+    if self.is_model_already_cached(model_weight, &cache_filesystem_path)? {
+      return Ok(cache_filesystem_path);
+    }
+
+    self.report_progress(format!("downloading model {}", &model_weight.token)).await;
+
+    let bucket_object_path = WeightFileBucketPath::from_object_hash(
+      &model_weight.public_bucket_hash,
+      model_weight.maybe_public_bucket_prefix.as_deref(),
+      model_weight.maybe_public_bucket_extension.as_deref(),
+    );
+
+    let bucket_object_path = bucket_object_path.to_full_object_pathbuf();
+
+    info!("Downloading model weight (type {:?} token {:?}) from bucket object path {:?}",
+        &model_weight.weights_type,
+        &model_weight.token,
+        &bucket_object_path);
+
+    // NB (1): Download to temp directory to stop concurrent writes and race conditions from other
+    //         workers writing to a shared volume.
+    // NB (2): TempDir exists until it goes out of scope, at which point it should delete from filesystem.
+    let temp_dir = self.scoped_tempdir_creator.new_tempdir("model_weight_download")
+        .map_err(|e| ProcessSingleJobError::from_io_error(e))?;
+
+    let temp_path = temp_dir.path().join("download.part");
+
+    self.public_bucket_client.download_file_to_disk(&bucket_object_path, &temp_path)
+        .await
+        .map_err(|e| {
+          safe_delete_temp_directory(&temp_dir);
+          ProcessSingleJobError::Other(e)
+        })?;
+
+    info!("Downloaded successfully from bucket!");
+
+    let original_size = file_size(&temp_path)
+        .map_err(|err| ProcessSingleJobError::from_anyhow_error(err))?;
+
+    info!("File size of temp download file {:?} is {original_size}", &temp_path);
+
+    info!("Renaming temp file from {:?} to {:?}!", &temp_path, &cache_filesystem_path);
+
+    rename_across_devices(&temp_path, &cache_filesystem_path)
+        .map_err(|err| {
+          error!("could not rename on disk: {:?}", err);
+          safe_delete_temp_file(&temp_path);
+          safe_delete_temp_directory(&temp_dir);
+          match err {
+            RenameError::StorageFull => ProcessSingleJobError::FilesystemFull,
+            RenameError::IoError(err) => ProcessSingleJobError::from_io_error(err),
+          }
+        })?;
+
+    info!("Finished downloading and moving file to {:?}", &cache_filesystem_path);
+
+    safe_delete_temp_file(&temp_path);
+    safe_delete_temp_directory(&temp_dir);
+
+    Ok(cache_filesystem_path)
+  }
+
+
+  fn is_model_already_cached(&self, model_weight: &RetrievedModelWeight, full_cache_filesystem_path: &PathBuf) -> Result<bool, ProcessSingleJobError> {
+    if !full_cache_filesystem_path.exists() {
+      warn!("Model weight (type {:?} token {:?}) does not already exist at path {:?}",
+        &model_weight.weights_type,
+        &model_weight.token,
+        &full_cache_filesystem_path);
+
+      return Ok(false);
+    }
+
+    info!("Model weight (type {:?} token {:?}) exists at path {:?}",
+      &model_weight.weights_type,
+      &model_weight.token,
+      &full_cache_filesystem_path);
+
+    if let Some(existing_file_minimum_size_required) = self.maybe_cached_file_minimum_size_required {
+      let size = file_size(full_cache_filesystem_path)
+          .map_err(|err| ProcessSingleJobError::from_anyhow_error(err))?;
+
+      info!("Model weight has file size = {size}");
+
+      if size < existing_file_minimum_size_required {
+        return Ok(false); // NB: We want to re-download the file.
+      }
+    }
+
+    // TODO(bt, 2022-07-15): Check signature of file as best proof of file validity
+    Ok(true)
+  }
+
+
+  async fn report_progress(&self, _description: String) {
+    if let Some(_job_progress_reporter) = &self.maybe_job_progress_reporter {
+      // TODO(bt,2024-02-10): Needs interior mutability. Meh. Also fail open.
+      //job_progress_reporter.log_status(&description)
+      //    .map_err(|e| ProcessSingleJobError::Other(e))?;
+    }
+  }
+}

crates/service/job/inference_job/src/util/model_weights_cache/model_weights_cache_filename.rs

@@ -0,0 +1,78 @@
+use std::path::PathBuf;
+use enums::by_table::model_weights::weights_types::WeightsType;
+use mysql_queries::queries::model_weights::get::get_weight::RetrievedModelWeight;
+use tokens::tokens::model_weights::ModelWeightToken;
+
+pub struct ModelWeightsCacheMapping {
+  pub token: ModelWeightToken,
+  pub weights_type: WeightsType,
+}
+
+impl ModelWeightsCacheMapping {
+
+  pub fn new_from_model(model_weight: &RetrievedModelWeight) -> Self {
+    Self::from_token_and_type(&model_weight.token, model_weight.weights_type)
+  }
+
+  pub fn from_token_and_type(token: &ModelWeightToken, weights_type: WeightsType) -> Self {
+    Self {
+      token: token.clone(),
+      weights_type,
+    }
+  }
+
+
+  pub fn to_path_buf(&self) -> PathBuf {
+    let extension = self.extension();
+    let filename = format!("{}.{}.{}", &self.token, self.weights_type, extension);
+    PathBuf::from(filename)
+  }
+
+
+  fn extension(&self) -> &str {
+    match self.weights_type {
+      // Pt
+      WeightsType::HifiganTacotron2 => "pt",
+      WeightsType::RvcV2 => "pt",
+      WeightsType::SoVitsSvc => "pt",
+      WeightsType::Tacotron2 => "pt",
+      WeightsType::VallE => "pt",
+      // Safetensors
+      WeightsType::StableDiffusion15 => "safetensors",
+      WeightsType::StableDiffusionXL => "safetensors",
+      WeightsType::LoRA => "safetensors",
+      // Json
+      WeightsType::ComfyUi => "json",
+    }
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use std::path::PathBuf;
+  use enums::by_table::model_weights::weights_types::WeightsType;
+  use tokens::tokens::model_weights::ModelWeightToken;
+  use crate::util::model_weights_cache::model_weights_cache_filename::ModelWeightsCacheMapping;
+
+  #[test]
+  fn stable_diffusion_checkpoints() {
+    let mapper = ModelWeightsCacheMapping::from_token_and_type(
+      &ModelWeightToken::new_from_str("test_token"),
+      WeightsType::StableDiffusion15
+    );
+    let path = mapper.to_path_buf();
+
+    assert_eq!(path, PathBuf::from("test_token.sd_1.5.safetensors"));
+  }
+
+  #[test]
+  fn lora_checkpoints() {
+    let mapper = ModelWeightsCacheMapping::from_token_and_type(
+      &ModelWeightToken::new_from_str("test_lora"),
+      WeightsType::LoRA
+    );
+    let path = mapper.to_path_buf();
+
+    assert_eq!(path, PathBuf::from("test_lora.loRA.safetensors"));
+  }
+}