initial
Browse files- CosyVoice-300M-25Hz-Music/VERSION_Vq0206Vocoder_Sing_0106 +2 -0
- CosyVoice-300M-25Hz-Music/campplus.onnx +3 -0
- CosyVoice-300M-25Hz-Music/cosyvoice.yaml +98 -0
- CosyVoice-300M-25Hz-Music/flow.pt +3 -0
- CosyVoice-300M-25Hz-Music/hift.pt +3 -0
- CosyVoice-300M-25Hz-Music/speech_tokenizer_v1.onnx +3 -0
- CosyVoice-300M-25Hz/VERSION_Vq0206Vocoder_1202 +0 -0
- CosyVoice-300M-25Hz/campplus.onnx +3 -0
- CosyVoice-300M-25Hz/cosyvoice.yaml +98 -0
- CosyVoice-300M-25Hz/flow.pt +3 -0
- CosyVoice-300M-25Hz/hift.pt +3 -0
- CosyVoice-300M-25Hz/speech_tokenizer_v1.onnx +3 -0
- config.json +22 -0
- configuration_stepaudio.py +41 -0
- model-00001.safetensors +3 -0
- model.safetensors.index.json +1 -0
- modeling_stepaudio.py +392 -0
- tokenizer.model +3 -0
- tokenizer_config.json +14 -0
CosyVoice-300M-25Hz-Music/VERSION_Vq0206Vocoder_Sing_0106
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
# Flow Model
|
2 |
+
/mnt/shared-storage/groups/tts/sfy/flowmatching/models_22k/flow_unet_22k_1node_vq0206_music_sft/model_epoch_20_whole.pt
|
CosyVoice-300M-25Hz-Music/campplus.onnx
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:a6ac6a63997761ae2997373e2ee1c47040854b4b759ea41ec48e4e42df0f4d73
|
3 |
+
size 28303423
|
CosyVoice-300M-25Hz-Music/cosyvoice.yaml
ADDED
@@ -0,0 +1,98 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# set random seed, so that you may reproduce your result.
|
2 |
+
__set_seed1: !apply:random.seed [1986]
|
3 |
+
__set_seed2: !apply:numpy.random.seed [1986]
|
4 |
+
__set_seed3: !apply:torch.manual_seed [1986]
|
5 |
+
__set_seed4: !apply:torch.cuda.manual_seed_all [1986]
|
6 |
+
|
7 |
+
# fixed params
|
8 |
+
sample_rate: 22050
|
9 |
+
spk_embed_dim: 192
|
10 |
+
|
11 |
+
# model params
|
12 |
+
# for all class/function included in this repo, we use !<name> or !<new> for intialization, so that user may find all corresponding class/function according to one single yaml.
|
13 |
+
# for system/third_party class/function, we do not require this.
|
14 |
+
flow: !new:cosyvoice.flow.flow.MaskedDiffWithXvec
|
15 |
+
input_size: 512
|
16 |
+
output_size: 80
|
17 |
+
spk_embed_dim: !ref <spk_embed_dim>
|
18 |
+
output_type: 'mel'
|
19 |
+
vocab_size: 5120
|
20 |
+
input_frame_rate: 41.667
|
21 |
+
only_mask_loss: True
|
22 |
+
encoder: !new:cosyvoice.transformer.encoder.ConformerEncoder
|
23 |
+
output_size: 512
|
24 |
+
attention_heads: 8
|
25 |
+
linear_units: 2048
|
26 |
+
num_blocks: 6
|
27 |
+
dropout_rate: 0.01
|
28 |
+
positional_dropout_rate: 0.01
|
29 |
+
attention_dropout_rate: 0.01
|
30 |
+
normalize_before: True
|
31 |
+
input_layer: 'linear'
|
32 |
+
pos_enc_layer_type: 'rel_pos_espnet'
|
33 |
+
selfattention_layer_type: 'rel_selfattn'
|
34 |
+
input_size: 512
|
35 |
+
use_cnn_module: False
|
36 |
+
macaron_style: False
|
37 |
+
length_regulator: !new:cosyvoice.flow.length_regulator.InterpolateRegulator
|
38 |
+
channels: 80
|
39 |
+
sampling_ratios: [1, 1, 1, 1]
|
40 |
+
decoder: !new:cosyvoice.flow.flow_matching.ConditionalCFM
|
41 |
+
in_channels: 240
|
42 |
+
n_spks: 1
|
43 |
+
spk_emb_dim: 80
|
44 |
+
cfm_params: !new:omegaconf.DictConfig
|
45 |
+
content:
|
46 |
+
sigma_min: 1e-06
|
47 |
+
solver: 'euler'
|
48 |
+
t_scheduler: 'cosine'
|
49 |
+
training_cfg_rate: 0.2
|
50 |
+
inference_cfg_rate: 0.7
|
51 |
+
reg_loss_type: 'l1'
|
52 |
+
estimator: !new:cosyvoice.flow.decoder.ConditionalDecoder
|
53 |
+
in_channels: 320
|
54 |
+
out_channels: 80
|
55 |
+
channels: [256, 256]
|
56 |
+
dropout: 0.0
|
57 |
+
attention_head_dim: 64
|
58 |
+
n_blocks: 4
|
59 |
+
num_mid_blocks: 12
|
60 |
+
num_heads: 8
|
61 |
+
act_fn: 'gelu'
|
62 |
+
|
63 |
+
hift: !new:cosyvoice.hifigan.generator.HiFTGenerator
|
64 |
+
in_channels: 80
|
65 |
+
base_channels: 512
|
66 |
+
nb_harmonics: 8
|
67 |
+
sampling_rate: !ref <sample_rate>
|
68 |
+
nsf_alpha: 0.1
|
69 |
+
nsf_sigma: 0.003
|
70 |
+
nsf_voiced_threshold: 10
|
71 |
+
upsample_rates: [8, 8]
|
72 |
+
upsample_kernel_sizes: [16, 16]
|
73 |
+
istft_params:
|
74 |
+
n_fft: 16
|
75 |
+
hop_len: 4
|
76 |
+
resblock_kernel_sizes: [3, 7, 11]
|
77 |
+
resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
|
78 |
+
source_resblock_kernel_sizes: [7, 11]
|
79 |
+
source_resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5]]
|
80 |
+
lrelu_slope: 0.1
|
81 |
+
audio_limit: 0.99
|
82 |
+
f0_predictor: !new:cosyvoice.hifigan.f0_predictor.ConvRNNF0Predictor
|
83 |
+
num_class: 1
|
84 |
+
in_channels: 80
|
85 |
+
cond_channels: 512
|
86 |
+
|
87 |
+
|
88 |
+
# processor functions
|
89 |
+
feat_extractor: !name:cosyvoice.utils.audio.mel_spectrogram # matcha.utils.audio.mel_spectrogram
|
90 |
+
n_fft: 1024
|
91 |
+
num_mels: 80
|
92 |
+
sampling_rate: !ref <sample_rate>
|
93 |
+
hop_size: 256
|
94 |
+
win_size: 1024
|
95 |
+
fmin: 0
|
96 |
+
fmax: 8000
|
97 |
+
center: False
|
98 |
+
|
CosyVoice-300M-25Hz-Music/flow.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:acdae6101bc558903b68f506a7034f5fc582f16801e6f8e4a416dce4c509bae7
|
3 |
+
size 422109962
|
CosyVoice-300M-25Hz-Music/hift.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:91e679b6ca1eff71187ffb4f3ab0444935594cdcc20a9bd12afad111ef8d6012
|
3 |
+
size 81896716
|
CosyVoice-300M-25Hz-Music/speech_tokenizer_v1.onnx
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:56285ddd4a83e883ee0cb9f8d69c1089b53a94b1f78ff7e4a0224a27eb4cb486
|
3 |
+
size 522625011
|
CosyVoice-300M-25Hz/VERSION_Vq0206Vocoder_1202
ADDED
File without changes
|
CosyVoice-300M-25Hz/campplus.onnx
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:a6ac6a63997761ae2997373e2ee1c47040854b4b759ea41ec48e4e42df0f4d73
|
3 |
+
size 28303423
|
CosyVoice-300M-25Hz/cosyvoice.yaml
ADDED
@@ -0,0 +1,98 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# set random seed, so that you may reproduce your result.
|
2 |
+
__set_seed1: !apply:random.seed [1986]
|
3 |
+
__set_seed2: !apply:numpy.random.seed [1986]
|
4 |
+
__set_seed3: !apply:torch.manual_seed [1986]
|
5 |
+
__set_seed4: !apply:torch.cuda.manual_seed_all [1986]
|
6 |
+
|
7 |
+
# fixed params
|
8 |
+
sample_rate: 22050
|
9 |
+
spk_embed_dim: 192
|
10 |
+
|
11 |
+
# model params
|
12 |
+
# for all class/function included in this repo, we use !<name> or !<new> for intialization, so that user may find all corresponding class/function according to one single yaml.
|
13 |
+
# for system/third_party class/function, we do not require this.
|
14 |
+
flow: !new:cosyvoice.flow.flow.MaskedDiffWithXvec
|
15 |
+
input_size: 512
|
16 |
+
output_size: 80
|
17 |
+
spk_embed_dim: !ref <spk_embed_dim>
|
18 |
+
output_type: 'mel'
|
19 |
+
vocab_size: 5120
|
20 |
+
input_frame_rate: 41.667
|
21 |
+
only_mask_loss: True
|
22 |
+
encoder: !new:cosyvoice.transformer.encoder.ConformerEncoder
|
23 |
+
output_size: 512
|
24 |
+
attention_heads: 8
|
25 |
+
linear_units: 2048
|
26 |
+
num_blocks: 6
|
27 |
+
dropout_rate: 0.01
|
28 |
+
positional_dropout_rate: 0.01
|
29 |
+
attention_dropout_rate: 0.01
|
30 |
+
normalize_before: True
|
31 |
+
input_layer: 'linear'
|
32 |
+
pos_enc_layer_type: 'rel_pos_espnet'
|
33 |
+
selfattention_layer_type: 'rel_selfattn'
|
34 |
+
input_size: 512
|
35 |
+
use_cnn_module: False
|
36 |
+
macaron_style: False
|
37 |
+
length_regulator: !new:cosyvoice.flow.length_regulator.InterpolateRegulator
|
38 |
+
channels: 80
|
39 |
+
sampling_ratios: [1, 1, 1, 1]
|
40 |
+
decoder: !new:cosyvoice.flow.flow_matching.ConditionalCFM
|
41 |
+
in_channels: 240
|
42 |
+
n_spks: 1
|
43 |
+
spk_emb_dim: 80
|
44 |
+
cfm_params: !new:omegaconf.DictConfig
|
45 |
+
content:
|
46 |
+
sigma_min: 1e-06
|
47 |
+
solver: 'euler'
|
48 |
+
t_scheduler: 'cosine'
|
49 |
+
training_cfg_rate: 0.2
|
50 |
+
inference_cfg_rate: 0.7
|
51 |
+
reg_loss_type: 'l1'
|
52 |
+
estimator: !new:cosyvoice.flow.decoder.ConditionalDecoder
|
53 |
+
in_channels: 320
|
54 |
+
out_channels: 80
|
55 |
+
channels: [256, 256]
|
56 |
+
dropout: 0.0
|
57 |
+
attention_head_dim: 64
|
58 |
+
n_blocks: 4
|
59 |
+
num_mid_blocks: 12
|
60 |
+
num_heads: 8
|
61 |
+
act_fn: 'gelu'
|
62 |
+
|
63 |
+
hift: !new:cosyvoice.hifigan.generator.HiFTGenerator
|
64 |
+
in_channels: 80
|
65 |
+
base_channels: 512
|
66 |
+
nb_harmonics: 8
|
67 |
+
sampling_rate: !ref <sample_rate>
|
68 |
+
nsf_alpha: 0.1
|
69 |
+
nsf_sigma: 0.003
|
70 |
+
nsf_voiced_threshold: 10
|
71 |
+
upsample_rates: [8, 8]
|
72 |
+
upsample_kernel_sizes: [16, 16]
|
73 |
+
istft_params:
|
74 |
+
n_fft: 16
|
75 |
+
hop_len: 4
|
76 |
+
resblock_kernel_sizes: [3, 7, 11]
|
77 |
+
resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
|
78 |
+
source_resblock_kernel_sizes: [7, 11]
|
79 |
+
source_resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5]]
|
80 |
+
lrelu_slope: 0.1
|
81 |
+
audio_limit: 0.99
|
82 |
+
f0_predictor: !new:cosyvoice.hifigan.f0_predictor.ConvRNNF0Predictor
|
83 |
+
num_class: 1
|
84 |
+
in_channels: 80
|
85 |
+
cond_channels: 512
|
86 |
+
|
87 |
+
|
88 |
+
# processor functions
|
89 |
+
feat_extractor: !name:cosyvoice.utils.audio.mel_spectrogram # matcha.utils.audio.mel_spectrogram
|
90 |
+
n_fft: 1024
|
91 |
+
num_mels: 80
|
92 |
+
sampling_rate: !ref <sample_rate>
|
93 |
+
hop_size: 256
|
94 |
+
win_size: 1024
|
95 |
+
fmin: 0
|
96 |
+
fmax: 8000
|
97 |
+
center: False
|
98 |
+
|
CosyVoice-300M-25Hz/flow.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:a0fefc7bc07c57cf4b8b13adec54585dd3ea4285b32f03b85fc9b7f2f34887f5
|
3 |
+
size 422117160
|
CosyVoice-300M-25Hz/hift.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:91e679b6ca1eff71187ffb4f3ab0444935594cdcc20a9bd12afad111ef8d6012
|
3 |
+
size 81896716
|
CosyVoice-300M-25Hz/speech_tokenizer_v1.onnx
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:56285ddd4a83e883ee0cb9f8d69c1089b53a94b1f78ff7e4a0224a27eb4cb486
|
3 |
+
size 522625011
|
config.json
ADDED
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"architectures": [
|
3 |
+
"StepAudioForCausalLM"
|
4 |
+
],
|
5 |
+
"auto_map": {
|
6 |
+
"AutoConfig": "configuration_stepaudio.StepAudioConfig",
|
7 |
+
"AutoModelForCausalLM": "modeling_stepaudio.StepAudioForCausalLM"
|
8 |
+
},
|
9 |
+
"model_type": "step_audio",
|
10 |
+
"bos_token_id": 1,
|
11 |
+
"pad_token_id": 0,
|
12 |
+
"eos_token_id": 3,
|
13 |
+
"hidden_size": 3072,
|
14 |
+
"intermediate_size": 8192,
|
15 |
+
"num_attention_heads": 48,
|
16 |
+
"num_attention_groups": 4,
|
17 |
+
"num_hidden_layers": 32,
|
18 |
+
"max_seq_len": 32768,
|
19 |
+
"vocab_size": 74752,
|
20 |
+
"rms_norm_eps": 1e-05,
|
21 |
+
"torch_dtype": "bfloat16"
|
22 |
+
}
|
configuration_stepaudio.py
ADDED
@@ -0,0 +1,41 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Optional, List, Any, Dict
|
2 |
+
from transformers.configuration_utils import PretrainedConfig
|
3 |
+
|
4 |
+
|
5 |
+
|
6 |
+
class StepAudioConfig(PretrainedConfig):
|
7 |
+
model_type = "step_audio"
|
8 |
+
keys_to_ignore_at_inference = ["past_key_values"]
|
9 |
+
|
10 |
+
def __init__(
|
11 |
+
self,
|
12 |
+
hidden_size: int = 5120,
|
13 |
+
intermediate_size: int = 13312,
|
14 |
+
num_attention_heads: int = 40,
|
15 |
+
num_attention_groups: int = 8,
|
16 |
+
num_hidden_layers: int = 48,
|
17 |
+
max_seq_len: int = 4096,
|
18 |
+
vocab_size: int = 65536,
|
19 |
+
rms_norm_eps: float = 1e-5,
|
20 |
+
bos_token_id: int = 1,
|
21 |
+
eos_token_id: int = 3,
|
22 |
+
pad_token_id: int = 0,
|
23 |
+
**kwargs,
|
24 |
+
) -> None:
|
25 |
+
self.hidden_size = hidden_size
|
26 |
+
self.intermediate_size = intermediate_size
|
27 |
+
self.num_attention_heads = num_attention_heads
|
28 |
+
self.num_attention_groups = num_attention_groups
|
29 |
+
self.num_hidden_layers = num_hidden_layers
|
30 |
+
self.max_seq_len = max_seq_len
|
31 |
+
self.vocab_size = vocab_size
|
32 |
+
self.rms_norm_eps = rms_norm_eps
|
33 |
+
super().__init__(
|
34 |
+
bos_token_id=bos_token_id,
|
35 |
+
pad_token_id=pad_token_id,
|
36 |
+
eos_token_id=eos_token_id,
|
37 |
+
**kwargs
|
38 |
+
)
|
39 |
+
|
40 |
+
|
41 |
+
__all__ = ["StepAudioConfig"]
|
model-00001.safetensors
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:99b74741d91dbe52844497fe5faf50cb195d28ffd51bb31777c632ae2eed3176
|
3 |
+
size 7059446656
|
model.safetensors.index.json
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
{"metadata": {"total_size": 7059412992}, "weight_map": {"model.embed_tokens.weight": "model-00001.safetensors", "model.layers.0.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.0.input_layernorm.weight": "model-00001.safetensors", "model.layers.0.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.0.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.0.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.0.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.0.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.0.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.0.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.1.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.1.input_layernorm.weight": "model-00001.safetensors", "model.layers.1.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.1.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.1.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.1.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.1.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.1.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.1.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.2.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.2.input_layernorm.weight": "model-00001.safetensors", "model.layers.2.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.2.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.2.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.2.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.2.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.2.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.2.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.3.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.3.input_layernorm.weight": "model-00001.safetensors", "model.layers.3.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.3.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.3.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.3.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.3.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.3.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.3.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.4.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.4.input_layernorm.weight": "model-00001.safetensors", "model.layers.4.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.4.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.4.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.4.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.4.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.4.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.4.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.5.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.5.input_layernorm.weight": "model-00001.safetensors", "model.layers.5.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.5.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.5.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.5.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.5.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.5.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.5.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.6.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.6.input_layernorm.weight": "model-00001.safetensors", "model.layers.6.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.6.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.6.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.6.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.6.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.6.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.6.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.7.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.7.input_layernorm.weight": "model-00001.safetensors", "model.layers.7.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.7.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.7.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.7.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.7.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.7.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.7.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.8.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.8.input_layernorm.weight": "model-00001.safetensors", "model.layers.8.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.8.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.8.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.8.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.8.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.8.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.8.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.9.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.9.input_layernorm.weight": "model-00001.safetensors", "model.layers.9.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.9.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.9.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.9.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.9.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.9.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.9.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.10.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.10.input_layernorm.weight": "model-00001.safetensors", "model.layers.10.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.10.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.10.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.10.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.10.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.10.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.10.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.11.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.11.input_layernorm.weight": "model-00001.safetensors", "model.layers.11.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.11.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.11.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.11.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.11.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.11.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.11.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.12.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.12.input_layernorm.weight": "model-00001.safetensors", "model.layers.12.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.12.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.12.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.12.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.12.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.12.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.12.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.13.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.13.input_layernorm.weight": "model-00001.safetensors", "model.layers.13.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.13.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.13.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.13.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.13.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.13.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.13.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.14.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.14.input_layernorm.weight": "model-00001.safetensors", "model.layers.14.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.14.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.14.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.14.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.14.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.14.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.14.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.15.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.15.input_layernorm.weight": "model-00001.safetensors", "model.layers.15.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.15.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.15.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.15.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.15.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.15.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.15.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.16.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.16.input_layernorm.weight": "model-00001.safetensors", "model.layers.16.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.16.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.16.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.16.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.16.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.16.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.16.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.17.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.17.input_layernorm.weight": "model-00001.safetensors", "model.layers.17.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.17.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.17.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.17.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.17.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.17.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.17.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.18.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.18.input_layernorm.weight": "model-00001.safetensors", "model.layers.18.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.18.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.18.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.18.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.18.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.18.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.18.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.19.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.19.input_layernorm.weight": "model-00001.safetensors", "model.layers.19.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.19.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.19.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.19.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.19.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.19.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.19.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.20.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.20.input_layernorm.weight": "model-00001.safetensors", "model.layers.20.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.20.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.20.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.20.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.20.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.20.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.20.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.21.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.21.input_layernorm.weight": "model-00001.safetensors", "model.layers.21.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.21.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.21.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.21.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.21.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.21.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.21.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.22.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.22.input_layernorm.weight": "model-00001.safetensors", "model.layers.22.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.22.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.22.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.22.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.22.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.22.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.22.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.23.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.23.input_layernorm.weight": "model-00001.safetensors", "model.layers.23.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.23.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.23.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.23.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.23.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.23.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.23.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.24.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.24.input_layernorm.weight": "model-00001.safetensors", "model.layers.24.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.24.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.24.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.24.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.24.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.24.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.24.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.25.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.25.input_layernorm.weight": "model-00001.safetensors", "model.layers.25.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.25.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.25.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.25.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.25.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.25.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.25.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.26.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.26.input_layernorm.weight": "model-00001.safetensors", "model.layers.26.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.26.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.26.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.26.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.26.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.26.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.26.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.27.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.27.input_layernorm.weight": "model-00001.safetensors", "model.layers.27.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.27.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.27.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.27.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.27.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.27.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.27.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.28.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.28.input_layernorm.weight": "model-00001.safetensors", "model.layers.28.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.28.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.28.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.28.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.28.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.28.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.28.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.29.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.29.input_layernorm.weight": "model-00001.safetensors", "model.layers.29.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.29.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.29.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.29.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.29.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.29.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.29.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.30.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.30.input_layernorm.weight": "model-00001.safetensors", "model.layers.30.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.30.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.30.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.30.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.30.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.30.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.30.mlp.up_proj.weight": "model-00001.safetensors", "model.layers.31.self_attn.o_proj.weight": "model-00001.safetensors", "model.layers.31.input_layernorm.weight": "model-00001.safetensors", "model.layers.31.mlp.down_proj.weight": "model-00001.safetensors", "model.layers.31.post_attention_layernorm.weight": "model-00001.safetensors", "model.layers.31.self_attn.q_proj.weight": "model-00001.safetensors", "model.layers.31.self_attn.k_proj.weight": "model-00001.safetensors", "model.layers.31.self_attn.v_proj.weight": "model-00001.safetensors", "model.layers.31.mlp.gate_proj.weight": "model-00001.safetensors", "model.layers.31.mlp.up_proj.weight": "model-00001.safetensors", "model.norm.weight": "model-00001.safetensors", "lm_head.weight": "model-00001.safetensors"}}
|
modeling_stepaudio.py
ADDED
@@ -0,0 +1,392 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import math
|
2 |
+
from typing import Optional, Tuple, Union, List
|
3 |
+
|
4 |
+
import torch
|
5 |
+
import torch.utils.checkpoint
|
6 |
+
from torch import nn
|
7 |
+
from transformers.generation import GenerationMixin
|
8 |
+
|
9 |
+
from transformers.modeling_utils import PreTrainedModel
|
10 |
+
from transformers.utils import logging
|
11 |
+
from .configuration_stepaudio import StepAudioConfig
|
12 |
+
from transformers.cache_utils import Cache, DynamicCache
|
13 |
+
from einops import rearrange
|
14 |
+
from transformers.modeling_outputs import (
|
15 |
+
BaseModelOutputWithPast,
|
16 |
+
CausalLMOutputWithPast,
|
17 |
+
)
|
18 |
+
|
19 |
+
logger = logging.get_logger(__name__)
|
20 |
+
|
21 |
+
|
22 |
+
def build_alibi_cache(block_size, n_heads, dtype, device):
|
23 |
+
# get slopes
|
24 |
+
n = 2 ** math.floor(math.log2(n_heads)) # nearest 2**n to n_heads
|
25 |
+
m0 = 2.0 ** (-8.0 / n)
|
26 |
+
# 2^(-8/n), 2^(-8*2/n), 2^(-8*3/n), ...
|
27 |
+
slopes = torch.pow(m0, torch.arange(1, n + 1))
|
28 |
+
if n < n_heads:
|
29 |
+
m1 = 2.0 ** (-4.0 / n)
|
30 |
+
# 2^(-8/(2n)), 2^(-8*3/(2n)), 2^(-8*5/(2n)), ...
|
31 |
+
mm = torch.pow(m1, torch.arange(1, 1 + 2 * (n_heads - n), 2))
|
32 |
+
slopes = torch.cat([slopes, mm])
|
33 |
+
slopes = slopes.to(device)
|
34 |
+
|
35 |
+
tril = torch.tril(torch.ones(1, 1, block_size, block_size, device=device))
|
36 |
+
|
37 |
+
bias_rows = torch.arange(block_size, device=device).view(1, -1)
|
38 |
+
bias_cols = torch.arange(block_size, device=device).view(-1, 1)
|
39 |
+
bias = -torch.sqrt(bias_cols - bias_rows)
|
40 |
+
bias = bias.view(1, block_size, block_size) * slopes.view(-1, 1, 1)
|
41 |
+
bias = bias.masked_fill(tril == 0, float("-inf"))
|
42 |
+
|
43 |
+
return bias.type(dtype)
|
44 |
+
|
45 |
+
|
46 |
+
class StepAudioRMSNorm(torch.nn.Module):
|
47 |
+
def __init__(self, hidden_size, eps=1e-5):
|
48 |
+
super().__init__()
|
49 |
+
self.weight = torch.nn.Parameter(torch.ones(hidden_size))
|
50 |
+
self.eps = eps
|
51 |
+
|
52 |
+
def forward(self, x: torch.Tensor):
|
53 |
+
var = x.float().pow(2).mean(-1, keepdim=True)
|
54 |
+
x = x * torch.rsqrt(var + self.eps).to(x.dtype)
|
55 |
+
x = x * self.weight
|
56 |
+
return x
|
57 |
+
|
58 |
+
|
59 |
+
class StepAudioAttention(torch.nn.Module):
|
60 |
+
def __init__(self, hidden_size, num_heads, num_groups, layer_idx: int):
|
61 |
+
super().__init__()
|
62 |
+
|
63 |
+
self.num_heads = num_heads
|
64 |
+
self.num_groups = num_groups
|
65 |
+
self.hidden_size = hidden_size
|
66 |
+
self.head_dim = hidden_size // num_heads
|
67 |
+
|
68 |
+
self.q_proj = torch.nn.Linear(hidden_size, hidden_size, bias=False)
|
69 |
+
self.k_proj = torch.nn.Linear(
|
70 |
+
hidden_size, num_groups * self.head_dim, bias=False
|
71 |
+
)
|
72 |
+
self.v_proj = torch.nn.Linear(
|
73 |
+
hidden_size, num_groups * self.head_dim, bias=False
|
74 |
+
)
|
75 |
+
self.o_proj = torch.nn.Linear(hidden_size, hidden_size, bias=False)
|
76 |
+
|
77 |
+
self.layer_idx = layer_idx
|
78 |
+
|
79 |
+
def forward(
|
80 |
+
self,
|
81 |
+
x: torch.Tensor,
|
82 |
+
past_key_value: Optional[Cache] = None,
|
83 |
+
attention_mask: Optional[torch.Tensor] = None,
|
84 |
+
cache_position: Optional[torch.LongTensor] = None,
|
85 |
+
):
|
86 |
+
|
87 |
+
q: torch.Tensor = self.q_proj(x)
|
88 |
+
k: torch.Tensor = self.k_proj(x)
|
89 |
+
v: torch.Tensor = self.v_proj(x)
|
90 |
+
if past_key_value is not None:
|
91 |
+
cache_kwargs = {"cache_position": cache_position}
|
92 |
+
k, v = past_key_value.update(k, v, self.layer_idx, cache_kwargs)
|
93 |
+
|
94 |
+
q = rearrange(q, "b s (h d) -> b s h d", h=self.num_heads)
|
95 |
+
k = rearrange(k, "b s (g d) -> b s g d", g=self.num_groups)
|
96 |
+
v = rearrange(v, "b s (g d) -> b s g d", g=self.num_groups)
|
97 |
+
|
98 |
+
k = k.repeat_interleave(self.num_heads // self.num_groups, dim=-2)
|
99 |
+
v = v.repeat_interleave(self.num_heads // self.num_groups, dim=-2)
|
100 |
+
|
101 |
+
attention_mask = build_alibi_cache(
|
102 |
+
k.size(1), self.num_heads, dtype=q.dtype, device=q.device
|
103 |
+
)[:, :, -q.size(1) :, :].contiguous()
|
104 |
+
|
105 |
+
q = q.transpose(1, 2)
|
106 |
+
k = k.transpose(1, 2)
|
107 |
+
v = v.transpose(1, 2)
|
108 |
+
|
109 |
+
o: torch.Tensor = torch.nn.functional.scaled_dot_product_attention(
|
110 |
+
q, k, v, attn_mask=attention_mask
|
111 |
+
)
|
112 |
+
o = o.transpose(1, 2).flatten(-2, -1)
|
113 |
+
|
114 |
+
o = self.o_proj(o)
|
115 |
+
return o
|
116 |
+
|
117 |
+
|
118 |
+
class StepAudioMLP(torch.nn.Module):
|
119 |
+
def __init__(self, hidden_size, intermediate_size):
|
120 |
+
super().__init__()
|
121 |
+
self.gate_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False)
|
122 |
+
self.up_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False)
|
123 |
+
self.down_proj = torch.nn.Linear(intermediate_size, hidden_size, bias=False)
|
124 |
+
|
125 |
+
def forward(self, x):
|
126 |
+
gate = self.gate_proj(x)
|
127 |
+
up = self.up_proj(x)
|
128 |
+
x = torch.nn.functional.silu(gate) * up
|
129 |
+
x = self.down_proj(x)
|
130 |
+
return x
|
131 |
+
|
132 |
+
|
133 |
+
class StepAudioLayer(torch.nn.Module):
|
134 |
+
def __init__(self, config: StepAudioConfig, layer_idx: int):
|
135 |
+
super().__init__()
|
136 |
+
self.layer_idx = layer_idx
|
137 |
+
self.self_attn = StepAudioAttention(
|
138 |
+
hidden_size=config.hidden_size,
|
139 |
+
num_heads=config.num_attention_heads,
|
140 |
+
num_groups=config.num_attention_groups,
|
141 |
+
layer_idx=layer_idx,
|
142 |
+
)
|
143 |
+
self.mlp = StepAudioMLP(
|
144 |
+
hidden_size=config.hidden_size,
|
145 |
+
intermediate_size=config.intermediate_size,
|
146 |
+
)
|
147 |
+
self.input_layernorm = StepAudioRMSNorm(
|
148 |
+
hidden_size=config.hidden_size, eps=config.rms_norm_eps
|
149 |
+
)
|
150 |
+
self.post_attention_layernorm = StepAudioRMSNorm(
|
151 |
+
hidden_size=config.hidden_size, eps=config.rms_norm_eps
|
152 |
+
)
|
153 |
+
|
154 |
+
def forward(
|
155 |
+
self,
|
156 |
+
x,
|
157 |
+
attention_mask: Optional[torch.Tensor] = None,
|
158 |
+
past_key_value: Optional[Cache] = None,
|
159 |
+
cache_position: Optional[torch.LongTensor] = None,
|
160 |
+
):
|
161 |
+
def f(x):
|
162 |
+
x = self.input_layernorm(x)
|
163 |
+
x = self.self_attn(x, past_key_value, attention_mask, cache_position)
|
164 |
+
return x
|
165 |
+
|
166 |
+
x = x + f(x)
|
167 |
+
|
168 |
+
def f(x):
|
169 |
+
x = self.post_attention_layernorm(x)
|
170 |
+
x = self.mlp(x)
|
171 |
+
return x
|
172 |
+
|
173 |
+
x = x + f(x)
|
174 |
+
|
175 |
+
return x
|
176 |
+
|
177 |
+
|
178 |
+
class StepAudioPreTrainedModel(PreTrainedModel):
|
179 |
+
config_class = StepAudioConfig
|
180 |
+
base_model_prefix = "model"
|
181 |
+
supports_gradient_checkpointing = True
|
182 |
+
_no_split_modules = ["StepAudioLayer"]
|
183 |
+
_skip_keys_device_placement = ["past_key_values"]
|
184 |
+
_supports_cache_class = True
|
185 |
+
_supports_static_cache = True
|
186 |
+
|
187 |
+
def _init_weights(self, module):
|
188 |
+
std = self.config.initializer_range
|
189 |
+
if isinstance(module, nn.Linear):
|
190 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
191 |
+
if module.bias is not None:
|
192 |
+
module.bias.data.zero_()
|
193 |
+
elif isinstance(module, nn.Embedding):
|
194 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
195 |
+
if module.padding_idx is not None:
|
196 |
+
module.weight.data[module.padding_idx].zero_()
|
197 |
+
|
198 |
+
|
199 |
+
class StepAudioModel(StepAudioPreTrainedModel):
|
200 |
+
"""
|
201 |
+
Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]
|
202 |
+
|
203 |
+
Args:
|
204 |
+
config: StepAudioConfig
|
205 |
+
"""
|
206 |
+
|
207 |
+
def __init__(self, config: StepAudioConfig):
|
208 |
+
super().__init__(config)
|
209 |
+
self.config = config
|
210 |
+
self.embed_tokens = torch.nn.Embedding(config.vocab_size, config.hidden_size)
|
211 |
+
|
212 |
+
self.layers = torch.nn.Sequential(
|
213 |
+
*[
|
214 |
+
StepAudioLayer(config, layer_idx)
|
215 |
+
for layer_idx in range(config.num_hidden_layers)
|
216 |
+
]
|
217 |
+
)
|
218 |
+
|
219 |
+
self.norm = StepAudioRMSNorm(
|
220 |
+
hidden_size=config.hidden_size, eps=config.rms_norm_eps
|
221 |
+
)
|
222 |
+
|
223 |
+
# Initialize weights and apply final processing
|
224 |
+
self.post_init()
|
225 |
+
|
226 |
+
def get_input_embeddings(self):
|
227 |
+
return self.embed_tokens
|
228 |
+
|
229 |
+
def set_input_embeddings(self, value):
|
230 |
+
self.embed_tokens = value
|
231 |
+
|
232 |
+
def forward(
|
233 |
+
self,
|
234 |
+
input_ids: torch.LongTensor = None,
|
235 |
+
attention_mask: Optional[torch.Tensor] = None,
|
236 |
+
past_key_values: Optional[Cache] = None,
|
237 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
238 |
+
use_cache: Optional[bool] = None,
|
239 |
+
output_attentions: Optional[bool] = None,
|
240 |
+
output_hidden_states: Optional[bool] = None,
|
241 |
+
return_dict: Optional[bool] = None,
|
242 |
+
cache_position: Optional[torch.LongTensor] = None,
|
243 |
+
) -> Union[Tuple, BaseModelOutputWithPast]:
|
244 |
+
output_attentions = False
|
245 |
+
output_hidden_states = False
|
246 |
+
|
247 |
+
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
248 |
+
return_dict = (
|
249 |
+
return_dict if return_dict is not None else self.config.use_return_dict
|
250 |
+
)
|
251 |
+
|
252 |
+
if (input_ids is None) ^ (inputs_embeds is not None):
|
253 |
+
raise ValueError(
|
254 |
+
"You must specify exactly one of input_ids or inputs_embeds"
|
255 |
+
)
|
256 |
+
|
257 |
+
if inputs_embeds is None:
|
258 |
+
inputs_embeds = self.embed_tokens(input_ids)
|
259 |
+
|
260 |
+
if use_cache and past_key_values is None:
|
261 |
+
past_key_values = DynamicCache()
|
262 |
+
|
263 |
+
if cache_position is None:
|
264 |
+
past_seen_tokens = (
|
265 |
+
past_key_values.get_seq_length() if past_key_values is not None else 0
|
266 |
+
)
|
267 |
+
cache_position = torch.arange(
|
268 |
+
past_seen_tokens,
|
269 |
+
past_seen_tokens + inputs_embeds.shape[1],
|
270 |
+
device=inputs_embeds.device,
|
271 |
+
)
|
272 |
+
|
273 |
+
causal_mask = attention_mask
|
274 |
+
|
275 |
+
hidden_states = inputs_embeds
|
276 |
+
|
277 |
+
for decoder_layer in self.layers[: self.config.num_hidden_layers]:
|
278 |
+
layer_outputs = decoder_layer(
|
279 |
+
hidden_states,
|
280 |
+
attention_mask=causal_mask,
|
281 |
+
past_key_value=past_key_values,
|
282 |
+
cache_position=cache_position,
|
283 |
+
)
|
284 |
+
|
285 |
+
hidden_states = layer_outputs
|
286 |
+
|
287 |
+
hidden_states = self.norm(hidden_states)
|
288 |
+
|
289 |
+
output = BaseModelOutputWithPast(
|
290 |
+
last_hidden_state=hidden_states,
|
291 |
+
past_key_values=past_key_values if use_cache else None,
|
292 |
+
hidden_states=hidden_states,
|
293 |
+
attentions=None,
|
294 |
+
)
|
295 |
+
return output if return_dict else output.to_tuple()
|
296 |
+
|
297 |
+
|
298 |
+
class StepAudioForCausalLM(StepAudioPreTrainedModel, GenerationMixin):
|
299 |
+
_tied_weights_keys = ["lm_head.weight"]
|
300 |
+
|
301 |
+
def __init__(self, config):
|
302 |
+
super().__init__(config)
|
303 |
+
self.model = StepAudioModel(config)
|
304 |
+
self.vocab_size = config.vocab_size
|
305 |
+
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
306 |
+
|
307 |
+
# Initialize weights and apply final processing
|
308 |
+
self.post_init()
|
309 |
+
|
310 |
+
def get_input_embeddings(self):
|
311 |
+
return self.model.embed_tokens
|
312 |
+
|
313 |
+
def set_input_embeddings(self, value):
|
314 |
+
self.model.embed_tokens = value
|
315 |
+
|
316 |
+
# def get_output_embeddings(self):
|
317 |
+
# return self.lm_head
|
318 |
+
|
319 |
+
# def set_output_embeddings(self, new_embeddings):
|
320 |
+
# self.lm_head = new_embeddings
|
321 |
+
|
322 |
+
def set_decoder(self, decoder):
|
323 |
+
self.model = decoder
|
324 |
+
|
325 |
+
def get_decoder(self):
|
326 |
+
return self.model
|
327 |
+
|
328 |
+
def forward(
|
329 |
+
self,
|
330 |
+
input_ids: torch.LongTensor = None,
|
331 |
+
attention_mask: Optional[torch.Tensor] = None,
|
332 |
+
position_ids: Optional[torch.LongTensor] = None,
|
333 |
+
past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
|
334 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
335 |
+
labels: Optional[torch.LongTensor] = None,
|
336 |
+
use_cache: Optional[bool] = None,
|
337 |
+
output_attentions: Optional[bool] = None,
|
338 |
+
output_hidden_states: Optional[bool] = None,
|
339 |
+
return_dict: Optional[bool] = None,
|
340 |
+
cache_position: Optional[torch.LongTensor] = None,
|
341 |
+
num_logits_to_keep: int = 0,
|
342 |
+
) -> Union[Tuple, CausalLMOutputWithPast]:
|
343 |
+
# output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
344 |
+
output_attentions = False
|
345 |
+
output_hidden_states = False
|
346 |
+
# output_hidden_states = (
|
347 |
+
# output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
348 |
+
# )
|
349 |
+
return_dict = (
|
350 |
+
return_dict if return_dict is not None else self.config.use_return_dict
|
351 |
+
)
|
352 |
+
|
353 |
+
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
354 |
+
outputs = self.model(
|
355 |
+
input_ids=input_ids,
|
356 |
+
attention_mask=attention_mask,
|
357 |
+
past_key_values=past_key_values,
|
358 |
+
inputs_embeds=inputs_embeds,
|
359 |
+
use_cache=use_cache,
|
360 |
+
output_attentions=output_attentions,
|
361 |
+
output_hidden_states=output_hidden_states,
|
362 |
+
return_dict=return_dict,
|
363 |
+
cache_position=cache_position,
|
364 |
+
)
|
365 |
+
|
366 |
+
hidden_states = outputs[0]
|
367 |
+
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
368 |
+
|
369 |
+
logits = self.lm_head(hidden_states)
|
370 |
+
|
371 |
+
# logits = torch.matmul(hidden_states, lm_stat)
|
372 |
+
|
373 |
+
loss = None
|
374 |
+
if labels is not None:
|
375 |
+
loss = self.loss_function(
|
376 |
+
logits=logits,
|
377 |
+
labels=labels,
|
378 |
+
vocab_size=self.config.vocab_size,
|
379 |
+
**kwargs
|
380 |
+
)
|
381 |
+
|
382 |
+
if not return_dict:
|
383 |
+
output = (logits,) + outputs[1:]
|
384 |
+
return (loss,) + output if loss is not None else output
|
385 |
+
|
386 |
+
return CausalLMOutputWithPast(
|
387 |
+
loss=loss,
|
388 |
+
logits=logits,
|
389 |
+
past_key_values=outputs.past_key_values,
|
390 |
+
hidden_states=outputs.hidden_states,
|
391 |
+
attentions=outputs.attentions,
|
392 |
+
)
|
tokenizer.model
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:25e122d9205d035033a9994c4d46a6a1b467a938654e4178fc0e5f4f5d610674
|
3 |
+
size 1264044
|
tokenizer_config.json
ADDED
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"bos_token": "<s>",
|
3 |
+
"clean_up_tokenization_spaces": false,
|
4 |
+
"eos_token": "</s>",
|
5 |
+
"legacy": false,
|
6 |
+
"model_max_length": 65536,
|
7 |
+
"pad_token": "<unk>",
|
8 |
+
"padding_side": "left",
|
9 |
+
"sp_model_kwargs": {},
|
10 |
+
"tokenizer_class": "LlamaTokenizer",
|
11 |
+
"unk_token": "<unk>",
|
12 |
+
"use_default_system_prompt": false
|
13 |
+
}
|
14 |
+
|