debug for amt model loading with safetensors

app.py

@@ -28,43 +28,61 @@ model_card = ModelCard(
 
 model_cache = {}
 
+'''
+def load_amt_model(model_choice):
+    """Loads and caches the AMT model inside the worker process."""
+    if model_choice in model_cache:
+        return model_cache[model_choice]  
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = AutoModelForCausalLM.from_pretrained(model_choice).to(device)
+    
+    model_cache[model_choice] = model  
+    return model
+'''
+
+import os
+import traceback
+from transformers import AutoModelForCausalLM, AutoConfig
+import torch
+from safetensors.torch import load_file
+
 def load_amt_model(model_choice):
     """Loads and caches the AMT model inside the worker process."""
     if model_choice in model_cache:
-        print(f"✅ Model {model_choice} loaded from cache")
         return model_cache[model_choice]  
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    cache_dir = "./models"
 
     try:
-        print(f"📌 Checking available model files for: {model_choice}")
+        print(f"📌 Loading model: {model_choice}")
 
-        # Detect model file format
-        model_path = os.path.join(cache_dir, model_choice)
-        bin_path = os.path.join(model_path, "pytorch_model.bin")
-        safetensor_path = os.path.join(model_path, "model.safetensors")
+        # For small and medium models, use the original loading method
+        if model_choice in ["stanford-crfm/music-small-800k", "stanford-crfm/music-medium-800k"]:
+            model = AutoModelForCausalLM.from_pretrained(model_choice).to(device)
+            print(f"✅ Successfully loaded {model_choice} using standard PyTorch method.")
 
-        if os.path.exists(bin_path):
-            print(f"✅ Detected pytorch_model.bin for {model_choice}. Loading normally...")
-            model = AutoModelForCausalLM.from_pretrained(model_choice, cache_dir=cache_dir).to(device)
+        # For large model, use SafeTensors method
+        elif model_choice == "stanford-crfm/music-large-800k":
+            print(f"📌 Detected SafeTensors format for {model_choice}, loading manually...")
 
-        elif os.path.exists(safetensor_path):
-            print(f"✅ Detected model.safetensors for {model_choice}. Loading using SafeTensors...")
-            
-            # Load model configuration
-            config = AutoConfig.from_pretrained(model_choice, cache_dir=cache_dir)
+            # Load model config
+            config = AutoConfig.from_pretrained(model_choice)
             
             # Load SafeTensors manually
-            state_dict = load_file(safetensor_path)  
+            safetensor_path = os.path.join("models", model_choice, "model.safetensors")
+
+            if not os.path.exists(safetensor_path):
+                raise FileNotFoundError(f"❌ SafeTensors file not found: {safetensor_path}")
+
+            state_dict = load_file(safetensor_path)
             model = AutoModelForCausalLM.from_config(config)  # Initialize model
             model.load_state_dict(state_dict)  # Load weights
             model.to(device)
+            print(f"✅ Successfully loaded {model_choice} using SafeTensors.")
 
         else:
-            raise ValueError(f"❌ No valid model file found for {model_choice}")
-
-        print(f"✅ Successfully loaded model: {model_choice}")
+            raise ValueError(f"❌ Unknown model choice: {model_choice}")
 
     except Exception as e:
         print(f"❌ Error loading model {model_choice}: {e}")
@@ -75,7 +93,6 @@ def load_amt_model(model_choice):
     return model
 
 
-
 @spaces.GPU
 def generate_accompaniment(midi_file, model_choice, selected_midi_program, history_length):
     """Generates accompaniment for the entire MIDI input, conditioned on the user-selected history length."""