c2f

requirements.txt

@@ -2,7 +2,7 @@ argbind>=0.3.1
 pytorch-ignite
 rich
 audiotools @ git+https://github.com/descriptinc/lyrebird-audiotools.git@hf/backup-info
-lac @ git+https://github.com/descriptinc/lyrebird-audio-codec.git@main
+lac @ git+https://github.com/descriptinc/lyrebird-audio-codec.git@hf/vampnet-temp
 tqdm
 tensorboard
 google-cloud-logging==2.2.0

scripts/exp/train.py

@@ -545,6 +545,30 @@ def train(
                     plot_fn=None,
                 )
 
+                # sample in 1 step (only for coarse2fine models)
+                if accel.unwrap(model).n_conditioning_codebooks > 0:
+                    sampled_argmax = accel.unwrap(model).sample(
+                        codec=codec,
+                        time_steps=z.shape[-1],
+                        start_tokens=z[i : i + 1],
+                        sample="argmax", 
+                        sampling_steps=1,
+                    )
+                    sampled_argmax.cpu().write_audio_to_tb(
+                        f"sampled_1step-argmax/{i}",
+                        self.writer,
+                        step=self.state.epoch,
+                        plot_fn=None,
+                    )
+                    conditioning = z[i:i+1, : accel.unwrap(model).n_conditioning_codebooks, :]
+                    conditioning = accel.unwrap(model).to_signal(conditioning, codec)
+                    conditioning.cpu().write_audio_to_tb(
+                        f"conditioning/{i}",
+                        self.writer,
+                        step=self.state.epoch,
+                        plot_fn=None,
+                    )
+
         def save_imputation(self, z: torch.Tensor):
             # imputations
             _prefix_amt = prefix_amt

scripts/utils/parallel-gpu.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+# Get the command to execute from the user
+command_to_execute="$1"
+
+# Get the maximum number of GPUs to use from the user
+max_gpus="$2"
+
+# Get the number of instances to start per GPU from the user
+instances_per_gpu="$3"
+
+# Set the CUDA_VISIBLE_DEVICES flag for each GPU
+for gpu_id in $(seq 0 $(($max_gpus - 1))); do
+    export CUDA_VISIBLE_DEVICES="$gpu_id"
+    # Start the specified number of instances for this GPU
+    for i in $(seq 1 "$instances_per_gpu"); do
+        # Run the command in the background
+        $command_to_execute &
+    done
+done
+
+# Wait for all instances to finish
+wait

scripts/utils/process_folder-c2f.py

@@ -6,6 +6,8 @@ import argbind
 from tqdm import tqdm
 import random
 
+from typing import List
+
 from collections import defaultdict
 
 def coarse2fine_infer(
@@ -15,14 +17,15 @@ def coarse2fine_infer(
         device,
         signal_window=3, 
         signal_hop=1.5,
-        max_excerpts=25, 
+        max_excerpts=20, 
     ):
     output = defaultdict(list)
 
     # split into 3 seconds
     windows = [s for s in signal.clone().windows(signal_window, signal_hop)]
+    windows = windows[1:] # skip first window since it's half zero padded
     random.shuffle(windows)
-    for w in windows[1:max_excerpts]: # skip the first window since it's mostly zero padded? 
+    for w in windows[:max_excerpts]: 
         # batch the signal into chunks of 3
         with torch.no_grad():
             # get codes
@@ -68,20 +71,21 @@ def coarse2fine_infer(
 @argbind.bind(without_prefix=True)
 def main(
         sources=[
-            "/home/hugo/data/spotdl/audio/val", "/home/hugo/data/spotdl/audio/test"
+            "/data/spotdl/audio/val", "/data/spotdl/audio/test"
         ], 
         audio_ext="mp3",
         exp_name="noise_mode",
         model_paths=[
-            "ckpt/mask/best/vampnet/weights.pth",
-            "ckpt/random/best/vampnet/weights.pth",
+            "runs/c2f-exp-03.22.23/ckpt/mask/best/vampnet/weights.pth",
+            "runs/c2f-exp-03.22.23/ckpt/random/best/vampnet/weights.pth",
         ],
         model_keys=[
-            "noise_mode=mask",
-            "noise_mode=random",
+            "mask",
+            "random",
         ],
-        vqvae_path="ckpt/wav2wav.pth",
-        device="cuda",
+        vqvae_path: str = "runs/codec-ckpt/codec.pth",
+        device: str = "cuda",
+        output_dir: str = ".",
     ):
     from vampnet.modules.transformer import VampNet
     from lac.model.lac import LAC
@@ -99,20 +103,28 @@ def main(
     vqvae.eval()
     print("Loaded VQVAE.")
 
-    audio_dict = defaultdict(list)
+    output_dir = Path(output_dir) / f"{exp_name}-samples"
+
     for source in sources:
         print(f"Processing {source}...")
-        for path in tqdm(list(Path(source).glob(f"**/*.{audio_ext}"))):
+        source_files = list(Path(source).glob(f"**/*.{audio_ext}"))
+        random.shuffle(source_files)
+        for path in tqdm(source_files):
             sig = AudioSignal(path)
             sig.resample(vqvae.sample_rate).normalize(-24).ensure_max_of_audio(1.0)
 
+            out_dir = output_dir / path.stem 
+            out_dir.mkdir(parents=True, exist_ok=True)
+            if out_dir.exists():
+                print(f"Skipping {path.stem} since {out_dir} already exists.")
+                continue
+
             for model_key, model in models.items():
                 out = coarse2fine_infer(sig, model, vqvae, device)
-                for k in out:
-                    audio_dict[f"{model_key}-{k}"].extend(out[k])
-
-    audio_zip(audio_dict, f"{exp_name}-results.zip")
-    
+                for k, sig_list in out.items():
+                    for i, s in enumerate(sig_list):
+                        s.write(out_dir / f"{model_key}-{k}-{i}.wav")
+            
 
 if __name__ == "__main__":
     args = argbind.parse_args()