初始化项目，由ModelHub XC社区提供模型

Model: Vikhrmodels/salt-asr_speech_1_wav_1_tts_speech_3_text-10k
Source: Original Platform

.gitattributes

@@ -0,0 +1,53 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+ 
+ 
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+ 
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+ 
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+ 
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
+pytorch_model-00001-of-00002.bin filter=lfs diff=lfs merge=lfs -text
+speechtokenizer/SpeechTokenizer.pt filter=lfs diff=lfs merge=lfs -text
+pytorch_model-00002-of-00002.bin filter=lfs diff=lfs merge=lfs -text
+wavtokenizer/WavTokenizer_small_600_24k_4096.ckpt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,152 @@
+---
+library_name: transformers
+datasets:
+- openslr/librispeech_asr
+- parler-tts/libritts-r-filtered-speaker-descriptions
+- llm-blender/mix-instruct
+language:
+- en
+base_model:
+- meta-llama/Llama-3.2-3B
+---
+
+**Inference**:<br>
+
+```python
+device = "cuda"
+
+n_codebooks_tts = 3
+n_codebooks_asr = 1
+
+start_audio_token = "<|start_of_audio|>"
+end_audio_token = "<|end_of_audio|>"
+end_sequence_token = "<|end_of_text|>"
+
+base_model = "Vikhrmodels/salt-asr_speech_1_wav_1_tts_speech_3_text-10k"
+
+
+def decode_tts(tokens, quantizer, n_codebooks, n_original_tokens, start_audio_token_id, end_audio_token_id):
+    # find start and end indices of audio tokens
+    start = torch.nonzero(tokens == start_audio_token_id)
+    end = torch.nonzero(tokens == end_audio_token_id)
+
+    start = start[0, -1] + 1 if len(start) else 0
+    end = end[0, -1] if len(end) else tokens.shape[-1]
+
+    # subtract length of original vocabulary -> tokens in range [0, 1024)
+    audio_tokens = tokens[start:end] % n_original_tokens
+    reminder = audio_tokens.shape[-1] % n_codebooks
+
+    if reminder:
+        # pad if last frame is incomplete
+        pad_tokens = torch.zeros(n_codebooks - reminder, device="cuda")
+        audio_tokens = torch.cat([audio_tokens, pad_tokens], dim=0)
+
+    transposed = audio_tokens.view(-1, n_codebooks).t()
+    codes = transposed.view(n_codebooks, 1, -1).to(device)
+
+    audio = quantizer.decode(codes).squeeze(0)
+
+    del tokens
+    del audio_tokens
+    torch.cuda.empty_cache()
+
+    return AudioSignal(audio.detach().cpu().numpy(), quantizer.sample_rate)
+
+
+def infer_text_to_audio(text, model, tokenizer, quantizer, max_seq_length=1024, top_k=20):
+    text_tokenized = tokenizer(text, return_tensors="pt")
+    text_input_tokens = text_tokenized["input_ids"].to(device)
+
+    soa = tokenizer(start_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    eoa = tokenizer(end_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+
+    text_tokens = torch.cat([text_input_tokens, soa], dim=1)
+    attention_mask = torch.ones(text_tokens.size(), device=device)
+
+    output_audio_tokens = model.generate(
+        text_tokens,
+        attention_mask=attention_mask,
+        max_new_tokens=max_seq_length,
+        top_k=top_k,
+        do_sample=True,
+        temperature=0.1,
+        repetition_penalty=1.1,
+        length_penalty=1.2,
+        no_repeat_ngram_size=3,
+    )
+
+    audio_signal = decode_tts(output_audio_tokens[0], quantizer, 3, len(tokenizer), soa, eoa)
+
+    return audio_signal
+
+
+def infer_audio_to_text(audio_path, model, tokenizer, quantizer_speech, quantizer_wav, max_seq_length=1024, top_k=20):
+    audio_data, sample_rate = torchaudio.load(audio_path)
+
+    audio = audio_data.view(1, -1).float().to(device)
+    bandwidth_id = torch.tensor([0])
+
+    codes_semantics = quantizer_speech.encode(audio.reshape(1, 1, -1))
+    raw_semantic_tokens = codes_semantics + len(tokenizer)
+    raw_semantic_tokens = raw_semantic_tokens[:1].view(1, -1)
+
+    _, codes = quantizer_wav.encode_infer(audio, bandwidth_id=bandwidth_id)
+    raw_acoustic_tokens = codes + len(tokenizer) + 1024
+    raw_acoustic_tokens = raw_acoustic_tokens.view(1, -1)
+
+    audio_tokens = torch.cat([raw_semantic_tokens, raw_acoustic_tokens], dim=1)
+
+    soa = tokenizer(start_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    eoa = tokenizer(end_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    audio_tokens = torch.cat([soa, audio_tokens, eoa], dim=1)
+    tokens = torch.cat([audio_tokens], dim=1)
+
+    attention_mask = torch.ones(tokens.size(), device=device)
+
+    output_text_tokens = model.generate(
+        tokens,
+        attention_mask=attention_mask,
+        max_new_tokens=max_seq_length,
+        do_sample=True,
+        temperature=0.1,
+        top_p=0.9,
+        top_k=top_k,
+    )
+
+    output_text_tokens = output_text_tokens.cpu()[0]
+    output_text_tokens = output_text_tokens[output_text_tokens < tokenizer(start_audio_token)["input_ids"][-1]]
+    decoded_text = tokenizer.decode(output_text_tokens, skip_special_tokens=True)
+
+    return decoded_text
+
+
+tokenizer = AutoTokenizer.from_pretrained(base_model, cache_dir=".")
+model = AutoModelForCausalLM.from_pretrained(
+    base_model,
+    cache_dir=".",
+    torch_dtype=torch.bfloat16,
+    attn_implementation="sdpa",
+    device_map={"": 0}
+)
+
+quantizer_speech = SpeechTokenizer.load_from_checkpoint("speechtokenizer/config.json",
+                                                        "speechtokenizer/SpeechTokenizer.pt")
+quantizer_speech = quantizer_speech.eval().to(device)
+codebook_size = quantizer_speech.quantizer.bins
+
+quantizer_wav = WavTokenizer.from_pretrained0802("wavtokenizer/config.yaml",
+                                                 "wavtokenizer/WavTokenizer_small_600_24k_4096.ckpt")
+quantizer_wav = quantizer_wav.to(device)
+
+text = ("Say 'COUNT NUMBERS FROM ONE TO TEN' with a male speaker delivers a very monotone and "
+        "low-pitched speech with a moderate speed in a setting with almost no noise, "
+        "creating a clear and quiet recording.")
+
+audio_signal = infer_text_to_audio(text, model, tokenizer, quantizer_speech, top_k=60)
+audio_signal.write("output.wav")
+
+audio_path = "./input.wav"
+generated_text = infer_audio_to_text(audio_path, model, tokenizer, quantizer_speech, quantizer_wav, top_k=10)
+print(generated_text)
+```

