debug for amt model loading with safetensors 2

app.py

@@ -41,14 +41,8 @@ def load_amt_model(model_choice):
     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:
         return model_cache[model_choice]  
 
@@ -56,30 +50,19 @@ def load_amt_model(model_choice):
 
     try:
         print(f"📌 Loading model: {model_choice}")
-
-        # 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.")
-
-        # For large model, use SafeTensors method
         elif model_choice == "stanford-crfm/music-large-800k":
-            print(f"📌 Detected SafeTensors format for {model_choice}, loading manually...")
-
-            # Load model config
             config = AutoConfig.from_pretrained(model_choice)
-            
-            # Load SafeTensors manually
             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}")
 
+            # Load SafeTensors model weights
             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"❌ Unknown model choice: {model_choice}")
@@ -94,6 +77,7 @@ def load_amt_model(model_choice):
 
 
 @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."""
     
@@ -129,6 +113,48 @@ def generate_accompaniment(midi_file, model_choice, selected_midi_program, histo
     mid.save(output_midi)
 
     return output_midi, None  
+'''
+def generate_accompaniment(midi_file, model_choice, selected_midi_program, history_length):
+    model = load_amt_model(model_choice)
+    # Ensure model loaded successfully before using it
+    if model is None:
+        print(" Model loading failed. Returning error.")
+        return None, "⚠️ Model failed to load."
+
+    print(f"Model loaded successfully: {model_choice}")
+
+    events = midi_to_events(midi_file.name)
+    total_time = round(ops.max_time(events, seconds=True))
+    events, melody = extract_instruments(events, [selected_midi_program])
+
+    if not melody:  
+        print("No melody detected. Please select a valid MIDI program.")
+        return None, "Please select a valid MIDI program that contains events."
+
+    history = ops.clip(events, 0, history_length, clip_duration=False)
+
+    # Generate accompaniment for the remaining duration
+    accompaniment = generate(
+        model,
+        history_length,
+        total_time,
+        inputs=history,
+        controls=melody,
+        top_p=0.95,
+        debug=True  # Enable debug mode if supported
+    )
+
+    # Combine accompaniment with the melody
+    output_events = ops.clip(ops.combine(accompaniment, melody), 0, total_time, clip_duration=True)
+
+    # Convert back to MIDI
+    output_midi = "generated_accompaniment_huggingface.mid"
+    mid = events_to_midi(output_events)
+    mid.save(output_midi)
+
+    print("✅ MIDI generation successful.")
+
+    return output_midi, None  
 
 
 def process_fn(input_midi, model_choice, selected_midi_program, history_length):