Step-Audio (models)

.gitattributes

@@ -55,3 +55,12 @@ models/Step-Audio-TTS-3B/lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-li
 models/Step-Audio-Chat/lib/liboptimus_ths-torch2.2-cu121.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
 models/Step-Audio-Chat/lib/liboptimus_ths-torch2.3-cu121.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
 models/Step-Audio-Chat/lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+models/Step-Audio-R1-NVFP4A16/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+models/StepAudio2mini_BPFT/checkpoint-12000/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+models/StepAudio2mini_BPFT/checkpoint-15000/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+models/StepAudio2mini_BPFT/checkpoint-18000/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+models/StepAudio2mini_BPFT/checkpoint-18850/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+models/StepAudio2mini_BPFT/checkpoint-3000/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+models/StepAudio2mini_BPFT/checkpoint-6000/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+models/StepAudio2mini_BPFT/checkpoint-9000/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+models/StepAudio2mini_BPFT/tokenizer.json filter=lfs diff=lfs merge=lfs -text

models/Step-Audio-EditX-bnb-4bit/.gitattributes

@@ -0,0 +1,35 @@
+*.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

models/Step-Audio-EditX-bnb-4bit/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

models/Step-Audio-EditX-bnb-4bit/CosyVoice-300M-25Hz/campplus.onnx

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

models/Step-Audio-EditX-bnb-4bit/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

models/Step-Audio-EditX-bnb-4bit/CosyVoice-300M-25Hz/flow.pt

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

models/Step-Audio-EditX-bnb-4bit/CosyVoice-300M-25Hz/hift.pt

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

models/Step-Audio-EditX-bnb-4bit/CosyVoice-300M-25Hz/speech_tokenizer_v1.onnx

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

models/Step-Audio-EditX-bnb-4bit/README.md

@@ -0,0 +1,3 @@
+---
+license: apache-2.0
+---

models/Step-Audio-EditX-bnb-4bit/config.json

@@ -0,0 +1,39 @@
+{
+  "architectures": [
+    "Step1ForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_step1.Step1Config",
+    "AutoModelForCausalLM": "modeling_step1.Step1ForCausalLM"
+  },
+  "bos_token_id": 1,
+  "eos_token_id": 3,
+  "hidden_size": 3072,
+  "intermediate_size": 8192,
+  "max_seq_len": 32768,
+  "model_type": "step1",
+  "num_attention_groups": 4,
+  "num_attention_heads": 48,
+  "num_hidden_layers": 32,
+  "pad_token_id": 0,
+  "quantization_config": {
+    "_load_in_4bit": true,
+    "_load_in_8bit": false,
+    "bnb_4bit_compute_dtype": "float16",
+    "bnb_4bit_quant_storage": "uint8",
+    "bnb_4bit_quant_type": "nf4",
+    "bnb_4bit_use_double_quant": true,
+    "llm_int8_enable_fp32_cpu_offload": false,
+    "llm_int8_has_fp16_weight": false,
+    "llm_int8_skip_modules": null,
+    "llm_int8_threshold": 6.0,
+    "load_in_4bit": true,
+    "load_in_8bit": false,
+    "quant_method": "bitsandbytes"
+  },
+  "rms_norm_eps": 1e-05,
+  "torch_dtype": "float16",
+  "transformers_version": "4.53.3",
+  "use_cache": true,
+  "vocab_size": 74752
+}

models/Step-Audio-EditX-bnb-4bit/configuration_step1.py

@@ -0,0 +1,43 @@
+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,
+        use_cache: bool = True,
+        **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
+        self.use_cache = use_cache
+        super().__init__(
+            bos_token_id=bos_token_id,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs
+        )
+
+
+__all__ = ["Step1Config"]

models/Step-Audio-EditX-bnb-4bit/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 3,
+  "pad_token_id": 0,
+  "transformers_version": "4.53.3"
+}

models/Step-Audio-EditX-bnb-4bit/model.safetensors

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

models/Step-Audio-EditX-bnb-4bit/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,
+        )

models/Step-Audio-EditX-bnb-4bit/quantization_config.json

@@ -0,0 +1,10 @@
+{
+  "bits": 4,
+  "compute_dtype": "float16",
+  "quantization_method": "bitsandbytes",
+  "load_in_4bit": true,
+  "load_in_8bit": false,
+  "bnb_4bit_compute_dtype": "float16",
+  "bnb_4bit_use_double_quant": true,
+  "bnb_4bit_quant_type": "nf4"
+}

models/Step-Audio-EditX-bnb-4bit/source.txt

@@ -0,0 +1 @@
+https://huggingface.co/SUP3RMASS1VE/Step-Audio-EditX-bnb-4bit

models/Step-Audio-EditX-bnb-4bit/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

models/Step-Audio-EditX-bnb-4bit/tokenizer.model

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

models/Step-Audio-EditX/.gitattributes

@@ -0,0 +1,35 @@
+*.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

models/Step-Audio-EditX/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

models/Step-Audio-EditX/CosyVoice-300M-25Hz/campplus.onnx

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

models/Step-Audio-EditX/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

models/Step-Audio-EditX/CosyVoice-300M-25Hz/flow.pt

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

models/Step-Audio-EditX/CosyVoice-300M-25Hz/hift.pt

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

models/Step-Audio-EditX/CosyVoice-300M-25Hz/speech_tokenizer_v1.onnx

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

models/Step-Audio-EditX/README.md

