buyun commited on
Commit
dd81d75
·
verified ·
1 Parent(s): 4f3efe4
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
+