Duplicate from stepfun-ai/Step-Audio-EditX

Co-authored-by: chao yan <yanchaomars@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,38 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/architechture.png filter=lfs diff=lfs merge=lfs -text
+assets/emotion-eval.png filter=lfs diff=lfs merge=lfs -text
+assets/test.wav filter=lfs diff=lfs merge=lfs -text

CosyVoice-300M-25Hz/FLOW_VERSION

@@ -0,0 +1,2 @@
+/mnt/wby-jfs/models/train/flow_matching/flow_v2_1node_vq0206_dit_v8_fullattn_exp0227_sft_exp0408_stepaudio_sft_exp0616/model_epoch_5_whole.pt
+fae53942e60310eb172b170396202069

CosyVoice-300M-25Hz/campplus.onnx

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

CosyVoice-300M-25Hz/cosyvoice.yaml

@@ -0,0 +1,72 @@
+mel_conf:
+    num_mels: 80
+    n_fft: 1920
+    hop_size: 480
+    win_size: 1920
+    sampling_rate: 24000
+    fmin: 0
+    fmax: 8000
+    
+
+flow: !new:stepvocoder.cosyvoice2.flow.flow.CausalMaskedDiffWithXvec
+    input_size: 512
+    output_size: 80
+    spk_embed_dim: 192
+    output_type: 'mel'
+    vocab_size: 5121    # 1024(vq02) + 4096(vq06) + 1(vq02-pad)
+    input_embedding: !new:stepvocoder.cosyvoice2.embedding.dual_codebook.DualCodebookEmbedding
+      vocab_size: 5121  # 1024(vq02) + 4096(vq06) + 1(vq02-pad)
+      input_size: 512
+    encoder: !new:stepvocoder.cosyvoice2.transformer.upsample_encoder_v2.UpsampleConformerEncoderV2
+        input_size: 512
+        output_size: 512
+        input_layer: 'linear'
+        pre_lookahead_len: 3
+        num_blocks: 6
+        num_up_blocks: 4
+        up_stride: 2
+        up_scale_factor: 2
+        attention_heads: 8
+        pos_enc_layer_type: 'rel_pos_espnet'
+        selfattention_layer_type: 'rel_selfattn'
+        key_bias: true
+        linear_units: 2048
+        dropout_rate: 0.1
+        positional_dropout_rate: 0.1
+        attention_dropout_rate: 0.1
+        normalize_before: True
+    decoder: !new:stepvocoder.cosyvoice2.flow.flow_matching.CausalConditionalCFM
+        inference_cfg_rate: 0.7
+        estimator: !new:stepvocoder.cosyvoice2.flow.decoder_dit.DiT
+            in_channels: 320
+            out_channels: 80
+            mlp_ratio: 4.0
+            depth: 16
+            num_heads: 8
+            head_dim: 64
+            hidden_size: 512
+
+
+hift: !new:stepvocoder.cosyvoice2.hifigan.generator.HiFTGenerator
+    in_channels: 80
+    base_channels: 512
+    nb_harmonics: 8
+    sampling_rate: 24000
+    nsf_alpha: 0.1
+    nsf_sigma: 0.003
+    nsf_voiced_threshold: 10
+    upsample_rates: [8, 5, 3]
+    upsample_kernel_sizes: [16, 11, 7]
+    istft_params:
+        n_fft: 16
+        hop_len: 4
+    resblock_kernel_sizes: [3, 7, 11]
+    resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+    source_resblock_kernel_sizes: [7, 7, 11]
+    source_resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+    lrelu_slope: 0.1
+    audio_limit: 0.99
+    f0_predictor: !new:stepvocoder.cosyvoice2.hifigan.f0_predictor.ConvRNNF0Predictor
+        num_class: 1
+        in_channels: 80
+        cond_channels: 512

CosyVoice-300M-25Hz/flow.pt

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

CosyVoice-300M-25Hz/hift.pt

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

CosyVoice-300M-25Hz/speech_tokenizer_v1.onnx

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

README.md

