first commit

config.json

@@ -0,0 +1,33 @@
+{
+  "architectures": [
+    "UltravoxModel"
+  ],
+  "audio_model_id": "AlexHung29629/wav2vec2-lv-60-espeak-cv-ft",
+  "auto_map": {
+    "AutoConfig": "ultravox_config.UltravoxConfig",
+    "AutoModel": "ultravox_model.UltravoxModel",
+    "AutoProcessor": "ultravox_processing.UltravoxProcessor"
+  },
+  "custom_pipelines": {
+    "ultravox-pipeline": {
+      "impl": "ultravox_pipeline.UltravoxPipeline",
+      "pt": [
+        "AutoModel"
+      ],
+      "tf": [],
+      "type": "multimodal"
+    }
+  },
+  "hidden_size": 4096,
+  "ignore_index": -100,
+  "initializer_range": 0.02,
+  "model_type": "ultravox",
+  "norm_init": 0.4,
+  "pad_token_id": 128009,
+  "projector_act": "swiglu",
+  "stack_factor": 8,
+  "text_model_id": "voidful/Llama-3.2-8B-Instruct",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.46.2",
+  "vocab_size": 128256
+}

generation_config.json

@@ -0,0 +1,11 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 128000,
+  "eos_token_id": [
+    128001,
+    128008,
+    128009
+  ],
+  "pad_token_id": 128009,
+  "transformers_version": "4.46.2"
+}

ultravox_config.py

