enable load large checkpoints using safetensors

app.py

@@ -7,6 +7,9 @@ from anticipation import ops
 from anticipation.tokenize import extract_instruments
 import torch
 from pyharp import *
+from safetensors.torch import load_file
+import os
+
 
 #Model Choices
 SMALL_MODEL = "stanford-crfm/music-small-800k"
@@ -25,6 +28,7 @@ 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:
@@ -35,6 +39,34 @@ def load_amt_model(model_choice):
     
     model_cache[model_choice] = model  
     return model
+    '''
+    
+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")
+
+    if model_choice == LARGE_MODEL:
+        # Large model uses safetensors
+        model_dir = "./tmp_music_large"
+        os.makedirs(model_dir, exist_ok=True)
+
+        print(f"Loading {LARGE_MODEL} from safetensors format...")
+        model = AutoModelForCausalLM.from_pretrained(
+            LARGE_MODEL, 
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            low_cpu_mem_usage=True
+        ).to(device)
+    else:
+        # Small and medium use standard PyTorch .bin format
+        print(f"Loading {model_choice} from standard format...")
+        model = AutoModelForCausalLM.from_pretrained(model_choice).to(device)
+    
+    model_cache[model_choice] = model  
+    return model
+