Update audiocraft/models/loaders.py

audiocraft/models/loaders.py

@@ -100,14 +100,19 @@ def _delete_param(cfg: DictConfig, full_name: str):
     OmegaConf.set_struct(cfg, True)
 
 
-def load_lm_model(file_or_url_or_id: tp.Union[Path, str], device='cpu', cache_dir: tp.Optional[str] = None):
+def load_lm_model(file_or_url_or_id: tp.Union[Path, str], device='cpu', depth='float32', cache_dir: tp.Optional[str] = None):
     pkg = load_lm_model_ckpt(file_or_url_or_id, cache_dir=cache_dir)
     cfg = OmegaConf.create(pkg['xp.cfg'])
     cfg.device = str(device)
     if cfg.device == 'cpu':
         cfg.dtype = 'float32'
     else:
-        cfg.dtype = 'float16'
+        if depth=='float32':
+            cfg.dtype = 'float32'
+        if depth=='bfloat16':
+            cfg.dtype = 'bfloat16'
+        if depth=='float16':
+            cfg.dtype = 'float16'
     _delete_param(cfg, 'conditioners.self_wav.chroma_stem.cache_path')
     _delete_param(cfg, 'conditioners.args.merge_text_conditions_p')
     _delete_param(cfg, 'conditioners.args.drop_desc_p')