@@ -0,0 +1,911 @@
+# Step-Audio-EditX
+<p align="center">
+  <img src="assets/logo.png"  height=100>
+</p>
+
+<div align="center">
+    <a href="https://stepaudiollm.github.io/step-audio-editx/"><img src="https://img.shields.io/static/v1?label=Demo%20Page&message=Web&color=green"></a> &ensp;
+  <a href="https://arxiv.org/abs/2511.03601"><img src="https://img.shields.io/static/v1?label=Tech%20Report&message=Arxiv&color=red"></a> &ensp;
+  <a href="https://huggingface.co/stepfun-ai/Step-Audio-EditX"><img src="https://img.shields.io/static/v1?label=Step-Audio-EditX&message=HuggingFace&color=yellow"></a> &ensp;
+    <a href="https://modelscope.cn/models/stepfun-ai/Step-Audio-EditX"><img src="https://img.shields.io/static/v1?label=Step-Audio-EditX&message=ModelScope&color=blue"></a> &ensp;
+  <a href="https://huggingface.co/spaces/stepfun-ai/Step-Audio-EditX"><img src="https://img.shields.io/static/v1?label=Space%20Playground&message=HuggingFace&color=yellow"></a> &ensp;
+</div>
+
+## 🔥🔥🔥 News!!！
+* Jan 23, 2026: 🌟 Training and inference for vLLM are now supported. Thanks to the vLLM team!
+* Jan 23, 2026: 💻 We release the GRPO training code.
+* Jan 23, 2026: 🧩 New Model Release: Now supporting more paralinguistic tags.
+* Nov 28, 2025: 🚀 New Model Release: Now supporting **`Japanese`** and **`Korean`** languages.
+* Nov 23, 2025: 📊 [Step-Audio-Edit-Benchmark](https://github.com/stepfun-ai/Step-Audio-Edit-Benchmark) Released!
+* Nov 19, 2025: ⚙️ We release a **new version** of our model, which **supports polyphonic pronunciation control** and improves the performance of emotion, speaking style, and paralinguistic editing.
+* Nov 12, 2025: 📦 We release the **optimized inference code** and **model weights** of **Step-Audio-EditX** ([HuggingFace](https://huggingface.co/stepfun-ai/Step-Audio-EditX);  [ModelScope](https://modelscope.cn/models/stepfun-ai/Step-Audio-EditX)) and **Step-Audio-Tokenizer**([HuggingFace](https://huggingface.co/stepfun-ai/Step-Audio-Tokenizer);  [ModelScope](https://modelscope.cn/models/stepfun-ai/Step-Audio-Tokenizer))
+* Nov 07, 2025: ✨ [Demo Page](https://stepaudiollm.github.io/step-audio-editx/) ; 🎮  [HF Space Playground](https://huggingface.co/spaces/stepfun-ai/Step-Audio-EditX)
+* Nov 06, 2025: 👋 We release the technical report of [Step-Audio-EditX](https://arxiv.org/abs/2511.03601).
+
+## Introduction
+We are open-sourcing Step-Audio-EditX, a powerful **3B-parameter** LLM-based **Reinforcement Learning** audio model specialized in expressive and iterative audio editing. It excels at editing emotion, speaking style, and paralinguistics, and also features robust zero-shot text-to-speech (TTS) capabilities. 
+
+## 📑 Open-source Plan
+- [x] Inference Code
+- [x] Online demo (Gradio)
+- [x] Step-Audio-Edit-Benchmark
+- [x] Model Checkpoints
+  - [x] Step-Audio-Tokenizer
+  - [x] Step-Audio-EditX
+  - [x] Step-Audio-EditX-Int4
+- [ ] Training Code
+  - [x] GRPO training
+  - [ ] SFT training
+  - [ ] PPO training
+- [ ] ⏳ Feature Support Plan
+  - [ ] Editing
+    - [x] Polyphone pronunciation control
+    - [x] More paralinguistic tags ([Cough, Crying, Stress, etc.])
+    - [ ] Filler word removal
+  - [ ] Other Languages
+    - [x] Japanese, Korean
+    - [ ] Arabic, French, Russian, Spanish, etc.
+
+## Online demonstration
+
+### StepFun Audio Studio
+
+- Both Step-Audio-EditX are available in our [StepFun Audio Studio](https://www.stepfun.com/studio/audio).
+- You will need an API key from the [StepFun Open Platform](https://platform.stepfun.com/).
+
+## WeChat group
+
+You can scan the following QR code to join our WeChat group for communication and discussion.
+<div align="center">
+  <img src="https://cdn-uploads.huggingface.co/production/uploads/66518fd07d8cb2629a514c18/DRdnp1SN-yxhlNOfy26mE.jpeg" width="200" alt="QR code">
+</div>
+
+## Features
+- **Zero-Shot TTS**
+  - Excellent zero-shot TTS cloning for Mandarin, English, Sichuanese, and Cantonese.
+  - To use dialect or other languages, just add a **`[Sichuanese]`** / **`[Cantonese]`** / **`[Japanese]`** / **`[Korean]`** tag before your text.
+  - 🔥 Polyphone pronunciation control, all you need to do is replace the polyphonic characters with pinyin.
+    - **[我也想过过过儿过过的生活]** -> **[我也想guo4guo4guo1儿guo4guo4的生活]**
+ 
+    
+- **Emotion and Speaking Style Editing**
+  - Remarkably effective iterative control over emotions and styles, supporting **dozens** of options for editing.
+    - Emotion Editing : [ *Angry*, *Happy*, *Sad*, *Excited*, *Fearful*, *Surprised*, *Disgusted*, etc. ]
+    - Speaking Style Editing: [ *Act_coy*, *Older*, *Child*, *Whisper*, *Serious*, *Generous*, *Exaggerated*, etc.]
+    - Editing with more emotion and more speaking styles is on the way. **Get Ready!** 🚀
+    
+
+- **Paralinguistic Editing**
+  -  Precise control over 10 types of paralinguistic features for more natural, human-like, and expressive synthetic audio.
+  - Supporting Tags:
+    - [ *Breathing*, *Laughter*, *Surprise-oh*, *Confirmation-en*, *Uhm*, *Surprise-ah*, *Surprise-wa*, *Sigh*, *Question-ei*, *Dissatisfaction-hnn* ]
+
+- **Available Tags**
+<table>
+  <tr>
+    <td rowspan="8" style="vertical-align: middle; text-align:center;" align="center">emotion</td>
+    <td align="center"><b>happy</b></td>
+    <td align="center">Expressing happiness</td>
+    <td align="center"><b>angry</b></td>
+    <td align="center">Expressing anger</td>
+  </tr>
+  <tr>
+    <td align="center"><b>sad</b></td>
+    <td align="center">Expressing sadness</td>
+    <td align="center"><b>fear</b></td>
+    <td align="center">Expressing fear</td>
+  </tr>
+  <tr>
+    <td align="center"><b>surprised</b></td>
+    <td align="center">Expressing surprise</td>
+    <td align="center"><b>confusion</b></td>
+    <td align="center">Expressing confusion</td>
+  </tr>
+  <tr>
+    <td align="center"><b>empathy</b></td>
+    <td align="center">Expressing empathy and understanding</td>
+    <td align="center"><b>embarrass</b></td>
+    <td align="center">Expressing embarrassment</td>
+  </tr>
+  <tr>
+    <td align="center"><b>excited</b></td>
+    <td align="center">Expressing excitement and enthusiasm</td>
+    <td align="center"><b>depressed</b></td>
+    <td align="center">Expressing a depressed or discouraged mood</td>
+  </tr>
+  <tr>
+    <td align="center"><b>admiration</b></td>
+    <td align="center">Expressing admiration or respect</td>
+    <td align="center"><b>coldness</b></td>
+    <td align="center">Expressing coldness and indifference</td>
+  </tr>
+  <tr>
+    <td align="center"><b>disgusted</b></td>
+    <td align="center">Expressing disgust or aversion</td>
+    <td align="center"><b>humour</b></td>
+    <td align="center">Expressing humor or playfulness</td>
+  </tr>
+  <tr>
+  </tr>
+  <tr>
+    <td rowspan="17" style="vertical-align: middle; text-align:center;" align="center">speaking style</td>
+    <td align="center"><b>serious</b></td>
+    <td align="center">Speaking in a serious or solemn manner</td>
+    <td align="center"><b>arrogant</b></td>
+    <td align="center">Speaking in an arrogant manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>child</b></td>
+    <td align="center">Speaking in a childlike manner</td>
+    <td align="center"><b>older</b></td>
+    <td align="center">Speaking in an elderly-sounding manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>girl</b></td>
+    <td align="center">Speaking in a light, youthful feminine manner</td>
+    <td align="center"><b>pure</b></td>
+    <td align="center">Speaking in a pure, innocent manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>sister</b></td>
+    <td align="center">Speaking in a mature, confident feminine manner</td>
+    <td align="center"><b>sweet</b></td>
+    <td align="center">Speaking in a sweet, lovely manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>exaggerated</b></td>
+    <td align="center">Speaking in an exaggerated, dramatic manner</td>
+    <td align="center"><b>ethereal</b></td>
+    <td align="center">Speaking in a soft, airy, dreamy manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>whisper</b></td>
+    <td align="center">Speaking in a whispering, very soft manner</td>
+    <td align="center"><b>generous</b></td>
+    <td align="center">Speaking in a hearty, outgoing, and straight-talking manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>recite</b></td>
+    <td align="center">Speaking in a clear, well-paced, poetry-reading manner</td>
+    <td align="center"><b>act_coy</b></td>
+    <td align="center">Speaking in a sweet, playful, and endearing manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>warm</b></td>
+    <td align="center">Speaking in a warm, friendly manner</td>
+    <td align="center"><b>shy</b></td>
+    <td align="center">Speaking in a shy, timid manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>comfort</b></td>
+    <td align="center">Speaking in a comforting, reassuring manner</td>
+    <td align="center"><b>authority</b></td>
+    <td align="center">Speaking in an authoritative, commanding manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>chat</b></td>
+    <td align="center">Speaking in a casual, conversational manner</td>
+    <td align="center"><b>radio</b></td>
+    <td align="center">Speaking in a radio-broadcast manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>soulful</b></td>
+    <td align="center">Speaking in a heartfelt, deeply emotional manner</td>
+    <td align="center"><b>gentle</b></td>
+    <td align="center">Speaking in a gentle, soft manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>story</b></td>
+    <td align="center">Speaking in a narrative, audiobook-style manner</td>
+    <td align="center"><b>vivid</b></td>
+    <td align="center">Speaking in a lively, expressive manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>program</b></td>
+    <td align="center">Speaking in a show-host/presenter manner</td>
+    <td align="center"><b>news</b></td>
+    <td align="center">Speaking in a news broadcasting manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>advertising</b></td>
+    <td align="center">Speaking in a polished, high-end commercial voiceover manner</td>
+    <td align="center"><b>roar</b></td>
+    <td align="center">Speaking in a loud, deep, roaring manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>murmur</b></td>
+    <td align="center">Speaking in a quiet, low manner</td>
+    <td align="center"><b>shout</b></td>
+    <td align="center">Speaking in a loud, sharp, shouting manner</td>
+  </tr>
+  <tr>
+    <td align="center"><b>deeply</b></td>
+    <td align="center">Speaking in a deep and low-pitched tone</td>
+    <td align="center"><b>loudly</b></td>
+    <td align="center">Speaking in a loud and high-pitched tone</td>
+  </tr>
+  <tr>
+  </tr>
+  <tr>
+  </tr>
+  <tr>
+  <td rowspan="11" style="vertical-align: middle; text-align:center;" align="center">paralinguistic</td>
+    <td align="center"><b>[sigh]</b></td>
+    <td align="center">Sighing sound</td>
+    <td align="center"><b>[inhale]</b></td>
+    <td align="center">Inhaling sound</td>
+  </tr>
+
+  <tr>
+    <td align="center"><b>[laugh]</b></td>
+    <td align="center">Laughter sound</td>
+    <td align="center"><b>[chuckle]</b></td>
+    <td align="center">Chuckling sound</td>
+  </tr>
+
+  <tr>
+    <td align="center"><b>[exhale]</b></td>
+    <td align="center">Exhaling sound</td>
+    <td align="center"><b>[clears throat]</b></td>
+    <td align="center">Throat clearing sound</td>
+  </tr>
+
+  <tr>
+    <td align="center"><b>[snort]</b></td>
+    <td align="center">Snorting sound</td>
+    <td align="center"><b>[giggle]</b></td>
+    <td align="center">Giggling sound</td>
+  </tr>
+
+  <tr>
+    <td align="center"><b>[cough]</b></td>
+    <td align="center">Coughing sound</td>
+    <td align="center"><b>[breath]</b></td>
+    <td align="center">Breathing sound</td>
+  </tr>
+
+  <tr>
+    <td align="center"><b>[uhm]</b></td>
+    <td align="center">Hesitation sound: "Uhm"</td>
+    <td align="center"><b>[Confirmation-en]</b></td>
+    <td align="center">Confirming: "En"</td>
+  </tr>
+
+  <tr>
+    <td align="center"><b>[Surprise-oh]</b></td>
+    <td align="center">Expressing surprise: "Oh"</td>
+    <td align="center"><b>[Surprise-ah]</b></td>
+    <td align="center">Expressing surprise: "Ah"</td>
+  </tr>
+
+  <tr>
+    <td align="center"><b>[Surprise-wa]</b></td>
+    <td align="center">Expressing surprise: "Wa"</td>
+    <td align="center"><b>[Surprise-yo]</b></td>
+    <td align="center">Expressing surprise: "Yo"</td>
+  </tr>
+
+  <tr>
+    <td align="center"><b>[Dissatisfaction-hnn]</b></td>
+    <td align="center">Dissatisfied sound: "Hnn"</td>
+    <td align="center"><b>[Question-ei]</b></td>
+    <td align="center">Questioning: "Ei"</td>
+  </tr>
+
+  <tr>
+    <td align="center"><b>[Question-ah]</b></td>
+    <td align="center">Questioning: "Ah"</td>
+    <td align="center"><b>[Question-en]</b></td>
+    <td align="center">Questioning: "En"</td>
+  </tr>
+
+  <tr>
+    <td align="center"><b>[Question-yi]</b></td>
+    <td align="center">Questioning: "Yi"</td>
+    <td align="center"><b>[Question-oh]</b></td>
+    <td align="center">Questioning: "Oh"</td>
+  </tr>
+</table>
+ 
+## Feature Requests & Wishlist
+💡 We welcome all ideas for new features! If you'd like to see a feature added to the project, please start a discussion in our [Discussions](https://github.com/stepfun-ai/Step-Audio-EditX/discussions) section.
+
+We'll be collecting community feedback here and will incorporate popular suggestions into our future development plans. Thank you for your contribution!
+
+## Demos
+
+<table>
+  <tr>
+    <th style="vertical-align : middle;text-align: center">Task</th>
+    <th style="vertical-align : middle;text-align: center">Text</th>
+    <th style="vertical-align : middle;text-align: center">Source</th>
+    <th style="vertical-align : middle;text-align: center">Edited</th>
+  </tr>
+
+  <tr>
+    <td align="center"> Emotion-Fear</td>
+    <td align="center"> 我总觉得，有人在跟着我，我能听到奇怪的脚步声。</td>
+    <td align="center">
+
+  [fear_zh_female_prompt.webm](https://github.com/user-attachments/assets/a088c059-032c-423f-81d6-3816ba347ff5) 
+  </td>
+    <td align="center">
+      
+  [fear_zh_female_output.webm](https://github.com/user-attachments/assets/917494ac-5913-4949-8022-46cf55ca05dd)
+  </td>
+  </tr>
+
+
+  <tr>
+    <td align="center"> Style-Whisper</td>
+    <td align="center"> 比如在工作间隙，做一些简单的伸展运动，放松一下身体，这样，会让你更有精力。</td>
+    <td align="center">
+      
+  [whisper_prompt.webm](https://github.com/user-attachments/assets/ed9e22f1-1bac-417b-913a-5f1db31f35c9)
+  </td>
+    <td align="center">
+      
+  [whisper_output.webm](https://github.com/user-attachments/assets/e0501050-40db-4d45-b380-8bcc309f0b5f)
+  </td>
+  </tr>
+
+  <tr>
+    <td align="center"> Style-Act_coy</td>
+    <td align="center"> 我今天想喝奶茶，可是不知道喝什么口味，你帮我选一下嘛，你选的都好喝～</td>
+    <td align="center">
+
+  [act_coy_prompt.webm](https://github.com/user-attachments/assets/74d60625-5b3c-4f45-becb-0d3fe7cc4b3f)
+  </td>
+    <td align="center"> 
+
+  [act_coy_output.webm](https://github.com/user-attachments/assets/b2f74577-56c2-4997-afd6-6bf47d15ea51)
+  </td>
+  </tr>
+
+
+  <tr>
+    <td align="center"> Paralinguistics</td>
+    <td align="center"> 你这次又忘记带钥匙了 [Dissatisfaction-hnn]，真是拿你没办法。</td>
+    <td align="center">
+      
+  [paralingustic_prompt.webm](https://github.com/user-attachments/assets/21e831a3-8110-4c64-a157-60e0cf6735f0)
+  </td>
+    <td align="center">
+      
+  [paralingustic_output.webm](https://github.com/user-attachments/assets/a82f5a40-c6a3-409b-bbe6-271180b20d7b)
+  </td>
+  </tr>
+
+
+  <tr>
+    <td align="center"> Denoising</td>
+    <td align="center"> Such legislation was clarified and extended from time to time thereafter. No, the man was not drunk, he wondered how we got tied up with this stranger. Suddenly, my reflexes had gone. It's healthier to cook without sugar.</td>
+    <td align="center">
+      
+  [denoising_prompt.webm](https://github.com/user-attachments/assets/70464bf4-ebde-44a3-b2a6-8c292333319b)
+  </td>
+    <td align="center">
+      
+  [denoising_output.webm](https://github.com/user-attachments/assets/7cd0ae8d-1bf0-40fc-9bcd-f419bd4b2d21)
+  </td>
+  </tr>
+
+  <tr>
+    <td align="center"> Speed-Faster</td>
+    <td align="center"> 上次你说鞋子有点磨脚，我给你买了一双软软的鞋垫。</td>
+    <td align="center">
+      
+  [speed_faster_prompt.webm](https://github.com/user-attachments/assets/db46609e-1b98-48d8-99c8-e166cfdfc6e3)
+  </td>
+    <td align="center">
+      
+  [speed_faster_output.webm](https://github.com/user-attachments/assets/0fbc14ca-dd4a-4362-aadc-afe0629f4c9f)
+  </td>
+  </tr>
+  
+</table>
+
+
+For more examples, see [demo page](https://stepaudiollm.github.io/step-audio-editx/).
+
+## Model Download
+
+| Models   | 🤗 Hugging Face | ModelScope |
+|-------|-------|-------|
+| Step-Audio-EditX | [stepfun-ai/Step-Audio-EditX](https://huggingface.co/stepfun-ai/Step-Audio-EditX) | [stepfun-ai/Step-Audio-EditX](https://modelscope.cn/models/stepfun-ai/Step-Audio-EditX) |
+| Step-Audio-EditX | [stepfun-ai/Step-Audio-EditX-AWQ-4bit](https://huggingface.co/stepfun-ai/Step-Audio-EditX-AWQ-4bit) | [stepfun-ai/Step-Audio-EditX-AWQ-4bit](https://modelscope.cn/models/stepfun-ai/Step-Audio-EditX-AWQ-4bit) |
+| Step-Audio-Tokenizer | [stepfun-ai/Step-Audio-Tokenizer](https://huggingface.co/stepfun-ai/Step-Audio-Tokenizer) | [stepfun-ai/Step-Audio-Tokenizer](https://modelscope.cn/models/stepfun-ai/Step-Audio-Tokenizer) |
+
+
+## Model Usage
+### 📜 Requirements
+The following table shows the requirements for running Step-Audio-EditX model (batch size = 1):
+
+|     Model    | Parameters |  Setting<br/>(sample frequency) | GPU Optimal Memory  |
+|------------|------------|--------------------------------|----------------|
+| Step-Audio-EditX   | 3B|         41.6Hz          |       12 GB        |
+
+* An NVIDIA GPU with CUDA support is required.
+  * The model is tested on a single L40S GPU.
+  * 12GB is just a critical value, and 16GB GPU memory shoule be safer. 
+* Tested operating system: Linux
+
+### 🔧 Dependencies and Installation
+- Python >= 3.12
+- [PyTorch >= 2.9.1](https://pytorch.org/)
+- [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads)
+
+```bash
+git clone https://github.com/stepfun-ai/Step-Audio-EditX.git
+
+cd Step-Audio-EditX
+uv sync --refresh
+source .venv/bin/activate
+
+git lfs install
+git clone https://huggingface.co/stepfun-ai/Step-Audio-Tokenizer
+git clone https://huggingface.co/stepfun-ai/Step-Audio-EditX
+git clone https://huggingface.co/stepfun-ai/Step-Audio-EditX-AWQ-4bit/
+
+```
+
+After downloading the models, where_you_download_dir should have the following structure:
+```
+where_you_download_dir
+├── Step-Audio-Tokenizer
+├── Step-Audio-EditX
+```
+
+#### Run with Docker
+
+You can set up the environment required for running Step-Audio-EditX using the provided Dockerfile.
+
+```bash
+# build docker
+docker build . -t step-audio-editx
+
+# run docker
+docker run --rm --gpus all \
+    -v /your/code/path:/app \
+    -v /your/model/path:/model \
+    -p 7860:7860 \
+    step-audio-editx
+```
+#### Local Inference Demo
+> [!TIP]
+> For optimal performance, keep audio under 30 seconds per inference.
+
+```bash
+# zero-shot cloning
+# The path of the generated audio file is output/fear_zh_female_prompt_cloned.wav
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --tokenizer-path where_you_download_dir \
+    --prompt-text "我总觉得，有人在跟着我，我能听到奇怪的脚步声。" \
+    --prompt-audio "examples/fear_zh_female_prompt.wav" \
+    --generated-text "可惜没有如果，已经发生的事情终究是发生了。" \
+    --edit-type "clone" \
+    --output-dir ./output 
+
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --tokenizer-path where_you_download_dir \
+    --prompt-text "His political stance was conservative, and he was particularly close to margaret thatcher." \
+    --prompt-audio "examples/zero_shot_en_prompt.wav" \
+    --generated-text "Underneath the courtyard is a large underground exhibition room which connects the two buildings.	" \
+    --edit-type "clone" \
+    --output-dir ./output 
+
+# edit
+# There will be one or multiple wave files corresponding to each edit iteration, for example: output/fear_zh_female_prompt_edited_iter1.wav, output/fear_zh_female_prompt_edited_iter2.wav, ...
+# emotion; fear
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --tokenizer-path where_you_download_dir \
+    --prompt-text "我总觉得，有人在跟着我，我能听到奇怪的脚步声。" \
+    --prompt-audio "examples/fear_zh_female_prompt.wav" \
+    --edit-type "emotion" \
+    --edit-info "fear" \
+    --output-dir ./output 
+
+# emotion; happy
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --tokenizer-path where_you_download_dir \
+    --prompt-text "You know, I just finished that big project and feel so relieved. Everything seems easier and more colorful, what a wonderful feeling!" \
+    --prompt-audio "examples/en_happy_prompt.wav" \
+    --edit-type "emotion" \
+    --edit-info "happy" \
+    --output-dir ./output 
+
+# style; whisper
+# for style whisper, the edit iteration num should be set bigger than 1 to get better results.
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --tokenizer-path where_you_download_dir \
+    --prompt-text "比如在工作间隙，做一些简单的伸展运动，放松一下身体，这样，会让你更有精力." \
+    --prompt-audio "examples/whisper_prompt.wav" \
+    --edit-type "style" \
+    --edit-info "whisper" \
+    --output-dir ./output 
+
+# paraliguistic 
+# supported tags, Breathing, Laughter, Surprise-oh, Confirmation-en, Uhm, Surprise-ah, Surprise-wa, Sigh, Question-ei, Dissatisfaction-hnn
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --tokenizer-path where_you_download_dir \
+    --prompt-text "我觉得这个计划大概是可行的，不过还需要再仔细考虑一下。" \
+    --prompt-audio "examples/paralingustic_prompt.wav" \
+    --generated-text "我觉得这个计划大概是可行的，[Uhm]不过还需要再仔细考虑一下。" \
+    --edit-type "paralinguistic" \
+    --output-dir ./output 
+
+# denoise
+# Prompt text is not needed.
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --tokenizer-path where_you_download_dir \
+    --prompt-audio "examples/denoise_prompt.wav"\
+    --edit-type "denoise" \
+    --output-dir ./output 
+
+# vad 
+# Prompt text is not needed.
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --tokenizer-path where_you_download_dir \
+    --prompt-audio "examples/vad_prompt.wav" \
+    --edit-type "vad" \
+    --output-dir ./output 
+
+# speed
+# supported edit-info: faster, slower, more faster, more slower
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --tokenizer-path where_you_download_dir \
+    --prompt-text "上次你说鞋子有点磨脚，我给你买了一双软软的鞋垫。" \
+    --prompt-audio "examples/speed_prompt.wav" \
+    --edit-type "speed" \
+    --edit-info "more faster" \
+    --output-dir ./output 
+
+```
+
+
+
+#### Launch Web Demo
+Start a local server for online inference.
+Assume you have one GPU with at least 12GB memory available and have already downloaded all the models.
+
+```bash
+# Standard launch
+python app.py --model-path where_you_download_dir --tokenizer-path where_you_download_dir --model-source local
+
+# Using pre-quantized AWQ 4-bit models, memory-efficient mode (for limited GPU memory, ~6-8GB usage)
+python app.py \
+    --model-path path/to/quantized/model \
+    --tokenizer-path where_you_download_dir \
+    --model-source local \
+    --gpu-memory-utilization 0.1 \
+    --enforce-eager \
+    --max-num-seqs 1 \
+    --cosyvoice-dtype bfloat16 \
+    --no-cosyvoice-cuda-graph
+
+```
+
+##### Available Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--model-path` | (required) | Path to the model directory |
+| `--model-source` | `auto` | Model source: `auto`, `local`, `modelscope`, `huggingface` |
+| `--gpu-memory-utilization` | `0.5` | GPU memory ratio for vLLM KV cache (0.0-1.0) |
+| `--max-model-len` | `3072` | Maximum sequence length, affects KV cache size |
+| `--enforce-eager` | `True` | Disable vLLM CUDA Graphs (saves ~0.5GB memory) |
+| `--max-num-seqs` | `1` | Maximum concurrent sequences (vLLM default: 256, lower = less memory) |
+| `--dtype` | `bfloat16` | Model dtype: `float16`, `bfloat16` |
+| `--quantization` | `None` | Quantization method: `awq`, `gptq`, `fp8` |
+| `--cosyvoice-dtype` | `bfloat16` | CosyVoice vocoder dtype: `float32`, `bfloat16`, `float16` |
+| `--no-cosyvoice-cuda-graph` | `False` | Disable CosyVoice CUDA Graphs (saves memory) |
+| `--enable-auto-transcribe` | `False` | Enable automatic audio transcription |
+
+##### Memory Usage Guide
+
+| Configuration | Estimated GPU Memory | Use Case |
+|--------------|---------------------|----------|
+| Standard (defaults) | ~12-15 GB | Best quality and speed |
+| Memory-efficient | ~6-8 GB | Limited GPU memory, some quality trade-off |
+| AWQ 4-bit quantized | ~8-10 GB | Good balance of quality and memory |
+
+## Training
+Please refer to script/ReadMe.md
+
+### 🔄 Model Quantization (Optional)
+
+For users with limited GPU memory, you can create quantized versions of the model to reduce memory requirements:
+
+```bash
+# Create an AWQ 4-bit quantized model
+python quantization/awq_quantize.py --model_path path/to/Step-Audio-EditX
+
+# Advanced quantization options
+python quantization/awq_quantize.py
+```
+
+For detailed quantization options and parameters, see [quantization/README.md](quantization/README.md).
+
+
+## Technical Details
+<img src="assets/architechture.png" width=900>
+Step-Audio-EditX comprises three primary components: 
+
+- A dual-codebook audio tokenizer, which converts reference or input audio into discrete tokens.
+- An audio LLM that generates dual-codebook token sequences.
+- An audio decoder, which converts the dual-codebook token sequences predicted by the audio LLM back into audio waveforms using a flow matching approach.
+
+Audio-Edit enables iterative control over emotion and speaking style across all voices, leveraging large-margin data during SFT and PPO training.
+
+## Evaluation
+
+### Comparison between Step-Audio-EditX and Closed-Source models.
+
+- Step-Audio-EditX demonstrates superior performance over Minimax and Doubao in both zero-shot cloning and emotion control.
+- Emotion editing of Step-Audio-EditX significantly improves the emotion-controlled audio outputs of all three models after just one iteration. With further iterations, their overall performance continues to improve.
+
+<div align="center">
+<img src="assets/emotion-eval.png" width=800 >
+</div>
+
+### Generalization on Closed-Source Models.
+- For emotion and speaking style editing, the built-in voices of leading closed-source systems possess considerable in-context capabilities, allowing them to partially convey the emotions in the text. After a single editing round with Step-Audio-EditX, the emotion and style accuracy across all voice models exhibited significant improvement. Further enhancement was observed over the next two iterations, robustly demonstrating our model's strong generalization.
+
+- For paralinguistic editing, after editing with Step-Audio-EditX, the performance of paralinguistic reproduction is comparable to that achieved by the built-in voices of closed-source models when synthesizing native paralinguistic content directly. (**sub** means replacement of paralinguistic tags with native words)
+
+
+<div align="center">
+
+  <table border="1" cellspacing="0" cellpadding="5" style="border-collapse: collapse; font-family: sans-serif; width: auto;">
+    <caption><b>Table: Generalization of Emotion, Speaking Style, and Paralinguistic Editing on Closed-Source Models.</b></caption>
+    <thead>
+      <tr>
+        <th rowspan="2" align="center" style="vertical-align: bottom;">Language</th>
+        <th rowspan="2" align="center" style="vertical-align: bottom;">Model</th>
+        <th colspan="4" style="border-bottom: 1px solid black;">Emotion &uarr;</th>
+        <th colspan="4" style="border-bottom: 1px solid black;">Speaking Style &uarr;</th>
+        <th colspan="3" style="border-bottom: 1px solid black; border-left: 1px solid black;">Paralinguistic &uarr;</th>
+      </tr>
+      <tr>
+        <th>Iter<sub>0</sub></th>
+        <th>Iter<sub>1</sub></th>
+        <th>Iter<sub>2</sub></th>
+        <th>Iter<sub>3</sub></th>
+        <th style="border-left: 1px solid #ccc;">Iter<sub>0</sub></th>
+        <th>Iter<sub>1</sub></th>
+        <th>Iter<sub>2</sub></th>
+        <th>Iter<sub>3</sub></th>
+        <th style="border-left: 1px solid black;">Iter<sub>0</sub></th>
+        <th>sub</th>
+        <th>Iter<sub>1</sub></th>
+      </tr>
+    </thead>
+    <tbody>
+      <tr>
+        <td rowspan="4" align="center" style="font-weight: bold; vertical-align: middle;">Chinese</td>
+        <td align="left">MiniMax-2.6-hd</td>
+        <td align="center">71.6</td>
+        <td align="center">78.6</td>
+        <td align="center">81.2</td>
+        <td align="center"><b>83.4</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">36.7</td>
+        <td align="center">58.8</td>
+        <td align="center">63.1</td>
+        <td align="center"><b>67.3</b></td>
+        <td align="center" style="border-left: 1px solid black;">1.73</td>
+        <td align="center">2.80</td>
+        <td align="center">2.90</td>
+      </tr>
+      <tr>
+        <td align="left">Doubao-Seed-TTS-2.0</td>
+        <td align="center">67.4</td>
+        <td align="center">77.8</td>
+        <td align="center">80.6</td>
+        <td align="center"><b>82.8</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">38.2</td>
+        <td align="center">60.2</td>
+        <td align="center"><b>65.0</b></td>
+        <td align="center">64.9</td>
+        <td align="center" style="border-left: 1px solid black;">1.67</td>
+        <td align="center">2.81</td>
+        <td align="center">2.90</td>
+      </tr>
+      <tr>
+        <td align="left">GPT-4o-mini-TTS</td>
+        <td align="center">62.6</td>
+        <td align="center">76.0</td>
+        <td align="center">77.0</td>
+        <td align="center"><b>81.8</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">45.9</td>
+        <td align="center">64.0</td>
+        <td align="center">65.7</td>
+        <td align="center"><b>69.7</b></td>
+        <td align="center" style="border-left: 1px solid black;">1.71</td>
+        <td align="center">2.88</td>
+        <td align="center">2.93</td>
+      </tr>
+      <tr style="border-bottom: 1px solid black;">
+        <td align="left">ElevenLabs-v2</td>
+        <td align="center">60.4</td>
+        <td align="center">74.6</td>
+        <td align="center">77.4</td>
+        <td align="center"><b>79.2</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">43.8</td>
+        <td align="center">63.3</td>
+        <td align="center">69.7</td>
+        <td align="center"><b>70.8</b></td>
+        <td align="center" style="border-left: 1px solid black;">1.70</td>
+        <td align="center">2.71</td>
+        <td align="center">2.92</td>
+      </tr>
+      <tr>
+        <td rowspan="4" align="center" style="font-weight: bold; vertical-align: middle;">English</td>
+        <td align="left">MiniMax-2.6-hd</td>
+        <td align="center">55.0</td>
+        <td align="center">64.0</td>
+        <td align="center">64.2</td>
+        <td align="center"><b>66.4</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">51.9</td>
+        <td align="center">60.3</td>
+        <td align="center">62.3</td>
+        <td align="center"><b>64.3</b></td>
+        <td align="center" style="border-left: 1px solid black;">1.72</td>
+        <td align="center">2.87</td>
+        <td align="center">2.88</td>
+      </tr>
+      <tr>
+        <td align="left">Doubao-Seed-TTS-2.0</td>
+        <td align="center">53.8</td>
+        <td align="center">65.8</td>
+        <td align="center">65.8</td>
+        <td align="center"><b>66.2</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">47.0</td>
+        <td align="center">62.0</td>
+        <td align="center"><b>62.7</b></td>
+        <td align="center">62.3</td>
+        <td align="center" style="border-left: 1px solid black;">1.72</td>
+        <td align="center">2.75</td>
+        <td align="center">2.92</td>
+      </tr>
+      <tr>
+        <td align="left">GPT-4o-mini-TTS</td>
+        <td align="center">56.8</td>
+        <td align="center">61.4</td>
+        <td align="center">64.8</td>
+        <td align="center"><b>65.2</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">52.3</td>
+        <td align="center">62.3</td>
+        <td align="center">62.4</td>
+        <td align="center"><b>63.4</b></td>
+        <td align="center" style="border-left: 1px solid black;">1.90</td>
+        <td align="center">2.90</td>
+        <td align="center">2.88</td>
+      </tr>
+      <tr style="border-bottom: 1px solid black;">
+        <td align="left">ElevenLabs-v2</td>
+        <td align="center">51.0</td>
+        <td align="center">61.2</td>
+        <td align="center">64.0</td>
+        <td align="center"><b>65.2</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">51.0</td>
+        <td align="center">62.1</td>
+        <td align="center">62.6</td>
+        <td align="center"><b>64.0</b></td>
+        <td align="center" style="border-left: 1px solid black;">1.93</td>
+        <td align="center">2.87</td>
+        <td align="center">2.88</td>
+      </tr>
+      <tr>
+        <td rowspan="4" align="center" style="font-weight: bold; vertical-align: middle;">Average</td>
+        <td align="left">MiniMax-2.6-hd</td>
+        <td align="center">63.3</td>
+        <td align="center">71.3</td>
+        <td align="center">72.7</td>
+        <td align="center"><b>74.9</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">44.2</td>
+        <td align="center">59.6</td>
+        <td align="center">62.7</td>
+        <td align="center"><b>65.8</b></td>
+        <td align="center" style="border-left: 1px solid black;">1.73</td>
+        <td align="center">2.84</td>
+        <td align="center">2.89</td>
+      </tr>
+      <tr>
+        <td align="left">Doubao-Seed-TTS-2.0</td>
+        <td align="center">60.6</td>
+        <td align="center">71.8</td>
+        <td align="center">73.2</td>
+        <td align="center"><b>74.5</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">42.6</td>
+        <td align="center">61.1</td>
+        <td align="center"><b>63.9</b></td>
+        <td align="center">63.6</td>
+        <td align="center" style="border-left: 1px solid black;">1.70</td>
+        <td align="center">2.78</td>
+        <td align="center">2.91</td>
+      </tr>
+      <tr>
+        <td align="left">GPT-4o-mini-TTS</td>
+        <td align="center">59.7</td>
+        <td align="center">68.7</td>
+        <td align="center">70.9</td>
+        <td align="center"><b>73.5</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">49.1</td>
+        <td align="center">63.2</td>
+        <td align="center">64.1</td>
+        <td align="center"><b>66.6</b></td>
+        <td align="center" style="border-left: 1px solid black;">1.81</td>
+        <td align="center">2.89</td>
+        <td align="center">2.90</td>
+      </tr>
+      <tr>
+        <td align="left">ElevenLabs-v2</td>
+        <td align="center">55.7</td>
+        <td align="center">67.9</td>
+        <td align="center">70.7</td>
+        <td align="center"><b>72.2</b></td>
+        <td align="center" style="border-left: 1px solid #ccc;">47.4</td>
+        <td align="center">62.7</td>
+        <td align="center">66.1</td>
+        <td align="center"><b>67.4</b></td>
+        <td align="center" style="border-left: 1px solid black;">1.82</td>
+        <td align="center">2.79</td>
+        <td align="center">2.90</td>
+      </tr>
+    </tbody>
+  </table>
+
+</div>
+
+
+## Acknowledgements
+
+Part of the code and data for this project comes from:
+* [CosyVoice](https://github.com/FunAudioLLM/CosyVoice)
+* [transformers](https://github.com/huggingface/transformers)
+* [FunASR](https://github.com/modelscope/FunASR)
+* [NVSpeech](https://huggingface.co/datasets/amphion/Emilia-NV)
+* [vllm](https://github.com/vllm-project/vllm)
+
+Thank you to all the open-source projects for their contributions to this project!
+
+## License Agreement
++ The code in this open-source repository is licensed under the [Apache 2.0](LICENSE) License.
+
+## Citation
+
+```
+@misc{yan2025stepaudioeditxtechnicalreport,
+      title={Step-Audio-EditX Technical Report}, 
+      author={Chao Yan and Boyong Wu and Peng Yang and Pengfei Tan and Guoqiang Hu and Yuxin Zhang and Xiangyu and Zhang and Fei Tian and Xuerui Yang and Xiangyu Zhang and Daxin Jiang and Gang Yu},
+      year={2025},
+      eprint={2511.03601},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2511.03601}, 
+}
+```
+
+
+## ⚠️ Usage Disclaimer
+- Do not use this model for any unauthorized activities, including but not limited to:
+  - Voice cloning without permission
+  - Identity impersonation
+  - Fraud
+  - Deepfakes or any other illegal purposes
+- Ensure compliance with local laws and regulations, and adhere to ethical guidelines when using this model.
+- The model developers are not responsible for any misuse or abuse of this technology.
+
+We advocate for responsible generative AI research and urge the community to uphold safety and ethical standards in AI development and application. If you have any concerns regarding the use of this model, please feel free to contact us.
+
+## Star History
+[![Star History Chart](https://api.star-history.com/svg?repos=stepfun-ai/Step-Audio-EditX&type=Date)](https://star-history.com/#stepfun-ai/Step-Audio-EditX&Date)

assets/test.wav

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

config.json

@@ -0,0 +1,24 @@
+{
+  "architectures": [
+    "Step1ForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_step1.Step1Config",
+    "AutoModelForCausalLM": "modeling_step1.Step1ForCausalLM"
+  },
+  "model_type": "step1",
+  "bos_token_id": 1,
+  "pad_token_id": 0,
+  "eos_token_id": 3,
+  "hidden_size": 3072,
+  "intermediate_size": 8192,
+  "num_attention_heads": 48,
+  "num_attention_groups": 4,
+  "num_hidden_layers": 32,
+  "max_seq_len": 32768,
+  "vocab_size": 74752,
+  "rms_norm_eps": 1e-05,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "use_cache": true 
+}

configuration_step1.py

@@ -0,0 +1,41 @@
+from typing import Optional, List, Any, Dict
+from transformers.configuration_utils import PretrainedConfig
+
+
+
+class Step1Config(PretrainedConfig):
+    model_type = "step1"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    
+    def __init__(
+        self,
+        hidden_size: int = 5120,
+        intermediate_size: int = 13312,
+        num_attention_heads: int = 40,
+        num_attention_groups: int = 8,
+        num_hidden_layers: int = 48,
+        max_seq_len: int = 4096,
+        vocab_size: int = 65536,
+        rms_norm_eps: float = 1e-5,
+        bos_token_id: int = 1,
+        eos_token_id: int = 3,
+        pad_token_id: int = 0,
+        **kwargs,
+    ) -> None:
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_attention_heads = num_attention_heads
+        self.num_attention_groups = num_attention_groups
+        self.num_hidden_layers = num_hidden_layers
+        self.max_seq_len = max_seq_len
+        self.vocab_size = vocab_size
+        self.rms_norm_eps = rms_norm_eps
+        super().__init__(
+            bos_token_id=bos_token_id,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs
+        )
+
+
+__all__ = ["Step1Config"]

model-00001.safetensors

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

model.safetensors.index.json

@@ -0,0 +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_step1.py

@@ -0,0 +1,414 @@
+import math
+from typing import Optional, Tuple, Union, List
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from transformers.generation import GenerationMixin
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+from .configuration_step1 import Step1Config
+from transformers.cache_utils import Cache, DynamicCache
+from einops import rearrange
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+
+logger = logging.get_logger(__name__)
+
+
+def build_alibi_cache(block_size, n_heads, dtype, device):
+    # get slopes
+    n = 2 ** math.floor(math.log2(n_heads))  # nearest 2**n to n_heads
+    m0 = 2.0 ** (-8.0 / n)
+    # 2^(-8/n), 2^(-8*2/n), 2^(-8*3/n), ...
+    slopes = torch.pow(m0, torch.arange(1, n + 1))
+    if n < n_heads:
+        m1 = 2.0 ** (-4.0 / n)
+        # 2^(-8/(2n)), 2^(-8*3/(2n)), 2^(-8*5/(2n)), ...
+        mm = torch.pow(m1, torch.arange(1, 1 + 2 * (n_heads - n), 2))
+        slopes = torch.cat([slopes, mm])
+    slopes = slopes.to(device)
+
+    tril = torch.tril(torch.ones(1, 1, block_size, block_size, device=device))
+
+    bias_rows = torch.arange(block_size, device=device).view(1, -1)
+    bias_cols = torch.arange(block_size, device=device).view(-1, 1)
+    bias = -torch.sqrt(bias_cols - bias_rows)
+    bias = bias.view(1, block_size, block_size) * slopes.view(-1, 1, 1)
+    bias = bias.masked_fill(tril == 0, float("-inf"))
+
+    return bias.type(dtype)
+
+
+class StepRMSNorm(torch.nn.Module):
+    def __init__(self, hidden_size, eps=1e-5):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor):
+        var = x.float().pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps).to(x.dtype)
+        x = x * self.weight
+        return x
+
+
+class StepAttention(torch.nn.Module):
+    def __init__(self, hidden_size, num_heads, num_groups, layer_idx: int):
+        super().__init__()
+
+        self.num_heads = num_heads
+        self.num_groups = num_groups
+        self.hidden_size = hidden_size
+        self.head_dim = hidden_size // num_heads
+
+        self.q_proj = torch.nn.Linear(hidden_size, hidden_size, bias=False)
+        self.k_proj = torch.nn.Linear(
+            hidden_size, num_groups * self.head_dim, bias=False
+        )
+        self.v_proj = torch.nn.Linear(
+            hidden_size, num_groups * self.head_dim, bias=False
+        )
+        self.o_proj = torch.nn.Linear(hidden_size, hidden_size, bias=False)
+
+        self.layer_idx = layer_idx
+
+    def flash_attn_func(self, q, k, v, dropout_p=0.0, softmax_scale=None, causal=True,
+                    return_attn_probs=False, tp_group_rank=0, tp_group_size=1):
+        softmax_scale = q.size(-1) ** (-0.5) if softmax_scale is None else softmax_scale
+        return torch.ops.Optimus.fwd(q, k, v, None, dropout_p, softmax_scale, causal, return_attn_probs, None, tp_group_rank, tp_group_size)[0]
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        past_key_value: Optional[Cache] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ):
+
+        q: torch.Tensor = self.q_proj(x)
+        k: torch.Tensor = self.k_proj(x)
+        v: torch.Tensor = self.v_proj(x)
+        if past_key_value is not None:
+            cache_kwargs = {"cache_position": cache_position}
+            k, v = past_key_value.update(k, v, self.layer_idx, cache_kwargs)
+
+        q = rearrange(q, "b s (h d) -> b s h d", h=self.num_heads)
+        k = rearrange(k, "b s (g d) -> b s g d", g=self.num_groups)
+        v = rearrange(v, "b s (g d) -> b s g d", g=self.num_groups)
+
+        try:
+            if self.head_dim not in (64, 128):
+                raise ValueError("head_dim must be 64 or 128")
+            attn_output = self.flash_attn_func(q, k, v)
+            attn_output = attn_output.flatten(-2, -1)
+        except:
+            k = k.repeat_interleave(self.num_heads // self.num_groups, dim=-2)
+            v = v.repeat_interleave(self.num_heads // self.num_groups, dim=-2)
+
+            attention_mask = build_alibi_cache(
+                k.size(1), self.num_heads, dtype=q.dtype, device=q.device
+            )[:, :, -q.size(1) :, :].contiguous()
+
+            q = q.transpose(1, 2)
+            k = k.transpose(1, 2)
+            v = v.transpose(1, 2)
+
+            attn_output: torch.Tensor = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=attention_mask
+            )
+
+            attn_output = attn_output.transpose(1, 2).flatten(-2, -1)
+
+        out = self.o_proj(attn_output)
+        return out, None  # attn weights are not returned
+
+
+class StepMLP(torch.nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.gate_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = torch.nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        gate = self.gate_proj(x)
+        up = self.up_proj(x)
+        x = torch.nn.functional.silu(gate) * up
+        x = self.down_proj(x)
+        return x
+
+
+class StepLayer(torch.nn.Module):
+    def __init__(self, config: Step1Config, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.self_attn = StepAttention(
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_groups=config.num_attention_groups,
+            layer_idx=layer_idx,
+        )
+        self.mlp = StepMLP(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+        )
+        self.input_layernorm = StepRMSNorm(
+            hidden_size=config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm = StepRMSNorm(
+            hidden_size=config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+    ):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, self_attn_weights = self.self_attn(hidden_states, past_key_value, attention_mask, cache_position)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, )
+        if output_attentions:
+            outputs += (self_attn_weights,)
+        return outputs
+
+
+class StepPreTrainedModel(PreTrainedModel):
+    config_class = Step1Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["StepLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_cache_class = True
+    _supports_static_cache = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class Step1Model(StepPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]
+
+    Args:
+        config: Step1Config
+    """
+
+    def __init__(self, config: Step1Config):
+        super().__init__(config)
+        self.config = config
+        self.embed_tokens = torch.nn.Embedding(config.vocab_size, config.hidden_size)
+
+        self.layers = torch.nn.Sequential(
+            *[
+                StepLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+
+        self.norm = StepRMSNorm(
+            hidden_size=config.hidden_size, eps=config.rms_norm_eps
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        causal_mask = attention_mask
+
+        hidden_states = inputs_embeds
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                past_key_value=past_key_values,
+                cache_position=cache_position,
+                output_attentions=output_attentions,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        output = BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=None,
+        )
+        return output if return_dict else output.to_tuple()
+
+
+class Step1ForCausalLM(StepPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Step1Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+            )
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

tokenizer.model

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

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": false,
+  "model_max_length": 65536,
+  "pad_token": "<unk>",
+  "padding_side": "left",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false,
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{{ '<s>' }}{% endif %}{% for message in messages %}{% if message['role'] == 'user' %}{% set role = 'human' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<|BOT|> ' + role + '\\n' }}{{ message['content'] }}{% if not loop.last or message['role'] != 'assistant' %}{{ '<|EOT|>' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|BOT|> assistant\\n' }}{% endif %}"
+}
+