@@ -0,0 +1,166 @@
+import dataclasses
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+import transformers
+
+
+@dataclasses.dataclass
+class LoraConfigSimplified:
+    """
+    Low Rank Approximation (LoRA) configuration.
+
+    Used for language and audio models separately.
+    """
+
+    # The rank of the approximation
+    r: int = 0
+    lora_alpha: float = 8
+    target_modules: Optional[List[str]] = dataclasses.field(
+        default_factory=lambda: ["k_proj", "q_proj", "linear_k", "linear_q"]
+    )
+
+
+class LossFunction(str, Enum):
+    CrossEntropy = "ce"
+    KL_Divergence = "kl"
+
+
+@dataclasses.dataclass
+class LossConfig:
+    loss_function: LossFunction = LossFunction.CrossEntropy
+    kl_temperature: float = 2.0
+
+    @property
+    def requires_alt_fields(self):
+        return self.loss_function == LossFunction.KL_Divergence
+
+
+class UltravoxConfig(transformers.PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`UltravoxForConditionalGeneration`]. It is used to instantiate an
+    Ultravox model according to the specified arguments, defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        audio_config (`Wav2Vec2Config`,  *optional*):
+            Custom audio config or dict
+        text_config (`Union[AutoConfig, dict]`, *optional*):
+            The config object of the text backbone. Can be any of `LlamaConfig` or `MistralConfig`.
+        ignore_index (`int`, *optional*, defaults to -100):
+            The ignore index for the loss function.
+        audio_token_index (`int`, *optional*, defaults to 32000):
+            The audio token index to encode the audio prompt.
+        stack_factor (`int`, *optional*, defaults to 8):
+            Audio downsampling factor for the multimodal projector.
+        norm_init (`float`, *optional*, defaults to 0.4):
+            The initialization value for the layer normalization.
+        projector_act (`str`, *optional*, defaults to `"swiglu"`):
+            The activation function used by the multimodal projector.
+        text_model_lora_config (`LoraConfigSimplified`, *optional*):
+            The LoRA configuration for finetuning the text model.
+        audio_model_lora_config (`LoraConfigSimplified`, *optional*):
+            The LoRA configuration for finetuning the audio model.
+
+
+    Example:
+
+    ```python
+    >>> from transformers import UltravoxForConditionalGeneration, Wav2Vec2Config, UltravoxConfig, LlamaConfig
+
+    >>> # Initializing an audio encoder config
+    >>> audio_config = Wav2Vec2Config()
+
+    >>> # Initializing a Llama config
+    >>> text_config = LlamaConfig()
+
+    >>> # Initializing a default configuration
+    >>> configuration = UltravoxConfig(audio_config, text_config)
+
+    >>> # Initializing a completely untrained model from the configuration
+    >>> model = UltravoxForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+
+    >>> # Initialize a model from pretrained checkpoints and random projector weights
+    >>> config = UltravoxConfig(audio_model_id="facebook/wav2vec2-base-960h", text_model_id="meta-llama/Llama-2-7b-chat-hf")
+    ```"""
+
+    model_type = "ultravox"
+    is_composition = False
+
+    def __init__(
+        self,
+        audio_config: Optional[Dict[str, Any]] = None,
+        text_config: Optional[Dict[str, Any]] = None,
+        audio_model_id: Optional[str] = None,
+        text_model_id: Optional[str] = None,
+        ignore_index: int = -100,
+        hidden_size: int = 4096,
+        stack_factor: int = 8,
+        norm_init: float = 0.4,
+        projector_act: str = "swiglu",
+        text_model_lora_config: Optional[LoraConfigSimplified] = None,
+        audio_model_lora_config: Optional[LoraConfigSimplified] = None,
+        **kwargs,
+    ):
+        self.ignore_index = ignore_index
+
+        self.audio_model_id = audio_model_id
+        self.text_model_id = text_model_id
+
+        self.hidden_size = hidden_size
+        self.stack_factor = stack_factor
+        self.norm_init = norm_init
+        self.projector_act = projector_act
+
+        if text_model_id is not None:
+            self.text_config: transformers.LlamaConfig = (
+                transformers.AutoConfig.from_pretrained(text_model_id)
+            )
+        else:
+            text_config = text_config or {}
+            self.text_config = transformers.CONFIG_MAPPING[
+                text_config.get("model_type", "llama")
+            ](**text_config)
+
+        if audio_model_id is not None:
+            self.audio_config: transformers.PretrainedConfig = (
+                transformers.AutoConfig.from_pretrained(audio_model_id)
+            )
+        else:
+            audio_config = audio_config or {}
+            self.audio_config = transformers.CONFIG_MAPPING[
+                audio_config.get("model_type", "wav2vec2")
+            ](**audio_config)
+
+        self.text_model_lora_config = (
+            text_model_lora_config
+            if isinstance(text_model_lora_config, dict)
+            else dataclasses.asdict(text_model_lora_config or LoraConfigSimplified())
+        )
+        self.audio_model_lora_config = (
+            audio_model_lora_config
+            if isinstance(audio_model_lora_config, dict)
+            else dataclasses.asdict(audio_model_lora_config or LoraConfigSimplified())
+        )
+
+        self.vocab_size = self.text_config.vocab_size
+
+        self.initializer_range = self.text_config.initializer_range
+
+        super().__init__(**kwargs)
+
+    def to_diff_dict(self) -> Dict[str, Any]:
+        diff_dict = super().to_diff_dict()
+
+        # remove text_config and audio_config if text_model_id and audio_model_id are present
+        if self.text_model_id is not None:
+            diff_dict.pop("text_config", None)
+        if self.audio_model_id is not None:
+            diff_dict.pop("audio_config", None)
+
+        return diff_dict

ultravox_model.py

@@ -0,0 +1,652 @@
+import logging
+from typing import Any, Dict, Optional, Set, Tuple, Union
+
+import peft
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+import transformers.activations
+import transformers.modeling_outputs
+import transformers.models
+from transformers.models.whisper import modeling_whisper as whisper
+
+# We must use relative import in this directory to allow uploading to HF Hub
+# Even "from . import X" pattern doesn't work (undocumented and unclear why)
+from .ultravox_config import LossConfig
+from .ultravox_config import LossFunction
+from .ultravox_config import UltravoxConfig
+
+
+class UltravoxModel(transformers.LlamaPreTrainedModel):
+    """
+    The Ultravox model which consists of an audio encoder and a language model.
+    Audio input is processed by the audio encoder, then every `stack_factor` frames are stacked together and
+    projected to the language model's embedding space using a few linear layers.
+    The text is embedded by the language model as usual and then the audio and text embeddings are merged together.
+    A special token `<|audio|>` is used to indicate the start of the audio embeddings in the merged embeddings.
+    Parameters:
+        config: Model configuration class with all the parameters of the model.
+    """
+
+    config_class = UltravoxConfig
+    config: UltravoxConfig  # for type hinting
+    # Usually we load encoder and LLM weights from a pretrained model separately, so they are allowed to be missing
+    _keys_to_ignore_on_load_missing = ["audio_tower.*", "language_model.*"]
+
+    def __init__(self, config: UltravoxConfig):
+        super().__init__(config)
+        self._register_load_state_dict_pre_hook(self._pre_load_state_dict_hook)
+
+        self.keep_params: Set[str] = set()
+        self.vocab_size = config.vocab_size
+
+        self.audio_tower = self._create_audio_tower(config)
+        self.multi_modal_projector = self._create_multi_modal_projector(config)
+        self.language_model = self._create_language_model(config)
+
+        # Determine no_split_modules dynamically to use with FSDP auto_wrap policy.
+        # FSDP throws an error if some of the layer types are not found in the model.
+        # This would be something like ["LlamaDecoderLayer", "WhisperEncoderLayer"]
+        self._no_split_modules = (self.language_model._no_split_modules or []) + (
+            self.audio_tower._no_split_modules or []
+        )
+
+        self.loss_config = LossConfig()
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def set_decoder(self, decoder):
+        self.language_model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.language_model.get_decoder()
+
+    def tie_weights(self):
+        return self.language_model.tie_weights()
+
+    def set_loss_config(self, loss_config: LossConfig):
+        self.loss_config = loss_config
+
+    def _setup_cache(
+        self, cache_cls, max_batch_size: int, max_cache_len: Optional[int] = None
+    ):
+        self.language_model._setup_cache(cache_cls, max_batch_size, max_cache_len)
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        return self.language_model._reorder_cache(past_key_values, beam_idx)
+
+    def resize_token_embeddings(
+        self,
+        new_num_tokens: Optional[int] = None,
+        pad_to_multiple_of: Optional[int] = None,
+    ) -> nn.Embedding:
+        model_embeds = self.language_model.resize_token_embeddings(
+            new_num_tokens, pad_to_multiple_of
+        )
+        # update vocab size
+        self.config.text_config.vocab_size = model_embeds.num_embeddings
+        self.config.vocab_size = model_embeds.num_embeddings
+        self.vocab_size = model_embeds.num_embeddings
+        return model_embeds
+
+    def _compute_kl_loss(
+        self,
+        lm_output: transformers.modeling_outputs.CausalLMOutputWithPast,
+        labels: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None,
+        alt_input_ids: Optional[torch.Tensor] = None,
+        alt_attention_mask: Optional[torch.Tensor] = None,
+        alt_labels: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        # disable gradient computation for the teacher model
+        with torch.no_grad():
+            # compute the teacher (text-only) model's distribution
+            alt_inputs_embeds = self.get_input_embeddings().forward(alt_input_ids)
+            alt_lm_output = self.language_model.forward(
+                inputs_embeds=alt_inputs_embeds,
+                labels=alt_labels,
+                attention_mask=alt_attention_mask,
+                past_key_values=past_key_values,
+                **kwargs,
+            )
+        # compute the KL divergence loss between the two models
+        kl_loss = F.kl_div(
+            F.log_softmax(
+                lm_output.logits[labels != -100] / self.loss_config.kl_temperature,
+                dim=-1,
+            ),
+            F.softmax(
+                alt_lm_output.logits[alt_labels != -100]
+                / self.loss_config.kl_temperature,
+                dim=-1,
+            ),
+            reduction="batchmean",
+        )
+        return {"loss": kl_loss}
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        audio_values: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        audio_token_start_idx: Optional[torch.Tensor] = None,
+        audio_len: Optional[torch.Tensor] = None,
+        audio_token_len: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None,
+        # the alt_* fields are needed for KL divergence loss
+        alt_input_ids: Optional[torch.Tensor] = None,
+        alt_attention_mask: Optional[torch.Tensor] = None,
+        alt_labels: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, transformers.modeling_outputs.CausalLMOutputWithPast]:
+        """
+        Forward pass for the Ultravox model.
+        `input_ids` are the tokenized text input. They are embedded by the language model as usual.
+        `audio_values` are processed by the audio encoder and then every `stack_factor` frames are stacked together and
+        projected to the language model's embedding space using a few linear layers.
+        The audio and text embeddings are merged together. A special token `<|audio|>` is used to indicate the start
+        of the audio embeddings in the merged embeddings.
+        Args:
+            input_ids: The tokenized text input.
+            audio_values: The processed audio values.
+            inputs_embeds: The embeddings for the input tokens.
+            labels: The tokenized text labels.
+            attention_mask: The attention mask for the input.
+            position_ids: The position ids for the input.
+            past_key_values: The past key value cache for the language model attention layers.
+            **kwargs: Additional keyword arguments. Passed directly to the language model.
+        """
+        if inputs_embeds is None:
+            # B x T  ->  B x T x D
+            inputs_embeds = self.get_input_embeddings().forward(input_ids)
+
+        if audio_values is not None:
+            assert (
+                audio_token_start_idx is not None and audio_token_len is not None
+            ), "audio_token_start_idx and audio_token_len must be provided if audio_values are provided."
+            assert (
+                len(audio_token_start_idx) == len(audio_token_len) == len(audio_values)
+            ), "audio_token_start_idx, audio_token_len, and audio_values must have the same batch size."
+
+            # B x A/3200 x D
+            audio_tower_output = self.audio_tower.forward(
+                audio_values.to(self.audio_tower.dtype),
+                audio_len = audio_len
+            ).last_hidden_state
+            audio_tower_output = audio_tower_output.to(inputs_embeds.dtype)
+
+            audio_embeds = self.multi_modal_projector.forward(audio_tower_output)
+
+            # combine audio and text embeddings
+            for i, (audio, start, length) in enumerate(
+                zip(audio_embeds, audio_token_start_idx, audio_token_len)
+            ):
+                assert length <= audio.shape[0]
+                inputs_embeds[i, start : start + length].copy_(audio[:length])
+                
+
+        lm_output = self.language_model.forward(
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            **kwargs,
+        )
+        if self.training:
+            if self.loss_config.loss_function == LossFunction.CrossEntropy:
+                return lm_output
+            elif self.loss_config.loss_function == LossFunction.KL_Divergence:
+                return self._compute_kl_loss(
+                    lm_output=lm_output,
+                    labels=labels,
+                    past_key_values=past_key_values,
+                    alt_input_ids=alt_input_ids,
+                    alt_attention_mask=alt_attention_mask,
+                    alt_labels=alt_labels,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported loss function: {self.loss_config.loss_function}"
+                )
+        else:
+            return lm_output
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.Tensor,
+        audio_values: Optional[torch.FloatTensor] = None,
+        audio_token_start_idx: Optional[torch.Tensor] = None,
+        audio_token_len: Optional[torch.Tensor] = None,
+        audio_len: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        model_input = self.language_model.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        # include audio information in model_input only when it is needed during prefilling
+        # audio_token_start_idx should always be relative to the current cache position
+        prefill_start_idx = 0 if cache_position is None else cache_position[0]
+        if (
+            audio_values is not None
+            and audio_token_start_idx is not None
+            and prefill_start_idx <= torch.max(audio_token_start_idx)
+        ):
+            model_input["audio_values"] = audio_values
+            model_input["audio_token_start_idx"] = (
+                audio_token_start_idx - prefill_start_idx
+            )
+            model_input["audio_token_len"] = audio_token_len
+            model_input["audio_len"] = audio_len
+
+        return model_input
+
+    @classmethod
+    def _create_multi_modal_projector(
+        cls, config: UltravoxConfig
+    ) -> "UltravoxProjector":
+        projector = UltravoxProjector(config)
+        projector.to(config.torch_dtype)
+        return projector
+
+    @classmethod
+    def _create_audio_tower(
+        cls, config: UltravoxConfig
+    ) -> Union[transformers.Wav2Vec2Model, "ModifiedWhisperEncoder"]:
+        if config.audio_model_id is not None:
+            if "whisper" in config.audio_model_id is not None:
+                audio_tower = ModifiedWhisperEncoder.from_pretrained(
+                    config.audio_model_id, torch_dtype=config.torch_dtype
+                )
+            else:
+                audio_tower = transformers.AutoModel.from_pretrained(
+                    config.audio_model_id, torch_dtype=config.torch_dtype
+                )
+        else:
+            if "whisper" in config.audio_config._name_or_path:
+                audio_tower = ModifiedWhisperEncoder(config.audio_config)
+            else:
+                with transformers.modeling_utils.no_init_weights():
+                    # we only ever use from_config if the weights are retrained, hence initializing is not
+                    # required. This makes the model quite creation faster since init on CPU is quite slow.
+                    audio_tower = transformers.AutoModel.from_config(
+                        config.audio_config
+                    )
+
+        if isinstance(
+            audio_tower,
+            (transformers.Wav2Vec2BertModel, transformers.WhisperModel),
+        ):
+            # For these models we only need the encoder part
+            # Wav2Vec2BertModel -> Wav2Vec2BertEncoder
+            # WhisperModel -> WhisperEncoder
+            audio_tower = audio_tower.encoder
+
+        audio_tower = apply_lora(audio_tower, config.audio_model_lora_config)
+        return audio_tower
+
+    @classmethod
+    def _create_language_model(
+        cls, config: UltravoxConfig
+    ) -> transformers.LlamaForCausalLM:
+        if config.text_model_id is not None:
+            language_model = transformers.AutoModelForCausalLM.from_pretrained(
+                config.text_model_id,
+                attn_implementation=config._attn_implementation,
+                torch_dtype=config.torch_dtype,
+            )
+        else:
+            with transformers.modeling_utils.no_init_weights():
+                # we only ever use from_config if the weights are retrained, hence initializing is not
+                # required. This makes the model quite creation faster since init on CPU is quite slow.
+                language_model = transformers.AutoModelForCausalLM.from_config(
+                    config.text_config,
+                    attn_implementation=config._attn_implementation,
+                    torch_dtype=config.torch_dtype,
+                )
+
+        language_model = apply_lora(language_model, config.text_model_lora_config)
+        return language_model
+
+    def merge_and_unload(self):
+        if isinstance(self.language_model, peft.PeftModel):
+            self.language_model = self.language_model.merge_and_unload()
+            # no need to download base language model weights anymore, so we can remove the id
+            self.config.text_model_id = None
+            self.keep_params.update(
+                set(
+                    [
+                        f"language_model.{name}"
+                        for name, _ in self.language_model.named_parameters()
+                    ]
+                )
+            )
+
+        if isinstance(self.audio_tower, peft.PeftModel):
+            self.audio_tower = self.audio_tower.merge_and_unload()
+            # no need to download base audio model weights anymore, so we can remove the id
+            self.config.audio_model_id = None
+            self.keep_params.update(
+                set(
+                    [
+                        f"audio_tower.{name}"
+                        for name, _ in self.audio_tower.named_parameters()
+                    ]
+                )
+            )
+
+        for param in ["text_model_lora_config", "audio_model_lora_config"]:
+            if hasattr(self.config, param):
+                delattr(self.config, param)
+
+    def push_to_hub(self, *args, **kwargs):
+        self.merge_and_unload()
+        self.to(self.language_model.dtype)
+        return super().push_to_hub(*args, **kwargs)
+
+    def save_pretrained(
+        self, *args, state_dict: Optional[Dict[str, Any]] = None, **kwargs
+    ):
+        if state_dict is None:
+            state_dict = super().state_dict()
+
+        named_params = dict(self.named_parameters())
+
+        state_dict = {
+            k: v
+            for k, v in state_dict.items()
+            if k in self.keep_params
+            or (k in named_params and named_params[k].requires_grad)
+        }
+
+        super().save_pretrained(*args, state_dict=state_dict, **kwargs)
+
+    def _pre_load_state_dict_hook(self, state_dict: Dict[str, Any], *args, **kwargs):
+        self.keep_params.update(set(state_dict.keys()))
+
+    def print_trainable_parameters(self):
+        """
+        Prints the number of trainable parameters in the model (reuses Peft model's method)
+        """
+        count_params = peft.peft_model.PeftModel.get_nb_trainable_parameters
+
+        trainable_params, all_param = count_params(self)
+
+        logging.info(
+            f"trainable params: {trainable_params:,d} || all params: {all_param:,d}"
+            f" || trainable%: {100 * trainable_params / all_param:.1f}%"
+        )
+
+        lm_trainable_params, lm_all_params = count_params(self.language_model)
+        audio_trainable_params, audio_all_params = count_params(self.audio_tower)
+
+        projector_trainable_params = (
+            trainable_params - lm_trainable_params - audio_trainable_params
+        )
+        projector_all_params = all_param - lm_all_params - audio_all_params
+
+        logging.info(
+            f"Trainable%:   "
+            f" LLM: {100 * lm_trainable_params / lm_all_params:.1f}%"
+            f" || Audio Encoder: {100 * audio_trainable_params / audio_all_params:.1f}%"
+            f" || Projector: {100 * projector_trainable_params / projector_all_params:.1f}%"
+        )
+
+
+def is_cache_empty(
+    past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]]
+) -> bool:
+    """
+    Check if the cache is empty.
+    """
+    if past_key_values is None:
+        return True
+    if isinstance(past_key_values, tuple):
+        return all(len(c) == 0 for c in past_key_values)
+    return past_key_values.get_seq_length() == 0
+
+
+def apply_lora(model: torch.nn.Module, lora_config: dict) -> torch.nn.Module:
+    """
+    Applies LoRA finetuning to the model. If the `r` parameter is set to 0, the model is frozen instead.
+    """
+    lora_config = peft.LoraConfig(**lora_config or {})
+
+    if lora_config.r == 0:
+        # freeze the model entirely
+        for param in model.parameters():
+            param.requires_grad = False
+    else:
+        model = peft.get_peft_model(model, lora_config)
+
+    return model
+
+
+class StackAudioFrames(nn.Module):
+    """
+    Stack the audio embedding frames to reduce the sequence length by a factor of `stack_factor`.
+    The number of output frames will be `ceil(T / stack_factor) + 1` where `T` is the number of input frames.
+    NOTE: the extra +1 is intentional: in case the number of audio tokens are over-estimated by the processor,
+    we want to make sure `processor.audio_token_replacement` (i.e. EOS) doesn't get leaked into the middle of embeddings.
+    In most cases this extra padding will get removed in the model's forward function so it has no effect.
+    """
+
+    def __init__(self, stack_factor: int = 8):
+        super().__init__()
+        self.stack_factor = stack_factor
+
+    def forward(self, audio_embeds: torch.Tensor) -> torch.Tensor:
+        B, T, C = audio_embeds.shape
+        T_pad = (T + self.stack_factor - 1) // self.stack_factor * self.stack_factor
+        audio_embeds = F.pad(audio_embeds, (0, 0, 0, T_pad - T + self.stack_factor))
+        B, T, C = audio_embeds.shape
+        audio_embeds = audio_embeds.view(
+            B, T // self.stack_factor, C * self.stack_factor
+        )
+        return audio_embeds
+
+
+class RMSNorm(transformers.models.llama.modeling_llama.LlamaRMSNorm):
+    def __init__(self, hidden_size: int, init: float = 1, eps: float = 1e-6):
+        super().__init__(hidden_size=hidden_size, eps=eps)
+        self.weight.data.fill_(init)
+
+
+class SwiGLU(nn.Module):
+    def forward(self, x):
+        x, gate = x.chunk(2, dim=-1)
+        return F.silu(gate) * x
+
+
+class UltravoxProjector(nn.Sequential):
+    def __init__(self, config: UltravoxConfig):
+        super().__init__()
+        self.hidden_dim = config.hidden_size
+        self._pad_and_stack = StackAudioFrames(config.stack_factor)
+        dim = config.audio_config.hidden_size * config.stack_factor
+        self.ln_pre = RMSNorm(dim, init=config.norm_init)
+        self.linear_1 = nn.Linear(dim, self.hidden_dim, bias=False)
+        dim = self.hidden_dim
+        self.act = transformers.activations.get_activation(config.projector_act)
+        dim = dim // 2 if config.projector_act == "swiglu" else dim
+        self.linear_2 = nn.Linear(dim, config.text_config.hidden_size, bias=False)
+        self.ln_post = RMSNorm(config.text_config.hidden_size, init=config.norm_init)
+
+    def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
+        audio_features = self._pad_and_stack(audio_features)
+        audio_features = self.ln_pre(audio_features)
+        hidden_states = self.linear_1(audio_features)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        hidden_states = self.ln_post(hidden_states)
+        return hidden_states
+
+
+class ModifiedWhisperEncoder(whisper.WhisperEncoder, transformers.modeling_utils.ModuleUtilsMixin):
+    """
+    Encoder portion of OpenAI's Whisper model.
+    This implementation is a slightly modified version of HF Transformers' Whisper Encoder, with only a few fixes:
+    1. base_model_prefix updated to allow for doing `.from_pretrained` directly on the encoder
+    2. allow less than 30 second of audio padding to be passed in:
+        - relaxed ValueError check for `input_features` length to be less than or equal to `expected_seq_length` instead of strictly equal
+        - embed_pos is now sliced to match the length of `inputs_embeds`
+    Original: https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/modeling_whisper.py
+    """
+
+    base_model_prefix = "model.encoder"
+    _no_split_modules = ["WhisperEncoderLayer"]
+
+    def forward(
+        self,
+        input_features,
+        audio_len=None,
+        head_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        expected_seq_length = (
+            self.config.max_source_positions
+            * self.conv1.stride[0]
+            * self.conv2.stride[0]
+        )
+        if input_features.shape[-1] > expected_seq_length:
+            raise ValueError(
+                f"Whisper expects the mel input features to be of length {expected_seq_length} or less, but found {input_features.shape[-1]}. Make sure to pad the input mel features to {expected_seq_length}."
+            )
+
+        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
+        )
+        inputs_embeds = nn.functional.gelu(self.conv1(input_features))
+        inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))
+
+        inputs_embeds = inputs_embeds.permute(0, 2, 1)
+        embed_pos = self.embed_positions.weight[: inputs_embeds.size(-2)]
+
+        hidden_states = inputs_embeds + embed_pos
+        hidden_states = nn.functional.dropout(
+            hidden_states, p=self.dropout, training=self.training
+        )
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        attention_mask = None
+        if audio_len != None:
+            audio_feature_len = self._get_feat_extract_output_lengths(audio_len)
+            batch_size = hidden_states.shape[0]
+            max_seq_len = hidden_states.shape[1]
+            attention_mask = (
+                torch.arange(max_seq_len, device=hidden_states.device)[None, :]
+                .expand(batch_size, -1)
+                .lt(audio_feature_len.view(batch_size, 1))
+            )
+            attention_mask = self.get_extended_attention_mask(
+                attention_mask,
+                None,
+                device=hidden_states.device,
+                dtype=hidden_states.dtype,
+            )
+
+        # check if head_mask has a correct number of layers specified if desired
+        if head_mask is not None:
+            assert head_mask.size()[0] == (
+                len(self.layers)
+            ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
+            to_drop = False
+            if self.training:
+                dropout_probability = torch.rand([])
+                if dropout_probability < self.layerdrop:  # skip the layer
+                    to_drop = True
+
+            if to_drop:
+                layer_outputs = (None, None)
+            else:
+                if self.gradient_checkpointing and self.training:
+                    layer_outputs = self._gradient_checkpointing_func(
+                        encoder_layer.__call__,
+                        hidden_states,
+                        attention_mask,
+                        (head_mask[idx] if head_mask is not None else None),
+                        output_attentions,
+                    )
+                else:
+                    layer_outputs = encoder_layer(
+                        hidden_states,
+                        attention_mask,
+                        layer_head_mask=(
+                            head_mask[idx] if head_mask is not None else None
+                        ),
+                        output_attentions=output_attentions,
+                    )
+
+                hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        hidden_states = self.layer_norm(hidden_states)
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, encoder_states, all_attentions]
+                if v is not None
+            )
+        return transformers.modeling_outputs.BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_states,
+            attentions=all_attentions,
+        )
+
+
+UltravoxConfig.register_for_auto_class()
+UltravoxModel.register_for_auto_class()
+
+transformers.AutoConfig.register("ultravox", UltravoxConfig)
+transformers.AutoModel.register(UltravoxConfig, UltravoxModel)
+
+transformers.activations.ACT2FN["swiglu"] = SwiGLU

ultravox_pipeline.py

@@ -0,0 +1,127 @@
+import logging
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import transformers
+
+# We must use relative import in this directory to allow uploading to HF Hub
+# Even "from . import X" pattern doesn't work (undocumented and unclear why)
+from .ultravox_model import UltravoxModel
+from .ultravox_processing import UltravoxProcessor
+
+
+class UltravoxPipeline(transformers.Pipeline):
+    def __init__(
+        self,
+        model: UltravoxModel,
+        tokenizer: Optional[transformers.PreTrainedTokenizerBase] = None,
+        audio_processor: Optional[transformers.ProcessorMixin] = None,
+        **kwargs
+    ):
+        if tokenizer is None:
+            try:
+                tokenizer = transformers.AutoTokenizer.from_pretrained(
+                    model.config._name_or_path
+                )
+            except:
+                tokenizer = transformers.AutoTokenizer.from_pretrained(
+                    model.config.text_model_id or model.config.text_config._name_or_path
+                )
+
+        if audio_processor is None:
+            audio_processor = transformers.AutoProcessor.from_pretrained(
+                model.config.audio_model_id or model.config.audio_config._name_or_path
+            )
+
+        super().__init__(model=model, tokenizer=tokenizer, **kwargs)
+        
+        self.processor = UltravoxProcessor(
+            audio_processor=audio_processor,
+            tokenizer=tokenizer,
+            stack_factor=model.config.stack_factor,
+        )
+
+    def _sanitize_parameters(self, **kwargs):
+        generation_keys = ["temperature", "max_new_tokens", "repetition_penalty"]
+        generation_kwargs = {k: kwargs[k] for k in kwargs if k in generation_keys}
+        return {}, generation_kwargs, {}
+
+    def preprocess(self, inputs: Dict[str, Any]):
+        turns: list = inputs.get("turns", [])
+
+        audio = inputs.get("audio", None)
+        # Convert to float32 if needed.
+        if isinstance(audio, np.ndarray):
+            if audio.dtype == np.float64:
+                audio = audio.astype(np.float32)
+            elif audio.dtype == np.int16:
+                audio = audio.astype(np.float32) / np.float32(32768.0)
+            elif audio.dtype == np.int32:
+                audio = audio.astype(np.float32) / np.float32(2147483648.0)
+
+        if audio is not None and (len(turns) == 0 or turns[-1]["role"] != "user"):
+            prompt = inputs.get("prompt", "<|audio|>")
+            if "<|audio|>" not in prompt:
+                logging.warning(
+                    "Prompt does not contain '<|audio|>', appending '<|audio|>' to the end of the prompt."
+                )
+
+                prompt += " <|audio|>"
+            turns.append({"role": "user", "content": prompt})
+
+        text = self.processor.tokenizer.apply_chat_template(
+            turns, add_generation_prompt=True, tokenize=False
+        )
+
+        if "sampling_rate" not in inputs and audio is not None:
+            logging.warning(
+                "No sampling rate provided, using default of 16kHz. We highly recommend providing the correct sampling rate."
+            )
+
+        output = self.processor(
+            text=text,
+            audio=audio,
+            sampling_rate=inputs.get("sampling_rate", 16000),
+        )
+        if "audio_values" in output:
+            output["audio_values"] = output["audio_values"].to(self.model.dtype)
+
+        return output
+
+    def _forward(
+        self,
+        model_inputs: Dict[str, Any],
+        temperature: Optional[float] = None,
+        max_new_tokens: Optional[int] = None,
+        repetition_penalty: float = 1.1,
+    ) -> List[int]:
+        temperature = temperature or None
+        do_sample = temperature is not None
+
+        terminators = [self.tokenizer.eos_token_id]
+        if "<|eot_id|>" in self.tokenizer.added_tokens_encoder:
+            terminators.append(self.tokenizer.convert_tokens_to_ids("<|eot_id|>"))
+
+        input_len = model_inputs["input_ids"].shape[1]
+
+        outputs = self.model.generate(
+            **model_inputs,
+            do_sample=do_sample,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            repetition_penalty=repetition_penalty,
+            eos_token_id=terminators
+        )
+        return outputs[0][input_len:]
+
+    def postprocess(self, model_outputs) -> str:
+        output_text = self.tokenizer.decode(model_outputs, skip_special_tokens=True)
+        return output_text
+
+
+transformers.pipelines.PIPELINE_REGISTRY.register_pipeline(
+    "ultravox-pipeline",
+    pipeline_class=UltravoxPipeline,
+    pt_model=transformers.AutoModel,
+    type="multimodal",
+)

ultravox_processing.py

@@ -0,0 +1,207 @@
+from typing import Optional, Union
+
+import numpy as np
+import torch
+import transformers
+
+from .ultravox_config import UltravoxConfig
+
+
+class UltravoxProcessor(transformers.ProcessorMixin):
+    """
+    Constructs an Ultravox processor which wraps an audio processor and a tokenizer into a single processor.
+    Args:
+        audio_processor: The audio processor for the audio encoder.
+        tokenizer: The tokenizer for the language model.
+    """
+
+    attributes = ["audio_processor", "tokenizer"]
+    audio_processor_class = (
+        "Wav2Vec2Processor",
+        "SeamlessM4TFeatureExtractor",
+        "WhisperProcessor",
+        "Wav2Vec2BertProcessor",
+    )
+    tokenizer_class = (
+        "PreTrainedTokenizer",
+        "PreTrainedTokenizerFast",
+    )
+
+    tokenizer: transformers.PreTrainedTokenizerBase
+    audio_processor: transformers.ProcessorMixin
+
+    def __init__(
+        self,
+        audio_processor=None,
+        tokenizer=None,
+        audio_padding: str = "longest",
+        encoder_ds_factor: int = 320,
+        stack_factor: int = 8,
+        audio_placeholder: str = "<|audio|>",
+    ):
+        """
+        Args:
+            audio_processor: The audio processor for the audio encoder.
+            tokenizer: The tokenizer for the language model.
+            audio_padding: The padding strategy for the audio encoder.
+            encoder_ds_factor: The downsample factor of the audio encoder.
+            stack_factor: The factor by which the audio encoder output is stacked in the multimodal projector.
+            audio_placeholder: The placeholder for the audio in the text.
+        """
+        self.audio_padding = audio_padding
+        self.encoder_ds_factor = encoder_ds_factor
+        self.stack_factor = stack_factor
+        self.audio_placeholder = audio_placeholder
+        self.audio_token_replacement = tokenizer.eos_token
+        assert (
+            self.audio_token_replacement is not None
+        ), "The tokenizer has no EOS token. Cannot recover."
+        if tokenizer.pad_token_id is None:
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+
+        super().__init__(audio_processor=audio_processor, tokenizer=tokenizer)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        config: UltravoxConfig = transformers.AutoConfig.from_pretrained(
+            pretrained_model_name_or_path, **kwargs
+        )
+        audio_processor = transformers.AutoProcessor.from_pretrained(
+            config.audio_model_id
+            or config.audio_config._name_or_path
+            or "facebook/wav2vec2-base-960h"
+        )
+
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path, **kwargs
+        )
+        tokenizer.padding_side = "left"
+        tokenizer.pad_token = tokenizer.eos_token
+
+        return cls(
+            audio_processor=audio_processor,
+            tokenizer=tokenizer,
+            stack_factor=config.stack_factor,
+        )
+
+    def __call__(
+        self,
+        text: Optional[str] = None,
+        audio: Optional[Union[np.ndarray, torch.Tensor]] = None,
+        sampling_rate: Optional[int] = None,
+        return_tensors: Optional[
+            Union[str, transformers.TensorType]
+        ] = transformers.TensorType.PYTORCH,
+        **kwargs,
+    ) -> transformers.BatchFeature:
+        """
+        Main method to prepare for the model one text sequence and audio. This method forwards the `text`
+        and `kwargs` arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizerFast.__call__`] if `text` is not `None` to encode
+        the text. To prepare the audio(s), this method forwards the `audio`, `sampling_rate` and `kwargs` arguments to
+        audio processor's [`~Wav2Vec2Processor.__call__`] if `audio` is not `None`. Please refer to the docstring
+        of the above two methods for more information.
+        Args:
+            text (`str`, `List[str]`):
+                The sequence to be encoded. Sequence can be a string or (pretokenized string).
+            audio (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
+                The audio to be prepared. Audio can be NumPy array or PyTorch tensor. In case of a
+                NumPy array/PyTorch tensor, each audio should be of shape (C, T), where C is a number of channels, and T the
+                sample length of the audio.
+            sampling_rate (`int`, *optional*, defaults to 16000):
+                Sampling rate of the input audio. We expect 16kHz audio. Don't change this value unless you know what
+                you are doing.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Acceptable values are:
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+              `None`).
+            - **audio_values** -- Processed audio values to be fed to a model. Returned when `audio` is not `None`.
+            - **audio_token_len** -- Predicted number of audio frames: this value is guaranteed to be a close upper bound.
+              Returned when `audio` is not `None`.
+            - **audio_token_start_idx** -- The index in the tokenized text where the audio starts. Returned when `audio` is not `None`.
+        """
+        # TODO: Add support for multiple audio and text inputs.
+        data = {}
+        audio_embed_frames = 0
+        if audio is not None and len(audio) > 0:
+            if self.audio_padding == "max_length":
+                # 30 seconds is the expected length for Whisper
+                assert sampling_rate is not None, "Sampling rate must be provided."
+                audio_len = [30 * sampling_rate] * len(audio)
+            else:
+                audio_len = [a.shape[-1] for a in audio]
+            # It's guaranteed that the number of frames is less than or equal to this amount.
+            # For Whisper this is exact AFAICT, but for Wav2Vec2 it's an upper bound.
+            # Currently, StackAudioFrames makes sure an over-estimation won't cause issues by padding the audio embeddings.
+            nb_encoder_frames = [int(round(a / self.encoder_ds_factor + 1e-4)) for a in audio_len]
+            audio_embed_frames = [int(np.ceil(n / self.stack_factor)) for n in nb_encoder_frames]
+            data["audio_token_len"] = audio_embed_frames
+
+            # Main audio processing. The processor is model-specific.
+            x = self.audio_processor(
+                audio,
+                sampling_rate=sampling_rate,
+                padding="longest",
+                max_length=max(audio_len),
+                return_attention_mask=True,
+                **kwargs,
+            )
+            if "input_features" in x:
+                data["audio_values"] = x.input_features
+            else:
+                data["audio_values"] = x.input_values
+            data["audio_len"] = x.attention_mask.sum(-1) - 1
+
+        if text is not None:
+            #assert isinstance(
+            #    text, str
+            #), "Text must be a string. Batch mode not supported yet."
+            data["audio_token_start_idx"] = []
+            for t in text:
+                assert self.audio_placeholder in t
+                if "audio_token_len" not in data:
+                    raise ValueError(
+                        f"audio must be provided when using audio placeholder ({self.audio_placeholder}) in text."
+                    )
+
+                start_idx = len(
+                    self.tokenizer.encode(
+                        t[: t.index(self.audio_placeholder)],
+                        add_special_tokens=False,
+                    )
+                )
+                data["audio_token_start_idx"].append(start_idx)
+
+            # Replace the audio placeholder with the audio token.
+            #   e.g. "Transcribe\n<|audio|>" -> "Transcribe </s></s></s></s></s></s></s></s>"
+            #        where the number of </s> is the number of audio frames.
+            text = [t.replace(self.audio_placeholder, self.audio_token_replacement * data["audio_token_len"][i]) for i, t in enumerate(text)]
+
+            # Special tokens like BOS should already have been added by the caller.
+            data.update(self.tokenizer(text, add_special_tokens=False, padding=True, **kwargs))
+
+        return transformers.BatchFeature(data=data, tensor_type=return_tensors)
+
+    def batch_decode(self, *args, **kwargs):
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        audio_processor_input_names = self.audio_processor.model_input_names
+        return list(set(tokenizer_input_names + audio_processor_input_names))
+
+
+UltravoxProcessor.register_for_auto_class()
+
+transformers.AutoProcessor.register(UltravoxConfig, UltravoxProcessor)