fix(tts): use Hugging Face cache for speaker encoder

nemo/collections/tts/models/audio_codec.py

@@ -20,6 +20,7 @@
 import torch
 import torch.nn.functional as F
 from einops import rearrange
+from huggingface_hub import hf_hub_download
 from lightning.pytorch import Trainer
 from omegaconf import DictConfig, OmegaConf
 
@@ -223,11 +224,11 @@ def __init__(self, cfg: DictConfig, trainer: Trainer = None):
         self.scl_loss_scale = cfg.get("scl_loss_scale", False)
         if self.use_scl_loss:
             self.speaker_encoder = ResNetSpeakerEncoder()
-            # load pretrained model
-            # self.speaker_encoder.load_checkpoint("https://github.com/coqui-ai/TTS/releases/download/speaker_encoder_model/model_se.pth.tar")
-            self.speaker_encoder.load_checkpoint(
-                "https://huggingface.co/Edresson/Speaker_Encoder_H_ASP/resolve/main/pytorch_model.bin", strict=False
+            speaker_encoder_checkpoint = hf_hub_download(
+                repo_id="Edresson/Speaker_Encoder_H_ASP",
+                filename="pytorch_model.bin",
             )
+            self.speaker_encoder.load_checkpoint(speaker_encoder_checkpoint, strict=False)
             # freeze the pretrained speaker encoder
             self.speaker_encoder.freeze()
             logging.info("Speaker encoder loaded and frozen !!")

tests/collections/tts/models/test_audio_codec.py

@@ -13,10 +13,13 @@
 # limitations under the License.
 
 
+from unittest.mock import Mock
+
 import pytest
 import torch
 from omegaconf import DictConfig
 
+import nemo.collections.tts.models.audio_codec as audio_codec_module
 from nemo.collections.tts.models import AudioCodecModel
 
 
@@ -96,6 +99,26 @@ def acoustic_codec_model():
 
 
 class TestAudioCodecModel:
+    @pytest.mark.unit
+    def test_speaker_encoder_checkpoint_uses_hf_cache(self, monkeypatch):
+        model_cfg = create_codec_config()
+        model_cfg.use_scl_loss = True
+        checkpoint_path = '/cache/Edresson/Speaker_Encoder_H_ASP/pytorch_model.bin'
+        mock_hf_hub_download = Mock(return_value=checkpoint_path)
+        speaker_encoder = torch.nn.Module()
+        speaker_encoder.audio_config = {'sample_rate': model_cfg.sample_rate}
+        speaker_encoder.load_checkpoint = Mock()
+        speaker_encoder.freeze = Mock()
+        monkeypatch.setattr(audio_codec_module, 'hf_hub_download', mock_hf_hub_download, raising=False)
+        monkeypatch.setattr(audio_codec_module, 'ResNetSpeakerEncoder', Mock(return_value=speaker_encoder))
+
+        AudioCodecModel(cfg=model_cfg)
+
+        mock_hf_hub_download.assert_called_once_with(
+            repo_id='Edresson/Speaker_Encoder_H_ASP', filename='pytorch_model.bin'
+        )
+        speaker_encoder.load_checkpoint.assert_called_once_with(checkpoint_path, strict=False)
+
     @pytest.mark.unit
     def test_forward(self, codec_model):
         batch_size = 2