config.json

@@ -0,0 +1,36 @@
+{
+  "_name_or_path": "meta-llama/Llama-3.2-3B",
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 128000,
+  "eos_token_id": 128001,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 3072,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "max_position_embeddings": 131072,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 24,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": {
+    "factor": 32.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "llama3"
+  },
+  "rope_theta": 500000.0,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.44.0",
+  "use_cache": true,
+  "vocab_size": 133379
+}

configuration.json

@@ -0,0 +1 @@
+{"framework": "pytorch", "task": "text-generation", "allow_remote": true}

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 128000,
+  "do_sample": true,
+  "eos_token_id": 128001,
+  "temperature": 0.6,
+  "top_p": 0.9,
+  "transformers_version": "4.44.0"
+}

inference.py

@@ -0,0 +1,154 @@
+import torchaudio
+import torch
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+)
+
+from speechtokenizer import SpeechTokenizer
+from WavTokenizer.decoder.pretrained import WavTokenizer
+from audiotools import AudioSignal
+
+
+def decode_tts(tokens, quantizer, n_codebooks, n_original_tokens, start_audio_token_id, end_audio_token_id):
+    # find start and end indices of audio tokens
+    start = torch.nonzero(tokens == start_audio_token_id)
+    end = torch.nonzero(tokens == end_audio_token_id)
+
+    start = start[0, -1] + 1 if len(start) else 0
+    end = end[0, -1] if len(end) else tokens.shape[-1]
+
+    # subtract length of original vocabulary -> tokens in range [0, 1024)
+    audio_tokens = tokens[start:end] % n_original_tokens
+    reminder = audio_tokens.shape[-1] % n_codebooks
+
+    if reminder:
+        # pad if last frame is incomplete
+        pad_tokens = torch.zeros(n_codebooks - reminder, device="cuda")
+        audio_tokens = torch.cat([audio_tokens, pad_tokens], dim=0)
+
+    transposed = audio_tokens.view(-1, n_codebooks).t()
+    codes = transposed.view(n_codebooks, 1, -1).to(device)
+    print(codes)
+
+    audio = quantizer.decode(codes).squeeze(0)
+
+    del tokens
+    del audio_tokens
+    torch.cuda.empty_cache()
+
+    return AudioSignal(audio.detach().cpu().numpy(), quantizer.sample_rate)
+
+
+def infer_text_to_audio(text, model, tokenizer, quantizer, max_seq_length=1024, top_k=20):
+    text_tokenized = tokenizer(text, return_tensors="pt")
+    text_input_tokens = text_tokenized["input_ids"].to(device)
+
+    soa = tokenizer(start_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    eoa = tokenizer(end_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+
+    text_tokens = torch.cat([text_input_tokens, soa], dim=1)
+    attention_mask = torch.ones(text_tokens.size(), device=device)
+
+    output_audio_tokens = model.generate(
+        text_tokens,
+        attention_mask=attention_mask,
+        max_new_tokens=max_seq_length,
+        top_k=top_k,
+        do_sample=True,
+        temperature=0.1,
+        repetition_penalty=1.1,
+        length_penalty=1.2,
+        no_repeat_ngram_size=3,
+    )
+
+    audio_signal = decode_tts(output_audio_tokens[0], quantizer, 3, len(tokenizer), soa, eoa)
+
+    return audio_signal
+
+
+def infer_audio_to_text(audio_path, model, tokenizer, quantizer_speech, quantizer_wav, max_seq_length=1024, top_k=20):
+    audio_data, sample_rate = torchaudio.load(audio_path)
+
+    audio = audio_data.view(1, -1).float().to(device)
+    bandwidth_id = torch.tensor([0])
+
+    codes_semantics = quantizer_speech.encode(audio.reshape(1, 1, -1))
+    raw_semantic_tokens = codes_semantics + len(tokenizer)
+    raw_semantic_tokens = raw_semantic_tokens[:1].view(1, -1)
+
+    _, codes = quantizer_wav.encode_infer(audio, bandwidth_id=bandwidth_id)
+    raw_acoustic_tokens = codes + len(tokenizer) + 1024
+    raw_acoustic_tokens = raw_acoustic_tokens.view(1, -1)
+
+    audio_tokens = torch.cat([raw_semantic_tokens, raw_acoustic_tokens], dim=1)
+
+    soa = tokenizer(start_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    eoa = tokenizer(end_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    audio_tokens = torch.cat([soa, audio_tokens, eoa], dim=1)
+    tokens = torch.cat([audio_tokens], dim=1)
+
+    attention_mask = torch.ones(tokens.size(), device=device)
+
+    output_text_tokens = model.generate(
+        tokens,
+        attention_mask=attention_mask,
+        max_new_tokens=max_seq_length,
+        do_sample=True,
+        temperature=0.1,
+        top_p=0.9,
+        top_k=top_k,
+    )
+
+    output_text_tokens = output_text_tokens.cpu()[0]
+    output_text_tokens = output_text_tokens[output_text_tokens < tokenizer(start_audio_token)["input_ids"][-1]]
+    decoded_text = tokenizer.decode(output_text_tokens, skip_special_tokens=True)
+
+    return decoded_text
+
+
+device = "cuda"
+
+n_codebooks_tts = 3
+n_codebooks_asr = 1
+
+start_audio_token = "<|start_of_audio|>"
+end_audio_token = "<|end_of_audio|>"
+end_sequence_token = "<|end_of_text|>"
+
+base_model = "Vikhrmodels/salt-asr_speech_1_wav_1_tts_speech_3_text-10k"
+
+
+if __name__ == "__main__":
+    tokenizer = AutoTokenizer.from_pretrained(base_model, cache_dir=".")
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model,
+        cache_dir=".",
+        torch_dtype=torch.bfloat16,
+        attn_implementation="sdpa",
+        device_map={"": 0}
+    )
+
+    quantizer_speech = SpeechTokenizer.load_from_checkpoint("speechtokenizer/config.json",
+                                                            "speechtokenizer/SpeechTokenizer.pt")
+    quantizer_speech = quantizer_speech.eval().to(device)
+    codebook_size = quantizer_speech.quantizer.bins
+
+    quantizer_wav = WavTokenizer.from_pretrained0802("wavtokenizer/config.yaml",
+                                                     "wavtokenizer/WavTokenizer_small_600_24k_4096.ckpt")
+    quantizer_wav = quantizer_wav.to(device)
+
+    text = ("Say 'COUNT NUMBERS FROM ONE TO TEN' with a male speaker delivers a very monotone and "
+            "low-pitched speech with a moderate speed in a setting with almost no noise, "
+            "creating a clear and quiet recording.")
+
+    audio_signal = infer_text_to_audio(text, model, tokenizer, quantizer_speech, top_k=60)
+    audio_signal.write("output.wav")
+
+    audio_path = "./input.wav"
+    generated_text = infer_audio_to_text(audio_path, model, tokenizer, quantizer_speech, quantizer_wav, top_k=10)
+    print(generated_text)
+
+
+

pytorch_model-00001-of-00002.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:227754cefd8940d4e35b590d77eedd0ac4f905bb3db6fe1a128b59d6460c3b6a
+size 4997319815

pytorch_model-00002-of-00002.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d73457999ff16e982d031bc6e06f8909060361c469aa96796f70314505967f7c
+size 1459746080

pytorch_model.bin.index.json

@@ -0,0 +1,262 @@
+{
+  "metadata": {
+    "total_size": 6456975360
+  },
+  "weight_map": {
+    "lm_head.weight": "pytorch_model-00001-of-00002.bin",
+    "model.embed_tokens.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.0.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.1.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.10.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.11.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.12.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.13.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.14.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.15.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.16.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.17.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.18.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.19.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.2.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.20.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.20.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.20.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.20.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.21.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.21.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.21.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.21.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.21.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.21.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.21.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.21.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.21.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.22.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.22.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.22.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.22.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.22.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.22.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.22.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.22.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.22.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.23.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.23.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.23.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.23.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.23.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.23.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.23.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.23.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.23.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.24.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.25.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.26.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.mlp.down_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.mlp.gate_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.mlp.up_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.k_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.o_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.q_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.27.self_attn.v_proj.weight": "pytorch_model-00002-of-00002.bin",
+    "model.layers.3.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.3.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.4.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.5.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.6.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.7.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.8.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.mlp.down_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.mlp.gate_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.mlp.up_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.k_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.o_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.q_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.layers.9.self_attn.v_proj.weight": "pytorch_model-00001-of-00002.bin",
+    "model.norm.weight": "pytorch_model-00002-of-00002.bin"
+  }
+}

special_tokens_map.json

@@ -0,0 +1,33 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<|start_of_audio|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|end_of_audio|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "bos_token": {
+    "content": "<|begin_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "[PAD]"
+}

speechtokenizer/SpeechTokenizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d04593b6c9a4b475f91ca481141a6ef5b23e6ac112f347dd2b2717f193c1c728
+size 481906997

speechtokenizer/config.json

@@ -0,0 +1,49 @@
+{
+    "resblock": "1",
+    "num_gpus": 3,
+    "batch_size": 60,
+    "learning_rate": 0.0001,
+    "adam_b1": 0.5,
+    "adam_b2": 0.9,
+    "lr_decay": 0.98,
+    "seed": 1234,
+    "lambda_distill": 0.15,
+
+    "n_filters": 64,
+    "strides": [8,5,4,2],
+    "dimension": 1024,
+    "semantic_dimension": 768,
+    "bidirectional": true,
+    "dilation_base": 2,
+    "residual_kernel_size": 3,
+    "n_residual_layers": 1,
+    "lstm_layers": 2,
+    "activation": "ELU",
+
+
+    "segment_size": 48000,
+    "num_mels": 80,
+    "num_freq": 1025,
+    "n_fft": 1024,
+    "hop_size": 240,
+    "win_size": 1024,
+
+    "sampling_rate": 16000,
+    "sample_rate": 16000,
+
+    "codebook_size": 1024,
+    "n_q": 8,
+
+    "fmin": 0,
+    "fmax": 8000,
+    "fmax_for_loss": null,
+
+    "num_workers": 12,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54322",
+        "world_size": 1
+    }
+}
+

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dc582d0849576e5e5308d8710b29a943b6d23a948280ac90523b8acbc64dded9
+size 9086227

wavtokenizer/WavTokenizer_small_600_24k_4096.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d44c40fbb83d2d42329ac098e252a31b5708fb7b3bf864d108dd3ed26911d004
+size 1589082492

wavtokenizer/config.yaml

@@ -0,0 +1,93 @@
+seed_everything: 3407
+
+data:
+  class_path: decoder.dataset.VocosDataModule
+  init_args:
+    train_params:
+      filelist_path: ./WavTokenizer/data/train/libritts_train
+      sampling_rate: 24000
+      num_samples: 72000
+      batch_size: 40  # 20
+      num_workers: 8
+
+    val_params:
+      filelist_path: ./WavTokenizer/data/infer/librttts_val
+      sampling_rate: 24000
+      num_samples: 72000
+      batch_size: 5   # 10
+      num_workers: 8
+
+model:
+  class_path: decoder.experiment.WavTokenizer
+  init_args:
+    sample_rate: 24000
+    initial_learning_rate: 2e-4
+    mel_loss_coeff: 45
+    mrd_loss_coeff: 1.0
+    num_warmup_steps: 0 # Optimizers warmup steps
+    pretrain_mel_steps: 0  # 0 means GAN objective from the first iteration
+
+    # automatic evaluation
+    evaluate_utmos: true
+    evaluate_pesq: true
+    evaluate_periodicty: true
+
+    resume: false
+    resume_config: ./WavTokenizer/configs/wavtokenizer_smalldata_frame40_3s_nq1_code16384_dim512_kmeans800_attn.yaml
+    resume_model: ./version_3/checkpoints/xxx.ckpt
+
+    feature_extractor:
+      class_path: decoder.feature_extractors.EncodecFeatures
+      init_args:
+        encodec_model: encodec_24khz
+        bandwidths: [6.6, 6.6, 6.6, 6.6]
+        train_codebooks: true
+        num_quantizers: 1  
+        dowmsamples: [6, 5, 5, 4]
+        vq_bins: 4096
+        vq_kmeans: 200
+
+    backbone:
+      class_path: decoder.models.VocosBackbone
+      init_args:
+        input_channels: 512
+        dim: 768
+        intermediate_dim: 2304
+        num_layers: 12
+        adanorm_num_embeddings: 4  
+
+    head:
+      class_path: decoder.heads.ISTFTHead
+      init_args:
+        dim: 768
+        n_fft: 2400 
+        hop_length: 600
+        padding: same
+
+trainer:
+  logger:
+    class_path: pytorch_lightning.loggers.TensorBoardLogger
+    init_args:
+      save_dir: ./WavTokenizer/result/train/wavtokenizer_smalldata_frame40_3s_nq1_code4096_dim512_kmeans200_attn/
+  callbacks:
+    - class_path: pytorch_lightning.callbacks.LearningRateMonitor
+    - class_path: pytorch_lightning.callbacks.ModelSummary
+      init_args:
+        max_depth: 2
+    - class_path: pytorch_lightning.callbacks.ModelCheckpoint
+      init_args:
+        monitor: val_loss
+        filename: wavtokenizer_checkpoint_{epoch}_{step}_{val_loss:.4f}
+        save_top_k: 10
+        save_last: true
+    - class_path: decoder.helpers.GradNormCallback
+
+  # Lightning calculates max_steps across all optimizer steps (rather than number of batches)
+  # This equals to 1M steps per generator and 1M per discriminator
+  max_steps: 20000000
+  # You might want to limit val batches when evaluating all the metrics, as they are time-consuming
+  limit_val_batches: 200
+  accelerator: gpu
+  strategy: ddp
+  devices: [0,1,2,3,4,5,6,7]
+  log_every_n_steps: 1000