@@ -0,0 +1,144 @@
+---
+license: apache-2.0
+pipeline_tag: text-to-speech
+library_name: transformers
+---
+## Step-Audio-EditX
+
+ ✨ [Demo Page](https://stepaudiollm.github.io/step-audio-editx/) 
+| 🌟 [GitHub](https://github.com/stepfun-ai/Step-Audio-EditX) 
+| 📑 [Paper](https://arxiv.org/abs/2511.03601) 
+
+Check our open-source repository https://github.com/stepfun-ai/Step-Audio-EditX for more details!
+
+## 🔥🔥🔥 News!!！
+* 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.
+
+We are open-sourcing **Step-Audio-EditX**, a powerful **3B parameters** LLM-based 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.
+
+## Features
+- **Zero-Shot TTS**
+  - Excellent zero-shot TTS cloning for `Mandarin`, `English`, `Sichuanese`, `Cantonese`, `Japanese` and `Korean`.
+  - To use a dialect, just add a **`[Sichuanese]`**, **`[Cantonese]`** ,**`[Japanese]`**,**`[Korean]`** tag before your text.
+ 
+- **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*, *Suprise-oh*, *Confirmation-en*, *Uhm*, *Suprise-ah*, *Suprise-wa*, *Sigh*, *Question-ei*, *Dissatisfaction-hnn* ]
+
+For more examples, see [demo page](https://stepaudiollm.github.io/step-audio-editx/).
+
+## Model Usage
+### 📜 Requirements
+The following table shows the requirements for running Step-Audio-EditX model:
+
+|     Model    | Parameters |  Setting<br/>(sample frequency) | GPU Optimal Memory  |
+|------------|------------|--------------------------------|----------------|
+| Step-Audio-EditX   | 3B|         41.6Hz          |       32 GB        |
+
+* An NVIDIA GPU with CUDA support is required.
+  * The model is tested on a single L40S GPU.
+* Tested operating system: Linux
+
+### 🔧 Dependencies and Installation
+- Python >= 3.10.0 (Recommend to use [Anaconda](https://www.anaconda.com/download/#linux) or [Miniconda](https://docs.conda.io/en/latest/miniconda.html))
+- [PyTorch >= 2.4.1-cu121](https://pytorch.org/)
+- [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads)
+
+```bash
+git clone https://github.com/stepfun-ai/Step-Audio-EditX.git
+conda create -n stepaudioedit python=3.10
+conda activate stepaudioedit
+
+cd Step-Audio-EditX
+pip install -r requirements.txt
+
+git lfs install
+git clone https://huggingface.co/stepfun-ai/Step-Audio-Tokenizer
+git clone https://huggingface.co/stepfun-ai/Step-Audio-EditX
+
+```
+
+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 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
+```
+
+
+#### Launch Web Demo
+Start a local server for online inference.
+Assume you have one GPU with at least 32GB memory available  and have already downloaded all the models.
+
+```bash
+# Step-Audio-EditX demo
+python app.py --model-path where_you_download_dir --model-source local
+```
+
+#### Local Inference Demo
+> [!TIP]
+> For optimal performance, keep audio under 30 seconds per inference.
+
+```bash
+# zero-shot cloning
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --output-dir ./output \
+    --prompt-text "your prompt text"\
+    --prompt-audio your_prompt_audio_path \
+    --generated-text "your target text" \
+    --edit-type "clone"
+
+# edit
+python3 tts_infer.py \
+    --model-path where_you_download_dir \
+    --output-dir ./output \
+    --prompt-text "your promt text" \
+    --prompt-audio your_prompt_audio_path \
+    --generated-text "" \ # for para-linguistic editing, you need to specify the generatedd text
+    --edit-type "emotion" \
+    --edit-info "sad" \
+    --n-edit-iter 2
+```
+
+
+## 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}, 
+}
+
+```

models/Step-Audio-EditX/config.json

@@ -0,0 +1,22 @@
+{
+  "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,
+  "torch_dtype": "bfloat16"
+}

models/Step-Audio-EditX/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"]

models/Step-Audio-EditX/model-00001.safetensors

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

models/Step-Audio-EditX/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"}}

models/Step-Audio-EditX/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,
+        )

models/Step-Audio-EditX/source.txt

@@ -0,0 +1 @@
+https://huggingface.co/stepfun-ai/Step-Audio-EditX

models/Step-Audio-EditX/tokenizer.model

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

models/Step-Audio-EditX/tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "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
+}
+

models/Step-Audio-R1-NVFP4A16/.gitattributes

@@ -0,0 +1,36 @@
+*.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
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

models/Step-Audio-R1-NVFP4A16/README.md

@@ -0,0 +1,216 @@
+---
+license: apache-2.0
+pipeline_tag: audio-text-to-text
+library_name: transformers
+tags:
+- audio-reasoning
+- chain-of-thought
+- multi-modal
+- step-audio-r1
+---
+## Step-Audio-R1-NVFP4A16 (Quantized)
+
+This is a **quantized version** of Step-Audio-R1 using NVFP4A16 quantization via [LLM Compressor](https://github.com/vllm-project/llm-compressor).
+
+### Quantization Details
+
+- **Scheme**: NVFP4A16 (FP4 weights with FP16 activations)
+- **Target layers**: All Linear layers (except `encoder`, `adapter`, `lm_head`)
+- **Group size**: 16
+- **Method**: Post-Training Quantization (PTQ)
+
+### Quantization Code
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from llmcompressor import oneshot
+from llmcompressor.modifiers.quantization import QuantizationModifier
+from llmcompressor.utils import dispatch_for_generation
+
+MODEL_ID = "stepfun-ai/Step-Audio-R1"
+
+# Load model
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+# Configure the quantization algorithm and scheme
+# Quantize weights to FP4 with per group 16 via PTQ
+recipe = QuantizationModifier(targets="Linear", scheme="NVFP4A16", ignore=["lm_head", "re:encoder.*", "re:adapter.*"])
+
+# Apply quantization
+oneshot(model=model, recipe=recipe)
+
+# Save to disk in compressed-tensors format
+SAVE_DIR = "Step-Audio-R1-NVFP4A16"
+model.save_pretrained(SAVE_DIR, save_compressed=True)
+tokenizer.save_pretrained(SAVE_DIR)
+```
+
+
+## Step-Audio-R1
+
+ ✨ [Demo Page](https://stepaudiollm.github.io/step-audio-r1/) 
+| 🎮 [Playground](https://huggingface.co/spaces/stepfun-ai/Step-Audio-R1) 
+| 🌟 [GitHub](https://github.com/stepfun-ai/Step-Audio-R1) 
+| 📑 [Paper](https://arxiv.org/abs/2511.15848) 
+
+Step-Audio-R1 is the **first audio language model to successfully unlock Chain-of-Thought (CoT) reasoning**. 
+It decisively solves the "inverted scaling" problem that plagues existing models, where performance degrades 
+with longer reasoning. Step-Audio-R1 is the first model to demonstrate that for audio, like text and vision, 
+allocating more compute at test-time predictably improves performance.
+
+We found the root cause of this anomaly: models were engaging in **textual surrogate reasoning** 
+(analyzing transcripts, not audio) due to a modality mismatch. To solve this, we introduce 
+**Modality-Grounded Reasoning Distillation (MGRD)**, an iterative training framework that shifts the model's 
+reasoning from textual abstractions to acoustic properties.
+
+This new approach allows us to create **Step-Audio-R1**, which:
+- Is the **first audio reasoning model** that successfully benefits from test-time compute scaling.
+- Surpasses **Gemini 2.5 Pro** and is comparable to **Gemini 3** across major audio reasoning tasks.
+- Transforms extended deliberation from a liability into a **powerful asset** for audio intelligence.
+
+## Features
+- **Chain-of-Thought (CoT) Reasoning**
+  - First audio language model to successfully unlock Chain-of-Thought reasoning capabilities.
+  - Generates audio-relevant reasoning chains that genuinely ground themselves in acoustic features.
+ 
+- **Modality-Grounded Reasoning Distillation (MGRD)**
+  - Innovative iterative training framework that shifts reasoning from textual abstractions to acoustic properties.
+  - Solves the modality mismatch problem that caused textual surrogate reasoning in previous models.
+    
+- **Superior Performance**
+  - Surpasses **Gemini 2.5 Pro** across comprehensive audio understanding and reasoning benchmarks.
+  - Comparable to **Gemini 3** across major audio reasoning tasks.
+  - Surpasses **Qwen3** in textual reasoning.
+  - Covers speech, environmental sounds, and music domains.
+
+
+For more examples, see [demo page](https://stepaudiollm.github.io/step-audio-r1/).
+
+## Model Usage
+### 📜 Requirements
+- **GPU**: NVIDIA GPUs with CUDA support (tested on 4×L40S/H100/H800/H20).
+- **Operating System**: Linux.
+- **Python**: >= 3.10.0.
+
+### ⬇️ Download Model
+First, you need to download the Step-Audio-R1 model weights.
+
+**Method A · Git LFS**
+   ```bash
+   git lfs install
+   git clone https://huggingface.co/stepfun-ai/Step-Audio-R1
+   ```
+
+**Method B · Hugging Face CLI**
+   ```bash
+   hf download stepfun-ai/Step-Audio-R1 --local-dir ./Step-Audio-R1
+   ```
+
+### 🚀 Deployment and Execution
+We provide two ways to serve the model: Docker (recommended) or compiling the customized vLLM backend.
+
+#### 🐳 Method 1 · Run with Docker (Recommended)
+
+A customized vLLM image is required.
+
+1.  **Pull the image**:
+```bash
+docker pull stepfun2025/vllm:step-audio-2-v20250909
+```
+2.  **Start the service**:
+    Assuming the model is downloaded in the `Step-Audio-R1` folder in the current directory.
+
+    ```bash
+    docker run --rm -ti --gpus all \
+        -v $(pwd)/Step-Audio-R1:/Step-Audio-R1 \
+        -p 9999:9999 \
+        stepfun2025/vllm:step-audio-2-v20250909 \
+        -- vllm serve /Step-Audio-R1 \
+        --served-model-name Step-Audio-R1 \
+        --port 9999 \
+        --max-model-len 16384 \
+        --max-num-seqs 32 \
+        --tensor-parallel-size 4 \
+        --chat-template '{%- macro render_content(content) -%}{%- if content is string -%}{{- content.replace("<audio_patch>\n", "<audio_patch>") -}}{%- elif content is mapping -%}{{- content['"'"'value'"'"'] if '"'"'value'"'"' in content else content['"'"'text'"'"'] -}}{%- elif content is iterable -%}{%- for item in content -%}{%- if item.type == '"'"'text'"'"' -%}{{- item['"'"'value'"'"'] if '"'"'value'"'"' in item else item['"'"'text'"'"'] -}}{%- elif item.type == '"'"'audio'"'"' -%}<audio_patch>{%- endif -%}{%- endfor -%}{%- endif -%}{%- endmacro -%}{%- if tools -%}{{- '"'"'<|BOT|>system\n'"'"' -}}{%- if messages[0]['"'"'role'"'"'] == '"'"'system'"'"' -%}{{- render_content(messages[0]['"'"'content'"'"']) + '"'"'<|EOT|>'"'"' -}}{%- endif -%}{{- '"'"'<|BOT|>tool_json_schemas\n'"'"' + tools|tojson + '"'"'<|EOT|>'"'"' -}}{%- else -%}{%- if messages[0]['"'"'role'"'"'] == '"'"'system'"'"' -%}{{- '"'"'<|BOT|>system\n'"'"' + render_content(messages[0]['"'"'content'"'"']) + '"'"'<|EOT|>'"'"' -}}{%- endif -%}{%- endif -%}{%- for message in messages -%}{%- if message["role"] == "user" -%}{{- '"'"'<|BOT|>human\n'"'"' + render_content(message["content"]) + '"'"'<|EOT|>'"'"' -}}{%- elif message["role"] == "assistant" -%}{{- '"'"'<|BOT|>assistant\n'"'"' + (render_content(message["content"]) if message["content"] else '"'"''"'"') -}}{%- set is_last_assistant = true -%}{%- for m in messages[loop.index:] -%}{%- if m["role"] == "assistant" -%}{%- set is_last_assistant = false -%}{%- endif -%}{%- endfor -%}{%- if not is_last_assistant -%}{{- '"'"'<|EOT|>'"'"' -}}{%- endif -%}{%- elif message["role"] == "function_output" -%}{%- else -%}{%- if not (loop.first and message["role"] == "system") -%}{{- '"'"'<|BOT|>'"'"' + message["role"] + '"'"'\n'"'"' + render_content(message["content"]) + '"'"'<|EOT|>'"'"' -}}{%- endif -%}{%- endif -%}{%- endfor -%}{%- if add_generation_prompt -%}{{- '"'"'<|BOT|>assistant\n<think>\n'"'"' -}}{%- endif -%}' \
+        --enable-log-requests \
+        --interleave-mm-strings \
+        --trust-remote-code
+    ```
+After the service starts, it will listen on `localhost:9999`.
+
+#### 🐳  Method 2 · Run from Source (Compile vLLM)
+Step-Audio-R1 requires a customized vLLM backend.
+
+1.  **Download Source Code**:
+    ```bash
+    git clone https://github.com/stepfun-ai/vllm.git
+    cd vllm
+    ```
+
+2.  **Prepare Environment**:
+    ```bash
+    python3 -m venv .venv
+    source .venv/bin/activate
+    ```
+
+3.  **Install and Compile**:
+    vLLM contains both C++ and Python code. We mainly modified the Python code, so the C++ part can use the pre-compiled version to speed up the process.
+    
+    ```bash
+    # Use pre-compiled C++ extensions (Recommended)
+    VLLM_USE_PRECOMPILED=1 pip install -e .
+    ```
+
+4.  **Switch Branch**:
+    After compilation, switch to the branch that supports Step-Audio.
+    ```bash
+    git checkout step-audio-2-mini
+    ```
+
+5.  **Start the Service**:
+    ```bash
+    # Ensure you are in the vllm directory and the virtual environment is activated
+    source .venv/bin/activate
+
+    python3 -m vllm.entrypoints.openai.api_server \
+        --model ../Step-Audio-R1 \
+        --served-model-name Step-Audio-R1 \
+        --port 9999 \
+        --host 0.0.0.0 \
+        --max-model-len 65536 \
+        --max-num-seqs 128 \
+        --tensor-parallel-size 4 \
+        --gpu-memory-utilization 0.85 \
+        --trust-remote-code \
+        --enable-log-requests \
+        --interleave-mm-strings \
+        --chat-template '{%- macro render_content(content) -%}{%- if content is string -%}{{- content.replace("<audio_patch>\n", "<audio_patch>") -}}{%- elif content is mapping -%}{{- content['"'"'value'"'"'] if '"'"'value'"'"' in content else content['"'"'text'"'"'] -}}{%- elif content is iterable -%}{%- for item in content -%}{%- if item.type == '"'"'text'"'"' -%}{{- item['"'"'value'"'"'] if '"'"'value'"'"' in item else item['"'"'text'"'"'] -}}{%- elif item.type == '"'"'audio'"'"' -%}<audio_patch>{%- endif -%}{%- endfor -%}{%- endif -%}{%- endmacro -%}{%- if tools -%}{{- '"'"'<|BOT|>system\n'"'"' -}}{%- if messages[0]['"'"'role'"'"'] == '"'"'system'"'"' -%}{{- render_content(messages[0]['"'"'content'"'"']) + '"'"'<|EOT|>'"'"' -}}{%- endif -%}{{- '"'"'<|BOT|>tool_json_schemas\n'"'"' + tools|tojson + '"'"'<|EOT|>'"'"' -}}{%- else -%}{%- if messages[0]['"'"'role'"'"'] == '"'"'system'"'"' -%}{{- '"'"'<|BOT|>system\n'"'"' + render_content(messages[0]['"'"'content'"'"']) + '"'"'<|EOT|>'"'"' -}}{%- endif -%}{%- endif -%}{%- for message in messages -%}{%- if message["role"] == "user" -%}{{- '"'"'<|BOT|>human\n'"'"' + render_content(message["content"]) + '"'"'<|EOT|>'"'"' -}}{%- elif message["role"] == "assistant" -%}{{- '"'"'<|BOT|>assistant\n'"'"' + (render_content(message["content"]) if message["content"] else '"'"''"'"') -}}{%- set is_last_assistant = true -%}{%- for m in messages[loop.index:] -%}{%- if m["role"] == "assistant" -%}{%- set is_last_assistant = false -%}{%- endif -%}{%- endfor -%}{%- if not is_last_assistant -%}{{- '"'"'<|EOT|>'"'"' -}}{%- endif -%}{%- elif message["role"] == "function_output" -%}{%- else -%}{%- if not (loop.first and message["role"] == "system") -%}{{- '"'"'<|BOT|>'"'"' + message["role"] + '"'"'\n'"'"' + render_content(message["content"]) + '"'"'<|EOT|>'"'"' -}}{%- endif -%}{%- endif -%}{%- endfor -%}{%- if add_generation_prompt -%}{{- '"'"'<|BOT|>assistant\n<think>\n'"'"' -}}{%- endif -%}'
+    ```
+
+After the service starts, it will listen on `localhost:9999`.
+
+
+### 🧪 Client Examples
+
+Get the example code and run it:
+```bash
+# Clone the repository containing example scripts
+git clone https://github.com/stepfun-ai/Step-Audio-R1.git r1-scripts
+
+# Run the example
+cd r1-scripts
+python examples-vllm_r1.py
+```
+
+
+## Citation
+
+```
+@article{tian2025step,
+  title={Step-Audio-R1 Technical Report},
+  author={Tian, Fei and Zhang, Xiangyu Tony and Zhang, Yuxin and Zhang, Haoyang and Li, Yuxin and Liu, Daijiao and Deng, Yayue and Wu, Donghang and Chen, Jun and Zhao, Liang and others},
+  journal={arXiv preprint arXiv:2511.15848},
+  year={2025}
+}
+
+```

models/Step-Audio-R1-NVFP4A16/chat_template.jinja

@@ -0,0 +1,74 @@
+{%- if tools %}
+    {{- '<|BOT|>system
+' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] + '<|EOT|>' }}
+    {%- else %}
+        {{- 'You are a helpful assistant. Please think step by step and provide your reasoning process within <think> </think> tags, followed by your final answer. Format: <think>your reasoning here</think>your final answer<|EOT|>' }}
+    {%- endif %}
+    {{- '<|BOT|>' }}
+    {{- "tool_json_schemas
+" }}
+    {{- tools | tojson }}
+    {{- '<|EOT|>' }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|BOT|>system
+' + messages[0]['content'] + '<|EOT|>' }}
+    {%- else %}
+        {{- '<|BOT|>system
+You are a helpful assistant. Please think step by step and provide your reasoning process within <think> </think> tags, followed by your final answer. Format: <think>your reasoning here</think>your final answer<|EOT|>' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if message["role"] == "user" %}
+        {{- '<|BOT|>human
+' + message["content"] + '<|EOT|>' }}
+    {%- elif (message["role"] == "system" and not loop.first) or (message["role"] == "assistant" and not message["tool_calls"]) %}
+        {{- '<|BOT|>' + message["role"] + '
+' + message["content"] + '<|EOT|>' }}
+    {%- elif message["role"] == "assistant" %}
+        {{- '<|BOT|>' + message["role"] + '
+' }}
+        {%- if message["content"] %}
+            {{- message["content"] }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call["function"] is defined %}
+                {%- set tool_call = tool_call["function"] %}
+            {%- endif %}
+            {{- '<|CALL_START|>' + 'function
+' + tool_call["name"] + '
+' }}
+            {{- tool_call["arguments"] | tojson }}
+            {{- '<|CALL_END|>' }}
+        {%- endfor %}
+        {{- '<|EOT|>' }}
+    {%- elif message["role"] == "tool" %}
+        {{- '<|BOT|>' }}
+        {%- set ns = namespace(function_name="tool") %}
+        {%- if message["tool_call_id"] %}
+            {%- for prev_msg in messages %}
+                {%- if prev_msg["role"] == "assistant" and prev_msg["tool_calls"] %}
+                    {%- for tool_call in prev_msg["tool_calls"] %}
+                        {%- if tool_call["id"] == message["tool_call_id"] %}
+                            {%- if tool_call["function"] is defined %}
+                                {%- set ns.function_name = tool_call["function"]["name"] %}
+                            {%- endif %}
+                        {%- endif %}
+                    {%- endfor %}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- 'function_output
+' + ns.function_name + '
+' }}
+        {{- message["content"] }}
+        {{- '<|EOT|>' }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|BOT|>assistant
+<think>
+' }}
+{%- endif %}

models/Step-Audio-R1-NVFP4A16/config.json

@@ -0,0 +1,348 @@
+{
+  "architectures": [
+    "StepAudio2ForCausalLM"
+  ],
+  "audio_encoder_config": {
+    "adapter_stride": 2,
+    "kernel_size": 3,
+    "llm_dim": 5120,
+    "model_type": "step_audio_2_encoder",
+    "n_audio_ctx": 1500,
+    "n_audio_head": 20,
+    "n_audio_layer": 32,
+    "n_audio_state": 1280,
+    "n_codebook_size": 4096,
+    "n_mels": 128
+  },
+  "auto_map": {
+    "AutoConfig": "configuration_step_audio_2.StepAudio2Config",
+    "AutoModelForCausalLM": "modeling_step_audio_2.StepAudio2ForCausalLM"
+  },
+  "dtype": "bfloat16",
+  "max_window_layers": null,
+  "model_type": "step_audio_2",
+  "quantization_config": {
+    "config_groups": {
+      "group_0": {
+        "format": "nvfp4-pack-quantized",
+        "input_activations": null,
+        "output_activations": null,
+        "targets": [
+          "Linear"
+        ],
+        "weights": {
+          "actorder": null,
+          "block_structure": null,
+          "dynamic": false,
+          "group_size": 16,
+          "num_bits": 4,
+          "observer": "minmax",
+          "observer_kwargs": {},
+          "strategy": "tensor_group",
+          "symmetric": true,
+          "type": "float"
+        }
+      }
+    },
+    "format": "nvfp4-pack-quantized",
+    "global_compression_ratio": null,
+    "ignore": [
+      "encoder.blocks.0.attn.query",
+      "encoder.blocks.0.attn.key",
+      "encoder.blocks.0.attn.value",
+      "encoder.blocks.0.attn.out",
+      "encoder.blocks.0.mlp.0",
+      "encoder.blocks.0.mlp.2",
+      "encoder.blocks.1.attn.query",
+      "encoder.blocks.1.attn.key",
+      "encoder.blocks.1.attn.value",
+      "encoder.blocks.1.attn.out",
+      "encoder.blocks.1.mlp.0",
+      "encoder.blocks.1.mlp.2",
+      "encoder.blocks.2.attn.query",
+      "encoder.blocks.2.attn.key",
+      "encoder.blocks.2.attn.value",
+      "encoder.blocks.2.attn.out",
+      "encoder.blocks.2.mlp.0",
+      "encoder.blocks.2.mlp.2",
+      "encoder.blocks.3.attn.query",
+      "encoder.blocks.3.attn.key",
+      "encoder.blocks.3.attn.value",
+      "encoder.blocks.3.attn.out",
+      "encoder.blocks.3.mlp.0",
+      "encoder.blocks.3.mlp.2",
+      "encoder.blocks.4.attn.query",
+      "encoder.blocks.4.attn.key",
+      "encoder.blocks.4.attn.value",
+      "encoder.blocks.4.attn.out",
+      "encoder.blocks.4.mlp.0",
+      "encoder.blocks.4.mlp.2",
+      "encoder.blocks.5.attn.query",
+      "encoder.blocks.5.attn.key",
+      "encoder.blocks.5.attn.value",
+      "encoder.blocks.5.attn.out",
+      "encoder.blocks.5.mlp.0",
+      "encoder.blocks.5.mlp.2",
+      "encoder.blocks.6.attn.query",
+      "encoder.blocks.6.attn.key",
+      "encoder.blocks.6.attn.value",
+      "encoder.blocks.6.attn.out",
+      "encoder.blocks.6.mlp.0",
+      "encoder.blocks.6.mlp.2",
+      "encoder.blocks.7.attn.query",
+      "encoder.blocks.7.attn.key",
+      "encoder.blocks.7.attn.value",
+      "encoder.blocks.7.attn.out",
+      "encoder.blocks.7.mlp.0",
+      "encoder.blocks.7.mlp.2",
+      "encoder.blocks.8.attn.query",
+      "encoder.blocks.8.attn.key",
+      "encoder.blocks.8.attn.value",
+      "encoder.blocks.8.attn.out",
+      "encoder.blocks.8.mlp.0",
+      "encoder.blocks.8.mlp.2",
+      "encoder.blocks.9.attn.query",
+      "encoder.blocks.9.attn.key",
+      "encoder.blocks.9.attn.value",
+      "encoder.blocks.9.attn.out",
+      "encoder.blocks.9.mlp.0",
+      "encoder.blocks.9.mlp.2",
+      "encoder.blocks.10.attn.query",
+      "encoder.blocks.10.attn.key",
+      "encoder.blocks.10.attn.value",
+      "encoder.blocks.10.attn.out",
+      "encoder.blocks.10.mlp.0",
+      "encoder.blocks.10.mlp.2",
+      "encoder.blocks.11.attn.query",
+      "encoder.blocks.11.attn.key",
+      "encoder.blocks.11.attn.value",
+      "encoder.blocks.11.attn.out",
+      "encoder.blocks.11.mlp.0",
+      "encoder.blocks.11.mlp.2",
+      "encoder.blocks.12.attn.query",
+      "encoder.blocks.12.attn.key",
+      "encoder.blocks.12.attn.value",
+      "encoder.blocks.12.attn.out",
+      "encoder.blocks.12.mlp.0",
+      "encoder.blocks.12.mlp.2",
+      "encoder.blocks.13.attn.query",
+      "encoder.blocks.13.attn.key",
+      "encoder.blocks.13.attn.value",
+      "encoder.blocks.13.attn.out",
+      "encoder.blocks.13.mlp.0",
+      "encoder.blocks.13.mlp.2",
+      "encoder.blocks.14.attn.query",
+      "encoder.blocks.14.attn.key",
+      "encoder.blocks.14.attn.value",
+      "encoder.blocks.14.attn.out",
+      "encoder.blocks.14.mlp.0",
+      "encoder.blocks.14.mlp.2",
+      "encoder.blocks.15.attn.query",
+      "encoder.blocks.15.attn.key",
+      "encoder.blocks.15.attn.value",
+      "encoder.blocks.15.attn.out",
+      "encoder.blocks.15.mlp.0",
+      "encoder.blocks.15.mlp.2",
+      "encoder.blocks.16.attn.query",
+      "encoder.blocks.16.attn.key",
+      "encoder.blocks.16.attn.value",
+      "encoder.blocks.16.attn.out",
+      "encoder.blocks.16.mlp.0",
+      "encoder.blocks.16.mlp.2",
+      "encoder.blocks.17.attn.query",
+      "encoder.blocks.17.attn.key",
+      "encoder.blocks.17.attn.value",
+      "encoder.blocks.17.attn.out",
+      "encoder.blocks.17.mlp.0",
+      "encoder.blocks.17.mlp.2",
+      "encoder.blocks.18.attn.query",
+      "encoder.blocks.18.attn.key",
+      "encoder.blocks.18.attn.value",
+      "encoder.blocks.18.attn.out",
+      "encoder.blocks.18.mlp.0",
+      "encoder.blocks.18.mlp.2",
+      "encoder.blocks.19.attn.query",
+      "encoder.blocks.19.attn.key",
+      "encoder.blocks.19.attn.value",
+      "encoder.blocks.19.attn.out",
+      "encoder.blocks.19.mlp.0",
+      "encoder.blocks.19.mlp.2",
+      "encoder.blocks.20.attn.query",
+      "encoder.blocks.20.attn.key",
+      "encoder.blocks.20.attn.value",
+      "encoder.blocks.20.attn.out",
+      "encoder.blocks.20.mlp.0",
+      "encoder.blocks.20.mlp.2",
+      "encoder.blocks.21.attn.query",
+      "encoder.blocks.21.attn.key",
+      "encoder.blocks.21.attn.value",
+      "encoder.blocks.21.attn.out",
+      "encoder.blocks.21.mlp.0",
+      "encoder.blocks.21.mlp.2",
+      "encoder.blocks.22.attn.query",
+      "encoder.blocks.22.attn.key",
+      "encoder.blocks.22.attn.value",
+      "encoder.blocks.22.attn.out",
+      "encoder.blocks.22.mlp.0",
+      "encoder.blocks.22.mlp.2",
+      "encoder.blocks.23.attn.query",
+      "encoder.blocks.23.attn.key",
+      "encoder.blocks.23.attn.value",
+      "encoder.blocks.23.attn.out",
+      "encoder.blocks.23.mlp.0",
+      "encoder.blocks.23.mlp.2",
+      "encoder.blocks.24.attn.query",
+      "encoder.blocks.24.attn.key",
+      "encoder.blocks.24.attn.value",
+      "encoder.blocks.24.attn.out",
+      "encoder.blocks.24.mlp.0",
+      "encoder.blocks.24.mlp.2",
+      "encoder.blocks.25.attn.query",
+      "encoder.blocks.25.attn.key",
+      "encoder.blocks.25.attn.value",
+      "encoder.blocks.25.attn.out",
+      "encoder.blocks.25.mlp.0",
+      "encoder.blocks.25.mlp.2",
+      "encoder.blocks.26.attn.query",
+      "encoder.blocks.26.attn.key",
+      "encoder.blocks.26.attn.value",
+      "encoder.blocks.26.attn.out",
+      "encoder.blocks.26.mlp.0",
+      "encoder.blocks.26.mlp.2",
+      "encoder.blocks.27.attn.query",
+      "encoder.blocks.27.attn.key",
+      "encoder.blocks.27.attn.value",
+      "encoder.blocks.27.attn.out",
+      "encoder.blocks.27.mlp.0",
+      "encoder.blocks.27.mlp.2",
+      "encoder.blocks.28.attn.query",
+      "encoder.blocks.28.attn.key",
+      "encoder.blocks.28.attn.value",
+      "encoder.blocks.28.attn.out",
+      "encoder.blocks.28.mlp.0",
+      "encoder.blocks.28.mlp.2",
+      "encoder.blocks.29.attn.query",
+      "encoder.blocks.29.attn.key",
+      "encoder.blocks.29.attn.value",
+      "encoder.blocks.29.attn.out",
+      "encoder.blocks.29.mlp.0",
+      "encoder.blocks.29.mlp.2",
+      "encoder.blocks.30.attn.query",
+      "encoder.blocks.30.attn.key",
+      "encoder.blocks.30.attn.value",
+      "encoder.blocks.30.attn.out",
+      "encoder.blocks.30.mlp.0",
+      "encoder.blocks.30.mlp.2",
+      "encoder.blocks.31.attn.query",
+      "encoder.blocks.31.attn.key",
+      "encoder.blocks.31.attn.value",
+      "encoder.blocks.31.attn.out",
+      "encoder.blocks.31.mlp.0",
+      "encoder.blocks.31.mlp.2",
+      "adapter.linear1",
+      "adapter.linear2",
+      "lm_head"
+    ],
+    "kv_cache_scheme": null,
+    "quant_method": "compressed-tensors",
+    "quantization_status": "compressed",
+    "sparsity_config": {},
+    "transform_config": {},
+    "version": "0.12.2"
+  },
+  "sliding_window": 2048,
+  "text_config": {
+    "architectures": [
+      "Qwen2ForCausalLM"
+    ],
+    "attention_dropout": 0.0,
+    "dtype": "bfloat16",
+    "hidden_act": "silu",
+    "hidden_size": 5120,
+    "initializer_range": 0.02,
+    "intermediate_size": 27648,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 65536,
+    "max_window_layers": 28,
+    "model_type": "qwen2",
+    "num_attention_heads": 40,
+    "num_hidden_layers": 64,
+    "num_attention_groups": 8,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": null,
+    "rope_theta": 1000000.0,
+    "sliding_window": null,
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 158720
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "4.56.2",
+  "use_sliding_window": false
+}

models/Step-Audio-R1-NVFP4A16/configuration_step_audio_2.py

@@ -0,0 +1,128 @@
+from typing import Optional, Union
+
+from transformers import Qwen2Config
+from transformers.configuration_utils import PretrainedConfig
+
+
+class StepAudio2EncoderConfig(PretrainedConfig):
+    model_type = "step_audio_2_encoder"
+
+    def __init__(
+        self,
+        n_mels=128,
+        n_audio_ctx=1500,
+        n_audio_state=512,
+        n_audio_head=8,
+        n_audio_layer=6,
+        llm_dim=4096,
+        kernel_size=3,
+        adapter_stride=2,
+        **kwargs,
+    ):
+        self.n_mels      = n_mels
+        self.n_audio_ctx = n_audio_ctx
+        self.n_audio_state = n_audio_state
+        self.n_audio_head = n_audio_head
+        self.n_audio_layer = n_audio_layer
+        self.llm_dim     = llm_dim
+        self.kernel_size = kernel_size
+        self.adapter_stride = adapter_stride
+        super().__init__(**kwargs)
+
+class StepAudio2TextConfig(PretrainedConfig):
+    model_type = "step_audio_2_text"
+
+    def __init__(
+        self,
+        vocab_size=64012,
+        hidden_size=4096,
+        intermediate_size=11008,
+        num_hidden_layers=48,
+        num_attention_heads=32,
+        num_attention_groups=4,
+        num_key_value_heads=4,
+        hidden_act="silu",
+        max_position_embeddings=8192,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        rope_theta=1000000.0,
+        rope_scaling=None,
+        eos_token_id=None,
+        **kwargs
+    ):
+
+        if eos_token_id is not None:
+            if isinstance(eos_token_id, list):
+                eos_token_id = list(set([151643, 151645, 151665] + eos_token_id))
+            else:
+                eos_token_id = [151643, 151645, 151665, eos_token_id]
+        else:
+            eos_token_id = [151643, 151645, 151665]
+
+        super().__init__(
+            eos_token_id=eos_token_id,
+            **kwargs)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_attention_groups = num_attention_groups
+        self.num_key_value_heads = num_key_value_heads
+        assert self.num_attention_groups == self.num_key_value_heads, "num_attention_groups must be equal to num_key_value_heads"
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+
+        self.text_config = Qwen2Config(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_hidden_layers=num_hidden_layers,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            hidden_act=hidden_act,
+            max_position_embeddings=max_position_embeddings,
+            initializer_range=initializer_range,
+            rms_norm_eps=rms_norm_eps,
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            architectures=["Qwen2ForCausalLM"],
+            torch_dtype=getattr(self, "torch_dtype", "bfloat16"),
+        )
+
+class StepAudio2Config(PretrainedConfig):
+    model_type = "step_audio_2"
+    architectures = ["StepAudio2ForCausalLM"]
+
+    def __init__(
+        self,
+        audio_encoder_config :Optional[Union[dict, StepAudio2EncoderConfig]] = None,
+        text_config: Optional[Union[dict, StepAudio2TextConfig]] = None,
+        use_sliding_window: bool = False,
+        sliding_window: Optional[int] = 2048,
+        max_window_layers: Optional[int] = None,
+        **kwargs
+    ):
+        kwargs.setdefault("use_sliding_window", use_sliding_window)
+        kwargs.setdefault("sliding_window", sliding_window)
+        if max_window_layers is None:
+            max_window_layers = kwargs.get("num_hidden_layers", None)
+        kwargs.setdefault("max_window_layers", max_window_layers)
+        super().__init__(**kwargs)
+
+        if text_config is None:
+            text_config = StepAudio2TextConfig().text_config
+        elif isinstance(text_config, dict):
+            text_config = StepAudio2TextConfig(**text_config).text_config
+
+        self.text_config = text_config
+
+        if audio_encoder_config is None:
+            self.audio_encoder_config = StepAudio2EncoderConfig()
+        elif isinstance(audio_encoder_config, dict):
+            self.audio_encoder_config = StepAudio2EncoderConfig(**audio_encoder_config)

models/Step-Audio-R1-NVFP4A16/generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 151643,
+  "eos_token_id": [
+    151643, 
+    151665
+  ],
+  "pad_token_id": 151643,
+  "transformers_version": "4.56.2"
+}

models/Step-Audio-R1-NVFP4A16/model-00001-of-00005.safetensors

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

models/Step-Audio-R1-NVFP4A16/model-00002-of-00005.safetensors

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

models/Step-Audio-R1-NVFP4A16/model-00003-of-00005.safetensors

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

models/Step-Audio-R1-NVFP4A16/model-00004-of-00005.safetensors

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

models/Step-Audio-R1-NVFP4A16/model-00005-of-00005.safetensors

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