Upload folder using huggingface_hub

README.md

@@ -51,7 +51,7 @@ base_model:
 
 >>> import torch
 >>> with torch.no_grad():
-...     model_inputs = processor(text=text, audio=audio)
+...     model_inputs = processor(text=text, audio=audio, sampling_rate=sr))
 ...     generation = model.generate(**model_inputs)
 ...     output = processor.batch_decode(generation, skip_special_tokens=True)
 

config.json

@@ -37,15 +37,10 @@
     "AutoConfig": "configuration_midashenglm.MiAudioLLMHFConfig",
     "AutoModelForCausalLM": "modeling_midashenglm.DashengQwen25OmniModelInstruct"
   },
-  "freeze": null,
-  "gradient_checkpoint_decoder": false,
-  "lora": null,
-  "model": "DashengQwen25OmniModelInstruct",
   "model_type": "miaudiollm",
   "resize_tokenizer": false,
   "subsample_factor": 5,
-  "text_model_config": {
-    "_attn_implementation_autoset": true,
+  "text_config": {
     "attention_dropout": 0.0,
     "hidden_act": "silu",
     "hidden_size": 2048,

configuration_midashenglm.py

@@ -1,5 +1,5 @@
 from ast import Dict
-from typing import Literal, Tuple, Union
+from typing import Tuple, Union
 
 from transformers import PretrainedConfig
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
@@ -66,22 +66,16 @@ class MiAudioLLMHFConfig(PretrainedConfig):
 
     def __init__(
         self,
-        model: str = "DashengQwen2ModelInstruct",
         audio_encoder_config: Dict = {},
-        freeze: Literal["audio", "text"] | str | None = None,
-        lora: Literal["encoder", "decoder"] | None = None,
         subsample_factor: int = 5,
-        text_model_config: Dict = None,
+        text_config: Dict = None,
         **kwargs,
     ):
-        self.model = model
         self.audio_encoder_config = DashengConfig(**audio_encoder_config)
-        self.freeze = freeze
-        self.lora = lora
         self.subsample_factor = subsample_factor
-        self.text_model_config = (
-            Qwen2_5OmniTextConfig(**text_model_config)
-            if text_model_config
+        self.text_config = (
+            Qwen2_5OmniTextConfig(**text_config)
+            if text_config
             else Qwen2_5OmniTextConfig()
         )
         super().__init__(**kwargs)

model.safetensors.index.json

@@ -390,20 +390,20 @@
     "audio_encoder.freq_pos_embed": "model-00001-of-00002.safetensors",
     "audio_encoder.front_end.0.mel_scale.fb": "model-00001-of-00002.safetensors",
     "audio_encoder.front_end.0.spectrogram.window": "model-00001-of-00002.safetensors",
-    "audio_encoder.init_bn.1.bias": "model-00001-of-00002.safetensors",
-    "audio_encoder.init_bn.1.num_batches_tracked": "model-00001-of-00002.safetensors",
-    "audio_encoder.init_bn.1.running_mean": "model-00001-of-00002.safetensors",
-    "audio_encoder.init_bn.1.running_var": "model-00001-of-00002.safetensors",
-    "audio_encoder.init_bn.1.weight": "model-00001-of-00002.safetensors",
+    "audio_encoder.init_bn.bias": "model-00001-of-00002.safetensors",
+    "audio_encoder.init_bn.num_batches_tracked": "model-00001-of-00002.safetensors",
+    "audio_encoder.init_bn.running_mean": "model-00001-of-00002.safetensors",
+    "audio_encoder.init_bn.running_var": "model-00001-of-00002.safetensors",
+    "audio_encoder.init_bn.weight": "model-00001-of-00002.safetensors",
     "audio_encoder.norm.bias": "model-00001-of-00002.safetensors",
     "audio_encoder.norm.weight": "model-00001-of-00002.safetensors",
     "audio_encoder.patch_embed.proj.bias": "model-00001-of-00002.safetensors",
     "audio_encoder.patch_embed.proj.weight": "model-00001-of-00002.safetensors",
     "audio_encoder.time_pos_embed": "model-00001-of-00002.safetensors",
-    "audio_projector.net.0.bias": "model-00002-of-00002.safetensors",
-    "audio_projector.net.0.weight": "model-00002-of-00002.safetensors",
-    "audio_projector.net.2.bias": "model-00002-of-00002.safetensors",
-    "audio_projector.net.2.weight": "model-00002-of-00002.safetensors",
+    "audio_projector.net.0.bias": "model-00001-of-00002.safetensors",
+    "audio_projector.net.0.weight": "model-00001-of-00002.safetensors",
+    "audio_projector.net.2.bias": "model-00001-of-00002.safetensors",
+    "audio_projector.net.2.weight": "model-00001-of-00002.safetensors",
     "decoder.lm_head.weight": "model-00002-of-00002.safetensors",
     "decoder.model.embed_tokens.weight": "model-00001-of-00002.safetensors",
     "decoder.model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors",
@@ -442,11 +442,11 @@
     "decoder.model.layers.10.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
     "decoder.model.layers.10.self_attn.v_proj.bias": "model-00001-of-00002.safetensors",
     "decoder.model.layers.10.self_attn.v_proj.weight": "model-00001-of-00002.safetensors",
-    "decoder.model.layers.11.input_layernorm.weight": "model-00001-of-00002.safetensors",
-    "decoder.model.layers.11.mlp.down_proj.weight": "model-00001-of-00002.safetensors",
+    "decoder.model.layers.11.input_layernorm.weight": "model-00002-of-00002.safetensors",
+    "decoder.model.layers.11.mlp.down_proj.weight": "model-00002-of-00002.safetensors",
     "decoder.model.layers.11.mlp.gate_proj.weight": "model-00001-of-00002.safetensors",
-    "decoder.model.layers.11.mlp.up_proj.weight": "model-00001-of-00002.safetensors",
-    "decoder.model.layers.11.post_attention_layernorm.weight": "model-00001-of-00002.safetensors",
+    "decoder.model.layers.11.mlp.up_proj.weight": "model-00002-of-00002.safetensors",
+    "decoder.model.layers.11.post_attention_layernorm.weight": "model-00002-of-00002.safetensors",
     "decoder.model.layers.11.self_attn.k_proj.bias": "model-00001-of-00002.safetensors",
     "decoder.model.layers.11.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "decoder.model.layers.11.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",

modeling_midashenglm.py

@@ -249,21 +249,12 @@ class Block(nn.Module):
         return x
 
 
-# TODO
-class RearranceReplace(nn.Module):
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # rearrange(x, "b c f t -> b f c t")
-        # or
-        # rearrange(x, "b f c t -> b c f t")
-        return torch.permute(x, (0, 2, 1, 3))
+class AudioTransformer(PreTrainedModel):
+    config_class = DashengConfig
 
+    def __init__(self, config: DashengConfig):
+        super().__init__(config)
 
-class AudioTransformer(nn.Module):
-    def __init__(
-        self,
-        config: DashengConfig,
-    ):
-        super().__init__()
         self.target_length = config.target_length
         self.embed_dim = config.embed_dim
         self.hop_length = config.hop_length
@@ -282,13 +273,7 @@ class AudioTransformer(nn.Module):
             audio_transforms.AmplitudeToDB(top_db=120),
         )
 
-        self.init_bn = nn.Sequential(
-            # Rearrange("b c f t -> b f c t"),
-            RearranceReplace(),
-            nn.BatchNorm2d(config.n_mels, momentum=0.01),
-            # Rearrange("b f c t -> b c f t"),
-            RearranceReplace(),
-        )
+        self.init_bn = nn.BatchNorm2d(config.n_mels, momentum=0.01)
 
         self.patch_embed = AudioPatchEmbed(
             input_size=(config.n_mels, config.target_length),
@@ -327,6 +312,8 @@ class AudioTransformer(nn.Module):
         )
         self.norm = norm_layer(config.embed_dim)
 
+        self.post_init()
+
     def forward_features(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
         t = x.shape[-1]
         x = x + self.time_pos_embed[:, :, :, :t]
@@ -357,7 +344,9 @@ class AudioTransformer(nn.Module):
         x = self.front_end(x)
         target_length_in_patches = self.target_length // 4
         x = x.unsqueeze(1)
+        x = torch.permute(x, (0, 2, 1, 3))
         x = self.init_bn(x)
+        x = torch.permute(x, (0, 2, 1, 3))
 
         x = self.patch_embed(x)
         t = x.shape[-1]
@@ -427,23 +416,21 @@ class DashengQwen25OmniModelInstructOutput(ModelOutput):
 
 class Decoder(PreTrainedModel, GenerationMixin):
     config_class = Qwen2_5OmniTextConfig
+    _supports_flash_attn_2 = Qwen2_5OmniThinkerTextModel._supports_flash_attn_2
+    _supports_sdpa = Qwen2_5OmniThinkerTextModel._supports_sdpa
+    _supports_flex_attn = Qwen2_5OmniThinkerTextModel._supports_flex_attn
+    _supports_cache_class = Qwen2_5OmniThinkerTextModel._supports_cache_class
+    _supports_static_cache = Qwen2_5OmniThinkerTextModel._supports_static_cache
+    _supports_quantized_cache = Qwen2_5OmniThinkerTextModel._supports_quantized_cache
 
     def __init__(self, config: Qwen2_5OmniTextConfig):
         super().__init__(config)
-        self.model = Qwen2_5OmniThinkerTextModel._from_config(
-            config,
-            attn_implementation="sdpa",  # TODO
-        )
+        self.model = Qwen2_5OmniThinkerTextModel._from_config(config)
         self.lm_head = nn.Linear(
             config.hidden_size,
             config.vocab_size,
             bias=False,
         )
-        # TODO fix dtype
-        self.lm_head.weight.data = self.lm_head.weight.data.to(
-            self.model.embed_tokens.weight.dtype
-        )
-        # TODO tie weight?
         self.post_init()
 
     def forward(
@@ -481,30 +468,25 @@ class Decoder(PreTrainedModel, GenerationMixin):
 
 class DashengQwen25OmniModelInstruct(PreTrainedModel):
     config_class = MiAudioLLMHFConfig
+    _supports_flash_attn_2 = Qwen2_5OmniThinkerTextModel._supports_flash_attn_2
+    _supports_sdpa = Qwen2_5OmniThinkerTextModel._supports_sdpa
+    _supports_flex_attn = Qwen2_5OmniThinkerTextModel._supports_flex_attn
+    _supports_cache_class = Qwen2_5OmniThinkerTextModel._supports_cache_class
+    _supports_static_cache = Qwen2_5OmniThinkerTextModel._supports_static_cache
+    _supports_quantized_cache = Qwen2_5OmniThinkerTextModel._supports_quantized_cache
 
     def __init__(self, config: MiAudioLLMHFConfig):
         super().__init__(config)
 
-        freeze = config.freeze
-        lora = config.lora
-        subsample_factor = config.subsample_factor
-
-        self.subsample_factor = subsample_factor
-        self.lora = lora
-        # Encoder part
-        self.audio_encoder = AudioTransformer(config.audio_encoder_config)
-        assert lora != "encoder"
-
-        # decoder
-        self.decoder = Decoder(config.text_model_config)
-        assert lora != "decoder"
-        assert freeze is None
-
-        # audio projector
+        self.audio_encoder = AudioTransformer._from_config(config.audio_encoder_config)
         self.audio_projector = AudioProjectorSubsample(
             self.audio_encoder.embed_dim,
-            config.text_model_config.hidden_size,
-            self.subsample_factor,
+            config.text_config.hidden_size,
+            config.subsample_factor,
+        )
+        self.decoder = Decoder._from_config(
+            config.text_config,
+            attn_implementation=config._attn_implementation,
         )
 
         self.post_init()

processing_midashenglm.py

@@ -55,32 +55,35 @@ class MiAudioLLMProcessor(ProcessorMixin):
         tokenizer: Qwen2Tokenizer | Qwen2TokenizerFast | None = None,
         model_subsampling: int = 5,
         chat_template: str | None = None,
-        # TODO 是否可以移除？
-        audio_token: str = "<|AUDIO|>",
-        audio_bos_token: str = "<|audio_bos|>",
-        audio_eos_token: str = "<|audio_eos|>",
+        audio_token: str | None = None,
+        audio_bos_token: str | None = None,
+        audio_eos_token: str | None = None,
     ):
-        if chat_template is None:
-            chat_template = self.default_chat_template
         assert tokenizer is not None, "Tokenizer Needs to be passed"
-        self.audio_token = (
-            tokenizer.audio_token if hasattr(tokenizer, "audio_token") else audio_token
+        assert audio_token is not None or hasattr(tokenizer, "audio_token"), (
+            "Either `audio_token` must be provided or tokenizer must have `audio_token` attribute."
         )
-        self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token)
-        self.audio_bos_token = (
-            tokenizer.audio_bos_token
-            if hasattr(tokenizer, "audio_bos_token")
-            else audio_bos_token
+        assert audio_bos_token is not None or hasattr(tokenizer, "audio_bos_token"), (
+            "Either `audio_bos_token` must be provided or tokenizer must have `audio_bos_token` attribute."
         )
-        self.audio_eos_token = (
-            tokenizer.audio_eos_token
-            if hasattr(tokenizer, "audio_eos_token")
-            else audio_eos_token
+        assert audio_eos_token is not None or hasattr(tokenizer, "audio_eos_token"), (
+            "Either `audio_eos_token` must be provided or tokenizer must have `audio_eos_token` attribute."
         )
+
+        if chat_template is None:
+            chat_template = self.default_chat_template
+
+        self.audio_token: str = audio_token or tokenizer.audio_token
+        self.audio_bos_token = audio_bos_token or tokenizer.audio_bos_token
+        self.audio_eos_token = audio_eos_token or tokenizer.audio_eos_token
+        self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token)
         self.model_subsampling = model_subsampling
-        # Fix Normalization
-        if feature_extractor is not None and feature_extractor.do_normalize is True:
-            feature_extractor.do_normalize = False
+
+        if feature_extractor is not None:
+            assert not feature_extractor.do_normalize, (
+                "This model does not use normalization. Please set `do_normalize=False` in the feature extractor."
+            )
+
         super().__init__(feature_extractor, tokenizer, chat_template=chat_template)
 
     def __call__(