Update audiocraft/models/musicgen.py

audiocraft/models/musicgen.py

@@ -100,7 +100,7 @@ class MusicGen:
         return self.compression_model.channels
 
     @staticmethod
-    def get_pretrained(name: str = 'facebook/musicgen-melody', device=None):
+    def get_pretrained(name: str = 'facebook/musicgen-melody', device=None, torch_dtype=torch.float32):
         """Return pretrained model, we provide four models:
         - facebook/musicgen-small (300M), text to music,
           # see: https://huggingface.co/facebook/musicgen-small
@@ -129,7 +129,7 @@ class MusicGen:
                 f"Please use full pre-trained id instead: facebook/musicgen-{name}")
             name = _HF_MODEL_CHECKPOINTS_MAP[name]
 
-        lm = load_lm_model(name, device=device)
+        lm = load_lm_model(name, device=device, torch_dtype=torch_dtype)
         compression_model = load_compression_model(name, device=device)
         if 'self_wav' in lm.condition_provider.conditioners:
             lm.condition_provider.conditioners['self_wav'].match_len_on_eval = True