Update demos/musicgen_app.py

demos/musicgen_app.py

@@ -93,6 +93,44 @@ def make_waveform(*args, **kwargs):
         return out
 
 
+def _delete_param(cfg: DictConfig, full_name: str):
+    parts = full_name.split(".")
+    for part in parts[:-1]:
+        if part in cfg:
+            cfg = cfg[part]
+        else:
+            return
+    OmegaConf.set_struct(cfg, False)
+    if parts[-1] in cfg:
+        del cfg[parts[-1]]
+    OmegaConf.set_struct(cfg, True)
+
+
+def load_lm_model(
+    file_or_url_or_id: tp.Union[Path, str],
+    device=None,
+):
+    pkg = torch.load(file_or_url_or_id, map_location=device)
+    cfg = OmegaConf.create(pkg["xp.cfg"])
+    cfg.device = str(device)
+    if cfg.device == "cpu":
+        cfg.dtype = "float32"
+    else:
+        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")
+    model = get_lm_model(cfg)
+    model.load_state_dict(pkg["best_state"])
+    model.eval()
+    model.cfg = cfg
+    return model
+
+
+def load_compression_model(file_or_url_or_id: tp.Union[Path, str], device=None):
+    return CompressionModel.get_pretrained(file_or_url_or_id, device=device)
+
+
 def load_model(version="facebook/musicgen-small"):
     global MODEL
     print("Loading Musivesal musicgen-small")  # , version
@@ -101,8 +139,25 @@ def load_model(version="facebook/musicgen-small"):
         del MODEL
         torch.cuda.empty_cache()
         MODEL = None  # in case loading would crash
-        MODEL = MusicGen.get_pretrained("data")
+        # MODEL = MusicGen.get_pretrained("/Users/ebenge/repos/audiocraft/data/")
+        lm = load_lm_model("../data/state_dict.bin", device="cudu")
+        compression_model = load_compression_model(
+            "facebook/encodec_32khz", device="cudu"
+        )
+        MODEL = MusicGen("musiversal/musicgen-small", compression_model, lm)
         print("Custom model loaded.")
+        
+
+# def load_model(version="facebook/musicgen-small"):
+#     global MODEL
+#     print("Loading Musivesal musicgen-small")  # , version
+#     if MODEL is None or MODEL.name != version:
+#         # Clear PyTorch CUDA cache and delete model
+#         del MODEL
+#         torch.cuda.empty_cache()
+#         MODEL = None  # in case loading would crash
+#         MODEL = MusicGen.get_pretrained("data")
+#         print("Custom model loaded.")
 
 
 def load_diffusion():