Merge pull request #298 from Harushii18/patch-1

Fix musigen app memory issue
facebookresearch · Nov 23, 2023 · 95f2747 · 95f2747
2 parents c65e9be + 201b542
commit 95f2747
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/demos/musicgen_app.py b/demos/musicgen_app.py
@@ -77,8 +77,7 @@ def _cleanup(self):
                 self.files.pop(0)
             else:
                 break
-
-
+
 file_cleaner = FileCleaner()
 
 
@@ -96,6 +95,9 @@ def load_model(version='facebook/musicgen-melody'):
     global MODEL
     print("Loading model", 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(version)