ok

README.md

@@ -79,16 +79,31 @@ pip install -r requirements.txt
                 --model speech_tokenizer_v2_25hz \
                 --device "cuda" \
                 --batch_size 64 \
-                --file_list /data/learnable-speech/speech/files_test.txt \
+                --file_list /speech/files_test.txt \
                 --skip_existing
    ```
 
 2. **Extracting DAC-VAE latent**
    ```bash
    cd dac-vae
-   python inference.py --checkpoint checkpoint.pt --config config.yml
+   python extract_dac_latents.py --checkpoint checkpoint.pt --config config.yml --root_path dataset --output_dir dataset/dac
    ```
 
+After processing you should have root folder with following files:
+
+```
+dataset_root/
+├── audio_name.wav
+├── audio_name.txt
+├── audio_name_fsq.pt
+├── audio_name_latent.pt
+├── another_audio.wav
+├── another_audio.txt
+├── another_audio_fsq.pt
+├── another_audio_latent.pt
+└── ...
+```
+
 3. **Stage 1: Auto Regressive Transformer**
    ```bash
    #!/bin/bash

speech/config.yaml

@@ -221,7 +221,7 @@ train_conf:
     scheduler_conf:
         warmup_steps: 500
     max_epoch: 2000
-    grad_clip: 5
+    grad_clip: 1
     accum_grad: 1
     log_interval: 5
     save_per_step: 2000

speech/train.py

@@ -54,6 +54,7 @@ def get_args():
         "--qwen_pretrain_path", required=False, help="qwen pretrain path"
     )
     parser.add_argument("--checkpoint", help="checkpoint model")
+    parser.add_argument("--pretrained_model", help="pretrained model")
     parser.add_argument("--model_dir", required=True, help="save model dir")
     parser.add_argument(
         "--tensorboard_dir", default="tensorboard", help="tensorboard log dir"
@@ -209,6 +210,13 @@ def main():
 
     model = configs[args.model]
     start_step, start_epoch = 0, -1
+
+    if args.pretrained_model is not None:
+        # load the pretrained model with some weights is ignore
+        logger.info(f"Load pretrained model from {args.pretrained_model}")
+        state_dict = torch.load(args.pretrained_model, map_location="cpu")
+        model.load_state_dict(state_dict, strict=False)
+
     if args.checkpoint is not None:
         if os.path.exists(args.checkpoint):
             logger.info(f"Load checkpoint from {args.checkpoint}")