stash save

tests/weight_test.rs

@@ -425,3 +425,30 @@ fn silero_vad_weight() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn moss_tts_nano_weight() -> Result<()> {
+    // cargo test -F cuda --test weight_test moss_tts_nano_weight -r -- --nocapture
+    let save_dir =
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    let model_path = format!("{}/openmoss/MOSS-TTS-Nano/pytorch_model.bin", save_dir);
+    let dict = read_all_with_key(&model_path, None)?;
+    for (k, v) in dict {
+        println!("key: {}, tensor shape: {:?}", k, v);
+    }
+    Ok(())
+}
+
+#[test]
+fn moss_audio_tokenizer_nano_weight() -> Result<()> {
+    // cargo test -F cuda --test weight_test moss_audio_tokenizer_nano_weight -r -- --nocapture
+    let save_dir =
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    let model_path = format!("{}/openmoss/MOSS-Audio-Tokenizer-Nano/model-00001-of-00001.safetensors", save_dir);
+    let device = get_device(None);
+    let weights = safetensors::load(model_path, &device)?;
+    for (key, tensor) in weights.iter() {
+        println!("=== {} === {:?}", key, tensor);
+    }
+    Ok(())
+}