fix chord mapping file error add sys and shutils

main.py

@@ -379,24 +379,25 @@ def patch_jasco_cache():
 patch_jasco_cache()
 
 def ensure_chord_mapping():
-    """Ensure chord mapping file exists and is properly located"""
-    # Create mapping in the cache directory
-    cache_mapping_path = os.path.join(CACHE_DIR, "chord_to_index_mapping.pkl")
-    
-    # Create mapping in the models directory
-    models_dir = os.path.join(os.path.dirname(__file__), "models")
-    os.makedirs(models_dir, exist_ok=True)
-    models_mapping_path = os.path.join(models_dir, "chord_to_index_mapping.pkl")
-    
-    # Create the mapping if it doesn't exist in either location
+    """Ensure chord mapping file exists in all necessary locations"""
     chord_mapping = create_default_chord_mapping()
     
-    # Save to both locations
-    for mapping_path in [cache_mapping_path, models_mapping_path]:
-        with open(mapping_path, "wb") as f:
+    # Define all necessary paths
+    paths = [
+        os.path.join(CACHE_DIR, "chord_to_index_mapping.pkl"),
+        os.path.join(os.path.dirname(__file__), "models", "chord_to_index_mapping.pkl"),
+        "/tmp/chord_to_index_mapping.pkl",  # Additional fallback location
+        os.path.join(CACHE_DIR, "models", "chord_to_index_mapping.pkl")
+    ]
+    
+    # Create directories and save mapping to all locations
+    for path in paths:
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        with open(path, "wb") as f:
             pickle.dump(chord_mapping, f)
+        print(f"Created chord mapping at: {path}")
     
-    return models_mapping_path
+    return paths[0]  # Return the primary mapping path
 
 # Create necessary directories
 os.makedirs("models", exist_ok=True)
@@ -408,45 +409,79 @@ def load_model(version='facebook/jasco-chords-drums-melody-400M'):
     if MODEL is None or MODEL.name != version:
         MODEL = None
         try:
-            # Set up custom cache paths
+            # Set up cache paths
             cache_path = os.path.join(CACHE_DIR, version.replace('/', '_'))
             os.makedirs(cache_path, exist_ok=True)
             
-            # Set environment variables for caching
+            # Ensure chord mapping exists and set environment
+            mapping_file = ensure_chord_mapping()
+            os.environ['AUDIOCRAFT_CHORD_MAPPING'] = mapping_file
+            
+            # Create symbolic links to ensure the mapping is found
+            target_paths = [
+                os.path.join(cache_path, "chord_to_index_mapping.pkl"),
+                os.path.join(os.path.dirname(__file__), "chord_to_index_mapping.pkl")
+            ]
+            
+            for target in target_paths:
+                try:
+                    if not os.path.exists(target):
+                        os.symlink(mapping_file, target)
+                except Exception as e:
+                    print(f"Warning: Could not create symlink at {target}: {e}")
+            
+            # Set environment variables
             os.environ['AUDIOCRAFT_CACHE_DIR'] = cache_path
             os.environ['TRANSFORMERS_CACHE'] = os.path.join(cache_path, 'transformers')
             
-            # Initialize chord mapping
-            mapping_file = initialize_chord_mapping()
-            os.environ['AUDIOCRAFT_CHORD_MAPPING'] = mapping_file
+            # Initialize JASCO with explicit paths
+            kwargs = {
+                'device': 'cuda',
+                'chords_mapping_path': mapping_file
+            }
             
-            # Load the model with only supported parameters
-            MODEL = JASCO.get_pretrained(
-                version,
-                device='cuda'
-            )
+            # Load the model
+            MODEL = JASCO.get_pretrained(version, **kwargs)
             MODEL.name = version
             
-            # Configure model paths after loading
-            if hasattr(MODEL, '_cache_dir'):
-                MODEL._cache_dir = cache_path
+            # Verify chord mapping is loaded
+            if not hasattr(MODEL, 'chord_to_index'):
+                with open(mapping_file, 'rb') as f:
+                    MODEL.chord_to_index = pickle.load(f)
+            
+            print("Model loaded successfully")
             
-            # Load the chord mapping
-            with open(mapping_file, 'rb') as f:
-                MODEL.chord_to_index = pickle.load(f)
-                
         except Exception as e:
+            print(f"Detailed error loading model: {str(e)}")
             raise gr.Error(f"Error loading model: {str(e)}")
         
         if MODEL is None:
             raise gr.Error("Failed to load model")
     return MODEL
 
-# Initialize chord mapping at startup
-mapping_file = ensure_chord_mapping()
-os.environ['AUDIOCRAFT_CHORD_MAPPING'] = mapping_file
-print(f"Initial chord mapping file location: {mapping_file}")
-print(f"Chord mapping exists: {os.path.exists(mapping_file)}")
+# Initialize environment and mappings at startup
+def initialize_environment():
+    """Initialize all necessary environment variables and files"""
+    # Create cache directory structure
+    for subdir in ['models', 'cache', 'huggingface', 'drum_cache', 'transformers']:
+        os.makedirs(os.path.join(CACHE_DIR, subdir), exist_ok=True)
+    
+    # Initialize chord mapping
+    mapping_file = ensure_chord_mapping()
+    os.environ['AUDIOCRAFT_CHORD_MAPPING'] = mapping_file
+    
+    # Set up environment variables
+    os.environ['AUDIOCRAFT_CACHE_DIR'] = CACHE_DIR
+    os.environ['TORCH_HOME'] = CACHE_DIR
+    os.environ['HF_HOME'] = os.path.join(CACHE_DIR, 'huggingface')
+    os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
+    os.environ['XDG_CACHE_HOME'] = CACHE_DIR
+    
+    print("Environment initialized successfully")
+    return mapping_file
+
+# Call initialization at startup
+mapping_file = initialize_environment()
 
 class ModelLoadingError(Exception):
     pass
@@ -635,6 +670,12 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     try:
+        # Initialize environment
+        mapping_file = initialize_environment()
+        print(f"Using chord mapping file: {mapping_file}")
+        print(f"Chord mapping exists: {os.path.exists(mapping_file)}")
+        
+        # Launch the demo
         demo.queue().launch(ssr_mode=False)
     except Exception as e:
         print(f"Error launching demo: {str(e)}")