fix directory arguments

main.py

@@ -333,19 +333,27 @@ def chords_string_to_list(chords: str):
     chrd_times = [x.split(',') for x in chords[1:-1].split('),(')]
     return [(x[0], float(x[1])) for x in chrd_times]
 
-# Add this before model loading
 def patch_jasco_cache():
     """Monkey patch JASCO cache initialization"""
     from audiocraft.modules import jasco_conditioners
     
-    original_init = jasco_conditioners.DrumsConditioner.__init__
+    if hasattr(jasco_conditioners, 'DrumConditioner'):
+        original_init = jasco_conditioners.DrumConditioner.__init__
+    elif hasattr(jasco_conditioners, 'DrumsConditioner'):
+        original_init = jasco_conditioners.DrumsConditioner.__init__
+    else:
+        print("Warning: Could not find DrumConditioner class")
+        return
     
     def new_init(self, *args, **kwargs):
         if 'cache_path' in kwargs:
             kwargs['cache_path'] = os.path.join(CACHE_DIR, 'drum_cache')
         return original_init(self, *args, **kwargs)
     
-    jasco_conditioners.DrumsConditioner.__init__ = new_init
+    if hasattr(jasco_conditioners, 'DrumConditioner'):
+        jasco_conditioners.DrumConditioner.__init__ = new_init
+    elif hasattr(jasco_conditioners, 'DrumsConditioner'):
+        jasco_conditioners.DrumsConditioner.__init__ = new_init
 
 # Apply the patch
 patch_jasco_cache()
@@ -364,32 +372,24 @@ def load_model(version='facebook/jasco-chords-drums-melody-400M'):
             cache_path = os.path.join(CACHE_DIR, version.replace('/', '_'))
             os.makedirs(cache_path, exist_ok=True)
             
-            # Set additional environment variables
+            # Set environment variables for caching
             os.environ['AUDIOCRAFT_CACHE_DIR'] = cache_path
             os.environ['TRANSFORMERS_CACHE'] = os.path.join(cache_path, 'transformers')
             
-            # Initialize model with custom cache configuration
-            model_kwargs = {
-                'device': 'cuda',
-                'cache_dir': cache_path,
-                'model_cache_dir': cache_path
-            }
-            
             # Initialize chord mapping
             mapping_file = initialize_chord_mapping()
             os.environ['AUDIOCRAFT_CHORD_MAPPING'] = mapping_file
             
-            # Load the model with specific cache paths
+            # Load the model with only supported parameters
             MODEL = JASCO.get_pretrained(
                 version,
-                device='cuda',
-                cache_dir=cache_path,
-                local_files_only=False
+                device='cuda'
             )
             MODEL.name = version
             
-            # Configure model paths
-            MODEL._cache_dir = cache_path
+            # Configure model paths after loading
+            if hasattr(MODEL, '_cache_dir'):
+                MODEL._cache_dir = cache_path
             
             # Load the chord mapping
             with open(mapping_file, 'rb') as f: