Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2025 The Nari Labs and HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch Dia model."""16 17from typing import Callable, Optional, Union18 19import torch20from torch import nn21 22from ...cache_utils import DynamicCache, EncoderDecoderCache23from ...masking_utils import create_causal_mask24from ...modeling_attn_mask_utils import (25    _prepare_4d_attention_mask,26    _prepare_4d_attention_mask_for_sdpa,27)28from ...modeling_flash_attention_utils import FlashAttentionKwargs29from ...modeling_layers import GradientCheckpointingLayer30from ...modeling_outputs import (31    BaseModelOutput,32    BaseModelOutputWithPastAndCrossAttentions,33    Seq2SeqLMOutput,34    Seq2SeqModelOutput,35)36from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel37from ...processing_utils import Unpack38from ...utils import auto_docstring, can_return_tuple, is_torch_flex_attn_available, is_torchdynamo_compiling, logging39from ..llama.modeling_llama import (40    LlamaAttention,41    LlamaRMSNorm,42    LlamaRotaryEmbedding,43    eager_attention_forward,44)45from ..phi3.modeling_phi3 import Phi3MLP46from .configuration_dia import DiaConfig, DiaDecoderConfig, DiaEncoderConfig47from .generation_dia import DiaGenerationMixin48 49 50if is_torch_flex_attn_available():51    from ...integrations.flex_attention import make_flex_block_causal_mask52 53 54logger = logging.get_logger(__name__)55 56 57@auto_docstring58class DiaPreTrainedModel(PreTrainedModel):59    config: DiaConfig60    base_model_prefix = "model"61    supports_gradient_checkpointing = True62    _supports_flash_attn = True63    _supports_sdpa = True64    _supports_flex_attn = True65    _can_compile_fullgraph = True66    main_input_name = "input_ids"67    _no_split_modules = ["DiaEncoderLayer", "DiaDecoderLayer"]68 69 70class DiaMultiChannelEmbedding(nn.Module):71    """In order to efficiently compute the audio embedding from the 9 different channels,72    we vectorize the embedding process by using a single embedding layer and an offset.73    Example:74    - num_embeds = 475    - vocab_size = 876    - num_channels = 377    We would have offsets = [0, 8, 16]78    If audio_codes = [0, 1, 2, 3], [1, 3, 4, 7], [5, 6, 7, 8],79    then tokens = audio_codes + offsets80                = [0, 1, 2, 3, 9, 11, 12, 15, 21, 22, 23, 24]81    This allows us to use a single embedding layer for all channels.82    """83 84    def __init__(self, config: DiaDecoderConfig):85        super().__init__()86        self.embed = nn.Embedding(config.vocab_size * config.num_channels, config.hidden_size)87        self.hidden_size = config.hidden_size88        self.num_channels = config.num_channels89        offsets = torch.arange(config.num_channels, dtype=torch.long) * config.vocab_size  # (C,)90        self.register_buffer("offsets", offsets, persistent=False)91 92    def forward(self, audio_codes: torch.Tensor) -> torch.Tensor:93        tokens = (audio_codes + self.offsets.to(audio_codes.device)).squeeze(1)94        embeds = self.embed(tokens).view(tokens.shape[0], audio_codes.shape[1], -1, self.hidden_size)95        return embeds.sum(dim=2)96 97 98class DiaMLP(Phi3MLP):99    pass100 101 102class DiaRMSNorm(LlamaRMSNorm):103    pass104 105 106class DiaRotaryEmbedding(LlamaRotaryEmbedding):107    pass108 109 110class DiaSelfAttention(LlamaAttention):111    """Multi-headed attention from 'Attention Is All You Need' paper"""112 113    def __init__(self, config: Union[DiaEncoderConfig, DiaDecoderConfig], layer_idx: int, is_causal: bool = False):114        nn.Module.__init__(self)115        self.config = config116        self.layer_idx = layer_idx117        self.hidden_size = config.hidden_size118        self.num_heads = self.config.num_attention_heads119        self.num_key_value_heads = self.config.num_key_value_heads or self.num_heads120        self.num_key_value_groups = self.num_heads // self.num_key_value_heads121        self.head_dim = getattr(config, "head_dim", config.hidden_size // self.num_heads)122        self.scaling = 1123        self.attention_dropout = 0.0124        self.is_causal = is_causal125 126        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)127        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)128        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)129        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)130 131 132class DiaCrossAttention(nn.Module):133    """Multi-headed attention from 'Attention Is All You Need' paper"""134 135    def __init__(self, config: DiaDecoderConfig, layer_idx: int):136        super().__init__()137        self.config = config138        self.layer_idx = layer_idx139        self.hidden_size = config.hidden_size140        self.cross_hidden_size = config.cross_hidden_size141        self.num_heads = self.config.cross_num_attention_heads142        self.num_key_value_heads = self.config.cross_num_key_value_heads143        self.num_key_value_groups = self.num_heads // self.num_key_value_heads144        self.head_dim = config.cross_head_dim145        self.scaling = 1146        self.attention_dropout = 0.0147        self.is_causal = False148 149        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)150        self.k_proj = nn.Linear(self.cross_hidden_size, self.num_key_value_heads * self.head_dim, bias=False)151        self.v_proj = nn.Linear(self.cross_hidden_size, self.num_key_value_heads * self.head_dim, bias=False)152        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)153 154    def forward(155        self,156        hidden_states: torch.Tensor,157        cross_attention_states: torch.Tensor,158        attention_mask: Optional[torch.Tensor] = None,159        past_key_values: Optional[EncoderDecoderCache] = None,160        **kwargs: Unpack[FlashAttentionKwargs],161    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:162        input_shape = hidden_states.shape[:-1]163        hidden_shape = (*input_shape, -1, self.head_dim)164        cross_shape = (*cross_attention_states.shape[:-1], -1, self.head_dim)165 166        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)167 168        is_updated = past_key_values.is_updated.get(self.layer_idx) if past_key_values is not None else False169        if past_key_values is not None and is_updated:170            # reuse k,v, cross_attentions171            key_states = past_key_values.cross_attention_cache.layers[self.layer_idx].keys172            value_states = past_key_values.cross_attention_cache.layers[self.layer_idx].values173        else:174            key_states = self.k_proj(cross_attention_states).view(cross_shape).transpose(1, 2)175            value_states = self.v_proj(cross_attention_states).view(cross_shape).transpose(1, 2)176 177            if past_key_values is not None:178                # save all states to the cache179                key_states, value_states = past_key_values.cross_attention_cache.update(180                    key_states,181                    value_states,182                    self.layer_idx,183                )184                # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls185                past_key_values.is_updated[self.layer_idx] = True186 187        attention_interface: Callable = eager_attention_forward188        if self.config._attn_implementation != "eager":189            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]190 191        attn_output, attn_weights = attention_interface(192            self,193            query_states,194            key_states,195            value_states,196            attention_mask,197            scaling=self.scaling,198            **kwargs,199        )200 201        attn_output = attn_output.reshape((*input_shape, -1)).contiguous()202        attn_output = self.o_proj(attn_output)203        return attn_output, attn_weights204 205 206class DiaEncoderLayer(GradientCheckpointingLayer):207    def __init__(self, config: DiaEncoderConfig, layer_idx: int):208        super().__init__()209        self.pre_sa_norm = DiaRMSNorm(config.hidden_size, eps=config.norm_eps)210        self.self_attention = DiaSelfAttention(config, layer_idx, is_causal=False)211        self.post_sa_norm = DiaRMSNorm(config.hidden_size, eps=config.norm_eps)212        self.mlp = DiaMLP(config)213 214    def forward(215        self,216        hidden_states: torch.Tensor,217        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC218        attention_mask: Optional[torch.Tensor] = None,219        **kwargs: Unpack[FlashAttentionKwargs],220    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:221        residual = hidden_states222        normed_states = self.pre_sa_norm(hidden_states)223        self_attn_output, self_attn_weights = self.self_attention(224            normed_states,225            position_embeddings=position_embeddings,226            attention_mask=attention_mask,227            **kwargs,228        )229        hidden_states = residual + self_attn_output230 231        residual = hidden_states232        normed_states = self.post_sa_norm(hidden_states)233        mlp_out = self.mlp(normed_states)234        hidden_states = residual + mlp_out235 236        return hidden_states, self_attn_weights237 238 239class DiaEncoder(DiaPreTrainedModel):240    def __init__(self, config: DiaEncoderConfig):241        super().__init__(config)242        self.config = config243 244        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)245        self.layers = nn.ModuleList(246            [DiaEncoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]247        )248        self.norm = DiaRMSNorm(config.hidden_size, eps=config.norm_eps)249        self.rotary_embeddings = DiaRotaryEmbedding(config)250 251    @auto_docstring252    @can_return_tuple253    def forward(254        self,255        input_ids: torch.Tensor,256        attention_mask: Optional[torch.Tensor] = None,257        output_attentions: Optional[bool] = False,258        output_hidden_states: Optional[bool] = False,259        **kwargs: Unpack[FlashAttentionKwargs],260    ) -> Union[BaseModelOutput, tuple]:261        hidden_states = self.embedding(input_ids)262 263        # RoPE264        # Note: We expect right padding and hence always generate265        # the position ids on the fly to reduce preparation overhead266        position_ids = torch.arange(input_ids.shape[-1], device=input_ids.device)[None, :]267        position_embeddings = self.rotary_embeddings(hidden_states, position_ids)268 269        attention_mask = self._update_full_mask(270            attention_mask,271            hidden_states,272        )273 274        encoder_states = () if output_hidden_states else None275        all_attentions = () if output_attentions else None276 277        for encoder_layer in self.layers:278            if output_hidden_states:279                encoder_states = encoder_states + (hidden_states,)280 281            layer_outputs = encoder_layer(282                hidden_states,283                position_embeddings=position_embeddings,284                attention_mask=attention_mask,285                **kwargs,286            )287            hidden_states = layer_outputs[0]288 289            if output_attentions:290                all_attentions = all_attentions + (layer_outputs[1],)291 292        hidden_states = self.norm(hidden_states)293 294        if output_hidden_states:295            encoder_states += (hidden_states,)296 297        return BaseModelOutput(298            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions299        )300 301    # Copied from transformers.models.bart.modeling_bart.BartPreTrainedModel._update_full_mask302    def _update_full_mask(303        self,304        attention_mask: Union[torch.Tensor, None],305        inputs_embeds: torch.Tensor,306    ):307        if attention_mask is not None:308            if self.config._attn_implementation == "flash_attention_2":309                attention_mask = attention_mask if 0 in attention_mask else None310            elif self.config._attn_implementation == "sdpa":311                # output_attentions=True & head_mask can not be supported when using SDPA, fall back to312                # the manual implementation that requires a 4D causal mask in all cases.313                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]314                attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, inputs_embeds.dtype)315            elif self.config._attn_implementation == "flex_attention":316                if isinstance(attention_mask, torch.Tensor):317                    attention_mask = make_flex_block_causal_mask(attention_mask, is_causal=False)318            else:319                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]320                attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)321 322        return attention_mask323 324 325class DiaDecoderLayer(GradientCheckpointingLayer):326    def __init__(self, config: DiaDecoderConfig, layer_idx: int):327        super().__init__()328        self.embed_dim = config.hidden_size329        self.self_attention = DiaSelfAttention(config, layer_idx, is_causal=True)330        self.cross_attention = DiaCrossAttention(config, layer_idx)331        self.pre_sa_norm = DiaRMSNorm(config.hidden_size, eps=config.norm_eps)332        self.pre_ca_norm = DiaRMSNorm(config.hidden_size, eps=config.norm_eps)333        self.pre_mlp_norm = DiaRMSNorm(config.hidden_size, eps=config.norm_eps)334        self.mlp = DiaMLP(config)335 336    def forward(337        self,338        hidden_states: torch.Tensor,339        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC340        attention_mask: Optional[torch.Tensor] = None,341        encoder_hidden_states: Optional[torch.Tensor] = None,342        encoder_attention_mask: Optional[torch.Tensor] = None,343        past_key_values: Optional[EncoderDecoderCache] = None,344        cache_position: Optional[torch.LongTensor] = None,345        **kwargs,346    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:347        self_attn_cache = past_key_values348        if isinstance(self_attn_cache, EncoderDecoderCache):349            self_attn_cache = self_attn_cache.self_attention_cache350 351        residual = hidden_states352        normed_states = self.pre_sa_norm(hidden_states)353        self_attn_output, self_attn_weights = self.self_attention(354            normed_states,355            position_embeddings,356            attention_mask,357            # Needs to be an arg in order to function properly358            # on inplace operations to be carried (e.g. compile)359            self_attn_cache,360            cache_position=cache_position,361            **kwargs,362        )363        hidden_states = residual + self_attn_output364 365        residual = hidden_states366        normed_states = self.pre_ca_norm(hidden_states)367        cross_states, cross_attn_weights = self.cross_attention(368            normed_states,369            encoder_hidden_states,370            attention_mask=encoder_attention_mask,371            past_key_values=past_key_values,372            **kwargs,373        )374        hidden_states = residual + cross_states375 376        residual = hidden_states377        normed_states = self.pre_mlp_norm(hidden_states)378        mlp_out = self.mlp(normed_states)379        hidden_states = residual + mlp_out380 381        return hidden_states, self_attn_weights, cross_attn_weights382 383 384class DiaDecoder(DiaPreTrainedModel):385    """Transformer Decoder Stack using DenseGeneral."""386 387    def __init__(self, config: DiaDecoderConfig):388        super().__init__(config)389        self.num_channels = config.num_channels390        self.vocab_size = config.vocab_size391        self.embeddings = DiaMultiChannelEmbedding(config)392        self.rotary_embeddings = DiaRotaryEmbedding(config)393        self.layers = nn.ModuleList(394            [DiaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]395        )396        self.norm = DiaRMSNorm(config.hidden_size, eps=config.norm_eps)397 398    @auto_docstring399    @can_return_tuple400    def forward(401        self,402        input_ids: torch.Tensor,403        position_ids: Optional[torch.LongTensor] = None,404        attention_mask: Optional[torch.Tensor] = None,405        encoder_hidden_states: Optional[torch.FloatTensor] = None,406        encoder_attention_mask: Optional[torch.LongTensor] = None,407        past_key_values: Optional[EncoderDecoderCache] = None,408        output_attentions: Optional[bool] = False,409        output_hidden_states: Optional[bool] = False,410        cache_position: Optional[torch.LongTensor] = None,411        **kwargs,412    ) -> Union[BaseModelOutputWithPastAndCrossAttentions, tuple]:413        r"""414        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks)`):415            The original `decoder_input_ids` in 3D shape to facilitate more efficient computations.416 417            [What are input IDs?](../glossary#input-ids)418        """419 420        batch_size, seq_length = input_ids.size()[:-1]421        past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0422        if cache_position is None:423            cache_position = torch.arange(424                past_key_values_length, past_key_values_length + seq_length, device=input_ids.device425            )426        if position_ids is None:427            position_ids = cache_position[None, :]428 429        # RoPE430        hidden_states = self.embeddings(input_ids)431        position_embeddings = self.rotary_embeddings(hidden_states, position_ids)432 433        if attention_mask is None and not is_torchdynamo_compiling():434            # required mask seq length can be calculated via length of past cache435            mask_seq_length = past_key_values_length + seq_length436            attention_mask = torch.ones(batch_size, mask_seq_length, device=input_ids.device)437 438        attention_mask = create_causal_mask(439            config=self.config,440            input_embeds=hidden_states,441            attention_mask=attention_mask,442            cache_position=cache_position,443            past_key_values=past_key_values,444            position_ids=position_ids,445        )446        encoder_attention_mask = self._update_cross_attn_mask(447            encoder_hidden_states,448            encoder_attention_mask,449            hidden_states.shape[:2],450            hidden_states,451        )452 453        all_hidden_states = () if output_hidden_states else None454        all_self_attns = () if output_attentions else None455        all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None456 457        for layer in self.layers:458            if output_hidden_states:459                all_hidden_states += (hidden_states,)460 461            layer_outputs = layer(462                hidden_states,463                position_embeddings,464                attention_mask,465                encoder_hidden_states,466                encoder_attention_mask=encoder_attention_mask,467                past_key_values=past_key_values,468                cache_position=cache_position,469                **kwargs,470            )471            hidden_states = layer_outputs[0]472 473            if output_attentions:474                all_self_attns = all_self_attns + (layer_outputs[1],)475 476                if encoder_hidden_states is not None:477                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)478 479        hidden_states = self.norm(hidden_states)480 481        if output_hidden_states:482            all_hidden_states += (hidden_states,)483 484        return BaseModelOutputWithPastAndCrossAttentions(485            last_hidden_state=hidden_states,486            past_key_values=past_key_values,487            hidden_states=all_hidden_states,488            attentions=all_self_attns,489            cross_attentions=all_cross_attentions,490        )491 492    # Copied from transformers.models.bart.modeling_bart.BartPreTrainedModel._update_cross_attn_mask493    def _update_cross_attn_mask(494        self,495        encoder_hidden_states: Union[torch.Tensor, None],496        encoder_attention_mask: Union[torch.Tensor, None],497        input_shape: torch.Size,498        inputs_embeds: torch.Tensor,499    ):500        # expand encoder attention mask501        if encoder_hidden_states is not None and encoder_attention_mask is not None:502            if self.config._attn_implementation == "flash_attention_2":503                encoder_attention_mask = encoder_attention_mask if 0 in encoder_attention_mask else None504            elif self.config._attn_implementation == "sdpa":505                # output_attentions=True & cross_attn_head_mask can not be supported when using SDPA, and we fall back on506                # the manual implementation that requires a 4D causal mask in all cases.507                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]508                encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(509                    encoder_attention_mask,510                    inputs_embeds.dtype,511                    tgt_len=input_shape[-1],512                )513            elif self.config._attn_implementation == "flex_attention":514                if isinstance(encoder_attention_mask, torch.Tensor):515                    encoder_attention_mask = make_flex_block_causal_mask(516                        encoder_attention_mask,517                        query_length=input_shape[-1],518                        is_causal=False,519                    )520            else:521                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]522                encoder_attention_mask = _prepare_4d_attention_mask(523                    encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]524                )525 526        return encoder_attention_mask527 528 529@auto_docstring(530    custom_intro="""531    The bare Dia model outputting raw hidden-states without any specific head on top.532    """533)534class DiaModel(DiaPreTrainedModel):535    def __init__(self, config: DiaConfig):536        super().__init__(config)537        self.config = config538        self.encoder = DiaEncoder(config.encoder_config)539        self.decoder = DiaDecoder(config.decoder_config)540        self.post_init()541 542    def get_encoder(self):543        return self.encoder544 545    @auto_docstring546    @can_return_tuple547    def forward(548        self,549        input_ids: Optional[torch.LongTensor] = None,550        attention_mask: Optional[torch.LongTensor] = None,551        decoder_input_ids: Optional[torch.LongTensor] = None,552        decoder_position_ids: Optional[torch.LongTensor] = None,553        decoder_attention_mask: Optional[torch.LongTensor] = None,554        encoder_outputs: Optional[Union[BaseModelOutput, tuple]] = None,555        past_key_values: Optional[EncoderDecoderCache] = None,556        use_cache: Optional[bool] = None,557        output_attentions: Optional[bool] = None,558        output_hidden_states: Optional[bool] = None,559        cache_position: Optional[torch.LongTensor] = None,560        **kwargs,561    ) -> Union[tuple, Seq2SeqModelOutput]:562        r"""563        decoder_input_ids (`torch.LongTensor` of shape `(batch_size * num_codebooks, target_sequence_length)564        or (batch_size, target_sequence_length, num_codebooks)`, *optional*):565            1. (batch_size * num_codebooks, target_sequence_length): corresponds to the general use case where566            the audio input codebooks are flattened into the batch dimension. This also aligns with the flat-567            tened audio logits which are used to calculate the loss.568 569            2. (batch_size, sequence_length, num_codebooks): corresponds to the internally used shape of570            Dia to calculate embeddings and subsequent steps more efficiently.571 572            If no `decoder_input_ids` are provided, it will create a tensor of `bos_token_id` with shape573            `(batch_size, 1, num_codebooks)`. Indices can be obtained using the [`DiaProcessor`]. See574            [`DiaProcessor.__call__`] for more details.575 576            [What are decoder input IDs?](../glossary#decoder-input-ids)577        decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`):578            Indices of positions of each input sequence tokens in the position embeddings.579            Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings`.580 581            [What are position IDs?](../glossary#position-ids)582        """583 584        if input_ids is None and encoder_outputs is None:585            raise ValueError(586                "You should either provide text ids or the cached text encodings. Neither has been found."587            )588 589        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions590        output_hidden_states = (591            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states592        )593        use_cache = use_cache if use_cache is not None else self.config.use_cache594 595        if self.is_gradient_checkpointing and self.training:596            if use_cache:597                logger.warning_once(598                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."599                )600                use_cache = False601 602        if use_cache and past_key_values is None:603            past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))604 605        if encoder_outputs is None:606            encoder_outputs = self.encoder(607                input_ids=input_ids,608                attention_mask=attention_mask,609                output_attentions=output_attentions,610                output_hidden_states=output_hidden_states,611                **kwargs,612            )613        # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput614        elif not isinstance(encoder_outputs, BaseModelOutput):615            encoder_outputs = BaseModelOutput(616                last_hidden_state=encoder_outputs[0],617                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,618                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,619            )620 621        # On default we initialize the decoder with bos tokens if nothing has been provided622        bsz, seq_len, channels = (encoder_outputs[0].shape[0], -1, self.config.decoder_config.num_channels)623        if decoder_input_ids is None:624            decoder_input_ids = torch.full(625                size=(bsz, 1, channels), fill_value=self.config.bos_token_id, device=self.device626            )627        # Ensure 3D628        if decoder_input_ids.ndim == 2:629            decoder_input_ids = decoder_input_ids.reshape(bsz, channels, seq_len).transpose(1, 2)630 631        decoder_outputs = self.decoder(632            input_ids=decoder_input_ids,633            position_ids=decoder_position_ids,634            attention_mask=decoder_attention_mask,635            encoder_hidden_states=encoder_outputs[0],636            encoder_attention_mask=attention_mask,637            past_key_values=past_key_values,638            output_attentions=output_attentions,639            output_hidden_states=output_hidden_states,640            use_cache=use_cache,641            cache_position=cache_position,642            **kwargs,643        )644 645        return Seq2SeqModelOutput(646            last_hidden_state=decoder_outputs.last_hidden_state,647            past_key_values=decoder_outputs.past_key_values,648            decoder_hidden_states=decoder_outputs.hidden_states,649            decoder_attentions=decoder_outputs.attentions,650            cross_attentions=decoder_outputs.cross_attentions,651            encoder_last_hidden_state=encoder_outputs[0],652            encoder_hidden_states=encoder_outputs.hidden_states,653            encoder_attentions=encoder_outputs.attentions,654        )655 656 657@auto_docstring(658    custom_intro="""659    The Dia model consisting of a (byte) text encoder and audio decoder with a prediction head on top.660    """661)662class DiaForConditionalGeneration(DiaPreTrainedModel, DiaGenerationMixin):663    base_model_prefix = "model"664 665    def __init__(self, config: DiaConfig):666        super().__init__(config)667        self.config = config668        self.model = DiaModel(config)669 670        self.num_channels = config.decoder_config.num_channels671        self.vocab_size = config.decoder_config.vocab_size672        self.logits_dense = nn.Linear(673            config.decoder_config.hidden_size, (self.num_channels * self.vocab_size), bias=False674        )675        self.loss_type = "ForMaskedLM"676 677        # Initialize weights and apply final processing678        self.post_init()679 680    def get_encoder(self):681        return self.model.get_encoder()682 683    def get_decoder(self):684        return self.model.get_decoder()685 686    @auto_docstring687    @can_return_tuple688    def forward(689        self,690        input_ids: Optional[torch.LongTensor] = None,691        attention_mask: Optional[torch.LongTensor] = None,692        decoder_input_ids: Optional[torch.LongTensor] = None,693        decoder_position_ids: Optional[torch.LongTensor] = None,694        decoder_attention_mask: Optional[torch.LongTensor] = None,695        encoder_outputs: Optional[Union[BaseModelOutput, tuple]] = None,696        past_key_values: Optional[EncoderDecoderCache] = None,697        use_cache: Optional[bool] = None,698        output_attentions: Optional[bool] = None,699        output_hidden_states: Optional[bool] = None,700        labels: Optional[torch.LongTensor] = None,701        cache_position: Optional[torch.LongTensor] = None,702        **kwargs,703    ) -> Union[tuple, Seq2SeqLMOutput]:704        r"""705        decoder_input_ids (`torch.LongTensor` of shape `(batch_size * num_codebooks, target_sequence_length)706        or (batch_size, target_sequence_length, num_codebooks)`, *optional*):707            1. (batch_size * num_codebooks, target_sequence_length): corresponds to the general use case where708            the audio input codebooks are flattened into the batch dimension. This also aligns with the flat-709            tened audio logits which are used to calculate the loss.710 711            2. (batch_size, sequence_length, num_codebooks): corresponds to the internally used shape of712            Dia to calculate embeddings and subsequent steps more efficiently.713 714            If no `decoder_input_ids` are provided, it will create a tensor of `bos_token_id` with shape715            `(batch_size, 1, num_codebooks)`. Indices can be obtained using the [`DiaProcessor`]. See716            [`DiaProcessor.__call__`] for more details.717 718            [What are decoder input IDs?](../glossary#decoder-input-ids)719        decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`):720            Indices of positions of each input sequence tokens in the position embeddings.721            Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings`.722 723            [What are position IDs?](../glossary#position-ids)724        labels (`torch.LongTensor` of shape `(batch_size * num_codebooks,)`, *optional*):725            Labels for computing the masked language modeling loss. Indices should either be in726            `[0, ..., config.decoder_config.vocab_size - 1]` or -100. Tokens with indices set to `-100`727            are ignored (masked).728        """729 730        outputs = self.model(731            input_ids=input_ids,732            attention_mask=attention_mask,733            decoder_input_ids=decoder_input_ids,734            decoder_position_ids=decoder_position_ids,735            decoder_attention_mask=decoder_attention_mask,736            encoder_outputs=encoder_outputs,737            past_key_values=past_key_values,738            use_cache=use_cache,739            output_attentions=output_attentions,740            output_hidden_states=output_hidden_states,741            cache_position=cache_position,742            **kwargs,743        )744 745        last_hidden_state = outputs[0]746        batch_size = last_hidden_state.shape[0]747        # 3D <-> 2D makes it necessary to prioritize channel dim748        audio_logits = (749            self.logits_dense(last_hidden_state)750            .view((batch_size, -1, self.num_channels, self.vocab_size))751            .transpose(1, 2)752            .contiguous()753            .view(batch_size * self.num_channels, -1, self.vocab_size)754        )755 756        loss = None757        if labels is not None:758            loss = self.loss_function(logits=audio_logits, labels=labels, vocab_size=self.vocab_size, **kwargs)759 760        return Seq2SeqLMOutput(761            loss=loss,762            logits=audio_logits,763            past_key_values=outputs.past_key_values,764            decoder_hidden_states=outputs.decoder_hidden_states,765            decoder_attentions=outputs.decoder_attentions,766            cross_attentions=outputs.cross_attentions,767            encoder_last_hidden_state=outputs.encoder_last_hidden_state,768            encoder_hidden_states=outputs.encoder_hidden_states,769            encoder_attentions=outputs.encoder_attentions,770        )771 772 773__all__ = ["DiaModel", "DiaPreTrainedModel", "DiaForConditionalGeneration"]774