Add .gitignore and implement model weight loading with quantization support

.gitignore

@@ -0,0 +1,8 @@
+env/
+.cache/
+squashfs-root/
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+*.AppImage

app.py

@@ -21,7 +21,7 @@ import traceback
 import os
 from smolagents import Tool
 from typing import Optional
-from weavemuse.models.notagen.inference import inference_patch, load_model_weights
+from weavemuse.models.notagen.inference import inference_patch, download_model_weights, model
 
 
 logging.basicConfig(level=logging.INFO)
@@ -37,6 +37,34 @@ device = "cuda"
 logger.info(f"Preparing NotaGen tool on device: {device}")
 
 
+def load_model_weights(model_id=None):
+    """Load model weights with intelligent quantization support."""
+    global model
+            
+    # Fall back to original weights
+    try:
+        logger.info("Loading original full-precision model...")
+        original_path = download_model_weights(repo_id="ElectricAlexis/NotaGen")
+        
+        checkpoint = torch.load(original_path, map_location=device, weights_only=False)
+        if 'model' in checkpoint:
+            state_dict = checkpoint['model']
+        else:
+            state_dict = checkpoint
+            
+        model.load_state_dict(state_dict)        
+        model.eval()
+        logger.info("✅ Original model loaded successfully!")
+        return True
+        
+    except Exception as e:
+        logger.error(f"Failed to load model weights: {e}")
+        import traceback
+        logger.error(traceback.format_exc())
+        raise
+
+
+
 class SimpleNotaGenTool(Tool):
     """
     Simple tool for symbolic music generation using NotaGen model.
@@ -120,12 +148,15 @@ class SimpleNotaGenTool(Tool):
             
         Returns:
             Path to generated ABC file or error message
-        """                                            
+        """
+        global model
+        global device
+
         logger.info(f"Generating music: {period}-{composer}-{instrumentation}")
         
         # Create prompt for NotaGen
         prompt = f"{period}-{composer}-{instrumentation}"
-        
+        model = model.to(device)
         # Use the inference function
         inference_fn = inference_patch
         if inference_fn is None: