dropping torch.compile for now

scripts/exp/train.py

@@ -29,6 +29,9 @@ from audiotools.ml.decorators import (
 
 import loralib as lora
 
+import torch._dynamo
+torch._dynamo.config.verbose=True
+
 
 # Enable cudnn autotuner to speed up training
 # (can be altered by the funcs.seed function)
@@ -510,14 +513,14 @@ def load(
 
     if args["fine_tune"]:
         assert fine_tune_checkpoint is not None, "Must provide a fine-tune checkpoint"
-        model = torch.compile(
+        model = (
             VampNet.load(location=Path(fine_tune_checkpoint), 
                          map_location="cpu", 
             )
         )
 
 
-    model = torch.compile(VampNet()) if model is None else model
+    model = VampNet() if model is None else model
     model = accel.prepare_model(model)
 
     # assert accel.unwrap(model).n_codebooks == codec.quantizer.n_codebooks
@@ -601,7 +604,7 @@ def train(
         accel=accel, 
         tracker=tracker, 
         save_path=save_path)
-
+    print("initialized state.")
 
     train_dataloader = accel.prepare_dataloader(
         state.train_data,
@@ -616,13 +619,15 @@ def train(
         num_workers=num_workers,
         batch_size=batch_size,
         collate_fn=state.val_data.collate,
-        persistent_workers=True,
+        persistent_workers=num_workers > 0,
     )
+    print("initialized dataloader.")
 
     
 
     if fine_tune:
         lora.mark_only_lora_as_trainable(state.model)
+        print("marked only lora as trainable.")
 
     # Wrap the functions so that they neatly track in TensorBoard + progress bars
     # and only run when specific conditions are met.
@@ -637,6 +642,7 @@ def train(
     save_samples = when(lambda: accel.local_rank == 0)(save_samples)
     checkpoint = when(lambda: accel.local_rank == 0)(checkpoint)
 
+    print("starting training loop.")
     with tracker.live:
         for tracker.step, batch in enumerate(train_dataloader, start=tracker.step):
             train_loop(state, batch, accel)

scripts/utils/split_long_audio_file.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+import argbind
+
+import audiotools as at
+import tqdm
+
+
+@argbind.bind(without_prefix=True)
+def split_long_audio_file(
+    file: str = None, 
+    max_chunk_size_s: int = 60*10
+):
+    file = Path(file)
+    output_dir = file.parent / file.stem
+    output_dir.mkdir()
+    
+    sig = at.AudioSignal(file)
+
+    # split into chunks
+    for i, sig in tqdm.tqdm(enumerate(sig.windows(
+        window_duration=max_chunk_size_s, hop_duration=max_chunk_size_s/2, 
+        preprocess=True))
+    ):
+        sig.write(output_dir / f"{i}.wav")
+
+    print(f"wrote {len(list(output_dir.glob('*.wav')))} files to {output_dir}")
+    
+    return output_dir
+
+if __name__ == "__main__":
+    args = argbind.parse_args()
+
+    with argbind.scope(args):
+        split_long_audio_file()