Update app.py

app.py

@@ -12,6 +12,7 @@ from audiocraft.models import MusicGen
 from torch.cuda.amp import autocast
 import warnings
 import random
+from huggingface_hub import login
 
 # Suppress warnings for cleaner output
 warnings.filterwarnings("ignore")
@@ -19,6 +20,17 @@ warnings.filterwarnings("ignore")
 # Set PYTORCH_CUDA_ALLOC_CONF to manage memory fragmentation
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
 
+# Authenticate with Hugging Face
+hf_token = os.getenv("HF_TOKEN")
+if not hf_token:
+    print("ERROR: HF_TOKEN environment variable not set. Please set it in the Hugging Face Space settings.")
+    sys.exit(1)
+try:
+    login(hf_token)
+except Exception as e:
+    print(f"ERROR: Failed to authenticate with Hugging Face: {e}")
+    sys.exit(1)
+
 # Check critical dependencies
 if np.__version__ != "1.23.5":
     print(f"WARNING: NumPy version {np.__version__} is being used. Tested with numpy==1.23.5.")
@@ -40,20 +52,15 @@ torch.cuda.synchronize()
 
 # 2) LOAD MUSICGEN INTO VRAM
 try:
-    print("Loading MusicGen medium model into VRAM...")
-    local_model_path = "./models/musicgen-medium"
-    if not os.path.exists(local_model_path):
-        print(f"ERROR: Local model path {local_model_path} does not exist.")
-        print("Please download the MusicGen medium model weights and place them in the correct directory.")
-        sys.exit(1)
-    musicgen_model = MusicGen.get_pretrained(local_model_path, device=device)
+    print("Loading MusicGen medium model from Hugging Face...")
+    musicgen_model = MusicGen.get_pretrained("facebook/musicgen-medium", device=device)
     musicgen_model.set_generation_params(
         duration=10,  # Default chunk duration
         two_step_cfg=False  # Disable two-step CFG for stability
     )
 except Exception as e:
     print(f"ERROR: Failed to load MusicGen model: {e}")
-    print("Ensure model weights are correctly placed and dependencies are installed.")
+    print("Ensure you have access to facebook/musicgen-medium and a valid HF_TOKEN.")
     sys.exit(1)
 
 # 3) RESOURCE MONITORING FUNCTION