Merge branch 'main' of github.com:descriptinc/lyrebird-vampnet into main

requirements.txt

@@ -2,12 +2,12 @@ argbind>=0.3.1
 pytorch-ignite
 rich
 audiotools @ git+https://github.com/descriptinc/lyrebird-audiotools.git@0.6.3
+lac @ git+https://github.com/descriptinc/lyrebird-audio-codec.git@main
 tqdm
 tensorboard
 google-cloud-logging==2.2.0
 pytest
 pytest-cov
-papaya_client @ git+https://github.com/descriptinc/lyrebird-papaya.git@master
 pynvml
 psutil
 pandas

scripts/exp/train.py

@@ -59,7 +59,7 @@ IGNORE_INDEX = -100
 @argbind.bind("train", "val", without_prefix=True)
 def build_transform():
     transform = transforms.Compose(
-        tfm.VolumeNorm(("uniform", -32, -14)),
+        tfm.VolumeNorm(("uniform", -32, -20)),
         tfm.VolumeChange(("uniform", -6, 3)),
         tfm.RescaleAudio(),
     )
@@ -250,6 +250,7 @@ def train(
     max_epochs: int = int(100e3),
     epoch_length: int = 1000,
     save_audio_epochs: int = 10,
+    save_epochs: list = [10, 50, 100, 200, 300, 400,],
     batch_size: int = 48,
     grad_acc_steps: int = 1,
     val_idx: list = [0, 1, 2, 3, 4],
@@ -506,6 +507,9 @@ def train(
             loss_key = "loss/val" if "loss/val" in metadata["logs"] else "loss/train"
             self.print(f"Saving to {str(Path('.').absolute())}")
 
+            if self.state.epoch in save_epochs:
+                tags.append(f"epoch={self.state.epoch}")
+
             if self.is_best(engine, loss_key):
                 self.print(f"Best model so far")
                 tags.append("best")

setup.py

@@ -30,11 +30,13 @@ setup(
         "argbind>=0.3.2",
         "pytorch-ignite",
         "rich",
-        "audiotools @ git+https://github.com/descriptinc/lyrebird-audiotools.git@0.6.0",
+        "audiotools @ git+https://github.com/descriptinc/lyrebird-audiotools.git@0.6.3",
+        "lac @ git+https://github.com/descriptinc/lyrebird-audio-codec.git@main",
         "tqdm",
         "tensorboard",
         "google-cloud-logging==2.2.0",
         "torchmetrics>=0.7.3",
         "einops",
+        "flash-attn",
     ],
 )

vampnet/modules/base.py

@@ -153,7 +153,7 @@ class VampBase(at.ml.BaseModel):
         sampling_steps: int = 12,
         start_tokens: Optional[torch.Tensor] = None,
         mask: Optional[torch.Tensor] = None,
-        temperature: Union[float, Tuple[float, float]] = 1.0,
+        temperature: Union[float, Tuple[float, float]] = 0.8,
         top_k: int = None,
         sample: str = "gumbel",
         renoise_mode: str = "start",
@@ -262,7 +262,7 @@ class VampBase(at.ml.BaseModel):
         sampling_steps: int = 24,
         start_tokens: Optional[torch.Tensor] = None,
         mask: Optional[torch.Tensor] = None,
-        temperature: Union[float, Tuple[float, float]] = 1.0,
+        temperature: Union[float, Tuple[float, float]] = 0.8,
         top_k: int = None,
         sample: str = "multinomial",
         typical_filtering=False,