diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/bloom/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/bloom/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ba39d13cedcb6ce497397c40fc85159f0bb0af16
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/bloom/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_bloom import *
+ from .modeling_bloom import *
+ from .tokenization_bloom import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/bloom/modeling_bloom.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/bloom/modeling_bloom.py
new file mode 100644
index 0000000000000000000000000000000000000000..71099024429aeb4fbb478804f877cf2ce7c21a8a
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/bloom/modeling_bloom.py
@@ -0,0 +1,987 @@
+# Copyright 2022 HuggingFace Inc. team and BigScience workshop.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch BLOOM model."""
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
+from torch.nn import functional as F
+
+from ...cache_utils import Cache, DynamicCache, StaticCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ QuestionAnsweringModelOutput,
+ SequenceClassifierOutputWithPast,
+ TokenClassifierOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ auto_docstring,
+ logging,
+)
+from .configuration_bloom import BloomConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+def build_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:
+ """
+ Link to paper: https://huggingface.co/papers/2108.12409 Alibi tensor is not causal as the original paper mentions, it
+ relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value
+ `softmax(l+a) = softmax(l)`. Based on
+ https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742
+ TODO @thomasw21 this doesn't work as nicely due to the masking strategy, and so masking varies slightly.
+
+ Args:
+ Returns tensor shaped (batch_size * num_heads, 1, max_seq_len)
+ attention_mask (`torch.Tensor`):
+ Token-wise attention mask, this should be of shape (batch_size, max_seq_len).
+ num_heads (`int`):
+ number of heads
+ dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`):
+ dtype of the output tensor
+ """
+ batch_size, seq_length = attention_mask.shape
+ closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
+ base = torch.tensor(
+ 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32
+ )
+ powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32)
+ slopes = torch.pow(base, powers)
+
+ if closest_power_of_2 != num_heads:
+ extra_base = torch.tensor(
+ 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32
+ )
+ num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
+ extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32)
+ slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
+
+ # Note: alibi will added to the attention bias that will be applied to the query, key product of attention
+ # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length)
+ # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length)
+ # => the query_length dimension will then be broadcasted correctly
+ # This is more or less identical to T5's relative position bias:
+ # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527
+ arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :]
+ alibi = slopes[..., None] * arange_tensor
+ return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype)
+
+
+def dropout_add(x: torch.Tensor, residual: torch.Tensor, prob: float, training: bool) -> torch.Tensor:
+ """
+ Dropout add function
+
+ Args:
+ x (`torch.tensor`):
+ input tensor
+ residual (`torch.tensor`):
+ residual tensor
+ prob (`float`):
+ dropout probability
+ training (`bool`):
+ training mode
+ """
+ out = F.dropout(x, p=prob, training=training)
+ out = residual + out
+ return out
+
+
+def bloom_gelu_forward(x: torch.Tensor) -> torch.Tensor:
+ """
+ Custom bias GELU function. Adapted from Megatron-DeepSpeed code. Here we use a simple implementation (inference) to
+ make the model jitable.
+
+ Args:
+ x (`torch.tensor`):
+ input hidden states
+ """
+ return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+
+
+def bloom_gelu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+ """
+ gradient of tanh approximation of gelu gradient of actual gelu is: 0.5 * (1. + torch.erf(x * 0.70710678)) +
+ 0.3989423 * x * torch.exp(-0.5 * x * x)
+
+ Args:
+ g (`torch.tensor`):
+ gradient output tensor
+ x (`torch.tensor`):
+ input tensor
+ """
+ x = x[0] # x is a tuple of 1 element, needs to unpack it first
+ tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
+ # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
+ ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
+ return ff * g
+
+
+class GeLUFunction(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, input: torch.Tensor) -> torch.Tensor:
+ ctx.save_for_backward(input)
+ return bloom_gelu_forward(input)
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+ input = ctx.saved_tensors
+ tmp = bloom_gelu_back(grad_output, input)
+ return tmp
+
+
+class BloomGelu(nn.Module):
+ """
+ Partly copied from Megatron-DeepSpeed code and adapted for our needs
+ """
+
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return GeLUFunction.apply(x)
+
+
+class BloomAttention(nn.Module):
+ def __init__(self, config: BloomConfig, layer_idx: int | None = None):
+ super().__init__()
+
+ self.pretraining_tp = config.pretraining_tp
+ self.slow_but_exact = config.slow_but_exact
+
+ self.hidden_size = config.hidden_size
+ self.num_heads = config.n_head
+ self.head_dim = self.hidden_size // self.num_heads
+ self.split_size = self.hidden_size
+ self.hidden_dropout = config.hidden_dropout
+
+ if self.head_dim * self.num_heads != self.hidden_size:
+ raise ValueError(
+ f"`hidden_size` must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+
+ # Layer-wise attention scaling
+ self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
+ self.beta = 1.0
+ self.layer_idx = layer_idx
+ if layer_idx is None:
+ logger.warning_once(
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ self.query_key_value = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=True)
+ self.dense = nn.Linear(self.hidden_size, self.hidden_size)
+ self.attention_dropout = nn.Dropout(config.attention_dropout)
+
+ def _reshape(self, fused_qkv: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Split the last dimension into (num_heads, head_dim) and reshapes to (bs, heads, len, dim) shape
+ without making any copies, results share same memory storage as `fused_qkv`
+
+ Args:
+ fused_qkv (`torch.tensor`): [batch_size, seq_length, num_heads * 3 * head_dim]
+
+ Returns:
+ query: [batch_size, num_heads, seq_length, head_dim]
+ key: [batch_size, num_heads, seq_length, head_dim]
+ value: [batch_size, num_heads, seq_length, head_dim]
+ """
+ batch_size, seq_length, three_times_hidden_size = fused_qkv.shape
+ fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads, 3, self.head_dim)
+ query_layer = fused_qkv[..., 0, :].transpose(1, 2)
+ key_layer = fused_qkv[..., 1, :].transpose(1, 2)
+ value_layer = fused_qkv[..., 2, :].transpose(1, 2)
+ return query_layer, key_layer, value_layer
+
+ def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Merge heads together over the last dimension
+
+ Args:
+ x (`torch.tensor`): [batch_size * num_heads, seq_length, head_dim]
+
+ Returns:
+ torch.tensor: [batch_size, seq_length, num_heads * head_dim]
+ """
+ # What we want to achieve is:
+ # batch_size * num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads * head_dim
+ batch_size_and_num_heads, seq_length, _ = x.shape
+ batch_size = batch_size_and_num_heads // self.num_heads
+
+ # First view to decompose the batch size
+ # batch_size * num_heads, seq_length, head_dim -> batch_size, num_heads, seq_length, head_dim
+ x = x.view(batch_size, self.num_heads, seq_length, self.head_dim)
+
+ # batch_size, num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads, head_dim
+ x = x.permute(0, 2, 1, 3)
+
+ # batch_size, seq_length, num_heads, head_dim -> batch_size, seq_length, num_heads * head_dim
+ return x.reshape(batch_size, seq_length, self.num_heads * self.head_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ residual: torch.Tensor,
+ alibi: torch.Tensor,
+ attention_mask: torch.Tensor,
+ layer_past: Cache | None = None,
+ use_cache: bool = False,
+ output_attentions: bool = False,
+ **kwargs,
+ ):
+ batch_size, q_length, _ = hidden_states.shape
+ fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size]
+ # 3 x [batch_size, num_heads, seq_length, head_dim]
+ query_layer, key_layer, value_layer = self._reshape(fused_qkv)
+
+ if layer_past is not None:
+ key_layer, value_layer = layer_past.update(key_layer, value_layer, self.layer_idx)
+
+ # reshape qkv for further computations
+ query_layer = query_layer.reshape(batch_size * self.num_heads, -1, self.head_dim)
+ key_layer = key_layer.reshape(batch_size * self.num_heads, -1, self.head_dim).transpose(-1, -2)
+ value_layer = value_layer.reshape(batch_size * self.num_heads, -1, self.head_dim)
+
+ # [batch_size * num_heads, q_length, kv_length]
+ attention_scores = alibi.baddbmm(
+ batch1=query_layer,
+ batch2=key_layer,
+ beta=self.beta,
+ alpha=self.inv_norm_factor,
+ )
+
+ # change view to [batch_size, num_heads, q_length, kv_length]
+ attn_weights = attention_scores.view(batch_size, self.num_heads, q_length, -1)
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype
+ attention_probs = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_layer.dtype)
+
+ # [batch_size, num_heads, q_length, kv_length]
+ attention_probs = self.attention_dropout(attention_probs)
+
+ # change view [batch_size x num_heads, q_length, kv_length]
+ attention_probs_reshaped = attention_probs.view(batch_size * self.num_heads, q_length, -1)
+
+ # matmul: [batch_size * num_heads, q_length, head_dim]
+ context_layer = torch.bmm(attention_probs_reshaped, value_layer)
+
+ # change view [batch_size, q_length, num_heads * head_dim]
+ context_layer = self._merge_heads(context_layer)
+
+ # aggregate results across tp ranks. See here: https://github.com/pytorch/pytorch/issues/76232
+ if self.pretraining_tp > 1 and self.slow_but_exact:
+ slices = self.hidden_size / self.pretraining_tp
+ output_tensor = torch.zeros_like(context_layer)
+ for i in range(self.pretraining_tp):
+ output_tensor = output_tensor + F.linear(
+ context_layer[:, :, int(i * slices) : int((i + 1) * slices)],
+ self.dense.weight[:, int(i * slices) : int((i + 1) * slices)],
+ )
+ else:
+ output_tensor = self.dense(context_layer)
+
+ output_tensor = dropout_add(output_tensor, residual, self.hidden_dropout, self.training)
+ return output_tensor, attention_probs
+
+
+class BloomMLP(nn.Module):
+ def __init__(self, config: BloomConfig):
+ super().__init__()
+ hidden_size = config.hidden_size
+
+ self.pretraining_tp = config.pretraining_tp
+ self.slow_but_exact = config.slow_but_exact
+ self.dense_h_to_4h = nn.Linear(hidden_size, 4 * hidden_size)
+ self.gelu_impl = BloomGelu()
+ self.dense_4h_to_h = nn.Linear(4 * hidden_size, hidden_size)
+ self.hidden_dropout = config.hidden_dropout
+
+ def forward(self, hidden_states: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.gelu_impl(self.dense_h_to_4h(hidden_states))
+
+ if self.pretraining_tp > 1 and self.slow_but_exact:
+ intermediate_output = torch.zeros_like(residual)
+ slices = self.dense_4h_to_h.weight.shape[-1] / self.pretraining_tp
+ for i in range(self.pretraining_tp):
+ intermediate_output = intermediate_output + F.linear(
+ hidden_states[:, :, int(i * slices) : int((i + 1) * slices)],
+ self.dense_4h_to_h.weight[:, int(i * slices) : int((i + 1) * slices)],
+ )
+ else:
+ intermediate_output = self.dense_4h_to_h(hidden_states)
+
+ output = dropout_add(intermediate_output, residual, self.hidden_dropout, self.training)
+
+ return output
+
+
+class BloomBlock(GradientCheckpointingLayer):
+ def __init__(self, config: BloomConfig, layer_idx: int | None = None):
+ super().__init__()
+ hidden_size = config.hidden_size
+
+ self.input_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+ self.num_heads = config.n_head
+ self.self_attention = BloomAttention(config, layer_idx)
+ self.post_attention_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+
+ self.mlp = BloomMLP(config)
+
+ self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
+ self.hidden_dropout = config.hidden_dropout
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ alibi: torch.Tensor,
+ attention_mask: torch.Tensor,
+ layer_past: Cache | None = None,
+ use_cache: bool = False,
+ output_attentions: bool = False,
+ **kwargs,
+ ):
+ # hidden_states: [batch_size, seq_length, hidden_size]
+
+ # Layer norm at the beginning of the transformer layer.
+ layernorm_output = self.input_layernorm(hidden_states)
+
+ # Layer norm post the self attention.
+ if self.apply_residual_connection_post_layernorm:
+ residual = layernorm_output
+ else:
+ residual = hidden_states
+
+ # Self attention.
+ attention_output, attn_weights = self.self_attention(
+ layernorm_output,
+ residual,
+ layer_past=layer_past,
+ attention_mask=attention_mask,
+ alibi=alibi,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ )
+
+ layernorm_output = self.post_attention_layernorm(attention_output)
+
+ # Get residual
+ if self.apply_residual_connection_post_layernorm:
+ residual = layernorm_output
+ else:
+ residual = attention_output
+
+ # MLP.
+ output = self.mlp(layernorm_output, residual)
+
+ return output, attn_weights # hidden_states, attentions
+
+
+@auto_docstring
+class BloomPreTrainedModel(PreTrainedModel):
+ config: BloomConfig
+ base_model_prefix = "transformer"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["BloomBlock"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _can_compile_fullgraph = True
+
+
+@auto_docstring
+class BloomModel(BloomPreTrainedModel):
+ def __init__(self, config: BloomConfig):
+ super().__init__(config)
+
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.n_head
+
+ # Embedding + LN Embedding
+ self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim)
+ self.word_embeddings_layernorm = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
+
+ # Transformer blocks
+ self.h = nn.ModuleList([BloomBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+
+ # Final Layer Norm
+ self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
+
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def build_alibi_tensor(self, attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:
+ return build_alibi_tensor(attention_mask, num_heads, dtype)
+
+ def get_input_embeddings(self):
+ return self.word_embeddings
+
+ def set_input_embeddings(self, new_embeddings: torch.Tensor):
+ self.word_embeddings = new_embeddings
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, ...] | BaseModelOutputWithPastAndCrossAttentions:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+ `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
+ (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+ If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+ `input_ids`.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ """
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if self.gradient_checkpointing and self.training and use_cache:
+ logger.warning_once(
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+ )
+ use_cache = False
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ batch_size, seq_length, _ = inputs_embeds.shape
+ past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+ seq_length_with_past = seq_length + past_length
+
+ hidden_states = self.word_embeddings_layernorm(inputs_embeds)
+
+ all_self_attentions = () if output_attentions else None
+ all_hidden_states = () if output_hidden_states else None
+
+ # Compute alibi tensor: check build_alibi_tensor documentation
+ if attention_mask is None:
+ attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
+ else:
+ attention_mask = attention_mask.to(hidden_states.device)
+
+ alibi = self.build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype)
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ )
+
+ for i, block in enumerate(self.h):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ outputs = block(
+ hidden_states,
+ layer_past=past_key_values,
+ attention_mask=causal_mask,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ alibi=alibi,
+ )
+
+ hidden_states = outputs[0]
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (outputs[1],)
+
+ # Add last hidden state
+ hidden_states = self.ln_f(hidden_states)
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(
+ v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attentions] if v is not None
+ )
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The Bloom Model transformer with a language modeling head on top (linear layer with weights tied to the input
+ embeddings).
+ """
+)
+class BloomForCausalLM(BloomPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "transformer.word_embeddings.weight"}
+
+ def __init__(self, config: BloomConfig):
+ super().__init__(config)
+ self.transformer = BloomModel(config)
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def set_output_embeddings(self, new_embeddings: torch.Tensor):
+ self.lm_head = new_embeddings
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ use_cache=True,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ # Overwritten because of the fixed-shape attention mask creation
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ # This part differs from other models because BLOOM needs a 2D mask to construct alibi tensor
+ # The only difference is the usage of 2D instead of 4D mask, but the shape will be static
+ if isinstance(past_key_values, StaticCache) and attention_mask is not None:
+ target_length = past_key_values.get_max_cache_shape()
+ batch_size, seq_length = attention_mask.shape
+ diff = target_length - seq_length
+
+ new_attn_mask = torch.zeros(batch_size, diff, device=attention_mask.device, dtype=attention_mask.dtype)
+ attention_mask = torch.cat([attention_mask, new_attn_mask], dim=-1)
+ model_inputs["attention_mask"] = attention_mask
+
+ return model_inputs
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | CausalLMOutputWithCrossAttentions:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+ `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
+ (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+ If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+ `input_ids`.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+ `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+ are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ transformer_outputs = self.transformer(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = transformer_outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits,
+ labels,
+ vocab_size=self.config.vocab_size,
+ num_items_in_batch=kwargs.get("num_items_in_batch"),
+ )
+
+ if not return_dict:
+ output = (logits,) + transformer_outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=logits,
+ past_key_values=transformer_outputs.past_key_values,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The Bloom Model transformer with a sequence classification head on top (linear layer).
+
+ [`BloomForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+ (e.g. GPT-1) do.
+
+ Since it does classification on the last token, it requires to know the position of the last token. If a
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+ each row of the batch).
+ """
+)
+class BloomForSequenceClassification(BloomPreTrainedModel):
+ def __init__(self, config: BloomConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.transformer = BloomModel(config)
+ self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | SequenceClassifierOutputWithPast:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+ `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
+ (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+ If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+ `input_ids`.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ transformer_outputs = self.transformer(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = transformer_outputs[0]
+ logits = self.score(hidden_states)
+
+ if input_ids is not None:
+ batch_size = input_ids.shape[0]
+ else:
+ batch_size = inputs_embeds.shape[0]
+
+ if self.config.pad_token_id is None and batch_size != 1:
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+ if self.config.pad_token_id is None:
+ last_non_pad_token = -1
+ elif input_ids is not None:
+ # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+ non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+ token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+ last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+ else:
+ last_non_pad_token = -1
+ logger.warning_once(
+ f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+ "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+ )
+
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+ loss = None
+ if labels is not None:
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(pooled_logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(pooled_logits, labels)
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(pooled_logits, labels)
+ if not return_dict:
+ output = (pooled_logits,) + transformer_outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return SequenceClassifierOutputWithPast(
+ loss=loss,
+ logits=pooled_logits,
+ past_key_values=transformer_outputs.past_key_values,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+@auto_docstring
+class BloomForTokenClassification(BloomPreTrainedModel):
+ def __init__(self, config: BloomConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.transformer = BloomModel(config)
+ if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
+ classifier_dropout = config.classifier_dropout
+ elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
+ classifier_dropout = config.hidden_dropout
+ else:
+ classifier_dropout = 0.1
+ self.dropout = nn.Dropout(classifier_dropout)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | TokenClassifierOutput:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+ `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
+ (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+ If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+ `input_ids`.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ transformer_outputs = self.transformer(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = transformer_outputs[0]
+ hidden_states = self.dropout(hidden_states)
+ logits = self.classifier(hidden_states)
+
+ loss = None
+ if labels is not None:
+ # move labels to correct device
+ labels = labels.to(logits.device)
+ batch_size, seq_length = labels.shape
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(
+ logits.view(batch_size * seq_length, self.num_labels), labels.view(batch_size * seq_length)
+ )
+
+ if not return_dict:
+ output = (logits,) + transformer_outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+@auto_docstring
+class BloomForQuestionAnswering(BloomPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.transformer = BloomModel(config)
+ self.qa_outputs = nn.Linear(config.hidden_size, 2)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ start_positions: torch.LongTensor | None = None,
+ end_positions: torch.LongTensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | QuestionAnsweringModelOutput:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+ `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
+ (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+ If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+ `input_ids`.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.transformer(
+ input_ids,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ sequence_output = outputs[0]
+
+ logits = self.qa_outputs(sequence_output)
+ start_logits, end_logits = logits.split(1, dim=-1)
+ start_logits = start_logits.squeeze(-1).contiguous()
+ end_logits = end_logits.squeeze(-1).contiguous()
+
+ total_loss = None
+ if start_positions is not None and end_positions is not None:
+ # If we are on multi-GPU, split add a dimension
+ if len(start_positions.size()) > 1:
+ start_positions = start_positions.squeeze(-1)
+ if len(end_positions.size()) > 1:
+ end_positions = end_positions.squeeze(-1)
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
+ ignored_index = start_logits.size(1)
+ start_positions = start_positions.clamp(0, ignored_index)
+ end_positions = end_positions.clamp(0, ignored_index)
+
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+ start_loss = loss_fct(start_logits, start_positions)
+ end_loss = loss_fct(end_logits, end_positions)
+ total_loss = (start_loss + end_loss) / 2
+
+ if not return_dict:
+ output = (start_logits, end_logits) + outputs[2:]
+ return ((total_loss,) + output) if total_loss is not None else output
+
+ return QuestionAnsweringModelOutput(
+ loss=total_loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "BloomForCausalLM",
+ "BloomModel",
+ "BloomPreTrainedModel",
+ "BloomForSequenceClassification",
+ "BloomForTokenClassification",
+ "BloomForQuestionAnswering",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..879b53709bc673bcf28553a51175f06fa1e362c0
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_d_fine import *
+ from .modeling_d_fine import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/configuration_d_fine.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/configuration_d_fine.py
new file mode 100644
index 0000000000000000000000000000000000000000..793f038102d3ef82fc4b4c4a3793eff18d670161
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/configuration_d_fine.py
@@ -0,0 +1,242 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/d_fine/modular_d_fine.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_d_fine.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 Baidu Inc and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from huggingface_hub.dataclasses import strict
+
+from ...backbone_utils import consolidate_backbone_kwargs_to_config
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+from ..auto import AutoConfig
+
+
+# TODO: Attribute map assignment logic should be fixed in modular
+# as well as super() call parsing because otherwise we cannot re-write args after initialization
+@auto_docstring(checkpoint="ustc-community/dfine-xlarge-coco")
+@strict
+class DFineConfig(PreTrainedConfig):
+ r"""
+ initializer_bias_prior_prob (`float`, *optional*):
+ The prior probability used by the bias initializer to initialize biases for `enc_score_head` and `class_embed`.
+ If `None`, `prior_prob` computed as `prior_prob = 1 / (num_labels + 1)` while initializing model weights.
+ freeze_backbone_batch_norms (`bool`, *optional*, defaults to `True`):
+ Whether to freeze the batch normalization layers in the backbone.
+ encoder_in_channels (`list`, *optional*, defaults to `[512, 1024, 2048]`):
+ Multi level features input for encoder.
+ feat_strides (`list[int]`, *optional*, defaults to `[8, 16, 32]`):
+ Strides used in each feature map.
+ encode_proj_layers (`list[int]`, *optional*, defaults to `[2]`):
+ Indexes of the projected layers to be used in the encoder.
+ positional_encoding_temperature (`int`, *optional*, defaults to 10000):
+ The temperature parameter used to create the positional encodings.
+ encoder_activation_function (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+ `"relu"`, `"silu"` and `"gelu_new"` are supported.
+ eval_size (`tuple[int, int]`, *optional*):
+ Height and width used to computes the effective height and width of the position embeddings after taking
+ into account the stride.
+ normalize_before (`bool`, *optional*, defaults to `False`):
+ Determine whether to apply layer normalization in the transformer encoder layer before self-attention and
+ feed-forward modules.
+ hidden_expansion (`float`, *optional*, defaults to 1.0):
+ Expansion ratio to enlarge the dimension size of RepVGGBlock and CSPRepLayer.
+ num_queries (`int`, *optional*, defaults to 300):
+ Number of object queries.
+ decoder_in_channels (`list`, *optional*, defaults to `[256, 256, 256]`):
+ Multi level features dimension for decoder
+ num_feature_levels (`int`, *optional*, defaults to 3):
+ The number of input feature levels.
+ decoder_n_points (`int`, *optional*, defaults to 4):
+ The number of sampled keys in each feature level for each attention head in the decoder.
+ decoder_activation_function (`str`, *optional*, defaults to `"relu"`):
+ The non-linear activation function (function or string) in the decoder. If string, `"gelu"`,
+ `"relu"`, `"silu"` and `"gelu_new"` are supported.
+ num_denoising (`int`, *optional*, defaults to 100):
+ The total number of denoising tasks or queries to be used for contrastive denoising.
+ label_noise_ratio (`float`, *optional*, defaults to 0.5):
+ The fraction of denoising labels to which random noise should be added.
+ box_noise_scale (`float`, *optional*, defaults to 1.0):
+ Scale or magnitude of noise to be added to the bounding boxes.
+ learn_initial_query (`bool`, *optional*, defaults to `False`):
+ Indicates whether the initial query embeddings for the decoder should be learned during training
+ anchor_image_size (`tuple[int, int]`, *optional*):
+ Height and width of the input image used during evaluation to generate the bounding box anchors. If None, automatic generate anchor is applied.
+ with_box_refine (`bool`, *optional*, defaults to `True`):
+ Whether to apply iterative bounding box refinement, where each decoder layer refines the bounding boxes
+ based on the predictions from the previous layer.
+ matcher_alpha (`float`, *optional*, defaults to 0.25):
+ Parameter alpha used by the Hungarian Matcher.
+ matcher_gamma (`float`, *optional*, defaults to 2.0):
+ Parameter gamma used by the Hungarian Matcher.
+ matcher_class_cost (`float`, *optional*, defaults to 2.0):
+ The relative weight of the class loss used by the Hungarian Matcher.
+ matcher_bbox_cost (`float`, *optional*, defaults to 5.0):
+ The relative weight of the bounding box loss used by the Hungarian Matcher.
+ matcher_giou_cost (`float`, *optional*, defaults to 2.0):
+ The relative weight of the giou loss of used by the Hungarian Matcher.
+ use_focal_loss (`bool`, *optional*, defaults to `True`):
+ Parameter informing if focal focal should be used.
+ focal_loss_alpha (`float`, *optional*, defaults to 0.75):
+ Parameter alpha used to compute the focal loss.
+ focal_loss_gamma (`float`, *optional*, defaults to 2.0):
+ Parameter gamma used to compute the focal loss.
+ weight_loss_vfl (`float`, *optional*, defaults to 1.0):
+ Relative weight of the varifocal loss in the object detection loss.
+ weight_loss_bbox (`float`, *optional*, defaults to 5.0):
+ Relative weight of the L1 bounding box loss in the object detection loss.
+ weight_loss_giou (`float`, *optional*, defaults to 2.0):
+ Relative weight of the generalized IoU loss in the object detection loss.
+ weight_loss_fgl (`float`, *optional*, defaults to 0.15):
+ Relative weight of the fine-grained localization loss in the object detection loss.
+ weight_loss_ddf (`float`, *optional*, defaults to 1.5):
+ Relative weight of the decoupled distillation focal loss in the object detection loss.
+ eval_idx (`int`, *optional*, defaults to -1):
+ Index of the decoder layer to use for evaluation. If negative, counts from the end
+ (e.g., -1 means use the last layer). This allows for early prediction in the decoder
+ stack while still training later layers.
+ layer_scale (`float`, *optional*, defaults to `1.0`):
+ Scaling factor for the hidden dimension in later decoder layers. Used to adjust the
+ model capacity after the evaluation layer.
+ max_num_bins (`int`, *optional*, defaults to 32):
+ Maximum number of bins for the distribution-guided bounding box refinement.
+ Higher values allow for more fine-grained localization but increase computation.
+ reg_scale (`float`, *optional*, defaults to 4.0):
+ Scale factor for the regression distribution. Controls the range and granularity
+ of the bounding box refinement process.
+ depth_mult (`float`, *optional*, defaults to 1.0):
+ Multiplier for the number of blocks in RepNCSPELAN4 layers. Used to scale the model's
+ depth while maintaining its architecture.
+ top_prob_values (`int`, *optional*, defaults to 4):
+ Number of top probability values to consider from each corner's distribution.
+ lqe_hidden_dim (`int`, *optional*, defaults to 64):
+ Hidden dimension size for the Location Quality Estimator (LQE) network.
+ lqe_layers (`int`, *optional*, defaults to 2):
+ Number of layers in the Location Quality Estimator MLP.
+ decoder_offset_scale (`float`, *optional*, defaults to 0.5):
+ Offset scale used in deformable attention.
+ decoder_method (`str`, *optional*, defaults to `"default"`):
+ The method to use for the decoder: `"default"` or `"discrete"`.
+ up (`float`, *optional*, defaults to 0.5):
+ Controls the upper bounds of the Weighting Function.
+ """
+
+ model_type = "d_fine"
+ sub_configs = {"backbone_config": AutoConfig}
+ layer_types = ["basic", "bottleneck"]
+ attribute_map = {
+ "hidden_size": "d_model",
+ "num_attention_heads": "encoder_attention_heads",
+ }
+
+ initializer_range: float = 0.01
+ initializer_bias_prior_prob: float | None = None
+ layer_norm_eps: float = 1e-5
+ batch_norm_eps: float = 1e-5
+ backbone_config: dict | PreTrainedConfig | None = None
+ freeze_backbone_batch_norms: bool = True
+
+ # encoder HybridEncoder
+ encoder_hidden_dim: int = 256
+ encoder_in_channels: list[int] | tuple[int, ...] = (512, 1024, 2048)
+ feat_strides: list[int] | tuple[int, ...] = (8, 16, 32)
+ encoder_layers: int = 1
+ encoder_ffn_dim: int = 1024
+ encoder_attention_heads: int = 8
+ dropout: float | int = 0.0
+ activation_dropout: float | int = 0.0
+ encode_proj_layers: list[int] | tuple[int, ...] = (2,)
+ positional_encoding_temperature: int = 10000
+ encoder_activation_function: str = "gelu"
+ activation_function: str = "silu"
+ eval_size: int | None = None
+ normalize_before: bool = False
+ hidden_expansion: float = 1.0
+
+ # decoder DFineTransformer
+ d_model: int = 256
+ num_queries: int = 300
+ decoder_in_channels: list[int] | tuple[int, ...] = (256, 256, 256)
+ decoder_ffn_dim: int = 1024
+ num_feature_levels: int = 3
+ decoder_n_points: int | list[int] = 4
+ decoder_layers: int = 6
+ decoder_attention_heads: int = 8
+ decoder_activation_function: str = "relu"
+ attention_dropout: float | int = 0.0
+ num_denoising: int = 100
+ label_noise_ratio: float = 0.5
+ box_noise_scale: float = 1.0
+ learn_initial_query: bool = False
+ anchor_image_size: int | list[int] | None = None
+ with_box_refine: bool = True
+
+ # Loss
+ matcher_alpha: float = 0.25
+ matcher_gamma: float = 2.0
+ matcher_class_cost: float = 2.0
+ matcher_bbox_cost: float = 5.0
+ matcher_giou_cost: float = 2.0
+ use_focal_loss: bool = True
+ auxiliary_loss: bool = True
+ focal_loss_alpha: float = 0.75
+ focal_loss_gamma: float = 2.0
+ weight_loss_vfl: float = 1.0
+ weight_loss_bbox: float = 5.0
+ weight_loss_giou: float = 2.0
+ weight_loss_fgl: float = 0.15
+ weight_loss_ddf: float = 1.5
+ eos_coefficient: float = 1e-4
+ eval_idx: int = -1
+ layer_scale: int | float = 1.0
+ max_num_bins: int = 32
+ reg_scale: float = 4.0
+ depth_mult: float = 1.0
+ top_prob_values: int = 4
+ lqe_hidden_dim: int = 64
+ lqe_layers: int = 2
+ decoder_offset_scale: float = 0.5
+ decoder_method: str = "default"
+ up: float = 0.5
+ tie_word_embeddings: bool = True
+ is_encoder_decoder: bool = True
+
+ def __post_init__(self, **kwargs):
+ self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config(
+ backbone_config=self.backbone_config,
+ default_config_type="hgnet_v2",
+ default_config_kwargs={"out_indices": [2, 3, 4]},
+ **kwargs,
+ )
+ self.head_dim = self.d_model // self.decoder_attention_heads
+ super().__post_init__(**kwargs)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if isinstance(self.decoder_n_points, list):
+ if len(self.decoder_n_points) != self.num_feature_levels:
+ raise ValueError(
+ f"Length of decoder_n_points list ({len(self.decoder_n_points)}) must match num_feature_levels ({self.num_feature_levels})."
+ )
+
+ if self.head_dim * self.decoder_attention_heads != self.d_model:
+ raise ValueError(
+ f"Embedded dimension {self.d_model} must be divisible by decoder_attention_heads {self.decoder_attention_heads}"
+ )
+
+
+__all__ = ["DFineConfig"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/modeling_d_fine.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/modeling_d_fine.py
new file mode 100644
index 0000000000000000000000000000000000000000..43eee035bdf60c6a6642cae7db30238ed742e42c
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/modeling_d_fine.py
@@ -0,0 +1,2109 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/d_fine/modular_d_fine.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_d_fine.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 Baidu Inc and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from ... import initialization as init
+from ...activations import ACT2CLS
+from ...backbone_utils import load_backbone
+from ...image_transforms import center_to_corners_format, corners_to_center_format
+from ...modeling_outputs import BaseModelOutput
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...pytorch_utils import compile_compatible_method_lru_cache
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, torch_compilable_check, torch_int
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_d_fine import DFineConfig
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for outputs of the DFineDecoder. This class adds two attributes to
+ BaseModelOutputWithCrossAttentions, namely:
+ - a stacked tensor of intermediate decoder hidden states (i.e. the output of each decoder layer)
+ - a stacked tensor of intermediate reference points.
+ """
+)
+@dataclass
+class DFineDecoderOutput(ModelOutput):
+ r"""
+ intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
+ Stacked intermediate hidden states (output of each layer of the decoder).
+ intermediate_logits (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, config.num_labels)`):
+ Stacked intermediate logits (logits of each layer of the decoder).
+ intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, hidden_size)`):
+ Stacked intermediate reference points (reference points of each layer of the decoder).
+ intermediate_predicted_corners (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+ Stacked intermediate predicted corners (predicted corners of each layer of the decoder).
+ initial_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+ Stacked initial reference points (initial reference points of each layer of the decoder).
+ cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`):
+ Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+ sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
+ used to compute the weighted average in the cross-attention heads.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ intermediate_hidden_states: torch.FloatTensor | None = None
+ intermediate_logits: torch.FloatTensor | None = None
+ intermediate_reference_points: torch.FloatTensor | None = None
+ intermediate_predicted_corners: torch.FloatTensor | None = None
+ initial_reference_points: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ cross_attentions: tuple[torch.FloatTensor] | None = None
+
+
+class DFineMLP(nn.Module):
+ def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int, act: str = "relu"):
+ super().__init__()
+ self.num_layers = num_layers
+ hidden_dims = [hidden_dim] * (num_layers - 1)
+ input_dims = [input_dim] + hidden_dims
+ output_dims = hidden_dims + [output_dim]
+ self.layers = nn.ModuleList(nn.Linear(in_dim, out_dim) for in_dim, out_dim in zip(input_dims, output_dims))
+ self.act = ACT2CLS[act]()
+
+ def forward(self, stat_features: torch.Tensor) -> torch.Tensor:
+ for i, layer in enumerate(self.layers):
+ stat_features = self.act(layer(stat_features)) if i < self.num_layers - 1 else layer(stat_features)
+ return stat_features
+
+
+class DFineGate(nn.Module):
+ def __init__(self, d_model: int):
+ super().__init__()
+ self.gate = nn.Linear(2 * d_model, 2 * d_model)
+ self.norm = nn.LayerNorm(d_model)
+
+ def forward(self, second_residual: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
+ gate_input = torch.cat([second_residual, hidden_states], dim=-1)
+ gates = torch.sigmoid(self.gate(gate_input))
+ gate1, gate2 = gates.chunk(2, dim=-1)
+ hidden_states = self.norm(gate1 * second_residual + gate2 * hidden_states)
+ return hidden_states
+
+
+class DFineFrozenBatchNorm2d(nn.Module):
+ """
+ BatchNorm2d where the batch statistics and the affine parameters are fixed.
+
+ Copy-paste from torchvision.misc.ops with added eps before rqsrt, without which any other models than
+ torchvision.models.resnet[18,34,50,101] produce nans.
+ """
+
+ def __init__(self, n):
+ super().__init__()
+ self.register_buffer("weight", torch.ones(n))
+ self.register_buffer("bias", torch.zeros(n))
+ self.register_buffer("running_mean", torch.zeros(n))
+ self.register_buffer("running_var", torch.ones(n))
+
+ def _load_from_state_dict(
+ self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+ ):
+ num_batches_tracked_key = prefix + "num_batches_tracked"
+ if num_batches_tracked_key in state_dict:
+ del state_dict[num_batches_tracked_key]
+
+ super()._load_from_state_dict(
+ state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+ )
+
+ def forward(self, x):
+ # move reshapes to the beginning
+ # to make it user-friendly
+ weight = self.weight.reshape(1, -1, 1, 1)
+ bias = self.bias.reshape(1, -1, 1, 1)
+ running_var = self.running_var.reshape(1, -1, 1, 1)
+ running_mean = self.running_mean.reshape(1, -1, 1, 1)
+ epsilon = 1e-5
+ scale = weight * (running_var + epsilon).rsqrt()
+ bias = bias - running_mean * scale
+ return x * scale + bias
+
+
+def multi_scale_deformable_attention_v2(
+ value: Tensor,
+ value_spatial_shapes: Tensor,
+ sampling_locations: Tensor,
+ attention_weights: Tensor,
+ num_points_list: list[int],
+ method="default",
+) -> Tensor:
+ batch_size, _, num_heads, hidden_dim = value.shape
+ _, num_queries, num_heads, num_levels, num_points = sampling_locations.shape
+ value_list = (
+ value.permute(0, 2, 3, 1)
+ .flatten(0, 1)
+ .split([height * width for height, width in value_spatial_shapes], dim=-1)
+ )
+ # sampling_offsets [8, 480, 8, 12, 2]
+ if method == "default":
+ sampling_grids = 2 * sampling_locations - 1
+ elif method == "discrete":
+ sampling_grids = sampling_locations
+ sampling_grids = sampling_grids.permute(0, 2, 1, 3, 4).flatten(0, 1)
+ sampling_grids = sampling_grids.split(num_points_list, dim=-2)
+ sampling_value_list = []
+ for level_id, (height, width) in enumerate(value_spatial_shapes):
+ # batch_size, height*width, num_heads, hidden_dim
+ # -> batch_size, height*width, num_heads*hidden_dim
+ # -> batch_size, num_heads*hidden_dim, height*width
+ # -> batch_size*num_heads, hidden_dim, height, width
+ value_l_ = value_list[level_id].reshape(batch_size * num_heads, hidden_dim, height, width)
+ # batch_size, num_queries, num_heads, num_points, 2
+ # -> batch_size, num_heads, num_queries, num_points, 2
+ # -> batch_size*num_heads, num_queries, num_points, 2
+ sampling_grid_l_ = sampling_grids[level_id]
+ # batch_size*num_heads, hidden_dim, num_queries, num_points
+ if method == "default":
+ sampling_value_l_ = nn.functional.grid_sample(
+ value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False
+ )
+ elif method == "discrete":
+ sampling_coord = (sampling_grid_l_ * torch.tensor([[width, height]], device=value.device) + 0.5).to(
+ torch.int64
+ )
+
+ # Separate clamping for x and y coordinates
+ sampling_coord_x = sampling_coord[..., 0].clamp(0, width - 1)
+ sampling_coord_y = sampling_coord[..., 1].clamp(0, height - 1)
+
+ # Combine the clamped coordinates
+ sampling_coord = torch.stack([sampling_coord_x, sampling_coord_y], dim=-1)
+ sampling_coord = sampling_coord.reshape(batch_size * num_heads, num_queries * num_points_list[level_id], 2)
+ sampling_idx = (
+ torch.arange(sampling_coord.shape[0], device=value.device)
+ .unsqueeze(-1)
+ .repeat(1, sampling_coord.shape[1])
+ )
+ sampling_value_l_ = value_l_[sampling_idx, :, sampling_coord[..., 1], sampling_coord[..., 0]]
+ sampling_value_l_ = sampling_value_l_.transpose(1, 2).reshape(
+ batch_size * num_heads, hidden_dim, num_queries, num_points_list[level_id]
+ )
+ sampling_value_list.append(sampling_value_l_)
+ # (batch_size, num_queries, num_heads, num_levels, num_points)
+ # -> (batch_size, num_heads, num_queries, num_levels, num_points)
+ # -> (batch_size, num_heads, 1, num_queries, num_levels*num_points)
+ attention_weights = attention_weights.permute(0, 2, 1, 3).reshape(
+ batch_size * num_heads, 1, num_queries, sum(num_points_list)
+ )
+ output = (
+ (torch.concat(sampling_value_list, dim=-1) * attention_weights)
+ .sum(-1)
+ .view(batch_size, num_heads * hidden_dim, num_queries)
+ )
+ return output.transpose(1, 2).contiguous()
+
+
+class DFineMultiscaleDeformableAttention(nn.Module):
+ def __init__(self, config: DFineConfig):
+ """
+ D-Fine version of multiscale deformable attention
+ """
+ super().__init__()
+ self.d_model = config.d_model
+ self.n_heads = config.decoder_attention_heads
+ self.n_levels = config.num_feature_levels
+ self.offset_scale = config.decoder_offset_scale
+ self.decoder_method = config.decoder_method
+ self.n_points = config.decoder_n_points
+
+ if isinstance(self.n_points, list):
+ num_points_list = self.n_points
+ else:
+ num_points_list = [self.n_points for _ in range(self.n_levels)]
+
+ self.num_points_list = num_points_list
+ num_points_scale = [1 / n for n in self.num_points_list for _ in range(n)]
+ self.register_buffer("num_points_scale", torch.tensor(num_points_scale, dtype=torch.float32))
+
+ self.total_points = self.n_heads * sum(self.num_points_list)
+
+ self.sampling_offsets = nn.Linear(self.d_model, self.total_points * 2)
+ self.attention_weights = nn.Linear(self.d_model, self.total_points)
+
+ self.ms_deformable_attn_core = multi_scale_deformable_attention_v2
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ reference_points=None,
+ encoder_hidden_states=None,
+ spatial_shapes=None,
+ spatial_shapes_list=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ batch_size, num_queries, _ = hidden_states.shape
+ batch_size, sequence_length, _ = encoder_hidden_states.shape
+
+ torch_compilable_check(
+ (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == sequence_length,
+ "Make sure to align the spatial shapes with the sequence length of the encoder hidden states",
+ )
+
+ # Reshape for multi-head attention
+ value = encoder_hidden_states.reshape(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads)
+ if attention_mask is not None:
+ value = value.masked_fill(~attention_mask[..., None], float(0))
+
+ sampling_offsets: torch.Tensor = self.sampling_offsets(hidden_states)
+ sampling_offsets = sampling_offsets.reshape(
+ batch_size, num_queries, self.n_heads, sum(self.num_points_list), 2
+ )
+
+ attention_weights = self.attention_weights(hidden_states).reshape(
+ batch_size, num_queries, self.n_heads, sum(self.num_points_list)
+ )
+ attention_weights = F.softmax(attention_weights, dim=-1)
+
+ if reference_points.shape[-1] == 2:
+ offset_normalizer = torch.tensor(spatial_shapes)
+ offset_normalizer = offset_normalizer.flip([1]).reshape(1, 1, 1, self.n_levels, 1, 2)
+ sampling_locations = (
+ reference_points.reshape(batch_size, sequence_length, 1, self.n_levels, 1, 2)
+ + sampling_offsets / offset_normalizer
+ )
+ elif reference_points.shape[-1] == 4:
+ # reference_points [8, 480, None, 1, 4]
+ # sampling_offsets [8, 480, 8, 12, 2]
+ num_points_scale = self.num_points_scale.to(dtype=hidden_states.dtype).unsqueeze(-1)
+ offset = sampling_offsets * num_points_scale * reference_points[:, :, None, :, 2:] * self.offset_scale
+ sampling_locations = reference_points[:, :, None, :, :2] + offset
+ else:
+ raise ValueError(
+ f"Last dim of reference_points must be 2 or 4, but get {reference_points.shape[-1]} instead."
+ )
+
+ output = self.ms_deformable_attn_core(
+ value,
+ spatial_shapes_list,
+ sampling_locations,
+ attention_weights,
+ self.num_points_list,
+ self.decoder_method,
+ )
+
+ return output, attention_weights
+
+
+class DFineConvNormLayer(nn.Module):
+ def __init__(
+ self,
+ config: DFineConfig,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int,
+ groups: int = 1,
+ padding: int | None = None,
+ activation: str | None = None,
+ ):
+ super().__init__()
+ self.conv = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride,
+ groups=groups,
+ padding=(kernel_size - 1) // 2 if padding is None else padding,
+ bias=False,
+ )
+ self.norm = nn.BatchNorm2d(out_channels, config.batch_norm_eps)
+ self.activation = nn.Identity() if activation is None else ACT2CLS[activation]()
+
+ def forward(self, hidden_state):
+ hidden_state = self.conv(hidden_state)
+ hidden_state = self.norm(hidden_state)
+ hidden_state = self.activation(hidden_state)
+ return hidden_state
+
+
+class DFineRepVggBlock(nn.Module):
+ """
+ RepVGG architecture block introduced by the work "RepVGG: Making VGG-style ConvNets Great Again".
+ """
+
+ def __init__(self, config: DFineConfig, in_channels: int, out_channels: int):
+ super().__init__()
+
+ activation = config.activation_function
+ hidden_channels = in_channels
+ self.conv1 = DFineConvNormLayer(config, hidden_channels, out_channels, 3, 1, padding=1)
+ self.conv2 = DFineConvNormLayer(config, hidden_channels, out_channels, 1, 1, padding=0)
+ self.activation = nn.Identity() if activation is None else ACT2CLS[activation]()
+
+ def forward(self, x):
+ y = self.conv1(x) + self.conv2(x)
+ return self.activation(y)
+
+
+class DFineCSPRepLayer(nn.Module):
+ """
+ Cross Stage Partial (CSP) network layer with RepVGG blocks.
+ """
+
+ def __init__(
+ self, config: DFineConfig, in_channels: int, out_channels: int, num_blocks: int, expansion: float = 1.0
+ ):
+ super().__init__()
+ activation = config.activation_function
+
+ hidden_channels = int(out_channels * expansion)
+ self.conv1 = DFineConvNormLayer(config, in_channels, hidden_channels, 1, 1, activation=activation)
+ self.conv2 = DFineConvNormLayer(config, in_channels, hidden_channels, 1, 1, activation=activation)
+ self.bottlenecks = nn.ModuleList(
+ [DFineRepVggBlock(config, hidden_channels, hidden_channels) for _ in range(num_blocks)]
+ )
+ if hidden_channels != out_channels:
+ self.conv3 = DFineConvNormLayer(config, hidden_channels, out_channels, 1, 1, activation=activation)
+ else:
+ self.conv3 = nn.Identity()
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state_1 = self.conv1(hidden_state)
+ for bottleneck in self.bottlenecks:
+ hidden_state_1 = bottleneck(hidden_state_1)
+ hidden_state_2 = self.conv2(hidden_state)
+ hidden_state_3 = self.conv3(hidden_state_1 + hidden_state_2)
+ return hidden_state_3
+
+
+class DFineRepNCSPELAN4(nn.Module):
+ def __init__(self, config: DFineConfig, act: str = "silu", numb_blocks: int = 3):
+ super().__init__()
+ conv1_dim = config.encoder_hidden_dim * 2
+ conv2_dim = config.encoder_hidden_dim
+ conv3_dim = config.encoder_hidden_dim * 2
+ conv4_dim = round(config.hidden_expansion * config.encoder_hidden_dim // 2)
+ self.conv_dim = conv3_dim // 2
+ self.conv1 = DFineConvNormLayer(config, conv1_dim, conv3_dim, 1, 1, activation=act)
+ self.csp_rep1 = DFineCSPRepLayer(config, conv3_dim // 2, conv4_dim, num_blocks=numb_blocks)
+ self.conv2 = DFineConvNormLayer(config, conv4_dim, conv4_dim, 3, 1, activation=act)
+ self.csp_rep2 = DFineCSPRepLayer(config, conv4_dim, conv4_dim, num_blocks=numb_blocks)
+ self.conv3 = DFineConvNormLayer(config, conv4_dim, conv4_dim, 3, 1, activation=act)
+ self.conv4 = DFineConvNormLayer(config, conv3_dim + (2 * conv4_dim), conv2_dim, 1, 1, activation=act)
+
+ def forward(self, input_features: torch.Tensor) -> torch.Tensor:
+ # Split initial features into two branches after first convolution
+ split_features = list(self.conv1(input_features).split((self.conv_dim, self.conv_dim), 1))
+
+ # Process branches sequentially
+ branch1 = self.csp_rep1(split_features[-1])
+ branch1 = self.conv2(branch1)
+ branch2 = self.csp_rep2(branch1)
+ branch2 = self.conv3(branch2)
+
+ split_features.extend([branch1, branch2])
+ merged_features = torch.cat(split_features, 1)
+ merged_features = self.conv4(merged_features)
+ return merged_features
+
+
+class DFineSCDown(nn.Module):
+ def __init__(self, config: DFineConfig, kernel_size: int, stride: int):
+ super().__init__()
+ self.conv1 = DFineConvNormLayer(config, config.encoder_hidden_dim, config.encoder_hidden_dim, 1, 1)
+ self.conv2 = DFineConvNormLayer(
+ config,
+ config.encoder_hidden_dim,
+ config.encoder_hidden_dim,
+ kernel_size,
+ stride,
+ config.encoder_hidden_dim,
+ )
+
+ def forward(self, input_features: torch.Tensor) -> torch.Tensor:
+ input_features = self.conv1(input_features)
+ input_features = self.conv2(input_features)
+ return input_features
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class DFineSelfAttention(nn.Module):
+ """
+ Multi-headed self-attention from 'Attention Is All You Need' paper.
+
+ In D_FINE, position embeddings are added to both queries and keys (but not values) in self-attention.
+ """
+
+ def __init__(
+ self,
+ config: DFineConfig,
+ hidden_size: int,
+ num_attention_heads: int,
+ dropout: float = 0.0,
+ bias: bool = True,
+ ):
+ super().__init__()
+ self.config = config
+ self.head_dim = hidden_size // num_attention_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = dropout
+ self.is_causal = False
+
+ self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+ self.v_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+ self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+ self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_embeddings: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Position embeddings are added to both queries and keys (but not values).
+ """
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_key_input = hidden_states + position_embeddings if position_embeddings is not None else hidden_states
+
+ query_states = self.q_proj(query_key_input).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(query_key_input).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class DFineEncoderLayer(nn.Module):
+ def __init__(self, config: DFineConfig):
+ super().__init__()
+ self.normalize_before = config.normalize_before
+ self.hidden_size = config.encoder_hidden_dim
+
+ # self-attention
+ self.self_attn = DFineSelfAttention(
+ config=config,
+ hidden_size=self.hidden_size,
+ num_attention_heads=config.num_attention_heads,
+ dropout=config.dropout,
+ )
+ self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = config.dropout
+ self.mlp = DFineMLP(
+ self.hidden_size, config.encoder_ffn_dim, self.hidden_size, 2, config.encoder_activation_function
+ )
+ self.final_layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ spatial_position_embeddings: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, hidden_size)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
+ values.
+ spatial_position_embeddings (`torch.FloatTensor`, *optional*):
+ Spatial position embeddings (2D positional encodings of image locations), to be added to both
+ the queries and keys in self-attention (but not to values).
+ """
+ residual = hidden_states
+ if self.normalize_before:
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_embeddings=spatial_position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ if not self.normalize_before:
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ if self.normalize_before:
+ hidden_states = self.final_layer_norm(hidden_states)
+ residual = hidden_states
+
+ hidden_states = self.mlp(hidden_states)
+
+ hidden_states = residual + hidden_states
+ if not self.normalize_before:
+ hidden_states = self.final_layer_norm(hidden_states)
+
+ if self.training:
+ if not torch.isfinite(hidden_states).all():
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+ hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+ return hidden_states
+
+
+def build_2d_sinusoidal_position_embedding(
+ height: int,
+ width: int,
+ embed_dim: int = 256,
+ temperature: float = 10000.0,
+ cls_token: bool = False,
+ device: torch.device | None = None,
+ dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+ """2D sinusoidal position embeddings for an image patch grid.
+
+ Each (h, w) position gets an ``embed_dim``-dimensional vector laid out as
+ ``[sin_h | cos_h | sin_w | cos_w]``, with row-major (H-outer) patch ordering.
+
+ Args:
+ height: Grid height in patches.
+ width: Grid width in patches.
+ embed_dim: Total embedding dimension; must be divisible by 4.
+ temperature: Base for the frequency decay.
+ cls_token: If `True`, prepend a zero row for a CLS token.
+ device: Target device; defaults to CPU.
+ dtype: Output dtype; frequency arithmetic uses float64 internally.
+
+ Returns:
+ Tensor of shape ``(height * width [+1], embed_dim)``.
+ """
+ if embed_dim % 4 != 0:
+ raise ValueError(f"`embed_dim` must be divisible by 4, got {embed_dim}")
+
+ pos_dim = embed_dim // 4
+ omega = torch.arange(pos_dim, dtype=torch.float64, device=device) / pos_dim
+ omega = 1.0 / temperature**omega # (D/4,)
+
+ grid_h = torch.arange(height, dtype=torch.float64, device=device)
+ grid_w = torch.arange(width, dtype=torch.float64, device=device)
+ grid_h, grid_w = torch.meshgrid(grid_h, grid_w, indexing="ij") # (H, W) each
+
+ emb_h = grid_h.flatten().outer(omega) # (H*W, D/4)
+ emb_w = grid_w.flatten().outer(omega) # (H*W, D/4)
+
+ pos_embed = torch.cat([emb_h.sin(), emb_h.cos(), emb_w.sin(), emb_w.cos()], dim=1)
+
+ if cls_token:
+ pos_embed = torch.cat([torch.zeros(1, embed_dim, dtype=torch.float64, device=device), pos_embed], dim=0)
+
+ return pos_embed.to(dtype)
+
+
+class DFineSinePositionEmbedding(nn.Module):
+ """
+ 2D sinusoidal position embedding used in RT-DETR hybrid encoder.
+ """
+
+ def __init__(self, embed_dim: int = 256, temperature: int = 10000):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.temperature = temperature
+
+ @staticmethod
+ @compile_compatible_method_lru_cache(maxsize=32)
+ def _cached_build_2d_sinusoidal_position_embedding(*args, **kwargs) -> torch.Tensor:
+ return build_2d_sinusoidal_position_embedding(*args, **kwargs)
+
+ def forward(
+ self,
+ width: int,
+ height: int,
+ device: torch.device | str,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ """
+ Generate 2D sinusoidal position embeddings.
+
+ Returns:
+ Position embeddings of shape (1, height*width, embed_dim)
+ """
+ return self._cached_build_2d_sinusoidal_position_embedding(
+ height=torch_int(height),
+ width=torch_int(width),
+ embed_dim=self.embed_dim,
+ temperature=self.temperature,
+ device=device,
+ dtype=dtype,
+ ).unsqueeze(0)
+
+
+class DFineAIFILayer(nn.Module):
+ """
+ AIFI (Attention-based Intra-scale Feature Interaction) layer used in RT-DETR hybrid encoder.
+ """
+
+ def __init__(self, config: DFineConfig):
+ super().__init__()
+ self.config = config
+ self.encoder_hidden_dim = config.encoder_hidden_dim
+ self.eval_size = config.eval_size
+
+ self.position_embedding = DFineSinePositionEmbedding(
+ embed_dim=self.encoder_hidden_dim,
+ temperature=config.positional_encoding_temperature,
+ )
+ self.layers = nn.ModuleList([DFineEncoderLayer(config) for _ in range(config.encoder_layers)])
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor` of shape `(batch_size, channels, height, width)`):
+ Feature map to process.
+ """
+ batch_size = hidden_states.shape[0]
+ height, width = hidden_states.shape[2:]
+
+ hidden_states = hidden_states.flatten(2).transpose(1, 2)
+
+ if self.training or self.eval_size is None:
+ pos_embed = self.position_embedding(
+ width=width,
+ height=height,
+ device=hidden_states.device,
+ dtype=hidden_states.dtype,
+ )
+ else:
+ pos_embed = None
+
+ for layer in self.layers:
+ hidden_states = layer(
+ hidden_states,
+ attention_mask=None,
+ spatial_position_embeddings=pos_embed,
+ **kwargs,
+ )
+
+ hidden_states = (
+ hidden_states.transpose(1, 2).reshape(batch_size, self.encoder_hidden_dim, height, width).contiguous()
+ )
+
+ return hidden_states
+
+
+class DFineIntegral(nn.Module):
+ """
+ A static layer that calculates integral results from a distribution.
+
+ This layer computes the target location using the formula: `sum{Pr(n) * W(n)}`,
+ where Pr(n) is the softmax probability vector representing the discrete
+ distribution, and W(n) is the non-uniform Weighting Function.
+
+ Args:
+ max_num_bins (int): Max number of the discrete bins. Default is 32.
+ It can be adjusted based on the dataset or task requirements.
+ """
+
+ def __init__(self, config: DFineConfig):
+ super().__init__()
+ self.max_num_bins = config.max_num_bins
+
+ def forward(self, pred_corners: torch.Tensor, project: torch.Tensor) -> torch.Tensor:
+ batch_size, num_queries, _ = pred_corners.shape
+ pred_corners = F.softmax(pred_corners.reshape(-1, self.max_num_bins + 1), dim=1)
+ pred_corners = F.linear(pred_corners, project.to(pred_corners.device)).reshape(-1, 4)
+ pred_corners = pred_corners.reshape(batch_size, num_queries, -1)
+ return pred_corners
+
+
+class DFineLQE(nn.Module):
+ def __init__(self, config: DFineConfig):
+ super().__init__()
+ self.top_prob_values = config.top_prob_values
+ self.max_num_bins = config.max_num_bins
+ self.reg_conf = DFineMLP(4 * (self.top_prob_values + 1), config.lqe_hidden_dim, 1, config.lqe_layers)
+
+ def forward(self, scores: torch.Tensor, pred_corners: torch.Tensor) -> torch.Tensor:
+ batch_size, length, _ = pred_corners.size()
+ prob = F.softmax(pred_corners.reshape(batch_size, length, 4, self.max_num_bins + 1), dim=-1)
+ prob_topk, _ = prob.topk(self.top_prob_values, dim=-1)
+ stat = torch.cat([prob_topk, prob_topk.mean(dim=-1, keepdim=True)], dim=-1)
+ quality_score = self.reg_conf(stat.reshape(batch_size, length, -1))
+ scores = scores + quality_score
+ return scores
+
+
+class DFineDecoderLayer(nn.Module):
+ def __init__(self, config: DFineConfig):
+ super().__init__()
+ self.hidden_size = config.d_model
+
+ # self-attention
+ self.self_attn = DFineSelfAttention(
+ config=config,
+ hidden_size=self.hidden_size,
+ num_attention_heads=config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ )
+ self.dropout = config.dropout
+
+ self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps)
+
+ # override the encoder attention module with d-fine version
+ self.encoder_attn = DFineMultiscaleDeformableAttention(config=config)
+ self.mlp = DFineMLP(
+ self.hidden_size, config.decoder_ffn_dim, self.hidden_size, 2, config.decoder_activation_function
+ )
+ self.final_layer_norm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps)
+ # gate
+ self.gateway = DFineGate(config.d_model)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: torch.Tensor | None = None,
+ reference_points=None,
+ spatial_shapes=None,
+ spatial_shapes_list=None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`):
+ Input to the layer of shape `(batch, seq_len, hidden_size)`.
+ object_queries_position_embeddings (`torch.FloatTensor`, *optional*):
+ Position embeddings for the object query slots. These are added to both queries and keys
+ in the self-attention layer (not values).
+ reference_points (`torch.FloatTensor`, *optional*):
+ Reference points.
+ spatial_shapes (`torch.LongTensor`, *optional*):
+ Spatial shapes.
+ level_start_index (`torch.LongTensor`, *optional*):
+ Level start index.
+ encoder_hidden_states (`torch.FloatTensor`):
+ cross attention input to the layer of shape `(batch, seq_len, hidden_size)`
+ encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
+ `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative
+ values.
+ """
+ residual = hidden_states
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=encoder_attention_mask,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ residual = hidden_states
+
+ # Cross-Attention
+ hidden_states = hidden_states if position_embeddings is None else hidden_states + position_embeddings
+ hidden_states, _ = self.encoder_attn(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ reference_points=reference_points,
+ spatial_shapes=spatial_shapes,
+ spatial_shapes_list=spatial_shapes_list,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = self.gateway(residual, hidden_states)
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ hidden_states = self.final_layer_norm(hidden_states.clamp(min=-65504, max=65504))
+
+ return hidden_states
+
+
+class DFineMLPPredictionHead(nn.Module):
+ """
+ Very simple multi-layer perceptron (MLP, also called FFN), used to predict the normalized center coordinates,
+ height and width of a bounding box w.r.t. an image.
+
+ """
+
+ def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+ super().__init__()
+ self.num_layers = num_layers
+ h = [hidden_dim] * (num_layers - 1)
+ self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+ def forward(self, x):
+ for i, layer in enumerate(self.layers):
+ x = nn.functional.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+ return x
+
+
+@auto_docstring
+class DFinePreTrainedModel(PreTrainedModel):
+ config: DFineConfig
+ base_model_prefix = "d_fine"
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ _no_split_modules = [r"DFineHybridEncoder", r"DFineDecoderLayer"]
+ _supports_sdpa = True
+ _supports_flash_attn = True
+ _supports_attention_backend = True
+ _supports_flex_attn = True
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ super()._init_weights(module)
+ # initialize linear layer bias value according to a given probability value.
+ if isinstance(module, (DFineForObjectDetection, DFineDecoder)):
+ if module.class_embed is not None:
+ for layer in module.class_embed:
+ prior_prob = self.config.initializer_bias_prior_prob or 1 / (self.config.num_labels + 1)
+ bias = float(-math.log((1 - prior_prob) / prior_prob))
+ init.xavier_uniform_(layer.weight)
+ init.constant_(layer.bias, bias)
+
+ if module.bbox_embed is not None:
+ for layer in module.bbox_embed:
+ init.constant_(layer.layers[-1].weight, 0)
+ init.constant_(layer.layers[-1].bias, 0)
+
+ if hasattr(module, "reg_scale"):
+ init.constant_(module.reg_scale, self.config.reg_scale)
+
+ if hasattr(module, "up"):
+ init.constant_(module.up, self.config.up)
+
+ if isinstance(module, DFineMultiscaleDeformableAttention):
+ init.constant_(module.sampling_offsets.weight, 0.0)
+ default_dtype = torch.get_default_dtype()
+ thetas = torch.arange(module.n_heads, dtype=torch.int64).to(default_dtype) * (
+ 2.0 * math.pi / module.n_heads
+ )
+ grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+ grid_init = grid_init / grid_init.abs().max(-1, keepdim=True).values
+ grid_init = grid_init.reshape(module.n_heads, 1, 2).tile([1, sum(module.num_points_list), 1])
+ scaling = torch.concat([torch.arange(1, n + 1) for n in module.num_points_list]).reshape(1, -1, 1)
+ grid_init *= scaling
+ init.copy_(module.sampling_offsets.bias, grid_init.flatten())
+
+ init.constant_(module.attention_weights.weight, 0.0)
+ init.constant_(module.attention_weights.bias, 0.0)
+
+ num_points_scale = [1 / n for n in module.num_points_list for _ in range(n)]
+ init.copy_(module.num_points_scale, torch.tensor(num_points_scale, dtype=torch.float32))
+
+ if isinstance(module, DFineModel):
+ prior_prob = self.config.initializer_bias_prior_prob or 1 / (self.config.num_labels + 1)
+ bias = float(-math.log((1 - prior_prob) / prior_prob))
+ init.xavier_uniform_(module.enc_score_head.weight)
+ init.constant_(module.enc_score_head.bias, bias)
+
+ if isinstance(module, DFineGate):
+ bias = float(-math.log((1 - 0.5) / 0.5))
+ init.constant_(module.gate.bias, bias)
+ init.constant_(module.gate.weight, 0)
+
+ if isinstance(module, DFineLQE):
+ init.constant_(module.reg_conf.layers[-1].bias, 0)
+ init.constant_(module.reg_conf.layers[-1].weight, 0)
+
+ if hasattr(module, "weight_embedding") and self.config.learn_initial_query:
+ init.xavier_uniform_(module.weight_embedding.weight)
+ if hasattr(module, "denoising_class_embed") and self.config.num_denoising > 0:
+ init.xavier_uniform_(module.denoising_class_embed.weight)
+
+
+class DFineHybridEncoder(DFinePreTrainedModel):
+ """
+ Hybrid encoder consisting of AIFI (Attention-based Intra-scale Feature Interaction) layers,
+ a top-down Feature Pyramid Network (FPN) and a bottom-up Path Aggregation Network (PAN).
+ More details on the paper: https://huggingface.co/papers/2304.08069
+
+ Args:
+ config: DFineConfig
+ """
+
+ _can_record_outputs = {
+ "hidden_states": DFineAIFILayer,
+ "attentions": DFineSelfAttention,
+ }
+
+ def __init__(self, config: DFineConfig):
+ super().__init__(config)
+ self.config = config
+ self.in_channels = config.encoder_in_channels
+ self.num_fpn_stages = len(self.in_channels) - 1
+ self.feat_strides = config.feat_strides
+ self.encoder_hidden_dim = config.encoder_hidden_dim
+ self.encode_proj_layers = config.encode_proj_layers
+ self.positional_encoding_temperature = config.positional_encoding_temperature
+ self.eval_size = config.eval_size
+ self.out_channels = [self.encoder_hidden_dim for _ in self.in_channels]
+ self.out_strides = self.feat_strides
+
+ # AIFI (Attention-based Intra-scale Feature Interaction) layers
+ self.aifi = nn.ModuleList([DFineAIFILayer(config) for _ in range(len(self.encode_proj_layers))])
+
+ # top-down fpn
+ self.lateral_convs = nn.ModuleList()
+ self.fpn_blocks = nn.ModuleList()
+ for _ in range(len(self.in_channels) - 1, 0, -1):
+ lateral_layer = DFineConvNormLayer(config, self.encoder_hidden_dim, self.encoder_hidden_dim, 1, 1)
+ self.lateral_convs.append(lateral_layer)
+ num_blocks = round(3 * config.depth_mult)
+ fpn_layer = DFineRepNCSPELAN4(config, numb_blocks=num_blocks)
+ self.fpn_blocks.append(fpn_layer)
+
+ # bottom-up pan
+ self.downsample_convs = nn.ModuleList()
+ self.pan_blocks = nn.ModuleList()
+ for _ in range(len(self.in_channels) - 1):
+ self.downsample_convs.append(DFineSCDown(config, 3, 2))
+ num_blocks = round(3 * config.depth_mult)
+ self.pan_blocks.append(DFineRepNCSPELAN4(config, numb_blocks=num_blocks))
+
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ def forward(
+ self,
+ inputs_embeds=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutput:
+ r"""
+ Args:
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ Flattened feature map (output of the backbone + projection layer) that is passed to the encoder.
+ """
+ feature_maps = inputs_embeds
+
+ # AIFI: Apply transformer encoder to specified feature levels
+ if self.config.encoder_layers > 0:
+ for i, enc_ind in enumerate(self.encode_proj_layers):
+ feature_maps[enc_ind] = self.aifi[i](feature_maps[enc_ind], **kwargs)
+
+ # top-down FPN
+ fpn_feature_maps = [feature_maps[-1]]
+ for idx, (lateral_conv, fpn_block) in enumerate(zip(self.lateral_convs, self.fpn_blocks)):
+ backbone_feature_map = feature_maps[self.num_fpn_stages - idx - 1]
+ top_fpn_feature_map = fpn_feature_maps[-1]
+ # apply lateral block
+ top_fpn_feature_map = lateral_conv(top_fpn_feature_map)
+ fpn_feature_maps[-1] = top_fpn_feature_map
+ # apply fpn block
+ top_fpn_feature_map = F.interpolate(top_fpn_feature_map, scale_factor=2.0, mode="nearest")
+ fused_feature_map = torch.concat([top_fpn_feature_map, backbone_feature_map], dim=1)
+ new_fpn_feature_map = fpn_block(fused_feature_map)
+ fpn_feature_maps.append(new_fpn_feature_map)
+
+ fpn_feature_maps.reverse()
+
+ # bottom-up PAN
+ pan_feature_maps = [fpn_feature_maps[0]]
+ for idx, (downsample_conv, pan_block) in enumerate(zip(self.downsample_convs, self.pan_blocks)):
+ top_pan_feature_map = pan_feature_maps[-1]
+ fpn_feature_map = fpn_feature_maps[idx + 1]
+ downsampled_feature_map = downsample_conv(top_pan_feature_map)
+ fused_feature_map = torch.concat([downsampled_feature_map, fpn_feature_map], dim=1)
+ new_pan_feature_map = pan_block(fused_feature_map)
+ pan_feature_maps.append(new_pan_feature_map)
+
+ return BaseModelOutput(last_hidden_state=pan_feature_maps)
+
+
+def inverse_sigmoid(x, eps=1e-5):
+ x = x.clamp(min=0, max=1)
+ x1 = x.clamp(min=eps)
+ x2 = (1 - x).clamp(min=eps)
+ return torch.log(x1 / x2)
+
+
+def weighting_function(max_num_bins: int, up: torch.Tensor, reg_scale: int) -> torch.Tensor:
+ """
+ Generates the non-uniform Weighting Function W(n) for bounding box regression.
+
+ Args:
+ max_num_bins (int): Max number of the discrete bins.
+ up (Tensor): Controls upper bounds of the sequence,
+ where maximum offset is ±up * H / W.
+ reg_scale (float): Controls the curvature of the Weighting Function.
+ Larger values result in flatter weights near the central axis W(max_num_bins/2)=0
+ and steeper weights at both ends.
+ Returns:
+ Tensor: Sequence of Weighting Function.
+ """
+ upper_bound1 = abs(up[0]) * abs(reg_scale)
+ upper_bound2 = abs(up[0]) * abs(reg_scale) * 2
+ step = (upper_bound1 + 1) ** (2 / (max_num_bins - 2))
+ left_values = [-((step) ** i) + 1 for i in range(max_num_bins // 2 - 1, 0, -1)]
+ right_values = [(step) ** i - 1 for i in range(1, max_num_bins // 2)]
+ values = [-upper_bound2] + left_values + [torch.zeros_like(up[0][None])] + right_values + [upper_bound2]
+ values = torch.cat(values, 0)
+ return values
+
+
+def distance2bbox(points, distance: torch.Tensor, reg_scale: float) -> torch.Tensor:
+ """
+ Decodes edge-distances into bounding box coordinates.
+
+ Args:
+ points (`torch.Tensor`):
+ (batch_size, num_boxes, 4) or (num_boxes, 4) format, representing [x_center, y_center, width, height]
+ distance (`torch.Tensor`):
+ (batch_size, num_boxes, 4) or (num_boxes, 4), representing distances from the point to the left, top, right, and bottom boundaries.
+ reg_scale (`float`):
+ Controls the curvature of the Weighting Function.
+ Returns:
+ `torch.Tensor`: Bounding boxes in (batch_size, num_boxes, 4) or (num_boxes, 4) format, representing [x_center, y_center, width, height]
+ """
+ reg_scale = abs(reg_scale)
+ top_left_x = points[..., 0] - (0.5 * reg_scale + distance[..., 0]) * (points[..., 2] / reg_scale)
+ top_left_y = points[..., 1] - (0.5 * reg_scale + distance[..., 1]) * (points[..., 3] / reg_scale)
+ bottom_right_x = points[..., 0] + (0.5 * reg_scale + distance[..., 2]) * (points[..., 2] / reg_scale)
+ bottom_right_y = points[..., 1] + (0.5 * reg_scale + distance[..., 3]) * (points[..., 3] / reg_scale)
+
+ bboxes = torch.stack([top_left_x, top_left_y, bottom_right_x, bottom_right_y], -1)
+
+ return corners_to_center_format(bboxes)
+
+
+class DFineDecoder(DFinePreTrainedModel):
+ """
+ D-FINE Decoder implementing Fine-grained Distribution Refinement (FDR).
+
+ This decoder refines object detection predictions through iterative updates across multiple layers,
+ utilizing attention mechanisms, location quality estimators, and distribution refinement techniques
+ to improve bounding box accuracy and robustness.
+ """
+
+ _can_record_outputs = {
+ "hidden_states": DFineDecoderLayer,
+ "attentions": DFineSelfAttention,
+ "cross_attentions": DFineMultiscaleDeformableAttention,
+ }
+
+ def __init__(self, config: DFineConfig):
+ super().__init__(config)
+ self.eval_idx = config.eval_idx if config.eval_idx >= 0 else config.decoder_layers + config.eval_idx
+
+ self.dropout = config.dropout
+ self.layers = nn.ModuleList(
+ [DFineDecoderLayer(config) for _ in range(config.decoder_layers)]
+ + [DFineDecoderLayer(config) for _ in range(config.decoder_layers - self.eval_idx - 1)]
+ )
+ self.query_pos_head = DFineMLPPredictionHead(4, 2 * config.d_model, config.d_model, num_layers=2)
+
+ # hack implementation for iterative bounding box refinement and two-stage Deformable DETR
+ self.bbox_embed = None
+ self.class_embed = None
+ self.reg_scale = nn.Parameter(torch.tensor([config.reg_scale]), requires_grad=False)
+ self.max_num_bins = config.max_num_bins
+ self.d_model = config.d_model
+ self.layer_scale = config.layer_scale
+ self.pre_bbox_head = DFineMLP(config.hidden_size, config.hidden_size, 4, 3)
+ self.integral = DFineIntegral(config)
+ self.num_head = config.decoder_attention_heads
+ self.up = nn.Parameter(torch.tensor([config.up]), requires_grad=False)
+ self.lqe_layers = nn.ModuleList([DFineLQE(config) for _ in range(config.decoder_layers)])
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self,
+ encoder_hidden_states: torch.Tensor,
+ reference_points: torch.Tensor,
+ inputs_embeds: torch.Tensor,
+ spatial_shapes,
+ level_start_index=None,
+ spatial_shapes_list=None,
+ encoder_attention_mask=None,
+ memory_mask=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> DFineDecoderOutput:
+ r"""
+ Args:
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`):
+ The query embeddings that are passed into the decoder.
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
+ of the decoder.
+ encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing cross-attention on padding pixel_values of the encoder. Mask values selected
+ in `[0, 1]`:
+ - 1 for pixels that are real (i.e. **not masked**),
+ - 0 for pixels that are padding (i.e. **masked**).
+ reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)` is `as_two_stage` else `(batch_size, num_queries, 2)` or , *optional*):
+ Reference point in range `[0, 1]`, top-left (0,0), bottom-right (1, 1), including padding area.
+ spatial_shapes (`torch.FloatTensor` of shape `(num_feature_levels, 2)`):
+ Spatial shapes of the feature maps.
+ level_start_index (`torch.LongTensor` of shape `(num_feature_levels)`, *optional*):
+ Indexes for the start of each feature level. In range `[0, sequence_length]`.
+ """
+ if inputs_embeds is not None:
+ hidden_states = inputs_embeds
+
+ # decoder layers
+ intermediate = ()
+ intermediate_reference_points = ()
+ intermediate_logits = ()
+ intermediate_predicted_corners = ()
+ initial_reference_points = ()
+
+ output_detach = pred_corners_undetach = 0
+
+ project = weighting_function(self.max_num_bins, self.up, self.reg_scale)
+ ref_points_detach = F.sigmoid(reference_points)
+
+ for i, decoder_layer in enumerate(self.layers):
+ ref_points_input = ref_points_detach.unsqueeze(2)
+ query_pos_embed = self.query_pos_head(ref_points_detach).clamp(min=-10, max=10)
+
+ hidden_states = decoder_layer(
+ hidden_states,
+ position_embeddings=query_pos_embed,
+ reference_points=ref_points_input,
+ spatial_shapes=spatial_shapes,
+ spatial_shapes_list=spatial_shapes_list,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ **kwargs,
+ )
+
+ if i == 0:
+ # Initial bounding box predictions with inverse sigmoid refinement
+ new_reference_points = F.sigmoid(
+ self.pre_bbox_head(hidden_states) + inverse_sigmoid(ref_points_detach)
+ )
+ ref_points_initial = new_reference_points.detach()
+
+ # Refine bounding box corners using FDR, integrating previous layer's corrections
+ if self.bbox_embed is not None:
+ pred_corners = self.bbox_embed[i](hidden_states + output_detach) + pred_corners_undetach
+ inter_ref_bbox = distance2bbox(
+ ref_points_initial, self.integral(pred_corners, project), self.reg_scale
+ )
+ pred_corners_undetach = pred_corners
+ ref_points_detach = inter_ref_bbox.detach()
+
+ output_detach = hidden_states.detach()
+
+ intermediate += (hidden_states,)
+
+ if self.class_embed is not None and (self.training or i == self.eval_idx):
+ scores = self.class_embed[i](hidden_states)
+ # Add initial logits and reference points with pre-bbox head
+ if i == 0:
+ intermediate_logits += (scores,)
+ intermediate_reference_points += (new_reference_points,)
+ # Lqe does not affect the performance here.
+ scores = self.lqe_layers[i](scores, pred_corners)
+ intermediate_logits += (scores,)
+ intermediate_reference_points += (inter_ref_bbox,)
+ initial_reference_points += (ref_points_initial,)
+ intermediate_predicted_corners += (pred_corners,)
+
+ # Keep batch_size as first dimension
+ intermediate = torch.stack(intermediate)
+ if self.class_embed is not None and self.bbox_embed is not None:
+ intermediate_logits = torch.stack(intermediate_logits, dim=1)
+ intermediate_predicted_corners = torch.stack(intermediate_predicted_corners, dim=1)
+ initial_reference_points = torch.stack(initial_reference_points, dim=1)
+ intermediate_reference_points = torch.stack(intermediate_reference_points, dim=1)
+
+ return DFineDecoderOutput(
+ last_hidden_state=hidden_states,
+ intermediate_hidden_states=intermediate,
+ intermediate_logits=intermediate_logits,
+ intermediate_reference_points=intermediate_reference_points,
+ intermediate_predicted_corners=intermediate_predicted_corners,
+ initial_reference_points=initial_reference_points,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for outputs of the RT-DETR encoder-decoder model.
+ """
+)
+@dataclass
+class DFineModelOutput(ModelOutput):
+ r"""
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`):
+ Sequence of hidden-states at the output of the last layer of the decoder of the model.
+ intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
+ Stacked intermediate hidden states (output of each layer of the decoder).
+ intermediate_logits (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, config.num_labels)`):
+ Stacked intermediate logits (logits of each layer of the decoder).
+ intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+ Stacked intermediate reference points (reference points of each layer of the decoder).
+ intermediate_predicted_corners (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+ Stacked intermediate predicted corners (predicted corners of each layer of the decoder).
+ initial_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
+ Initial reference points used for the first decoder layer.
+ init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
+ Initial reference points sent through the Transformer decoder.
+ enc_topk_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`):
+ Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
+ picked as region proposals in the encoder stage. Output of bounding box binary classification (i.e.
+ foreground and background).
+ enc_topk_bboxes (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`):
+ Logits of predicted bounding boxes coordinates in the encoder stage.
+ enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+ Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
+ picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
+ foreground and background).
+ enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+ Logits of predicted bounding boxes coordinates in the first stage.
+ denoising_meta_values (`dict`):
+ Extra dictionary for the denoising related values.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ intermediate_hidden_states: torch.FloatTensor | None = None
+ intermediate_logits: torch.FloatTensor | None = None
+ intermediate_reference_points: torch.FloatTensor | None = None
+ intermediate_predicted_corners: torch.FloatTensor | None = None
+ initial_reference_points: torch.FloatTensor | None = None
+ decoder_hidden_states: tuple[torch.FloatTensor] | None = None
+ decoder_attentions: tuple[torch.FloatTensor] | None = None
+ cross_attentions: tuple[torch.FloatTensor] | None = None
+ encoder_last_hidden_state: torch.FloatTensor | None = None
+ encoder_hidden_states: tuple[torch.FloatTensor] | None = None
+ encoder_attentions: tuple[torch.FloatTensor] | None = None
+ init_reference_points: torch.FloatTensor | None = None
+ enc_topk_logits: torch.FloatTensor | None = None
+ enc_topk_bboxes: torch.FloatTensor | None = None
+ enc_outputs_class: torch.FloatTensor | None = None
+ enc_outputs_coord_logits: torch.FloatTensor | None = None
+ denoising_meta_values: dict | None = None
+
+
+def replace_batch_norm(model):
+ r"""
+ Recursively replace all `torch.nn.BatchNorm2d` with `DFineFrozenBatchNorm2d`.
+
+ Args:
+ model (torch.nn.Module):
+ input model
+ """
+ for name, module in model.named_children():
+ if isinstance(module, nn.BatchNorm2d):
+ new_module = DFineFrozenBatchNorm2d(module.num_features)
+
+ if module.weight.device != torch.device("meta"):
+ new_module.weight.copy_(module.weight)
+ new_module.bias.copy_(module.bias)
+ new_module.running_mean.copy_(module.running_mean)
+ new_module.running_var.copy_(module.running_var)
+
+ model._modules[name] = new_module
+
+ if len(list(module.children())) > 0:
+ replace_batch_norm(module)
+
+
+class DFineConvEncoder(nn.Module):
+ """
+ Convolutional backbone using the modeling_d_fine_resnet.py.
+
+ nn.BatchNorm2d layers are replaced by DFineFrozenBatchNorm2d as defined above.
+ https://github.com/lyuwenyu/RT-DETR/blob/main/DFine_pytorch/src/nn/backbone/presnet.py#L142
+ """
+
+ def __init__(self, config):
+ super().__init__()
+
+ backbone = load_backbone(config)
+
+ if config.freeze_backbone_batch_norms:
+ # replace batch norm by frozen batch norm
+ with torch.no_grad():
+ replace_batch_norm(backbone)
+ self.model = backbone
+ self.intermediate_channel_sizes = self.model.channels
+
+ def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor):
+ # send pixel_values through the model to get list of feature maps
+ features = self.model(pixel_values).feature_maps
+
+ out = []
+ for feature_map in features:
+ # downsample pixel_mask to match shape of corresponding feature_map
+ mask = nn.functional.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0]
+ out.append((feature_map, mask))
+ return out
+
+
+def get_contrastive_denoising_training_group(
+ targets,
+ num_classes,
+ num_queries,
+ class_embed,
+ num_denoising_queries=100,
+ label_noise_ratio=0.5,
+ box_noise_scale=1.0,
+):
+ """
+ Creates a contrastive denoising training group using ground-truth samples. It adds noise to labels and boxes.
+
+ Args:
+ targets (`list[dict]`):
+ The target objects, each containing 'class_labels' and 'boxes' for objects in an image.
+ num_classes (`int`):
+ Total number of classes in the dataset.
+ num_queries (`int`):
+ Number of query slots in the transformer.
+ class_embed (`callable`):
+ A function or a model layer to embed class labels.
+ num_denoising_queries (`int`, *optional*, defaults to 100):
+ Number of denoising queries.
+ label_noise_ratio (`float`, *optional*, defaults to 0.5):
+ Ratio of noise applied to labels.
+ box_noise_scale (`float`, *optional*, defaults to 1.0):
+ Scale of noise applied to bounding boxes.
+ Returns:
+ `tuple` comprising various elements:
+ - **input_query_class** (`torch.FloatTensor`) --
+ Class queries with applied label noise.
+ - **input_query_bbox** (`torch.FloatTensor`) --
+ Bounding box queries with applied box noise.
+ - **attn_mask** (`torch.FloatTensor`) --
+ Attention mask for separating denoising and reconstruction queries.
+ - **denoising_meta_values** (`dict`) --
+ Metadata including denoising positive indices, number of groups, and split sizes.
+ """
+
+ if num_denoising_queries <= 0:
+ return None, None, None, None
+
+ num_ground_truths = [len(t["class_labels"]) for t in targets]
+ device = targets[0]["class_labels"].device
+
+ max_gt_num = max(num_ground_truths)
+ if max_gt_num == 0:
+ return None, None, None, None
+
+ num_groups_denoising_queries = num_denoising_queries // max_gt_num
+ num_groups_denoising_queries = 1 if num_groups_denoising_queries == 0 else num_groups_denoising_queries
+ # pad gt to max_num of a batch
+ batch_size = len(num_ground_truths)
+
+ input_query_class = torch.full([batch_size, max_gt_num], num_classes, dtype=torch.int32, device=device)
+ input_query_bbox = torch.zeros([batch_size, max_gt_num, 4], device=device)
+ pad_gt_mask = torch.zeros([batch_size, max_gt_num], dtype=torch.bool, device=device)
+
+ for i in range(batch_size):
+ num_gt = num_ground_truths[i]
+ if num_gt > 0:
+ input_query_class[i, :num_gt] = targets[i]["class_labels"]
+ input_query_bbox[i, :num_gt] = targets[i]["boxes"]
+ pad_gt_mask[i, :num_gt] = 1
+ # each group has positive and negative queries.
+ input_query_class = input_query_class.tile([1, 2 * num_groups_denoising_queries])
+ input_query_bbox = input_query_bbox.tile([1, 2 * num_groups_denoising_queries, 1])
+ pad_gt_mask = pad_gt_mask.tile([1, 2 * num_groups_denoising_queries])
+ # positive and negative mask
+ negative_gt_mask = torch.zeros([batch_size, max_gt_num * 2, 1], device=device)
+ negative_gt_mask[:, max_gt_num:] = 1
+ negative_gt_mask = negative_gt_mask.tile([1, num_groups_denoising_queries, 1])
+ positive_gt_mask = 1 - negative_gt_mask
+ # contrastive denoising training positive index
+ positive_gt_mask = positive_gt_mask.squeeze(-1) * pad_gt_mask
+ denoise_positive_idx = torch.nonzero(positive_gt_mask)[:, 1]
+ denoise_positive_idx = torch.split(
+ denoise_positive_idx, [n * num_groups_denoising_queries for n in num_ground_truths]
+ )
+ # total denoising queries
+ num_denoising_queries = torch_int(max_gt_num * 2 * num_groups_denoising_queries)
+
+ if label_noise_ratio > 0:
+ mask = torch.rand_like(input_query_class, dtype=torch.float) < (label_noise_ratio * 0.5)
+ # randomly put a new one here
+ new_label = torch.randint_like(mask, 0, num_classes, dtype=input_query_class.dtype)
+ input_query_class = torch.where(mask & pad_gt_mask, new_label, input_query_class)
+
+ if box_noise_scale > 0:
+ known_bbox = center_to_corners_format(input_query_bbox)
+ diff = torch.tile(input_query_bbox[..., 2:] * 0.5, [1, 1, 2]) * box_noise_scale
+ rand_sign = torch.randint_like(input_query_bbox, 0, 2) * 2.0 - 1.0
+ rand_part = torch.rand_like(input_query_bbox)
+ rand_part = (rand_part + 1.0) * negative_gt_mask + rand_part * (1 - negative_gt_mask)
+ rand_part *= rand_sign
+ known_bbox += rand_part * diff
+ known_bbox.clip_(min=0.0, max=1.0)
+ input_query_bbox = corners_to_center_format(known_bbox)
+ input_query_bbox = inverse_sigmoid(input_query_bbox)
+
+ input_query_class = class_embed(input_query_class)
+
+ target_size = num_denoising_queries + num_queries
+ attn_mask = torch.full([target_size, target_size], 0, dtype=torch.float, device=device)
+ # match query cannot see the reconstruction
+ attn_mask[num_denoising_queries:, :num_denoising_queries] = -torch.inf
+
+ # reconstructions cannot see each other
+ for i in range(num_groups_denoising_queries):
+ idx_block_start = max_gt_num * 2 * i
+ idx_block_end = max_gt_num * 2 * (i + 1)
+ attn_mask[idx_block_start:idx_block_end, :idx_block_start] = -torch.inf
+ attn_mask[idx_block_start:idx_block_end, idx_block_end:num_denoising_queries] = -torch.inf
+
+ denoising_meta_values = {
+ "dn_positive_idx": denoise_positive_idx,
+ "dn_num_group": num_groups_denoising_queries,
+ "dn_num_split": [num_denoising_queries, num_queries],
+ }
+
+ return input_query_class, input_query_bbox, attn_mask, denoising_meta_values
+
+
+@auto_docstring(
+ custom_intro="""
+ RT-DETR Model (consisting of a backbone and encoder-decoder) outputting raw hidden states without any head on top.
+ """
+)
+class DFineModel(DFinePreTrainedModel):
+ def __init__(self, config: DFineConfig):
+ super().__init__(config)
+
+ # Create backbone
+ self.backbone = DFineConvEncoder(config)
+ intermediate_channel_sizes = self.backbone.intermediate_channel_sizes
+ num_backbone_outs = len(config.decoder_in_channels)
+ encoder_input_proj_list = []
+ for i in range(num_backbone_outs):
+ in_channels = intermediate_channel_sizes[i]
+ encoder_input_proj_list.append(
+ nn.Sequential(
+ nn.Conv2d(in_channels, config.encoder_hidden_dim, kernel_size=1, bias=False),
+ nn.BatchNorm2d(config.encoder_hidden_dim),
+ )
+ )
+ self.encoder_input_proj = nn.ModuleList(encoder_input_proj_list)
+ self.encoder = DFineHybridEncoder(config=config)
+
+ # denoising part
+ if config.num_denoising > 0:
+ self.denoising_class_embed = nn.Embedding(
+ config.num_labels + 1, config.d_model, padding_idx=config.num_labels
+ )
+
+ # decoder embedding
+ if config.learn_initial_query:
+ self.weight_embedding = nn.Embedding(config.num_queries, config.d_model)
+
+ # encoder head
+ self.enc_output = nn.Sequential(
+ nn.Linear(config.d_model, config.d_model),
+ nn.LayerNorm(config.d_model, eps=config.layer_norm_eps),
+ )
+ self.enc_score_head = nn.Linear(config.d_model, config.num_labels)
+ self.enc_bbox_head = DFineMLPPredictionHead(config.d_model, config.d_model, 4, num_layers=3)
+
+ # init encoder output anchors and valid_mask
+ if config.anchor_image_size:
+ self.anchors, self.valid_mask = self.generate_anchors(dtype=self.dtype)
+ num_backbone_outs = len(config.decoder_in_channels)
+ decoder_input_proj_list = []
+ for i in range(num_backbone_outs):
+ in_channels = config.decoder_in_channels[i]
+ decoder_input_proj_list.append(
+ nn.Sequential(
+ nn.Conv2d(in_channels, config.d_model, kernel_size=1, bias=False),
+ nn.BatchNorm2d(config.d_model, config.batch_norm_eps),
+ )
+ )
+ for _ in range(config.num_feature_levels - num_backbone_outs):
+ decoder_input_proj_list.append(
+ nn.Sequential(
+ nn.Conv2d(in_channels, config.d_model, kernel_size=3, stride=2, padding=1, bias=False),
+ nn.BatchNorm2d(config.d_model, config.batch_norm_eps),
+ )
+ )
+ in_channels = config.d_model
+ self.decoder = DFineDecoder(config)
+ decoder_input_proj = []
+ in_channels = config.decoder_in_channels[-1]
+ for _ in range(num_backbone_outs):
+ if config.hidden_size == config.decoder_in_channels[-1]:
+ decoder_input_proj.append(nn.Identity())
+ else:
+ conv = nn.Conv2d(in_channels, config.d_model, kernel_size=1, bias=False)
+ batchnorm = nn.BatchNorm2d(config.d_model, config.batch_norm_eps)
+ decoder_input_proj.append(nn.Sequential(conv, batchnorm))
+ for _ in range(config.num_feature_levels - num_backbone_outs):
+ if config.hidden_size == config.decoder_in_channels[-1]:
+ decoder_input_proj.append(nn.Identity())
+ else:
+ conv = nn.Conv2d(in_channels, config.d_model, kernel_size=3, stride=2, padding=1, bias=False)
+ batchnorm = nn.BatchNorm2d(config.d_model, config.batch_norm_eps)
+ decoder_input_proj.append(nn.Sequential(conv, batchnorm))
+ self.decoder_input_proj = nn.ModuleList(decoder_input_proj)
+
+ self.post_init()
+
+ def freeze_backbone(self):
+ for param in self.backbone.parameters():
+ param.requires_grad_(False)
+
+ def unfreeze_backbone(self):
+ for param in self.backbone.parameters():
+ param.requires_grad_(True)
+
+ @staticmethod
+ @compile_compatible_method_lru_cache(maxsize=32)
+ def _cached_generate_anchors(
+ spatial_shapes: tuple[tuple[int, int], ...],
+ grid_size: float,
+ device: torch.device | str = "cpu",
+ dtype: torch.dtype = torch.float32,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ anchors = []
+ for level, (height, width) in enumerate(spatial_shapes):
+ grid_y, grid_x = torch.meshgrid(
+ torch.arange(end=height, device=device).to(dtype),
+ torch.arange(end=width, device=device).to(dtype),
+ indexing="ij",
+ )
+ grid_xy = torch.stack([grid_x, grid_y], -1)
+ grid_xy = grid_xy.unsqueeze(0) + 0.5
+ grid_xy[..., 0] /= width
+ grid_xy[..., 1] /= height
+ wh = torch.ones_like(grid_xy) * grid_size * (2.0**level)
+ anchors.append(torch.concat([grid_xy, wh], -1).reshape(-1, height * width, 4))
+ # define the valid range for anchor coordinates
+ eps = 1e-2
+ anchors = torch.concat(anchors, 1)
+ valid_mask = ((anchors > eps) * (anchors < 1 - eps)).all(-1, keepdim=True)
+ anchors = torch.log(anchors / (1 - anchors))
+ anchors = torch.where(valid_mask, anchors, torch.tensor(torch.finfo(dtype).max, dtype=dtype, device=device))
+
+ return anchors, valid_mask
+
+ def generate_anchors(self, spatial_shapes=None, grid_size=0.05, device="cpu", dtype=torch.float32):
+ if spatial_shapes is None:
+ spatial_shapes = (
+ (int(self.config.anchor_image_size[0] / s), int(self.config.anchor_image_size[1] / s))
+ for s in self.config.feat_strides
+ )
+ return self._cached_generate_anchors(spatial_shapes, grid_size, device, dtype)
+
+ @auto_docstring
+ @can_return_tuple
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ pixel_mask: torch.LongTensor | None = None,
+ encoder_outputs: torch.FloatTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: list[dict] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.FloatTensor] | DFineModelOutput:
+ r"""
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you
+ can choose to directly pass a flattened representation of an image.
+ labels (`list[Dict]` of len `(batch_size,)`, *optional*):
+ Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
+ following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch
+ respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes
+ in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, DFineModel
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> image_processor = AutoImageProcessor.from_pretrained("PekingU/DFine_r50vd")
+ >>> model = DFineModel.from_pretrained("PekingU/DFine_r50vd")
+
+ >>> inputs = image_processor(images=image, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+
+ >>> last_hidden_states = outputs.last_hidden_state
+ >>> list(last_hidden_states.shape)
+ [1, 300, 256]
+ ```"""
+ if pixel_values is None and inputs_embeds is None:
+ raise ValueError("You have to specify either pixel_values or inputs_embeds")
+
+ if inputs_embeds is None:
+ batch_size, num_channels, height, width = pixel_values.shape
+ device = pixel_values.device
+ if pixel_mask is None:
+ pixel_mask = torch.ones(((batch_size, height, width)), device=device)
+ features = self.backbone(pixel_values, pixel_mask)
+ proj_feats = [self.encoder_input_proj[level](source) for level, (source, mask) in enumerate(features)]
+ else:
+ batch_size = inputs_embeds.shape[0]
+ device = inputs_embeds.device
+ proj_feats = inputs_embeds
+
+ if encoder_outputs is None:
+ encoder_outputs = self.encoder(
+ proj_feats,
+ **kwargs,
+ )
+ # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput
+ elif not isinstance(encoder_outputs, BaseModelOutput):
+ encoder_outputs = BaseModelOutput(
+ last_hidden_state=encoder_outputs[0],
+ hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+ attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+ )
+
+ # Equivalent to def _get_encoder_input
+ # https://github.com/lyuwenyu/RT-DETR/blob/94f5e16708329d2f2716426868ec89aa774af016/DFine_pytorch/src/zoo/DFine/DFine_decoder.py#L412
+ sources = []
+ for level, source in enumerate(encoder_outputs.last_hidden_state):
+ sources.append(self.decoder_input_proj[level](source))
+
+ # Lowest resolution feature maps are obtained via 3x3 stride 2 convolutions on the final stage
+ if self.config.num_feature_levels > len(sources):
+ _len_sources = len(sources)
+ sources.append(self.decoder_input_proj[_len_sources](encoder_outputs.last_hidden_state)[-1])
+ for i in range(_len_sources + 1, self.config.num_feature_levels):
+ sources.append(self.decoder_input_proj[i](encoder_outputs.last_hidden_state[-1]))
+
+ # Prepare encoder inputs (by flattening)
+ source_flatten = []
+ spatial_shapes_list = []
+ spatial_shapes = torch.empty((len(sources), 2), device=device, dtype=torch.long)
+ for level, source in enumerate(sources):
+ height, width = source.shape[-2:]
+ spatial_shapes[level, 0] = height
+ spatial_shapes[level, 1] = width
+ spatial_shapes_list.append((height, width))
+ source = source.flatten(2).transpose(1, 2)
+ source_flatten.append(source)
+ source_flatten = torch.cat(source_flatten, 1)
+ level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
+
+ # prepare denoising training
+ if self.training and self.config.num_denoising > 0 and labels is not None:
+ (
+ denoising_class,
+ denoising_bbox_unact,
+ attention_mask,
+ denoising_meta_values,
+ ) = get_contrastive_denoising_training_group(
+ targets=labels,
+ num_classes=self.config.num_labels,
+ num_queries=self.config.num_queries,
+ class_embed=self.denoising_class_embed,
+ num_denoising_queries=self.config.num_denoising,
+ label_noise_ratio=self.config.label_noise_ratio,
+ box_noise_scale=self.config.box_noise_scale,
+ )
+ else:
+ denoising_class, denoising_bbox_unact, attention_mask, denoising_meta_values = None, None, None, None
+
+ batch_size = len(source_flatten)
+ device = source_flatten.device
+ dtype = source_flatten.dtype
+
+ # prepare input for decoder
+ if self.training or self.config.anchor_image_size is None:
+ # Pass spatial_shapes as tuple to make it hashable and make sure
+ # lru_cache is working for generate_anchors()
+ spatial_shapes_tuple = tuple(spatial_shapes_list)
+ anchors, valid_mask = self.generate_anchors(spatial_shapes_tuple, device=device, dtype=dtype)
+ else:
+ anchors, valid_mask = self.anchors, self.valid_mask
+ anchors, valid_mask = anchors.to(device, dtype), valid_mask.to(device, dtype)
+
+ # use the valid_mask to selectively retain values in the feature map where the mask is `True`
+ memory = valid_mask.to(source_flatten.dtype) * source_flatten
+
+ output_memory = self.enc_output(memory)
+
+ enc_outputs_class = self.enc_score_head(output_memory)
+ enc_outputs_coord_logits = self.enc_bbox_head(output_memory) + anchors
+
+ _, topk_ind = torch.topk(enc_outputs_class.max(-1).values, self.config.num_queries, dim=1)
+
+ reference_points_unact = enc_outputs_coord_logits.gather(
+ dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, enc_outputs_coord_logits.shape[-1])
+ )
+
+ enc_topk_bboxes = F.sigmoid(reference_points_unact)
+ if denoising_bbox_unact is not None:
+ reference_points_unact = torch.concat([denoising_bbox_unact, reference_points_unact], 1)
+
+ enc_topk_logits = enc_outputs_class.gather(
+ dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, enc_outputs_class.shape[-1])
+ )
+
+ # extract region features
+ if self.config.learn_initial_query:
+ target = self.weight_embedding.tile([batch_size, 1, 1])
+ else:
+ target = output_memory.gather(dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, output_memory.shape[-1]))
+ target = target.detach()
+
+ if denoising_class is not None:
+ target = torch.concat([denoising_class, target], 1)
+
+ init_reference_points = reference_points_unact.detach()
+
+ # decoder
+ decoder_outputs = self.decoder(
+ inputs_embeds=target,
+ encoder_hidden_states=source_flatten,
+ encoder_attention_mask=attention_mask,
+ reference_points=init_reference_points,
+ spatial_shapes=spatial_shapes,
+ spatial_shapes_list=spatial_shapes_list,
+ level_start_index=level_start_index,
+ **kwargs,
+ )
+
+ return DFineModelOutput(
+ last_hidden_state=decoder_outputs.last_hidden_state,
+ intermediate_hidden_states=decoder_outputs.intermediate_hidden_states,
+ intermediate_logits=decoder_outputs.intermediate_logits,
+ intermediate_reference_points=decoder_outputs.intermediate_reference_points,
+ intermediate_predicted_corners=decoder_outputs.intermediate_predicted_corners,
+ initial_reference_points=decoder_outputs.initial_reference_points,
+ decoder_hidden_states=decoder_outputs.hidden_states,
+ decoder_attentions=decoder_outputs.attentions,
+ cross_attentions=decoder_outputs.cross_attentions,
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+ encoder_hidden_states=encoder_outputs.hidden_states,
+ encoder_attentions=encoder_outputs.attentions,
+ init_reference_points=init_reference_points,
+ enc_topk_logits=enc_topk_logits,
+ enc_topk_bboxes=enc_topk_bboxes,
+ enc_outputs_class=enc_outputs_class,
+ enc_outputs_coord_logits=enc_outputs_coord_logits,
+ denoising_meta_values=denoising_meta_values,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Output type of [`DFineForObjectDetection`].
+ """
+)
+@dataclass
+class DFineObjectDetectionOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
+ Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a
+ bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
+ scale-invariant IoU loss.
+ loss_dict (`Dict`, *optional*):
+ A dictionary containing the individual losses. Useful for logging.
+ logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
+ Classification logits (including no-object) for all queries.
+ pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
+ Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
+ values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
+ possible padding). You can use [`~DFineImageProcessor.post_process_object_detection`] to retrieve the
+ unnormalized (absolute) bounding boxes.
+ auxiliary_outputs (`list[Dict]`, *optional*):
+ Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
+ and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
+ `pred_boxes`) for each decoder layer.
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`):
+ Sequence of hidden-states at the output of the last layer of the decoder of the model.
+ intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
+ Stacked intermediate hidden states (output of each layer of the decoder).
+ intermediate_logits (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, config.num_labels)`):
+ Stacked intermediate logits (logits of each layer of the decoder).
+ intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+ Stacked intermediate reference points (reference points of each layer of the decoder).
+ intermediate_predicted_corners (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+ Stacked intermediate predicted corners (predicted corners of each layer of the decoder).
+ initial_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+ Stacked initial reference points (initial reference points of each layer of the decoder).
+ init_reference_points (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
+ Initial reference points sent through the Transformer decoder.
+ enc_topk_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+ Logits of predicted bounding boxes coordinates in the encoder.
+ enc_topk_bboxes (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+ Logits of predicted bounding boxes coordinates in the encoder.
+ enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+ Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
+ picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
+ foreground and background).
+ enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+ Logits of predicted bounding boxes coordinates in the first stage.
+ denoising_meta_values (`dict`):
+ Extra dictionary for the denoising related values
+ """
+
+ loss: torch.FloatTensor | None = None
+ loss_dict: dict | None = None
+ logits: torch.FloatTensor | None = None
+ pred_boxes: torch.FloatTensor | None = None
+ auxiliary_outputs: list[dict] | None = None
+ last_hidden_state: torch.FloatTensor | None = None
+ intermediate_hidden_states: torch.FloatTensor | None = None
+ intermediate_logits: torch.FloatTensor | None = None
+ intermediate_reference_points: torch.FloatTensor | None = None
+ intermediate_predicted_corners: torch.FloatTensor | None = None
+ initial_reference_points: torch.FloatTensor | None = None
+ decoder_hidden_states: tuple[torch.FloatTensor] | None = None
+ decoder_attentions: tuple[torch.FloatTensor] | None = None
+ cross_attentions: tuple[torch.FloatTensor] | None = None
+ encoder_last_hidden_state: torch.FloatTensor | None = None
+ encoder_hidden_states: tuple[torch.FloatTensor] | None = None
+ encoder_attentions: tuple[torch.FloatTensor] | None = None
+ init_reference_points: tuple[torch.FloatTensor] | None = None
+ enc_topk_logits: torch.FloatTensor | None = None
+ enc_topk_bboxes: torch.FloatTensor | None = None
+ enc_outputs_class: torch.FloatTensor | None = None
+ enc_outputs_coord_logits: torch.FloatTensor | None = None
+ denoising_meta_values: dict | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ RT-DETR Model (consisting of a backbone and encoder-decoder) outputting bounding boxes and logits to be further
+ decoded into scores and classes.
+ """
+)
+class DFineForObjectDetection(DFinePreTrainedModel):
+ # When using clones, all layers > 0 will be clones, but layer 0 *is* required
+ # We can't initialize the model on meta device as some weights are modified during the initialization
+ _no_split_modules = None
+ _tied_weights_keys = {
+ r"bbox_embed.(?![0])\d+": r"bbox_embed.0",
+ r"class_embed.(?![0])\d+": r"^class_embed.0",
+ "class_embed": "model.decoder.class_embed",
+ "bbox_embed": "model.decoder.bbox_embed",
+ }
+
+ def __init__(self, config: DFineConfig):
+ super().__init__(config)
+
+ # D-FINE encoder-decoder model
+ self.eval_idx = config.eval_idx if config.eval_idx >= 0 else config.decoder_layers + config.eval_idx
+ self.model = DFineModel(config)
+ scaled_dim = round(config.layer_scale * config.hidden_size)
+ num_pred = config.decoder_layers
+ self.class_embed = nn.ModuleList([nn.Linear(config.d_model, config.num_labels) for _ in range(num_pred)])
+ self.bbox_embed = nn.ModuleList(
+ [
+ DFineMLP(config.hidden_size, config.hidden_size, 4 * (config.max_num_bins + 1), 3)
+ for _ in range(self.eval_idx + 1)
+ ]
+ + [
+ DFineMLP(scaled_dim, scaled_dim, 4 * (config.max_num_bins + 1), 3)
+ for _ in range(config.decoder_layers - self.eval_idx - 1)
+ ]
+ )
+
+ self.model.decoder.class_embed = self.class_embed
+ self.model.decoder.bbox_embed = self.bbox_embed
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def _set_aux_loss(self, outputs_class, outputs_coord):
+ return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class, outputs_coord)]
+
+ @auto_docstring
+ @can_return_tuple
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ pixel_mask: torch.LongTensor | None = None,
+ encoder_outputs: torch.FloatTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: list[dict] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.FloatTensor] | DFineObjectDetectionOutput:
+ r"""
+ Example:
+
+ ```python
+ >>> import torch
+ >>> from transformers.image_utils import load_image
+ >>> from transformers import AutoImageProcessor, DFineForObjectDetection
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = load_image(url)
+
+ >>> image_processor = AutoImageProcessor.from_pretrained("ustc-community/dfine-xlarge-coco")
+ >>> model = DFineForObjectDetection.from_pretrained("ustc-community/dfine-xlarge-coco")
+
+ >>> # prepare image for the model
+ >>> inputs = image_processor(images=image, return_tensors="pt")
+
+ >>> # forward pass
+ >>> outputs = model(**inputs)
+
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 300, 80]
+
+ >>> boxes = outputs.pred_boxes
+ >>> list(boxes.shape)
+ [1, 300, 4]
+
+ >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
+ >>> target_sizes = torch.tensor([image.size[::-1]])
+ >>> results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)
+ >>> result = results[0] # first image in batch
+
+ >>> for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
+ ... box = [round(i, 2) for i in box.tolist()]
+ ... print(
+ ... f"Detected {model.config.id2label[label.item()]} with confidence "
+ ... f"{round(score.item(), 3)} at location {box}"
+ ... )
+ Detected cat with confidence 0.958 at location [344.49, 23.4, 639.84, 374.27]
+ Detected cat with confidence 0.956 at location [11.71, 53.52, 316.64, 472.33]
+ Detected remote with confidence 0.947 at location [40.46, 73.7, 175.62, 117.57]
+ Detected sofa with confidence 0.918 at location [0.59, 1.88, 640.25, 474.74]
+ ```
+ """
+ outputs = self.model(
+ pixel_values,
+ pixel_mask=pixel_mask,
+ encoder_outputs=encoder_outputs,
+ inputs_embeds=inputs_embeds,
+ labels=labels,
+ **kwargs,
+ )
+
+ denoising_meta_values = outputs.denoising_meta_values if self.training else None
+
+ outputs_class = outputs.intermediate_logits
+ outputs_coord = outputs.intermediate_reference_points
+ predicted_corners = outputs.intermediate_predicted_corners
+ initial_reference_points = outputs.initial_reference_points
+
+ logits = outputs_class[:, -1]
+ pred_boxes = outputs_coord[:, -1]
+
+ loss, loss_dict, auxiliary_outputs, enc_topk_logits, enc_topk_bboxes = None, None, None, None, None
+ if labels is not None:
+ enc_topk_logits = outputs.enc_topk_logits
+ enc_topk_bboxes = outputs.enc_topk_bboxes
+ loss, loss_dict, auxiliary_outputs = self.loss_function(
+ logits,
+ labels,
+ self.device,
+ pred_boxes,
+ self.config,
+ outputs_class,
+ outputs_coord,
+ enc_topk_logits=enc_topk_logits,
+ enc_topk_bboxes=enc_topk_bboxes,
+ denoising_meta_values=denoising_meta_values,
+ predicted_corners=predicted_corners,
+ initial_reference_points=initial_reference_points,
+ **kwargs,
+ )
+
+ return DFineObjectDetectionOutput(
+ loss=loss,
+ loss_dict=loss_dict,
+ logits=logits,
+ pred_boxes=pred_boxes,
+ auxiliary_outputs=auxiliary_outputs,
+ last_hidden_state=outputs.last_hidden_state,
+ intermediate_hidden_states=outputs.intermediate_hidden_states,
+ intermediate_logits=outputs.intermediate_logits,
+ intermediate_reference_points=outputs.intermediate_reference_points,
+ intermediate_predicted_corners=outputs.intermediate_predicted_corners,
+ initial_reference_points=outputs.initial_reference_points,
+ decoder_hidden_states=outputs.decoder_hidden_states,
+ decoder_attentions=outputs.decoder_attentions,
+ cross_attentions=outputs.cross_attentions,
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+ encoder_hidden_states=outputs.encoder_hidden_states,
+ encoder_attentions=outputs.encoder_attentions,
+ init_reference_points=outputs.init_reference_points,
+ enc_topk_logits=outputs.enc_topk_logits,
+ enc_topk_bboxes=outputs.enc_topk_bboxes,
+ enc_outputs_class=outputs.enc_outputs_class,
+ enc_outputs_coord_logits=outputs.enc_outputs_coord_logits,
+ denoising_meta_values=outputs.denoising_meta_values,
+ )
+
+
+__all__ = ["DFineModel", "DFinePreTrainedModel", "DFineForObjectDetection"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/modular_d_fine.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/modular_d_fine.py
new file mode 100644
index 0000000000000000000000000000000000000000..49289f0750378142e32ee4cd74fc642b20751f3d
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/d_fine/modular_d_fine.py
@@ -0,0 +1,1016 @@
+# Copyright 2025 Baidu Inc and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub.dataclasses import strict
+
+from ... import initialization as init
+from ...activations import ACT2CLS
+from ...backbone_utils import consolidate_backbone_kwargs_to_config
+from ...configuration_utils import PreTrainedConfig
+from ...image_transforms import corners_to_center_format
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, logging, torch_compilable_check
+from ..auto import AutoConfig
+from ..rt_detr.modeling_rt_detr import (
+ RTDetrAIFILayer,
+ RTDetrConvNormLayer,
+ RTDetrDecoder,
+ RTDetrDecoderLayer,
+ RTDetrDecoderOutput,
+ RTDetrEncoderLayer,
+ RTDetrForObjectDetection,
+ RTDetrFrozenBatchNorm2d,
+ RTDetrHybridEncoder,
+ RTDetrMLPPredictionHead,
+ RTDetrModel,
+ RTDetrPreTrainedModel,
+ RTDetrRepVggBlock,
+ inverse_sigmoid,
+)
+from ..rt_detr_v2.modeling_rt_detr_v2 import multi_scale_deformable_attention_v2
+
+
+logger = logging.get_logger(__name__)
+
+
+# TODO: Attribute map assignment logic should be fixed in modular
+# as well as super() call parsing because otherwise we cannot re-write args after initialization
+@auto_docstring(checkpoint="ustc-community/dfine-xlarge-coco")
+@strict
+class DFineConfig(PreTrainedConfig):
+ r"""
+ initializer_bias_prior_prob (`float`, *optional*):
+ The prior probability used by the bias initializer to initialize biases for `enc_score_head` and `class_embed`.
+ If `None`, `prior_prob` computed as `prior_prob = 1 / (num_labels + 1)` while initializing model weights.
+ freeze_backbone_batch_norms (`bool`, *optional*, defaults to `True`):
+ Whether to freeze the batch normalization layers in the backbone.
+ encoder_in_channels (`list`, *optional*, defaults to `[512, 1024, 2048]`):
+ Multi level features input for encoder.
+ feat_strides (`list[int]`, *optional*, defaults to `[8, 16, 32]`):
+ Strides used in each feature map.
+ encode_proj_layers (`list[int]`, *optional*, defaults to `[2]`):
+ Indexes of the projected layers to be used in the encoder.
+ positional_encoding_temperature (`int`, *optional*, defaults to 10000):
+ The temperature parameter used to create the positional encodings.
+ encoder_activation_function (`str`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+ `"relu"`, `"silu"` and `"gelu_new"` are supported.
+ eval_size (`tuple[int, int]`, *optional*):
+ Height and width used to computes the effective height and width of the position embeddings after taking
+ into account the stride.
+ normalize_before (`bool`, *optional*, defaults to `False`):
+ Determine whether to apply layer normalization in the transformer encoder layer before self-attention and
+ feed-forward modules.
+ hidden_expansion (`float`, *optional*, defaults to 1.0):
+ Expansion ratio to enlarge the dimension size of RepVGGBlock and CSPRepLayer.
+ num_queries (`int`, *optional*, defaults to 300):
+ Number of object queries.
+ decoder_in_channels (`list`, *optional*, defaults to `[256, 256, 256]`):
+ Multi level features dimension for decoder
+ num_feature_levels (`int`, *optional*, defaults to 3):
+ The number of input feature levels.
+ decoder_n_points (`int`, *optional*, defaults to 4):
+ The number of sampled keys in each feature level for each attention head in the decoder.
+ decoder_activation_function (`str`, *optional*, defaults to `"relu"`):
+ The non-linear activation function (function or string) in the decoder. If string, `"gelu"`,
+ `"relu"`, `"silu"` and `"gelu_new"` are supported.
+ num_denoising (`int`, *optional*, defaults to 100):
+ The total number of denoising tasks or queries to be used for contrastive denoising.
+ label_noise_ratio (`float`, *optional*, defaults to 0.5):
+ The fraction of denoising labels to which random noise should be added.
+ box_noise_scale (`float`, *optional*, defaults to 1.0):
+ Scale or magnitude of noise to be added to the bounding boxes.
+ learn_initial_query (`bool`, *optional*, defaults to `False`):
+ Indicates whether the initial query embeddings for the decoder should be learned during training
+ anchor_image_size (`tuple[int, int]`, *optional*):
+ Height and width of the input image used during evaluation to generate the bounding box anchors. If None, automatic generate anchor is applied.
+ with_box_refine (`bool`, *optional*, defaults to `True`):
+ Whether to apply iterative bounding box refinement, where each decoder layer refines the bounding boxes
+ based on the predictions from the previous layer.
+ matcher_alpha (`float`, *optional*, defaults to 0.25):
+ Parameter alpha used by the Hungarian Matcher.
+ matcher_gamma (`float`, *optional*, defaults to 2.0):
+ Parameter gamma used by the Hungarian Matcher.
+ matcher_class_cost (`float`, *optional*, defaults to 2.0):
+ The relative weight of the class loss used by the Hungarian Matcher.
+ matcher_bbox_cost (`float`, *optional*, defaults to 5.0):
+ The relative weight of the bounding box loss used by the Hungarian Matcher.
+ matcher_giou_cost (`float`, *optional*, defaults to 2.0):
+ The relative weight of the giou loss of used by the Hungarian Matcher.
+ use_focal_loss (`bool`, *optional*, defaults to `True`):
+ Parameter informing if focal focal should be used.
+ focal_loss_alpha (`float`, *optional*, defaults to 0.75):
+ Parameter alpha used to compute the focal loss.
+ focal_loss_gamma (`float`, *optional*, defaults to 2.0):
+ Parameter gamma used to compute the focal loss.
+ weight_loss_vfl (`float`, *optional*, defaults to 1.0):
+ Relative weight of the varifocal loss in the object detection loss.
+ weight_loss_bbox (`float`, *optional*, defaults to 5.0):
+ Relative weight of the L1 bounding box loss in the object detection loss.
+ weight_loss_giou (`float`, *optional*, defaults to 2.0):
+ Relative weight of the generalized IoU loss in the object detection loss.
+ weight_loss_fgl (`float`, *optional*, defaults to 0.15):
+ Relative weight of the fine-grained localization loss in the object detection loss.
+ weight_loss_ddf (`float`, *optional*, defaults to 1.5):
+ Relative weight of the decoupled distillation focal loss in the object detection loss.
+ eval_idx (`int`, *optional*, defaults to -1):
+ Index of the decoder layer to use for evaluation. If negative, counts from the end
+ (e.g., -1 means use the last layer). This allows for early prediction in the decoder
+ stack while still training later layers.
+ layer_scale (`float`, *optional*, defaults to `1.0`):
+ Scaling factor for the hidden dimension in later decoder layers. Used to adjust the
+ model capacity after the evaluation layer.
+ max_num_bins (`int`, *optional*, defaults to 32):
+ Maximum number of bins for the distribution-guided bounding box refinement.
+ Higher values allow for more fine-grained localization but increase computation.
+ reg_scale (`float`, *optional*, defaults to 4.0):
+ Scale factor for the regression distribution. Controls the range and granularity
+ of the bounding box refinement process.
+ depth_mult (`float`, *optional*, defaults to 1.0):
+ Multiplier for the number of blocks in RepNCSPELAN4 layers. Used to scale the model's
+ depth while maintaining its architecture.
+ top_prob_values (`int`, *optional*, defaults to 4):
+ Number of top probability values to consider from each corner's distribution.
+ lqe_hidden_dim (`int`, *optional*, defaults to 64):
+ Hidden dimension size for the Location Quality Estimator (LQE) network.
+ lqe_layers (`int`, *optional*, defaults to 2):
+ Number of layers in the Location Quality Estimator MLP.
+ decoder_offset_scale (`float`, *optional*, defaults to 0.5):
+ Offset scale used in deformable attention.
+ decoder_method (`str`, *optional*, defaults to `"default"`):
+ The method to use for the decoder: `"default"` or `"discrete"`.
+ up (`float`, *optional*, defaults to 0.5):
+ Controls the upper bounds of the Weighting Function.
+ """
+
+ model_type = "d_fine"
+ sub_configs = {"backbone_config": AutoConfig}
+ layer_types = ["basic", "bottleneck"]
+ attribute_map = {
+ "hidden_size": "d_model",
+ "num_attention_heads": "encoder_attention_heads",
+ }
+
+ initializer_range: float = 0.01
+ initializer_bias_prior_prob: float | None = None
+ layer_norm_eps: float = 1e-5
+ batch_norm_eps: float = 1e-5
+ backbone_config: dict | PreTrainedConfig | None = None
+ freeze_backbone_batch_norms: bool = True
+
+ # encoder HybridEncoder
+ encoder_hidden_dim: int = 256
+ encoder_in_channels: list[int] | tuple[int, ...] = (512, 1024, 2048)
+ feat_strides: list[int] | tuple[int, ...] = (8, 16, 32)
+ encoder_layers: int = 1
+ encoder_ffn_dim: int = 1024
+ encoder_attention_heads: int = 8
+ dropout: float | int = 0.0
+ activation_dropout: float | int = 0.0
+ encode_proj_layers: list[int] | tuple[int, ...] = (2,)
+ positional_encoding_temperature: int = 10000
+ encoder_activation_function: str = "gelu"
+ activation_function: str = "silu"
+ eval_size: int | None = None
+ normalize_before: bool = False
+ hidden_expansion: float = 1.0
+
+ # decoder DFineTransformer
+ d_model: int = 256
+ num_queries: int = 300
+ decoder_in_channels: list[int] | tuple[int, ...] = (256, 256, 256)
+ decoder_ffn_dim: int = 1024
+ num_feature_levels: int = 3
+ decoder_n_points: int | list[int] = 4
+ decoder_layers: int = 6
+ decoder_attention_heads: int = 8
+ decoder_activation_function: str = "relu"
+ attention_dropout: float | int = 0.0
+ num_denoising: int = 100
+ label_noise_ratio: float = 0.5
+ box_noise_scale: float = 1.0
+ learn_initial_query: bool = False
+ anchor_image_size: int | list[int] | None = None
+ with_box_refine: bool = True
+
+ # Loss
+ matcher_alpha: float = 0.25
+ matcher_gamma: float = 2.0
+ matcher_class_cost: float = 2.0
+ matcher_bbox_cost: float = 5.0
+ matcher_giou_cost: float = 2.0
+ use_focal_loss: bool = True
+ auxiliary_loss: bool = True
+ focal_loss_alpha: float = 0.75
+ focal_loss_gamma: float = 2.0
+ weight_loss_vfl: float = 1.0
+ weight_loss_bbox: float = 5.0
+ weight_loss_giou: float = 2.0
+ weight_loss_fgl: float = 0.15
+ weight_loss_ddf: float = 1.5
+ eos_coefficient: float = 1e-4
+ eval_idx: int = -1
+ layer_scale: int | float = 1.0
+ max_num_bins: int = 32
+ reg_scale: float = 4.0
+ depth_mult: float = 1.0
+ top_prob_values: int = 4
+ lqe_hidden_dim: int = 64
+ lqe_layers: int = 2
+ decoder_offset_scale: float = 0.5
+ decoder_method: str = "default"
+ up: float = 0.5
+ tie_word_embeddings: bool = True
+ is_encoder_decoder: bool = True
+
+ def __post_init__(self, **kwargs):
+ self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config(
+ backbone_config=self.backbone_config,
+ default_config_type="hgnet_v2",
+ default_config_kwargs={"out_indices": [2, 3, 4]},
+ **kwargs,
+ )
+ self.head_dim = self.d_model // self.decoder_attention_heads
+ super().__post_init__(**kwargs)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if isinstance(self.decoder_n_points, list):
+ if len(self.decoder_n_points) != self.num_feature_levels:
+ raise ValueError(
+ f"Length of decoder_n_points list ({len(self.decoder_n_points)}) must match num_feature_levels ({self.num_feature_levels})."
+ )
+
+ if self.head_dim * self.decoder_attention_heads != self.d_model:
+ raise ValueError(
+ f"Embedded dimension {self.d_model} must be divisible by decoder_attention_heads {self.decoder_attention_heads}"
+ )
+
+
+class DFineDecoderOutput(RTDetrDecoderOutput):
+ pass
+
+
+def weighting_function(max_num_bins: int, up: torch.Tensor, reg_scale: int) -> torch.Tensor:
+ """
+ Generates the non-uniform Weighting Function W(n) for bounding box regression.
+
+ Args:
+ max_num_bins (int): Max number of the discrete bins.
+ up (Tensor): Controls upper bounds of the sequence,
+ where maximum offset is ±up * H / W.
+ reg_scale (float): Controls the curvature of the Weighting Function.
+ Larger values result in flatter weights near the central axis W(max_num_bins/2)=0
+ and steeper weights at both ends.
+ Returns:
+ Tensor: Sequence of Weighting Function.
+ """
+ upper_bound1 = abs(up[0]) * abs(reg_scale)
+ upper_bound2 = abs(up[0]) * abs(reg_scale) * 2
+ step = (upper_bound1 + 1) ** (2 / (max_num_bins - 2))
+ left_values = [-((step) ** i) + 1 for i in range(max_num_bins // 2 - 1, 0, -1)]
+ right_values = [(step) ** i - 1 for i in range(1, max_num_bins // 2)]
+ values = [-upper_bound2] + left_values + [torch.zeros_like(up[0][None])] + right_values + [upper_bound2]
+ values = torch.cat(values, 0)
+ return values
+
+
+def distance2bbox(points, distance: torch.Tensor, reg_scale: float) -> torch.Tensor:
+ """
+ Decodes edge-distances into bounding box coordinates.
+
+ Args:
+ points (`torch.Tensor`):
+ (batch_size, num_boxes, 4) or (num_boxes, 4) format, representing [x_center, y_center, width, height]
+ distance (`torch.Tensor`):
+ (batch_size, num_boxes, 4) or (num_boxes, 4), representing distances from the point to the left, top, right, and bottom boundaries.
+ reg_scale (`float`):
+ Controls the curvature of the Weighting Function.
+ Returns:
+ `torch.Tensor`: Bounding boxes in (batch_size, num_boxes, 4) or (num_boxes, 4) format, representing [x_center, y_center, width, height]
+ """
+ reg_scale = abs(reg_scale)
+ top_left_x = points[..., 0] - (0.5 * reg_scale + distance[..., 0]) * (points[..., 2] / reg_scale)
+ top_left_y = points[..., 1] - (0.5 * reg_scale + distance[..., 1]) * (points[..., 3] / reg_scale)
+ bottom_right_x = points[..., 0] + (0.5 * reg_scale + distance[..., 2]) * (points[..., 2] / reg_scale)
+ bottom_right_y = points[..., 1] + (0.5 * reg_scale + distance[..., 3]) * (points[..., 3] / reg_scale)
+
+ bboxes = torch.stack([top_left_x, top_left_y, bottom_right_x, bottom_right_y], -1)
+
+ return corners_to_center_format(bboxes)
+
+
+class DFineMLP(nn.Module):
+ def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int, act: str = "relu"):
+ super().__init__()
+ self.num_layers = num_layers
+ hidden_dims = [hidden_dim] * (num_layers - 1)
+ input_dims = [input_dim] + hidden_dims
+ output_dims = hidden_dims + [output_dim]
+ self.layers = nn.ModuleList(nn.Linear(in_dim, out_dim) for in_dim, out_dim in zip(input_dims, output_dims))
+ self.act = ACT2CLS[act]()
+
+ def forward(self, stat_features: torch.Tensor) -> torch.Tensor:
+ for i, layer in enumerate(self.layers):
+ stat_features = self.act(layer(stat_features)) if i < self.num_layers - 1 else layer(stat_features)
+ return stat_features
+
+
+class DFineGate(nn.Module):
+ def __init__(self, d_model: int):
+ super().__init__()
+ self.gate = nn.Linear(2 * d_model, 2 * d_model)
+ self.norm = nn.LayerNorm(d_model)
+
+ def forward(self, second_residual: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
+ gate_input = torch.cat([second_residual, hidden_states], dim=-1)
+ gates = torch.sigmoid(self.gate(gate_input))
+ gate1, gate2 = gates.chunk(2, dim=-1)
+ hidden_states = self.norm(gate1 * second_residual + gate2 * hidden_states)
+ return hidden_states
+
+
+class DFineFrozenBatchNorm2d(RTDetrFrozenBatchNorm2d):
+ pass
+
+
+class DFineMultiscaleDeformableAttention(nn.Module):
+ def __init__(self, config: DFineConfig):
+ """
+ D-Fine version of multiscale deformable attention
+ """
+ super().__init__()
+ self.d_model = config.d_model
+ self.n_heads = config.decoder_attention_heads
+ self.n_levels = config.num_feature_levels
+ self.offset_scale = config.decoder_offset_scale
+ self.decoder_method = config.decoder_method
+ self.n_points = config.decoder_n_points
+
+ if isinstance(self.n_points, list):
+ num_points_list = self.n_points
+ else:
+ num_points_list = [self.n_points for _ in range(self.n_levels)]
+
+ self.num_points_list = num_points_list
+ num_points_scale = [1 / n for n in self.num_points_list for _ in range(n)]
+ self.register_buffer("num_points_scale", torch.tensor(num_points_scale, dtype=torch.float32))
+
+ self.total_points = self.n_heads * sum(self.num_points_list)
+
+ self.sampling_offsets = nn.Linear(self.d_model, self.total_points * 2)
+ self.attention_weights = nn.Linear(self.d_model, self.total_points)
+
+ self.ms_deformable_attn_core = multi_scale_deformable_attention_v2
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ reference_points=None,
+ encoder_hidden_states=None,
+ spatial_shapes=None,
+ spatial_shapes_list=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ batch_size, num_queries, _ = hidden_states.shape
+ batch_size, sequence_length, _ = encoder_hidden_states.shape
+
+ torch_compilable_check(
+ (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == sequence_length,
+ "Make sure to align the spatial shapes with the sequence length of the encoder hidden states",
+ )
+
+ # Reshape for multi-head attention
+ value = encoder_hidden_states.reshape(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads)
+ if attention_mask is not None:
+ value = value.masked_fill(~attention_mask[..., None], float(0))
+
+ sampling_offsets: torch.Tensor = self.sampling_offsets(hidden_states)
+ sampling_offsets = sampling_offsets.reshape(
+ batch_size, num_queries, self.n_heads, sum(self.num_points_list), 2
+ )
+
+ attention_weights = self.attention_weights(hidden_states).reshape(
+ batch_size, num_queries, self.n_heads, sum(self.num_points_list)
+ )
+ attention_weights = F.softmax(attention_weights, dim=-1)
+
+ if reference_points.shape[-1] == 2:
+ offset_normalizer = torch.tensor(spatial_shapes)
+ offset_normalizer = offset_normalizer.flip([1]).reshape(1, 1, 1, self.n_levels, 1, 2)
+ sampling_locations = (
+ reference_points.reshape(batch_size, sequence_length, 1, self.n_levels, 1, 2)
+ + sampling_offsets / offset_normalizer
+ )
+ elif reference_points.shape[-1] == 4:
+ # reference_points [8, 480, None, 1, 4]
+ # sampling_offsets [8, 480, 8, 12, 2]
+ num_points_scale = self.num_points_scale.to(dtype=hidden_states.dtype).unsqueeze(-1)
+ offset = sampling_offsets * num_points_scale * reference_points[:, :, None, :, 2:] * self.offset_scale
+ sampling_locations = reference_points[:, :, None, :, :2] + offset
+ else:
+ raise ValueError(
+ f"Last dim of reference_points must be 2 or 4, but get {reference_points.shape[-1]} instead."
+ )
+
+ output = self.ms_deformable_attn_core(
+ value,
+ spatial_shapes_list,
+ sampling_locations,
+ attention_weights,
+ self.num_points_list,
+ self.decoder_method,
+ )
+
+ return output, attention_weights
+
+
+class DFineConvNormLayer(RTDetrConvNormLayer):
+ def __init__(
+ self,
+ config: DFineConfig,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int,
+ groups: int = 1,
+ padding: int | None = None,
+ activation: str | None = None,
+ ):
+ super().__init__(config, in_channels, out_channels, kernel_size, stride, padding=None, activation=activation)
+ self.conv = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride,
+ groups=groups,
+ padding=(kernel_size - 1) // 2 if padding is None else padding,
+ bias=False,
+ )
+
+
+class DFineRepVggBlock(RTDetrRepVggBlock):
+ def __init__(self, config: DFineConfig, in_channels: int, out_channels: int):
+ super().__init__(config)
+ hidden_channels = in_channels
+ self.conv1 = DFineConvNormLayer(config, hidden_channels, out_channels, 3, 1, padding=1)
+ self.conv2 = DFineConvNormLayer(config, hidden_channels, out_channels, 1, 1, padding=0)
+
+
+class DFineCSPRepLayer(nn.Module):
+ """
+ Cross Stage Partial (CSP) network layer with RepVGG blocks.
+ """
+
+ def __init__(
+ self, config: DFineConfig, in_channels: int, out_channels: int, num_blocks: int, expansion: float = 1.0
+ ):
+ super().__init__()
+ activation = config.activation_function
+
+ hidden_channels = int(out_channels * expansion)
+ self.conv1 = DFineConvNormLayer(config, in_channels, hidden_channels, 1, 1, activation=activation)
+ self.conv2 = DFineConvNormLayer(config, in_channels, hidden_channels, 1, 1, activation=activation)
+ self.bottlenecks = nn.ModuleList(
+ [DFineRepVggBlock(config, hidden_channels, hidden_channels) for _ in range(num_blocks)]
+ )
+ if hidden_channels != out_channels:
+ self.conv3 = DFineConvNormLayer(config, hidden_channels, out_channels, 1, 1, activation=activation)
+ else:
+ self.conv3 = nn.Identity()
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state_1 = self.conv1(hidden_state)
+ for bottleneck in self.bottlenecks:
+ hidden_state_1 = bottleneck(hidden_state_1)
+ hidden_state_2 = self.conv2(hidden_state)
+ hidden_state_3 = self.conv3(hidden_state_1 + hidden_state_2)
+ return hidden_state_3
+
+
+class DFineRepNCSPELAN4(nn.Module):
+ def __init__(self, config: DFineConfig, act: str = "silu", numb_blocks: int = 3):
+ super().__init__()
+ conv1_dim = config.encoder_hidden_dim * 2
+ conv2_dim = config.encoder_hidden_dim
+ conv3_dim = config.encoder_hidden_dim * 2
+ conv4_dim = round(config.hidden_expansion * config.encoder_hidden_dim // 2)
+ self.conv_dim = conv3_dim // 2
+ self.conv1 = DFineConvNormLayer(config, conv1_dim, conv3_dim, 1, 1, activation=act)
+ self.csp_rep1 = DFineCSPRepLayer(config, conv3_dim // 2, conv4_dim, num_blocks=numb_blocks)
+ self.conv2 = DFineConvNormLayer(config, conv4_dim, conv4_dim, 3, 1, activation=act)
+ self.csp_rep2 = DFineCSPRepLayer(config, conv4_dim, conv4_dim, num_blocks=numb_blocks)
+ self.conv3 = DFineConvNormLayer(config, conv4_dim, conv4_dim, 3, 1, activation=act)
+ self.conv4 = DFineConvNormLayer(config, conv3_dim + (2 * conv4_dim), conv2_dim, 1, 1, activation=act)
+
+ def forward(self, input_features: torch.Tensor) -> torch.Tensor:
+ # Split initial features into two branches after first convolution
+ split_features = list(self.conv1(input_features).split((self.conv_dim, self.conv_dim), 1))
+
+ # Process branches sequentially
+ branch1 = self.csp_rep1(split_features[-1])
+ branch1 = self.conv2(branch1)
+ branch2 = self.csp_rep2(branch1)
+ branch2 = self.conv3(branch2)
+
+ split_features.extend([branch1, branch2])
+ merged_features = torch.cat(split_features, 1)
+ merged_features = self.conv4(merged_features)
+ return merged_features
+
+
+class DFineSCDown(nn.Module):
+ def __init__(self, config: DFineConfig, kernel_size: int, stride: int):
+ super().__init__()
+ self.conv1 = DFineConvNormLayer(config, config.encoder_hidden_dim, config.encoder_hidden_dim, 1, 1)
+ self.conv2 = DFineConvNormLayer(
+ config,
+ config.encoder_hidden_dim,
+ config.encoder_hidden_dim,
+ kernel_size,
+ stride,
+ config.encoder_hidden_dim,
+ )
+
+ def forward(self, input_features: torch.Tensor) -> torch.Tensor:
+ input_features = self.conv1(input_features)
+ input_features = self.conv2(input_features)
+ return input_features
+
+
+class DFineEncoderLayer(RTDetrEncoderLayer):
+ def __init__(self, config: DFineConfig):
+ super().__init__(config)
+ self.mlp = DFineMLP(
+ self.hidden_size, config.encoder_ffn_dim, self.hidden_size, 2, config.encoder_activation_function
+ )
+
+
+class DFineAIFILayer(RTDetrAIFILayer):
+ pass
+
+
+class DFineIntegral(nn.Module):
+ """
+ A static layer that calculates integral results from a distribution.
+
+ This layer computes the target location using the formula: `sum{Pr(n) * W(n)}`,
+ where Pr(n) is the softmax probability vector representing the discrete
+ distribution, and W(n) is the non-uniform Weighting Function.
+
+ Args:
+ max_num_bins (int): Max number of the discrete bins. Default is 32.
+ It can be adjusted based on the dataset or task requirements.
+ """
+
+ def __init__(self, config: DFineConfig):
+ super().__init__()
+ self.max_num_bins = config.max_num_bins
+
+ def forward(self, pred_corners: torch.Tensor, project: torch.Tensor) -> torch.Tensor:
+ batch_size, num_queries, _ = pred_corners.shape
+ pred_corners = F.softmax(pred_corners.reshape(-1, self.max_num_bins + 1), dim=1)
+ pred_corners = F.linear(pred_corners, project.to(pred_corners.device)).reshape(-1, 4)
+ pred_corners = pred_corners.reshape(batch_size, num_queries, -1)
+ return pred_corners
+
+
+class DFineLQE(nn.Module):
+ def __init__(self, config: DFineConfig):
+ super().__init__()
+ self.top_prob_values = config.top_prob_values
+ self.max_num_bins = config.max_num_bins
+ self.reg_conf = DFineMLP(4 * (self.top_prob_values + 1), config.lqe_hidden_dim, 1, config.lqe_layers)
+
+ def forward(self, scores: torch.Tensor, pred_corners: torch.Tensor) -> torch.Tensor:
+ batch_size, length, _ = pred_corners.size()
+ prob = F.softmax(pred_corners.reshape(batch_size, length, 4, self.max_num_bins + 1), dim=-1)
+ prob_topk, _ = prob.topk(self.top_prob_values, dim=-1)
+ stat = torch.cat([prob_topk, prob_topk.mean(dim=-1, keepdim=True)], dim=-1)
+ quality_score = self.reg_conf(stat.reshape(batch_size, length, -1))
+ scores = scores + quality_score
+ return scores
+
+
+class DFineDecoderLayer(RTDetrDecoderLayer):
+ def __init__(self, config: DFineConfig):
+ super().__init__(config)
+
+ # override the encoder attention module with d-fine version
+ self.encoder_attn = DFineMultiscaleDeformableAttention(config=config)
+ # gate
+ self.gateway = DFineGate(config.d_model)
+ self.mlp = DFineMLP(
+ self.hidden_size, config.decoder_ffn_dim, self.hidden_size, 2, config.decoder_activation_function
+ )
+
+ del self.encoder_attn_layer_norm
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: torch.Tensor | None = None,
+ reference_points=None,
+ spatial_shapes=None,
+ spatial_shapes_list=None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=encoder_attention_mask,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ residual = hidden_states
+
+ # Cross-Attention
+ hidden_states = hidden_states if position_embeddings is None else hidden_states + position_embeddings
+ hidden_states, _ = self.encoder_attn(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ reference_points=reference_points,
+ spatial_shapes=spatial_shapes,
+ spatial_shapes_list=spatial_shapes_list,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = self.gateway(residual, hidden_states)
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ hidden_states = self.final_layer_norm(hidden_states.clamp(min=-65504, max=65504))
+
+ return hidden_states
+
+
+class DFineMLPPredictionHead(RTDetrMLPPredictionHead):
+ pass
+
+
+class DFinePreTrainedModel(RTDetrPreTrainedModel):
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ PreTrainedModel._init_weights(self, module)
+ # initialize linear layer bias value according to a given probability value.
+ if isinstance(module, (DFineForObjectDetection, DFineDecoder)):
+ if module.class_embed is not None:
+ for layer in module.class_embed:
+ prior_prob = self.config.initializer_bias_prior_prob or 1 / (self.config.num_labels + 1)
+ bias = float(-math.log((1 - prior_prob) / prior_prob))
+ init.xavier_uniform_(layer.weight)
+ init.constant_(layer.bias, bias)
+
+ if module.bbox_embed is not None:
+ for layer in module.bbox_embed:
+ init.constant_(layer.layers[-1].weight, 0)
+ init.constant_(layer.layers[-1].bias, 0)
+
+ if hasattr(module, "reg_scale"):
+ init.constant_(module.reg_scale, self.config.reg_scale)
+
+ if hasattr(module, "up"):
+ init.constant_(module.up, self.config.up)
+
+ if isinstance(module, DFineMultiscaleDeformableAttention):
+ init.constant_(module.sampling_offsets.weight, 0.0)
+ default_dtype = torch.get_default_dtype()
+ thetas = torch.arange(module.n_heads, dtype=torch.int64).to(default_dtype) * (
+ 2.0 * math.pi / module.n_heads
+ )
+ grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+ grid_init = grid_init / grid_init.abs().max(-1, keepdim=True).values
+ grid_init = grid_init.reshape(module.n_heads, 1, 2).tile([1, sum(module.num_points_list), 1])
+ scaling = torch.concat([torch.arange(1, n + 1) for n in module.num_points_list]).reshape(1, -1, 1)
+ grid_init *= scaling
+ init.copy_(module.sampling_offsets.bias, grid_init.flatten())
+
+ init.constant_(module.attention_weights.weight, 0.0)
+ init.constant_(module.attention_weights.bias, 0.0)
+
+ num_points_scale = [1 / n for n in module.num_points_list for _ in range(n)]
+ init.copy_(module.num_points_scale, torch.tensor(num_points_scale, dtype=torch.float32))
+
+ if isinstance(module, DFineModel):
+ prior_prob = self.config.initializer_bias_prior_prob or 1 / (self.config.num_labels + 1)
+ bias = float(-math.log((1 - prior_prob) / prior_prob))
+ init.xavier_uniform_(module.enc_score_head.weight)
+ init.constant_(module.enc_score_head.bias, bias)
+
+ if isinstance(module, DFineGate):
+ bias = float(-math.log((1 - 0.5) / 0.5))
+ init.constant_(module.gate.bias, bias)
+ init.constant_(module.gate.weight, 0)
+
+ if isinstance(module, DFineLQE):
+ init.constant_(module.reg_conf.layers[-1].bias, 0)
+ init.constant_(module.reg_conf.layers[-1].weight, 0)
+
+ if hasattr(module, "weight_embedding") and self.config.learn_initial_query:
+ init.xavier_uniform_(module.weight_embedding.weight)
+ if hasattr(module, "denoising_class_embed") and self.config.num_denoising > 0:
+ init.xavier_uniform_(module.denoising_class_embed.weight)
+
+
+class DFineHybridEncoder(RTDetrHybridEncoder):
+ def __init__(self, config: DFineConfig):
+ DFinePreTrainedModel.__init__(config)
+ self.config = config
+ self.in_channels = config.encoder_in_channels
+ self.num_fpn_stages = len(self.in_channels) - 1
+ self.feat_strides = config.feat_strides
+ self.encoder_hidden_dim = config.encoder_hidden_dim
+ self.encode_proj_layers = config.encode_proj_layers
+ self.positional_encoding_temperature = config.positional_encoding_temperature
+ self.eval_size = config.eval_size
+ self.out_channels = [self.encoder_hidden_dim for _ in self.in_channels]
+ self.out_strides = self.feat_strides
+
+ # AIFI (Attention-based Intra-scale Feature Interaction) layers
+ self.aifi = nn.ModuleList([DFineAIFILayer(config) for _ in range(len(self.encode_proj_layers))])
+
+ # top-down fpn
+ self.lateral_convs = nn.ModuleList()
+ self.fpn_blocks = nn.ModuleList()
+ for _ in range(len(self.in_channels) - 1, 0, -1):
+ lateral_layer = DFineConvNormLayer(config, self.encoder_hidden_dim, self.encoder_hidden_dim, 1, 1)
+ self.lateral_convs.append(lateral_layer)
+ num_blocks = round(3 * config.depth_mult)
+ fpn_layer = DFineRepNCSPELAN4(config, numb_blocks=num_blocks)
+ self.fpn_blocks.append(fpn_layer)
+
+ # bottom-up pan
+ self.downsample_convs = nn.ModuleList()
+ self.pan_blocks = nn.ModuleList()
+ for _ in range(len(self.in_channels) - 1):
+ self.downsample_convs.append(DFineSCDown(config, 3, 2))
+ num_blocks = round(3 * config.depth_mult)
+ self.pan_blocks.append(DFineRepNCSPELAN4(config, numb_blocks=num_blocks))
+
+ self.post_init()
+
+
+class DFineDecoder(RTDetrDecoder):
+ """
+ D-FINE Decoder implementing Fine-grained Distribution Refinement (FDR).
+
+ This decoder refines object detection predictions through iterative updates across multiple layers,
+ utilizing attention mechanisms, location quality estimators, and distribution refinement techniques
+ to improve bounding box accuracy and robustness.
+ """
+
+ def __init__(self, config: DFineConfig):
+ self.eval_idx = config.eval_idx if config.eval_idx >= 0 else config.decoder_layers + config.eval_idx
+ super().__init__(config=config)
+ self.reg_scale = nn.Parameter(torch.tensor([config.reg_scale]), requires_grad=False)
+ self.max_num_bins = config.max_num_bins
+ self.d_model = config.d_model
+ self.layer_scale = config.layer_scale
+ self.pre_bbox_head = DFineMLP(config.hidden_size, config.hidden_size, 4, 3)
+ self.integral = DFineIntegral(config)
+ self.num_head = config.decoder_attention_heads
+ self.up = nn.Parameter(torch.tensor([config.up]), requires_grad=False)
+ self.lqe_layers = nn.ModuleList([DFineLQE(config) for _ in range(config.decoder_layers)])
+ self.layers = nn.ModuleList(
+ [DFineDecoderLayer(config) for _ in range(config.decoder_layers)]
+ + [DFineDecoderLayer(config) for _ in range(config.decoder_layers - self.eval_idx - 1)]
+ )
+
+ def forward(
+ self,
+ encoder_hidden_states: torch.Tensor,
+ reference_points: torch.Tensor,
+ inputs_embeds: torch.Tensor,
+ spatial_shapes,
+ level_start_index=None,
+ spatial_shapes_list=None,
+ encoder_attention_mask=None,
+ memory_mask=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> DFineDecoderOutput:
+ if inputs_embeds is not None:
+ hidden_states = inputs_embeds
+
+ # decoder layers
+ intermediate = ()
+ intermediate_reference_points = ()
+ intermediate_logits = ()
+ intermediate_predicted_corners = ()
+ initial_reference_points = ()
+
+ output_detach = pred_corners_undetach = 0
+
+ project = weighting_function(self.max_num_bins, self.up, self.reg_scale)
+ ref_points_detach = F.sigmoid(reference_points)
+
+ for i, decoder_layer in enumerate(self.layers):
+ ref_points_input = ref_points_detach.unsqueeze(2)
+ query_pos_embed = self.query_pos_head(ref_points_detach).clamp(min=-10, max=10)
+
+ hidden_states = decoder_layer(
+ hidden_states,
+ position_embeddings=query_pos_embed,
+ reference_points=ref_points_input,
+ spatial_shapes=spatial_shapes,
+ spatial_shapes_list=spatial_shapes_list,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ **kwargs,
+ )
+
+ if i == 0:
+ # Initial bounding box predictions with inverse sigmoid refinement
+ new_reference_points = F.sigmoid(
+ self.pre_bbox_head(hidden_states) + inverse_sigmoid(ref_points_detach)
+ )
+ ref_points_initial = new_reference_points.detach()
+
+ # Refine bounding box corners using FDR, integrating previous layer's corrections
+ if self.bbox_embed is not None:
+ pred_corners = self.bbox_embed[i](hidden_states + output_detach) + pred_corners_undetach
+ inter_ref_bbox = distance2bbox(
+ ref_points_initial, self.integral(pred_corners, project), self.reg_scale
+ )
+ pred_corners_undetach = pred_corners
+ ref_points_detach = inter_ref_bbox.detach()
+
+ output_detach = hidden_states.detach()
+
+ intermediate += (hidden_states,)
+
+ if self.class_embed is not None and (self.training or i == self.eval_idx):
+ scores = self.class_embed[i](hidden_states)
+ # Add initial logits and reference points with pre-bbox head
+ if i == 0:
+ intermediate_logits += (scores,)
+ intermediate_reference_points += (new_reference_points,)
+ # Lqe does not affect the performance here.
+ scores = self.lqe_layers[i](scores, pred_corners)
+ intermediate_logits += (scores,)
+ intermediate_reference_points += (inter_ref_bbox,)
+ initial_reference_points += (ref_points_initial,)
+ intermediate_predicted_corners += (pred_corners,)
+
+ # Keep batch_size as first dimension
+ intermediate = torch.stack(intermediate)
+ if self.class_embed is not None and self.bbox_embed is not None:
+ intermediate_logits = torch.stack(intermediate_logits, dim=1)
+ intermediate_predicted_corners = torch.stack(intermediate_predicted_corners, dim=1)
+ initial_reference_points = torch.stack(initial_reference_points, dim=1)
+ intermediate_reference_points = torch.stack(intermediate_reference_points, dim=1)
+
+ return DFineDecoderOutput(
+ last_hidden_state=hidden_states,
+ intermediate_hidden_states=intermediate,
+ intermediate_logits=intermediate_logits,
+ intermediate_reference_points=intermediate_reference_points,
+ intermediate_predicted_corners=intermediate_predicted_corners,
+ initial_reference_points=initial_reference_points,
+ )
+
+
+class DFineModel(RTDetrModel):
+ def __init__(self, config: DFineConfig):
+ super().__init__(config)
+ del self.decoder_input_proj
+ self.encoder = DFineHybridEncoder(config=config)
+ num_backbone_outs = len(config.decoder_in_channels)
+ decoder_input_proj = []
+ in_channels = config.decoder_in_channels[-1]
+ for _ in range(num_backbone_outs):
+ if config.hidden_size == config.decoder_in_channels[-1]:
+ decoder_input_proj.append(nn.Identity())
+ else:
+ conv = nn.Conv2d(in_channels, config.d_model, kernel_size=1, bias=False)
+ batchnorm = nn.BatchNorm2d(config.d_model, config.batch_norm_eps)
+ decoder_input_proj.append(nn.Sequential(conv, batchnorm))
+ for _ in range(config.num_feature_levels - num_backbone_outs):
+ if config.hidden_size == config.decoder_in_channels[-1]:
+ decoder_input_proj.append(nn.Identity())
+ else:
+ conv = nn.Conv2d(in_channels, config.d_model, kernel_size=3, stride=2, padding=1, bias=False)
+ batchnorm = nn.BatchNorm2d(config.d_model, config.batch_norm_eps)
+ decoder_input_proj.append(nn.Sequential(conv, batchnorm))
+ self.decoder_input_proj = nn.ModuleList(decoder_input_proj)
+ self.decoder = DFineDecoder(config)
+
+
+class DFineForObjectDetection(RTDetrForObjectDetection):
+ # When using clones, all layers > 0 will be clones, but layer 0 *is* required
+ # We can't initialize the model on meta device as some weights are modified during the initialization
+ _no_split_modules = None
+ _tied_weights_keys = {
+ r"bbox_embed.(?![0])\d+": r"bbox_embed.0",
+ r"class_embed.(?![0])\d+": r"^class_embed.0",
+ "class_embed": "model.decoder.class_embed",
+ "bbox_embed": "model.decoder.bbox_embed",
+ }
+
+ def __init__(self, config: DFineConfig):
+ DFinePreTrainedModel.__init__(self, config)
+
+ # D-FINE encoder-decoder model
+ self.eval_idx = config.eval_idx if config.eval_idx >= 0 else config.decoder_layers + config.eval_idx
+ self.model = DFineModel(config)
+ scaled_dim = round(config.layer_scale * config.hidden_size)
+ num_pred = config.decoder_layers
+ self.class_embed = nn.ModuleList([nn.Linear(config.d_model, config.num_labels) for _ in range(num_pred)])
+ self.bbox_embed = nn.ModuleList(
+ [
+ DFineMLP(config.hidden_size, config.hidden_size, 4 * (config.max_num_bins + 1), 3)
+ for _ in range(self.eval_idx + 1)
+ ]
+ + [
+ DFineMLP(scaled_dim, scaled_dim, 4 * (config.max_num_bins + 1), 3)
+ for _ in range(config.decoder_layers - self.eval_idx - 1)
+ ]
+ )
+
+ self.model.decoder.class_embed = self.class_embed
+ self.model.decoder.bbox_embed = self.bbox_embed
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def forward(**super_kwargs):
+ r"""
+ Example:
+
+ ```python
+ >>> import torch
+ >>> from transformers.image_utils import load_image
+ >>> from transformers import AutoImageProcessor, DFineForObjectDetection
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = load_image(url)
+
+ >>> image_processor = AutoImageProcessor.from_pretrained("ustc-community/dfine-xlarge-coco")
+ >>> model = DFineForObjectDetection.from_pretrained("ustc-community/dfine-xlarge-coco")
+
+ >>> # prepare image for the model
+ >>> inputs = image_processor(images=image, return_tensors="pt")
+
+ >>> # forward pass
+ >>> outputs = model(**inputs)
+
+ >>> logits = outputs.logits
+ >>> list(logits.shape)
+ [1, 300, 80]
+
+ >>> boxes = outputs.pred_boxes
+ >>> list(boxes.shape)
+ [1, 300, 4]
+
+ >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
+ >>> target_sizes = torch.tensor([image.size[::-1]])
+ >>> results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)
+ >>> result = results[0] # first image in batch
+
+ >>> for score, label, box in zip(result["scores"], result["labels"], result["boxes"]):
+ ... box = [round(i, 2) for i in box.tolist()]
+ ... print(
+ ... f"Detected {model.config.id2label[label.item()]} with confidence "
+ ... f"{round(score.item(), 3)} at location {box}"
+ ... )
+ Detected cat with confidence 0.958 at location [344.49, 23.4, 639.84, 374.27]
+ Detected cat with confidence 0.956 at location [11.71, 53.52, 316.64, 472.33]
+ Detected remote with confidence 0.947 at location [40.46, 73.7, 175.62, 117.57]
+ Detected sofa with confidence 0.918 at location [0.59, 1.88, 640.25, 474.74]
+ ```
+ """
+ super().forward(**super_kwargs)
+
+
+__all__ = [
+ "DFineConfig",
+ "DFineModel",
+ "DFinePreTrainedModel",
+ "DFineForObjectDetection",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..39a2e4ce07d9e197920f97b428ede3b77deb7eb6
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/__init__.py
@@ -0,0 +1,32 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_got_ocr2 import *
+ from .image_processing_got_ocr2 import *
+ from .image_processing_pil_got_ocr2 import *
+ from .modeling_got_ocr2 import *
+ from .processing_got_ocr2 import *
+
+
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/configuration_got_ocr2.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/configuration_got_ocr2.py
new file mode 100644
index 0000000000000000000000000000000000000000..66ef97abfdc3b1a20af26a65ce93e911a74e66b5
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/configuration_got_ocr2.py
@@ -0,0 +1,132 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/got_ocr2/modular_got_ocr2.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_got_ocr2.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2024 HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+@auto_docstring(checkpoint="facebook/sam-vit-huge")
+@strict
+class GotOcr2VisionConfig(PreTrainedConfig):
+ r"""
+ output_channels (`int`, *optional*, defaults to 256):
+ Dimensionality of the output channels in the Patch Encoder.
+ use_abs_pos (`bool`, *optional*, defaults to `True`):
+ Whether to use absolute position embedding.
+ use_rel_pos (`bool`, *optional*, defaults to `True`):
+ Whether to use relative position embedding.
+ window_size (`int`, *optional*, defaults to 14):
+ Window size for relative position.
+ global_attn_indexes (`list[int]`, *optional*, defaults to `[2, 5, 8, 11]`):
+ The indexes of the global attention layers.
+ mlp_dim (`int`, *optional*, defaults to 3072):
+ The dimensionality of the MLP layer in the Transformer encoder.
+ """
+
+ base_config_key = "vision_config"
+ hidden_size: int = 768
+ output_channels: int = 256
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ num_channels: int = 3
+ image_size: int | list[int] | tuple[int, int] = 1024
+ patch_size: int | list[int] | tuple[int, int] = 16
+ hidden_act: str = "gelu"
+ layer_norm_eps: float = 1e-06
+ attention_dropout: float | int = 0.0
+ initializer_range: float = 1e-10
+ qkv_bias: bool = True
+ use_abs_pos: bool = True
+ use_rel_pos: bool = True
+ window_size: int = 14
+ global_attn_indexes: list[int] | tuple[int, ...] = (2, 5, 8, 11)
+ mlp_dim: int = 3072
+
+
+@auto_docstring(checkpoint="facebook/sam-vit-huge")
+@strict
+class GotOcr2Config(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import GotOcr2ForConditionalGeneration, GotOcr2Config
+
+ >>> # Initializing a GotOcr2 style configuration
+ >>> configuration = GotOcr2Config()
+
+ >>> # Initializing a model from the Qwen2-VL-7B style configuration
+ >>> model = GotOcr2ForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "got_ocr2"
+ attribute_map = {
+ "image_token_id": "image_token_index",
+ }
+ sub_configs = {"text_config": AutoConfig, "vision_config": GotOcr2VisionConfig}
+
+ vision_config: dict | PreTrainedConfig | None = None
+ text_config: dict | PreTrainedConfig | None = None
+ image_token_index: int = 151859
+ image_seq_length: int = 576
+ tie_word_embeddings: bool = True
+
+ def __post_init__(self, **kwargs):
+ if self.vision_config is None:
+ self.vision_config = GotOcr2VisionConfig()
+ elif isinstance(self.vision_config, dict):
+ self.vision_config = GotOcr2VisionConfig(**self.vision_config)
+
+ if isinstance(self.text_config, dict):
+ self.text_config["model_type"] = self.text_config.get("model_type", "qwen2")
+ self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config)
+ elif self.text_config is None:
+ self.text_config = CONFIG_MAPPING["qwen2"](
+ vocab_size=151860,
+ hidden_size=1024,
+ intermediate_size=2816,
+ num_hidden_layers=24,
+ num_attention_heads=16,
+ num_key_value_heads=16,
+ hidden_act="silu",
+ max_position_embeddings=32768,
+ initializer_range=0.02,
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ tie_word_embeddings=self.tie_word_embeddings,
+ rope_theta=1000000.0,
+ rope_parameters=None,
+ use_sliding_window=False,
+ sliding_window=4096,
+ max_window_layers=21,
+ attention_dropout=0.0,
+ )
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["GotOcr2VisionConfig", "GotOcr2Config"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/image_processing_got_ocr2.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/image_processing_got_ocr2.py
new file mode 100644
index 0000000000000000000000000000000000000000..d035921873db5a08cd306fee5fc3ec50b3c723eb
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/image_processing_got_ocr2.py
@@ -0,0 +1,307 @@
+# Copyright 2025 HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Image processor class for Got-OCR-2."""
+
+from functools import lru_cache
+
+import torch
+from torchvision.transforms.v2 import functional as tvF
+
+from ...image_processing_backends import TorchvisionBackend
+from ...image_processing_utils import BatchFeature
+from ...image_transforms import group_images_by_shape, reorder_images
+from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling, SizeDict
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import (
+ TensorType,
+ auto_docstring,
+)
+
+
+class GotOcr2ImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ crop_to_patches (`bool`, *optional*, defaults to `self.crop_to_patches`):
+ Whether to crop the image to patches. Can be overridden by the `crop_to_patches` parameter in the
+ `preprocess` method.
+ min_patches (`int`, *optional*, defaults to `self.min_patches`):
+ The minimum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is
+ set to `True`. Can be overridden by the `min_patches` parameter in the `preprocess` method.
+ max_patches (`int`, *optional*, defaults to `self.max_patches`):
+ The maximum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is
+ set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method.
+ """
+
+ crop_to_patches: bool
+ min_patches: int
+ max_patches: int
+
+
+@lru_cache(maxsize=10)
+def get_all_supported_aspect_ratios(min_image_tiles: int, max_image_tiles: int) -> list[tuple[int, int]]:
+ """
+ Computes all allowed aspect ratios for a given minimum and maximum number of input tiles.
+
+ This function calculates all possible arrangements of tiles that can be formed
+ within the constraint of the minimum and maximum number of tiles. Each arrangement is
+ represented by its aspect ratio (width/height) and the corresponding tile configuration.
+
+ Args:
+ min_image_tiles (`int`):
+ The minimum number of tiles allowed.
+ max_image_tiles (`int`):
+ The maximum number of tiles allowed.
+
+ Returns:
+ `list[tuple[int, int]]`: A list of tuples, each tuple representing a valid (width, height)
+ configuration in terms of number of tiles.
+
+ Example:
+ >>> get_all_supported_aspect_ratios(1, 4)
+ [(1, 1), (1, 2), (2, 1), (1, 3), (3, 1), (1, 4), (2, 2), (4, 1)]
+
+ """
+ aspect_ratios = []
+ for width in range(1, max_image_tiles + 1):
+ for height in range(1, max_image_tiles + 1):
+ if width * height <= max_image_tiles and width * height >= min_image_tiles:
+ aspect_ratios.append((width, height))
+
+ aspect_ratios = sorted(aspect_ratios, key=lambda x: x[0] * x[1])
+
+ return aspect_ratios
+
+
+@lru_cache(maxsize=100)
+def get_optimal_tiled_canvas(
+ original_image_size: tuple[int, int],
+ target_tile_size: tuple[int, int],
+ min_image_tiles: int,
+ max_image_tiles: int,
+) -> tuple[int, int]:
+ """
+ Given a minimum and maximum number of tiles, find the canvas with the closest aspect ratio to the
+ original image aspect ratio.
+ In case of tie-breaking condition when two canvases have the same aspect ratio difference, we favor the canvas with
+ more tiles, until the area covered by the tiles is more than twice the target area, in order to avoid unnecessarily
+ excessive tiling.
+ """
+ possible_tile_arrangements = get_all_supported_aspect_ratios(min_image_tiles, max_image_tiles)
+
+ original_height, original_width = original_image_size
+ target_tile_height, target_tile_width = target_tile_size
+ aspect_ratio = original_width / original_height
+ area = original_width * original_height
+
+ # find the grid with the best aspect ratio
+ best_ratio_diff = float("inf")
+ best_grid = (1, 1)
+ for grid in possible_tile_arrangements:
+ grid_aspect_ratio = grid[0] / grid[1]
+ ratio_diff = abs(aspect_ratio - grid_aspect_ratio)
+ if ratio_diff < best_ratio_diff:
+ best_ratio_diff = ratio_diff
+ best_grid = grid
+ elif ratio_diff == best_ratio_diff:
+ # if the aspect ratio difference is the same, we favor the grid with more patches
+ # until the area covered by the patches is more than twice the original image area
+ if area > 0.5 * target_tile_height * target_tile_width * grid[0] * grid[1]:
+ best_grid = grid
+
+ return best_grid
+
+
+@auto_docstring
+class GotOcr2ImageProcessor(TorchvisionBackend):
+ valid_kwargs = GotOcr2ImageProcessorKwargs
+ resample = PILImageResampling.BICUBIC
+ image_mean = OPENAI_CLIP_MEAN
+ image_std = OPENAI_CLIP_STD
+ size = {"height": 384, "width": 384}
+ do_resize = True
+ do_rescale = True
+ do_normalize = True
+ do_convert_rgb = True
+ crop_to_patches = False
+ min_patches = 1
+ max_patches = 12
+
+ def __init__(self, **kwargs: Unpack[GotOcr2ImageProcessorKwargs]):
+ super().__init__(**kwargs)
+
+ def crop_image_to_patches(
+ self,
+ images: "torch.Tensor",
+ min_patches: int,
+ max_patches: int,
+ use_thumbnail: bool = True,
+ patch_size: SizeDict | None = None,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None,
+ ):
+ """
+ Crop the images to patches and return a list of cropped images.
+ The number of patches and their grid arrangement are determined by the original image size,
+ the target patch size and the minimum and maximum number of patches.
+ The aspect ratio of the patches grid is chosen to be the closest to the original image aspect ratio.
+
+ Args:
+ images (`torch.Tensor`):
+ The images to be cropped.
+ min_patches (`int`):
+ The minimum number of patches to be extracted from the image.
+ max_patches (`int`):
+ The maximum number of patches to be extracted from the image.
+ use_thumbnail (`bool`, *optional*, defaults to `True`):
+ Whether to add a thumbnail image to the list of cropped patches.
+ patch_size (`SizeDict`, *optional*):
+ The size of the output patches.
+ resample (`PILImageResampling | tvF.InterpolationMode | int | None`, *optional*):
+ Resampling filter to use when resizing.
+ """
+ patch_size_height, patch_size_width = patch_size.height, patch_size.width
+ original_height, original_width = images.shape[-2:]
+ # find the closest aspect ratio to the target
+ num_columns, num_rows = get_optimal_tiled_canvas(
+ (original_height, original_width), (patch_size_height, patch_size_width), min_patches, max_patches
+ )
+
+ # calculate the target width and height
+ target_width = patch_size_width * num_columns
+ target_height = patch_size_height * num_rows
+ num_blocks = num_columns * num_rows
+
+ # resize the image so that each patch is of patch_size
+ resized_image = self.resize(images, SizeDict(height=target_height, width=target_width), resample=resample)
+ # split the image into patches
+ processed_images = []
+ for i in range(num_blocks):
+ column = i % num_columns
+ row = i // num_columns
+ box = (
+ column * patch_size_width,
+ row * patch_size_height,
+ (column + 1) * patch_size_width,
+ (row + 1) * patch_size_height,
+ )
+ # split the image
+ patch_image = resized_image[..., box[1] : box[3], box[0] : box[2]]
+ processed_images.append(patch_image)
+
+ if use_thumbnail and len(processed_images) != 1:
+ thumbnail_img = self.resize(images, patch_size, resample=resample)
+ processed_images.append(thumbnail_img)
+
+ processed_images = torch.stack(processed_images, dim=0).transpose(0, 1).contiguous()
+
+ return processed_images
+
+ def _preprocess(
+ self,
+ images: list["torch.Tensor"],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ disable_grouping: bool | None,
+ return_tensors: str | TensorType | None,
+ crop_to_patches: bool = False,
+ min_patches: int = 1,
+ max_patches: int = 12,
+ **kwargs,
+ ) -> BatchFeature:
+ if crop_to_patches:
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ processed_images_grouped = {}
+ num_patches = {}
+ for shape, stacked_images in grouped_images.items():
+ stacked_images = self.crop_image_to_patches(
+ stacked_images,
+ min_patches,
+ max_patches,
+ patch_size=size,
+ resample=resample,
+ )
+ processed_images_grouped[shape] = stacked_images
+ num_patches[shape] = [stacked_images.shape[1]] * stacked_images.shape[0]
+ images = reorder_images(processed_images_grouped, grouped_images_index)
+ images = [image for images_list in images for image in images_list]
+ num_patches = reorder_images(num_patches, grouped_images_index)
+ else:
+ num_patches = [1] * len(images)
+
+ # Group images by size for batched resizing
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ resized_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_resize:
+ stacked_images = self.resize(image=stacked_images, size=size, resample=resample)
+ resized_images_grouped[shape] = stacked_images
+ resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+ # Group images by size for further processing
+ # Needed in case do_resize is False, or resize returns images with different sizes
+ grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
+ processed_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ stacked_images = self.rescale_and_normalize(
+ stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+ )
+ processed_images_grouped[shape] = stacked_images
+
+ processed_images = reorder_images(processed_images_grouped, grouped_images_index)
+
+ return BatchFeature(
+ data={"pixel_values": processed_images, "num_patches": num_patches}, tensor_type=return_tensors
+ )
+
+ def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None):
+ """
+ A utility that returns number patches for a given image size.
+
+ Args:
+ height (`int`):
+ Height of the input image.
+ width (`int`):
+ Width of the input image.
+ images_kwargs (`dict`, *optional*)
+ Any kwargs to override defaults of the image processor.
+ Returns:
+ `int`: Number of patches per image.
+ """
+ min_patches = images_kwargs.get("min_patches", self.min_patches) if images_kwargs else self.min_patches
+ max_patches = images_kwargs.get("max_patches", self.max_patches) if images_kwargs else self.max_patches
+ patch_size = images_kwargs.get("patch_size", self.size) if images_kwargs else self.size
+ crop_to_patches = (
+ images_kwargs.get("crop_to_patches", self.crop_to_patches) if images_kwargs else self.crop_to_patches
+ )
+
+ num_patches = 1
+ if crop_to_patches and max_patches > 1:
+ if isinstance(patch_size, dict):
+ patch_height, patch_width = patch_size["height"], patch_size["width"]
+ else:
+ patch_height, patch_width = patch_size.height, patch_size.width
+ num_columns, num_rows = get_optimal_tiled_canvas(
+ (height, width), (patch_height, patch_width), min_patches, max_patches
+ )
+ if num_columns * num_rows > 1:
+ num_patches += num_columns * num_rows
+
+ return num_patches
+
+
+__all__ = ["GotOcr2ImageProcessor"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/modeling_got_ocr2.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/modeling_got_ocr2.py
new file mode 100644
index 0000000000000000000000000000000000000000..add27ed822a5a7e7d5025174fcd9703e36eb60f4
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/got_ocr2/modeling_got_ocr2.py
@@ -0,0 +1,771 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/got_ocr2/modular_got_ocr2.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_got_ocr2.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2024 HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import collections
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache
+from ...generation import GenerationMixin
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from ..auto import AutoModel
+from .configuration_got_ocr2 import GotOcr2Config, GotOcr2VisionConfig
+
+
+class GotOcr2MLPBlock(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.lin1 = nn.Linear(config.hidden_size, config.mlp_dim)
+ self.lin2 = nn.Linear(config.mlp_dim, config.hidden_size)
+ self.act = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.lin1(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.lin2(hidden_states)
+ return hidden_states
+
+
+class GotOcr2VisionAttention(nn.Module):
+ """Multi-head Attention block with relative position embeddings."""
+
+ def __init__(self, config, window_size):
+ super().__init__()
+ input_size = (
+ (config.image_size // config.patch_size, config.image_size // config.patch_size)
+ if window_size == 0
+ else (window_size, window_size)
+ )
+
+ self.num_attention_heads = config.num_attention_heads
+ head_dim = config.hidden_size // config.num_attention_heads
+ self.scale = head_dim**-0.5
+ self.dropout = config.attention_dropout
+
+ self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.qkv_bias)
+ self.proj = nn.Linear(config.hidden_size, config.hidden_size)
+
+ self.use_rel_pos = config.use_rel_pos
+ if self.use_rel_pos:
+ if input_size is None:
+ raise ValueError("Input size must be provided if using relative positional encoding.")
+
+ # initialize relative positional embeddings
+ self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))
+ self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))
+
+ def get_rel_pos(self, q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
+ """
+ Get relative positional embeddings according to the relative positions of
+ query and key sizes.
+
+ Args:
+ q_size (int):
+ size of the query.
+ k_size (int):
+ size of key k.
+ rel_pos (`torch.Tensor`):
+ relative position embeddings (L, channel).
+
+ Returns:
+ Extracted positional embeddings according to relative positions.
+ """
+ max_rel_dist = int(2 * max(q_size, k_size) - 1)
+ # Interpolate rel pos.
+ rel_pos_resized = F.interpolate(
+ rel_pos.reshape(1, rel_pos.shape[0], -1).transpose(1, 2),
+ size=max_rel_dist,
+ mode="linear",
+ )
+ rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
+
+ # Scale the coords with short length if shapes for q and k are different.
+ q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
+ k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
+ relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
+
+ return rel_pos_resized[relative_coords.long()]
+
+ def get_decomposed_rel_pos(
+ self,
+ query: torch.Tensor,
+ rel_pos_h: torch.Tensor,
+ rel_pos_w: torch.Tensor,
+ q_size: tuple[int, int],
+ k_size: tuple[int, int],
+ ) -> torch.Tensor:
+ """
+ Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
+ https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py
+
+ Args:
+ query (`torch.Tensor`):
+ query q in the attention layer with shape (batch_size, query_height * query_width, channel).
+ rel_pos_h (`torch.Tensor`):
+ relative position embeddings (Lh, channel) for height axis.
+ rel_pos_w (`torch.Tensor`):
+ relative position embeddings (Lw, channel) for width axis.
+ q_size (tuple):
+ spatial sequence size of query q with (query_height, query_width).
+ k_size (tuple):
+ spatial sequence size of key k with (key_height, key_width).
+
+ Returns:
+ decomposed_rel_pos (`torch.Tensor`):
+ decomposed relative position embeddings.
+ """
+ query_height, query_width = q_size
+ key_height, key_width = k_size
+ relative_position_height = self.get_rel_pos(query_height, key_height, rel_pos_h)
+ relative_position_width = self.get_rel_pos(query_width, key_width, rel_pos_w)
+
+ batch_size, _, dim = query.shape
+ reshaped_query = query.reshape(batch_size, query_height, query_width, dim)
+ rel_h = torch.einsum("bhwc,hkc->bhwk", reshaped_query, relative_position_height)
+ rel_w = torch.einsum("bhwc,wkc->bhwk", reshaped_query, relative_position_width)
+
+ decomposed_rel_pos = rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]
+
+ return decomposed_rel_pos
+
+ def forward(self, hidden_states: torch.Tensor, output_attentions=None) -> tuple[torch.Tensor, torch.Tensor]:
+ batch_size, height, width, _ = hidden_states.shape
+ # qkv with shape (3, batch_size, nHead, height * width, channel)
+ qkv = (
+ self.qkv(hidden_states)
+ .reshape(batch_size, height * width, 3, self.num_attention_heads, -1)
+ .permute(2, 0, 3, 1, 4)
+ )
+ # q, k, v with shape (batch_size * nHead, height * width, channel)
+ query, key, value = qkv.reshape(3, batch_size * self.num_attention_heads, height * width, -1).unbind(0)
+
+ attn_weights = (query * self.scale) @ key.transpose(-2, -1)
+
+ if self.use_rel_pos:
+ decomposed_rel_pos = self.get_decomposed_rel_pos(
+ query, self.rel_pos_h, self.rel_pos_w, (height, width), (height, width)
+ )
+ decomposed_rel_pos = decomposed_rel_pos.reshape_as(attn_weights)
+ attn_weights = attn_weights + decomposed_rel_pos
+
+ attn_weights = torch.nn.functional.softmax(attn_weights, dtype=torch.float32, dim=-1).to(query.dtype)
+
+ attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+
+ attn_output = (attn_probs @ value).reshape(batch_size, self.num_attention_heads, height, width, -1)
+ attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(batch_size, height, width, -1)
+
+ attn_output = self.proj(attn_output)
+ return attn_output, attn_weights
+
+
+class GotOcr2VisionLayer(GradientCheckpointingLayer):
+ def __init__(self, config, window_size):
+ super().__init__()
+ self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.attn = GotOcr2VisionAttention(config, window_size)
+ self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.mlp = GotOcr2MLPBlock(config)
+ self.window_size = window_size
+
+ def window_partition(self, hidden_states: torch.Tensor, window_size: int) -> tuple[torch.Tensor, tuple[int, int]]:
+ """
+ Args:
+ Partition into non-overlapping windows with padding if needed.
+ hidden_states (tensor): input tokens with [batch_size, height, width, channel]. window_size (int): window
+ size.
+
+ Returns:
+ windows: windows after partition with [batch_size * num_windows, window_size, window_size, channel].
+ (pad_height, pad_width): padded height and width before partition
+ """
+ batch_size, height, width, channel = hidden_states.shape
+
+ pad_h = (window_size - height % window_size) % window_size
+ pad_w = (window_size - width % window_size) % window_size
+ hidden_states = F.pad(hidden_states, (0, 0, 0, pad_w, 0, pad_h))
+ pad_height, pad_width = height + pad_h, width + pad_w
+
+ hidden_states = hidden_states.reshape(
+ batch_size, pad_height // window_size, window_size, pad_width // window_size, window_size, channel
+ )
+ windows = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(-1, window_size, window_size, channel)
+ return windows, (pad_height, pad_width)
+
+ def window_unpartition(
+ self, windows: torch.Tensor, window_size: int, padding_shape: tuple[int, int], original_shape: tuple[int, int]
+ ) -> torch.Tensor:
+ """
+ Args:
+ Window unpartition into original sequences and removing padding.
+ hidden_states (tensor):
+ input tokens with [batch_size * num_windows, window_size, window_size, channel].
+ window_size (int):
+ window size.
+ padding_shape (Tuple):
+ padded height and width (pad_height, pad_width).
+ original_shape (Tuple): original height and width (height, width) before padding.
+
+ Returns:
+ hidden_states: unpartitioned sequences with [batch_size, height, width, channel].
+ """
+ pad_height, pad_width = padding_shape
+ height, width = original_shape
+ batch_size = windows.shape[0] // (pad_height * pad_width // window_size // window_size)
+ hidden_states = windows.reshape(
+ batch_size, pad_height // window_size, pad_width // window_size, window_size, window_size, -1
+ )
+ hidden_states = (
+ hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(batch_size, pad_height, pad_width, -1)
+ )
+
+ hidden_states = hidden_states[:, :height, :width, :].contiguous()
+ return hidden_states
+
+ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.FloatTensor]:
+ residual = hidden_states
+ hidden_states = self.layer_norm1(hidden_states)
+ # Window partition
+ if self.window_size > 0:
+ height, width = hidden_states.shape[1], hidden_states.shape[2]
+ hidden_states, padding_shape = self.window_partition(hidden_states, self.window_size)
+
+ hidden_states, attn_weights = self.attn(
+ hidden_states=hidden_states,
+ )
+ # Reverse window partition
+ if self.window_size > 0:
+ hidden_states = self.window_unpartition(hidden_states, self.window_size, padding_shape, (height, width))
+
+ hidden_states = residual + hidden_states
+ layernorm_output = self.layer_norm2(hidden_states)
+ hidden_states = hidden_states + self.mlp(layernorm_output)
+ return hidden_states
+
+
+@auto_docstring
+class GotOcr2PreTrainedModel(PreTrainedModel):
+ config: GotOcr2Config
+ base_model_prefix = "model"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = True
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = False
+ _supports_sdpa = False
+
+ _can_compile_fullgraph = True
+ _supports_flex_attn = False
+ _supports_attention_backend = True
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, GotOcr2VisionAttention):
+ if module.use_rel_pos:
+ init.zeros_(module.rel_pos_h)
+ init.zeros_(module.rel_pos_w)
+ elif isinstance(module, GotOcr2VisionEncoder):
+ if module.pos_embed is not None:
+ init.zeros_(module.pos_embed)
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for got_ocr2 vision model's outputs that also contains image embeddings obtained by applying the projection
+ layer to the pooler_output.
+ """
+)
+@dataclass
+class GotOcr2VisionEncoderOutput(ModelOutput):
+ r"""
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
+ The image embeddings obtained by applying the projection layer to the pooler_output.
+ """
+
+ image_embeds: torch.FloatTensor | None = None
+ last_hidden_state: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
+class GotOcr2PatchEmbeddings(nn.Module):
+ """
+ This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
+ `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
+ Transformer.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ image_size, patch_size = config.image_size, config.patch_size
+ num_channels, hidden_size = config.num_channels, config.hidden_size
+ image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
+ patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
+ num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.num_channels = num_channels
+ self.num_patches = num_patches
+
+ self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
+
+ def forward(self, pixel_values):
+ batch_size, num_channels, height, width = pixel_values.shape
+ if num_channels != self.num_channels:
+ raise ValueError(
+ "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+ )
+ if height != self.image_size[0] or width != self.image_size[1]:
+ raise ValueError(
+ f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})."
+ )
+ embeddings = self.projection(pixel_values).permute(0, 2, 3, 1)
+ return embeddings
+
+
+class GotOcr2LayerNorm(nn.LayerNorm):
+ r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
+ The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height,
+ width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width).
+ """
+
+ def __init__(self, normalized_shape, *, eps=1e-6, data_format="channels_last", **kwargs):
+ super().__init__(normalized_shape, eps=eps, **kwargs)
+ if data_format not in ["channels_last", "channels_first"]:
+ raise NotImplementedError(f"Unsupported data format: {data_format}")
+ self.data_format = data_format
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels)
+ """
+ if self.data_format == "channels_first":
+ features = features.permute(0, 2, 3, 1)
+ features = super().forward(features)
+ features = features.permute(0, 3, 1, 2)
+ else:
+ features = super().forward(features)
+ return features
+
+
+class GotOcr2VisionNeck(nn.Module):
+ def __init__(self, config: GotOcr2VisionConfig):
+ super().__init__()
+ self.config = config
+
+ self.conv1 = nn.Conv2d(config.hidden_size, config.output_channels, kernel_size=1, bias=False)
+ self.layer_norm1 = GotOcr2LayerNorm(config.output_channels, data_format="channels_first")
+ self.conv2 = nn.Conv2d(config.output_channels, config.output_channels, kernel_size=3, padding=1, bias=False)
+ self.layer_norm2 = GotOcr2LayerNorm(config.output_channels, data_format="channels_first")
+
+ def forward(self, hidden_states):
+ hidden_states = hidden_states.permute(0, 3, 1, 2)
+ hidden_states = self.conv1(hidden_states)
+ hidden_states = self.layer_norm1(hidden_states)
+
+ hidden_states = self.conv2(hidden_states)
+ hidden_states = self.layer_norm2(hidden_states)
+ return hidden_states
+
+
+class GotOcr2VisionEncoder(GotOcr2PreTrainedModel):
+ _can_record_outputs = {"hidden_states": GotOcr2VisionLayer, "attentions": GotOcr2VisionAttention}
+ input_modalities = ("image",)
+
+ def __init__(self, config: GotOcr2VisionConfig):
+ super().__init__(config)
+ self.config = config
+ self.image_size = config.image_size
+ self.patch_embed = GotOcr2PatchEmbeddings(config)
+
+ self.pos_embed = None
+ if config.use_abs_pos:
+ # Initialize absolute positional embedding with pretrain image size.
+ self.pos_embed = nn.Parameter(
+ torch.zeros(
+ 1,
+ config.image_size // config.patch_size,
+ config.image_size // config.patch_size,
+ config.hidden_size,
+ )
+ )
+
+ self.layers = nn.ModuleList()
+ for i in range(config.num_hidden_layers):
+ layer = GotOcr2VisionLayer(
+ config,
+ window_size=config.window_size if i not in config.global_attn_indexes else 0,
+ )
+ self.layers.append(layer)
+
+ self.neck = GotOcr2VisionNeck(config)
+
+ self.gradient_checkpointing = False
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.patch_embed
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ def forward(
+ self, pixel_values: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs]
+ ) -> tuple | GotOcr2VisionEncoderOutput:
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ hidden_states = self.patch_embed(pixel_values)
+ if self.pos_embed is not None:
+ hidden_states = hidden_states + self.pos_embed
+ for layer_module in self.layers:
+ hidden_states = layer_module(hidden_states)
+ hidden_states = self.neck(hidden_states)
+ return GotOcr2VisionEncoderOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+class GotOcr2MultiModalProjector(nn.Module):
+ def __init__(self, config: GotOcr2Config):
+ super().__init__()
+ vision_output_channels = config.vision_config.output_channels
+ language_hidden_size = config.text_config.hidden_size
+ self.conv_upsampler1 = nn.Conv2d(
+ vision_output_channels, vision_output_channels * 2, kernel_size=3, stride=2, padding=1, bias=False
+ )
+ self.conv_upsampler2 = nn.Conv2d(
+ vision_output_channels * 2, language_hidden_size, kernel_size=3, stride=2, padding=1, bias=False
+ )
+ self.multimodal_projector = nn.Linear(language_hidden_size, language_hidden_size)
+
+ def forward(self, vision_embeddings: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.conv_upsampler1(vision_embeddings)
+ hidden_state = self.conv_upsampler2(hidden_state)
+ hidden_state = hidden_state.flatten(2).permute(0, 2, 1)
+ hidden_state = self.multimodal_projector(hidden_state)
+ return hidden_state
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for GotOcr2 causal language model (or autoregressive) outputs.
+ """
+)
+@dataclass
+class GotOcr2CausalLMOutputWithPast(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ image_hidden_states (`torch.FloatTensor`, *optional*):
+ A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+ image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ image_hidden_states: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for GotOcr2 outputs, with hidden states and attentions.
+ """
+)
+@dataclass
+class GotOcr2ModelOutputWithPast(BaseModelOutputWithPast):
+ r"""
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ image_hidden_states (`torch.FloatTensor`, *optional*):
+ A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+ image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+ """
+
+ image_hidden_states: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ The GotOcr2 model which consists of a vision backbone and a language model, without a language modeling head.
+ """
+)
+class GotOcr2Model(GotOcr2PreTrainedModel):
+ def __init__(self, config: GotOcr2Config):
+ super().__init__(config)
+ self.vision_tower = GotOcr2VisionEncoder(config.vision_config)
+
+ self.multi_modal_projector = GotOcr2MultiModalProjector(config)
+ self.language_model = AutoModel.from_config(config.text_config)
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring(
+ custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection."
+ )
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs)
+ last_hidden_state = image_outputs.last_hidden_state
+ image_outputs.pooler_output = self.multi_modal_projector(last_hidden_state)
+
+ return image_outputs
+
+ def get_placeholder_mask(
+ self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ n_image_features = image_features.shape[0] * image_features.shape[1]
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ torch_compilable_check(
+ inputs_embeds[special_image_mask].numel() == image_features.numel(),
+ f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}",
+ )
+ return special_image_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | GotOcr2ModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_features = self.get_image_features(
+ pixel_values=pixel_values.to(inputs_embeds.dtype), return_dict=True
+ ).pooler_output
+ image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+ special_image_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_features
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+ outputs = self.language_model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ return_dict=True,
+ **kwargs,
+ )
+
+ return GotOcr2ModelOutputWithPast(
+ last_hidden_state=outputs.last_hidden_state,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=image_features if pixel_values is not None else None,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The GOT_OCR2 model which consists of a vision backbone and a language model.
+ """
+)
+class GotOcr2ForConditionalGeneration(GotOcr2PreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+
+ def __init__(self, config: GotOcr2Config):
+ super().__init__(config)
+ self.model = GotOcr2Model(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+ self.post_init()
+
+ def get_output_embeddings(self) -> nn.Module:
+ return self.lm_head
+
+ @auto_docstring
+ def get_image_features(
+ self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> tuple | BaseModelOutputWithPooling:
+ return self.model.get_image_features(pixel_values=pixel_values, **kwargs)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | GotOcr2CausalLMOutputWithPast:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, GotOcr2ForConditionalGeneration, TextStreamer
+
+ >>> model = GotOcr2ForConditionalGeneration.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf").to("cuda")
+ >>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf")
+
+ >>> url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = processor(image, return_tensors="pt", color="green").to("cuda")
+
+ >>> # Generate
+ >>> streamer = TextStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
+ >>> generate_ids = model.generate(
+ ... **inputs,
+ ... do_sample=False,
+ ... tokenizer = processor.tokenizer,
+ ... stop_strings='<|im_end|>',
+ ... streamer=streamer,
+ ... max_new_tokens=4096,
+ ... )
+ "You should keep in mind what features from the module should be used, especially
+ when you're planning to sell a template."
+ ```"""
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ return_dict=True,
+ logits_to_keep=logits_to_keep,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+ )
+
+ return GotOcr2CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=outputs.image_hidden_states,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ inputs_embeds=None,
+ pixel_values=None,
+ attention_mask=None,
+ logits_to_keep=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ logits_to_keep=logits_to_keep,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ if is_first_iteration or not kwargs.get("use_cache", True):
+ # Pixel values are used only in the first iteration if available
+ # In subsequent iterations, they are already merged with text and cached
+ # NOTE: first iteration doesn't have to be prefill, it can be the first
+ # iteration with a question and cached system prompt (continue generate from cache)
+ model_inputs["pixel_values"] = pixel_values
+
+ return model_inputs
+
+
+__all__ = ["GotOcr2PreTrainedModel", "GotOcr2Model", "GotOcr2ForConditionalGeneration"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/moonshine/configuration_moonshine.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/moonshine/configuration_moonshine.py
new file mode 100644
index 0000000000000000000000000000000000000000..9cb7adbfe59428b92e2ee66282159f1bdd492640
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/moonshine/configuration_moonshine.py
@@ -0,0 +1,116 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/moonshine/modular_moonshine.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_moonshine.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="UsefulSensors/moonshine-tiny")
+@strict
+class MoonshineConfig(PreTrainedConfig):
+ r"""
+ encoder_num_key_value_heads (`int`, *optional*):
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+ `encoder_num_key_value_heads=encoder_num_attention_heads`, the model will use Multi Head Attention (MHA), if
+ `encoder_num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+ by meanpooling all the original heads within that group. For more details, check out [this
+ paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
+ `num_attention_heads`.
+ decoder_num_key_value_heads (`int`, *optional*):
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+ `decoder_num_key_value_heads=decoder_num_attention_heads`, the model will use Multi Head Attention (MHA), if
+ `decoder_num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+ by meanpooling all the original heads within that group. For more details, check out [this
+ paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
+ `decoder_num_attention_heads`.
+ pad_head_dim_to_multiple_of (`int`, *optional*):
+ Pad head dimension in encoder and decoder to the next multiple of this value. Necessary for using certain
+ optimized attention implementations.
+ encoder_hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function (function or string) in the encoder.
+ decoder_hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function (function or string) in the decoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import MoonshineModel, MoonshineConfig
+
+ >>> # Initializing a Moonshine style configuration
+ >>> configuration = MoonshineConfig().from_pretrained("UsefulSensors/moonshine-tiny")
+
+ >>> # Initializing a model from the configuration
+ >>> model = MoonshineModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "moonshine"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ attribute_map = {
+ "num_key_value_heads": "decoder_num_key_value_heads",
+ "num_attention_heads": "decoder_num_attention_heads",
+ "num_hidden_layers": "decoder_num_hidden_layers",
+ "hidden_act": "decoder_hidden_act",
+ }
+
+ vocab_size: int = 32768
+ hidden_size: int = 288
+ intermediate_size: int = 1152
+ encoder_num_hidden_layers: int = 6
+ decoder_num_hidden_layers: int = 6
+ encoder_num_attention_heads: int = 8
+ decoder_num_attention_heads: int = 8
+ encoder_num_key_value_heads: int | None = None
+ decoder_num_key_value_heads: int | None = None
+ pad_head_dim_to_multiple_of: int | None = None
+ encoder_hidden_act: str = "gelu"
+ decoder_hidden_act: str = "silu"
+ max_position_embeddings: int = 512
+ initializer_range: float = 0.02
+ decoder_start_token_id: int = 1
+ use_cache: bool = True
+ rope_parameters: RopeParameters | dict | None = None
+ is_encoder_decoder: bool = True
+ attention_bias: bool = False
+ attention_dropout: float | int = 0.0
+ bos_token_id: int | None = 1
+ eos_token_id: int | list[int] | None = 2
+ pad_token_id: int | None = None
+ tie_word_embeddings: bool = True
+
+ def __post_init__(self, **kwargs):
+ if self.encoder_num_key_value_heads is None:
+ self.encoder_num_key_value_heads = self.encoder_num_attention_heads
+
+ if self.decoder_num_key_value_heads is None:
+ self.decoder_num_key_value_heads = self.decoder_num_attention_heads
+
+ kwargs.setdefault("partial_rotary_factor", 0.9) # assign default for BC
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["MoonshineConfig"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/moonshine/modeling_moonshine.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/moonshine/modeling_moonshine.py
new file mode 100644
index 0000000000000000000000000000000000000000..c19a3f19b5745e7f9c337f0c76f45af0b2eaf0a7
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/moonshine/modeling_moonshine.py
@@ -0,0 +1,953 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/moonshine/modular_moonshine.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_moonshine.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernelized_func
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPast,
+ BaseModelOutputWithPastAndCrossAttentions,
+ Seq2SeqLMOutput,
+ Seq2SeqModelOutput,
+)
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
+from ...utils.generic import maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_moonshine import MoonshineConfig
+
+
+@auto_docstring(
+ custom_intro="""
+ Extends [~modeling_outputs.BaseModelOutput] to include the output attention mask since sequence length is not preserved in the model's forward.
+ """
+)
+@dataclass
+class MoonshineEncoderModelOutput(BaseModelOutput):
+ r"""
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices after sequence compression. Returned because the
+ sequence length may differ from the input sequence length. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ """
+
+ attention_mask: torch.Tensor | None = None
+
+
+class MoonshineEncoderMLP(nn.Module):
+ def __init__(self, config, hidden_act):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[hidden_act]
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+class MoonshineDecoderMLP(nn.Module):
+ def __init__(self, config, hidden_act):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[hidden_act]
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size * 2)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states, gate = hidden_states.chunk(2, dim=-1)
+ hidden_states = self.activation_fn(gate) * hidden_states
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+class MoonshineRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: MoonshineConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: MoonshineConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+ """
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., 0::2]
+ x2 = x[..., 1::2]
+ return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+
+ # Interleave them instead of usual shape
+ cos = cos[..., : cos.shape[-1] // 2].repeat_interleave(2, dim=-1)
+ sin = sin[..., : sin.shape[-1] // 2].repeat_interleave(2, dim=-1)
+
+ # Keep half or full tensor for later concatenation
+ rotary_dim = cos.shape[-1]
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+ # Apply rotary embeddings on the first half or full tensor
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+ # Concatenate back to full shape
+ q_embed = torch.cat([q_embed, q_pass], dim=-1)
+ k_embed = torch.cat([k_embed, k_pass], dim=-1)
+ return q_embed, k_embed
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class MoonshineAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ config: MoonshineConfig,
+ layer_idx: int,
+ is_causal: bool,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ ):
+ super().__init__()
+ config.update({"num_attention_heads": num_attention_heads, "num_key_value_heads": num_key_value_heads})
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = is_causal
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+
+ # Pad head dimension to the next specified multiple.
+ if self.config.pad_head_dim_to_multiple_of is not None:
+ target_multiple = self.config.pad_head_dim_to_multiple_of
+ target_head_dim = target_multiple * ((self.head_dim + target_multiple - 1) // target_multiple)
+ self.head_dim_padding = target_head_dim - self.head_dim
+ else:
+ self.head_dim_padding = 0
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ key_value_states: torch.Tensor | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+ bsz, q_len = hidden_states.shape[:-1]
+
+ query_states = (
+ self.q_proj(hidden_states).view(bsz, q_len, self.config.num_key_value_heads, self.head_dim).transpose(1, 2)
+ )
+
+ is_cross_attention = key_value_states is not None
+ if past_key_values is not None:
+ is_updated = past_key_values.is_updated.get(self.layer_idx)
+ if is_cross_attention:
+ # after the first generated id, we can subsequently re-use all key/value_states from cache
+ past_key_values.is_updated[self.layer_idx] = True
+ past_key_values = past_key_values.cross_attention_cache
+ else:
+ past_key_values = past_key_values.self_attention_cache
+
+ # use key_value_states if cross attention
+ current_states = key_value_states if key_value_states is not None else hidden_states
+ if is_cross_attention and past_key_values and is_updated:
+ key_states = past_key_values.layers[self.layer_idx].keys
+ value_states = past_key_values.layers[self.layer_idx].values
+ else:
+ key_states = (
+ self.k_proj(current_states)
+ .view(bsz, -1, self.config.num_key_value_heads, self.head_dim)
+ .transpose(1, 2)
+ )
+ value_states = (
+ self.v_proj(current_states)
+ .view(bsz, -1, self.config.num_key_value_heads, self.head_dim)
+ .transpose(1, 2)
+ )
+ if is_cross_attention and past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ if not is_cross_attention:
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ is_causal = self.is_causal and attention_mask is None and q_len > 1
+
+ if self.head_dim_padding > 0:
+ query_states = torch.nn.functional.pad(query_states, (0, self.head_dim_padding))
+ key_states = torch.nn.functional.pad(key_states, (0, self.head_dim_padding))
+ value_states = torch.nn.functional.pad(value_states, (0, self.head_dim_padding))
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ is_causal=is_causal,
+ **kwargs,
+ )
+
+ if self.head_dim_padding > 0:
+ attn_output = attn_output[..., : -self.head_dim_padding]
+
+ attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class MoonshineEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: MoonshineConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = MoonshineAttention(
+ config=config,
+ layer_idx=layer_idx,
+ is_causal=False,
+ num_attention_heads=config.encoder_num_attention_heads,
+ num_key_value_heads=config.encoder_num_key_value_heads,
+ )
+
+ self.mlp = MoonshineEncoderMLP(config, config.encoder_hidden_act)
+ self.input_layernorm = nn.LayerNorm(config.hidden_size, bias=False)
+ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, bias=False)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+class MoonshineDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: MoonshineConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = MoonshineAttention(
+ config=config,
+ layer_idx=layer_idx,
+ is_causal=True,
+ num_attention_heads=config.num_attention_heads,
+ num_key_value_heads=config.num_key_value_heads,
+ )
+ self.encoder_attn = MoonshineAttention(
+ config=config,
+ layer_idx=layer_idx,
+ is_causal=False,
+ num_attention_heads=config.num_attention_heads,
+ num_key_value_heads=config.num_key_value_heads,
+ )
+
+ self.mlp = MoonshineDecoderMLP(config, config.hidden_act)
+ self.input_layernorm = nn.LayerNorm(config.hidden_size, bias=False)
+ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, bias=False)
+ self.final_layernorm = nn.LayerNorm(config.hidden_size, bias=False)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ encoder_position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ encoder_position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ if encoder_hidden_states is not None:
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states, _ = self.encoder_attn(
+ hidden_states=hidden_states,
+ key_value_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.final_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+@auto_docstring
+class MoonshinePreTrainedModel(PreTrainedModel):
+ config: MoonshineConfig
+ base_model_prefix = "model"
+ main_input_name = "input_values"
+ input_modalities = "audio"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["MoonshineEncoderLayer", "MoonshineDecoderLayer"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+
+ _can_compile_fullgraph = True
+ # TODO arthur, how do we separate when it cross / self coming from different layer?
+
+ def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
+ """
+ Computes the output length of the convolutional layers
+ """
+ output_conv1_length = int((input_lengths - 127) / 64 + 1)
+ output_conv2_length = int((output_conv1_length - 7) / 3 + 1)
+ output_conv3_length = int((output_conv2_length - 3) / 2 + 1)
+
+ return output_conv3_length
+
+
+class MoonshineEncoder(MoonshinePreTrainedModel):
+ """
+ Transformer encoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MoonshineEncoderLayer`]
+
+ Args:
+ config: MoonshineConfig
+ """
+
+ main_input_name = "input_values"
+ _can_record_outputs = {
+ "attentions": MoonshineAttention,
+ "hidden_states": MoonshineEncoderLayer,
+ }
+
+ def __init__(self, config: MoonshineConfig):
+ super().__init__(config)
+ self.config = config
+ embed_dim = config.hidden_size
+
+ self.conv1 = nn.Conv1d(1, embed_dim, kernel_size=127, stride=64, bias=False)
+ self.conv2 = nn.Conv1d(embed_dim, 2 * embed_dim, kernel_size=7, stride=3)
+ self.conv3 = nn.Conv1d(2 * embed_dim, embed_dim, kernel_size=3, stride=2)
+ self.groupnorm = nn.GroupNorm(num_groups=1, num_channels=embed_dim, eps=1e-5)
+
+ self.layers = nn.ModuleList(
+ [MoonshineEncoderLayer(config, idx) for idx in range(config.encoder_num_hidden_layers)]
+ )
+ self.layer_norm = nn.LayerNorm(embed_dim, bias=False)
+ self.rotary_emb = MoonshineRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.conv1
+
+ def set_input_embeddings(self, value: nn.Module):
+ self.conv1 = value
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self,
+ input_values: torch.FloatTensor,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPast:
+ r"""
+ Args:
+ input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):
+ Float values of the raw speech waveform. Raw speech waveform can be
+ obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a
+ `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or
+ the soundfile library (`pip install soundfile`). To prepare the array into
+ `input_values`, the [`AutoFeatureExtractor`] should be used for padding
+ and conversion into a tensor of type `torch.FloatTensor`.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding indices in `input_values`. Mask values selected in `[0, 1]`:
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ [What are attention masks?](../glossary#attention-mask)
+ """
+ input_values = input_values.unsqueeze(1)
+ hidden_states = nn.functional.tanh(self.conv1(input_values))
+ hidden_states = self.groupnorm(hidden_states)
+ hidden_states = nn.functional.gelu(self.conv2(hidden_states))
+ hidden_states = nn.functional.gelu(self.conv3(hidden_states))
+ hidden_states = hidden_states.permute(0, 2, 1)
+
+ # attention mask downsampling
+ output_attention_mask = None
+ if attention_mask is not None:
+ mask_len = self._get_feat_extract_output_lengths(attention_mask.shape[-1])
+ downsample_stride = 64 * 3 * 2 # conv strides
+ attention_mask = attention_mask[..., ::downsample_stride][..., :mask_len]
+ output_attention_mask = attention_mask
+
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=hidden_states,
+ attention_mask=attention_mask,
+ encoder_hidden_states=hidden_states,
+ )
+
+ position_ids = torch.arange(0, hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for encoder_layer in self.layers:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.layer_norm(hidden_states)
+
+ return MoonshineEncoderModelOutput(
+ last_hidden_state=hidden_states,
+ attention_mask=output_attention_mask.int() if output_attention_mask is not None else None,
+ )
+
+
+@auto_docstring
+class MoonshineDecoder(MoonshinePreTrainedModel):
+ main_input_name = "input_ids"
+ _can_record_outputs = {
+ "attentions": OutputRecorder(MoonshineAttention, index=1, layer_name="self_attn"),
+ "hidden_states": MoonshineDecoderLayer,
+ "cross_attentions": OutputRecorder(MoonshineAttention, index=1, layer_name="encoder_attn"),
+ }
+
+ def __init__(self, config: MoonshineConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList([MoonshineDecoderLayer(config, idx) for idx in range(config.num_hidden_layers)])
+ self.norm = nn.LayerNorm(config.hidden_size, bias=False)
+ self.rotary_emb = MoonshineRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPast:
+ r"""
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
+ of the decoder.
+ encoder_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding indices in `encoder_hidden_states`. Mask values selected in `[0, 1]`:
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ [What are attention masks?](../glossary#attention-mask)
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers:
+ hidden_states = decoder_layer(
+ hidden_states,
+ causal_mask,
+ encoder_hidden_states, # as a positional argument for gradient checkpointing
+ encoder_attention_mask=encoder_attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ )
+
+
+@auto_docstring
+class MoonshineModel(MoonshinePreTrainedModel):
+ def __init__(self, config: MoonshineConfig):
+ super().__init__(config)
+
+ self.encoder = MoonshineEncoder(config)
+ self.decoder = MoonshineDecoder(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.decoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.decoder.embed_tokens = value
+
+ def freeze_encoder(self):
+ """
+ Calling this function will disable the gradient computation for the Moonshine encoder so that its parameters will
+ not be updated during training.
+ """
+ self.encoder._freeze_parameters()
+
+ def _mask_input_features(self):
+ """
+ Masks extracted features along time axis and/or along feature axis according to
+ [SpecAugment](https://huggingface.co/papers/1904.08779).
+ """
+ raise AttributeError("Not needed for Moonshine")
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_values: torch.FloatTensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None,
+ past_key_values: EncoderDecoderCache | None = None,
+ decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None,
+ decoder_position_ids: tuple[torch.LongTensor] | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Seq2SeqModelOutput:
+ r"""
+ input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):
+ Float values of the raw speech waveform. Raw speech waveform can be
+ obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a
+ `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or
+ the soundfile library (`pip install soundfile`). To prepare the array into
+ `input_values`, the [`AutoFeatureExtractor`] should be used for padding
+ and conversion into a tensor of type `torch.FloatTensor`.
+ decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`):
+ Indices of positions of each input sequence tokens in the position embeddings.
+ Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings`
+
+ Example:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoFeatureExtractor, MoonshineModel
+ >>> from datasets import load_dataset
+
+ >>> model = MoonshineModel.from_pretrained("UsefulSensors/moonshine-tiny")
+ >>> feature_extractor = AutoFeatureExtractor.from_pretrained("UsefulSensors/moonshine-tiny")
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt")
+ >>> input_values = inputs.input_values
+ >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id
+ >>> last_hidden_state = model(input_values, decoder_input_ids=decoder_input_ids).last_hidden_state
+ >>> list(last_hidden_state.shape)
+ [1, 2, 288]
+ ```
+ """
+ if encoder_outputs is None:
+ encoder_outputs: BaseModelOutput = self.encoder(input_values, attention_mask=attention_mask, **kwargs)
+
+ decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder(
+ input_ids=decoder_input_ids,
+ attention_mask=decoder_attention_mask,
+ encoder_hidden_states=encoder_outputs.last_hidden_state,
+ encoder_attention_mask=encoder_outputs.attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=decoder_inputs_embeds,
+ position_ids=decoder_position_ids,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return Seq2SeqModelOutput(
+ last_hidden_state=decoder_outputs.last_hidden_state,
+ past_key_values=decoder_outputs.past_key_values,
+ decoder_hidden_states=decoder_outputs.hidden_states,
+ decoder_attentions=decoder_outputs.attentions,
+ cross_attentions=decoder_outputs.cross_attentions,
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+ encoder_hidden_states=encoder_outputs.hidden_states,
+ encoder_attentions=encoder_outputs.attentions,
+ )
+
+
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+ """
+ Shift input ids one token to the right.
+ """
+ shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+ shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+ shifted_input_ids[:, 0] = decoder_start_token_id
+
+ if pad_token_id is None:
+ raise ValueError("self.model.config.pad_token_id has to be defined.")
+ # replace possible -100 values in labels by `pad_token_id`
+ shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+ return shifted_input_ids
+
+
+@auto_docstring(
+ custom_intro="""
+ The Moonshine Model with a language modeling head. Can be used for automatic speech recognition.
+ """
+)
+class MoonshineForConditionalGeneration(MoonshinePreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"proj_out.weight": "model.decoder.embed_tokens.weight"}
+
+ def __init__(self, config: MoonshineConfig):
+ super().__init__(config)
+ self.model = MoonshineModel(config)
+ self.proj_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.proj_out
+
+ def set_output_embeddings(self, new_embeddings):
+ self.proj_out = new_embeddings
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.model.get_input_embeddings()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_values: torch.FloatTensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None,
+ past_key_values: EncoderDecoderCache | None = None,
+ decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None,
+ decoder_position_ids: tuple[torch.LongTensor] | None = None,
+ use_cache: bool | None = None,
+ labels: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Seq2SeqLMOutput:
+ r"""
+ input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):
+ Float values of the raw speech waveform. Raw speech waveform can be
+ obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a
+ `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or
+ the soundfile library (`pip install soundfile`). To prepare the array into
+ `input_values`, the [`AutoFeatureExtractor`] should be used for padding
+ and conversion into a tensor of type `torch.FloatTensor`.
+ decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`):
+ Indices of positions of each input sequence tokens in the position embeddings.
+ Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings`
+
+ Example:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoProcessor, MoonshineForConditionalGeneration
+ >>> from datasets import load_dataset
+
+ >>> processor = AutoProcessor.from_pretrained("UsefulSensors/moonshine-tiny")
+ >>> model = MoonshineForConditionalGeneration.from_pretrained("UsefulSensors/moonshine-tiny")
+
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+ >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt")
+ >>> input_values = inputs.input_values
+
+ >>> generated_ids = model.generate(input_values, max_new_tokens=100)
+
+ >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+ >>> transcription
+ 'Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.'
+ ```"""
+
+ if labels is not None:
+ if decoder_input_ids is None and decoder_inputs_embeds is None:
+ decoder_input_ids = shift_tokens_right(
+ labels, self.config.pad_token_id, self.config.decoder_start_token_id
+ )
+
+ outputs: Seq2SeqModelOutput = self.model(
+ input_values,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ encoder_outputs=encoder_outputs,
+ decoder_attention_mask=decoder_attention_mask,
+ past_key_values=past_key_values,
+ decoder_inputs_embeds=decoder_inputs_embeds,
+ decoder_position_ids=decoder_position_ids,
+ use_cache=use_cache,
+ **kwargs,
+ )
+ logits = self.proj_out(outputs.last_hidden_state)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)
+
+ return Seq2SeqLMOutput(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ decoder_hidden_states=outputs.decoder_hidden_states,
+ decoder_attentions=outputs.decoder_attentions,
+ cross_attentions=outputs.cross_attentions,
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+ encoder_hidden_states=outputs.encoder_hidden_states,
+ encoder_attentions=outputs.encoder_attentions,
+ )
+
+
+__all__ = ["MoonshineModel", "MoonshinePreTrainedModel", "MoonshineForConditionalGeneration"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/moonshine/modular_moonshine.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/moonshine/modular_moonshine.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f2a6bcdaa424fdc2d0d11dea253cbe402d7298d
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/moonshine/modular_moonshine.py
@@ -0,0 +1,795 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from huggingface_hub.dataclasses import strict
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...configuration_utils import PreTrainedConfig
+from ...generation import GenerationMixin
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPast,
+ BaseModelOutputWithPastAndCrossAttentions,
+ Seq2SeqLMOutput,
+ Seq2SeqModelOutput,
+)
+from ...modeling_rope_utils import RopeParameters
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from ..glm.modeling_glm import GlmAttention, GlmRotaryEmbedding, apply_rotary_pos_emb
+from ..llama.modeling_llama import LlamaDecoderLayer, LlamaModel, eager_attention_forward
+from ..whisper.modeling_whisper import WhisperModel, shift_tokens_right
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="UsefulSensors/moonshine-tiny")
+@strict
+class MoonshineConfig(PreTrainedConfig):
+ r"""
+ encoder_num_key_value_heads (`int`, *optional*):
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+ `encoder_num_key_value_heads=encoder_num_attention_heads`, the model will use Multi Head Attention (MHA), if
+ `encoder_num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+ by meanpooling all the original heads within that group. For more details, check out [this
+ paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
+ `num_attention_heads`.
+ decoder_num_key_value_heads (`int`, *optional*):
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+ `decoder_num_key_value_heads=decoder_num_attention_heads`, the model will use Multi Head Attention (MHA), if
+ `decoder_num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+ by meanpooling all the original heads within that group. For more details, check out [this
+ paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
+ `decoder_num_attention_heads`.
+ pad_head_dim_to_multiple_of (`int`, *optional*):
+ Pad head dimension in encoder and decoder to the next multiple of this value. Necessary for using certain
+ optimized attention implementations.
+ encoder_hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
+ The non-linear activation function (function or string) in the encoder.
+ decoder_hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function (function or string) in the decoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import MoonshineModel, MoonshineConfig
+
+ >>> # Initializing a Moonshine style configuration
+ >>> configuration = MoonshineConfig().from_pretrained("UsefulSensors/moonshine-tiny")
+
+ >>> # Initializing a model from the configuration
+ >>> model = MoonshineModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "moonshine"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ attribute_map = {
+ "num_key_value_heads": "decoder_num_key_value_heads",
+ "num_attention_heads": "decoder_num_attention_heads",
+ "num_hidden_layers": "decoder_num_hidden_layers",
+ "hidden_act": "decoder_hidden_act",
+ }
+
+ vocab_size: int = 32768
+ hidden_size: int = 288
+ intermediate_size: int = 1152
+ encoder_num_hidden_layers: int = 6
+ decoder_num_hidden_layers: int = 6
+ encoder_num_attention_heads: int = 8
+ decoder_num_attention_heads: int = 8
+ encoder_num_key_value_heads: int | None = None
+ decoder_num_key_value_heads: int | None = None
+ pad_head_dim_to_multiple_of: int | None = None
+ encoder_hidden_act: str = "gelu"
+ decoder_hidden_act: str = "silu"
+ max_position_embeddings: int = 512
+ initializer_range: float = 0.02
+ decoder_start_token_id: int = 1
+ use_cache: bool = True
+ rope_parameters: RopeParameters | dict | None = None
+ is_encoder_decoder: bool = True
+ attention_bias: bool = False
+ attention_dropout: float | int = 0.0
+ bos_token_id: int | None = 1
+ eos_token_id: int | list[int] | None = 2
+ pad_token_id: int | None = None
+ tie_word_embeddings: bool = True
+
+ def __post_init__(self, **kwargs):
+ if self.encoder_num_key_value_heads is None:
+ self.encoder_num_key_value_heads = self.encoder_num_attention_heads
+
+ if self.decoder_num_key_value_heads is None:
+ self.decoder_num_key_value_heads = self.decoder_num_attention_heads
+
+ kwargs.setdefault("partial_rotary_factor", 0.9) # assign default for BC
+ super().__post_init__(**kwargs)
+
+
+@auto_docstring(
+ custom_intro="""
+ Extends [~modeling_outputs.BaseModelOutput] to include the output attention mask since sequence length is not preserved in the model's forward.
+ """
+)
+@dataclass
+class MoonshineEncoderModelOutput(BaseModelOutput):
+ r"""
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices after sequence compression. Returned because the
+ sequence length may differ from the input sequence length. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ """
+
+ attention_mask: torch.Tensor | None = None
+
+
+class MoonshineEncoderMLP(nn.Module):
+ def __init__(self, config, hidden_act):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[hidden_act]
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+class MoonshineDecoderMLP(nn.Module):
+ def __init__(self, config, hidden_act):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[hidden_act]
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size * 2)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states, gate = hidden_states.chunk(2, dim=-1)
+ hidden_states = self.activation_fn(gate) * hidden_states
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+class MoonshineRotaryEmbedding(GlmRotaryEmbedding):
+ pass
+
+
+class MoonshineAttention(GlmAttention):
+ def __init__(
+ self,
+ config: MoonshineConfig,
+ layer_idx: int,
+ is_causal: bool,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ ):
+ config.update({"num_attention_heads": num_attention_heads, "num_key_value_heads": num_key_value_heads})
+ super().__init__(config, layer_idx)
+ self.is_causal = is_causal
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+
+ # Pad head dimension to the next specified multiple.
+ if self.config.pad_head_dim_to_multiple_of is not None:
+ target_multiple = self.config.pad_head_dim_to_multiple_of
+ target_head_dim = target_multiple * ((self.head_dim + target_multiple - 1) // target_multiple)
+ self.head_dim_padding = target_head_dim - self.head_dim
+ else:
+ self.head_dim_padding = 0
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ key_value_states: torch.Tensor | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+ bsz, q_len = hidden_states.shape[:-1]
+
+ query_states = (
+ self.q_proj(hidden_states).view(bsz, q_len, self.config.num_key_value_heads, self.head_dim).transpose(1, 2)
+ )
+
+ is_cross_attention = key_value_states is not None
+ if past_key_values is not None:
+ is_updated = past_key_values.is_updated.get(self.layer_idx)
+ if is_cross_attention:
+ # after the first generated id, we can subsequently re-use all key/value_states from cache
+ past_key_values.is_updated[self.layer_idx] = True
+ past_key_values = past_key_values.cross_attention_cache
+ else:
+ past_key_values = past_key_values.self_attention_cache
+
+ # use key_value_states if cross attention
+ current_states = key_value_states if key_value_states is not None else hidden_states
+ if is_cross_attention and past_key_values and is_updated:
+ key_states = past_key_values.layers[self.layer_idx].keys
+ value_states = past_key_values.layers[self.layer_idx].values
+ else:
+ key_states = (
+ self.k_proj(current_states)
+ .view(bsz, -1, self.config.num_key_value_heads, self.head_dim)
+ .transpose(1, 2)
+ )
+ value_states = (
+ self.v_proj(current_states)
+ .view(bsz, -1, self.config.num_key_value_heads, self.head_dim)
+ .transpose(1, 2)
+ )
+ if is_cross_attention and past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ if not is_cross_attention:
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ is_causal = self.is_causal and attention_mask is None and q_len > 1
+
+ if self.head_dim_padding > 0:
+ query_states = torch.nn.functional.pad(query_states, (0, self.head_dim_padding))
+ key_states = torch.nn.functional.pad(key_states, (0, self.head_dim_padding))
+ value_states = torch.nn.functional.pad(value_states, (0, self.head_dim_padding))
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ is_causal=is_causal,
+ **kwargs,
+ )
+
+ if self.head_dim_padding > 0:
+ attn_output = attn_output[..., : -self.head_dim_padding]
+
+ attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class MoonshineEncoderLayer(LlamaDecoderLayer):
+ def __init__(self, config: MoonshineConfig, layer_idx: int):
+ super().__init__(config, layer_idx)
+
+ self.self_attn = MoonshineAttention(
+ config=config,
+ layer_idx=layer_idx,
+ is_causal=False,
+ num_attention_heads=config.encoder_num_attention_heads,
+ num_key_value_heads=config.encoder_num_key_value_heads,
+ )
+
+ self.mlp = MoonshineEncoderMLP(config, config.encoder_hidden_act)
+ self.input_layernorm = nn.LayerNorm(config.hidden_size, bias=False)
+ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, bias=False)
+
+
+class MoonshineDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: MoonshineConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = MoonshineAttention(
+ config=config,
+ layer_idx=layer_idx,
+ is_causal=True,
+ num_attention_heads=config.num_attention_heads,
+ num_key_value_heads=config.num_key_value_heads,
+ )
+ self.encoder_attn = MoonshineAttention(
+ config=config,
+ layer_idx=layer_idx,
+ is_causal=False,
+ num_attention_heads=config.num_attention_heads,
+ num_key_value_heads=config.num_key_value_heads,
+ )
+
+ self.mlp = MoonshineDecoderMLP(config, config.hidden_act)
+ self.input_layernorm = nn.LayerNorm(config.hidden_size, bias=False)
+ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, bias=False)
+ self.final_layernorm = nn.LayerNorm(config.hidden_size, bias=False)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ encoder_position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ encoder_position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ if encoder_hidden_states is not None:
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states, _ = self.encoder_attn(
+ hidden_states=hidden_states,
+ key_value_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.final_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+@auto_docstring
+class MoonshinePreTrainedModel(PreTrainedModel):
+ config: MoonshineConfig
+ base_model_prefix = "model"
+ main_input_name = "input_values"
+ input_modalities = "audio"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["MoonshineEncoderLayer", "MoonshineDecoderLayer"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+
+ _can_compile_fullgraph = True
+ # TODO arthur, how do we separate when it cross / self coming from different layer?
+
+ def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
+ """
+ Computes the output length of the convolutional layers
+ """
+ output_conv1_length = int((input_lengths - 127) / 64 + 1)
+ output_conv2_length = int((output_conv1_length - 7) / 3 + 1)
+ output_conv3_length = int((output_conv2_length - 3) / 2 + 1)
+
+ return output_conv3_length
+
+
+class MoonshineEncoder(MoonshinePreTrainedModel):
+ """
+ Transformer encoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MoonshineEncoderLayer`]
+
+ Args:
+ config: MoonshineConfig
+ """
+
+ main_input_name = "input_values"
+ _can_record_outputs = {
+ "attentions": MoonshineAttention,
+ "hidden_states": MoonshineEncoderLayer,
+ }
+
+ def __init__(self, config: MoonshineConfig):
+ super().__init__(config)
+ self.config = config
+ embed_dim = config.hidden_size
+
+ self.conv1 = nn.Conv1d(1, embed_dim, kernel_size=127, stride=64, bias=False)
+ self.conv2 = nn.Conv1d(embed_dim, 2 * embed_dim, kernel_size=7, stride=3)
+ self.conv3 = nn.Conv1d(2 * embed_dim, embed_dim, kernel_size=3, stride=2)
+ self.groupnorm = nn.GroupNorm(num_groups=1, num_channels=embed_dim, eps=1e-5)
+
+ self.layers = nn.ModuleList(
+ [MoonshineEncoderLayer(config, idx) for idx in range(config.encoder_num_hidden_layers)]
+ )
+ self.layer_norm = nn.LayerNorm(embed_dim, bias=False)
+ self.rotary_emb = MoonshineRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.conv1
+
+ def set_input_embeddings(self, value: nn.Module):
+ self.conv1 = value
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self,
+ input_values: torch.FloatTensor,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPast:
+ r"""
+ Args:
+ input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):
+ Float values of the raw speech waveform. Raw speech waveform can be
+ obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a
+ `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or
+ the soundfile library (`pip install soundfile`). To prepare the array into
+ `input_values`, the [`AutoFeatureExtractor`] should be used for padding
+ and conversion into a tensor of type `torch.FloatTensor`.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding indices in `input_values`. Mask values selected in `[0, 1]`:
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ [What are attention masks?](../glossary#attention-mask)
+ """
+ input_values = input_values.unsqueeze(1)
+ hidden_states = nn.functional.tanh(self.conv1(input_values))
+ hidden_states = self.groupnorm(hidden_states)
+ hidden_states = nn.functional.gelu(self.conv2(hidden_states))
+ hidden_states = nn.functional.gelu(self.conv3(hidden_states))
+ hidden_states = hidden_states.permute(0, 2, 1)
+
+ # attention mask downsampling
+ output_attention_mask = None
+ if attention_mask is not None:
+ mask_len = self._get_feat_extract_output_lengths(attention_mask.shape[-1])
+ downsample_stride = 64 * 3 * 2 # conv strides
+ attention_mask = attention_mask[..., ::downsample_stride][..., :mask_len]
+ output_attention_mask = attention_mask
+
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=hidden_states,
+ attention_mask=attention_mask,
+ encoder_hidden_states=hidden_states,
+ )
+
+ position_ids = torch.arange(0, hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for encoder_layer in self.layers:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.layer_norm(hidden_states)
+
+ return MoonshineEncoderModelOutput(
+ last_hidden_state=hidden_states,
+ attention_mask=output_attention_mask.int() if output_attention_mask is not None else None,
+ )
+
+
+class MoonshineDecoder(LlamaModel):
+ main_input_name = "input_ids"
+ _can_record_outputs = {
+ "attentions": OutputRecorder(MoonshineAttention, index=1, layer_name="self_attn"),
+ "hidden_states": MoonshineDecoderLayer,
+ "cross_attentions": OutputRecorder(MoonshineAttention, index=1, layer_name="encoder_attn"),
+ }
+
+ def __init__(self, config: MoonshineConfig):
+ super().__init__(config)
+ self.norm = nn.LayerNorm(config.hidden_size, bias=False)
+ self.layers = nn.ModuleList([MoonshineDecoderLayer(config, idx) for idx in range(config.num_hidden_layers)])
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPast:
+ r"""
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
+ of the decoder.
+ encoder_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding indices in `encoder_hidden_states`. Mask values selected in `[0, 1]`:
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ [What are attention masks?](../glossary#attention-mask)
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers:
+ hidden_states = decoder_layer(
+ hidden_states,
+ causal_mask,
+ encoder_hidden_states, # as a positional argument for gradient checkpointing
+ encoder_attention_mask=encoder_attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ )
+
+
+class MoonshineModel(WhisperModel):
+ def _mask_input_features(self):
+ raise AttributeError("Not needed for Moonshine")
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_values: torch.FloatTensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None,
+ past_key_values: EncoderDecoderCache | None = None,
+ decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None,
+ decoder_position_ids: tuple[torch.LongTensor] | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Seq2SeqModelOutput:
+ r"""
+ input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):
+ Float values of the raw speech waveform. Raw speech waveform can be
+ obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a
+ `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or
+ the soundfile library (`pip install soundfile`). To prepare the array into
+ `input_values`, the [`AutoFeatureExtractor`] should be used for padding
+ and conversion into a tensor of type `torch.FloatTensor`.
+ decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`):
+ Indices of positions of each input sequence tokens in the position embeddings.
+ Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings`
+
+ Example:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoFeatureExtractor, MoonshineModel
+ >>> from datasets import load_dataset
+
+ >>> model = MoonshineModel.from_pretrained("UsefulSensors/moonshine-tiny")
+ >>> feature_extractor = AutoFeatureExtractor.from_pretrained("UsefulSensors/moonshine-tiny")
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt")
+ >>> input_values = inputs.input_values
+ >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id
+ >>> last_hidden_state = model(input_values, decoder_input_ids=decoder_input_ids).last_hidden_state
+ >>> list(last_hidden_state.shape)
+ [1, 2, 288]
+ ```
+ """
+ if encoder_outputs is None:
+ encoder_outputs: BaseModelOutput = self.encoder(input_values, attention_mask=attention_mask, **kwargs)
+
+ decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder(
+ input_ids=decoder_input_ids,
+ attention_mask=decoder_attention_mask,
+ encoder_hidden_states=encoder_outputs.last_hidden_state,
+ encoder_attention_mask=encoder_outputs.attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=decoder_inputs_embeds,
+ position_ids=decoder_position_ids,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return Seq2SeqModelOutput(
+ last_hidden_state=decoder_outputs.last_hidden_state,
+ past_key_values=decoder_outputs.past_key_values,
+ decoder_hidden_states=decoder_outputs.hidden_states,
+ decoder_attentions=decoder_outputs.attentions,
+ cross_attentions=decoder_outputs.cross_attentions,
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+ encoder_hidden_states=encoder_outputs.hidden_states,
+ encoder_attentions=encoder_outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The Moonshine Model with a language modeling head. Can be used for automatic speech recognition.
+ """
+)
+class MoonshineForConditionalGeneration(MoonshinePreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"proj_out.weight": "model.decoder.embed_tokens.weight"}
+
+ def __init__(self, config: MoonshineConfig):
+ super().__init__(config)
+ self.model = MoonshineModel(config)
+ self.proj_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.proj_out
+
+ def set_output_embeddings(self, new_embeddings):
+ self.proj_out = new_embeddings
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.model.get_input_embeddings()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_values: torch.FloatTensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None,
+ past_key_values: EncoderDecoderCache | None = None,
+ decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None,
+ decoder_position_ids: tuple[torch.LongTensor] | None = None,
+ use_cache: bool | None = None,
+ labels: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Seq2SeqLMOutput:
+ r"""
+ input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):
+ Float values of the raw speech waveform. Raw speech waveform can be
+ obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a
+ `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or
+ the soundfile library (`pip install soundfile`). To prepare the array into
+ `input_values`, the [`AutoFeatureExtractor`] should be used for padding
+ and conversion into a tensor of type `torch.FloatTensor`.
+ decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`):
+ Indices of positions of each input sequence tokens in the position embeddings.
+ Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings`
+
+ Example:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoProcessor, MoonshineForConditionalGeneration
+ >>> from datasets import load_dataset
+
+ >>> processor = AutoProcessor.from_pretrained("UsefulSensors/moonshine-tiny")
+ >>> model = MoonshineForConditionalGeneration.from_pretrained("UsefulSensors/moonshine-tiny")
+
+ >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+ >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt")
+ >>> input_values = inputs.input_values
+
+ >>> generated_ids = model.generate(input_values, max_new_tokens=100)
+
+ >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+ >>> transcription
+ 'Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.'
+ ```"""
+
+ if labels is not None:
+ if decoder_input_ids is None and decoder_inputs_embeds is None:
+ decoder_input_ids = shift_tokens_right(
+ labels, self.config.pad_token_id, self.config.decoder_start_token_id
+ )
+
+ outputs: Seq2SeqModelOutput = self.model(
+ input_values,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ encoder_outputs=encoder_outputs,
+ decoder_attention_mask=decoder_attention_mask,
+ past_key_values=past_key_values,
+ decoder_inputs_embeds=decoder_inputs_embeds,
+ decoder_position_ids=decoder_position_ids,
+ use_cache=use_cache,
+ **kwargs,
+ )
+ logits = self.proj_out(outputs.last_hidden_state)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)
+
+ return Seq2SeqLMOutput(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ decoder_hidden_states=outputs.decoder_hidden_states,
+ decoder_attentions=outputs.decoder_attentions,
+ cross_attentions=outputs.cross_attentions,
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+ encoder_hidden_states=outputs.encoder_hidden_states,
+ encoder_attentions=outputs.encoder_attentions,
+ )
+
+
+__all__ = [
+ "MoonshineConfig",
+ "MoonshineModel",
+ "MoonshinePreTrainedModel",
+ "MoonshineForConditionalGeneration",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..880274309cbab486a90df05548a0e4d3f2ea0925
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_musicgen import *
+ from .modeling_musicgen import *
+ from .processing_musicgen import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/configuration_musicgen.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/configuration_musicgen.py
new file mode 100644
index 0000000000000000000000000000000000000000..1d496b9b3557c7566db898f312a8fd7489b33709
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/configuration_musicgen.py
@@ -0,0 +1,156 @@
+# Copyright 2023 Meta AI and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MusicGen model configuration"""
+
+from typing import ClassVar
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+from ..auto.configuration_auto import AutoConfig
+
+
+@auto_docstring(checkpoint="facebook/musicgen-small")
+@strict
+class MusicgenDecoderConfig(PreTrainedConfig):
+ model_type = "musicgen_decoder"
+ base_config_key = "decoder_config"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ vocab_size: int = 2048
+ max_position_embeddings: int = 2048
+ num_hidden_layers: int = 24
+ ffn_dim: int = 4096
+ num_attention_heads: int = 16
+ layerdrop: float | int = 0.0
+ use_cache: bool = True
+ activation_function: str = "gelu"
+ hidden_size: int = 1024
+ dropout: float | int = 0.1
+ attention_dropout: float | int = 0.0
+ activation_dropout: float | int = 0.0
+ initializer_factor: float = 0.02
+ scale_embedding: bool = False
+ num_codebooks: int = 4
+ audio_channels: int = 1
+ pad_token_id: int | None = 2048
+ bos_token_id: int | None = 2048
+ eos_token_id: int | list[int] | None = None
+ tie_word_embeddings: bool = False
+ is_decoder: bool = False
+ add_cross_attention: bool = False
+ cross_attention_hidden_size: int | None = None
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if self.audio_channels not in [1, 2]:
+ raise ValueError(f"Expected 1 (mono) or 2 (stereo) audio channels, got {self.audio_channels} channels.")
+
+
+@auto_docstring(checkpoint="facebook/musicgen-small")
+@strict
+class MusicgenConfig(PreTrainedConfig):
+ r"""
+ text_encoder (`Union[dict, `PretrainedConfig`]`):
+ An instance of a configuration object that defines the text encoder config.
+ audio_encoder (`Union[dict, `PretrainedConfig`]`):
+ An instance of a configuration object that defines the audio encoder config.
+ decoder (`Union[dict, `PretrainedConfig`]`):
+ An instance of a configuration object that defines the decoder config.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... MusicgenConfig,
+ ... MusicgenDecoderConfig,
+ ... T5Config,
+ ... EncodecConfig,
+ ... MusicgenForConditionalGeneration,
+ ... )
+
+ >>> # Initializing text encoder, audio encoder, and decoder model configurations
+ >>> text_encoder_config = T5Config()
+ >>> audio_encoder_config = EncodecConfig()
+ >>> decoder_config = MusicgenDecoderConfig()
+
+ >>> configuration = MusicgenConfig(
+ ... text_encoder=text_encoder_config,
+ ... audio_encoder=audio_encoder_config,
+ ... decoder=decoder_config,
+ ... )
+
+ >>> # Initializing a MusicgenForConditionalGeneration (with random weights) from the facebook/musicgen-small style configuration
+ >>> model = MusicgenForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ >>> config_text_encoder = model.config.text_encoder
+ >>> config_audio_encoder = model.config.audio_encoder
+ >>> config_decoder = model.config.decoder
+
+ >>> # Saving the model, including its configuration
+ >>> model.save_pretrained("musicgen-model")
+
+ >>> # loading model and config from pretrained folder
+ >>> musicgen_config = MusicgenConfig.from_pretrained("musicgen-model")
+ >>> model = MusicgenForConditionalGeneration.from_pretrained("musicgen-model", config=musicgen_config)
+ ```"""
+
+ model_type: ClassVar[str] = "musicgen"
+ sub_configs: ClassVar[dict[str, type[PreTrainedConfig]]] = {
+ "text_encoder": AutoConfig,
+ "audio_encoder": AutoConfig,
+ "decoder": MusicgenDecoderConfig,
+ }
+ has_no_defaults_at_init: ClassVar[bool] = True
+
+ text_encoder: dict | PreTrainedConfig = None
+ audio_encoder: dict | PreTrainedConfig = None
+ decoder: dict | PreTrainedConfig = None
+ initializer_factor: float = 0.02
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.text_encoder, dict):
+ text_encoder_model_type = self.text_encoder.pop("model_type")
+ self.text_encoder = AutoConfig.for_model(text_encoder_model_type, **self.text_encoder)
+ elif self.text_encoder is None:
+ raise ValueError(
+ f"A configuration of type {self.model_type} cannot be instantiated because text_encoder is not passed"
+ )
+
+ if isinstance(self.audio_encoder, dict):
+ audio_encoder_model_type = self.audio_encoder.pop("model_type")
+ self.audio_encoder = AutoConfig.for_model(audio_encoder_model_type, **self.audio_encoder)
+ elif self.audio_encoder is None:
+ raise ValueError(
+ f"A configuration of type {self.model_type} cannot be instantiated because audio_encoder is not passed"
+ )
+
+ if isinstance(self.decoder, dict):
+ self.decoder = MusicgenDecoderConfig(**self.decoder)
+ elif self.decoder is None:
+ self.decoder = MusicgenDecoderConfig()
+
+ self.is_encoder_decoder = True
+ super().__post_init__(**kwargs)
+
+ @property
+ # This is a property because you might want to change the codec model on the fly
+ def sampling_rate(self):
+ return self.audio_encoder.sampling_rate
+
+
+__all__ = ["MusicgenConfig", "MusicgenDecoderConfig"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/modeling_musicgen.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/modeling_musicgen.py
new file mode 100644
index 0000000000000000000000000000000000000000..33a02805a38afdb9d04fb9284aa56eb7b5e0588e
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/modeling_musicgen.py
@@ -0,0 +1,2173 @@
+# Copyright 2023 Meta AI and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Musicgen model."""
+
+import copy
+import inspect
+import math
+import random
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import (
+ ClassifierFreeGuidanceLogitsProcessor,
+ GenerationConfig,
+ GenerationMixin,
+ GenerationMode,
+ LogitsProcessorList,
+ StoppingCriteriaList,
+)
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_flash_attention_utils import (
+ FlashAttentionKwargs,
+)
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ ModelOutput,
+ Seq2SeqLMOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, logging
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from ..auto.configuration_auto import AutoConfig
+from ..auto.modeling_auto import AutoModel
+from .configuration_musicgen import MusicgenConfig, MusicgenDecoderConfig
+
+
+if TYPE_CHECKING:
+ from ...generation.streamers import BaseStreamer
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring
+@dataclass
+class MusicgenUnconditionalInput(ModelOutput):
+ r"""
+ encoder_outputs (`tuple[torch.FloatTensor]` of length 1, with tensor shape `(batch_size, sequence_length, hidden_size)`):
+ Sequence of hidden-states at the output of the last layer of the text encoder model.
+ attention_mask (`torch.LongTensor`) of shape `(batch_size, sequence_length)`, *optional*):
+ Encoder attention mask to avoid performing attention on padding token indices. Mask values selected in `[0,
+ 1]`: 1 for tokens that are **not masked**, 0 for tokens that are **masked**.
+ guidance_scale (`float`, *optional*):
+ Guidance scale for classifier free guidance, setting the balance between the conditional logits (predicted
+ from the prompts) and the unconditional logits (predicted without prompts).
+ """
+
+ encoder_outputs: tuple[torch.FloatTensor] | None = None
+ attention_mask: torch.LongTensor | None = None
+ guidance_scale: float | None = None
+
+
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+ """
+ Shift input ids one token to the right.
+ """
+ # transpose to get (bsz, num_codebooks, seq_len)
+ input_ids = input_ids.transpose(1, 2)
+ shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+ shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()
+ if decoder_start_token_id is None:
+ raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")
+ shifted_input_ids[..., 0] = decoder_start_token_id
+
+ if pad_token_id is None:
+ raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")
+ # replace possible -100 values in labels by `pad_token_id`
+ shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+ return shifted_input_ids
+
+
+class MusicgenSinusoidalPositionalEmbedding(nn.Module):
+ """This module produces sinusoidal positional embeddings of any length."""
+
+ def __init__(self, num_positions: int, embedding_dim: int):
+ super().__init__()
+ self.embedding_dim = embedding_dim
+ self.num_positions = num_positions
+ self.make_weights(num_positions, embedding_dim)
+
+ def make_weights(self, num_embeddings: int, embedding_dim: int):
+ emb_weights = self.get_embedding(num_embeddings, embedding_dim)
+ if hasattr(self, "weights"):
+ # in forward put the weights on the correct dtype and device of the param
+ emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device)
+
+ self.register_buffer("weights", emb_weights, persistent=False)
+
+ @staticmethod
+ def get_embedding(num_embeddings: int, embedding_dim: int):
+ """
+ Build sinusoidal embeddings. This matches the implementation in tensor2tensor, but differs slightly from the
+ description in Section 3.5 of "Attention Is All You Need".
+ """
+ half_dim = embedding_dim // 2
+ emb = math.log(10000) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb)
+ emb = torch.arange(num_embeddings, dtype=torch.int64).float().unsqueeze(1) * emb.unsqueeze(0)
+ emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=1).view(num_embeddings, -1)
+ if embedding_dim % 2 == 1:
+ # zero pad
+ emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
+ return emb.to(torch.get_default_dtype())
+
+ @torch.no_grad()
+ def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0):
+ bsz, codebooks, seq_len = input_ids.size()
+ # Create the position ids from the input token ids.
+ position_ids = (torch.arange(seq_len) + past_key_values_length).to(input_ids.device)
+ # expand embeddings if needed
+ if seq_len > self.weights.size(0):
+ self.make_weights(seq_len, self.embedding_dim)
+ return self.weights.index_select(0, position_ids.view(-1)).detach()
+
+
+# Copied from transformers.models.bert.modeling_bert.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class MusicgenAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float | None = 0.0,
+ is_decoder: bool | None = False,
+ bias: bool | None = True,
+ is_causal: bool | None = False,
+ config: MusicgenConfig | None = None,
+ layer_idx: int | None = None,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.head_dim = embed_dim // num_heads
+ self.config = config
+
+ if (self.head_dim * num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.is_decoder = is_decoder
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ key_value_states: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool | None = False,
+ # TODO: we need a refactor so that the different attention modules can get their specific kwargs
+ # ATM, we have mixed things encoder, decoder, and encoder-decoder attn
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ # if key_value_states are provided this layer is used as a cross-attention layer
+ # for the decoder
+ is_cross_attention = key_value_states is not None
+
+ # determine input shapes
+ input_shape = hidden_states.shape[:-1]
+
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ # get query proj
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ is_updated = False
+ if past_key_values is not None:
+ if isinstance(past_key_values, EncoderDecoderCache):
+ is_updated = past_key_values.is_updated.get(self.layer_idx)
+ if is_cross_attention:
+ # after the first generated id, we can subsequently re-use all key/value_layer from cache
+ curr_past_key_values = past_key_values.cross_attention_cache
+ else:
+ curr_past_key_values = past_key_values.self_attention_cache
+ else:
+ curr_past_key_values = past_key_values
+
+ current_states = key_value_states if is_cross_attention else hidden_states
+ if is_cross_attention and past_key_values is not None and is_updated:
+ # reuse k,v, cross_attentions
+ key_states = curr_past_key_values.layers[self.layer_idx].keys
+ value_states = curr_past_key_values.layers[self.layer_idx].values
+ else:
+ kv_shape = (*current_states.shape[:-1], -1, self.head_dim)
+ key_states = self.k_proj(current_states).view(kv_shape).transpose(1, 2)
+ value_states = self.v_proj(current_states).view(kv_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ # save all key/value_states to cache to be re-used for fast auto-regressive generation
+ key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
+ # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+ if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):
+ past_key_values.is_updated[self.layer_idx] = True
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout,
+ scaling=self.scaling,
+ output_attentions=output_attentions,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+class MusicgenDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: MusicgenDecoderConfig, layer_idx=None):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+
+ self.self_attn = MusicgenAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.num_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ bias=False,
+ is_causal=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.encoder_attn = MusicgenAttention(
+ self.embed_dim,
+ config.num_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ bias=False,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.fc1 = nn.Linear(self.embed_dim, config.ffn_dim, bias=False)
+ self.fc2 = nn.Linear(config.ffn_dim, self.embed_dim, bias=False)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = True,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ encoder_hidden_states (`torch.FloatTensor`):
+ cross attention input to the layer of shape `(batch, seq_len, embed_dim)`
+ encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ past_key_values (`Cache`): cached past key and value projection states
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # Cross-Attention Block
+ if encoder_hidden_states is not None:
+ residual = hidden_states
+ hidden_states = self.encoder_attn_layer_norm(hidden_states)
+
+ hidden_states, _ = self.encoder_attn(
+ hidden_states,
+ key_value_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring
+class MusicgenPreTrainedModel(PreTrainedModel):
+ config: MusicgenDecoderConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["MusicgenDecoderLayer", "MusicgenAttention"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ std = self.config.initializer_factor
+ if isinstance(module, nn.Linear):
+ init.normal_(module.weight, mean=0.0, std=std)
+ if module.bias is not None:
+ init.zeros_(module.bias)
+ elif isinstance(module, nn.LayerNorm):
+ init.ones_(module.weight)
+ init.zeros_(module.bias)
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight, mean=0.0, std=std)
+ # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+ if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+ init.zeros_(module.weight[module.padding_idx])
+ elif isinstance(module, MusicgenSinusoidalPositionalEmbedding):
+ emb_weights = module.get_embedding(module.num_positions, module.embedding_dim)
+ init.copy_(module.weights, emb_weights)
+
+
+class MusicgenDecoder(MusicgenPreTrainedModel):
+ """
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MusicgenDecoderLayer`]
+ """
+
+ _can_record_outputs = {
+ "hidden_states": MusicgenDecoderLayer,
+ "attentions": OutputRecorder(MusicgenAttention, index=1, layer_name="self_attn"),
+ "cross_attentions": OutputRecorder(MusicgenAttention, index=1, layer_name="encoder_attn"),
+ }
+
+ def __init__(self, config: MusicgenDecoderConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+ self.layerdrop = config.layerdrop
+ self.max_target_positions = config.max_position_embeddings
+ self.d_model = config.hidden_size
+ self.num_codebooks = config.num_codebooks
+ self.embed_scale = math.sqrt(config.hidden_size) if config.scale_embedding else 1.0
+
+ embed_dim = config.vocab_size + 1
+ self.embed_tokens = nn.ModuleList(
+ [nn.Embedding(embed_dim, config.hidden_size) for _ in range(config.num_codebooks)]
+ )
+
+ self.embed_positions = MusicgenSinusoidalPositionalEmbedding(
+ config.max_position_embeddings,
+ config.hidden_size,
+ )
+
+ self.layers = nn.ModuleList(
+ [MusicgenDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)]
+ )
+ self.layer_norm = nn.LayerNorm(config.hidden_size)
+ self.attn_implementation = config._attn_implementation
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPastAndCrossAttentions:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size * num_codebooks, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary, corresponding to the sequence of audio codes.
+
+ Indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes,
+ such as with the [`EncodecModel`]. See [`EncodecModel.encode`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+
+
+
+ The `input_ids` will automatically be converted from shape `(batch_size * num_codebooks,
+ target_sequence_length)` to `(batch_size, num_codebooks, target_sequence_length)` in the forward pass. If
+ you obtain audio codes from an audio encoding model, such as [`EncodecModel`], ensure that the number of
+ frames is equal to 1, and that you reshape the audio codes from `(frames, batch_size, num_codebooks,
+ target_sequence_length)` to `(batch_size * num_codebooks, target_sequence_length)` prior to passing them as
+ `input_ids`.
+
+
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of
+ the decoder.
+ encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
+ Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values
+ selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+
+ [What are attention masks?](../glossary#attention-mask)
+ """
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+ elif input_ids is not None:
+ # (bsz * codebooks, seq_len) -> (bsz, codebooks, seq_len)
+ input = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1])
+ bsz, num_codebooks, seq_len = input.shape
+ elif inputs_embeds is not None:
+ input = inputs_embeds[:, :, -1:]
+ else:
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+ if use_cache and past_key_values is None:
+ past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+ if inputs_embeds is None:
+ inputs_embeds = sum(self.embed_tokens[codebook](input[:, codebook]) for codebook in range(num_codebooks))
+
+ attention_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ )
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ # embed positions
+ positions = self.embed_positions(input, past_key_values_length)
+ hidden_states = inputs_embeds + positions.to(inputs_embeds.device)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ for idx, decoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ dropout_probability = random.uniform(0, 1)
+ if self.training and (dropout_probability < self.layerdrop):
+ continue
+
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask,
+ encoder_hidden_states, # as a positional argument for gradient checkpointing
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.layer_norm(hidden_states)
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class MusicgenModel(MusicgenPreTrainedModel):
+ def __init__(self, config: MusicgenDecoderConfig):
+ super().__init__(config)
+ self.decoder = MusicgenDecoder(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.decoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.decoder.embed_tokens = value
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPastAndCrossAttentions:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size * num_codebooks, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary, corresponding to the sequence of audio codes.
+
+ Indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes,
+ such as with the [`EncodecModel`]. See [`EncodecModel.encode`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+
+
+
+ The `input_ids` will automatically be converted from shape `(batch_size * num_codebooks,
+ target_sequence_length)` to `(batch_size, num_codebooks, target_sequence_length)` in the forward pass. If
+ you obtain audio codes from an audio encoding model, such as [`EncodecModel`], ensure that the number of
+ frames is equal to 1, and that you reshape the audio codes from `(frames, batch_size, num_codebooks,
+ target_sequence_length)` to `(batch_size * num_codebooks, target_sequence_length)` prior to passing them as
+ `input_ids`.
+
+
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of
+ the decoder.
+ encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
+ Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values
+ selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+
+ [What are attention masks?](../glossary#attention-mask)
+ """
+ decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ encoder_attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return decoder_outputs
+
+
+@auto_docstring(
+ custom_intro="""
+ The MusicGen decoder model with a language modelling head on top.
+ """
+)
+class MusicgenForCausalLM(MusicgenPreTrainedModel, GenerationMixin):
+ output_modalities = ("audio",)
+
+ def __init__(self, config: MusicgenDecoderConfig):
+ super().__init__(config)
+
+ self.model = MusicgenModel(config)
+
+ self.num_codebooks = config.num_codebooks
+ self.lm_heads = nn.ModuleList(
+ [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_codebooks)]
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.decoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.decoder.embed_tokens = value
+
+ def get_output_embeddings(self):
+ return self.lm_heads
+
+ def set_output_embeddings(self, new_embeddings):
+ self.lm_heads = new_embeddings
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithCrossAttentions:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size * num_codebooks, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary, corresponding to the sequence of audio codes.
+
+ Indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes,
+ such as with the [`EncodecModel`]. See [`EncodecModel.encode`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+
+
+
+ The `input_ids` will automatically be converted from shape `(batch_size * num_codebooks,
+ target_sequence_length)` to `(batch_size, num_codebooks, target_sequence_length)` in the forward pass. If
+ you obtain audio codes from an audio encoding model, such as [`EncodecModel`], ensure that the number of
+ frames is equal to 1, and that you reshape the audio codes from `(frames, batch_size, num_codebooks,
+ target_sequence_length)` to `(batch_size * num_codebooks, target_sequence_length)` prior to passing them as
+ `input_ids`.
+
+
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of
+ the decoder.
+ encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
+ Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values
+ selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+
+ [What are attention masks?](../glossary#attention-mask)
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks)`, *optional*):
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+ `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+ are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+ """
+
+ if (labels is not None) and (input_ids is None and inputs_embeds is None):
+ input_ids = shift_tokens_right(labels, self.config.pad_token_id, self.config.bos_token_id)
+
+ outputs: BaseModelOutputWithPastAndCrossAttentions = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+
+ lm_logits = torch.stack([head(hidden_states) for head in self.lm_heads], dim=1)
+
+ loss = None
+ if labels is not None:
+ # since encoder hidden states have been concatenated to the decoder hidden states,
+ # we take the last timestamps corresponding to labels
+ logits = lm_logits[:, :, -labels.shape[1] :]
+
+ loss_fct = CrossEntropyLoss()
+ loss = torch.zeros([], device=self.device)
+
+ # per codebook cross-entropy
+ # -100 labels are ignored
+ labels = labels.masked_fill(labels == self.config.pad_token_id, -100)
+
+ # per codebook cross-entropy
+ # ref: https://github.com/facebookresearch/audiocraft/blob/69fea8b290ad1b4b40d28f92d1dfc0ab01dbab85/audiocraft/solvers/musicgen.py#L242-L243
+ for codebook in range(self.config.num_codebooks):
+ codebook_logits = logits[:, codebook].contiguous().view(-1, logits.shape[-1])
+ codebook_labels = labels[..., codebook].contiguous().view(-1)
+ loss += loss_fct(codebook_logits, codebook_labels)
+
+ loss = loss / self.config.num_codebooks
+
+ # (bsz, num_codebooks, seq_len, vocab_size) -> (bsz * num_codebooks, seq_len, vocab_size)
+ lm_logits = lm_logits.reshape(-1, *lm_logits.shape[2:])
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=lm_logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ attention_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ use_cache=True,
+ delay_pattern_mask=None,
+ guidance_scale=None,
+ **kwargs,
+ ):
+ # Overwritten -- MusicGen has custom processing
+ if delay_pattern_mask is None:
+ input_ids, delay_pattern_mask = self.build_delay_pattern_mask(
+ input_ids,
+ pad_token_id=self.generation_config.pad_token_id,
+ max_length=self.generation_config.max_length,
+ )
+
+ # apply the delay pattern mask
+ input_ids = self.apply_delay_pattern_mask(input_ids, delay_pattern_mask)
+
+ if guidance_scale is not None and guidance_scale > 1:
+ # for classifier free guidance we need to replicate the decoder args across the batch dim (we'll split these
+ # before sampling)
+ input_ids = input_ids.repeat((2, 1))
+ if attention_mask is not None:
+ attention_mask = attention_mask.repeat((2, 1))
+
+ if past_key_values is not None:
+ input_ids = input_ids[:, -1:]
+
+ return {
+ "input_ids": input_ids,
+ "attention_mask": attention_mask,
+ "encoder_hidden_states": encoder_hidden_states,
+ "encoder_attention_mask": encoder_attention_mask,
+ "past_key_values": past_key_values,
+ "use_cache": use_cache,
+ }
+
+ def build_delay_pattern_mask(self, input_ids: torch.LongTensor, pad_token_id: int, max_length: int | None = None):
+ """Build a delayed pattern mask to the input_ids. Each codebook is offset by the previous codebook by
+ one, giving a delayed pattern mask at the start of sequence and end of sequence. Take the example where there
+ are 4 codebooks and a max sequence length of 8, we have the delayed pattern mask of shape `(codebooks,
+ seq_len)`:
+ - [P, -1, -1, -1, -1, P, P, P]
+ - [P, P, -1, -1, -1, -1, P, P]
+ - [P, P, P, -1, -1, -1, -1, P]
+ - [P, P, P, P, -1, -1, -1, -1]
+ where P is the special padding token id and -1 indicates that the token is valid for prediction. If we include
+ a prompt (decoder input ids), the -1 positions indicate where new tokens should be predicted. Otherwise, the
+ mask is set to the value in the prompt:
+ - [P, a, b, -1, -1, P, P, P]
+ - [P, P, c, d, -1, -1, P, P]
+ - [P, P, P, e, f, -1, -1, P]
+ - [P, P, P, P, g, h, -1, -1]
+ where a-h indicate the input prompt (decoder input ids) that are offset by 1. Now, we only override the -1
+ tokens in our prediction.
+ """
+ # (bsz * num_codebooks, seq_len) -> (bsz, num_codebooks, seq_len)
+ input_ids = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1])
+ bsz, num_codebooks, seq_len = input_ids.shape
+
+ max_length = max_length if max_length is not None else self.generation_config.max_length
+ input_ids_shifted = (
+ torch.ones((bsz, num_codebooks, max_length), dtype=torch.long, device=input_ids.device) * -1
+ )
+
+ channel_codebooks = num_codebooks // 2 if self.config.audio_channels == 2 else num_codebooks
+ # we only apply the mask if we have a large enough seq len - otherwise we return as is
+ if max_length < 2 * channel_codebooks - 1:
+ return input_ids.reshape(bsz * num_codebooks, -1), input_ids_shifted.reshape(bsz * num_codebooks, -1)
+
+ # fill the shifted ids with the prompt entries, offset by the codebook idx
+ for codebook in range(channel_codebooks):
+ if self.config.audio_channels == 1:
+ # mono channel - loop over the codebooks one-by-one
+ input_ids_shifted[:, codebook, codebook : seq_len + codebook] = input_ids[:, codebook]
+ else:
+ # left/right channels are interleaved in the generated codebooks, so handle one then the other
+ input_ids_shifted[:, 2 * codebook, codebook : seq_len + codebook] = input_ids[:, 2 * codebook]
+ input_ids_shifted[:, 2 * codebook + 1, codebook : seq_len + codebook] = input_ids[:, 2 * codebook + 1]
+
+ # construct a pattern mask that indicates the positions of padding tokens for each codebook
+ # first fill the upper triangular part (the EOS padding)
+ delay_pattern = torch.triu(
+ torch.ones((channel_codebooks, max_length), dtype=torch.bool), diagonal=max_length - channel_codebooks + 1
+ )
+ # then fill the lower triangular part (the BOS padding)
+ delay_pattern = delay_pattern + torch.tril(torch.ones((channel_codebooks, max_length), dtype=torch.bool))
+
+ if self.config.audio_channels == 2:
+ # for left/right channel we need to duplicate every row of the pattern mask in an interleaved fashion
+ delay_pattern = delay_pattern.repeat_interleave(2, dim=0)
+
+ mask = ~delay_pattern.to(input_ids.device)
+ input_ids = mask * input_ids_shifted + ~mask * pad_token_id
+
+ # find the first position to start generating - this is the first place we have the -1 token
+ # and will always be in the first codebook (since it has no codebook offset)
+ first_codebook_ids = input_ids[:, 0, :]
+ start_ids = (first_codebook_ids == -1).nonzero()[:, 1]
+ if len(start_ids) > 0:
+ first_start_id = min(start_ids)
+ else:
+ # we have no tokens that need to be filled - return entire matrix of input ids
+ first_start_id = seq_len
+
+ # (bsz * num_codebooks, seq_len) -> (bsz, num_codebooks, seq_len)
+ pattern_mask = input_ids.reshape(bsz * num_codebooks, -1)
+ input_ids = input_ids[..., :first_start_id].reshape(bsz * num_codebooks, -1)
+ return input_ids, pattern_mask
+
+ @staticmethod
+ def apply_delay_pattern_mask(input_ids, decoder_pad_token_mask):
+ """Apply a delay pattern mask to the decoder input ids, only preserving predictions where
+ the mask is set to -1, and otherwise setting to the value detailed in the mask."""
+ seq_len = input_ids.shape[-1]
+ decoder_pad_token_mask = decoder_pad_token_mask[..., :seq_len]
+ input_ids = torch.where(decoder_pad_token_mask == -1, input_ids, decoder_pad_token_mask)
+ return input_ids
+
+ @torch.no_grad()
+ def generate(
+ self,
+ inputs: torch.Tensor | None = None,
+ generation_config: GenerationConfig | None = None,
+ logits_processor: LogitsProcessorList | None = None,
+ stopping_criteria: StoppingCriteriaList | None = None,
+ synced_gpus: bool | None = None,
+ streamer: Optional["BaseStreamer"] = None,
+ **kwargs,
+ ):
+ """
+
+ Generates sequences of token ids for models with a language modeling head.
+
+
+
+ Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the
+ model's default generation configuration. You can override any `generation_config` by passing the corresponding
+ parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`.
+
+ For an overview of generation strategies and code examples, check out the [following
+ guide](./generation_strategies).
+
+
+
+ Parameters:
+ inputs (`torch.Tensor` of varying shape depending on the modality, *optional*):
+ The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the
+ method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`
+ should be in the format `input_ids`. For encoder-decoder models *inputs* can represent any of
+ `input_ids`, `input_values`, `input_features`, or `pixel_values`.
+ generation_config (`~generation.GenerationConfig`, *optional*):
+ The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+ passed to generate matching the attributes of `generation_config` will override them. If
+ `generation_config` is not provided, the default will be used, which had the following loading
+ priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+ configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+ default values, whose documentation should be checked to parameterize generation.
+ logits_processor (`LogitsProcessorList`, *optional*):
+ Custom logits processors that complement the default logits processors built from arguments and
+ generation config. If a logit processor is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ stopping_criteria (`StoppingCriteriaList`, *optional*):
+ Custom stopping criteria that complement the default stopping criteria built from arguments and a
+ generation config. If a stopping criteria is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ synced_gpus (`bool`, *optional*, defaults to `False`):
+ Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+ `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+ streamer (`BaseStreamer`, *optional*):
+ Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+ through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+ kwargs (`dict[str, Any]`, *optional*):
+ Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
+ forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder
+ specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.
+
+ Return:
+ [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True`
+ or when `config.return_dict_in_generate=True`) or a `torch.FloatTensor`.
+
+ If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateDecoderOnlyOutput`],
+ - [`~generation.GenerateBeamDecoderOnlyOutput`]
+
+ If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateEncoderDecoderOutput`],
+ - [`~generation.GenerateBeamEncoderDecoderOutput`]
+ """
+ # 1. Handle `generation_config` and kwargs that might update it, and validate the resulting objects
+ if generation_config is None:
+ generation_config = self.generation_config
+
+ generation_config = copy.deepcopy(generation_config)
+ model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs
+ generation_config.validate()
+ self._validate_model_kwargs(model_kwargs.copy())
+
+ # 2. Set generation parameters if not already defined
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+
+ requires_attention_mask = "encoder_outputs" not in model_kwargs
+ kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None
+
+ # 3. Define model inputs`
+ input_ids, model_input_name, model_kwargs = self._prepare_model_inputs(
+ inputs, generation_config.bos_token_id, model_kwargs
+ )
+ batch_size = input_ids.shape[0] // self.num_codebooks
+ self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=input_ids.device)
+
+ # 4. Define other model kwargs
+ model_kwargs["use_cache"] = generation_config.use_cache
+ model_kwargs["guidance_scale"] = generation_config.guidance_scale
+
+ if model_kwargs.get("attention_mask", None) is None and requires_attention_mask:
+ model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
+ input_ids, generation_config, model_kwargs
+ )
+
+ # 5. Prepare `max_length` depending on other stopping criteria.
+ input_ids_length = input_ids.shape[-1]
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+ has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
+ generation_config = self._prepare_generated_length(
+ generation_config=generation_config,
+ has_default_max_length=has_default_max_length,
+ has_default_min_length=has_default_min_length,
+ model_input_name=model_input_name,
+ inputs_tensor=input_ids,
+ input_ids_length=input_ids_length,
+ )
+
+ self._validate_generated_length(generation_config, input_ids_length, has_default_max_length)
+
+ # 6. Prepare the cache.
+ # - `model_kwargs` may be updated in place with a cache as defined by the parameters in `generation_config`.
+ # - different models have a different cache name expected by the model (default = "past_key_values")
+ # - `max_length`, prepared above, is used to determine the maximum cache length
+ max_cache_length = generation_config.max_length - 1
+ if (
+ input_ids_length.shape[1] != input_ids_length
+ and model_input_name == "inputs_embeds"
+ and not self.config.is_encoder_decoder
+ ):
+ max_cache_length += input_ids_length.shape[1]
+ self._prepare_cache_for_generation(
+ generation_config,
+ model_kwargs,
+ generation_mode=None,
+ batch_size=batch_size,
+ max_cache_length=max_cache_length,
+ )
+
+ # 7. Prepare `input_ids` which will be used for auto-regressive generation
+ # Build the delay pattern mask for offsetting each codebook prediction by 1 (this behaviour is specific to MusicGen)
+ input_ids, delay_pattern_mask = self.build_delay_pattern_mask(
+ input_ids,
+ pad_token_id=generation_config._decoder_start_token_tensor,
+ max_length=generation_config.max_length,
+ )
+
+ if streamer is not None:
+ streamer.put(input_ids.cpu())
+
+ # stash the delay mask so that we don't have to recompute it in each forward pass
+ model_kwargs["delay_pattern_mask"] = delay_pattern_mask
+
+ # 8. determine generation mode
+ generation_mode = generation_config.get_generation_mode()
+
+ # 9. prepare batched CFG externally (to enable coexistence with the unbatched CFG)
+ if generation_config.guidance_scale is not None and generation_config.guidance_scale > 1:
+ logits_processor.append(ClassifierFreeGuidanceLogitsProcessor(generation_config.guidance_scale))
+ generation_config.guidance_scale = None
+
+ # 10. prepare distribution pre_processing samplers
+ logits_processor = self._get_logits_processor(
+ generation_config=generation_config,
+ input_ids_seq_length=input_ids_length,
+ encoder_input_ids=input_ids,
+ prefix_allowed_tokens_fn=None,
+ logits_processor=logits_processor,
+ device=input_ids.device,
+ )
+
+ # 10. prepare stopping criteria
+ stopping_criteria = self._get_stopping_criteria(
+ generation_config=generation_config, stopping_criteria=stopping_criteria
+ )
+
+ if generation_mode in (GenerationMode.SAMPLE, GenerationMode.GREEDY_SEARCH):
+ # expand input_ids with `num_return_sequences` additional sequences per batch
+ input_ids, model_kwargs = self._expand_inputs_for_generation(
+ input_ids=input_ids,
+ expand_size=generation_config.num_return_sequences,
+ **model_kwargs,
+ )
+
+ # 11. run sample
+ outputs = self._sample(
+ input_ids,
+ logits_processor=logits_processor,
+ stopping_criteria=stopping_criteria,
+ generation_config=generation_config,
+ synced_gpus=synced_gpus,
+ streamer=streamer,
+ **model_kwargs,
+ )
+
+ else:
+ raise ValueError(
+ "Got incompatible mode for generation, should be one of greedy or sampling. "
+ "Ensure that beam search is de-activated by setting `num_beams=1`."
+ )
+
+ if generation_config.return_dict_in_generate:
+ output_ids = outputs.sequences
+ else:
+ output_ids = outputs
+
+ # apply the pattern mask to the final ids
+ output_ids = self.apply_delay_pattern_mask(output_ids, model_kwargs["delay_pattern_mask"])
+
+ # revert the pattern delay mask by filtering the pad token id
+ output_ids = output_ids[output_ids != generation_config._pad_token_tensor].reshape(
+ batch_size, self.num_codebooks, -1
+ )
+
+ if generation_config.return_dict_in_generate:
+ outputs.sequences = output_ids
+ return outputs
+ else:
+ return output_ids
+
+
+@auto_docstring(
+ custom_intro="""
+ The composite MusicGen model with a text encoder, audio encoder and Musicgen decoder,
+ """
+)
+class MusicgenForConditionalGeneration(MusicgenPreTrainedModel, GenerationMixin):
+ config: MusicgenConfig
+ output_modalities = ("audio",)
+ base_model_prefix = "encoder_decoder"
+ main_input_name = "input_ids"
+ supports_gradient_checkpointing = True
+
+ def __init__(
+ self,
+ config: MusicgenConfig | None = None,
+ text_encoder: PreTrainedModel | None = None,
+ audio_encoder: PreTrainedModel | None = None,
+ decoder: MusicgenForCausalLM | None = None,
+ ):
+ r"""
+ text_encoder (`PreTrainedModel`, *optional*):
+ The text encoder model that encodes text into hidden states for conditioning.
+ audio_encoder (`PreTrainedModel`, *optional*):
+ The audio encoder model that encodes audio into hidden states for conditioning.
+ decoder (`MusicgenForCausalLM`, *optional*):
+ The decoder model that generates audio tokens based on conditioning signals.
+ """
+ if config is None and (text_encoder is None or audio_encoder is None or decoder is None):
+ raise ValueError(
+ "Either a configuration has to be provided, or all three of text encoder, audio encoder and MusicGen decoder."
+ )
+ if config is None:
+ config = MusicgenConfig(
+ text_encoder=text_encoder.config, audio_encoder=audio_encoder.config, decoder=decoder.config
+ )
+ else:
+ if not isinstance(config, self.config_class):
+ raise ValueError(f"Config: {config} has to be of type {self.config_class}")
+
+ if config.decoder.cross_attention_hidden_size is not None:
+ if config.decoder.cross_attention_hidden_size != config.text_encoder.hidden_size:
+ raise ValueError(
+ "If `cross_attention_hidden_size` is specified in the MusicGen decoder's configuration, it has to be equal"
+ f" to the text encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"
+ f" `config.decoder.cross_attention_hidden_size` and {config.text_encoder.hidden_size} for"
+ " `config.text_encoder.hidden_size`."
+ )
+
+ # initialize with config
+ super().__init__(config)
+
+ if text_encoder is None:
+ from ..auto.modeling_auto import AutoModelForTextEncoding
+
+ text_encoder = AutoModelForTextEncoding.from_config(config.text_encoder)
+
+ if audio_encoder is None:
+ from ..auto.modeling_auto import AutoModel
+
+ audio_encoder = AutoModel.from_config(config.audio_encoder)
+
+ if decoder is None:
+ decoder = MusicgenForCausalLM._from_config(config.decoder)
+
+ self.text_encoder = text_encoder
+ self.audio_encoder = audio_encoder
+ self.decoder = decoder
+
+ if self.text_encoder.config.to_dict() != self.config.text_encoder.to_dict():
+ logger.warning(
+ f"Config of the text_encoder: {self.text_encoder.__class__} is overwritten by shared text_encoder config:"
+ f" {self.config.text_encoder}"
+ )
+ if self.audio_encoder.config.to_dict() != self.config.audio_encoder.to_dict():
+ logger.warning(
+ f"Config of the audio_encoder: {self.audio_encoder.__class__} is overwritten by shared audio_encoder config:"
+ f" {self.config.audio_encoder}"
+ )
+ if self.decoder.config.to_dict() != self.config.decoder.to_dict():
+ logger.warning(
+ f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config:"
+ f" {self.config.decoder}"
+ )
+
+ # make sure that the individual model's config refers to the shared config
+ # so that the updates to the config will be synced
+ self.config.text_encoder._attn_implementation = self.text_encoder.config._attn_implementation
+ self.config.audio_encoder._attn_implementation = self.audio_encoder.config._attn_implementation
+ self.config.decoder._attn_implementation = self.decoder.config._attn_implementation
+ self.text_encoder.config = self.config.text_encoder
+ self.audio_encoder.config = self.config.audio_encoder
+ self.decoder.config = self.config.decoder
+
+ # text encoder outputs might need to be projected to different dimension for decoder
+ if (
+ self.text_encoder.config.hidden_size != self.decoder.config.hidden_size
+ and self.decoder.config.cross_attention_hidden_size is None
+ ):
+ self.enc_to_dec_proj = nn.Linear(self.text_encoder.config.hidden_size, self.decoder.config.hidden_size)
+
+ if self.text_encoder.get_output_embeddings() is not None:
+ raise ValueError(
+ f"The encoder {self.text_encoder} should not have a LM Head. Please use a model without and LM Head"
+ )
+
+ decoder_signature = set(inspect.signature(self.decoder.forward).parameters.keys())
+ if "encoder_hidden_states" not in decoder_signature:
+ raise ValueError(
+ "The selected decoder is not prepared for the encoder hidden states to be passed. Please see the "
+ "following discussion on GitHub: https://github.com/huggingface/transformers/issues/23350"
+ )
+
+ # tie text encoder, decoder weights if config set accordingly
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.text_encoder.get_input_embeddings()
+
+ def get_output_embeddings(self):
+ return self.decoder.get_output_embeddings()
+
+ def set_output_embeddings(self, new_embeddings):
+ return self.decoder.set_output_embeddings(new_embeddings)
+
+ @classmethod
+ def from_sub_models_pretrained(
+ cls,
+ text_encoder_pretrained_model_name_or_path: str | None = None,
+ audio_encoder_pretrained_model_name_or_path: str | None = None,
+ decoder_pretrained_model_name_or_path: str | None = None,
+ *model_args,
+ **kwargs,
+ ) -> PreTrainedModel:
+ r"""
+ Instantiate a text encoder, an audio encoder, and a MusicGen decoder from one, two or three base classes of the
+ library from pretrained model checkpoints.
+
+
+ The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train
+ the model, you need to first set it back in training mode with `model.train()`.
+
+ Params:
+ text_encoder_pretrained_model_name_or_path (`str`, *optional*):
+ Information necessary to initiate the text encoder. Can be either:
+
+ - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
+ - A path to a *directory* containing model weights saved using
+ [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+
+ audio_encoder_pretrained_model_name_or_path (`str`, *optional*):
+ Information necessary to initiate the audio encoder. Can be either:
+
+ - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
+ - A path to a *directory* containing model weights saved using
+ [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+
+ decoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`):
+ Information necessary to initiate the decoder. Can be either:
+
+ - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
+ - A path to a *directory* containing model weights saved using
+ [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+
+ model_args (remaining positional arguments, *optional*):
+ All remaining positional arguments will be passed to the underlying model's `__init__` method.
+
+ kwargs (remaining dictionary of keyword arguments, *optional*):
+ Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
+ `output_attentions=True`).
+
+ - To update the text encoder configuration, use the prefix *text_encoder_* for each configuration
+ parameter.
+ - To update the audio encoder configuration, use the prefix *audio_encoder_* for each configuration
+ parameter.
+ - To update the decoder configuration, use the prefix *decoder_* for each configuration parameter.
+ - To update the parent model configuration, do not use a prefix for each configuration parameter.
+
+ Behaves differently depending on whether a `config` is provided or automatically loaded.
+
+ Example:
+
+ ```python
+ >>> from transformers import MusicgenForConditionalGeneration
+
+ >>> # initialize a musicgen model from a t5 text encoder, encodec audio encoder, and musicgen decoder
+ >>> model = MusicgenForConditionalGeneration.from_sub_models_pretrained(
+ ... text_encoder_pretrained_model_name_or_path="google-t5/t5-base",
+ ... audio_encoder_pretrained_model_name_or_path="facebook/encodec_24khz",
+ ... decoder_pretrained_model_name_or_path="facebook/musicgen-small",
+ ... )
+ >>> # saving model after fine-tuning
+ >>> model.save_pretrained("./musicgen-ft")
+ >>> # load fine-tuned model
+ >>> model = MusicgenForConditionalGeneration.from_pretrained("./musicgen-ft")
+ ```"""
+
+ kwargs_text_encoder = {
+ argument[len("text_encoder_") :]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("text_encoder_")
+ }
+
+ kwargs_audio_encoder = {
+ argument[len("audio_encoder_") :]: value
+ for argument, value in kwargs.items()
+ if argument.startswith("audio_encoder_")
+ }
+
+ kwargs_decoder = {
+ argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+ }
+
+ # remove text encoder, audio encoder and decoder kwargs from kwargs
+ for key in kwargs_text_encoder:
+ del kwargs["text_encoder_" + key]
+ for key in kwargs_audio_encoder:
+ del kwargs["audio_encoder_" + key]
+ for key in kwargs_decoder:
+ del kwargs["decoder_" + key]
+
+ # Load and initialize the encoder and decoder
+ # The distinction between encoder and decoder at the model level is made
+ # by the value of the flag `is_decoder` that we need to set correctly.
+ text_encoder = kwargs_text_encoder.pop("model", None)
+ if text_encoder is None:
+ if text_encoder_pretrained_model_name_or_path is None:
+ raise ValueError(
+ "If `text_encoder_model` is not defined as an argument, a `text_encoder_pretrained_model_name_or_path` has "
+ "to be defined."
+ )
+
+ if "config" not in kwargs_text_encoder:
+ encoder_config, kwargs_text_encoder = AutoConfig.from_pretrained(
+ text_encoder_pretrained_model_name_or_path, **kwargs_text_encoder, return_unused_kwargs=True
+ )
+
+ if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
+ logger.info(
+ f"Initializing {text_encoder_pretrained_model_name_or_path} as a text_encoder model "
+ "from a decoder model. Cross-attention and causal mask are disabled."
+ )
+ encoder_config.is_decoder = False
+ encoder_config.add_cross_attention = False
+
+ kwargs_text_encoder["config"] = encoder_config
+
+ text_encoder = AutoModel.from_pretrained(
+ text_encoder_pretrained_model_name_or_path, *model_args, **kwargs_text_encoder
+ )
+
+ audio_encoder = kwargs_audio_encoder.pop("model", None)
+ if audio_encoder is None:
+ if audio_encoder_pretrained_model_name_or_path is None:
+ raise ValueError(
+ "If `audio_encoder_model` is not defined as an argument, an `audio_encoder_pretrained_model_name_or_path` has "
+ "to be defined."
+ )
+
+ if "config" not in kwargs_audio_encoder:
+ encoder_config, kwargs_audio_encoder = AutoConfig.from_pretrained(
+ audio_encoder_pretrained_model_name_or_path, **kwargs_audio_encoder, return_unused_kwargs=True
+ )
+
+ if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
+ logger.info(
+ f"Initializing {audio_encoder_pretrained_model_name_or_path} as an audio_encoder model "
+ "from a decoder model. Cross-attention and causal mask are disabled."
+ )
+ encoder_config.is_decoder = False
+ encoder_config.add_cross_attention = False
+
+ kwargs_audio_encoder["config"] = encoder_config
+
+ audio_encoder = AutoModel.from_pretrained(
+ audio_encoder_pretrained_model_name_or_path, *model_args, **kwargs_audio_encoder
+ )
+
+ decoder = kwargs_decoder.pop("model", None)
+ if decoder is None:
+ if decoder_pretrained_model_name_or_path is None:
+ raise ValueError(
+ "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
+ "to be defined."
+ )
+
+ if "config" not in kwargs_decoder:
+ decoder_config, kwargs_decoder = AutoConfig.from_pretrained(
+ decoder_pretrained_model_name_or_path, **kwargs_decoder, return_unused_kwargs=True
+ )
+
+ if isinstance(decoder_config, MusicgenConfig):
+ decoder_config = decoder_config.decoder
+
+ if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:
+ logger.info(
+ f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"
+ f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"
+ f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."
+ )
+ decoder_config.is_decoder = True
+ decoder_config.add_cross_attention = True
+
+ kwargs_decoder["config"] = decoder_config
+
+ if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:
+ logger.warning(
+ f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "
+ f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "
+ "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "
+ "passed to `.from_sub_models_pretrained(...)` are set to `True` or do not pass a "
+ "`decoder_config` to `.from_sub_models_pretrained(...)`"
+ )
+
+ decoder = MusicgenForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)
+
+ # instantiate config with corresponding kwargs
+ config = MusicgenConfig(
+ text_encoder=text_encoder.config, audio_encoder=audio_encoder.config, decoder=decoder.config, **kwargs
+ )
+ return cls(text_encoder=text_encoder, audio_encoder=audio_encoder, decoder=decoder, config=config)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.BoolTensor | None = None,
+ input_values: torch.FloatTensor | None = None,
+ padding_mask: torch.BoolTensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.BoolTensor | None = None,
+ encoder_outputs: tuple[torch.FloatTensor] | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | Seq2SeqLMOutput:
+ r"""
+ padding_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+
+ [What are attention masks?](../glossary#attention-mask)
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size * num_codebooks, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary, corresponding to the sequence of audio codes.
+
+ Indices can be obtained by encoding an audio prompt with an audio encoder model to predict audio codes,
+ such as with the [`EncodecModel`]. See [`EncodecModel.encode`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+
+
+ The `decoder_input_ids` will automatically be converted from shape `(batch_size * num_codebooks,
+ target_sequence_length)` to `(batch_size, num_codebooks, target_sequence_length)` in the forward pass. If
+ you obtain audio codes from an audio encoding model, such as [`EncodecModel`], ensure that the number of
+ frames is equal to 1, and that you reshape the audio codes from `(frames, batch_size, num_codebooks,
+ target_sequence_length)` to `(batch_size * num_codebooks, target_sequence_length)` prior to passing them as
+ `decoder_input_ids`.
+
+
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks)`, *optional*):
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+ `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+ are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+
+ Examples:
+ ```python
+ >>> from transformers import AutoProcessor, MusicgenForConditionalGeneration
+ >>> import torch
+
+ >>> processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+ >>> model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+
+ >>> inputs = processor(
+ ... text=["80s pop track with bassy drums and synth", "90s rock song with loud guitars and heavy drums"],
+ ... padding=True,
+ ... return_tensors="pt",
+ ... )
+
+ >>> pad_token_id = model.generation_config.pad_token_id
+ >>> decoder_input_ids = (
+ ... torch.ones((inputs.input_ids.shape[0] * model.decoder.num_codebooks, 1), dtype=torch.long)
+ ... * pad_token_id
+ ... )
+
+ >>> logits = model(**inputs, decoder_input_ids=decoder_input_ids).logits
+ >>> logits.shape # (bsz * num_codebooks, tgt_len, vocab_size)
+ torch.Size([8, 1, 2048])
+ ```"""
+ kwargs_text_encoder = {}
+ kwargs_audio_encoder = {}
+ kwargs_decoder = {}
+ common_kwargs = {}
+ for key, value in kwargs.items():
+ if key.startswith("text_encoder_"):
+ kwargs_text_encoder[key[len("text_encoder_") :]] = value
+ elif key.startswith("audio_encoder_"):
+ kwargs_audio_encoder[key[len("audio_encoder_") :]] = value
+ elif key.startswith("decoder_"):
+ kwargs_decoder[key[len("decoder_") :]] = value
+ else:
+ common_kwargs[key] = value
+
+ if encoder_outputs is None:
+ encoder_outputs = self.text_encoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ **kwargs_text_encoder,
+ **common_kwargs,
+ )
+ elif isinstance(encoder_outputs, tuple):
+ encoder_outputs = BaseModelOutput(*encoder_outputs)
+
+ encoder_hidden_states = encoder_outputs[0]
+
+ # optionally project encoder_hidden_states
+ if (
+ self.text_encoder.config.hidden_size != self.decoder.config.hidden_size
+ and self.decoder.config.cross_attention_hidden_size is None
+ ):
+ encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
+
+ if attention_mask is not None:
+ encoder_hidden_states = encoder_hidden_states * attention_mask[..., None]
+
+ if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):
+ decoder_input_ids = shift_tokens_right(
+ labels, self.config.decoder.pad_token_id, self.config.decoder.decoder_start_token_id
+ )
+
+ elif decoder_input_ids is None and decoder_inputs_embeds is None:
+ audio_encoder_outputs = self.audio_encoder(
+ input_values=input_values,
+ padding_mask=padding_mask,
+ **kwargs_audio_encoder,
+ )
+ audio_codes = audio_encoder_outputs.audio_codes
+ frames, bsz, codebooks, seq_len = audio_codes.shape
+ if frames != 1:
+ raise ValueError(
+ f"Expected 1 frame in the audio code outputs, got {frames} frames. Ensure chunking is "
+ "disabled by setting `chunk_length=None` in the audio encoder."
+ )
+
+ if self.config.decoder.audio_channels == 2 and audio_codes.shape[2] == self.decoder.num_codebooks // 2:
+ # mono input through encodec that we convert to stereo
+ audio_codes = audio_codes.repeat_interleave(2, dim=2)
+
+ decoder_input_ids = audio_codes[0, ...].reshape(bsz * self.decoder.num_codebooks, seq_len)
+
+ # Decode
+ decoder_outputs: CausalLMOutputWithCrossAttentions = self.decoder(
+ input_ids=decoder_input_ids,
+ attention_mask=decoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=attention_mask,
+ inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ past_key_values=past_key_values,
+ labels=labels,
+ **kwargs_decoder,
+ **common_kwargs,
+ )
+
+ return Seq2SeqLMOutput(
+ loss=decoder_outputs.loss,
+ logits=decoder_outputs.logits,
+ past_key_values=decoder_outputs.past_key_values,
+ decoder_hidden_states=decoder_outputs.hidden_states,
+ decoder_attentions=decoder_outputs.attentions,
+ cross_attentions=decoder_outputs.cross_attentions,
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+ encoder_hidden_states=encoder_outputs.hidden_states,
+ encoder_attentions=encoder_outputs.attentions,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ decoder_input_ids,
+ next_sequence_length: int | None = None,
+ past_key_values=None,
+ attention_mask=None,
+ decoder_attention_mask=None,
+ use_cache=None,
+ encoder_outputs=None,
+ decoder_delay_pattern_mask=None,
+ guidance_scale=None,
+ **kwargs,
+ ):
+ # Overwritten -- MusicGen has custom processing
+ if decoder_delay_pattern_mask is None:
+ decoder_input_ids, decoder_delay_pattern_mask = self.decoder.build_delay_pattern_mask(
+ decoder_input_ids,
+ self.generation_config.pad_token_id,
+ max_length=self.generation_config.max_length,
+ )
+
+ # apply the delay pattern mask
+ decoder_input_ids = self.decoder.apply_delay_pattern_mask(decoder_input_ids, decoder_delay_pattern_mask)
+
+ if guidance_scale is not None and guidance_scale > 1:
+ # for classifier free guidance we need to replicate the decoder args across the batch dim (we'll split these
+ # before sampling)
+ decoder_input_ids = decoder_input_ids.repeat((2, 1))
+ if decoder_attention_mask is not None:
+ decoder_attention_mask = decoder_attention_mask.repeat((2, 1))
+
+ if past_key_values is not None:
+ decoder_input_ids = (
+ decoder_input_ids[:, -next_sequence_length:] if next_sequence_length is not None else decoder_input_ids
+ )
+
+ return {
+ "input_ids": None, # encoder_outputs is defined. input_ids not needed
+ "encoder_outputs": encoder_outputs,
+ "past_key_values": past_key_values,
+ "decoder_input_ids": decoder_input_ids,
+ "attention_mask": attention_mask,
+ "decoder_attention_mask": decoder_attention_mask,
+ "use_cache": use_cache,
+ }
+
+ def _prepare_decoder_input_ids_for_generation(
+ self,
+ batch_size: int,
+ model_input_name: str,
+ model_kwargs: dict[str, torch.Tensor],
+ decoder_start_token_id: int | None = None,
+ bos_token_id: int | None = None,
+ device: torch.device | None = None,
+ ) -> tuple[torch.LongTensor, dict[str, torch.Tensor]]:
+ """Prepares `decoder_input_ids` for generation with encoder-decoder models"""
+
+ # 1. Check whether the user has defined `decoder_input_ids` manually. To facilitate in terms of input naming,
+ # we also allow the user to pass it under `input_ids`, if the encoder does not use it as the main input.
+ if model_kwargs is not None and "decoder_input_ids" in model_kwargs:
+ decoder_input_ids = model_kwargs.pop("decoder_input_ids")
+ elif "input_ids" in model_kwargs and model_input_name != "input_ids":
+ decoder_input_ids = model_kwargs.pop("input_ids")
+ else:
+ decoder_input_ids = None
+
+ # 2. Encoder-decoder models expect the `decoder_input_ids` to start with a special token. Let's ensure that.
+ decoder_start_token_id = self._get_decoder_start_token_id(decoder_start_token_id, bos_token_id)
+ if device is None:
+ device = self.device
+ decoder_input_ids_start = (
+ torch.ones((batch_size * self.decoder.num_codebooks, 1), dtype=torch.long, device=device)
+ * decoder_start_token_id
+ )
+
+ # no user input -> use decoder_start_token_id as decoder_input_ids
+ if decoder_input_ids is None:
+ decoder_input_ids = decoder_input_ids_start
+
+ # user input but doesn't start with decoder_start_token_id -> prepend decoder_start_token_id (and adjust
+ # decoder_attention_mask if provided)
+ elif (decoder_input_ids[..., 0] != decoder_start_token_id).all().item():
+ decoder_input_ids = torch.cat([decoder_input_ids_start, decoder_input_ids], dim=-1)
+ if "decoder_attention_mask" in model_kwargs:
+ decoder_attention_mask = model_kwargs["decoder_attention_mask"]
+ decoder_attention_mask = torch.cat(
+ (torch.ones_like(decoder_attention_mask)[:, :1], decoder_attention_mask),
+ dim=-1,
+ )
+ model_kwargs["decoder_attention_mask"] = decoder_attention_mask
+
+ return decoder_input_ids, model_kwargs
+
+ def _prepare_text_encoder_kwargs_for_generation(
+ self,
+ inputs_tensor: torch.Tensor,
+ model_kwargs,
+ model_input_name: str | None,
+ generation_config: GenerationConfig,
+ ) -> dict[str, Any]:
+ # 1. get text encoder
+ encoder = self.get_encoder()
+ # Compatibility with Accelerate big model inference: we need the encoder to outputs stuff on the same device
+ # as the inputs.
+ if hasattr(encoder, "_hf_hook"):
+ encoder._hf_hook.io_same_device = True
+
+ # 2. Prepare encoder args and encoder kwargs from model kwargs.
+ irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"]
+ encoder_kwargs = {
+ argument: value
+ for argument, value in model_kwargs.items()
+ if not any(argument.startswith(p) for p in irrelevant_prefix)
+ }
+ encoder_signature = set(inspect.signature(encoder.forward).parameters)
+ encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature
+ if not encoder_accepts_wildcard:
+ encoder_kwargs = {
+ argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature
+ }
+ encoder_kwargs["output_attentions"] = generation_config.output_attentions
+ encoder_kwargs["output_hidden_states"] = generation_config.output_hidden_states
+ guidance_scale = generation_config.guidance_scale
+
+ # 3. make sure that encoder returns `ModelOutput`
+ model_input_name = model_input_name if model_input_name is not None else self.text_encoder.main_input_name
+ encoder_kwargs["return_dict"] = True
+ encoder_kwargs[model_input_name] = inputs_tensor
+ last_hidden_state = encoder(**encoder_kwargs).last_hidden_state
+
+ # for classifier free guidance we need to add a 'null' input to our encoder hidden states
+ if guidance_scale is not None and guidance_scale > 1:
+ last_hidden_state = torch.concatenate([last_hidden_state, torch.zeros_like(last_hidden_state)], dim=0)
+ if "attention_mask" in model_kwargs:
+ model_kwargs["attention_mask"] = torch.concatenate(
+ [model_kwargs["attention_mask"], torch.zeros_like(model_kwargs["attention_mask"])], dim=0
+ )
+
+ model_kwargs["encoder_outputs"] = BaseModelOutput(last_hidden_state=last_hidden_state)
+
+ return model_kwargs
+
+ def _prepare_audio_encoder_kwargs_for_generation(
+ self, input_values, model_kwargs, model_input_name: str | None = None
+ ):
+ # 1. get audio encoder
+ encoder = self.get_encoder(modality="audio")
+ # Compatibility with Accelerate big model inference: we need the encoder to outputs stuff on the same device
+ # as the inputs.
+ if hasattr(encoder, "_hf_hook"):
+ encoder._hf_hook.io_same_device = True
+
+ # 2. Prepare encoder args and encoder kwargs from model kwargs.
+ irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"]
+ encoder_kwargs = {
+ argument: value
+ for argument, value in model_kwargs.items()
+ if not any(argument.startswith(p) for p in irrelevant_prefix)
+ }
+ encoder_signature = set(inspect.signature(encoder.forward).parameters)
+ encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature
+ if not encoder_accepts_wildcard:
+ encoder_kwargs = {
+ argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature
+ }
+
+ # 3. make sure that encoder returns `ModelOutput`
+ model_input_name = model_input_name if model_input_name is not None else self.audio_encoder.main_input_name
+ encoder_kwargs["return_dict"] = True
+
+ if self.decoder.config.audio_channels == 1:
+ encoder_kwargs[model_input_name] = input_values
+ audio_encoder_outputs = encoder.encode(**encoder_kwargs)
+ audio_codes = audio_encoder_outputs.audio_codes
+ audio_scales = audio_encoder_outputs.audio_scales
+
+ frames, bsz, codebooks, seq_len = audio_codes.shape
+
+ else:
+ if input_values.shape[1] != 2:
+ raise ValueError(
+ f"Expected stereo audio (2-channels) but example has {input_values.shape[1]} channel."
+ )
+
+ encoder_kwargs[model_input_name] = input_values[:, :1, :]
+ audio_encoder_outputs_left = encoder.encode(**encoder_kwargs)
+ audio_codes_left = audio_encoder_outputs_left.audio_codes
+ audio_scales_left = audio_encoder_outputs_left.audio_scales
+
+ encoder_kwargs[model_input_name] = input_values[:, 1:, :]
+ audio_encoder_outputs_right = encoder.encode(**encoder_kwargs)
+ audio_codes_right = audio_encoder_outputs_right.audio_codes
+ audio_scales_right = audio_encoder_outputs_right.audio_scales
+
+ frames, bsz, codebooks, seq_len = audio_codes_left.shape
+ # copy alternating left/right channel codes into stereo codebook
+ audio_codes = audio_codes_left.new_ones((frames, bsz, 2 * codebooks, seq_len))
+
+ audio_codes[:, :, ::2, :] = audio_codes_left
+ audio_codes[:, :, 1::2, :] = audio_codes_right
+
+ if audio_scales_left != [None] or audio_scales_right != [None]:
+ audio_scales = torch.stack([audio_scales_left, audio_scales_right], dim=1)
+ else:
+ audio_scales = [None] * bsz
+
+ if frames != 1:
+ raise ValueError(
+ f"Expected 1 frame in the audio code outputs, got {frames} frames. Ensure chunking is "
+ "disabled by setting `chunk_length=None` in the audio encoder."
+ )
+
+ decoder_input_ids = audio_codes[0, ...].reshape(bsz * self.decoder.num_codebooks, seq_len)
+
+ model_kwargs["decoder_input_ids"] = decoder_input_ids
+ model_kwargs["audio_scales"] = audio_scales
+ return model_kwargs
+
+ def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+ return shift_tokens_right(labels, self.config.decoder.pad_token_id, self.config.decoder.bos_token_id)
+
+ def resize_token_embeddings(self, *args, **kwargs):
+ raise NotImplementedError(
+ "Resizing the embedding layers via the EncoderDecoderModel directly is not supported. Please use the"
+ " respective methods of the wrapped objects (model.encoder.resize_token_embeddings(...) or"
+ " model.decoder.resize_token_embeddings(...))"
+ )
+
+ def freeze_audio_encoder(self):
+ """
+ Freeze the audio encoder weights.
+ """
+ for param in self.audio_encoder.parameters():
+ param.requires_grad = False
+ self.audio_encoder._requires_grad = False
+
+ def freeze_text_encoder(self):
+ """
+ Freeze the text encoder weights.
+ """
+ for param in self.text_encoder.parameters():
+ param.requires_grad = False
+ self.text_encoder._requires_grad = False
+
+ def _maybe_initialize_input_ids_for_generation(
+ self,
+ inputs: torch.Tensor | None,
+ bos_token_id: int | None,
+ model_kwargs: dict[str, torch.Tensor],
+ ) -> torch.LongTensor:
+ """Initializes input ids for generation, if necessary."""
+ if inputs is not None:
+ return inputs
+
+ encoder_outputs = model_kwargs.get("encoder_outputs")
+ if encoder_outputs is not None:
+ # make dummy input_ids with value -100, as a sanity check ensuring that they won't be used for encoding
+ shape = encoder_outputs[0].size()[:-1]
+ return torch.ones(shape, dtype=torch.long, device=self.device) * -100
+
+ if bos_token_id is None:
+ raise ValueError("`bos_token_id` has to be defined when no `input_ids` are provided.")
+
+ # If there is some tensor in `model_kwargs`, we can infer the batch size from it. This is helpful with
+ # soft-prompting or in multimodal implementations built on top of decoder-only language models.
+ batch_size = 1
+ for value in model_kwargs.values():
+ if isinstance(value, torch.Tensor):
+ batch_size = value.shape[0]
+ break
+ return torch.ones((batch_size, 1), dtype=torch.long, device=self.device) * bos_token_id
+
+ def _get_decoder_start_token_id(
+ self, decoder_start_token_id: int | list[int] | None = None, bos_token_id: int | None = None
+ ) -> int:
+ decoder_start_token_id = (
+ decoder_start_token_id
+ if decoder_start_token_id is not None
+ else self.generation_config.decoder_start_token_id
+ )
+ bos_token_id = bos_token_id if bos_token_id is not None else self.generation_config.bos_token_id
+
+ if decoder_start_token_id is not None:
+ return decoder_start_token_id
+ elif bos_token_id is not None:
+ return bos_token_id
+ raise ValueError(
+ "`decoder_start_token_id` or `bos_token_id` has to be defined for encoder-decoder generation."
+ )
+
+ @torch.no_grad()
+ def generate(
+ self,
+ inputs: torch.Tensor | None = None,
+ generation_config: GenerationConfig | None = None,
+ logits_processor: LogitsProcessorList | None = None,
+ stopping_criteria: StoppingCriteriaList | None = None,
+ synced_gpus: bool | None = None,
+ streamer: Optional["BaseStreamer"] = None,
+ **kwargs,
+ ):
+ """
+
+ Generates sequences of token ids for models with a language modeling head.
+
+
+
+ Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the
+ model's default generation configuration. You can override any `generation_config` by passing the corresponding
+ parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`.
+
+ For an overview of generation strategies and code examples, check out the [following
+ guide](./generation_strategies).
+
+
+
+ Parameters:
+ inputs (`torch.Tensor` of varying shape depending on the modality, *optional*):
+ The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the
+ method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`
+ should be in the format `input_ids`. For encoder-decoder models *inputs* can represent any of
+ `input_ids`, `input_values`, `input_features`, or `pixel_values`.
+ generation_config (`~generation.GenerationConfig`, *optional*):
+ The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+ passed to generate matching the attributes of `generation_config` will override them. If
+ `generation_config` is not provided, the default will be used, which had the following loading
+ priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+ configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+ default values, whose documentation should be checked to parameterize generation.
+ logits_processor (`LogitsProcessorList`, *optional*):
+ Custom logits processors that complement the default logits processors built from arguments and
+ generation config. If a logit processor is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ stopping_criteria (`StoppingCriteriaList`, *optional*):
+ Custom stopping criteria that complement the default stopping criteria built from arguments and a
+ generation config. If a stopping criteria is passed that is already created with the arguments or a
+ generation config an error is thrown. This feature is intended for advanced users.
+ synced_gpus (`bool`, *optional*, defaults to `False`):
+ Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+ `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+ streamer (`BaseStreamer`, *optional*):
+ Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+ through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+ kwargs (`dict[str, Any]`, *optional*):
+ Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
+ forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder
+ specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.
+
+ Return:
+ [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True`
+ or when `config.return_dict_in_generate=True`) or a `torch.FloatTensor`.
+
+ If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateDecoderOnlyOutput`],
+ - [`~generation.GenerateBeamDecoderOnlyOutput`]
+
+ If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible
+ [`~utils.ModelOutput`] types are:
+
+ - [`~generation.GenerateEncoderDecoderOutput`],
+ - [`~generation.GenerateBeamEncoderDecoderOutput`]
+ """
+ # 1. Handle `generation_config` and kwargs that might update it, and validate the resulting objects
+ generation_mode_kwargs = self._extract_generation_mode_kwargs(None, kwargs, False, None, None)
+ generation_config, model_kwargs = self._prepare_generation_config(generation_config, **kwargs)
+ generation_mode = generation_config.get_generation_mode()
+ if generation_mode not in [GenerationMode.SAMPLE, GenerationMode.GREEDY_SEARCH]:
+ raise ValueError(
+ "Got incompatible mode for generation, should be one of greedy or sampling. "
+ "Ensure that beam search is de-activated by setting `num_beams=1`."
+ )
+
+ self._validate_model_kwargs(model_kwargs.copy())
+ self._validate_generation_mode(generation_mode, generation_config, generation_mode_kwargs)
+
+ if model_kwargs.get("encoder_outputs") is not None and type(model_kwargs["encoder_outputs"]) is tuple:
+ # wrap the unconditional outputs as a BaseModelOutput for compatibility with the rest of generate
+ model_kwargs["encoder_outputs"] = BaseModelOutput(last_hidden_state=model_kwargs["encoder_outputs"][0])
+
+ # 2. Set generation parameters if not already defined
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+
+ requires_attention_mask = "encoder_outputs" not in model_kwargs
+ kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None
+
+ # 3. Define model inputs
+ inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
+ inputs, generation_config.bos_token_id, model_kwargs
+ )
+ batch_size = inputs_tensor.shape[0]
+ self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=inputs_tensor.device)
+
+ # 4. Define other model kwargs
+ model_kwargs["use_cache"] = generation_config.use_cache
+ model_kwargs["guidance_scale"] = generation_config.guidance_scale
+
+ if model_kwargs.get("attention_mask", None) is None and requires_attention_mask:
+ model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
+ inputs_tensor, generation_config, model_kwargs
+ )
+
+ if "encoder_outputs" not in model_kwargs:
+ # encoder_outputs are created and added to `model_kwargs`
+ model_kwargs = self._prepare_text_encoder_kwargs_for_generation(
+ inputs_tensor, model_kwargs, model_input_name, generation_config
+ )
+
+ if "decoder_input_ids" not in model_kwargs and "input_values" in model_kwargs:
+ model_kwargs = self._prepare_audio_encoder_kwargs_for_generation(
+ model_kwargs["input_values"],
+ model_kwargs,
+ )
+
+ # 5. Prepare `input_ids` which will be used for auto-regressive generation
+ input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation(
+ batch_size=batch_size,
+ model_input_name=model_input_name,
+ model_kwargs=model_kwargs,
+ decoder_start_token_id=generation_config._decoder_start_token_tensor,
+ bos_token_id=generation_config._bos_token_tensor,
+ device=inputs_tensor.device,
+ )
+
+ # 6. Prepare `max_length` depending on other stopping criteria.
+ input_ids_length = input_ids.shape[-1]
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+ has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
+ generation_config = self._prepare_generated_length(
+ generation_config=generation_config,
+ has_default_max_length=has_default_max_length,
+ has_default_min_length=has_default_min_length,
+ model_input_name=model_input_name,
+ inputs_tensor=inputs_tensor,
+ input_ids_length=input_ids_length,
+ )
+
+ # build the delay pattern mask for offsetting each codebook prediction by 1 (this behaviour is specific to MusicGen)
+ input_ids, decoder_delay_pattern_mask = self.decoder.build_delay_pattern_mask(
+ input_ids,
+ pad_token_id=generation_config._decoder_start_token_tensor,
+ max_length=generation_config.max_length,
+ )
+ # stash the delay mask so that we don't have to recompute in each forward pass
+ model_kwargs["decoder_delay_pattern_mask"] = decoder_delay_pattern_mask
+
+ # input_ids are ready to be placed on the streamer (if used)
+ if streamer is not None:
+ streamer.put(input_ids.cpu())
+
+ # 7. determine generation mode
+ generation_mode = generation_config.get_generation_mode()
+
+ # 8. prepare batched CFG externally (to enable coexistence with the unbatched CFG)
+ if generation_config.guidance_scale is not None and generation_config.guidance_scale > 1:
+ logits_processor.append(ClassifierFreeGuidanceLogitsProcessor(generation_config.guidance_scale))
+ generation_config.guidance_scale = None
+
+ # 9. prepare distribution pre_processing samplers
+ logits_processor = self._get_logits_processor(
+ generation_config=generation_config,
+ input_ids_seq_length=input_ids_length,
+ encoder_input_ids=inputs_tensor,
+ prefix_allowed_tokens_fn=None,
+ logits_processor=logits_processor,
+ device=input_ids.device,
+ )
+
+ # 10. prepare stopping criteria
+ stopping_criteria = self._get_stopping_criteria(
+ generation_config=generation_config, stopping_criteria=stopping_criteria
+ )
+
+ # expand input_ids with `num_return_sequences` additional sequences per batch
+ input_ids, model_kwargs = self._expand_inputs_for_generation(
+ input_ids=input_ids,
+ expand_size=generation_config.num_return_sequences,
+ is_encoder_decoder=self.config.is_encoder_decoder,
+ **model_kwargs,
+ )
+
+ # 10b. prepare prefill outputs
+ generation_mode_kwargs["prefill_outputs"] = self._prefill(input_ids, generation_config, model_kwargs)
+
+ # 11. run sample
+ outputs = self._sample(
+ input_ids,
+ logits_processor=logits_processor,
+ stopping_criteria=stopping_criteria,
+ generation_config=generation_config,
+ **generation_mode_kwargs,
+ **model_kwargs,
+ )
+
+ if generation_config.return_dict_in_generate:
+ output_ids = outputs.sequences
+ else:
+ output_ids = outputs
+
+ # apply the pattern mask to the final ids
+ output_ids = self.decoder.apply_delay_pattern_mask(output_ids, model_kwargs["decoder_delay_pattern_mask"])
+
+ # revert the pattern delay mask by filtering the pad token id
+ output_ids = output_ids[output_ids != generation_config._pad_token_tensor].reshape(
+ batch_size, self.decoder.num_codebooks, -1
+ )
+
+ # append the frame dimension back to the audio codes
+ output_ids = output_ids[None, ...]
+
+ audio_scales = model_kwargs.get("audio_scales")
+ if audio_scales is None:
+ audio_scales = [None] * batch_size
+
+ if self.decoder.config.audio_channels == 1:
+ output_values = self.audio_encoder.decode(
+ output_ids,
+ audio_scales=audio_scales,
+ ).audio_values
+ else:
+ codec_outputs_left = self.audio_encoder.decode(output_ids[:, :, ::2, :], audio_scales=audio_scales)
+ output_values_left = codec_outputs_left.audio_values
+
+ codec_outputs_right = self.audio_encoder.decode(output_ids[:, :, 1::2, :], audio_scales=audio_scales)
+ output_values_right = codec_outputs_right.audio_values
+
+ output_values = torch.cat([output_values_left, output_values_right], dim=1)
+
+ if generation_config.return_dict_in_generate:
+ outputs.sequences = output_values
+ return outputs
+ else:
+ return output_values
+
+ def get_unconditional_inputs(self, num_samples=1):
+ """
+ Helper function to get null inputs for unconditional generation, enabling the model to be used without the
+ feature extractor or tokenizer.
+
+ Args:
+ num_samples (int, *optional*):
+ Number of audio samples to unconditionally generate.
+ max_new_tokens (int, *optional*):
+ Number of tokens to generate for each sample. More tokens means longer audio samples, at the expense of
+ longer inference (since more audio tokens need to be generated per sample).
+
+ Example:
+ ```python
+ >>> from transformers import MusicgenForConditionalGeneration
+
+ >>> model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+
+ >>> # get the unconditional (or 'null') inputs for the model
+ >>> unconditional_inputs = model.get_unconditional_inputs(num_samples=1)
+ >>> audio_samples = model.generate(**unconditional_inputs, max_new_tokens=256)
+ ```"""
+ last_hidden_state = torch.zeros(
+ (num_samples, 1, self.config.text_encoder.hidden_size), device=self.device, dtype=self.dtype
+ )
+
+ attention_mask = torch.zeros((num_samples, 1), device=self.device, dtype=torch.long)
+
+ return MusicgenUnconditionalInput(
+ encoder_outputs=(last_hidden_state,),
+ attention_mask=attention_mask,
+ guidance_scale=1.0,
+ )
+
+
+__all__ = ["MusicgenForConditionalGeneration", "MusicgenForCausalLM", "MusicgenModel", "MusicgenPreTrainedModel"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/processing_musicgen.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/processing_musicgen.py
new file mode 100644
index 0000000000000000000000000000000000000000..180f0a2a95eecc662f71a5b53d553b832fd8ad88
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/musicgen/processing_musicgen.py
@@ -0,0 +1,86 @@
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Text/audio processor class for MusicGen
+"""
+
+from typing import Any
+
+import numpy as np
+
+from ...processing_utils import ProcessorMixin
+from ...utils import auto_docstring, to_numpy
+
+
+@auto_docstring
+class MusicgenProcessor(ProcessorMixin):
+ def __init__(self, feature_extractor, tokenizer):
+ super().__init__(feature_extractor, tokenizer)
+
+ def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
+ return self.tokenizer.get_decoder_prompt_ids(task=task, language=language, no_timestamps=no_timestamps)
+
+ @auto_docstring
+ def __call__(self, *args, **kwargs):
+ if len(args) > 0:
+ kwargs["audio"] = args[0]
+ return super().__call__(*args, **kwargs)
+
+ def batch_decode(self, *args, **kwargs):
+ """
+ This method is used to decode either batches of audio outputs from the MusicGen model, or batches of token ids
+ from the tokenizer. In the case of decoding token ids, this method forwards all its arguments to T5Tokenizer's
+ [`~PreTrainedTokenizer.batch_decode`]. Please refer to the docstring of this method for more information.
+ """
+ audio_values = kwargs.pop("audio", None)
+ padding_mask = kwargs.pop("padding_mask", None)
+
+ if len(args) > 0:
+ audio_values = args[0]
+ args = args[1:]
+
+ if audio_values is not None:
+ return self._decode_audio(audio_values, padding_mask=padding_mask)
+ else:
+ return self.tokenizer.batch_decode(*args, **kwargs)
+
+ def _decode_audio(self, audio_values, padding_mask: Any = None) -> list[np.ndarray]:
+ """
+ This method strips any padding from the audio values to return a list of numpy audio arrays.
+ """
+ audio_values = to_numpy(audio_values)
+ bsz, channels, seq_len = audio_values.shape
+
+ if padding_mask is None:
+ return list(audio_values)
+
+ padding_mask = to_numpy(padding_mask)
+
+ # match the sequence length of the padding mask to the generated audio arrays by padding with the **non-padding**
+ # token (so that the generated audio values are **not** treated as padded tokens)
+ difference = seq_len - padding_mask.shape[-1]
+ padding_value = 1 - self.feature_extractor.padding_value
+ padding_mask = np.pad(padding_mask, ((0, 0), (0, difference)), "constant", constant_values=padding_value)
+
+ audio_values = audio_values.tolist()
+ for i in range(bsz):
+ sliced_audio = np.asarray(audio_values[i])[
+ padding_mask[i][None, :] != self.feature_extractor.padding_value
+ ]
+ audio_values[i] = sliced_audio.reshape(channels, -1)
+
+ return audio_values
+
+
+__all__ = ["MusicgenProcessor"]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen2_5_omni/__init__.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen2_5_omni/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d7ddae0da7e1b2c16b95e94e6678f9511cb716f
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen2_5_omni/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_qwen2_5_omni import *
+ from .modeling_qwen2_5_omni import *
+ from .processing_qwen2_5_omni import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b163af19a54ee7b048a905b17e4461300d8a05e
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py
@@ -0,0 +1,4028 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_qwen2_5_omni.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import warnings
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import Parameter
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub
+from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+ TransformersKwargs,
+ auto_docstring,
+ can_return_tuple,
+ check_torch_load_is_safe,
+ logging,
+ torch_compilable_check,
+)
+from ...utils.deprecation import deprecate_kwarg
+from ...utils.generic import (
+ accepts_precomputed_kwargs,
+ is_flash_attention_requested,
+ maybe_autocast,
+ merge_with_config_defaults,
+)
+from ...utils.hub import cached_file
+from ...utils.output_capturing import capture_outputs
+from ...vision_utils import get_vision_cu_seqlens, get_vision_position_ids, get_vision_window_index
+from .configuration_qwen2_5_omni import (
+ Qwen2_5OmniAudioEncoderConfig,
+ Qwen2_5OmniBigVGANConfig,
+ Qwen2_5OmniConfig,
+ Qwen2_5OmniDiTConfig,
+ Qwen2_5OmniTalkerConfig,
+ Qwen2_5OmniTextConfig,
+ Qwen2_5OmniThinkerConfig,
+ Qwen2_5OmniToken2WavConfig,
+ Qwen2_5OmniVisionEncoderConfig,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
+ """Generates a 1D Kaiser-windowed sinc filter.
+
+ Args:
+ cutoff (float): Normalized cutoff frequency (0 to 0.5).
+ half_width (float): Transition bandwidth.
+ kernel_size (int): Number of filter taps.
+
+ Returns:
+ torch.Tensor: A tensor of shape (1, 1, kernel_size) representing the filter.
+ """
+ is_even = kernel_size % 2 == 0
+ half_size = kernel_size // 2
+
+ # Compute Kaiser window parameters
+ delta_f = 4 * half_width
+ attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+
+ if attenuation > 50.0:
+ beta = 0.1102 * (attenuation - 8.7)
+ elif attenuation >= 21.0:
+ beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0)
+ else:
+ beta = 0.0
+
+ kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
+
+ # Compute time indices
+ if is_even:
+ time_indices = torch.arange(-half_size, half_size) + 0.5
+ else:
+ time_indices = torch.arange(kernel_size) - half_size
+
+ # Compute sinc filter
+ if cutoff == 0:
+ return torch.zeros((1, 1, kernel_size), dtype=torch.float32) # Ensures correct shape
+
+ sinc_filter = torch.sinc(2 * cutoff * time_indices)
+ normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
+
+ # Normalize to ensure sum = 1 (avoid leakage of constant component)
+ normalized_filter /= normalized_filter.sum()
+
+ return normalized_filter.view(1, 1, kernel_size)
+
+
+@auto_docstring
+class Qwen2_5OmniPreTrainedModel(PreTrainedModel):
+ config: Qwen2_5OmniConfig
+ base_model_prefix = "model"
+ input_modalities = ("image", "video", "audio", "text")
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Qwen2_5OmniDecoderLayer", "Qwen2_5OmniVisionBlock"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _can_compile_fullgraph = False
+ _supports_attention_backend = True
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, SinusoidsPositionEmbedding):
+ log_timescale_increment = np.log(module.max_timescale) / (module.channels // 2 - 1)
+ inv_timescales = torch.exp(-log_timescale_increment * torch.arange(module.channels // 2).float())
+ scaled_time = torch.arange(module.length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+ init.copy_(module.positional_embedding, torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1))
+ elif isinstance(module, UpSample1d):
+ filter_tensor = kaiser_sinc_filter1d(0.5 / module.ratio, 0.6 / module.ratio, module.kernel_size)
+ init.copy_(module.filter, filter_tensor)
+ elif isinstance(module, DownSample1d):
+ filter_tensor = kaiser_sinc_filter1d(module.cutoff, module.half_width, module.kernel_size)
+ init.copy_(module.filter, filter_tensor)
+ elif isinstance(module, Qwen2_5_VisionRotaryEmbedding):
+ inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim))
+ init.copy_(module.inv_freq, inv_freq)
+
+
+class Qwen2_5OmniPreTrainedModelForConditionalGeneration(Qwen2_5OmniPreTrainedModel):
+ input_modalities = ("image", "video", "audio", "text")
+
+ def get_llm_pos_ids_for_vision(
+ self,
+ start_idx: int,
+ vision_idx: int,
+ spatial_merge_size: int,
+ t_index: list[int],
+ grid_hs: list[int],
+ grid_ws: list[int],
+ ):
+ llm_pos_ids_list = []
+ llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
+ llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten()
+ t_index = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().long()
+ _llm_pos_ids = torch.stack([t_index, h_index, w_index])
+ llm_pos_ids_list.append(_llm_pos_ids + start_idx) # + 1 ) # 12.09 by malinhan
+ llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
+ return llm_pos_ids
+
+ def get_chunked_index(
+ self, token_indices: torch.Tensor, tokens_per_chunk: int, remove_index: int
+ ) -> list[tuple[int, int]]:
+ """
+ Splits token index list into chunks based on token value ranges.
+
+ Given a list of token indices, returns a list of (start, end) index tuples representing
+ slices of the list where the token values fall within successive ranges of `t_ntoken_per_chunk`.
+
+ For example, if `t_ntoken_per_chunk` is 1000, the function will create chunks such that:
+ - the first chunk contains token values < 1000,
+ - the second chunk contains values >= 1000 and < 2000, and so on.
+
+ Parameters:
+ token_indices (`torch.Tensor` of shape `(seq_len, )`): A monotonically increasing list of
+ token index values.
+ t_ntoken_per_chunk (`int`): Number of tokens per chunk (used as the chunk size threshold).
+ remove_index (`int`) An index id to subtract from `token_indices` before chunking
+
+ Returns:
+ `list[tuple[int, int]]`: A list of tuples, each representing the start (inclusive)
+ and end (exclusive) indices of a chunk in `token_indices`.
+ """
+
+ def _iter():
+ i, start_idx = 0, 0 # skip bos token
+ current_chunk = 1
+ while i < len(token_indices): # skip eos token
+ if token_indices[i] - remove_index >= current_chunk * tokens_per_chunk:
+ yield (start_idx, i)
+ start_idx = i
+ current_chunk += 1
+ i += 1
+ yield (start_idx, len(token_indices))
+
+ return list(_iter())
+
+ def get_rope_index(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ use_audio_in_video: bool = False,
+ audio_seqlens: torch.LongTensor | None = None,
+ second_per_grids: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
+
+ Explanation:
+ Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
+
+ For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
+ Examples:
+ input_ids: [T T T T T], here T is for text.
+ temporal position_ids: [0, 1, 2, 3, 4]
+ height position_ids: [0, 1, 2, 3, 4]
+ width position_ids: [0, 1, 2, 3, 4]
+
+ For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
+ and 1D rotary position embedding for text part.
+ Examples:
+ Temporal (Time): 3 patches, representing different segments of the video in time.
+ Height: 2 patches, dividing each frame vertically.
+ Width: 2 patches, dividing each frame horizontally.
+ We also have some important parameters:
+ fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.
+ tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity.
+ temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.
+ interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.
+ input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
+ vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]
+ vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
+ vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
+ text temporal position_ids: [101, 102, 103, 104, 105]
+ text height position_ids: [101, 102, 103, 104, 105]
+ text width position_ids: [101, 102, 103, 104, 105]
+ Here we calculate the text start position_ids as the max vision position_ids plus 1.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+ it.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ use_audio_in_video (`bool`, *optional*):
+ If set to `True`, use the audio in video.
+ audio_seqlens (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ second_per_grids (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+ """
+ spatial_merge_size = self.spatial_merge_size
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ audio_token_id = self.config.audio_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+ audio_start_token_id = self.config.audio_start_token_id
+ position_id_per_seconds = self.config.position_id_per_seconds
+ seconds_per_chunk = self.config.seconds_per_chunk
+
+ mrope_position_deltas = []
+ if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+ total_input_ids = input_ids
+ if attention_mask is not None:
+ attention_mask = attention_mask == 1
+ position_ids = torch.ones(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ image_idx, video_idx, audio_idx = 0, 0, 0
+ for i, input_ids in enumerate(total_input_ids):
+ if attention_mask is not None:
+ input_ids = input_ids[attention_mask[i]]
+ image_nums, video_nums, audio_nums = 0, 0, 0
+ vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+ vision_tokens = input_ids[vision_start_indices + 1]
+ audio_nums = torch.sum(input_ids == audio_start_token_id)
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (
+ (vision_tokens == audio_start_token_id).sum()
+ if use_audio_in_video
+ else (vision_tokens == video_token_id).sum()
+ )
+ input_tokens = input_ids.tolist()
+ llm_pos_ids_list: list = []
+ st = 0
+ remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums
+ multimodal_nums = (
+ image_nums + audio_nums if use_audio_in_video else image_nums + video_nums + audio_nums
+ )
+ for _ in range(multimodal_nums):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ if image_token_id in input_tokens and remain_images > 0:
+ ed_image = input_tokens.index(image_token_id, st)
+ else:
+ ed_image = len(input_tokens) + 1
+ if video_token_id in input_tokens and remain_videos > 0:
+ ed_video = input_tokens.index(video_token_id, st)
+ else:
+ ed_video = len(input_tokens) + 1
+ if audio_token_id in input_tokens and remain_audios > 0:
+ ed_audio = input_tokens.index(audio_token_id, st)
+ else:
+ ed_audio = len(input_tokens) + 1
+ min_ed = min(ed_image, ed_video, ed_audio)
+ if min_ed == ed_audio:
+ text_len = min_ed - st - 1
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ audio_len = ((audio_seqlens[audio_idx] - 1) // 2 + 1 - 2) // 2 + 1
+ llm_pos_ids = torch.arange(audio_len).view(1, -1).expand(3, -1) + st_idx
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st += text_len + bos_len + audio_len + eos_len
+ audio_idx += 1
+ remain_audios -= 1
+
+ elif min_ed == ed_image:
+ text_len = min_ed - st - 1
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ grid_t = image_grid_thw[image_idx][0]
+ grid_hs = image_grid_thw[:, 1]
+ grid_ws = image_grid_thw[:, 2]
+ t_index = (torch.arange(grid_t) * 1 * position_id_per_seconds).long()
+ llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st += text_len + bos_len + image_len + eos_len
+ image_idx += 1
+ remain_images -= 1
+
+ elif min_ed == ed_video and not use_audio_in_video:
+ text_len = min_ed - st - 1
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ grid_t = video_grid_thw[video_idx][0]
+ grid_hs = video_grid_thw[:, 1]
+ grid_ws = video_grid_thw[:, 2]
+ t_index = (
+ torch.arange(grid_t) * second_per_grids[video_idx].cpu().float() * position_id_per_seconds
+ ).long()
+ llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st += text_len + bos_len + video_len + eos_len
+ video_idx += 1
+ remain_videos -= 1
+
+ elif min_ed == ed_video and use_audio_in_video:
+ text_len = min_ed - st - 2
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ audio_len = ((audio_seqlens[audio_idx] - 1) // 2 + 1 - 2) // 2 + 1
+ audio_llm_pos_ids = torch.arange(audio_len).view(1, -1).expand(3, -1) + st_idx
+ grid_t = video_grid_thw[video_idx][0]
+ grid_hs = video_grid_thw[:, 1]
+ grid_ws = video_grid_thw[:, 2]
+
+ t_index = (
+ torch.arange(grid_t) * second_per_grids[video_idx].cpu().float() * position_id_per_seconds
+ ).long()
+ video_llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+
+ t_ntoken_per_chunk = int(position_id_per_seconds * seconds_per_chunk)
+ video_chunk_indexes = self.get_chunked_index(video_llm_pos_ids[0], t_ntoken_per_chunk, st_idx)
+ audio_chunk_indexes = self.get_chunked_index(audio_llm_pos_ids[0], t_ntoken_per_chunk, st_idx)
+ sub_len = 0
+ for j in range(max(len(video_chunk_indexes), len(audio_chunk_indexes))):
+ video_chunk_index = video_chunk_indexes[j] if j < len(video_chunk_indexes) else None
+ audio_chunk_index = audio_chunk_indexes[j] if j < len(audio_chunk_indexes) else None
+ if video_chunk_index is not None:
+ sub_len += video_chunk_index[1] - video_chunk_index[0]
+
+ llm_pos_ids_list.append(
+ video_llm_pos_ids[:, video_chunk_index[0] : video_chunk_index[1]]
+ )
+ if audio_chunk_index is not None:
+ sub_len += audio_chunk_index[1] - audio_chunk_index[0]
+
+ llm_pos_ids_list.append(
+ audio_llm_pos_ids[:, audio_chunk_index[0] : audio_chunk_index[1]]
+ )
+ video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st += text_len + bos_len * 2 + audio_len + video_len + eos_len * 2
+
+ audio_idx += 1
+ video_idx += 1
+ remain_videos -= 1
+ remain_audios -= 1
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+
+ if attention_mask is not None:
+ position_ids[..., i, attention_mask[i]] = llm_positions.to(position_ids.device)
+ else:
+ position_ids[..., i, :] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(input_ids))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas).unsqueeze(1).to(device=input_ids.device)
+
+ return position_ids, mrope_position_deltas
+ else:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+ max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+ mrope_position_deltas = max_position_ids + 1 - torch.sum(attention_mask, dim=-1, keepdim=True)
+
+ return position_ids, mrope_position_deltas
+
+
+############################
+# Start Thinker #
+############################
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Qwen2.5OmniThinker causal language model (or autoregressive) outputs.
+ """
+)
+@dataclass
+class Qwen2_5OmniThinkerCausalLMOutputWithPast(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`, *optional*):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ rope_deltas: torch.LongTensor | None = None
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+ """
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class Qwen2_5OmniAudioAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ config: Qwen2_5OmniAudioEncoderConfig,
+ ):
+ super().__init__()
+ self.embed_dim = config.d_model
+ self.num_heads = config.encoder_attention_heads
+ self.dropout = config.attention_dropout
+ self.head_dim = self.embed_dim // self.num_heads
+ self.num_key_value_groups = 1 # needed for eager attention
+ self.config = config
+
+ if (self.head_dim * self.num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {self.num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = 0.0
+ self.is_decoder = False
+ self.is_causal = False
+
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
+ self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+ self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ **kwargs,
+ ) -> torch.Tensor:
+ """Input shape: Batch x Time x Channel"""
+
+ seq_length, _ = hidden_states.size()
+
+ query_states = self.q_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
+ key_states = self.k_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
+ value_states = self.v_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
+
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ if is_flash_attention_requested(self.config):
+ # Flash Attention: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output
+
+
+class Qwen2_5OmniAudioEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Qwen2_5OmniAudioEncoderConfig):
+ super().__init__()
+ self.embed_dim = config.d_model
+ self.self_attn = Qwen2_5OmniAudioAttention(config)
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+ self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+ self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ **kwargs,
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+ hidden_states = self.self_attn(
+ hidden_states=hidden_states,
+ cu_seqlens=cu_seqlens,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = residual + hidden_states
+
+ if hidden_states.dtype == torch.float16:
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+ hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+ outputs = (hidden_states,)
+
+ return outputs
+
+
+class SinusoidsPositionEmbedding(nn.Module):
+ def __init__(self, length, channels, max_timescale=10000):
+ super().__init__()
+ self.length = length
+ self.channels = channels
+ self.max_timescale = max_timescale
+ if channels % 2 != 0:
+ raise ValueError("SinusoidsPositionEmbedding needs even channels input")
+ log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
+ inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2).float())
+ scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+ self.register_buffer(
+ "positional_embedding",
+ torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1),
+ persistent=False,
+ )
+
+ def forward(self, seqlen: int):
+ return self.positional_embedding[:seqlen, :]
+
+
+def chunk_and_pad_features(
+ input_features: torch.Tensor, feature_lens: torch.Tensor, n_window: int, kwargs: dict | None = None
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Split audio features into fixed-size chunks and pad to uniform length, or pop precomputed pair from `kwargs`.
+
+ Each audio sample is split into chunks of `n_window * 2` frames (the last
+ chunk may be shorter), then all chunks are right-padded to the longest chunk.
+
+ Args:
+ input_features: `(feature_dim, total_frames)` concatenated audio features.
+ feature_lens: `(batch_size,)` per-sample frame counts.
+ n_window: half the target chunk size in frames.
+ kwargs: optional caller kwargs — if it contains both `"padded_feature"` and `"chunk_lengths"` they are popped and returned.
+
+ Returns:
+ `padded_feature`: `(num_chunks, feature_dim, max_chunk_len)` padded chunks.
+ `chunk_lengths`: `(num_chunks,)` actual length of each chunk before padding.
+ """
+ if kwargs is not None:
+ padded_feature = kwargs.pop("padded_feature", None)
+ chunk_lengths = kwargs.pop("chunk_lengths", None)
+ if padded_feature is not None and chunk_lengths is not None:
+ return padded_feature, chunk_lengths
+
+ chunk_num = torch.ceil(feature_lens / (n_window * 2)).long()
+ chunk_lengths = torch.full((chunk_num.sum(),), n_window * 2, dtype=torch.long, device=feature_lens.device)
+ tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
+ chunk_lengths[tail_chunk_index] = feature_lens % (n_window * 2)
+ chunk_lengths = torch.where(chunk_lengths == 0, n_window * 2, chunk_lengths)
+
+ chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
+ padded_feature = nn.utils.rnn.pad_sequence(chunk_list, batch_first=True).transpose(1, 2)
+ return padded_feature, chunk_lengths
+
+
+def get_audio_cu_seqlens(chunk_lengths: torch.Tensor, kwargs: dict | None = None) -> torch.Tensor:
+ """Compute cumulative sequence lengths for audio attention, or pop `"cu_seqlens"` from `kwargs` if precomputed.
+
+ Applies one stride-2 convolution length reduction, then returns cumulative
+ boundaries for flash-attention-style sequence packing.
+
+ Args:
+ chunk_lengths: `(num_chunks,)` pre-CNN chunk lengths.
+ kwargs: optional caller kwargs — if it contains `"cu_seqlens"` it is popped and returned.
+
+ Returns:
+ `(num_chunks + 1,)` int32 cumulative sequence boundaries.
+ """
+ if kwargs is not None and (cu_seqlens := kwargs.pop("cu_seqlens", None)) is not None:
+ return cu_seqlens
+ after_conv1 = (chunk_lengths - 1) // 2 + 1
+ return F.pad(after_conv1.cumsum(0), (1, 0), value=0).to(torch.int32)
+
+
+def get_valid_indices(chunk_lengths: torch.Tensor, kwargs: dict | None = None) -> torch.Tensor:
+ """Compute flat indices of valid (non-padding) positions after one stride-2 conv, or pop `"valid_indices"` from `kwargs` if precomputed.
+
+ Args:
+ chunk_lengths: `(num_chunks,)` pre-CNN chunk lengths.
+ kwargs: optional caller kwargs — if it contains `"valid_indices"` it is popped and returned.
+
+ Returns:
+ `(total_valid,)` flat indices into the `(num_chunks * max_len_after_conv)` grid.
+ """
+ if kwargs is not None and (valid_indices := kwargs.pop("valid_indices", None)) is not None:
+ return valid_indices
+ after_conv1 = (chunk_lengths - 1) // 2 + 1
+ max_len = after_conv1.max().item()
+ mask = torch.arange(max_len, device=chunk_lengths.device) < after_conv1.unsqueeze(1)
+ return mask.flatten().nonzero().squeeze(-1)
+
+
+def get_pool_indices(feature_lens: torch.Tensor, kwargs: dict | None = None) -> torch.Tensor:
+ """Compute indices for post-encoder stride-2 average pooling, or pop `"pool_indices"` from `kwargs` if precomputed.
+
+ Args:
+ feature_lens: `(batch_size,)` mel spectrogram lengths.
+ kwargs: optional caller kwargs — if it contains `"pool_indices"` it is popped and returned.
+
+ Returns:
+ `(total_pooled,)` flat index of first element of each stride-2 pair.
+ """
+ if kwargs is not None and (pool_indices := kwargs.pop("pool_indices", None)) is not None:
+ return pool_indices
+ after_conv1 = (feature_lens - 1) // 2 + 1
+ num_pooled = (after_conv1 - 2) // 2 + 1
+ offsets = F.pad(after_conv1[:-1].cumsum(0), (1, 0), value=0)
+ pair_offsets = torch.repeat_interleave(offsets, num_pooled)
+ local_indices = torch.arange(num_pooled.sum(), device=feature_lens.device)
+ local_indices -= torch.repeat_interleave(F.pad(num_pooled[:-1].cumsum(0), (1, 0), value=0), num_pooled)
+ return pair_offsets + local_indices * 2
+
+
+@auto_docstring(
+ custom_intro="""
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+ [`Qwen2_5OmniAudioEncoderLayer`].
+ """
+)
+class Qwen2_5OmniAudioEncoder(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniAudioEncoderConfig
+ main_input_name = "input_features"
+ input_modalities = "audio"
+ _no_split_modules = ["Qwen2_5OmniAudioEncoderLayer"]
+ _supports_sdpa = True
+ _can_record_outputs = {
+ "hidden_states": Qwen2_5OmniAudioEncoderLayer,
+ "attentions": Qwen2_5OmniAudioAttention,
+ }
+
+ def __init__(self, config: Qwen2_5OmniAudioEncoderConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+
+ embed_dim = config.d_model
+ self.num_mel_bins = config.num_mel_bins
+ self.max_source_positions = config.max_source_positions
+ self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+ self.n_window = config.n_window
+ self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1)
+ self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)
+ self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, embed_dim)
+ self.audio_bos_eos_token = nn.Embedding(2, config.output_dim)
+ self.layers = nn.ModuleList([Qwen2_5OmniAudioEncoderLayer(config) for _ in range(config.encoder_layers)])
+ self.ln_post = nn.LayerNorm(config.d_model)
+ self.avg_pooler = nn.AvgPool1d(2, stride=2)
+ self.proj = nn.Linear(config.d_model, config.output_dim)
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def _freeze_parameters(self):
+ for param in self.parameters():
+ param.requires_grad = False
+ self._requires_grad = False
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.conv1
+
+ def set_input_embeddings(self, value: nn.Module):
+ self.conv1 = value
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ @auto_docstring
+ def forward(
+ self, input_features=None, feature_lens=None, aftercnn_lens=None, **kwargs: Unpack[TransformersKwargs]
+ ):
+ r"""
+ feature_lens (`torch.LongTensor` of shape `(batch_size,)`):
+ mel length
+ aftercnn_lens (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ mel length after cnn
+ """
+ padded_feature, chunk_lengths = chunk_and_pad_features(
+ input_features, feature_lens, self.n_window, kwargs=kwargs
+ )
+ valid_indices = get_valid_indices(chunk_lengths, kwargs=kwargs)
+ pool_indices = get_pool_indices(feature_lens, kwargs=kwargs)
+ cu_seqlens = get_audio_cu_seqlens(chunk_lengths, kwargs=kwargs)
+
+ # Derive masks from chunk_lengths (traceable arithmetic + arange broadcasting)
+ padded_feature = padded_feature.to(self.conv1.weight.dtype)
+ padded_mask = (
+ (torch.arange(padded_feature.shape[2], device=padded_feature.device) < chunk_lengths.unsqueeze(1))
+ .unsqueeze(1)
+ .long()
+ )
+ padded_embed = nn.functional.gelu(self.conv1(padded_feature)) * padded_mask
+ padded_embed = nn.functional.gelu(self.conv2(padded_embed)).transpose(1, 2)
+ padded_embed = padded_embed + self.positional_embedding.positional_embedding[
+ : padded_embed.shape[1], :
+ ].unsqueeze(0).to(padded_embed.dtype)
+ hidden_states = torch.index_select(padded_embed.reshape(-1, padded_embed.shape[-1]), 0, valid_indices)
+
+ for encoder_layer in self.layers:
+ layer_outputs = encoder_layer(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ **kwargs,
+ )
+ hidden_states = layer_outputs[0]
+
+ # Post-process: stride-2 average pooling using precomputed indices, then project
+ hidden_states = (hidden_states[pool_indices] + hidden_states[pool_indices + 1]) / 2
+ hidden_states = self.proj(self.ln_post(hidden_states))
+ return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
+
+ # Ignore copy
+ def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
+ """
+ Computes the output length of the convolutional layers and the output length of the audio encoder
+ """
+ input_lengths = (input_lengths - 1) // 2 + 1
+ output_lengths = (input_lengths - 2) // 2 + 1
+ return input_lengths, output_lengths
+
+ def padded_and_mask_function(self, tensor_list, tensor_len, padding_value=0, padding_side="right"):
+ """
+ Pads a sequence of tensors to their maximum length on indicated `padding_side`.
+ Then prepares a mask so that pad tokens are not attended to.
+ """
+ warnings.warn(
+ f"`{self.__class__.__name__}.padded_and_mask_function` is deprecated and will be removed in v5.11. Use `chunk_and_pad_features` and `get_audio_cu_seqlens` helpers instead.",
+ FutureWarning,
+ stacklevel=2,
+ )
+ max_len = tensor_len.max()
+ dim = tensor_list[0].shape[0]
+ padded_tensor = torch.full(
+ size=(len(tensor_list), dim, max_len),
+ fill_value=padding_value,
+ dtype=self.dtype,
+ device=tensor_list[0].device,
+ )
+
+ batch_mask = torch.zeros(
+ (len(tensor_len), max_len),
+ dtype=torch.long,
+ device=padded_tensor.device,
+ )
+ for i, length in enumerate(tensor_len):
+ batch_mask[i, :length] = 1
+ padded_tensor[i, :, :length] = tensor_list[i]
+
+ feature_lens_after_cnn = (tensor_len - 1) // 2 + 1
+ max_len_after_cnn = feature_lens_after_cnn.max()
+ batch_mask_after_cnn = torch.zeros(
+ (len(tensor_len), max_len_after_cnn),
+ dtype=torch.long,
+ device=padded_tensor.device,
+ )
+ for i, length in enumerate(feature_lens_after_cnn):
+ batch_mask_after_cnn[i, :length] = 1
+ return (
+ padded_tensor,
+ batch_mask.unsqueeze(1),
+ batch_mask_after_cnn.bool(),
+ )
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb_vision(tensor: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+ orig_dtype = tensor.dtype
+ tensor = tensor.float()
+ cos = freqs.cos()
+ sin = freqs.sin()
+ cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
+ sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
+ output = (tensor * cos) + (rotate_half(tensor) * sin)
+ output = output.to(orig_dtype)
+ return output
+
+
+class Qwen2_5OmniVisionAttention(nn.Module):
+ def __init__(self, config: Qwen2_5OmniVisionEncoderConfig = None) -> None:
+ super().__init__()
+ self.dim = config.hidden_size
+ self.num_heads = config.num_heads
+ self.head_dim = self.dim // self.num_heads
+ self.q = nn.Linear(self.dim, self.dim, bias=True)
+ self.k = nn.Linear(self.dim, self.dim, bias=True)
+ self.v = nn.Linear(self.dim, self.dim, bias=True)
+ self.proj = nn.Linear(self.dim, self.dim)
+ self.scaling = self.head_dim**-0.5
+ self.num_key_value_groups = 1 # needed for eager attention
+ self.config = config
+ self.attention_dropout = 0.0
+ self.is_causal = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: torch.Tensor | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ query_states = self.q(hidden_states).reshape(seq_length, self.num_heads, -1)
+ key_states = self.k(hidden_states).reshape(seq_length, self.num_heads, -1)
+ value_states = self.v(hidden_states).reshape(seq_length, self.num_heads, -1)
+ query_states = apply_rotary_pos_emb_vision(query_states.unsqueeze(0), rotary_pos_emb).squeeze(0)
+ key_states = apply_rotary_pos_emb_vision(key_states.unsqueeze(0), rotary_pos_emb).squeeze(0)
+
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ if is_flash_attention_requested(self.config):
+ # Flash Attention 2: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Qwen2_5OmniRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ Qwen2_5OmniRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(torch.float32)
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+ return self.weight * hidden_states.to(input_dtype)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Qwen2_5OmniMLP(nn.Module):
+ def __init__(self, config, bias: bool = False):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
+
+
+class Qwen2_5OmniVisionBlock(GradientCheckpointingLayer):
+ def __init__(self, config: Qwen2_5OmniVisionEncoderConfig) -> None:
+ super().__init__()
+ self.norm1 = Qwen2_5OmniRMSNorm(config.hidden_size, eps=1e-6)
+ self.norm2 = Qwen2_5OmniRMSNorm(config.hidden_size, eps=1e-6)
+ self.attn = Qwen2_5OmniVisionAttention(config=config)
+ self.mlp = Qwen2_5OmniMLP(config, bias=True)
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: torch.Tensor | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ r"""
+ cu_seqlens (`torch.Tensor`):
+ Cumulative sequence lengths used for packed variable-length attention in Flash Attention kernels.
+ rotary_pos_emb (`torch.Tensor`, *optional*):
+ Precomputed rotary positional embeddings applied to the vision attention query/key states.
+ """
+ hidden_states = hidden_states + self.attn(
+ self.norm1(hidden_states),
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ **kwargs,
+ )
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+ return hidden_states
+
+
+class Qwen2_5_VisionRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
+ super().__init__()
+ self.dim = dim
+ self.theta = theta
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
+ return (position_ids.unsqueeze(-1) * self.inv_freq).flatten(1)
+
+
+class Qwen2_5_VisionPatchEmbed(nn.Module):
+ def __init__(
+ self,
+ patch_size: int | list[int] | tuple[int, int] = 14,
+ temporal_patch_size: int | list[int] | tuple[int, int] = 2,
+ in_channels: int = 3,
+ embed_dim: int = 1152,
+ ) -> None:
+ super().__init__()
+ self.patch_size = patch_size
+ self.temporal_patch_size = temporal_patch_size
+ self.in_channels = in_channels
+ self.embed_dim = embed_dim
+
+ kernel_size = [temporal_patch_size, patch_size, patch_size]
+ self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ target_dtype = self.proj.weight.dtype
+ hidden_states = hidden_states.view(
+ -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
+ )
+ hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+ return hidden_states
+
+
+class Qwen2_5OmniPatchMerger(nn.Module):
+ def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None:
+ super().__init__()
+ self.hidden_size = context_dim * (spatial_merge_size**2)
+ self.ln_q = Qwen2_5OmniRMSNorm(context_dim, eps=1e-6)
+ self.mlp = nn.Sequential(
+ nn.Linear(self.hidden_size, self.hidden_size),
+ nn.GELU(),
+ nn.Linear(self.hidden_size, dim),
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.mlp(self.ln_q(x).view(-1, self.hidden_size))
+ return x
+
+
+class Qwen2_5OmniVisionEncoder(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniVisionEncoderConfig
+ _no_split_modules = ["Qwen2_5OmniVisionBlock"]
+ _input_embed_layer = "patch_embed"
+ _can_record_outputs = {
+ "hidden_states": Qwen2_5OmniVisionBlock,
+ "attentions": Qwen2_5OmniVisionAttention,
+ }
+ input_modalities = ("image", "video")
+
+ def __init__(self, config: Qwen2_5OmniVisionEncoderConfig, *inputs, **kwargs) -> None:
+ super().__init__(config, *inputs, **kwargs)
+ self.spatial_merge_size = config.spatial_merge_size
+ self.patch_size = config.patch_size
+ self.fullatt_block_indexes = config.fullatt_block_indexes
+ self.window_size = config.window_size
+ self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
+
+ self.patch_embed = Qwen2_5_VisionPatchEmbed(
+ patch_size=config.patch_size,
+ temporal_patch_size=config.temporal_patch_size,
+ in_channels=config.in_channels,
+ embed_dim=config.hidden_size,
+ )
+
+ head_dim = config.hidden_size // config.num_heads
+ self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
+ self.blocks = nn.ModuleList([Qwen2_5OmniVisionBlock(config) for _ in range(config.depth)])
+ self.merger = Qwen2_5OmniPatchMerger(
+ dim=config.out_hidden_size,
+ context_dim=config.hidden_size,
+ spatial_merge_size=config.spatial_merge_size,
+ )
+ self.gradient_checkpointing = False
+
+ self.post_init()
+
+ def rot_pos_emb(self, grid_thw):
+ warnings.warn(
+ f"`{self.__class__.__name__}.rot_pos_emb` is deprecated and will be removed in v5.11. Use `get_vision_position_ids` from `transformers.vision_utils` and apply the rotary embedding module.",
+ FutureWarning,
+ stacklevel=2,
+ )
+ position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size)
+ rotary_pos_emb = self.rotary_pos_emb(position_ids)
+ return rotary_pos_emb
+
+ def get_window_index(self, grid_thw):
+ warnings.warn(
+ f"`{self.__class__.__name__}.get_window_index` is deprecated and will be removed in v5.11. Use `get_vision_window_index` from `transformers.vision_utils` instead.",
+ FutureWarning,
+ stacklevel=2,
+ )
+ window_index, cu_window_seqlens = get_vision_window_index(
+ grid_thw,
+ spatial_merge_size=self.spatial_merge_size,
+ window_size=self.window_size,
+ patch_size=self.patch_size,
+ )
+ return window_index, cu_window_seqlens.tolist()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> tuple | BaseModelOutputWithPooling:
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+ The final hidden states of the model.
+ grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+ The temporal, height and width of feature shape of each image in LLM.
+
+ Returns:
+ `torch.Tensor`: hidden_states.
+ """
+ position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs)
+ cu_seqlens = get_vision_cu_seqlens(grid_thw, kwargs=kwargs)
+ window_index, cu_window_seqlens = get_vision_window_index(
+ grid_thw,
+ spatial_merge_size=self.spatial_merge_size,
+ window_size=self.window_size,
+ patch_size=self.patch_size,
+ kwargs=kwargs,
+ )
+
+ hidden_states = self.patch_embed(hidden_states)
+
+ seq_len, _ = hidden_states.size()
+ reverse_indices = torch.argsort(window_index)
+ hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+ hidden_states = hidden_states[window_index, :, :]
+ hidden_states = hidden_states.reshape(seq_len, -1)
+
+ rotary_pos_emb = self.rotary_pos_emb(position_ids)
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+ rotary_pos_emb = rotary_pos_emb[window_index, :, :]
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+
+ # Modification here
+ for layer_num, blk in enumerate(self.blocks):
+ if layer_num in self.fullatt_block_indexes:
+ cu_seqlens_now = cu_seqlens
+ else:
+ cu_seqlens_now = cu_window_seqlens
+
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens_now,
+ rotary_pos_emb=rotary_pos_emb,
+ **kwargs,
+ )
+
+ merged_hidden_states = self.merger(hidden_states)
+ merged_hidden_states = merged_hidden_states[reverse_indices, :]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=hidden_states,
+ pooler_output=merged_hidden_states,
+ )
+
+
+class Qwen2_5OmniRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Qwen2_5OmniThinkerConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: Qwen2_5OmniConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ # Ignore copy
+ def forward(self, x, position_ids):
+ # In contrast to other models, Qwen2_5Omni has different position ids for the grids
+ # So we expand the inv_freq to shape (3, ...)
+ inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
+ position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions)
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/).
+
+ Explanation:
+ Multimodal 3D rotary position embedding is an extension to 1D rotary position embedding. The input embedding
+ sequence contains vision (images / videos) embedding and text embedding or just contains text embedding. For
+ vision embedding part, we apply rotary position embedding on temporal, height and width dimension separately.
+ Here we split the channel dimension to 3 chunks for the temporal, height and width rotary position embedding.
+ For text embedding part, we just apply 1D rotary position embedding. The three rotary position index (temporal,
+ height and width) of text embedding is always the same, so the text embedding rotary position embedding has no
+ difference with modern LLMs.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ position_ids (`torch.Tensor`):
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
+ used to pass offsetted position ids when working with a KV-cache.
+ mrope_section(`List(int)`):
+ Multimodal rope section is for channel dimension of temporal, height and width in rope calculation.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ mrope_section = mrope_section * 2
+ cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
+ unsqueeze_dim
+ )
+ sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
+ unsqueeze_dim
+ )
+
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class Qwen2_5OmniAttention(nn.Module):
+ """
+ Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
+ and "Generating Long Sequences with Sparse Transformers".
+ """
+
+ def __init__(self, config: Qwen2_5OmniConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ if layer_idx is None:
+ logger.warning_once(
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ self.hidden_size = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+ self.num_key_value_heads = config.num_key_value_heads
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+ self.is_causal = True
+ self.attention_dropout = config.attention_dropout
+ self.rope_parameters = config.rope_parameters
+ self.scaling = self.head_dim**-0.5
+
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+ self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None
+ self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ output_attentions: bool = False,
+ use_cache: bool = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+ bsz, q_len, _ = hidden_states.size()
+
+ query_states = self.q_proj(hidden_states)
+ key_states = self.k_proj(hidden_states)
+ value_states = self.v_proj(hidden_states)
+
+ query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_multimodal_rotary_pos_emb(
+ query_states, key_states, cos, sin, self.config.rope_parameters["mrope_section"]
+ )
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ sliding_window=self.sliding_window,
+ position_ids=position_ids, # pass positions for FA2
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class Qwen2MLP(nn.Module):
+ def __init__(self, config, bias: bool = False):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
+
+
+class Qwen2_5OmniDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Qwen2_5OmniTextConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ if config.use_sliding_window and not is_flash_attention_requested(config):
+ logger.warning_once(
+ f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "
+ "unexpected results may be encountered."
+ )
+ self.self_attn = Qwen2_5OmniAttention(config, layer_idx)
+
+ self.mlp = Qwen2MLP(config)
+ self.input_layernorm = Qwen2_5OmniRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Qwen2_5OmniRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+ `(batch, sequence_length)` where padding elements are indicated by 0.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+ (see `past_key_values`).
+ past_key_values (`Cache`, *optional*): cached past key and value projection states
+ position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+ Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+ with `head_dim` being the embedding dimension of each attention head.
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring
+class Qwen2_5OmniThinkerTextModel(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniTextConfig
+ input_modalities = ("text",)
+ _can_record_outputs = {
+ "hidden_states": Qwen2_5OmniDecoderLayer,
+ "attentions": Qwen2_5OmniAttention,
+ }
+ _no_split_modules = ["Qwen2_5OmniDecoderLayer"]
+
+ def __init__(self, config: Qwen2_5OmniTextConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [Qwen2_5OmniDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self._attn_implementation = config._attn_implementation
+ self.norm = Qwen2_5OmniRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.has_sliding_layers = "sliding_attention" in self.config.layer_types
+ self.rotary_emb = Qwen2_5OmniRotaryEmbedding(config=config)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ # torch.jit.trace() doesn't support cache objects in the output
+ if use_cache and past_key_values is None and not torch.jit.is_tracing():
+ past_key_values = DynamicCache(config=self.config)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ # the hard coded `3` is for temporal, height and width.
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+ elif position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+ # NOTE: we need to pass text position ids for packing. Qwen2-VL uses 3D positions
+ # where each dim indicates visual spatial positions for temporal/height/width grids.
+ # There are two scenarios when FA2-like packed masking might be activated.
+ # 1. User specifically passed packed `position_ids` and no attention mask.
+ # In this case we expect the useer to create correct position ids for all 3 grids
+ # and prepend text-only position ids to it. The final tensor will be [4, bs, seq-len]
+ # 2. User runs forward with no attention mask and no position ids. In this case, position ids
+ # are prepared by the model (`get_rope_index`) as `[4, bs, seq-len]` tensor. Text-only positions are
+ # prepended by us when creating positions so that the mask is constructed correctly. NOTE: failing to pass
+ # text-only positions will cause incorrect mask construction, do not change `prepare_input_for_generation`
+ if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ text_position_ids = position_ids[0]
+ position_ids = position_ids[1:]
+ else:
+ # If inputs are not packed (usual 3D positions), do not prepare mask from position_ids
+ text_position_ids = None
+
+ # It may already have been prepared by e.g. `generate`
+ if not isinstance(causal_mask_mapping := attention_mask, dict):
+ # Prepare mask arguments
+ mask_kwargs = {
+ "config": self.config,
+ "inputs_embeds": inputs_embeds,
+ "attention_mask": attention_mask,
+ "past_key_values": past_key_values,
+ "position_ids": text_position_ids,
+ }
+ # Create the masks
+ causal_mask_mapping = {
+ "full_attention": create_causal_mask(**mask_kwargs),
+ }
+ # The sliding window alternating layers are not always activated depending on the config
+ if self.has_sliding_layers:
+ causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ for i, decoder_layer in enumerate(self.layers):
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask_mapping[self.config.layer_types[i]],
+ position_embeddings=position_embeddings,
+ position_ids=text_position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The Qwen2.5OmniThinker model which consists of a audio backbone and a language model.
+ """
+)
+class Qwen2_5OmniThinkerForConditionalGeneration(Qwen2_5OmniPreTrainedModelForConditionalGeneration, GenerationMixin):
+ config: Qwen2_5OmniThinkerConfig
+ base_model_prefix = "thinker"
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _no_split_modules = ["Qwen2_5OmniAudioEncoder", "Qwen2_5OmniVisionEncoder"]
+
+ def __init__(self, config: Qwen2_5OmniThinkerConfig):
+ super().__init__(config)
+ self.audio_tower = Qwen2_5OmniAudioEncoder._from_config(config.audio_config)
+ self.visual = Qwen2_5OmniVisionEncoder._from_config(config.vision_config)
+ self.vocab_size = config.text_config.vocab_size
+ self.model = Qwen2_5OmniThinkerTextModel._from_config(config.text_config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+ self.spatial_merge_size = config.vision_config.spatial_merge_size
+ self.rope_deltas = None
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ @accepts_precomputed_kwargs(modality="video")
+ @can_return_tuple
+ @auto_docstring
+ def get_video_features(
+ self,
+ pixel_values_videos: torch.FloatTensor,
+ video_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input videos.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ """
+ pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
+ return self.visual(pixel_values_videos, grid_thw=video_grid_thw, **kwargs)
+
+ @accepts_precomputed_kwargs(modality="image")
+ @can_return_tuple
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ pixel_values = pixel_values.type(self.visual.dtype)
+ return self.visual(pixel_values, grid_thw=image_grid_thw, **kwargs)
+
+ @can_return_tuple
+ @auto_docstring
+ def get_audio_features(
+ self,
+ input_features: torch.FloatTensor,
+ feature_attention_mask: torch.LongTensor | None = None,
+ audio_feature_lengths: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ input_features (`torch.FloatTensor`):
+ The tensors corresponding to the input audios.
+ feature_attention_mask (`torch.LongTensor`, *optional*):
+ Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`:
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ """
+ if feature_attention_mask is not None:
+ audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
+ input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
+ else:
+ audio_feature_lengths = None
+
+ audio_feat_lengths, audio_output_lengths = self.audio_tower._get_feat_extract_output_lengths(
+ audio_feature_lengths if audio_feature_lengths is not None else feature_attention_mask.sum(-1)
+ )
+ feature_lens = audio_feature_lengths if audio_feature_lengths is not None else feature_attention_mask.sum(-1)
+ audio_outputs = self.audio_tower(
+ input_features, feature_lens=feature_lens, aftercnn_lens=audio_feat_lengths, **kwargs
+ )
+ if audio_outputs.last_hidden_state.shape[0] != sum(audio_output_lengths.tolist()):
+ raise ValueError("length of audio_features should match audio_output_lengths")
+
+ return audio_outputs
+
+ def get_placeholder_mask(
+ self,
+ input_ids: torch.LongTensor,
+ inputs_embeds: torch.FloatTensor,
+ image_features: torch.FloatTensor | None = None,
+ video_features: torch.FloatTensor | None = None,
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ special_video_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_video_mask = special_video_mask.all(-1)
+ special_audio_mask = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(self.config.audio_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ ).all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+ special_video_mask = input_ids == self.config.video_token_id
+ special_audio_mask = input_ids == self.config.audio_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if image_features is not None:
+ torch_compilable_check(
+ inputs_embeds[special_image_mask].numel() == image_features.numel(),
+ f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}",
+ )
+
+ n_video_tokens = special_video_mask.sum()
+ special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if video_features is not None:
+ torch_compilable_check(
+ inputs_embeds[special_video_mask].numel() == video_features.numel(),
+ f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}",
+ )
+
+ special_audio_mask = special_audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ return special_image_mask, special_video_mask, special_audio_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ input_features: torch.FloatTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ pixel_values_videos: torch.FloatTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ feature_attention_mask: torch.Tensor | None = None,
+ audio_feature_lengths: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ rope_deltas: torch.LongTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ use_audio_in_video: bool | None = None,
+ video_second_per_grid: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | Qwen2_5OmniThinkerCausalLMOutputWithPast:
+ r"""
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ feature_attention_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+ use_audio_in_video (`bool`, *optional*):
+ Whether or not use audio track in video, should same as the parameter in `process_audio_info`.
+ video_second_per_grid (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ Number of seconds per grid for each video, used for temporal feature mapping.
+
+ Example:
+
+ ```python
+ >>> from io import BytesIO
+ >>> from urllib.request import urlopen
+ >>> import librosa
+ >>> from qwen_vl_utils import process_vision_info
+ >>> from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration
+
+ >>> thinker = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B")
+ >>> processor = Qwen2_5OmniProcessor.from_pretrained("Qwen/Qwen2.5-Omni-7B")
+
+ >>> conversations = [
+ >>> {'role': 'system', 'content': 'You are a helpful voice chat bot, and please respond to me in a casual conversation manner using random voice.'},
+ >>> {"role": "user", "content": [
+ >>> {"type": "image", "image_url": "https://www.ilankelman.org/stopsigns/australia.jpg"},
+ >>> {"type": "audio", "audio_url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/glass-breaking-151256.mp3"},
+ >>> ]},
+ >>> ]
+
+ >>> text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
+ >>> audios = [ librosa.load(BytesIO(urlopen( conversations[1]['content'][1]['audio_url'] ).read()), sr=self.processor.feature_extractor.sampling_rate) ]
+ >>> images, videos = process_vision_info(conversations)
+ >>> inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors="pt", padding=True)
+
+ >>> # Generate
+ >>> inputs['use_audio_in_video'] = `True` or `False`
+ >>> generation = thinker.generate(**inputs, max_new_tokens=2048)
+ >>> generate_ids = generation[:, inputs.input_ids.size(1):]
+
+ >>> response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ ```"""
+
+ if inputs_embeds is None:
+ # 1. Extract the input embeddings
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ # 2. Merge text , audios , image and video
+ if input_features is not None:
+ audio_features = self.get_audio_features(
+ input_features, feature_attention_mask, audio_feature_lengths, return_dict=True, **kwargs
+ ).last_hidden_state
+ audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
+ _, _, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_features)
+
+ if pixel_values is not None:
+ image_embeds = self.get_image_features(
+ pixel_values, image_grid_thw, return_dict=True, **kwargs
+ ).pooler_output
+ image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ image_mask, _, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ if pixel_values_videos is not None:
+ video_embeds = self.get_video_features(
+ pixel_values_videos, video_grid_thw, return_dict=True, **kwargs
+ ).pooler_output
+ video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ _, video_mask, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ if feature_attention_mask is not None:
+ audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
+ else:
+ audio_feature_lengths = None
+
+ if attention_mask is not None and position_ids is None:
+ past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+ if past_key_values_length == 0 or self.rope_deltas is None:
+ delta0 = (1 - attention_mask).sum(dim=-1).unsqueeze(1)
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask,
+ use_audio_in_video,
+ audio_feature_lengths,
+ video_second_per_grid,
+ )
+ rope_deltas = rope_deltas - delta0
+ self.rope_deltas = rope_deltas
+ else:
+ batch_size, seq_length = input_ids.shape
+ delta = (past_key_values_length + self.rope_deltas).to(input_ids.device)
+ position_ids = torch.arange(seq_length, device=input_ids.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs = self.model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ return_dict=True,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ logits = self.lm_head(hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits, labels=labels, vocab_size=self.config.get_text_config().vocab_size
+ )
+
+ return Qwen2_5OmniThinkerCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=self.rope_deltas,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ input_features=None,
+ feature_attention_mask=None,
+ use_audio_in_video=False,
+ video_second_per_grid=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ use_cache=use_cache,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ input_features=input_features,
+ feature_attention_mask=feature_attention_mask,
+ use_audio_in_video=use_audio_in_video,
+ video_second_per_grid=video_second_per_grid,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ model_inputs["position_ids"] = None
+
+ if not is_first_iteration and use_cache:
+ model_inputs["pixel_values"] = None
+ model_inputs["pixel_values_videos"] = None
+ model_inputs["input_features"] = None
+
+ return model_inputs
+
+
+############################
+# Start Talker #
+############################
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Qwen2.5OmniTalker causal language model (or autoregressive) outputs.
+ """
+)
+@dataclass
+class Qwen2_5OmniTalkerCausalLMOutputWithPast(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ thinker_reply_part (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ Hidden states from the thinker model that are used as input for the talker model. These represent the encoded
+ response that the talker model will use to generate speech tokens.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ rope_deltas: torch.LongTensor | None = None
+ thinker_reply_part: torch.FloatTensor | None = None
+
+
+@auto_docstring
+class Qwen2_5OmniTalkerModel(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniTalkerConfig
+ input_modalities = ("image", "video", "audio", "text")
+ _can_record_outputs = {
+ "hidden_states": Qwen2_5OmniDecoderLayer,
+ "attentions": Qwen2_5OmniAttention,
+ }
+
+ _no_split_modules = ["Qwen2_5OmniTalkerDecoderLayer"]
+
+ def __init__(self, config: Qwen2_5OmniTalkerConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.embedding_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [Qwen2_5OmniDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self._attn_implementation = config._attn_implementation
+ self.norm = Qwen2_5OmniRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.has_sliding_layers = "sliding_attention" in self.config.layer_types
+ self.rotary_emb = Qwen2_5OmniRotaryEmbedding(config=config)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ # torch.jit.trace() doesn't support cache objects in the output
+ if use_cache and past_key_values is None and not torch.jit.is_tracing():
+ past_key_values = DynamicCache(config=self.config)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ # the hard coded `3` is for temporal, height and width.
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+ elif position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+ # NOTE: we need to pass text position ids for packing. Qwen2-VL uses 3D positions
+ # where each dim indicates visual spatial positions for temporal/height/width grids.
+ # There are two scenarios when FA2-like packed masking might be activated.
+ # 1. User specifically passed packed `position_ids` and no attention mask.
+ # In this case we expect the useer to create correct position ids for all 3 grids
+ # and prepend text-only position ids to it. The final tensor will be [4, bs, seq-len]
+ # 2. User runs forward with no attention mask and no position ids. In this case, position ids
+ # are prepared by the model (`get_rope_index`) as `[4, bs, seq-len]` tensor. Text-only positions are
+ # prepended by us when creating positions so that the mask is constructed correctly. NOTE: failing to pass
+ # text-only positions will cause incorrect mask construction, do not change `prepare_input_for_generation`
+ if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ text_position_ids = position_ids[0]
+ position_ids = position_ids[1:]
+ else:
+ # If inputs are not packed (usual 3D positions), do not prepare mask from position_ids
+ text_position_ids = None
+
+ # It may already have been prepared by e.g. `generate`
+ if not isinstance(causal_mask_mapping := attention_mask, dict):
+ # Prepare mask arguments
+ mask_kwargs = {
+ "config": self.config,
+ "inputs_embeds": inputs_embeds,
+ "attention_mask": attention_mask,
+ "past_key_values": past_key_values,
+ "position_ids": text_position_ids,
+ }
+ # Create the masks
+ causal_mask_mapping = {
+ "full_attention": create_causal_mask(**mask_kwargs),
+ }
+ # The sliding window alternating layers are not always activated depending on the config
+ if self.has_sliding_layers:
+ causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ for i, decoder_layer in enumerate(self.layers):
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask_mapping[self.config.layer_types[i]],
+ position_embeddings=position_embeddings,
+ position_ids=text_position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+class Qwen2_5OmniTalkerForConditionalGeneration(Qwen2_5OmniPreTrainedModelForConditionalGeneration, GenerationMixin):
+ config: Qwen2_5OmniTalkerConfig
+ base_model_prefix = "talker"
+ output_modalities = ("audio",)
+
+ def __init__(self, config: Qwen2_5OmniTalkerConfig):
+ super().__init__(config)
+
+ self.thinker_to_talker_proj = nn.Linear(config.embedding_size, config.hidden_size)
+
+ self.model = Qwen2_5OmniTalkerModel(config)
+ self.codebook_size = config.vocab_size
+ self.codec_head = nn.Linear(config.hidden_size, self.codebook_size, bias=False)
+
+ self.codec_bos_token = config.tts_codec_start_token_id
+ self.codec_eos_token = config.tts_codec_end_token_id
+ self.codec_pad_token = config.tts_codec_pad_token_id
+ self.codec_mask_token = config.tts_codec_mask_token_id
+
+ self.text_bos_token = config.tts_text_start_token_id
+ self.text_eos_token = config.tts_text_end_token_id
+ self.text_pad_token = config.tts_text_pad_token_id
+
+ self.spatial_merge_size = self.config.spatial_merge_size
+ self.rope_deltas = None
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ thinker_reply_part: torch.FloatTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ rope_deltas: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ input_text_ids: torch.LongTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ use_audio_in_video: bool | None = None,
+ audio_feature_lengths: torch.LongTensor | None = None,
+ video_second_per_grid: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> tuple | Qwen2_5OmniTalkerCausalLMOutputWithPast:
+ r"""
+ thinker_reply_part (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Hidden states from the thinker model's output that represent the text reply part to be processed.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ input_text_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Input token IDs for text-only content, used for position calculation in multimodal contexts.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ use_audio_in_video (`bool`, *optional*):
+ Whether or not use audio track in video, should same as the parameter in `process_audio_info`.
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ video_second_per_grid (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ Number of seconds per grid for each video, used for temporal feature mapping.
+
+ Example:
+
+ ```python
+ >>> from io import BytesIO
+ >>> from urllib.request import urlopen
+ >>> import librosa
+ >>> from transformers import AutoProcessor, Qwen2_5OmniTalkerForConditionalGeneration
+
+ >>> model = Qwen2_5OmniTalkerForConditionalGeneration.from_pretrained("Qwen/Qwen2-Audio-7B")
+ >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2-Audio-7B")
+
+ >>> prompt = "<|audio_bos|><|AUDIO|><|audio_eos|>Generate the caption in English:"
+ >>> url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/glass-breaking-151256.mp3"
+ >>> audio, _ = librosa.load(BytesIO(urlopen(url).read()), sr=self.processor.feature_extractor.sampling_rate)
+
+ >>> inputs = processor(text=prompt, audio=audio, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(**inputs, max_length=30)
+ >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Generate the caption in English: Glass is breaking."
+ ```"""
+
+ if attention_mask is not None and position_ids is None:
+ past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+ if past_key_values_length == 0 or self.rope_deltas is None:
+ position_ids, rope_deltas = self.get_rope_index(
+ input_text_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask,
+ use_audio_in_video,
+ audio_feature_lengths,
+ video_second_per_grid,
+ )
+
+ if inputs_embeds is not None:
+ inputs_embeds[:, -1, :] += self.get_input_embeddings()(
+ torch.tensor([self.codec_bos_token], dtype=torch.long, device=inputs_embeds.device)
+ )
+ inputs_embeds[:, -2, :] += self.get_input_embeddings()(
+ torch.tensor([self.codec_pad_token], dtype=torch.long, device=inputs_embeds.device)
+ )
+ self.rope_deltas = rope_deltas
+
+ else:
+ if inputs_embeds is not None:
+ batch_size, seq_length, _ = inputs_embeds.shape
+ else:
+ batch_size, seq_length = input_ids.shape
+
+ delta = (past_key_values_length + self.rope_deltas).to(input_ids.device)
+ position_ids = torch.arange(seq_length, device=input_ids.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ if inputs_embeds is None:
+ # 1. Inference tokens after second token
+ codec_embeds = self.get_input_embeddings()(input_ids)
+ inputs_embeds = codec_embeds + thinker_reply_part[:, :1, :]
+ if thinker_reply_part.shape[1] > 1:
+ thinker_reply_part = thinker_reply_part[:, 1:, :]
+
+ talker_lm_input = self.thinker_to_talker_proj(inputs_embeds)
+
+ if attention_mask is not None:
+ attention_mask = attention_mask.to(inputs_embeds.device)
+
+ outputs = self.model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=talker_lm_input,
+ use_cache=use_cache,
+ return_dict=True,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ logits = self.codec_head(hidden_states)
+ logits = logits.float()
+
+ loss = None
+
+ return Qwen2_5OmniTalkerCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=self.rope_deltas,
+ thinker_reply_part=thinker_reply_part,
+ )
+
+ # prepare inputs for talker lm generation
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ input_text_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ thinker_reply_part=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ input_audio_features=None,
+ audio_feature_attention_mask=None,
+ audio_feature_lengths=None,
+ use_audio_in_video=False,
+ video_second_per_grid=None,
+ **kwargs,
+ ):
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ thinker_reply_part=thinker_reply_part,
+ input_text_ids=input_text_ids,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ use_audio_in_video=use_audio_in_video,
+ audio_feature_lengths=audio_feature_lengths,
+ video_second_per_grid=video_second_per_grid,
+ **kwargs,
+ )
+
+ model_inputs["position_ids"] = None
+
+ return model_inputs
+
+ def _update_model_kwargs_for_generation(
+ self,
+ outputs: ModelOutput,
+ model_kwargs: dict[str, Any],
+ is_encoder_decoder: bool = False,
+ num_new_tokens: int = 1,
+ ) -> dict[str, Any]:
+ model_kwargs = super()._update_model_kwargs_for_generation(
+ outputs, model_kwargs, is_encoder_decoder, num_new_tokens
+ )
+
+ if getattr(outputs, "thinker_reply_part", None) is not None:
+ model_kwargs["thinker_reply_part"] = outputs.thinker_reply_part
+
+ return model_kwargs
+
+
+class Qwen2_5OmniDiTRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Qwen2_5OmniDiTConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: Qwen2_5OmniDiTConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class TimeDelayNetBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ kernel_size,
+ dilation,
+ ):
+ super().__init__()
+ self.conv = nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ dilation=dilation,
+ padding="same",
+ padding_mode="reflect",
+ )
+ self.activation = nn.ReLU()
+
+ def forward(self, hidden_states: torch.Tensor):
+ return self.activation(self.conv(hidden_states))
+
+
+class Res2NetBlock(torch.nn.Module):
+ def __init__(self, in_channels, out_channels, scale=8, kernel_size=3, dilation=1):
+ super().__init__()
+
+ in_channel = in_channels // scale
+ hidden_channel = out_channels // scale
+
+ self.blocks = nn.ModuleList(
+ [
+ TimeDelayNetBlock(
+ in_channel,
+ hidden_channel,
+ kernel_size=kernel_size,
+ dilation=dilation,
+ )
+ for i in range(scale - 1)
+ ]
+ )
+ self.scale = scale
+
+ def forward(self, hidden_states):
+ outputs = []
+ for i, hidden_part in enumerate(torch.chunk(hidden_states, self.scale, dim=1)):
+ if i == 0:
+ output_part = hidden_part
+ elif i == 1:
+ output_part = self.blocks[i - 1](hidden_part)
+ else:
+ output_part = self.blocks[i - 1](hidden_part + output_part)
+ outputs.append(output_part)
+ output = torch.cat(outputs, dim=1)
+ return output
+
+
+class SqueezeExcitationBlock(nn.Module):
+ def __init__(self, in_channels, se_channels, out_channels):
+ super().__init__()
+
+ self.conv1 = nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=se_channels,
+ kernel_size=1,
+ padding="same",
+ padding_mode="reflect",
+ )
+ self.relu = nn.ReLU(inplace=True)
+ self.conv2 = nn.Conv1d(
+ in_channels=se_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ padding="same",
+ padding_mode="reflect",
+ )
+ self.sigmoid = nn.Sigmoid()
+
+ def forward(self, hidden_states):
+ hidden_states_mean = hidden_states.mean(dim=2, keepdim=True)
+
+ hidden_states_mean = self.relu(self.conv1(hidden_states_mean))
+ hidden_states_mean = self.sigmoid(self.conv2(hidden_states_mean))
+
+ return hidden_states * hidden_states_mean
+
+
+class AttentiveStatisticsPooling(nn.Module):
+ """This class implements an attentive statistic pooling layer for each channel.
+ It returns the concatenated mean and std of the input tensor.
+ """
+
+ def __init__(self, channels, attention_channels=128):
+ super().__init__()
+
+ self.eps = 1e-12
+ self.tdnn = TimeDelayNetBlock(channels * 3, attention_channels, 1, 1)
+ self.tanh = nn.Tanh()
+ self.conv = nn.Conv1d(
+ in_channels=attention_channels,
+ out_channels=channels,
+ kernel_size=1,
+ padding="same",
+ padding_mode="reflect",
+ )
+
+ def _length_to_mask(self, length, max_len=None, dtype=None, device=None):
+ """Creates a binary mask for each sequence.
+
+ Reference: https://discuss.pytorch.org/t/how-to-generate-variable-length-mask/23397/3
+
+ Arguments
+ ---------
+ length : torch.LongTensor
+ Containing the length of each sequence in the batch. Must be 1D.
+ max_len : int
+ Max length for the mask, also the size of the second dimension.
+ dtype : torch.dtype, default: None
+ The dtype of the generated mask.
+ device: torch.device, default: None
+ The device to put the mask variable.
+
+ Returns
+ -------
+ mask : tensor
+ The binary mask.
+ """
+
+ if max_len is None:
+ max_len = length.max().long().item() # using arange to generate mask
+ mask = torch.arange(max_len, device=length.device, dtype=length.dtype).expand(
+ len(length), max_len
+ ) < length.unsqueeze(1)
+
+ mask = torch.as_tensor(mask, dtype=dtype, device=device)
+ return mask
+
+ def _compute_statistics(self, x, m, dim=2):
+ mean = (m * x).sum(dim)
+ std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps))
+ return mean, std
+
+ def forward(self, hidden_states):
+ seq_length = hidden_states.shape[-1]
+ lengths = torch.ones(hidden_states.shape[0], device=hidden_states.device)
+
+ # Make binary mask of shape [N, 1, L]
+ mask = self._length_to_mask(
+ lengths * seq_length, max_len=seq_length, dtype=hidden_states.dtype, device=hidden_states.device
+ )
+ mask = mask.unsqueeze(1)
+
+ # Expand the temporal context of the pooling layer by allowing the
+ # self-attention to look at global properties of the utterance.
+ total = mask.sum(dim=2, keepdim=True)
+
+ mean, std = self._compute_statistics(hidden_states, mask / total)
+ mean = mean.unsqueeze(2).repeat(1, 1, seq_length)
+ std = std.unsqueeze(2).repeat(1, 1, seq_length)
+ attention = torch.cat([hidden_states, mean, std], dim=1)
+
+ # Apply layers
+ attention = self.conv(self.tanh(self.tdnn(attention)))
+
+ # Filter out zero-paddings
+ attention = attention.masked_fill(mask == 0, float("-inf"))
+
+ attention = F.softmax(attention, dim=2)
+ mean, std = self._compute_statistics(hidden_states, attention)
+ # Append mean and std of the batch
+ pooled_stats = torch.cat((mean, std), dim=1)
+ pooled_stats = pooled_stats.unsqueeze(2)
+
+ return pooled_stats
+
+
+class SqueezeExcitationRes2NetBlock(nn.Module):
+ """An implementation of building block in ECAPA-TDNN, i.e.,
+ TDNN-Res2Net-TDNN-SqueezeExcitationBlock.
+ """
+
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ res2net_scale=8,
+ se_channels=128,
+ kernel_size=1,
+ dilation=1,
+ ):
+ super().__init__()
+ self.out_channels = out_channels
+ self.tdnn1 = TimeDelayNetBlock(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ dilation=1,
+ )
+ self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation)
+ self.tdnn2 = TimeDelayNetBlock(
+ out_channels,
+ out_channels,
+ kernel_size=1,
+ dilation=1,
+ )
+ self.se_block = SqueezeExcitationBlock(out_channels, se_channels, out_channels)
+
+ def forward(self, hidden_state):
+ residual = hidden_state
+
+ hidden_state = self.tdnn1(hidden_state)
+ hidden_state = self.res2net_block(hidden_state)
+ hidden_state = self.tdnn2(hidden_state)
+ hidden_state = self.se_block(hidden_state)
+
+ return hidden_state + residual
+
+
+class ECAPA_TimeDelayNet(torch.nn.Module):
+ """An implementation of the speaker embedding model in a paper.
+ "ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in
+ TDNN Based Speaker Verification" (https://huggingface.co/papers/2005.07143).
+ """
+
+ def __init__(self, config: Qwen2_5OmniDiTConfig):
+ super().__init__()
+ if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(config.enc_channels) != len(
+ config.enc_dilations
+ ):
+ raise ValueError("enc_channels, enc_kernel_sizes and enc_dilations should have same length")
+ self.channels = config.enc_channels
+ self.blocks = nn.ModuleList()
+
+ # The initial TDNN layer
+ self.blocks.append(
+ TimeDelayNetBlock(
+ config.mel_dim,
+ config.enc_channels[0],
+ config.enc_kernel_sizes[0],
+ config.enc_dilations[0],
+ )
+ )
+
+ # SE-Res2Net layers
+ for i in range(1, len(config.enc_channels) - 1):
+ self.blocks.append(
+ SqueezeExcitationRes2NetBlock(
+ config.enc_channels[i - 1],
+ config.enc_channels[i],
+ res2net_scale=config.enc_res2net_scale,
+ se_channels=config.enc_se_channels,
+ kernel_size=config.enc_kernel_sizes[i],
+ dilation=config.enc_dilations[i],
+ )
+ )
+
+ # Multi-layer feature aggregation
+ self.mfa = TimeDelayNetBlock(
+ config.enc_channels[-1],
+ config.enc_channels[-1],
+ config.enc_kernel_sizes[-1],
+ config.enc_dilations[-1],
+ )
+
+ # Attentive Statistical Pooling
+ self.asp = AttentiveStatisticsPooling(
+ config.enc_channels[-1],
+ attention_channels=config.enc_attention_channels,
+ )
+
+ # Final linear transformation
+ self.fc = nn.Conv1d(
+ in_channels=config.enc_channels[-1] * 2,
+ out_channels=config.enc_dim,
+ kernel_size=1,
+ padding="same",
+ padding_mode="reflect",
+ )
+
+ def forward(self, hidden_states):
+ # Minimize transpose for efficiency
+ hidden_states = hidden_states.transpose(1, 2)
+
+ hidden_states_list = []
+ for layer in self.blocks:
+ hidden_states = layer(hidden_states)
+ hidden_states_list.append(hidden_states)
+
+ # Multi-layer feature aggregation
+ hidden_states = torch.cat(hidden_states_list[1:], dim=1)
+ hidden_states = self.mfa(hidden_states)
+
+ # Attentive Statistical Pooling
+ hidden_states = self.asp(hidden_states)
+
+ # Final linear transformation
+ hidden_states = self.fc(hidden_states)
+
+ hidden_states = hidden_states.squeeze(-1)
+ return hidden_states
+
+
+class DiTInputEmbedding(nn.Module):
+ def __init__(self, config: Qwen2_5OmniDiTConfig):
+ super().__init__()
+ self.proj = nn.Linear(
+ config.mel_dim + config.enc_dim + config.enc_emb_dim + config.emb_dim,
+ config.hidden_size,
+ )
+ self.spk_encoder = ECAPA_TimeDelayNet(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ speaker_embedding: torch.Tensor,
+ condition_vector: torch.Tensor,
+ code_embed: torch.Tensor,
+ drop_audio_cond: bool | None = False,
+ code_embed_uncond: bool | None = None,
+ apply_cfg: bool | None = True,
+ ):
+ if apply_cfg:
+ hidden_states = torch.cat([hidden_states, hidden_states], dim=0)
+ speaker_embedding = torch.cat([speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0)
+ condition_vector = torch.cat([condition_vector, torch.zeros_like(condition_vector)], dim=0)
+ code_embed = torch.cat([code_embed, code_embed_uncond], dim=0)
+ elif drop_audio_cond: # cfg for cond audio
+ condition_vector = torch.zeros_like(condition_vector)
+ speaker_embedding = torch.zeros_like(speaker_embedding)
+ condition_vector = self.spk_encoder(condition_vector).unsqueeze(1).repeat(1, hidden_states.size(1), 1)
+ hidden_states = self.proj(torch.cat((hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1))
+
+ return hidden_states
+
+
+# Transformer backbone using DiT blocks
+class DiTCodecEmbedding(nn.Module):
+ def __init__(self, codec_num_embeds, codec_dim, repeats):
+ super().__init__()
+ self.repeats = repeats
+ self.codec_embed = nn.Embedding(codec_num_embeds + 1, codec_dim)
+
+ def forward(self, code, drop_code=False):
+ if drop_code:
+ code = torch.zeros_like(code)
+ code_embed = self.codec_embed(code)
+
+ code_embed = torch.repeat_interleave(code_embed, repeats=self.repeats, dim=1)
+ return code_embed
+
+
+# AdaLayerNormZero
+# return with modulated x for attn input, and params for later mlp modulation
+class Qwen2_5_OmniAdaLayerNormZero(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+
+ self.silu = nn.SiLU()
+ self.linear = nn.Linear(dim, dim * 6)
+
+ self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+ def forward(self, hidden_states, emb=None):
+ emb = self.linear(self.silu(emb))
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+
+ hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+ return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+# AdaLayerNormZero for final layer
+# return only with modulated x for attn input, cuz no more mlp modulation
+class Qwen2_5_OmniAdaLayerNormZero_Final(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+
+ self.silu = nn.SiLU()
+ self.linear = nn.Linear(dim, dim * 2)
+
+ self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+ def forward(self, hidden_states, emb):
+ emb = self.linear(self.silu(emb))
+ scale, shift = torch.chunk(emb, 2, dim=1)
+
+ hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
+ return hidden_states
+
+
+# FeedForward
+class DiTMLP(nn.Module):
+ def __init__(self, dim, mult=4, dropout=0.0):
+ super().__init__()
+ inner_dim = int(dim * mult)
+
+ self.ff = nn.ModuleList(
+ [
+ nn.Linear(dim, inner_dim),
+ nn.GELU(approximate="tanh"),
+ nn.Dropout(dropout),
+ nn.Linear(inner_dim, dim),
+ ]
+ )
+
+ def forward(self, hidden_states):
+ for layer in self.ff:
+ hidden_states = layer(hidden_states)
+ return hidden_states
+
+
+# Modified from Llama with a different rotate function, will fixed in next release
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+
+ def rotate_half_codec(x):
+ # x = rearrange(x, "... (d r) -> ... d r", r=2)
+ x = x.reshape(*x.shape[:-1], -1, 2)
+ x1, x2 = x.unbind(dim=-1)
+ x = torch.stack((-x2, x1), dim=-1)
+ return x.reshape(*x.shape[:-2], -1)
+
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half_codec(q) * sin)
+ k_embed = (k * cos) + (rotate_half_codec(k) * sin)
+ return q_embed, k_embed
+
+
+class DiTAttention(nn.Module):
+ def __init__(self, config: Qwen2_5OmniDiTConfig):
+ super().__init__()
+
+ self.config = config
+ self.dim = config.hidden_size
+ self.heads = config.num_attention_heads
+ self.inner_dim = config.head_dim * config.num_attention_heads
+ self.dropout = config.dropout
+ self.is_causal = False
+
+ self.to_q = nn.Linear(config.hidden_size, self.inner_dim)
+ self.to_k = nn.Linear(config.hidden_size, self.inner_dim)
+ self.to_v = nn.Linear(config.hidden_size, self.inner_dim)
+
+ self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)])
+
+ def forward(
+ self,
+ hidden_states, # noised input x
+ position_embeddings=None, # rotary position embedding for x
+ attention_mask=None,
+ ) -> torch.Tensor:
+ batch_size = hidden_states.shape[0]
+
+ # `sample` projections.
+ query = self.to_q(hidden_states)
+ key = self.to_k(hidden_states)
+ value = self.to_v(hidden_states)
+
+ # attention
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // self.heads
+ query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+
+ # apply rotary position embedding
+ # Due to training process, only first head is applied with RoPE, will be fixed at next release
+ cos, sin = position_embeddings
+ query[:, :1], key[:, :1] = apply_rotary_pos_emb(query[:, :1], key[:, :1], cos, sin)
+
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+ attention_weights, _ = attention_interface(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=attention_mask,
+ is_causal=False,
+ )
+
+ # mask. e.g. inference got a batch with different target durations, mask out the padding
+ attention_weights = attention_weights.reshape(batch_size, -1, self.heads * head_dim)
+ attention_weights = attention_weights.to(query.dtype)
+
+ # linear proj
+ attention_output = self.to_out[0](attention_weights)
+ attention_output = self.to_out[1](attention_output)
+
+ return attention_output
+
+
+# time step conditioning embedding
+class SinusPositionEmbedding(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ self.dim = dim
+
+ def forward(self, hidden_states, scale=1000):
+ device = hidden_states.device
+ half_dim = self.dim // 2
+ emb = math.log(10000) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+ emb = scale * hidden_states.unsqueeze(1) * emb.unsqueeze(0)
+ emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+ return emb.type_as(hidden_states)
+
+
+class DiTTimestepEmbedding(nn.Module):
+ def __init__(self, dim, freq_embed_dim=256):
+ super().__init__()
+ self.time_embed = SinusPositionEmbedding(freq_embed_dim)
+ self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)])
+
+ def forward(self, timestep):
+ time_hidden = self.time_embed(timestep)
+ time_hidden = time_hidden.to(timestep.dtype)
+ for layer in self.time_mlp:
+ time_hidden = layer(time_hidden) # b d
+ return time_hidden
+
+
+class DiTDecoderLayer(nn.Module):
+ def __init__(self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0):
+ super().__init__()
+ self.attn_norm = Qwen2_5_OmniAdaLayerNormZero(config.hidden_size)
+
+ self.attn = DiTAttention(config)
+ self.look_ahead_block = look_ahead_block
+ self.look_backward_block = look_backward_block
+ self.ff_norm = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=1e-6)
+ self.ff = DiTMLP(dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout)
+
+ def forward(
+ self, hidden_states, timestep, position_embeddings=None, block_diff=None
+ ): # x: noised input, t: time embedding
+ # pre-norm & modulation for attention input
+ norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(hidden_states, emb=timestep)
+
+ # attention
+ attn_output = self.attn(
+ hidden_states=norm,
+ position_embeddings=position_embeddings,
+ attention_mask=(block_diff >= -float(self.look_backward_block))
+ & (block_diff <= float(self.look_ahead_block)),
+ )
+
+ # process attention output for input x
+ hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_output
+
+ norm = self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+ ff_output = self.ff(norm)
+ hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
+
+ return hidden_states
+
+
+class SnakeBeta(nn.Module):
+ """
+ A modified Snake function which uses separate parameters for the magnitude of the periodic components
+ Shape:
+ - Input: (B, C, T)
+ - Output: (B, C, T), same shape as the input
+ Parameters:
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ References:
+ - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+ https://huggingface.co/papers/2006.08195
+ """
+
+ def __init__(self, in_features, alpha=1.0):
+ super().__init__()
+ self.in_features = in_features
+
+ # initialize alpha
+ self.alpha = Parameter(torch.zeros(in_features) * alpha)
+ self.beta = Parameter(torch.zeros(in_features) * alpha)
+
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, hidden_states):
+ """
+ Forward pass of the function.
+ Applies the function to the input elementwise.
+ SnakeBeta ∶= x + 1/b * sin^2 (xa)
+ """
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+ beta = self.beta.unsqueeze(0).unsqueeze(-1)
+ alpha = torch.exp(alpha)
+ beta = torch.exp(beta)
+ hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
+ torch.sin(hidden_states * alpha), 2
+ )
+
+ return hidden_states
+
+
+class UpSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ self.stride = ratio
+ self.pad = self.kernel_size // ratio - 1
+ self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+ self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+
+ filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
+ self.register_buffer("filter", filter, persistent=False)
+
+ def forward(self, hidden_states):
+ channels = hidden_states.shape[1]
+
+ hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate")
+ hidden_states = self.ratio * F.conv_transpose1d(
+ hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels
+ )
+ hidden_states = hidden_states[..., self.pad_left : -self.pad_right]
+
+ return hidden_states
+
+
+class DownSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ cutoff = 0.5 / ratio
+ half_width = 0.6 / ratio
+ self.cutoff = cutoff
+ self.half_width = half_width
+ self.kernel_size = kernel_size
+
+ if cutoff < 0.0:
+ raise ValueError("Minimum cutoff must be larger than zero.")
+ if cutoff > 0.5:
+ raise ValueError("A cutoff above 0.5 does not make sense.")
+
+ self.even = kernel_size % 2 == 0
+ self.pad_left = kernel_size // 2 - int(self.even)
+ self.pad_right = kernel_size // 2
+ self.stride = ratio
+ filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+ self.register_buffer("filter", filter, persistent=False)
+
+ def forward(self, hidden_states):
+ channels = hidden_states.shape[1]
+ hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate")
+ out = F.conv1d(hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels)
+ return out
+
+
+class TorchActivation1d(nn.Module):
+ def __init__(
+ self,
+ activation,
+ up_ratio: int = 2,
+ down_ratio: int = 2,
+ up_kernel_size: int = 12,
+ down_kernel_size: int = 12,
+ ):
+ super().__init__()
+ if not callable(activation):
+ raise TypeError("Activation function must be callable")
+ self.act = activation
+ self.upsample = UpSample1d(up_ratio, up_kernel_size)
+ self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+ def forward(self, hidden_states):
+ hidden_states = self.upsample(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.downsample(hidden_states)
+
+ return hidden_states
+
+
+class AMPBlock(torch.nn.Module):
+ def __init__(
+ self,
+ channels,
+ kernel_size=3,
+ dilation=(1, 3, 5),
+ ):
+ super().__init__()
+
+ self.convs1 = nn.ModuleList(
+ [
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[0],
+ padding=self._get_padding(kernel_size, dilation[0]),
+ ),
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[1],
+ padding=self._get_padding(kernel_size, dilation[1]),
+ ),
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[2],
+ padding=self._get_padding(kernel_size, dilation[2]),
+ ),
+ ]
+ )
+
+ self.convs2 = nn.ModuleList(
+ [
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding=self._get_padding(kernel_size, 1),
+ ),
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding=self._get_padding(kernel_size, 1),
+ ),
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding=self._get_padding(kernel_size, 1),
+ ),
+ ]
+ )
+
+ self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers
+
+ self.activations = nn.ModuleList(
+ [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
+ )
+
+ def _get_padding(self, kernel_size, dilation=1):
+ return int((kernel_size * dilation - dilation) / 2)
+
+ def forward(self, hidden_states):
+ acts1, acts2 = self.activations[::2], self.activations[1::2]
+ for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2):
+ residual = hidden_states
+ hidden_states = act1(hidden_states)
+ hidden_states = conv1(hidden_states)
+ hidden_states = act2(hidden_states)
+ hidden_states = conv2(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring(
+ custom_intro="""
+ The full Qwen2.5Omni Token2WavBigVGAN model. Which take mel spectrogram as input and predict waveform.
+ """
+)
+class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniBigVGANConfig
+ input_modalities = "audio"
+
+ def __init__(self, config: Qwen2_5OmniBigVGANConfig):
+ super().__init__(config)
+ self.num_residual_blocks = len(config.resblock_kernel_sizes)
+ self.num_upsample_layers = len(config.upsample_rates)
+
+ self.conv_pre = nn.Conv1d(config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3)
+
+ # Removing extra ModuleList breaks official state dict
+ ups = [
+ nn.ModuleList(
+ [
+ nn.ConvTranspose1d(
+ config.upsample_initial_channel // (2**layer_idx),
+ config.upsample_initial_channel // (2 ** (layer_idx + 1)),
+ kernel_size,
+ stride,
+ padding=(kernel_size - stride) // 2,
+ )
+ ]
+ )
+ for layer_idx, (stride, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes))
+ ]
+ self.ups = nn.ModuleList(ups)
+
+ self.resblocks = nn.ModuleList(
+ [
+ AMPBlock(config.upsample_initial_channel // (2 ** (layer_idx + 1)), kernel_size, dilation)
+ for layer_idx in range(self.num_upsample_layers)
+ for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes)
+ ]
+ )
+
+ self.activation_post = TorchActivation1d(
+ activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers))
+ )
+ self.conv_post = nn.Conv1d(
+ config.upsample_initial_channel // (2**self.num_upsample_layers), 1, 7, 1, padding=3, bias=False
+ )
+
+ self.post_init()
+
+ def normalize_spectrogram(self, spectrogram, max_value, min_db):
+ return torch.clamp((2 * max_value) * ((spectrogram - min_db) / (-min_db)) - max_value, -max_value, max_value)
+
+ def amplitude_to_db(self, amplitude, min_db_level):
+ min_level = torch.exp(
+ torch.tensor(min_db_level / 20.0 * np.log(10), device=amplitude.device, dtype=amplitude.dtype)
+ )
+ return 20 * torch.log10(torch.clamp(amplitude, min=min_level))
+
+ def process_mel_spectrogram(self, mel_spectrogram):
+ amplitude_spectrum = torch.exp(mel_spectrogram)
+ decibel_spectrum = self.amplitude_to_db(amplitude_spectrum, -115) - 20
+ return self.normalize_spectrogram(decibel_spectrum, 1, -115)
+
+ def forward(self, mel_spectrogram, **kwargs):
+ processed_spectrogram = self.process_mel_spectrogram(mel_spectrogram)
+ hidden_representation = self.conv_pre(processed_spectrogram)
+
+ for layer_index in range(self.num_upsample_layers):
+ hidden_representation = self.ups[layer_index][0](hidden_representation)
+ residual_output = sum(
+ self.resblocks[layer_index * self.num_residual_blocks + block_index](hidden_representation)
+ for block_index in range(self.num_residual_blocks)
+ )
+ residual_output = residual_output / self.num_residual_blocks
+ hidden_representation = residual_output
+
+ hidden_representation = self.activation_post(hidden_representation)
+ output_waveform = self.conv_post(hidden_representation)
+ return torch.clamp(output_waveform, min=-1.0, max=1.0).squeeze().cpu()
+
+
+class RungeKutta4ODESolver:
+ def __init__(self, function, initial_value):
+ self.function = function
+ self.initial_value = initial_value
+
+ self._one_third = 1 / 3
+ self._two_thirds = 2 / 3
+
+ def _rk4_step(self, function, time_start, time_step, time_end, value_start, function_value_start=None):
+ k1 = function_value_start if function_value_start is not None else function(time_start, value_start)
+ k2 = function(time_start + time_step * self._one_third, value_start + time_step * k1 * self._one_third)
+ k3 = function(time_start + time_step * self._two_thirds, value_start + time_step * (k2 - k1 * self._one_third))
+ k4 = function(time_end, value_start + time_step * (k1 - k2 + k3))
+ return (k1 + 3 * (k2 + k3) + k4) * time_step / 8
+
+ def _compute_step(self, function, time_start, time_step, time_end, value_start):
+ function_value_start = function(time_start, value_start)
+ return self._rk4_step(
+ function, time_start, time_step, time_end, value_start, function_value_start=function_value_start
+ ), function_value_start
+
+ def _linear_interpolation(self, time_start, time_end, value_start, value_end, time_point):
+ if time_point == time_start:
+ return value_start
+ if time_point == time_end:
+ return value_end
+ weight = (time_point - time_start) / (time_end - time_start)
+ return value_start + weight * (value_end - value_start)
+
+ def integrate(self, time_points):
+ solution = torch.empty(
+ len(time_points),
+ *self.initial_value.shape,
+ dtype=self.initial_value.dtype,
+ device=self.initial_value.device,
+ )
+ solution[0] = self.initial_value
+
+ current_index = 1
+ current_value = self.initial_value
+ for time_start, time_end in zip(time_points[:-1], time_points[1:]):
+ time_step = time_end - time_start
+ delta_value, _ = self._compute_step(self.function, time_start, time_step, time_end, current_value)
+ next_value = current_value + delta_value
+
+ while current_index < len(time_points) and time_end >= time_points[current_index]:
+ solution[current_index] = self._linear_interpolation(
+ time_start, time_end, current_value, next_value, time_points[current_index]
+ )
+ current_index += 1
+
+ current_value = next_value
+
+ return solution
+
+
+@auto_docstring(
+ custom_intro="""
+ The full Qwen2.5Omni Token2WavDiT model. Which take speech tokens as input and predict mel spectrogram.
+ """
+)
+class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniDiTConfig
+ input_modalities = "audio"
+ _no_split_modules = ["DiTDecoderLayer"]
+
+ def __init__(self, config: Qwen2_5OmniDiTConfig):
+ super().__init__(config)
+ self.mel_dim = config.mel_dim
+ self.repeats = config.repeats
+ self.time_embed = DiTTimestepEmbedding(config.hidden_size)
+
+ self.text_embed = DiTCodecEmbedding(config.num_embeds, config.emb_dim, config.repeats)
+ self.input_embed = DiTInputEmbedding(config)
+
+ self.rotary_embed = Qwen2_5OmniDiTRotaryEmbedding(config=config)
+
+ self.hidden_size = config.hidden_size
+ self.layers = config.num_hidden_layers
+ self.block_size = config.block_size
+ self.num_attention_heads = config.num_attention_heads
+
+ self.transformer_blocks = nn.ModuleList()
+ for i in range(config.num_hidden_layers):
+ self.transformer_blocks.append(
+ DiTDecoderLayer(
+ config,
+ look_ahead_block=1 if i in config.look_ahead_layers else 0,
+ look_backward_block=1 if i in config.look_backward_layers else 0,
+ )
+ )
+
+ self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(config.hidden_size) # final modulation
+ self.proj_out = nn.Linear(config.hidden_size, config.mel_dim)
+
+ self.post_init()
+
+ def _create_block_diff(self, hidden_states):
+ batch, seq_len = hidden_states.shape[0], hidden_states.shape[1]
+ block_indices = torch.arange(seq_len, device=hidden_states.device) // self.block_size # [seq_length]
+
+ block_i = block_indices.unsqueeze(1) # [seq_length, 1]
+ block_j = block_indices.unsqueeze(0) # [1, seq_length]
+ block_diff = block_j - block_i # (n, n)
+
+ return block_diff.expand(batch, self.num_attention_heads, seq_len, seq_len)
+
+ def forward(
+ self,
+ hidden_states,
+ condition_vector,
+ speaker_embedding,
+ quantized_code,
+ time_step,
+ drop_audio_conditioning=False,
+ drop_code=False,
+ apply_cfg=True,
+ **kwargs,
+ ):
+ batch_size = hidden_states.shape[0]
+ if time_step.ndim == 0:
+ time_step = time_step.repeat(batch_size)
+
+ # Compute embeddings
+ time_embedding = self.time_embed(time_step)
+ text_embedding = self.text_embed(quantized_code, drop_code=False if apply_cfg else drop_code)
+ text_embedding_unconditioned = self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
+
+ hidden_states = self.input_embed(
+ hidden_states,
+ speaker_embedding,
+ condition_vector,
+ text_embedding,
+ drop_audio_cond=drop_audio_conditioning,
+ code_embed_uncond=text_embedding_unconditioned,
+ apply_cfg=apply_cfg,
+ )
+
+ # Compute positional encodings
+ position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device)
+ position_ids = position_ids[None, :].repeat(batch_size, 1)
+ position_embeddings = self.rotary_embed(hidden_states, position_ids)
+ blockwise_difference = self._create_block_diff(hidden_states)
+
+ # Transformer blocks
+ for transformer_block in self.transformer_blocks:
+ hidden_states = transformer_block(
+ hidden_states,
+ time_embedding,
+ position_embeddings=position_embeddings,
+ block_diff=blockwise_difference,
+ )
+
+ hidden_states = self.norm_out(hidden_states, time_embedding)
+ output = self.proj_out(hidden_states)
+
+ return output
+
+ @torch.no_grad()
+ def sample(
+ self,
+ conditioning_vector,
+ reference_mel_spectrogram,
+ quantized_code,
+ num_steps=10,
+ guidance_scale=0.5,
+ sway_coefficient=-1.0,
+ ):
+ maximum_duration = quantized_code.shape[1] * self.repeats
+ batch_size = reference_mel_spectrogram.shape[0]
+ if batch_size != 1:
+ raise ValueError("Only batch size = 1 is currently supported")
+
+ if maximum_duration > self.config.max_position_embeddings:
+ raise ValueError(
+ f"Requested mel length ({maximum_duration}) exceeds `dit_config.max_position_embeddings` "
+ f"({self.config.max_position_embeddings}). Provide shorter `quantized_code`."
+ )
+
+ initial_state = torch.randn(
+ [batch_size, maximum_duration, self.mel_dim],
+ dtype=reference_mel_spectrogram.dtype,
+ device=quantized_code.device,
+ )
+ conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, maximum_duration, 1)
+
+ def ode_function(time_step, hidden_states):
+ if guidance_scale < 1e-5:
+ prediction = self(
+ hidden_states=hidden_states,
+ speaker_embedding=conditioning_vector,
+ condition_vector=reference_mel_spectrogram,
+ quantized_code=quantized_code,
+ time_step=time_step,
+ drop_audio_conditioning=False,
+ drop_code=False,
+ apply_cfg=False,
+ )
+ return prediction
+
+ model_output = self(
+ hidden_states=hidden_states,
+ quantized_code=quantized_code,
+ speaker_embedding=conditioning_vector,
+ condition_vector=reference_mel_spectrogram,
+ time_step=time_step,
+ apply_cfg=True,
+ )
+ guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
+ return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
+
+ initial_time = 0
+ time_embedding = torch.linspace(
+ initial_time, 1, num_steps, device=quantized_code.device, dtype=conditioning_vector.dtype
+ )
+
+ if sway_coefficient is not None:
+ time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
+
+ ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
+ solution_trajectory = ode_solver.integrate(time_embedding)
+
+ generated_waveform = solution_trajectory[-1]
+ generated_mel_spectrogram = generated_waveform.permute(0, 2, 1)
+ return generated_mel_spectrogram
+
+
+@auto_docstring(
+ custom_intro="""
+ The full Qwen2.5Omni Token2Wav model. Consists a DiT model take speech tokens as input and predict mel spectrogram and a BigVGAN vocoder take mel spectrogram as input and predict waveform.
+ """
+)
+class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniToken2WavConfig
+ base_model_prefix = "model"
+ input_modalities = "audio"
+ _no_split_modules = ["Qwen2_5OmniToken2WavDiTModel", "Qwen2_5OmniToken2WavBigVGANModel"]
+
+ def __init__(self, config: Qwen2_5OmniToken2WavConfig):
+ super().__init__(config)
+ attn_impl = config._attn_implementation
+ if is_flash_attention_requested(config):
+ logger.warning_once(
+ "Qwen2_5OmniToken2WavModel must inference with fp32, but Flash Attention only supports fp16 and bf16, "
+ "attention implementation of Qwen2_5OmniToken2WavModel will fallback to sdpa."
+ )
+ attn_impl = "sdpa"
+ elif config._attn_implementation == "eager":
+ logger.warning_once(
+ "Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa"
+ )
+ attn_impl = "sdpa"
+ self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel._from_config(
+ config.dit_config, attn_implementation=attn_impl
+ )
+ self.code2wav_bigvgan_model = Qwen2_5OmniToken2WavBigVGANModel._from_config(
+ config.bigvgan_config, attn_implementation=attn_impl
+ )
+
+ self.post_init()
+
+ def forward(
+ self,
+ code,
+ conditioning,
+ reference_mel,
+ num_steps=10,
+ guidance_scale=0.5,
+ sway_coefficient=-1.0,
+ **kwargs,
+ ):
+ """Generates a waveform from input code and conditioning parameters."""
+
+ mel_spectrogram = self.code2wav_dit_model.sample(
+ conditioning,
+ reference_mel,
+ code,
+ num_steps=num_steps,
+ guidance_scale=guidance_scale,
+ sway_coefficient=sway_coefficient,
+ )
+
+ waveform = self.code2wav_bigvgan_model(mel_spectrogram)
+
+ return waveform
+
+
+############################
+# Start Qwen2.5Omni #
+############################
+
+
+@auto_docstring(
+ custom_intro="""
+ The full Qwen2.5Omni model, a multimodal model composed of 3 sub-models:
+ - [`Qwen2_5OmniThinkerForConditionalGeneration`]:
+ a causal auto-regressive transformer takes text, audio, image, video as input and predict text tokens.
+ - [`Qwen2_5OmniTalkerForConditionalGeneration`]:
+ a causal auto-regressive transformer takes thinker hidden states and response as input and predict speech tokens.
+ - [`Qwen2_5OmniToken2WavModel`]:
+ a DiT model take speech tokens as input and predict mel spectrogram and a BigVGAN vocoder take mel spectrogram as input and predict waveform.
+ """
+)
+class Qwen2_5OmniForConditionalGeneration(Qwen2_5OmniPreTrainedModel, GenerationMixin):
+ config: Qwen2_5OmniConfig
+ output_modalities = ("audio", "text")
+ _no_split_modules = [
+ "Qwen2_5OmniTalkerForConditionalGeneration",
+ "Qwen2_5OmniToken2WavModel",
+ ]
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.thinker = Qwen2_5OmniThinkerForConditionalGeneration(config.thinker_config)
+
+ self.has_talker = config.enable_audio_output
+ self.speaker_map = {}
+ if config.enable_audio_output:
+ self.enable_talker()
+ self.post_init()
+
+ def enable_talker(self):
+ self.talker = Qwen2_5OmniTalkerForConditionalGeneration(self.config.talker_config)
+ self.token2wav = Qwen2_5OmniToken2WavModel(self.config.token2wav_config)
+ self.token2wav.float()
+ self.has_talker = True
+
+ def load_speakers(self, path):
+ check_torch_load_is_safe()
+ for key, value in torch.load(path, weights_only=True).items():
+ self.speaker_map[key] = value
+ logger.info(f"Speaker {list(self.speaker_map.keys())} loaded")
+
+ def disable_talker(self):
+ if hasattr(self, "talker"):
+ del self.talker
+ if hasattr(self, "token2wav"):
+ del self.token2wav
+ self.has_talker = False
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ pretrained_model_name_or_path,
+ *model_args,
+ config=None,
+ cache_dir=None,
+ ignore_mismatched_sizes=False,
+ force_download=False,
+ local_files_only=False,
+ token=None,
+ revision="main",
+ use_safetensors=None,
+ weights_only=True,
+ **kwargs,
+ ):
+ model = super().from_pretrained(
+ pretrained_model_name_or_path,
+ *model_args,
+ config=config,
+ cache_dir=cache_dir,
+ ignore_mismatched_sizes=ignore_mismatched_sizes,
+ force_download=force_download,
+ local_files_only=local_files_only,
+ token=token,
+ revision=revision,
+ use_safetensors=use_safetensors,
+ weights_only=weights_only,
+ **kwargs,
+ )
+ spk_path = cached_file(
+ pretrained_model_name_or_path,
+ "spk_dict.pt",
+ subfolder=kwargs.pop("subfolder", None),
+ cache_dir=kwargs.pop("cache_dir", None),
+ force_download=kwargs.pop("force_download", False),
+ proxies=kwargs.pop("proxies", None),
+ local_files_only=kwargs.pop("local_files_only", False),
+ token=token,
+ revision=kwargs.pop("revision", None),
+ )
+ if spk_path is None:
+ raise ValueError(f"""{pretrained_model_name_or_path}/{spk_path} not exists""")
+ model.load_speakers(spk_path)
+
+ return model
+
+ @torch.no_grad()
+ @deprecate_kwarg("return_audio", version="v5", new_name="generation_mode")
+ # TODO: raushan, defaults should be saved in generation config
+ def generate(
+ self,
+ input_ids: torch.Tensor | None = None,
+ speaker: str = "Chelsie",
+ use_audio_in_video: bool = False,
+ thinker_max_new_tokens: int = 1024,
+ talker_max_new_tokens: int = 4096,
+ talker_do_sample: bool = True,
+ talker_top_k: int = 40,
+ talker_top_p: float = 0.8,
+ talker_temperature: float = 0.9,
+ talker_eos_token_id: list[int] = [8292, 8294],
+ talker_repetition_penalty: float = 1.05,
+ **kwargs,
+ ):
+ r"""
+ Generate text response and audio from input.
+
+ Args:
+ input_ids (`Optional[torch.Tensor]`, *optional*):
+ Input ids, should obtain from processor.
+ speaker (`str` , defaults to "Chelsie"):
+ Which speaker should be used in audio response.
+ use_audio_in_video (`bool`, defaults to False):
+ Whether or not use audio track in video, should same as the parameter in `process_audio_info`.
+ generation_mode (`Optional[str]`, *optional*):
+ Whether or not return response in audio format. When `generation_mode="audio"`, this parameter is same as `config.enable_audio_output`.
+ kwargs (*optional*):
+ - Without a prefix, they will be entered as `**kwargs` for the `generate` method of each sub-model.
+ - With a *thinker_*, *talker_*, *token2wav_* prefix, they will be input for the `generate` method of the
+ thinker, talker and token2wav respectively. It has the priority over the keywords without a prefix.
+ Returns:
+ When `return_audio=False`:
+ - **Text** (`torch.Tensor`): Generated text token sequence.
+ When `return_audio=True`:
+ - **Text** (`torch.Tensor`): Generated text token sequence.
+ - **Audio waveform** (`torch.Tensor`): Generated audio waveform.
+ """
+ # check `False` on purpose because the parameter can be `str/bool`. This is needed for BC
+ generation_mode = kwargs.pop("generation_mode", None)
+ return_audio = generation_mode != "text" and generation_mode is not False
+
+ if speaker not in self.speaker_map:
+ raise ValueError(f"{speaker} is not available, available speakers: {self.speaker_map.keys()}")
+ if return_audio and not self.has_talker:
+ raise ValueError(
+ "Cannot use talker when talker module not initialized. Use `enable_talker` method or set enable_talker in config to enable talker."
+ )
+ if return_audio is None:
+ return_audio = self.has_talker
+ if input_ids.shape[0] != 1 and return_audio:
+ raise NotImplementedError("Qwen2.5-Omni currently does not support batched inference with audio output")
+
+ shared_kwargs = {"use_audio_in_video": use_audio_in_video}
+ thinker_kwargs = {
+ "max_new_tokens": thinker_max_new_tokens,
+ }
+ talker_kwargs = {
+ "max_new_tokens": talker_max_new_tokens,
+ "do_sample": talker_do_sample,
+ "top_k": talker_top_k,
+ "top_p": talker_top_p,
+ "temperature": talker_temperature,
+ "eos_token_id": talker_eos_token_id,
+ "repetition_penalty": talker_repetition_penalty,
+ }
+ token2wav_kwargs = {}
+
+ for key, value in kwargs.items():
+ if key.startswith("thinker_"):
+ thinker_kwargs[key[len("thinker_") :]] = value
+ elif key.startswith("talker_"):
+ talker_kwargs[key[len("talker_") :]] = value
+ elif key.startswith("token2wav_"):
+ token2wav_kwargs[key[len("token2wav_") :]] = value
+ # Process special input values
+ elif key == "feature_attention_mask":
+ thinker_kwargs[key] = value
+ talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
+ elif key == "input_features" or key == "attention_mask":
+ thinker_kwargs[key] = value
+ # Put other key to shared kwargs
+ else:
+ shared_kwargs[key] = value
+
+ # Merge kwargs
+ for key, value in shared_kwargs.items():
+ if key not in thinker_kwargs:
+ thinker_kwargs[key] = value
+ if key not in talker_kwargs:
+ talker_kwargs[key] = value
+ if key not in token2wav_kwargs:
+ token2wav_kwargs[key] = value
+ speaker_params = self.speaker_map[speaker]
+
+ # 1. Generate from thinker module
+ generate_audio = return_audio and self.has_talker
+ if generate_audio:
+ thinker_kwargs["output_hidden_states"] = True
+ thinker_kwargs["return_dict_in_generate"] = True
+
+ thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
+
+ if not generate_audio:
+ return thinker_result
+
+ # 2. Generate speech tokens from talker module
+ embeds_to_talker = thinker_result.hidden_states[0][0].clone().to(input_ids.device)
+ if thinker_kwargs.get("input_features") is not None:
+ audio_ids_mask = input_ids == self.config.thinker_config.audio_token_index
+ audio_mask = audio_ids_mask.unsqueeze(-1).expand_as(embeds_to_talker)
+ audio_mask_tensor = torch.zeros(
+ [audio_ids_mask.sum(), embeds_to_talker.shape[-1]],
+ dtype=embeds_to_talker.dtype,
+ device=input_ids.device,
+ )
+ embeds_to_talker.masked_scatter_(audio_mask, audio_mask_tensor)
+ if thinker_kwargs.get("pixel_values") is not None:
+ image_ids_mask = input_ids == self.config.thinker_config.image_token_index
+ image_mask = image_ids_mask.unsqueeze(-1).expand_as(embeds_to_talker)
+ image_mask_tensor = torch.zeros(
+ [image_ids_mask.sum(), embeds_to_talker.shape[-1]],
+ dtype=embeds_to_talker.dtype,
+ device=input_ids.device,
+ )
+ embeds_to_talker.masked_scatter_(image_mask, image_mask_tensor)
+ if thinker_kwargs.get("pixel_values_videos") is not None:
+ video_ids_mask = input_ids == self.config.thinker_config.video_token_index
+ video_mask = video_ids_mask.unsqueeze(-1).expand_as(embeds_to_talker)
+ video_mask_tensor = torch.zeros(
+ [video_ids_mask.sum(), embeds_to_talker.shape[-1]],
+ dtype=embeds_to_talker.dtype,
+ device=input_ids.device,
+ )
+ embeds_to_talker.masked_scatter_(video_mask, video_mask_tensor)
+
+ processed_thinker_hidden = (
+ (embeds_to_talker,) + thinker_result.hidden_states[0][1:],
+ ) + thinker_result.hidden_states[1:]
+ thinker_generate_ids = thinker_result.sequences[:, input_ids.size(1) :].to(input_ids.device)
+ thinker_token_embeds = [
+ token_hidden_states[0].to(input_ids.device) for token_hidden_states in processed_thinker_hidden
+ ]
+ thinker_hidden_states = [
+ token_hidden_states[-1].to(input_ids.device) for token_hidden_states in processed_thinker_hidden
+ ]
+
+ talker_text_bos_token = speaker_params["bos_token"]
+ talker_input_text_ids = torch.cat(
+ [
+ input_ids,
+ torch.tensor([[talker_text_bos_token]], dtype=torch.long, device=input_ids.device),
+ thinker_generate_ids[:, :1],
+ ],
+ dim=-1,
+ )
+
+ talker_input_ids = torch.cat(
+ [
+ torch.full_like(input_ids, fill_value=self.talker.codec_mask_token),
+ torch.tensor([[self.talker.codec_pad_token]], dtype=torch.long, device=input_ids.device),
+ torch.tensor([[self.talker.codec_bos_token]], dtype=torch.long, device=input_ids.device),
+ ],
+ dim=1,
+ )
+
+ thinker_embed_tokens = self.thinker.get_input_embeddings()
+ thinker_reply_part = torch.cat(thinker_hidden_states[1:], dim=1) + torch.cat(thinker_token_embeds[1:], dim=1)
+ talker_inputs_embeds = thinker_hidden_states[0] + thinker_token_embeds[0]
+ talker_text_bos_token = torch.tensor([[talker_text_bos_token]], dtype=torch.long, device=input_ids.device)
+ talker_text_bos_embed = thinker_embed_tokens(talker_text_bos_token).to(input_ids.device)
+ talker_inputs_embeds = torch.cat(
+ [
+ talker_inputs_embeds,
+ talker_text_bos_embed,
+ thinker_reply_part[:, :1, :],
+ ],
+ dim=1,
+ )
+
+ eos_token = torch.tensor([[self.talker.text_eos_token]], dtype=torch.long, device=input_ids.device)
+ eos_embedding = thinker_embed_tokens(eos_token).to(input_ids.device)
+
+ pad_token = torch.tensor([[self.talker.text_pad_token]], dtype=torch.long, device=input_ids.device)
+ pad_embedding = thinker_embed_tokens(pad_token).to(input_ids.device)
+
+ thinker_reply_part = torch.cat(
+ [
+ thinker_reply_part[:, 1:, :],
+ eos_embedding,
+ pad_embedding,
+ ],
+ dim=1,
+ )
+
+ talker_attention_mask = None
+ if "attention_mask" in kwargs:
+ talker_attention_mask = torch.cat(
+ [kwargs["attention_mask"], kwargs["attention_mask"].new_ones((1, 2))], dim=1
+ ).to(input_ids.device)
+
+ talker_result = self.talker.generate(
+ input_ids=talker_input_ids,
+ input_text_ids=talker_input_text_ids,
+ thinker_reply_part=thinker_reply_part,
+ inputs_embeds=talker_inputs_embeds,
+ attention_mask=talker_attention_mask,
+ suppress_tokens=[self.talker.codec_bos_token],
+ **{k: (v.to(input_ids.device) if torch.is_tensor(v) else v) for k, v in talker_kwargs.items()},
+ )
+ talker_generate_codes = talker_result[:, talker_input_ids.shape[1] : -1]
+
+ # 3. Generate wavs from code
+ if self.token2wav.dtype != torch.float:
+ self.token2wav.float()
+
+ wav = self.token2wav(
+ talker_generate_codes.to(input_ids.device),
+ conditioning=speaker_params["cond"].to(input_ids.device).float(),
+ reference_mel=speaker_params["ref_mel"].to(input_ids.device).float(),
+ **token2wav_kwargs,
+ )
+
+ return thinker_result.sequences, wav.float()
+
+
+__all__ = [
+ "Qwen2_5OmniForConditionalGeneration",
+ "Qwen2_5OmniThinkerTextModel",
+ "Qwen2_5OmniThinkerForConditionalGeneration",
+ "Qwen2_5OmniTalkerModel",
+ "Qwen2_5OmniTalkerForConditionalGeneration",
+ "Qwen2_5OmniToken2WavDiTModel",
+ "Qwen2_5OmniToken2WavBigVGANModel",
+ "Qwen2_5OmniToken2WavModel",
+ "Qwen2_5OmniPreTrainedModel",
+ "Qwen2_5OmniPreTrainedModelForConditionalGeneration",
+]
diff --git a/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
new file mode 100644
index 0000000000000000000000000000000000000000..acebc79a3c9483fd359c64dd057c3b1a58bdf8d1
--- /dev/null
+++ b/LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
@@ -0,0 +1,3916 @@
+# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Qwen2.5Omni model (Audio, Image, Video)."""
+
+import math
+import warnings
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from huggingface_hub.dataclasses import strict
+from torch import nn
+from torch.nn import Parameter
+
+from ... import initialization as init
+from ...cache_utils import Cache
+from ...configuration_utils import PreTrainedConfig
+from ...generation import GenerationMixin
+from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput
+from ...modeling_rope_utils import RopeParameters
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+ TransformersKwargs,
+ auto_docstring,
+ can_return_tuple,
+ check_torch_load_is_safe,
+ logging,
+ torch_compilable_check,
+)
+from ...utils.deprecation import deprecate_kwarg
+from ...utils.generic import accepts_precomputed_kwargs, is_flash_attention_requested, merge_with_config_defaults
+from ...utils.hub import cached_file
+from ...utils.output_capturing import capture_outputs
+from ...vision_utils import get_vision_cu_seqlens, get_vision_position_ids, get_vision_window_index
+from ..llama.modeling_llama import LlamaRotaryEmbedding, rotate_half
+from ..qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig
+from ..qwen2_5_vl.modeling_qwen2_5_vl import (
+ Qwen2_5_VisionRotaryEmbedding,
+ Qwen2_5_VisionTransformerPretrainedModel,
+ Qwen2_5_VLAttention,
+ Qwen2_5_VLMLP,
+ Qwen2_5_VLPreTrainedModel,
+ Qwen2_5_VLTextModel,
+ Qwen2_5_VLVisionBlock,
+ eager_attention_forward,
+)
+from ..qwen2_audio.configuration_qwen2_audio import Qwen2AudioEncoderConfig
+from ..qwen2_audio.modeling_qwen2_audio import Qwen2AudioEncoderLayer
+from ..qwen2_vl.modeling_qwen2_vl import Qwen2VLRotaryEmbedding
+
+
+logger = logging.get_logger(__name__)
+
+
+def chunk_and_pad_features(
+ input_features: torch.Tensor, feature_lens: torch.Tensor, n_window: int, kwargs: dict | None = None
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Split audio features into fixed-size chunks and pad to uniform length, or pop precomputed pair from `kwargs`.
+
+ Each audio sample is split into chunks of `n_window * 2` frames (the last
+ chunk may be shorter), then all chunks are right-padded to the longest chunk.
+
+ Args:
+ input_features: `(feature_dim, total_frames)` concatenated audio features.
+ feature_lens: `(batch_size,)` per-sample frame counts.
+ n_window: half the target chunk size in frames.
+ kwargs: optional caller kwargs — if it contains both `"padded_feature"` and `"chunk_lengths"` they are popped and returned.
+
+ Returns:
+ `padded_feature`: `(num_chunks, feature_dim, max_chunk_len)` padded chunks.
+ `chunk_lengths`: `(num_chunks,)` actual length of each chunk before padding.
+ """
+ if kwargs is not None:
+ padded_feature = kwargs.pop("padded_feature", None)
+ chunk_lengths = kwargs.pop("chunk_lengths", None)
+ if padded_feature is not None and chunk_lengths is not None:
+ return padded_feature, chunk_lengths
+
+ chunk_num = torch.ceil(feature_lens / (n_window * 2)).long()
+ chunk_lengths = torch.full((chunk_num.sum(),), n_window * 2, dtype=torch.long, device=feature_lens.device)
+ tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
+ chunk_lengths[tail_chunk_index] = feature_lens % (n_window * 2)
+ chunk_lengths = torch.where(chunk_lengths == 0, n_window * 2, chunk_lengths)
+
+ chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
+ padded_feature = nn.utils.rnn.pad_sequence(chunk_list, batch_first=True).transpose(1, 2)
+ return padded_feature, chunk_lengths
+
+
+def get_audio_cu_seqlens(chunk_lengths: torch.Tensor, kwargs: dict | None = None) -> torch.Tensor:
+ """Compute cumulative sequence lengths for audio attention, or pop `"cu_seqlens"` from `kwargs` if precomputed.
+
+ Applies one stride-2 convolution length reduction, then returns cumulative
+ boundaries for flash-attention-style sequence packing.
+
+ Args:
+ chunk_lengths: `(num_chunks,)` pre-CNN chunk lengths.
+ kwargs: optional caller kwargs — if it contains `"cu_seqlens"` it is popped and returned.
+
+ Returns:
+ `(num_chunks + 1,)` int32 cumulative sequence boundaries.
+ """
+ if kwargs is not None and (cu_seqlens := kwargs.pop("cu_seqlens", None)) is not None:
+ return cu_seqlens
+ after_conv1 = (chunk_lengths - 1) // 2 + 1
+ return F.pad(after_conv1.cumsum(0), (1, 0), value=0).to(torch.int32)
+
+
+def get_valid_indices(chunk_lengths: torch.Tensor, kwargs: dict | None = None) -> torch.Tensor:
+ """Compute flat indices of valid (non-padding) positions after one stride-2 conv, or pop `"valid_indices"` from `kwargs` if precomputed.
+
+ Args:
+ chunk_lengths: `(num_chunks,)` pre-CNN chunk lengths.
+ kwargs: optional caller kwargs — if it contains `"valid_indices"` it is popped and returned.
+
+ Returns:
+ `(total_valid,)` flat indices into the `(num_chunks * max_len_after_conv)` grid.
+ """
+ if kwargs is not None and (valid_indices := kwargs.pop("valid_indices", None)) is not None:
+ return valid_indices
+ after_conv1 = (chunk_lengths - 1) // 2 + 1
+ max_len = after_conv1.max().item()
+ mask = torch.arange(max_len, device=chunk_lengths.device) < after_conv1.unsqueeze(1)
+ return mask.flatten().nonzero().squeeze(-1)
+
+
+def get_pool_indices(feature_lens: torch.Tensor, kwargs: dict | None = None) -> torch.Tensor:
+ """Compute indices for post-encoder stride-2 average pooling, or pop `"pool_indices"` from `kwargs` if precomputed.
+
+ Args:
+ feature_lens: `(batch_size,)` mel spectrogram lengths.
+ kwargs: optional caller kwargs — if it contains `"pool_indices"` it is popped and returned.
+
+ Returns:
+ `(total_pooled,)` flat index of first element of each stride-2 pair.
+ """
+ if kwargs is not None and (pool_indices := kwargs.pop("pool_indices", None)) is not None:
+ return pool_indices
+ after_conv1 = (feature_lens - 1) // 2 + 1
+ num_pooled = (after_conv1 - 2) // 2 + 1
+ offsets = F.pad(after_conv1[:-1].cumsum(0), (1, 0), value=0)
+ pair_offsets = torch.repeat_interleave(offsets, num_pooled)
+ local_indices = torch.arange(num_pooled.sum(), device=feature_lens.device)
+ local_indices -= torch.repeat_interleave(F.pad(num_pooled[:-1].cumsum(0), (1, 0), value=0), num_pooled)
+ return pair_offsets + local_indices * 2
+
+
+@auto_docstring(checkpoint="Qwen/Qwen2.5-Omni-7B")
+@strict
+class Qwen2_5OmniVisionEncoderConfig(Qwen2_5_VLVisionConfig):
+ r"""
+ window_size (`int`, *optional*, defaults to 11):
+ Size of windows.
+ out_hidden_size (`int`, *optional*, defaults to 3584):
+ The output hidden size of the vision model.
+ fullatt_block_indexes (`int`, *optional*, defaults to `[7, 15, 23, 31]`):
+ Indices of layers with full attention
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen2_5OmniVisionEncoderConfig, Qwen2_5OmniVisionEncoder
+
+ >>> # Initializing a Qwen2_5OmniVisionEncoderConfig
+ >>> configuration = Qwen2_5OmniVisionEncoderConfig()
+
+ >>> # Initializing a Qwen2_5OmniVisionEncoder (with random weights)
+ >>> model = Qwen2_5OmniVisionEncoder(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen2_5_omni_vision_encoder"
+
+ tokens_per_second = AttributeError()
+
+
+@auto_docstring(checkpoint="Qwen/Qwen2.5-Omni-7B")
+@strict
+class Qwen2_5OmniAudioEncoderConfig(Qwen2AudioEncoderConfig):
+ r"""
+ max_source_positions (`int`, *optional*, defaults to 1500):
+ The maximum sequence length of log-mel filter-bank features that this model might ever be used with.
+ n_window (`int`, *optional*, defaults to 100):
+ The chunk for conv and flash attn in AudioEncoder.
+ output_dim (`int`, *optional*, defaults to 3584):
+ The output dimension of AudioEncoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen2_5OmniAudioEncoderConfig, Qwen2_5OmniAudioEncoder
+
+ >>> # Initializing a Qwen2_5OmniAudioEncoderConfig
+ >>> configuration = Qwen2_5OmniAudioEncoderConfig()
+
+ >>> # Initializing a Qwen2_5OmniAudioEncoder (with random weights)
+ >>> model = Qwen2_5OmniAudioEncoder(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen2_5_omni_audio_encoder"
+
+ n_window: int = 100
+ output_dim: int = 3584
+ encoder_layerdrop = AttributeError()
+
+
+@auto_docstring(checkpoint="Qwen/Qwen2.5-Omni-7B")
+@strict
+class Qwen2_5OmniTextConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import Qwen2_5OmniThinkerForConditionalGeneration, Qwen2_5OmniThinkerConfig, Qwen2_5OmniAudioEncoderConfig, Qwen2_5OmniVisionEncoderConfig
+
+ >>> # Initializing a Qwen2_5OmniAudioEncoder config
+ >>> audio_config = Qwen2_5OmniAudioEncoderConfig()
+
+ >>> # Initializing a Qwen2_5OmniVisionEncoder config
+ >>> vision_config = Qwen2_5OmniVisionEncoderConfig()
+
+ >>> # Initializing a Qwen2.5OmniThinker configuration
+ >>> configuration = Qwen2_5OmniThinkerConfig(audio_config, vision_config)
+
+ >>> # Initializing a model from the Qwen-Omni style configuration
+ >>> model = Qwen2_5OmniThinkerForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen2_5_omni_text"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ default_theta = 1000000.0
+
+ # Default tensor parallel plan for base model `Qwen25OmniText`
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+ ignore_keys_at_rope_validation = {"mrope_section"}
+
+ vocab_size: int = 152064
+ hidden_size: int = 3584
+ intermediate_size: int = 18944
+ num_hidden_layers: int = 28
+ num_attention_heads: int = 28
+ num_key_value_heads: int | None = 4
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 32768
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-6
+ use_cache: bool = True
+ rope_parameters: RopeParameters | dict | None = None
+ use_sliding_window: bool = False
+ sliding_window: int | None = 32768
+ max_window_layers: int = 28
+ layer_types: list[str] | None = None
+ attention_dropout: float | int = 0.0
+ pad_token_id: int | None = None
+ bos_token_id: int | None = None
+ eos_token_id: int | list[int] | None = None
+ tie_word_embeddings: bool = True
+
+ def __post_init__(self, **kwargs):
+ self.sliding_window = self.sliding_window if self.use_sliding_window else None
+ if self.num_key_value_heads is None:
+ self.num_key_value_heads = self.num_attention_heads
+
+ if self.layer_types is None:
+ self.layer_types = [
+ "sliding_attention"
+ if self.sliding_window is not None and i >= self.max_window_layers
+ else "full_attention"
+ for i in range(self.num_hidden_layers)
+ ]
+
+ super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="Qwen/Qwen2.5-Omni-7B")
+@strict
+class Qwen2_5OmniThinkerConfig(PreTrainedConfig):
+ r"""
+ position_id_per_seconds (`int`, *optional*, defaults to 25):
+ The increment of position id per second.
+ seconds_per_chunk (`int`, *optional*, defaults to 2):
+ The duration in seconds of the chunk of audio and video data.
+ audio_start_token_id (`int`, *optional*, defaults to 151647):
+ The audio start token index to encode the audio prompt.
+ audio_end_token_id (`int`, *optional*, defaults to 151648):
+ The audio end token index to encode the audio prompt.
+ user_token_id (`int, *optional*, defaults to 872):
+ The user token index to encode the user token.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen2_5OmniThinkerForConditionalGeneration, Qwen2_5OmniThinkerConfig, Qwen2_5OmniAudioEncoderConfig, Qwen2_5OmniVisionEncoderConfig
+
+ >>> # Initializing a Qwen2_5OmniAudioEncoder config
+ >>> audio_config = Qwen2_5OmniAudioEncoderConfig()
+
+ >>> # Initializing a Qwen2_5OmniVisionEncoder config
+ >>> vision_config = Qwen2_5OmniVisionEncoderConfig()
+
+ >>> # Initializing a Qwen2_5OmniTextConfig config
+ >>> text_config = Qwen2_5OmniTextConfig()
+
+ >>> # Initializing a Qwen2.5OmniThinker configuration
+ >>> configuration = Qwen2_5OmniThinkerConfig(audio_config, vision_config, text_config)
+
+ >>> # Initializing a model from the Qwen-Omni style configuration
+ >>> model = Qwen2_5OmniThinkerForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen2_5_omni_thinker"
+ attribute_map = {
+ "image_token_id": "image_token_index",
+ "video_token_id": "video_token_index",
+ "audio_token_id": "audio_token_index",
+ }
+ sub_configs = {
+ "audio_config": Qwen2_5OmniAudioEncoderConfig,
+ "vision_config": Qwen2_5OmniVisionEncoderConfig,
+ "text_config": Qwen2_5OmniTextConfig,
+ }
+
+ audio_config: dict | PreTrainedConfig | None = None
+ vision_config: dict | PreTrainedConfig | None = None
+ text_config: dict | PreTrainedConfig | None = None
+ audio_token_index: int = 151646
+ image_token_index: int = 151655
+ video_token_index: int = 151656
+ position_id_per_seconds: int = 25
+ seconds_per_chunk: int = 2
+ audio_start_token_id: int = 151647
+ audio_end_token_id: int = 151648
+ user_token_id: int = 872
+ initializer_range: float = 0.02
+ tie_word_embeddings: bool = False
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.vision_config, dict):
+ self.vision_config = Qwen2_5OmniVisionEncoderConfig(**self.vision_config)
+ elif self.vision_config is None:
+ self.vision_config = Qwen2_5OmniVisionEncoderConfig()
+
+ if isinstance(self.audio_config, dict):
+ self.audio_config = Qwen2_5OmniAudioEncoderConfig(**self.audio_config)
+ elif self.audio_config is None:
+ self.audio_config = Qwen2_5OmniAudioEncoderConfig()
+
+ if isinstance(self.text_config, dict):
+ self.text_config = Qwen2_5OmniTextConfig(**self.text_config)
+ elif self.text_config is None:
+ self.text_config = Qwen2_5OmniTextConfig()
+
+ super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="Qwen/Qwen2.5-Omni-7B")
+@strict
+class Qwen2_5OmniTalkerConfig(PreTrainedConfig):
+ r"""
+ tts_text_start_token_id (`int`, *optional*, defaults to 151860):
+ The tts text start token index to encode the start of tts text.
+ tts_text_end_token_id (`int`, *optional*, defaults to 151861):
+ The tts text end token index to encode the end of tts text.
+ tts_text_pad_token_id (`int`, *optional*, defaults to 151859):
+ The tts text pad token index to encode the pad of tts text.
+ tts_codec_start_token_id (`int`, *optional*, defaults to 8293):
+ The tts codec start token index to encode the start of tts codec.
+ tts_codec_end_token_id (`int`, *optional*, defaults to 8294):
+ The tts codec end token index to encode the end of tts codec.
+ tts_codec_pad_token_id (`int`, *optional*, defaults to 8292):
+ The tts codec pad token index to encode the pad of tts codec.
+ tts_codec_mask_token_id (`int`, *optional*, defaults to 8296):
+ The tts codec mask token index to encode the mask of tts codec.
+ position_id_per_seconds (`int`, *optional*, defaults to 25):
+ The increment of position id per second.
+ seconds_per_chunk (`int`, *optional*, defaults to 2):
+ The duration in seconds of the chunk of audio and video data.
+ audio_start_token_id (`int`, *optional*, defaults to 151647):
+ The audio start token index to encode the audio prompt.
+ audio_end_token_id (`int`, *optional*, defaults to 151648):
+ The audio end token index to encode the audio prompt.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen2_5OmniTalkerForConditionalGeneration, Qwen2_5OmniThinkerConfig, Qwen2_5OmniAudioEncoderConfig, Qwen2_5OmniVisionEncoderConfig
+
+ >>> # Initializing a Qwen2_5OmniAudioEncoder config
+ >>> audio_config = Qwen2_5OmniAudioEncoderConfig()
+
+ >>> # Initializing a Qwen2 config
+ >>> text_config = Qwen2Config()
+
+ >>> # Initializing a Qwen2_5Omni configuration
+ >>> configuration = Qwen2_5OmniThinkerConfig(audio_config, text_config)
+
+ >>> # Initializing a model from the qwen2-audio style configuration
+ >>> model = Qwen2_5OmniTalkerForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "qwen2_5_omni_talker"
+ default_theta = 1000000.0
+ attribute_map = {
+ "image_token_id": "image_token_index",
+ "video_token_id": "video_token_index",
+ "audio_token_id": "audio_token_index",
+ }
+ ignore_keys_at_rope_validation = {"mrope_section"}
+
+ audio_token_index: int = 151646
+ image_token_index: int = 151655
+ video_token_index: int = 151656
+ vocab_size: int = 8448
+ tts_text_start_token_id: int = 151860
+ tts_text_end_token_id: int = 151861
+ tts_text_pad_token_id: int = 151859
+ tts_codec_start_token_id: int = 8293
+ tts_codec_end_token_id: int = 8294
+ tts_codec_pad_token_id: int = 8292
+ tts_codec_mask_token_id: int = 8296
+ vision_start_token_id: int = 151652
+ vision_end_token_id: int = 151653
+ embedding_size: int = 3584
+ hidden_size: int = 3584
+ intermediate_size: int = 18944
+ num_hidden_layers: int = 28
+ num_attention_heads: int = 28
+ num_key_value_heads: int = 4
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 32768
+ rms_norm_eps: float = 1e-06
+ head_dim: int = 128
+ use_cache: bool = True
+ tie_word_embeddings: bool = False
+ use_sliding_window: bool = False
+ sliding_window: int | None = 32768
+ max_window_layers: int = 28
+ attention_dropout: float | int = 0.0
+ rope_parameters: RopeParameters | dict | None = None
+ position_id_per_seconds: int = 25
+ seconds_per_chunk: int = 2
+ audio_start_token_id: int = 151647
+ audio_end_token_id: int = 151648
+ initializer_range: float = 0.02
+ spatial_merge_size: int = 2
+ layer_types: list[str] | None = None
+ pad_token_id: int | None = None
+
+ def __post_init__(self, **kwargs):
+ self.sliding_window = self.sliding_window if self.use_sliding_window else None
+
+ if self.num_key_value_heads is None:
+ self.num_key_value_heads = self.num_attention_heads
+
+ if self.layer_types is None:
+ self.layer_types = [
+ "sliding_attention"
+ if self.sliding_window is not None and i >= self.max_window_layers
+ else "full_attention"
+ for i in range(self.num_hidden_layers)
+ ]
+
+ super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="Qwen/Qwen2.5-Omni-7B")
+@strict
+class Qwen2_5OmniDiTConfig(PreTrainedConfig):
+ r"""
+ ff_mult (`int`, *optional*, defaults to 2):
+ The multiplier for the feedforward layer in each transformer block.
+ emb_dim (`int`, *optional*, defaults to 512):
+ The dimension of the embedding layer.
+ block_size (`int`, *optional*, defaults to 64):
+ Number of tokens (frames) in each processing block.
+ look_ahead_layers (`list[int]`, *optional*, defaults to `[10]`):
+ Number of transformer layers that are permitted to attend to future blocks
+ look_backward_layers (`list[int]`, *optional*, defaults to `[0, 20]`):
+ Number of transformer layers that attend to past blocks beyond the current block boundary
+ repeats (`int`, *optional*, defaults to 2):
+ The number of times the codec embeddings are repeated.
+ num_embeds (`int`, *optional*, defaults to 8193):
+ The number of unique embeddings in the codec.
+ mel_dim (`int`, *optional*, defaults to 80):
+ The dimension of the mel-spectrogram.
+ enc_emb_dim (`int`, *optional*, defaults to 192):
+ The dimension of the pre-trained speaker embedding.
+ enc_dim (`int`, *optional*, defaults to 128):
+ The dimension of the encoder output.
+ enc_channels (`list[int]`, *optional*, defaults to `[256, 256, 256, 256, 768]`):
+ A list of output channels for each TDNN/SERes2Net layer in the encoder.
+ enc_kernel_sizes (`list[int]`, *optional*, defaults to `[5, 3, 3, 3, 1]`):
+ A list of kernel sizes for each layer in the encoder.
+ enc_dilations (`list[int]`, *optional*, defaults to `[1, 2, 3, 4, 1]`):
+ A list of dilations for each layer in the encoder.
+ enc_attention_channels (`int`, *optional*, defaults to 64):
+ The number of attention channels in the SqueezeExcitationBlock.
+ enc_res2net_scale (`int`, *optional*, defaults to 2):
+ The scale of the Res2Net block in the encoder.
+ enc_se_channels (`int`, *optional*, defaults to 64):
+ The number of output channels after squeeze in the SqueezeExcitationBlock.
+ """
+
+ model_type = "qwen2_5_omni_dit"
+
+ hidden_size: int = 1024
+ num_hidden_layers: int = 22
+ num_attention_heads: int = 16
+ ff_mult: int = 2
+ emb_dim: int = 512
+ head_dim: int = 64
+ rope_parameters: RopeParameters | dict | None = None
+ max_position_embeddings: int = 32768
+ block_size: int = 24
+ look_ahead_layers: list[int] | tuple[int, ...] = (10,)
+ look_backward_layers: list[int] | tuple[int, ...] = (0, 20)
+ repeats: int = 2
+ num_embeds: int = 8193
+ mel_dim: int = 80
+ dropout: float | int = 0.1
+ enc_emb_dim: int = 192
+ enc_dim: int = 128
+ enc_channels: list[int] | tuple[int, ...] = (256, 256, 256, 256, 768)
+ enc_kernel_sizes: list[int] | tuple[int, ...] = (5, 3, 3, 3, 1)
+ enc_dilations: list[int] | tuple[int, ...] = (1, 2, 3, 4, 1)
+ enc_attention_channels: int = 64
+ enc_res2net_scale: int = 2
+ enc_se_channels: int = 64
+
+
+@auto_docstring(checkpoint="Qwen/Qwen2.5-Omni-7B")
+@strict
+class Qwen2_5OmniBigVGANConfig(PreTrainedConfig):
+ r"""
+ mel_dim (`int`, *optional*, defaults to 80):
+ The dimension of the mel-spectrogram.
+ upsample_initial_channel (`int`, *optional*, defaults to 1536):
+ The number of channels in the initial upsampling layer.
+ resblock_kernel_sizes (`list[int]`, *optional*, defaults to `[3, 7, 11]`):
+ A list of kernel sizes for each residual block.
+ resblock_dilation_sizes (`list[list[int]]`, *optional*, defaults to `[[1, 3, 5], [1, 3, 5], [1, 3, 5]]`):
+ A list of dilation sizes for each residual block.
+ upsample_rates (`list[int]`, *optional*, defaults to `[5, 3, 2, 2, 2, 2]`):
+ A list of upsampling rates for each upsampling layer.
+ upsample_kernel_sizes (`list[int]`, *optional*, defaults to `[11, 7, 4, 4, 4, 4]`):
+ A list of kernel sizes for each upsampling layer.
+ """
+
+ model_type = "qwen2_5_omni_bigvgan"
+
+ mel_dim: int = 80
+ upsample_initial_channel: int = 1536
+ resblock_kernel_sizes: list[int] | tuple[int, ...] = (3, 7, 11)
+ resblock_dilation_sizes: list | tuple = ((1, 3, 5), (1, 3, 5), (1, 3, 5))
+ upsample_rates: list[int] | tuple[int, ...] = (5, 3, 2, 2, 2, 2)
+ upsample_kernel_sizes: list[int] | tuple[int, ...] = (11, 7, 4, 4, 4, 4)
+
+
+@auto_docstring(checkpoint="Qwen/Qwen2.5-Omni-7B")
+@strict
+class Qwen2_5OmniToken2WavConfig(PreTrainedConfig):
+ r"""
+ dit_config ([`DiT_Args`], *optional*):
+ Configuration class for the Diffusion Transformer (DiT) module responsible for generating mel-spectrograms.
+ bigvgan_config ([`BigVGAN_Args`], *optional*):
+ Configuration class for the BigVGAN module responsible for converting mel-spectrograms to waveforms.
+
+ Example:
+
+ ```python
+ >>> from transformers import Qwen2_5OmniToken2WavModel, DiT_Args, BigVGAN_Args
+
+ >>> # Initialize DiT configuration
+ >>> dit_config = DiT_Args(
+ ... dim=1024,
+ ... depth=22,
+ ... heads=16,
+ ... ff_mult=2
+ ... )
+
+ >>> # Initialize BigVGAN configuration
+ >>> bigvgan_config = BigVGAN_Args(
+ ... mel_dim=80,
+ ... upsample_rates=[5,3,2,2,2,2]
+ ... )
+
+ >>> # Initialize main configuration
+ >>> config = Qwen2_5OmniToken2WavConfig(dit_config, bigvgan_config)
+
+ >>> # Initialize model with config
+ >>> model = Qwen2_5OmniToken2Wav(config)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+ """
+
+ model_type = "qwen2_5_omni_token2wav"
+ sub_configs = {
+ "dit_config": Qwen2_5OmniDiTConfig,
+ "bigvgan_config": Qwen2_5OmniBigVGANConfig,
+ }
+
+ dit_config: dict | PreTrainedConfig | None = None
+ bigvgan_config: dict | PreTrainedConfig | None = None
+
+ def __post_init__(self, **kwargs):
+ if self.dit_config is None:
+ self.dit_config = Qwen2_5OmniDiTConfig()
+ elif isinstance(self.dit_config, dict):
+ self.dit_config = Qwen2_5OmniDiTConfig(**self.dit_config)
+
+ if self.bigvgan_config is None:
+ self.bigvgan_config = Qwen2_5OmniBigVGANConfig()
+ elif isinstance(self.bigvgan_config, dict):
+ self.bigvgan_config = Qwen2_5OmniBigVGANConfig(**self.bigvgan_config)
+
+ super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="Qwen/Qwen2.5-Omni-7B")
+@strict
+class Qwen2_5OmniConfig(PreTrainedConfig):
+ r"""
+ thinker_config (`dict`, *optional*):
+ Configuration of the underlying thinker sub-model.
+ talker_config (`dict`, *optional*):
+ Configuration of the underlying talker sub-model.
+ token2wav_config (`dict`, *optional*):
+ Configuration of the underlying codec sub-model.
+ enable_audio_output (`bool`, *optional*, defaults to `True`):
+ Whether enable audio output and load talker and token2wav module.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... Qwen2_5OmniThinkerConfig,
+ ... Qwen2_5OmniTalkerConfig,
+ ... Qwen2_5OmniToken2WavConfig,
+ ... Qwen2_5OmniForConditionalGeneration,
+ ... Qwen2_5OmniConfig,
+ ... )
+
+ >>> # Initializing sub-modules configurations.
+ >>> thinker_config = Qwen2_5OmniThinkerConfig()
+ >>> talker_config = Qwen2_5OmniTalkerConfig()
+ >>> token2wav_config = Qwen2_5OmniToken2WavConfig()
+
+
+ >>> # Initializing a module style configuration
+ >>> configuration = Qwen2_5OmniConfig(
+ ... thinker_config, talker_config, token2wav_config
+ ... )
+
+ >>> # Initializing a model (with random weights)
+ >>> model = Qwen2_5OmniForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+ """
+
+ model_type = "qwen2_5_omni"
+ sub_configs = {
+ "thinker_config": Qwen2_5OmniThinkerConfig,
+ "talker_config": Qwen2_5OmniTalkerConfig,
+ "token2wav_config": Qwen2_5OmniToken2WavConfig,
+ }
+
+ thinker_config: dict | PreTrainedConfig | None = None
+ talker_config: dict | PreTrainedConfig | None = None
+ token2wav_config: dict | PreTrainedConfig | None = None
+ enable_audio_output: bool = True
+
+ def __post_init__(self, **kwargs):
+ if self.thinker_config is None:
+ self.thinker_config = Qwen2_5OmniThinkerConfig()
+ logger.info("thinker_config is None. Initializing thinker model with default values")
+ elif isinstance(self.thinker_config, dict):
+ self.thinker_config = Qwen2_5OmniThinkerConfig(**self.thinker_config)
+
+ if self.talker_config is None:
+ self.talker_config = Qwen2_5OmniTalkerConfig()
+ logger.info("talker_config is None. Initializing talker model with default values")
+ elif isinstance(self.talker_config, dict):
+ self.talker_config = Qwen2_5OmniTalkerConfig(**self.talker_config)
+
+ if self.token2wav_config is None:
+ self.token2wav_config = Qwen2_5OmniToken2WavConfig()
+ logger.info("token2wav_config is None. Initializing token2wav model with default values")
+ elif isinstance(self.token2wav_config, dict):
+ self.token2wav_config = Qwen2_5OmniToken2WavConfig(**self.token2wav_config)
+
+ super().__post_init__(**kwargs)
+
+ def get_text_config(self, *args, **kwargs):
+ """
+ Returns the config that is meant to be used with text IO. On most models, it is the original config instance
+ itself. On specific composite models, it is under a set of valid names.
+
+ Args:
+ decoder (`Optional[bool]`, *optional*, defaults to `False`):
+ If set to `True`, then only search for decoder config names.
+ """
+ # Overridden for deeply nested config like Qwen2-Omni. We don't have any omni model
+ # except for Qwen yet. This has to be generalized if more deeply nested configs are
+ # added. NOTE: currently method used only by vLLM
+ return self.thinker_config.get_text_config(*args, **kwargs)
+
+
+class Qwen2_5OmniPreTrainedModel(Qwen2_5_VLPreTrainedModel):
+ config: Qwen2_5OmniConfig
+ input_modalities = ("image", "video", "audio", "text")
+ _can_compile_fullgraph = False
+
+ def _init_weights(self, module):
+ PreTrainedModel._init_weights(self, module)
+ if isinstance(module, SinusoidsPositionEmbedding):
+ log_timescale_increment = np.log(module.max_timescale) / (module.channels // 2 - 1)
+ inv_timescales = torch.exp(-log_timescale_increment * torch.arange(module.channels // 2).float())
+ scaled_time = torch.arange(module.length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+ init.copy_(module.positional_embedding, torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1))
+ elif isinstance(module, UpSample1d):
+ filter_tensor = kaiser_sinc_filter1d(0.5 / module.ratio, 0.6 / module.ratio, module.kernel_size)
+ init.copy_(module.filter, filter_tensor)
+ elif isinstance(module, DownSample1d):
+ filter_tensor = kaiser_sinc_filter1d(module.cutoff, module.half_width, module.kernel_size)
+ init.copy_(module.filter, filter_tensor)
+ elif isinstance(module, Qwen2_5_VisionRotaryEmbedding):
+ inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim))
+ init.copy_(module.inv_freq, inv_freq)
+
+
+class Qwen2_5OmniPreTrainedModelForConditionalGeneration(Qwen2_5OmniPreTrainedModel):
+ input_modalities = ("image", "video", "audio", "text")
+
+ def get_llm_pos_ids_for_vision(
+ self,
+ start_idx: int,
+ vision_idx: int,
+ spatial_merge_size: int,
+ t_index: list[int],
+ grid_hs: list[int],
+ grid_ws: list[int],
+ ):
+ llm_pos_ids_list = []
+ llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
+ llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
+ h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten()
+ w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten()
+ t_index = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().long()
+ _llm_pos_ids = torch.stack([t_index, h_index, w_index])
+ llm_pos_ids_list.append(_llm_pos_ids + start_idx) # + 1 ) # 12.09 by malinhan
+ llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
+ return llm_pos_ids
+
+ def get_chunked_index(
+ self, token_indices: torch.Tensor, tokens_per_chunk: int, remove_index: int
+ ) -> list[tuple[int, int]]:
+ """
+ Splits token index list into chunks based on token value ranges.
+
+ Given a list of token indices, returns a list of (start, end) index tuples representing
+ slices of the list where the token values fall within successive ranges of `t_ntoken_per_chunk`.
+
+ For example, if `t_ntoken_per_chunk` is 1000, the function will create chunks such that:
+ - the first chunk contains token values < 1000,
+ - the second chunk contains values >= 1000 and < 2000, and so on.
+
+ Parameters:
+ token_indices (`torch.Tensor` of shape `(seq_len, )`): A monotonically increasing list of
+ token index values.
+ t_ntoken_per_chunk (`int`): Number of tokens per chunk (used as the chunk size threshold).
+ remove_index (`int`) An index id to subtract from `token_indices` before chunking
+
+ Returns:
+ `list[tuple[int, int]]`: A list of tuples, each representing the start (inclusive)
+ and end (exclusive) indices of a chunk in `token_indices`.
+ """
+
+ def _iter():
+ i, start_idx = 0, 0 # skip bos token
+ current_chunk = 1
+ while i < len(token_indices): # skip eos token
+ if token_indices[i] - remove_index >= current_chunk * tokens_per_chunk:
+ yield (start_idx, i)
+ start_idx = i
+ current_chunk += 1
+ i += 1
+ yield (start_idx, len(token_indices))
+
+ return list(_iter())
+
+ def get_rope_index(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ use_audio_in_video: bool = False,
+ audio_seqlens: torch.LongTensor | None = None,
+ second_per_grids: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
+
+ Explanation:
+ Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
+
+ For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
+ Examples:
+ input_ids: [T T T T T], here T is for text.
+ temporal position_ids: [0, 1, 2, 3, 4]
+ height position_ids: [0, 1, 2, 3, 4]
+ width position_ids: [0, 1, 2, 3, 4]
+
+ For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
+ and 1D rotary position embedding for text part.
+ Examples:
+ Temporal (Time): 3 patches, representing different segments of the video in time.
+ Height: 2 patches, dividing each frame vertically.
+ Width: 2 patches, dividing each frame horizontally.
+ We also have some important parameters:
+ fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.
+ tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity.
+ temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.
+ interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.
+ input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
+ vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]
+ vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
+ vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
+ text temporal position_ids: [101, 102, 103, 104, 105]
+ text height position_ids: [101, 102, 103, 104, 105]
+ text width position_ids: [101, 102, 103, 104, 105]
+ Here we calculate the text start position_ids as the max vision position_ids plus 1.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+ it.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ use_audio_in_video (`bool`, *optional*):
+ If set to `True`, use the audio in video.
+ audio_seqlens (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ second_per_grids (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+ """
+ spatial_merge_size = self.spatial_merge_size
+ image_token_id = self.config.image_token_id
+ video_token_id = self.config.video_token_id
+ audio_token_id = self.config.audio_token_id
+ vision_start_token_id = self.config.vision_start_token_id
+ audio_start_token_id = self.config.audio_start_token_id
+ position_id_per_seconds = self.config.position_id_per_seconds
+ seconds_per_chunk = self.config.seconds_per_chunk
+
+ mrope_position_deltas = []
+ if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+ total_input_ids = input_ids
+ if attention_mask is not None:
+ attention_mask = attention_mask == 1
+ position_ids = torch.ones(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ image_idx, video_idx, audio_idx = 0, 0, 0
+ for i, input_ids in enumerate(total_input_ids):
+ if attention_mask is not None:
+ input_ids = input_ids[attention_mask[i]]
+ image_nums, video_nums, audio_nums = 0, 0, 0
+ vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+ vision_tokens = input_ids[vision_start_indices + 1]
+ audio_nums = torch.sum(input_ids == audio_start_token_id)
+ image_nums = (vision_tokens == image_token_id).sum()
+ video_nums = (
+ (vision_tokens == audio_start_token_id).sum()
+ if use_audio_in_video
+ else (vision_tokens == video_token_id).sum()
+ )
+ input_tokens = input_ids.tolist()
+ llm_pos_ids_list: list = []
+ st = 0
+ remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums
+ multimodal_nums = (
+ image_nums + audio_nums if use_audio_in_video else image_nums + video_nums + audio_nums
+ )
+ for _ in range(multimodal_nums):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ if image_token_id in input_tokens and remain_images > 0:
+ ed_image = input_tokens.index(image_token_id, st)
+ else:
+ ed_image = len(input_tokens) + 1
+ if video_token_id in input_tokens and remain_videos > 0:
+ ed_video = input_tokens.index(video_token_id, st)
+ else:
+ ed_video = len(input_tokens) + 1
+ if audio_token_id in input_tokens and remain_audios > 0:
+ ed_audio = input_tokens.index(audio_token_id, st)
+ else:
+ ed_audio = len(input_tokens) + 1
+ min_ed = min(ed_image, ed_video, ed_audio)
+ if min_ed == ed_audio:
+ text_len = min_ed - st - 1
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ audio_len = ((audio_seqlens[audio_idx] - 1) // 2 + 1 - 2) // 2 + 1
+ llm_pos_ids = torch.arange(audio_len).view(1, -1).expand(3, -1) + st_idx
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st += text_len + bos_len + audio_len + eos_len
+ audio_idx += 1
+ remain_audios -= 1
+
+ elif min_ed == ed_image:
+ text_len = min_ed - st - 1
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ grid_t = image_grid_thw[image_idx][0]
+ grid_hs = image_grid_thw[:, 1]
+ grid_ws = image_grid_thw[:, 2]
+ t_index = (torch.arange(grid_t) * 1 * position_id_per_seconds).long()
+ llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st += text_len + bos_len + image_len + eos_len
+ image_idx += 1
+ remain_images -= 1
+
+ elif min_ed == ed_video and not use_audio_in_video:
+ text_len = min_ed - st - 1
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ grid_t = video_grid_thw[video_idx][0]
+ grid_hs = video_grid_thw[:, 1]
+ grid_ws = video_grid_thw[:, 2]
+ t_index = (
+ torch.arange(grid_t) * second_per_grids[video_idx].cpu().float() * position_id_per_seconds
+ ).long()
+ llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+ video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+ llm_pos_ids_list.append(llm_pos_ids)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st += text_len + bos_len + video_len + eos_len
+ video_idx += 1
+ remain_videos -= 1
+
+ elif min_ed == ed_video and use_audio_in_video:
+ text_len = min_ed - st - 2
+ if text_len != 0:
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ bos_len = 1
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+ llm_pos_ids_list.append(torch.arange(bos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ audio_len = ((audio_seqlens[audio_idx] - 1) // 2 + 1 - 2) // 2 + 1
+ audio_llm_pos_ids = torch.arange(audio_len).view(1, -1).expand(3, -1) + st_idx
+ grid_t = video_grid_thw[video_idx][0]
+ grid_hs = video_grid_thw[:, 1]
+ grid_ws = video_grid_thw[:, 2]
+
+ t_index = (
+ torch.arange(grid_t) * second_per_grids[video_idx].cpu().float() * position_id_per_seconds
+ ).long()
+ video_llm_pos_ids = self.get_llm_pos_ids_for_vision(
+ st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+ )
+
+ t_ntoken_per_chunk = int(position_id_per_seconds * seconds_per_chunk)
+ video_chunk_indexes = self.get_chunked_index(video_llm_pos_ids[0], t_ntoken_per_chunk, st_idx)
+ audio_chunk_indexes = self.get_chunked_index(audio_llm_pos_ids[0], t_ntoken_per_chunk, st_idx)
+ sub_len = 0
+ for j in range(max(len(video_chunk_indexes), len(audio_chunk_indexes))):
+ video_chunk_index = video_chunk_indexes[j] if j < len(video_chunk_indexes) else None
+ audio_chunk_index = audio_chunk_indexes[j] if j < len(audio_chunk_indexes) else None
+ if video_chunk_index is not None:
+ sub_len += video_chunk_index[1] - video_chunk_index[0]
+
+ llm_pos_ids_list.append(
+ video_llm_pos_ids[:, video_chunk_index[0] : video_chunk_index[1]]
+ )
+ if audio_chunk_index is not None:
+ sub_len += audio_chunk_index[1] - audio_chunk_index[0]
+
+ llm_pos_ids_list.append(
+ audio_llm_pos_ids[:, audio_chunk_index[0] : audio_chunk_index[1]]
+ )
+ video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
+
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ eos_len = 1
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+ llm_pos_ids_list.append(torch.arange(eos_len).view(1, -1).expand(3, -1) + st_idx)
+
+ st += text_len + bos_len * 2 + audio_len + video_len + eos_len * 2
+
+ audio_idx += 1
+ video_idx += 1
+ remain_videos -= 1
+ remain_audios -= 1
+
+ if st < len(input_tokens):
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+ text_len = len(input_tokens) - st
+ llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+
+ if attention_mask is not None:
+ position_ids[..., i, attention_mask[i]] = llm_positions.to(position_ids.device)
+ else:
+ position_ids[..., i, :] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(input_ids))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas).unsqueeze(1).to(device=input_ids.device)
+
+ return position_ids, mrope_position_deltas
+ else:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+ max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+ mrope_position_deltas = max_position_ids + 1 - torch.sum(attention_mask, dim=-1, keepdim=True)
+
+ return position_ids, mrope_position_deltas
+
+
+############################
+# Start Thinker #
+############################
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Qwen2.5OmniThinker causal language model (or autoregressive) outputs.
+ """
+)
+@dataclass
+class Qwen2_5OmniThinkerCausalLMOutputWithPast(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`, *optional*):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ rope_deltas: torch.LongTensor | None = None
+
+
+class Qwen2_5OmniAudioAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ config: Qwen2_5OmniAudioEncoderConfig,
+ ):
+ super().__init__()
+ self.embed_dim = config.d_model
+ self.num_heads = config.encoder_attention_heads
+ self.dropout = config.attention_dropout
+ self.head_dim = self.embed_dim // self.num_heads
+ self.num_key_value_groups = 1 # needed for eager attention
+ self.config = config
+
+ if (self.head_dim * self.num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {self.num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = 0.0
+ self.is_decoder = False
+ self.is_causal = False
+
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
+ self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+ self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ **kwargs,
+ ) -> torch.Tensor:
+ """Input shape: Batch x Time x Channel"""
+
+ seq_length, _ = hidden_states.size()
+
+ query_states = self.q_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
+ key_states = self.k_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
+ value_states = self.v_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
+
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ if is_flash_attention_requested(self.config):
+ # Flash Attention: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output
+
+
+class Qwen2_5OmniAudioEncoderLayer(Qwen2AudioEncoderLayer):
+ def __init__(self, config: Qwen2_5OmniAudioEncoderConfig):
+ super().__init__(config)
+ self.self_attn = Qwen2_5OmniAudioAttention(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ **kwargs,
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+ hidden_states = self.self_attn(
+ hidden_states=hidden_states,
+ cu_seqlens=cu_seqlens,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = residual + hidden_states
+
+ if hidden_states.dtype == torch.float16:
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+ hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+ outputs = (hidden_states,)
+
+ return outputs
+
+
+class SinusoidsPositionEmbedding(nn.Module):
+ def __init__(self, length, channels, max_timescale=10000):
+ super().__init__()
+ self.length = length
+ self.channels = channels
+ self.max_timescale = max_timescale
+ if channels % 2 != 0:
+ raise ValueError("SinusoidsPositionEmbedding needs even channels input")
+ log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
+ inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2).float())
+ scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+ self.register_buffer(
+ "positional_embedding",
+ torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1),
+ persistent=False,
+ )
+
+ def forward(self, seqlen: int):
+ return self.positional_embedding[:seqlen, :]
+
+
+@auto_docstring(
+ custom_intro="""
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+ [`Qwen2_5OmniAudioEncoderLayer`].
+ """
+)
+class Qwen2_5OmniAudioEncoder(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniAudioEncoderConfig
+ main_input_name = "input_features"
+ input_modalities = "audio"
+ _no_split_modules = ["Qwen2_5OmniAudioEncoderLayer"]
+ _supports_sdpa = True
+ _can_record_outputs = {
+ "hidden_states": Qwen2_5OmniAudioEncoderLayer,
+ "attentions": Qwen2_5OmniAudioAttention,
+ }
+
+ def __init__(self, config: Qwen2_5OmniAudioEncoderConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+
+ embed_dim = config.d_model
+ self.num_mel_bins = config.num_mel_bins
+ self.max_source_positions = config.max_source_positions
+ self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+ self.n_window = config.n_window
+ self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1)
+ self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)
+ self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, embed_dim)
+ self.audio_bos_eos_token = nn.Embedding(2, config.output_dim)
+ self.layers = nn.ModuleList([Qwen2_5OmniAudioEncoderLayer(config) for _ in range(config.encoder_layers)])
+ self.ln_post = nn.LayerNorm(config.d_model)
+ self.avg_pooler = nn.AvgPool1d(2, stride=2)
+ self.proj = nn.Linear(config.d_model, config.output_dim)
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def _freeze_parameters(self):
+ for param in self.parameters():
+ param.requires_grad = False
+ self._requires_grad = False
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.conv1
+
+ def set_input_embeddings(self, value: nn.Module):
+ self.conv1 = value
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ @auto_docstring
+ def forward(
+ self, input_features=None, feature_lens=None, aftercnn_lens=None, **kwargs: Unpack[TransformersKwargs]
+ ):
+ r"""
+ feature_lens (`torch.LongTensor` of shape `(batch_size,)`):
+ mel length
+ aftercnn_lens (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ mel length after cnn
+ """
+ padded_feature, chunk_lengths = chunk_and_pad_features(
+ input_features, feature_lens, self.n_window, kwargs=kwargs
+ )
+ valid_indices = get_valid_indices(chunk_lengths, kwargs=kwargs)
+ pool_indices = get_pool_indices(feature_lens, kwargs=kwargs)
+ cu_seqlens = get_audio_cu_seqlens(chunk_lengths, kwargs=kwargs)
+
+ # Derive masks from chunk_lengths (traceable arithmetic + arange broadcasting)
+ padded_feature = padded_feature.to(self.conv1.weight.dtype)
+ padded_mask = (
+ (torch.arange(padded_feature.shape[2], device=padded_feature.device) < chunk_lengths.unsqueeze(1))
+ .unsqueeze(1)
+ .long()
+ )
+ padded_embed = nn.functional.gelu(self.conv1(padded_feature)) * padded_mask
+ padded_embed = nn.functional.gelu(self.conv2(padded_embed)).transpose(1, 2)
+ padded_embed = padded_embed + self.positional_embedding.positional_embedding[
+ : padded_embed.shape[1], :
+ ].unsqueeze(0).to(padded_embed.dtype)
+ hidden_states = torch.index_select(padded_embed.reshape(-1, padded_embed.shape[-1]), 0, valid_indices)
+
+ for encoder_layer in self.layers:
+ layer_outputs = encoder_layer(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ **kwargs,
+ )
+ hidden_states = layer_outputs[0]
+
+ # Post-process: stride-2 average pooling using precomputed indices, then project
+ hidden_states = (hidden_states[pool_indices] + hidden_states[pool_indices + 1]) / 2
+ hidden_states = self.proj(self.ln_post(hidden_states))
+ return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
+
+ # Ignore copy
+ def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
+ """
+ Computes the output length of the convolutional layers and the output length of the audio encoder
+ """
+ input_lengths = (input_lengths - 1) // 2 + 1
+ output_lengths = (input_lengths - 2) // 2 + 1
+ return input_lengths, output_lengths
+
+ def padded_and_mask_function(self, tensor_list, tensor_len, padding_value=0, padding_side="right"):
+ """
+ Pads a sequence of tensors to their maximum length on indicated `padding_side`.
+ Then prepares a mask so that pad tokens are not attended to.
+ """
+ warnings.warn(
+ f"`{self.__class__.__name__}.padded_and_mask_function` is deprecated and will be removed in v5.11. Use `chunk_and_pad_features` and `get_audio_cu_seqlens` helpers instead.",
+ FutureWarning,
+ stacklevel=2,
+ )
+ max_len = tensor_len.max()
+ dim = tensor_list[0].shape[0]
+ padded_tensor = torch.full(
+ size=(len(tensor_list), dim, max_len),
+ fill_value=padding_value,
+ dtype=self.dtype,
+ device=tensor_list[0].device,
+ )
+
+ batch_mask = torch.zeros(
+ (len(tensor_len), max_len),
+ dtype=torch.long,
+ device=padded_tensor.device,
+ )
+ for i, length in enumerate(tensor_len):
+ batch_mask[i, :length] = 1
+ padded_tensor[i, :, :length] = tensor_list[i]
+
+ feature_lens_after_cnn = (tensor_len - 1) // 2 + 1
+ max_len_after_cnn = feature_lens_after_cnn.max()
+ batch_mask_after_cnn = torch.zeros(
+ (len(tensor_len), max_len_after_cnn),
+ dtype=torch.long,
+ device=padded_tensor.device,
+ )
+ for i, length in enumerate(feature_lens_after_cnn):
+ batch_mask_after_cnn[i, :length] = 1
+ return (
+ padded_tensor,
+ batch_mask.unsqueeze(1),
+ batch_mask_after_cnn.bool(),
+ )
+
+
+def apply_rotary_pos_emb_vision(tensor: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+ orig_dtype = tensor.dtype
+ tensor = tensor.float()
+ cos = freqs.cos()
+ sin = freqs.sin()
+ cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
+ sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
+ output = (tensor * cos) + (rotate_half(tensor) * sin)
+ output = output.to(orig_dtype)
+ return output
+
+
+class Qwen2_5OmniVisionAttention(nn.Module):
+ def __init__(self, config: Qwen2_5OmniVisionEncoderConfig = None) -> None:
+ super().__init__()
+ self.dim = config.hidden_size
+ self.num_heads = config.num_heads
+ self.head_dim = self.dim // self.num_heads
+ self.q = nn.Linear(self.dim, self.dim, bias=True)
+ self.k = nn.Linear(self.dim, self.dim, bias=True)
+ self.v = nn.Linear(self.dim, self.dim, bias=True)
+ self.proj = nn.Linear(self.dim, self.dim)
+ self.scaling = self.head_dim**-0.5
+ self.num_key_value_groups = 1 # needed for eager attention
+ self.config = config
+ self.attention_dropout = 0.0
+ self.is_causal = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: torch.Tensor | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ query_states = self.q(hidden_states).reshape(seq_length, self.num_heads, -1)
+ key_states = self.k(hidden_states).reshape(seq_length, self.num_heads, -1)
+ value_states = self.v(hidden_states).reshape(seq_length, self.num_heads, -1)
+ query_states = apply_rotary_pos_emb_vision(query_states.unsqueeze(0), rotary_pos_emb).squeeze(0)
+ key_states = apply_rotary_pos_emb_vision(key_states.unsqueeze(0), rotary_pos_emb).squeeze(0)
+
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ if is_flash_attention_requested(self.config):
+ # Flash Attention 2: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class Qwen2_5OmniVisionBlock(Qwen2_5_VLVisionBlock):
+ def __init__(self, config: Qwen2_5OmniVisionEncoderConfig) -> None:
+ super().__init__(config, config._attn_implementation)
+ self.attn = Qwen2_5OmniVisionAttention(config=config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: torch.Tensor | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ hidden_states = hidden_states + self.attn(
+ self.norm1(hidden_states),
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ **kwargs,
+ )
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+ return hidden_states
+
+
+class Qwen2_5_VisionRotaryEmbedding(Qwen2_5_VisionRotaryEmbedding):
+ pass
+
+
+class Qwen2_5OmniVisionEncoder(Qwen2_5_VisionTransformerPretrainedModel):
+ config: Qwen2_5OmniVisionEncoderConfig
+ input_modalities = ("image", "video")
+ _no_split_modules = ["Qwen2_5OmniVisionBlock"]
+ _input_embed_layer = "patch_embed"
+ _can_record_outputs = {
+ "hidden_states": Qwen2_5OmniVisionBlock,
+ "attentions": Qwen2_5OmniVisionAttention,
+ }
+
+ def __init__(self, config: Qwen2_5OmniVisionEncoderConfig, *inputs, **kwargs) -> None:
+ super().__init__(config, *inputs, **kwargs)
+ self.blocks = nn.ModuleList([Qwen2_5OmniVisionBlock(config) for _ in range(config.depth)])
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> tuple | BaseModelOutputWithPooling:
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+ The final hidden states of the model.
+ grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+ The temporal, height and width of feature shape of each image in LLM.
+
+ Returns:
+ `torch.Tensor`: hidden_states.
+ """
+ position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs)
+ cu_seqlens = get_vision_cu_seqlens(grid_thw, kwargs=kwargs)
+ window_index, cu_window_seqlens = get_vision_window_index(
+ grid_thw,
+ spatial_merge_size=self.spatial_merge_size,
+ window_size=self.window_size,
+ patch_size=self.patch_size,
+ kwargs=kwargs,
+ )
+
+ hidden_states = self.patch_embed(hidden_states)
+
+ seq_len, _ = hidden_states.size()
+ reverse_indices = torch.argsort(window_index)
+ hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+ hidden_states = hidden_states[window_index, :, :]
+ hidden_states = hidden_states.reshape(seq_len, -1)
+
+ rotary_pos_emb = self.rotary_pos_emb(position_ids)
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+ rotary_pos_emb = rotary_pos_emb[window_index, :, :]
+ rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+
+ # Modification here
+ for layer_num, blk in enumerate(self.blocks):
+ if layer_num in self.fullatt_block_indexes:
+ cu_seqlens_now = cu_seqlens
+ else:
+ cu_seqlens_now = cu_window_seqlens
+
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens_now,
+ rotary_pos_emb=rotary_pos_emb,
+ **kwargs,
+ )
+
+ merged_hidden_states = self.merger(hidden_states)
+ merged_hidden_states = merged_hidden_states[reverse_indices, :]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=hidden_states,
+ pooler_output=merged_hidden_states,
+ )
+
+
+class Qwen2_5OmniRotaryEmbedding(Qwen2VLRotaryEmbedding):
+ def __init__(self, config: Qwen2_5OmniThinkerConfig, device=None):
+ super().__init__(config, device)
+
+
+# It's same as `Qwen2_5_VLAttention`, but talker model's hidden_size isn't divisible by num_heads.
+# Removes the value error as a workaround.
+class Qwen2_5OmniAttention(Qwen2_5_VLAttention):
+ def __init__(self, config: Qwen2_5OmniConfig, layer_idx: int | None = None):
+ nn.Module.__init__(self)
+ self.config = config
+ self.layer_idx = layer_idx
+ if layer_idx is None:
+ logger.warning_once(
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ self.hidden_size = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+ self.num_key_value_heads = config.num_key_value_heads
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+ self.is_causal = True
+ self.attention_dropout = config.attention_dropout
+ self.rope_parameters = config.rope_parameters
+ self.scaling = self.head_dim**-0.5
+
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+ self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None
+ self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None
+
+
+class Qwen2MLP(Qwen2_5_VLMLP):
+ pass
+
+
+class Qwen2_5OmniThinkerTextModel(Qwen2_5_VLTextModel):
+ config: Qwen2_5OmniTextConfig
+ _no_split_modules = ["Qwen2_5OmniDecoderLayer"]
+
+ def __init__(self, config: Qwen2_5OmniTextConfig):
+ super().__init__(config)
+
+
+@auto_docstring(
+ custom_intro="""
+ The Qwen2.5OmniThinker model which consists of a audio backbone and a language model.
+ """
+)
+class Qwen2_5OmniThinkerForConditionalGeneration(Qwen2_5OmniPreTrainedModelForConditionalGeneration, GenerationMixin):
+ config: Qwen2_5OmniThinkerConfig
+ base_model_prefix = "thinker"
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _no_split_modules = ["Qwen2_5OmniAudioEncoder", "Qwen2_5OmniVisionEncoder"]
+
+ def __init__(self, config: Qwen2_5OmniThinkerConfig):
+ super().__init__(config)
+ self.audio_tower = Qwen2_5OmniAudioEncoder._from_config(config.audio_config)
+ self.visual = Qwen2_5OmniVisionEncoder._from_config(config.vision_config)
+ self.vocab_size = config.text_config.vocab_size
+ self.model = Qwen2_5OmniThinkerTextModel._from_config(config.text_config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+ self.spatial_merge_size = config.vision_config.spatial_merge_size
+ self.rope_deltas = None
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ @accepts_precomputed_kwargs(modality="video")
+ @can_return_tuple
+ @auto_docstring
+ def get_video_features(
+ self,
+ pixel_values_videos: torch.FloatTensor,
+ video_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input videos.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ """
+ pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
+ return self.visual(pixel_values_videos, grid_thw=video_grid_thw, **kwargs)
+
+ @accepts_precomputed_kwargs(modality="image")
+ @can_return_tuple
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ pixel_values = pixel_values.type(self.visual.dtype)
+ return self.visual(pixel_values, grid_thw=image_grid_thw, **kwargs)
+
+ @can_return_tuple
+ @auto_docstring
+ def get_audio_features(
+ self,
+ input_features: torch.FloatTensor,
+ feature_attention_mask: torch.LongTensor | None = None,
+ audio_feature_lengths: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ input_features (`torch.FloatTensor`):
+ The tensors corresponding to the input audios.
+ feature_attention_mask (`torch.LongTensor`, *optional*):
+ Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`:
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ """
+ if feature_attention_mask is not None:
+ audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
+ input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
+ else:
+ audio_feature_lengths = None
+
+ audio_feat_lengths, audio_output_lengths = self.audio_tower._get_feat_extract_output_lengths(
+ audio_feature_lengths if audio_feature_lengths is not None else feature_attention_mask.sum(-1)
+ )
+ feature_lens = audio_feature_lengths if audio_feature_lengths is not None else feature_attention_mask.sum(-1)
+ audio_outputs = self.audio_tower(
+ input_features, feature_lens=feature_lens, aftercnn_lens=audio_feat_lengths, **kwargs
+ )
+ if audio_outputs.last_hidden_state.shape[0] != sum(audio_output_lengths.tolist()):
+ raise ValueError("length of audio_features should match audio_output_lengths")
+
+ return audio_outputs
+
+ def get_placeholder_mask(
+ self,
+ input_ids: torch.LongTensor,
+ inputs_embeds: torch.FloatTensor,
+ image_features: torch.FloatTensor | None = None,
+ video_features: torch.FloatTensor | None = None,
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ special_video_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_video_mask = special_video_mask.all(-1)
+ special_audio_mask = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(self.config.audio_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ ).all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+ special_video_mask = input_ids == self.config.video_token_id
+ special_audio_mask = input_ids == self.config.audio_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if image_features is not None:
+ torch_compilable_check(
+ inputs_embeds[special_image_mask].numel() == image_features.numel(),
+ f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}",
+ )
+
+ n_video_tokens = special_video_mask.sum()
+ special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if video_features is not None:
+ torch_compilable_check(
+ inputs_embeds[special_video_mask].numel() == video_features.numel(),
+ f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}",
+ )
+
+ special_audio_mask = special_audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ return special_image_mask, special_video_mask, special_audio_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ input_features: torch.FloatTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ pixel_values_videos: torch.FloatTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ feature_attention_mask: torch.Tensor | None = None,
+ audio_feature_lengths: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ rope_deltas: torch.LongTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ use_audio_in_video: bool | None = None,
+ video_second_per_grid: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | Qwen2_5OmniThinkerCausalLMOutputWithPast:
+ r"""
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ feature_attention_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+ use_audio_in_video (`bool`, *optional*):
+ Whether or not use audio track in video, should same as the parameter in `process_audio_info`.
+ video_second_per_grid (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ Number of seconds per grid for each video, used for temporal feature mapping.
+
+ Example:
+
+ ```python
+ >>> from io import BytesIO
+ >>> from urllib.request import urlopen
+ >>> import librosa
+ >>> from qwen_vl_utils import process_vision_info
+ >>> from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration
+
+ >>> thinker = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B")
+ >>> processor = Qwen2_5OmniProcessor.from_pretrained("Qwen/Qwen2.5-Omni-7B")
+
+ >>> conversations = [
+ >>> {'role': 'system', 'content': 'You are a helpful voice chat bot, and please respond to me in a casual conversation manner using random voice.'},
+ >>> {"role": "user", "content": [
+ >>> {"type": "image", "image_url": "https://www.ilankelman.org/stopsigns/australia.jpg"},
+ >>> {"type": "audio", "audio_url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/glass-breaking-151256.mp3"},
+ >>> ]},
+ >>> ]
+
+ >>> text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
+ >>> audios = [ librosa.load(BytesIO(urlopen( conversations[1]['content'][1]['audio_url'] ).read()), sr=self.processor.feature_extractor.sampling_rate) ]
+ >>> images, videos = process_vision_info(conversations)
+ >>> inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors="pt", padding=True)
+
+ >>> # Generate
+ >>> inputs['use_audio_in_video'] = `True` or `False`
+ >>> generation = thinker.generate(**inputs, max_new_tokens=2048)
+ >>> generate_ids = generation[:, inputs.input_ids.size(1):]
+
+ >>> response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ ```"""
+
+ if inputs_embeds is None:
+ # 1. Extract the input embeddings
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ # 2. Merge text , audios , image and video
+ if input_features is not None:
+ audio_features = self.get_audio_features(
+ input_features, feature_attention_mask, audio_feature_lengths, return_dict=True, **kwargs
+ ).last_hidden_state
+ audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
+ _, _, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_features)
+
+ if pixel_values is not None:
+ image_embeds = self.get_image_features(
+ pixel_values, image_grid_thw, return_dict=True, **kwargs
+ ).pooler_output
+ image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ image_mask, _, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ if pixel_values_videos is not None:
+ video_embeds = self.get_video_features(
+ pixel_values_videos, video_grid_thw, return_dict=True, **kwargs
+ ).pooler_output
+ video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ _, video_mask, _ = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ if feature_attention_mask is not None:
+ audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
+ else:
+ audio_feature_lengths = None
+
+ if attention_mask is not None and position_ids is None:
+ past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+ if past_key_values_length == 0 or self.rope_deltas is None:
+ delta0 = (1 - attention_mask).sum(dim=-1).unsqueeze(1)
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask,
+ use_audio_in_video,
+ audio_feature_lengths,
+ video_second_per_grid,
+ )
+ rope_deltas = rope_deltas - delta0
+ self.rope_deltas = rope_deltas
+ else:
+ batch_size, seq_length = input_ids.shape
+ delta = (past_key_values_length + self.rope_deltas).to(input_ids.device)
+ position_ids = torch.arange(seq_length, device=input_ids.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ outputs = self.model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ return_dict=True,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ logits = self.lm_head(hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits, labels=labels, vocab_size=self.config.get_text_config().vocab_size
+ )
+
+ return Qwen2_5OmniThinkerCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=self.rope_deltas,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ input_features=None,
+ feature_attention_mask=None,
+ use_audio_in_video=False,
+ video_second_per_grid=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ use_cache=use_cache,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ input_features=input_features,
+ feature_attention_mask=feature_attention_mask,
+ use_audio_in_video=use_audio_in_video,
+ video_second_per_grid=video_second_per_grid,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ model_inputs["position_ids"] = None
+
+ if not is_first_iteration and use_cache:
+ model_inputs["pixel_values"] = None
+ model_inputs["pixel_values_videos"] = None
+ model_inputs["input_features"] = None
+
+ return model_inputs
+
+
+############################
+# Start Talker #
+############################
+
+
+@auto_docstring(
+ custom_intro="""
+ Base class for Qwen2.5OmniTalker causal language model (or autoregressive) outputs.
+ """
+)
+@dataclass
+class Qwen2_5OmniTalkerCausalLMOutputWithPast(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+ `past_key_values` input) to speed up sequential decoding.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ thinker_reply_part (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ Hidden states from the thinker model that are used as input for the talker model. These represent the encoded
+ response that the talker model will use to generate speech tokens.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ rope_deltas: torch.LongTensor | None = None
+ thinker_reply_part: torch.FloatTensor | None = None
+
+
+class Qwen2_5OmniTalkerModel(Qwen2_5_VLTextModel):
+ config: Qwen2_5OmniTalkerConfig
+ input_modalities = ("image", "video", "audio", "text")
+
+ _no_split_modules = ["Qwen2_5OmniTalkerDecoderLayer"]
+
+ def __init__(self, config: Qwen2_5OmniTalkerConfig):
+ super().__init__(config)
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.embedding_size, self.padding_idx)
+
+
+class Qwen2_5OmniTalkerForConditionalGeneration(Qwen2_5OmniPreTrainedModelForConditionalGeneration, GenerationMixin):
+ config: Qwen2_5OmniTalkerConfig
+ base_model_prefix = "talker"
+ output_modalities = ("audio",)
+
+ def __init__(self, config: Qwen2_5OmniTalkerConfig):
+ super().__init__(config)
+
+ self.thinker_to_talker_proj = nn.Linear(config.embedding_size, config.hidden_size)
+
+ self.model = Qwen2_5OmniTalkerModel(config)
+ self.codebook_size = config.vocab_size
+ self.codec_head = nn.Linear(config.hidden_size, self.codebook_size, bias=False)
+
+ self.codec_bos_token = config.tts_codec_start_token_id
+ self.codec_eos_token = config.tts_codec_end_token_id
+ self.codec_pad_token = config.tts_codec_pad_token_id
+ self.codec_mask_token = config.tts_codec_mask_token_id
+
+ self.text_bos_token = config.tts_text_start_token_id
+ self.text_eos_token = config.tts_text_end_token_id
+ self.text_pad_token = config.tts_text_pad_token_id
+
+ self.spatial_merge_size = self.config.spatial_merge_size
+ self.rope_deltas = None
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ thinker_reply_part: torch.FloatTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ rope_deltas: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ input_text_ids: torch.LongTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ use_audio_in_video: bool | None = None,
+ audio_feature_lengths: torch.LongTensor | None = None,
+ video_second_per_grid: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> tuple | Qwen2_5OmniTalkerCausalLMOutputWithPast:
+ r"""
+ thinker_reply_part (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Hidden states from the thinker model's output that represent the text reply part to be processed.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ input_text_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Input token IDs for text-only content, used for position calculation in multimodal contexts.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ use_audio_in_video (`bool`, *optional*):
+ Whether or not use audio track in video, should same as the parameter in `process_audio_info`.
+ audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
+ The length of feature shape of each audio in LLM.
+ video_second_per_grid (`torch.LongTensor` of shape `(num_videos)`, *optional*):
+ Number of seconds per grid for each video, used for temporal feature mapping.
+
+ Example:
+
+ ```python
+ >>> from io import BytesIO
+ >>> from urllib.request import urlopen
+ >>> import librosa
+ >>> from transformers import AutoProcessor, Qwen2_5OmniTalkerForConditionalGeneration
+
+ >>> model = Qwen2_5OmniTalkerForConditionalGeneration.from_pretrained("Qwen/Qwen2-Audio-7B")
+ >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2-Audio-7B")
+
+ >>> prompt = "<|audio_bos|><|AUDIO|><|audio_eos|>Generate the caption in English:"
+ >>> url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/glass-breaking-151256.mp3"
+ >>> audio, _ = librosa.load(BytesIO(urlopen(url).read()), sr=self.processor.feature_extractor.sampling_rate)
+
+ >>> inputs = processor(text=prompt, audio=audio, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(**inputs, max_length=30)
+ >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Generate the caption in English: Glass is breaking."
+ ```"""
+
+ if attention_mask is not None and position_ids is None:
+ past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+ if past_key_values_length == 0 or self.rope_deltas is None:
+ position_ids, rope_deltas = self.get_rope_index(
+ input_text_ids,
+ image_grid_thw,
+ video_grid_thw,
+ attention_mask,
+ use_audio_in_video,
+ audio_feature_lengths,
+ video_second_per_grid,
+ )
+
+ if inputs_embeds is not None:
+ inputs_embeds[:, -1, :] += self.get_input_embeddings()(
+ torch.tensor([self.codec_bos_token], dtype=torch.long, device=inputs_embeds.device)
+ )
+ inputs_embeds[:, -2, :] += self.get_input_embeddings()(
+ torch.tensor([self.codec_pad_token], dtype=torch.long, device=inputs_embeds.device)
+ )
+ self.rope_deltas = rope_deltas
+
+ else:
+ if inputs_embeds is not None:
+ batch_size, seq_length, _ = inputs_embeds.shape
+ else:
+ batch_size, seq_length = input_ids.shape
+
+ delta = (past_key_values_length + self.rope_deltas).to(input_ids.device)
+ position_ids = torch.arange(seq_length, device=input_ids.device)
+ position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+ position_ids = position_ids.add(delta)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ if inputs_embeds is None:
+ # 1. Inference tokens after second token
+ codec_embeds = self.get_input_embeddings()(input_ids)
+ inputs_embeds = codec_embeds + thinker_reply_part[:, :1, :]
+ if thinker_reply_part.shape[1] > 1:
+ thinker_reply_part = thinker_reply_part[:, 1:, :]
+
+ talker_lm_input = self.thinker_to_talker_proj(inputs_embeds)
+
+ if attention_mask is not None:
+ attention_mask = attention_mask.to(inputs_embeds.device)
+
+ outputs = self.model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=talker_lm_input,
+ use_cache=use_cache,
+ return_dict=True,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ logits = self.codec_head(hidden_states)
+ logits = logits.float()
+
+ loss = None
+
+ return Qwen2_5OmniTalkerCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=self.rope_deltas,
+ thinker_reply_part=thinker_reply_part,
+ )
+
+ # prepare inputs for talker lm generation
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ input_text_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ thinker_reply_part=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ input_audio_features=None,
+ audio_feature_attention_mask=None,
+ audio_feature_lengths=None,
+ use_audio_in_video=False,
+ video_second_per_grid=None,
+ **kwargs,
+ ):
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ thinker_reply_part=thinker_reply_part,
+ input_text_ids=input_text_ids,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ use_audio_in_video=use_audio_in_video,
+ audio_feature_lengths=audio_feature_lengths,
+ video_second_per_grid=video_second_per_grid,
+ **kwargs,
+ )
+
+ model_inputs["position_ids"] = None
+
+ return model_inputs
+
+ def _update_model_kwargs_for_generation(
+ self,
+ outputs: ModelOutput,
+ model_kwargs: dict[str, Any],
+ is_encoder_decoder: bool = False,
+ num_new_tokens: int = 1,
+ ) -> dict[str, Any]:
+ model_kwargs = super()._update_model_kwargs_for_generation(
+ outputs, model_kwargs, is_encoder_decoder, num_new_tokens
+ )
+
+ if getattr(outputs, "thinker_reply_part", None) is not None:
+ model_kwargs["thinker_reply_part"] = outputs.thinker_reply_part
+
+ return model_kwargs
+
+
+############################
+# Start Token2Wav #
+############################
+
+
+class Qwen2_5OmniDiTRotaryEmbedding(LlamaRotaryEmbedding):
+ def __init__(self, config: Qwen2_5OmniDiTConfig, device=None):
+ super().__init__(config, device=device)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: Qwen2_5OmniDiTConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ return super().compute_default_rope_parameters(
+ config,
+ device=device,
+ seq_len=seq_len,
+ )
+
+
+# Modified from Llama with a different rotate function, will fixed in next release
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+
+ def rotate_half_codec(x):
+ # x = rearrange(x, "... (d r) -> ... d r", r=2)
+ x = x.reshape(*x.shape[:-1], -1, 2)
+ x1, x2 = x.unbind(dim=-1)
+ x = torch.stack((-x2, x1), dim=-1)
+ return x.reshape(*x.shape[:-2], -1)
+
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half_codec(q) * sin)
+ k_embed = (k * cos) + (rotate_half_codec(k) * sin)
+ return q_embed, k_embed
+
+
+class TimeDelayNetBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ kernel_size,
+ dilation,
+ ):
+ super().__init__()
+ self.conv = nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ dilation=dilation,
+ padding="same",
+ padding_mode="reflect",
+ )
+ self.activation = nn.ReLU()
+
+ def forward(self, hidden_states: torch.Tensor):
+ return self.activation(self.conv(hidden_states))
+
+
+class Res2NetBlock(torch.nn.Module):
+ def __init__(self, in_channels, out_channels, scale=8, kernel_size=3, dilation=1):
+ super().__init__()
+
+ in_channel = in_channels // scale
+ hidden_channel = out_channels // scale
+
+ self.blocks = nn.ModuleList(
+ [
+ TimeDelayNetBlock(
+ in_channel,
+ hidden_channel,
+ kernel_size=kernel_size,
+ dilation=dilation,
+ )
+ for i in range(scale - 1)
+ ]
+ )
+ self.scale = scale
+
+ def forward(self, hidden_states):
+ outputs = []
+ for i, hidden_part in enumerate(torch.chunk(hidden_states, self.scale, dim=1)):
+ if i == 0:
+ output_part = hidden_part
+ elif i == 1:
+ output_part = self.blocks[i - 1](hidden_part)
+ else:
+ output_part = self.blocks[i - 1](hidden_part + output_part)
+ outputs.append(output_part)
+ output = torch.cat(outputs, dim=1)
+ return output
+
+
+class SqueezeExcitationBlock(nn.Module):
+ def __init__(self, in_channels, se_channels, out_channels):
+ super().__init__()
+
+ self.conv1 = nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=se_channels,
+ kernel_size=1,
+ padding="same",
+ padding_mode="reflect",
+ )
+ self.relu = nn.ReLU(inplace=True)
+ self.conv2 = nn.Conv1d(
+ in_channels=se_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ padding="same",
+ padding_mode="reflect",
+ )
+ self.sigmoid = nn.Sigmoid()
+
+ def forward(self, hidden_states):
+ hidden_states_mean = hidden_states.mean(dim=2, keepdim=True)
+
+ hidden_states_mean = self.relu(self.conv1(hidden_states_mean))
+ hidden_states_mean = self.sigmoid(self.conv2(hidden_states_mean))
+
+ return hidden_states * hidden_states_mean
+
+
+class AttentiveStatisticsPooling(nn.Module):
+ """This class implements an attentive statistic pooling layer for each channel.
+ It returns the concatenated mean and std of the input tensor.
+ """
+
+ def __init__(self, channels, attention_channels=128):
+ super().__init__()
+
+ self.eps = 1e-12
+ self.tdnn = TimeDelayNetBlock(channels * 3, attention_channels, 1, 1)
+ self.tanh = nn.Tanh()
+ self.conv = nn.Conv1d(
+ in_channels=attention_channels,
+ out_channels=channels,
+ kernel_size=1,
+ padding="same",
+ padding_mode="reflect",
+ )
+
+ def _length_to_mask(self, length, max_len=None, dtype=None, device=None):
+ """Creates a binary mask for each sequence.
+
+ Reference: https://discuss.pytorch.org/t/how-to-generate-variable-length-mask/23397/3
+
+ Arguments
+ ---------
+ length : torch.LongTensor
+ Containing the length of each sequence in the batch. Must be 1D.
+ max_len : int
+ Max length for the mask, also the size of the second dimension.
+ dtype : torch.dtype, default: None
+ The dtype of the generated mask.
+ device: torch.device, default: None
+ The device to put the mask variable.
+
+ Returns
+ -------
+ mask : tensor
+ The binary mask.
+ """
+
+ if max_len is None:
+ max_len = length.max().long().item() # using arange to generate mask
+ mask = torch.arange(max_len, device=length.device, dtype=length.dtype).expand(
+ len(length), max_len
+ ) < length.unsqueeze(1)
+
+ mask = torch.as_tensor(mask, dtype=dtype, device=device)
+ return mask
+
+ def _compute_statistics(self, x, m, dim=2):
+ mean = (m * x).sum(dim)
+ std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps))
+ return mean, std
+
+ def forward(self, hidden_states):
+ seq_length = hidden_states.shape[-1]
+ lengths = torch.ones(hidden_states.shape[0], device=hidden_states.device)
+
+ # Make binary mask of shape [N, 1, L]
+ mask = self._length_to_mask(
+ lengths * seq_length, max_len=seq_length, dtype=hidden_states.dtype, device=hidden_states.device
+ )
+ mask = mask.unsqueeze(1)
+
+ # Expand the temporal context of the pooling layer by allowing the
+ # self-attention to look at global properties of the utterance.
+ total = mask.sum(dim=2, keepdim=True)
+
+ mean, std = self._compute_statistics(hidden_states, mask / total)
+ mean = mean.unsqueeze(2).repeat(1, 1, seq_length)
+ std = std.unsqueeze(2).repeat(1, 1, seq_length)
+ attention = torch.cat([hidden_states, mean, std], dim=1)
+
+ # Apply layers
+ attention = self.conv(self.tanh(self.tdnn(attention)))
+
+ # Filter out zero-paddings
+ attention = attention.masked_fill(mask == 0, float("-inf"))
+
+ attention = F.softmax(attention, dim=2)
+ mean, std = self._compute_statistics(hidden_states, attention)
+ # Append mean and std of the batch
+ pooled_stats = torch.cat((mean, std), dim=1)
+ pooled_stats = pooled_stats.unsqueeze(2)
+
+ return pooled_stats
+
+
+class SqueezeExcitationRes2NetBlock(nn.Module):
+ """An implementation of building block in ECAPA-TDNN, i.e.,
+ TDNN-Res2Net-TDNN-SqueezeExcitationBlock.
+ """
+
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ res2net_scale=8,
+ se_channels=128,
+ kernel_size=1,
+ dilation=1,
+ ):
+ super().__init__()
+ self.out_channels = out_channels
+ self.tdnn1 = TimeDelayNetBlock(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ dilation=1,
+ )
+ self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation)
+ self.tdnn2 = TimeDelayNetBlock(
+ out_channels,
+ out_channels,
+ kernel_size=1,
+ dilation=1,
+ )
+ self.se_block = SqueezeExcitationBlock(out_channels, se_channels, out_channels)
+
+ def forward(self, hidden_state):
+ residual = hidden_state
+
+ hidden_state = self.tdnn1(hidden_state)
+ hidden_state = self.res2net_block(hidden_state)
+ hidden_state = self.tdnn2(hidden_state)
+ hidden_state = self.se_block(hidden_state)
+
+ return hidden_state + residual
+
+
+class ECAPA_TimeDelayNet(torch.nn.Module):
+ """An implementation of the speaker embedding model in a paper.
+ "ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in
+ TDNN Based Speaker Verification" (https://huggingface.co/papers/2005.07143).
+ """
+
+ def __init__(self, config: Qwen2_5OmniDiTConfig):
+ super().__init__()
+ if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(config.enc_channels) != len(
+ config.enc_dilations
+ ):
+ raise ValueError("enc_channels, enc_kernel_sizes and enc_dilations should have same length")
+ self.channels = config.enc_channels
+ self.blocks = nn.ModuleList()
+
+ # The initial TDNN layer
+ self.blocks.append(
+ TimeDelayNetBlock(
+ config.mel_dim,
+ config.enc_channels[0],
+ config.enc_kernel_sizes[0],
+ config.enc_dilations[0],
+ )
+ )
+
+ # SE-Res2Net layers
+ for i in range(1, len(config.enc_channels) - 1):
+ self.blocks.append(
+ SqueezeExcitationRes2NetBlock(
+ config.enc_channels[i - 1],
+ config.enc_channels[i],
+ res2net_scale=config.enc_res2net_scale,
+ se_channels=config.enc_se_channels,
+ kernel_size=config.enc_kernel_sizes[i],
+ dilation=config.enc_dilations[i],
+ )
+ )
+
+ # Multi-layer feature aggregation
+ self.mfa = TimeDelayNetBlock(
+ config.enc_channels[-1],
+ config.enc_channels[-1],
+ config.enc_kernel_sizes[-1],
+ config.enc_dilations[-1],
+ )
+
+ # Attentive Statistical Pooling
+ self.asp = AttentiveStatisticsPooling(
+ config.enc_channels[-1],
+ attention_channels=config.enc_attention_channels,
+ )
+
+ # Final linear transformation
+ self.fc = nn.Conv1d(
+ in_channels=config.enc_channels[-1] * 2,
+ out_channels=config.enc_dim,
+ kernel_size=1,
+ padding="same",
+ padding_mode="reflect",
+ )
+
+ def forward(self, hidden_states):
+ # Minimize transpose for efficiency
+ hidden_states = hidden_states.transpose(1, 2)
+
+ hidden_states_list = []
+ for layer in self.blocks:
+ hidden_states = layer(hidden_states)
+ hidden_states_list.append(hidden_states)
+
+ # Multi-layer feature aggregation
+ hidden_states = torch.cat(hidden_states_list[1:], dim=1)
+ hidden_states = self.mfa(hidden_states)
+
+ # Attentive Statistical Pooling
+ hidden_states = self.asp(hidden_states)
+
+ # Final linear transformation
+ hidden_states = self.fc(hidden_states)
+
+ hidden_states = hidden_states.squeeze(-1)
+ return hidden_states
+
+
+class DiTInputEmbedding(nn.Module):
+ def __init__(self, config: Qwen2_5OmniDiTConfig):
+ super().__init__()
+ self.proj = nn.Linear(
+ config.mel_dim + config.enc_dim + config.enc_emb_dim + config.emb_dim,
+ config.hidden_size,
+ )
+ self.spk_encoder = ECAPA_TimeDelayNet(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ speaker_embedding: torch.Tensor,
+ condition_vector: torch.Tensor,
+ code_embed: torch.Tensor,
+ drop_audio_cond: bool | None = False,
+ code_embed_uncond: bool | None = None,
+ apply_cfg: bool | None = True,
+ ):
+ if apply_cfg:
+ hidden_states = torch.cat([hidden_states, hidden_states], dim=0)
+ speaker_embedding = torch.cat([speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0)
+ condition_vector = torch.cat([condition_vector, torch.zeros_like(condition_vector)], dim=0)
+ code_embed = torch.cat([code_embed, code_embed_uncond], dim=0)
+ elif drop_audio_cond: # cfg for cond audio
+ condition_vector = torch.zeros_like(condition_vector)
+ speaker_embedding = torch.zeros_like(speaker_embedding)
+ condition_vector = self.spk_encoder(condition_vector).unsqueeze(1).repeat(1, hidden_states.size(1), 1)
+ hidden_states = self.proj(torch.cat((hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1))
+
+ return hidden_states
+
+
+# Transformer backbone using DiT blocks
+class DiTCodecEmbedding(nn.Module):
+ def __init__(self, codec_num_embeds, codec_dim, repeats):
+ super().__init__()
+ self.repeats = repeats
+ self.codec_embed = nn.Embedding(codec_num_embeds + 1, codec_dim)
+
+ def forward(self, code, drop_code=False):
+ if drop_code:
+ code = torch.zeros_like(code)
+ code_embed = self.codec_embed(code)
+
+ code_embed = torch.repeat_interleave(code_embed, repeats=self.repeats, dim=1)
+ return code_embed
+
+
+# AdaLayerNormZero
+# return with modulated x for attn input, and params for later mlp modulation
+class Qwen2_5_OmniAdaLayerNormZero(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+
+ self.silu = nn.SiLU()
+ self.linear = nn.Linear(dim, dim * 6)
+
+ self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+ def forward(self, hidden_states, emb=None):
+ emb = self.linear(self.silu(emb))
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+
+ hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+ return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+# AdaLayerNormZero for final layer
+# return only with modulated x for attn input, cuz no more mlp modulation
+class Qwen2_5_OmniAdaLayerNormZero_Final(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+
+ self.silu = nn.SiLU()
+ self.linear = nn.Linear(dim, dim * 2)
+
+ self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+ def forward(self, hidden_states, emb):
+ emb = self.linear(self.silu(emb))
+ scale, shift = torch.chunk(emb, 2, dim=1)
+
+ hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
+ return hidden_states
+
+
+# FeedForward
+class DiTMLP(nn.Module):
+ def __init__(self, dim, mult=4, dropout=0.0):
+ super().__init__()
+ inner_dim = int(dim * mult)
+
+ self.ff = nn.ModuleList(
+ [
+ nn.Linear(dim, inner_dim),
+ nn.GELU(approximate="tanh"),
+ nn.Dropout(dropout),
+ nn.Linear(inner_dim, dim),
+ ]
+ )
+
+ def forward(self, hidden_states):
+ for layer in self.ff:
+ hidden_states = layer(hidden_states)
+ return hidden_states
+
+
+class DiTAttention(nn.Module):
+ def __init__(self, config: Qwen2_5OmniDiTConfig):
+ super().__init__()
+
+ self.config = config
+ self.dim = config.hidden_size
+ self.heads = config.num_attention_heads
+ self.inner_dim = config.head_dim * config.num_attention_heads
+ self.dropout = config.dropout
+ self.is_causal = False
+
+ self.to_q = nn.Linear(config.hidden_size, self.inner_dim)
+ self.to_k = nn.Linear(config.hidden_size, self.inner_dim)
+ self.to_v = nn.Linear(config.hidden_size, self.inner_dim)
+
+ self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)])
+
+ def forward(
+ self,
+ hidden_states, # noised input x
+ position_embeddings=None, # rotary position embedding for x
+ attention_mask=None,
+ ) -> torch.Tensor:
+ batch_size = hidden_states.shape[0]
+
+ # `sample` projections.
+ query = self.to_q(hidden_states)
+ key = self.to_k(hidden_states)
+ value = self.to_v(hidden_states)
+
+ # attention
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // self.heads
+ query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+
+ # apply rotary position embedding
+ # Due to training process, only first head is applied with RoPE, will be fixed at next release
+ cos, sin = position_embeddings
+ query[:, :1], key[:, :1] = apply_rotary_pos_emb(query[:, :1], key[:, :1], cos, sin)
+
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+ attention_weights, _ = attention_interface(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=attention_mask,
+ is_causal=False,
+ )
+
+ # mask. e.g. inference got a batch with different target durations, mask out the padding
+ attention_weights = attention_weights.reshape(batch_size, -1, self.heads * head_dim)
+ attention_weights = attention_weights.to(query.dtype)
+
+ # linear proj
+ attention_output = self.to_out[0](attention_weights)
+ attention_output = self.to_out[1](attention_output)
+
+ return attention_output
+
+
+# time step conditioning embedding
+class SinusPositionEmbedding(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ self.dim = dim
+
+ def forward(self, hidden_states, scale=1000):
+ device = hidden_states.device
+ half_dim = self.dim // 2
+ emb = math.log(10000) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+ emb = scale * hidden_states.unsqueeze(1) * emb.unsqueeze(0)
+ emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+ return emb.type_as(hidden_states)
+
+
+class DiTTimestepEmbedding(nn.Module):
+ def __init__(self, dim, freq_embed_dim=256):
+ super().__init__()
+ self.time_embed = SinusPositionEmbedding(freq_embed_dim)
+ self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)])
+
+ def forward(self, timestep):
+ time_hidden = self.time_embed(timestep)
+ time_hidden = time_hidden.to(timestep.dtype)
+ for layer in self.time_mlp:
+ time_hidden = layer(time_hidden) # b d
+ return time_hidden
+
+
+class DiTDecoderLayer(nn.Module):
+ def __init__(self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0):
+ super().__init__()
+ self.attn_norm = Qwen2_5_OmniAdaLayerNormZero(config.hidden_size)
+
+ self.attn = DiTAttention(config)
+ self.look_ahead_block = look_ahead_block
+ self.look_backward_block = look_backward_block
+ self.ff_norm = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=1e-6)
+ self.ff = DiTMLP(dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout)
+
+ def forward(
+ self, hidden_states, timestep, position_embeddings=None, block_diff=None
+ ): # x: noised input, t: time embedding
+ # pre-norm & modulation for attention input
+ norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(hidden_states, emb=timestep)
+
+ # attention
+ attn_output = self.attn(
+ hidden_states=norm,
+ position_embeddings=position_embeddings,
+ attention_mask=(block_diff >= -float(self.look_backward_block))
+ & (block_diff <= float(self.look_ahead_block)),
+ )
+
+ # process attention output for input x
+ hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_output
+
+ norm = self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+ ff_output = self.ff(norm)
+ hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
+
+ return hidden_states
+
+
+class SnakeBeta(nn.Module):
+ """
+ A modified Snake function which uses separate parameters for the magnitude of the periodic components
+ Shape:
+ - Input: (B, C, T)
+ - Output: (B, C, T), same shape as the input
+ Parameters:
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ References:
+ - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+ https://huggingface.co/papers/2006.08195
+ """
+
+ def __init__(self, in_features, alpha=1.0):
+ super().__init__()
+ self.in_features = in_features
+
+ # initialize alpha
+ self.alpha = Parameter(torch.zeros(in_features) * alpha)
+ self.beta = Parameter(torch.zeros(in_features) * alpha)
+
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, hidden_states):
+ """
+ Forward pass of the function.
+ Applies the function to the input elementwise.
+ SnakeBeta ∶= x + 1/b * sin^2 (xa)
+ """
+ alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+ beta = self.beta.unsqueeze(0).unsqueeze(-1)
+ alpha = torch.exp(alpha)
+ beta = torch.exp(beta)
+ hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
+ torch.sin(hidden_states * alpha), 2
+ )
+
+ return hidden_states
+
+
+def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
+ """Generates a 1D Kaiser-windowed sinc filter.
+
+ Args:
+ cutoff (float): Normalized cutoff frequency (0 to 0.5).
+ half_width (float): Transition bandwidth.
+ kernel_size (int): Number of filter taps.
+
+ Returns:
+ torch.Tensor: A tensor of shape (1, 1, kernel_size) representing the filter.
+ """
+ is_even = kernel_size % 2 == 0
+ half_size = kernel_size // 2
+
+ # Compute Kaiser window parameters
+ delta_f = 4 * half_width
+ attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+
+ if attenuation > 50.0:
+ beta = 0.1102 * (attenuation - 8.7)
+ elif attenuation >= 21.0:
+ beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0)
+ else:
+ beta = 0.0
+
+ kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
+
+ # Compute time indices
+ if is_even:
+ time_indices = torch.arange(-half_size, half_size) + 0.5
+ else:
+ time_indices = torch.arange(kernel_size) - half_size
+
+ # Compute sinc filter
+ if cutoff == 0:
+ return torch.zeros((1, 1, kernel_size), dtype=torch.float32) # Ensures correct shape
+
+ sinc_filter = torch.sinc(2 * cutoff * time_indices)
+ normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
+
+ # Normalize to ensure sum = 1 (avoid leakage of constant component)
+ normalized_filter /= normalized_filter.sum()
+
+ return normalized_filter.view(1, 1, kernel_size)
+
+
+class UpSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ self.stride = ratio
+ self.pad = self.kernel_size // ratio - 1
+ self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+ self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+
+ filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
+ self.register_buffer("filter", filter, persistent=False)
+
+ def forward(self, hidden_states):
+ channels = hidden_states.shape[1]
+
+ hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate")
+ hidden_states = self.ratio * F.conv_transpose1d(
+ hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels
+ )
+ hidden_states = hidden_states[..., self.pad_left : -self.pad_right]
+
+ return hidden_states
+
+
+class DownSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ cutoff = 0.5 / ratio
+ half_width = 0.6 / ratio
+ self.cutoff = cutoff
+ self.half_width = half_width
+ self.kernel_size = kernel_size
+
+ if cutoff < 0.0:
+ raise ValueError("Minimum cutoff must be larger than zero.")
+ if cutoff > 0.5:
+ raise ValueError("A cutoff above 0.5 does not make sense.")
+
+ self.even = kernel_size % 2 == 0
+ self.pad_left = kernel_size // 2 - int(self.even)
+ self.pad_right = kernel_size // 2
+ self.stride = ratio
+ filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+ self.register_buffer("filter", filter, persistent=False)
+
+ def forward(self, hidden_states):
+ channels = hidden_states.shape[1]
+ hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate")
+ out = F.conv1d(hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels)
+ return out
+
+
+class TorchActivation1d(nn.Module):
+ def __init__(
+ self,
+ activation,
+ up_ratio: int = 2,
+ down_ratio: int = 2,
+ up_kernel_size: int = 12,
+ down_kernel_size: int = 12,
+ ):
+ super().__init__()
+ if not callable(activation):
+ raise TypeError("Activation function must be callable")
+ self.act = activation
+ self.upsample = UpSample1d(up_ratio, up_kernel_size)
+ self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+ def forward(self, hidden_states):
+ hidden_states = self.upsample(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.downsample(hidden_states)
+
+ return hidden_states
+
+
+class AMPBlock(torch.nn.Module):
+ def __init__(
+ self,
+ channels,
+ kernel_size=3,
+ dilation=(1, 3, 5),
+ ):
+ super().__init__()
+
+ self.convs1 = nn.ModuleList(
+ [
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[0],
+ padding=self._get_padding(kernel_size, dilation[0]),
+ ),
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[1],
+ padding=self._get_padding(kernel_size, dilation[1]),
+ ),
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[2],
+ padding=self._get_padding(kernel_size, dilation[2]),
+ ),
+ ]
+ )
+
+ self.convs2 = nn.ModuleList(
+ [
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding=self._get_padding(kernel_size, 1),
+ ),
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding=self._get_padding(kernel_size, 1),
+ ),
+ nn.Conv1d(
+ channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding=self._get_padding(kernel_size, 1),
+ ),
+ ]
+ )
+
+ self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers
+
+ self.activations = nn.ModuleList(
+ [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
+ )
+
+ def _get_padding(self, kernel_size, dilation=1):
+ return int((kernel_size * dilation - dilation) / 2)
+
+ def forward(self, hidden_states):
+ acts1, acts2 = self.activations[::2], self.activations[1::2]
+ for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2):
+ residual = hidden_states
+ hidden_states = act1(hidden_states)
+ hidden_states = conv1(hidden_states)
+ hidden_states = act2(hidden_states)
+ hidden_states = conv2(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring(
+ custom_intro="""
+ The full Qwen2.5Omni Token2WavBigVGAN model. Which take mel spectrogram as input and predict waveform.
+ """
+)
+class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniBigVGANConfig
+ input_modalities = "audio"
+
+ def __init__(self, config: Qwen2_5OmniBigVGANConfig):
+ super().__init__(config)
+ self.num_residual_blocks = len(config.resblock_kernel_sizes)
+ self.num_upsample_layers = len(config.upsample_rates)
+
+ self.conv_pre = nn.Conv1d(config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3)
+
+ # Removing extra ModuleList breaks official state dict
+ ups = [
+ nn.ModuleList(
+ [
+ nn.ConvTranspose1d(
+ config.upsample_initial_channel // (2**layer_idx),
+ config.upsample_initial_channel // (2 ** (layer_idx + 1)),
+ kernel_size,
+ stride,
+ padding=(kernel_size - stride) // 2,
+ )
+ ]
+ )
+ for layer_idx, (stride, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes))
+ ]
+ self.ups = nn.ModuleList(ups)
+
+ self.resblocks = nn.ModuleList(
+ [
+ AMPBlock(config.upsample_initial_channel // (2 ** (layer_idx + 1)), kernel_size, dilation)
+ for layer_idx in range(self.num_upsample_layers)
+ for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes)
+ ]
+ )
+
+ self.activation_post = TorchActivation1d(
+ activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers))
+ )
+ self.conv_post = nn.Conv1d(
+ config.upsample_initial_channel // (2**self.num_upsample_layers), 1, 7, 1, padding=3, bias=False
+ )
+
+ self.post_init()
+
+ def normalize_spectrogram(self, spectrogram, max_value, min_db):
+ return torch.clamp((2 * max_value) * ((spectrogram - min_db) / (-min_db)) - max_value, -max_value, max_value)
+
+ def amplitude_to_db(self, amplitude, min_db_level):
+ min_level = torch.exp(
+ torch.tensor(min_db_level / 20.0 * np.log(10), device=amplitude.device, dtype=amplitude.dtype)
+ )
+ return 20 * torch.log10(torch.clamp(amplitude, min=min_level))
+
+ def process_mel_spectrogram(self, mel_spectrogram):
+ amplitude_spectrum = torch.exp(mel_spectrogram)
+ decibel_spectrum = self.amplitude_to_db(amplitude_spectrum, -115) - 20
+ return self.normalize_spectrogram(decibel_spectrum, 1, -115)
+
+ def forward(self, mel_spectrogram, **kwargs):
+ processed_spectrogram = self.process_mel_spectrogram(mel_spectrogram)
+ hidden_representation = self.conv_pre(processed_spectrogram)
+
+ for layer_index in range(self.num_upsample_layers):
+ hidden_representation = self.ups[layer_index][0](hidden_representation)
+ residual_output = sum(
+ self.resblocks[layer_index * self.num_residual_blocks + block_index](hidden_representation)
+ for block_index in range(self.num_residual_blocks)
+ )
+ residual_output = residual_output / self.num_residual_blocks
+ hidden_representation = residual_output
+
+ hidden_representation = self.activation_post(hidden_representation)
+ output_waveform = self.conv_post(hidden_representation)
+ return torch.clamp(output_waveform, min=-1.0, max=1.0).squeeze().cpu()
+
+
+class RungeKutta4ODESolver:
+ def __init__(self, function, initial_value):
+ self.function = function
+ self.initial_value = initial_value
+
+ self._one_third = 1 / 3
+ self._two_thirds = 2 / 3
+
+ def _rk4_step(self, function, time_start, time_step, time_end, value_start, function_value_start=None):
+ k1 = function_value_start if function_value_start is not None else function(time_start, value_start)
+ k2 = function(time_start + time_step * self._one_third, value_start + time_step * k1 * self._one_third)
+ k3 = function(time_start + time_step * self._two_thirds, value_start + time_step * (k2 - k1 * self._one_third))
+ k4 = function(time_end, value_start + time_step * (k1 - k2 + k3))
+ return (k1 + 3 * (k2 + k3) + k4) * time_step / 8
+
+ def _compute_step(self, function, time_start, time_step, time_end, value_start):
+ function_value_start = function(time_start, value_start)
+ return self._rk4_step(
+ function, time_start, time_step, time_end, value_start, function_value_start=function_value_start
+ ), function_value_start
+
+ def _linear_interpolation(self, time_start, time_end, value_start, value_end, time_point):
+ if time_point == time_start:
+ return value_start
+ if time_point == time_end:
+ return value_end
+ weight = (time_point - time_start) / (time_end - time_start)
+ return value_start + weight * (value_end - value_start)
+
+ def integrate(self, time_points):
+ solution = torch.empty(
+ len(time_points),
+ *self.initial_value.shape,
+ dtype=self.initial_value.dtype,
+ device=self.initial_value.device,
+ )
+ solution[0] = self.initial_value
+
+ current_index = 1
+ current_value = self.initial_value
+ for time_start, time_end in zip(time_points[:-1], time_points[1:]):
+ time_step = time_end - time_start
+ delta_value, _ = self._compute_step(self.function, time_start, time_step, time_end, current_value)
+ next_value = current_value + delta_value
+
+ while current_index < len(time_points) and time_end >= time_points[current_index]:
+ solution[current_index] = self._linear_interpolation(
+ time_start, time_end, current_value, next_value, time_points[current_index]
+ )
+ current_index += 1
+
+ current_value = next_value
+
+ return solution
+
+
+@auto_docstring(
+ custom_intro="""
+ The full Qwen2.5Omni Token2WavDiT model. Which take speech tokens as input and predict mel spectrogram.
+ """
+)
+class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniDiTConfig
+ input_modalities = "audio"
+ _no_split_modules = ["DiTDecoderLayer"]
+
+ def __init__(self, config: Qwen2_5OmniDiTConfig):
+ super().__init__(config)
+ self.mel_dim = config.mel_dim
+ self.repeats = config.repeats
+ self.time_embed = DiTTimestepEmbedding(config.hidden_size)
+
+ self.text_embed = DiTCodecEmbedding(config.num_embeds, config.emb_dim, config.repeats)
+ self.input_embed = DiTInputEmbedding(config)
+
+ self.rotary_embed = Qwen2_5OmniDiTRotaryEmbedding(config=config)
+
+ self.hidden_size = config.hidden_size
+ self.layers = config.num_hidden_layers
+ self.block_size = config.block_size
+ self.num_attention_heads = config.num_attention_heads
+
+ self.transformer_blocks = nn.ModuleList()
+ for i in range(config.num_hidden_layers):
+ self.transformer_blocks.append(
+ DiTDecoderLayer(
+ config,
+ look_ahead_block=1 if i in config.look_ahead_layers else 0,
+ look_backward_block=1 if i in config.look_backward_layers else 0,
+ )
+ )
+
+ self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(config.hidden_size) # final modulation
+ self.proj_out = nn.Linear(config.hidden_size, config.mel_dim)
+
+ self.post_init()
+
+ def _create_block_diff(self, hidden_states):
+ batch, seq_len = hidden_states.shape[0], hidden_states.shape[1]
+ block_indices = torch.arange(seq_len, device=hidden_states.device) // self.block_size # [seq_length]
+
+ block_i = block_indices.unsqueeze(1) # [seq_length, 1]
+ block_j = block_indices.unsqueeze(0) # [1, seq_length]
+ block_diff = block_j - block_i # (n, n)
+
+ return block_diff.expand(batch, self.num_attention_heads, seq_len, seq_len)
+
+ def forward(
+ self,
+ hidden_states,
+ condition_vector,
+ speaker_embedding,
+ quantized_code,
+ time_step,
+ drop_audio_conditioning=False,
+ drop_code=False,
+ apply_cfg=True,
+ **kwargs,
+ ):
+ batch_size = hidden_states.shape[0]
+ if time_step.ndim == 0:
+ time_step = time_step.repeat(batch_size)
+
+ # Compute embeddings
+ time_embedding = self.time_embed(time_step)
+ text_embedding = self.text_embed(quantized_code, drop_code=False if apply_cfg else drop_code)
+ text_embedding_unconditioned = self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
+
+ hidden_states = self.input_embed(
+ hidden_states,
+ speaker_embedding,
+ condition_vector,
+ text_embedding,
+ drop_audio_cond=drop_audio_conditioning,
+ code_embed_uncond=text_embedding_unconditioned,
+ apply_cfg=apply_cfg,
+ )
+
+ # Compute positional encodings
+ position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device)
+ position_ids = position_ids[None, :].repeat(batch_size, 1)
+ position_embeddings = self.rotary_embed(hidden_states, position_ids)
+ blockwise_difference = self._create_block_diff(hidden_states)
+
+ # Transformer blocks
+ for transformer_block in self.transformer_blocks:
+ hidden_states = transformer_block(
+ hidden_states,
+ time_embedding,
+ position_embeddings=position_embeddings,
+ block_diff=blockwise_difference,
+ )
+
+ hidden_states = self.norm_out(hidden_states, time_embedding)
+ output = self.proj_out(hidden_states)
+
+ return output
+
+ @torch.no_grad()
+ def sample(
+ self,
+ conditioning_vector,
+ reference_mel_spectrogram,
+ quantized_code,
+ num_steps=10,
+ guidance_scale=0.5,
+ sway_coefficient=-1.0,
+ ):
+ maximum_duration = quantized_code.shape[1] * self.repeats
+ batch_size = reference_mel_spectrogram.shape[0]
+ if batch_size != 1:
+ raise ValueError("Only batch size = 1 is currently supported")
+
+ if maximum_duration > self.config.max_position_embeddings:
+ raise ValueError(
+ f"Requested mel length ({maximum_duration}) exceeds `dit_config.max_position_embeddings` "
+ f"({self.config.max_position_embeddings}). Provide shorter `quantized_code`."
+ )
+
+ initial_state = torch.randn(
+ [batch_size, maximum_duration, self.mel_dim],
+ dtype=reference_mel_spectrogram.dtype,
+ device=quantized_code.device,
+ )
+ conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, maximum_duration, 1)
+
+ def ode_function(time_step, hidden_states):
+ if guidance_scale < 1e-5:
+ prediction = self(
+ hidden_states=hidden_states,
+ speaker_embedding=conditioning_vector,
+ condition_vector=reference_mel_spectrogram,
+ quantized_code=quantized_code,
+ time_step=time_step,
+ drop_audio_conditioning=False,
+ drop_code=False,
+ apply_cfg=False,
+ )
+ return prediction
+
+ model_output = self(
+ hidden_states=hidden_states,
+ quantized_code=quantized_code,
+ speaker_embedding=conditioning_vector,
+ condition_vector=reference_mel_spectrogram,
+ time_step=time_step,
+ apply_cfg=True,
+ )
+ guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
+ return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
+
+ initial_time = 0
+ time_embedding = torch.linspace(
+ initial_time, 1, num_steps, device=quantized_code.device, dtype=conditioning_vector.dtype
+ )
+
+ if sway_coefficient is not None:
+ time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
+
+ ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
+ solution_trajectory = ode_solver.integrate(time_embedding)
+
+ generated_waveform = solution_trajectory[-1]
+ generated_mel_spectrogram = generated_waveform.permute(0, 2, 1)
+ return generated_mel_spectrogram
+
+
+@auto_docstring(
+ custom_intro="""
+ The full Qwen2.5Omni Token2Wav model. Consists a DiT model take speech tokens as input and predict mel spectrogram and a BigVGAN vocoder take mel spectrogram as input and predict waveform.
+ """
+)
+class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
+ config: Qwen2_5OmniToken2WavConfig
+ base_model_prefix = "model"
+ input_modalities = "audio"
+ _no_split_modules = ["Qwen2_5OmniToken2WavDiTModel", "Qwen2_5OmniToken2WavBigVGANModel"]
+
+ def __init__(self, config: Qwen2_5OmniToken2WavConfig):
+ super().__init__(config)
+ attn_impl = config._attn_implementation
+ if is_flash_attention_requested(config):
+ logger.warning_once(
+ "Qwen2_5OmniToken2WavModel must inference with fp32, but Flash Attention only supports fp16 and bf16, "
+ "attention implementation of Qwen2_5OmniToken2WavModel will fallback to sdpa."
+ )
+ attn_impl = "sdpa"
+ elif config._attn_implementation == "eager":
+ logger.warning_once(
+ "Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa"
+ )
+ attn_impl = "sdpa"
+ self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel._from_config(
+ config.dit_config, attn_implementation=attn_impl
+ )
+ self.code2wav_bigvgan_model = Qwen2_5OmniToken2WavBigVGANModel._from_config(
+ config.bigvgan_config, attn_implementation=attn_impl
+ )
+
+ self.post_init()
+
+ def forward(
+ self,
+ code,
+ conditioning,
+ reference_mel,
+ num_steps=10,
+ guidance_scale=0.5,
+ sway_coefficient=-1.0,
+ **kwargs,
+ ):
+ """Generates a waveform from input code and conditioning parameters."""
+
+ mel_spectrogram = self.code2wav_dit_model.sample(
+ conditioning,
+ reference_mel,
+ code,
+ num_steps=num_steps,
+ guidance_scale=guidance_scale,
+ sway_coefficient=sway_coefficient,
+ )
+
+ waveform = self.code2wav_bigvgan_model(mel_spectrogram)
+
+ return waveform
+
+
+############################
+# Start Qwen2.5Omni #
+############################
+
+
+@auto_docstring(
+ custom_intro="""
+ The full Qwen2.5Omni model, a multimodal model composed of 3 sub-models:
+ - [`Qwen2_5OmniThinkerForConditionalGeneration`]:
+ a causal auto-regressive transformer takes text, audio, image, video as input and predict text tokens.
+ - [`Qwen2_5OmniTalkerForConditionalGeneration`]:
+ a causal auto-regressive transformer takes thinker hidden states and response as input and predict speech tokens.
+ - [`Qwen2_5OmniToken2WavModel`]:
+ a DiT model take speech tokens as input and predict mel spectrogram and a BigVGAN vocoder take mel spectrogram as input and predict waveform.
+ """
+)
+class Qwen2_5OmniForConditionalGeneration(Qwen2_5OmniPreTrainedModel, GenerationMixin):
+ config: Qwen2_5OmniConfig
+ output_modalities = ("audio", "text")
+ _no_split_modules = [
+ "Qwen2_5OmniTalkerForConditionalGeneration",
+ "Qwen2_5OmniToken2WavModel",
+ ]
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.thinker = Qwen2_5OmniThinkerForConditionalGeneration(config.thinker_config)
+
+ self.has_talker = config.enable_audio_output
+ self.speaker_map = {}
+ if config.enable_audio_output:
+ self.enable_talker()
+ self.post_init()
+
+ def enable_talker(self):
+ self.talker = Qwen2_5OmniTalkerForConditionalGeneration(self.config.talker_config)
+ self.token2wav = Qwen2_5OmniToken2WavModel(self.config.token2wav_config)
+ self.token2wav.float()
+ self.has_talker = True
+
+ def load_speakers(self, path):
+ check_torch_load_is_safe()
+ for key, value in torch.load(path, weights_only=True).items():
+ self.speaker_map[key] = value
+ logger.info(f"Speaker {list(self.speaker_map.keys())} loaded")
+
+ def disable_talker(self):
+ if hasattr(self, "talker"):
+ del self.talker
+ if hasattr(self, "token2wav"):
+ del self.token2wav
+ self.has_talker = False
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ pretrained_model_name_or_path,
+ *model_args,
+ config=None,
+ cache_dir=None,
+ ignore_mismatched_sizes=False,
+ force_download=False,
+ local_files_only=False,
+ token=None,
+ revision="main",
+ use_safetensors=None,
+ weights_only=True,
+ **kwargs,
+ ):
+ model = super().from_pretrained(
+ pretrained_model_name_or_path,
+ *model_args,
+ config=config,
+ cache_dir=cache_dir,
+ ignore_mismatched_sizes=ignore_mismatched_sizes,
+ force_download=force_download,
+ local_files_only=local_files_only,
+ token=token,
+ revision=revision,
+ use_safetensors=use_safetensors,
+ weights_only=weights_only,
+ **kwargs,
+ )
+ spk_path = cached_file(
+ pretrained_model_name_or_path,
+ "spk_dict.pt",
+ subfolder=kwargs.pop("subfolder", None),
+ cache_dir=kwargs.pop("cache_dir", None),
+ force_download=kwargs.pop("force_download", False),
+ proxies=kwargs.pop("proxies", None),
+ local_files_only=kwargs.pop("local_files_only", False),
+ token=token,
+ revision=kwargs.pop("revision", None),
+ )
+ if spk_path is None:
+ raise ValueError(f"""{pretrained_model_name_or_path}/{spk_path} not exists""")
+ model.load_speakers(spk_path)
+
+ return model
+
+ @torch.no_grad()
+ @deprecate_kwarg("return_audio", version="v5", new_name="generation_mode")
+ # TODO: raushan, defaults should be saved in generation config
+ def generate(
+ self,
+ input_ids: torch.Tensor | None = None,
+ speaker: str = "Chelsie",
+ use_audio_in_video: bool = False,
+ thinker_max_new_tokens: int = 1024,
+ talker_max_new_tokens: int = 4096,
+ talker_do_sample: bool = True,
+ talker_top_k: int = 40,
+ talker_top_p: float = 0.8,
+ talker_temperature: float = 0.9,
+ talker_eos_token_id: list[int] = [8292, 8294],
+ talker_repetition_penalty: float = 1.05,
+ **kwargs,
+ ):
+ r"""
+ Generate text response and audio from input.
+
+ Args:
+ input_ids (`Optional[torch.Tensor]`, *optional*):
+ Input ids, should obtain from processor.
+ speaker (`str` , defaults to "Chelsie"):
+ Which speaker should be used in audio response.
+ use_audio_in_video (`bool`, defaults to False):
+ Whether or not use audio track in video, should same as the parameter in `process_audio_info`.
+ generation_mode (`Optional[str]`, *optional*):
+ Whether or not return response in audio format. When `generation_mode="audio"`, this parameter is same as `config.enable_audio_output`.
+ kwargs (*optional*):
+ - Without a prefix, they will be entered as `**kwargs` for the `generate` method of each sub-model.
+ - With a *thinker_*, *talker_*, *token2wav_* prefix, they will be input for the `generate` method of the
+ thinker, talker and token2wav respectively. It has the priority over the keywords without a prefix.
+ Returns:
+ When `return_audio=False`:
+ - **Text** (`torch.Tensor`): Generated text token sequence.
+ When `return_audio=True`:
+ - **Text** (`torch.Tensor`): Generated text token sequence.
+ - **Audio waveform** (`torch.Tensor`): Generated audio waveform.
+ """
+ # check `False` on purpose because the parameter can be `str/bool`. This is needed for BC
+ generation_mode = kwargs.pop("generation_mode", None)
+ return_audio = generation_mode != "text" and generation_mode is not False
+
+ if speaker not in self.speaker_map:
+ raise ValueError(f"{speaker} is not available, available speakers: {self.speaker_map.keys()}")
+ if return_audio and not self.has_talker:
+ raise ValueError(
+ "Cannot use talker when talker module not initialized. Use `enable_talker` method or set enable_talker in config to enable talker."
+ )
+ if return_audio is None:
+ return_audio = self.has_talker
+ if input_ids.shape[0] != 1 and return_audio:
+ raise NotImplementedError("Qwen2.5-Omni currently does not support batched inference with audio output")
+
+ shared_kwargs = {"use_audio_in_video": use_audio_in_video}
+ thinker_kwargs = {
+ "max_new_tokens": thinker_max_new_tokens,
+ }
+ talker_kwargs = {
+ "max_new_tokens": talker_max_new_tokens,
+ "do_sample": talker_do_sample,
+ "top_k": talker_top_k,
+ "top_p": talker_top_p,
+ "temperature": talker_temperature,
+ "eos_token_id": talker_eos_token_id,
+ "repetition_penalty": talker_repetition_penalty,
+ }
+ token2wav_kwargs = {}
+
+ for key, value in kwargs.items():
+ if key.startswith("thinker_"):
+ thinker_kwargs[key[len("thinker_") :]] = value
+ elif key.startswith("talker_"):
+ talker_kwargs[key[len("talker_") :]] = value
+ elif key.startswith("token2wav_"):
+ token2wav_kwargs[key[len("token2wav_") :]] = value
+ # Process special input values
+ elif key == "feature_attention_mask":
+ thinker_kwargs[key] = value
+ talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
+ elif key == "input_features" or key == "attention_mask":
+ thinker_kwargs[key] = value
+ # Put other key to shared kwargs
+ else:
+ shared_kwargs[key] = value
+
+ # Merge kwargs
+ for key, value in shared_kwargs.items():
+ if key not in thinker_kwargs:
+ thinker_kwargs[key] = value
+ if key not in talker_kwargs:
+ talker_kwargs[key] = value
+ if key not in token2wav_kwargs:
+ token2wav_kwargs[key] = value
+ speaker_params = self.speaker_map[speaker]
+
+ # 1. Generate from thinker module
+ generate_audio = return_audio and self.has_talker
+ if generate_audio:
+ thinker_kwargs["output_hidden_states"] = True
+ thinker_kwargs["return_dict_in_generate"] = True
+
+ thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
+
+ if not generate_audio:
+ return thinker_result
+
+ # 2. Generate speech tokens from talker module
+ embeds_to_talker = thinker_result.hidden_states[0][0].clone().to(input_ids.device)
+ if thinker_kwargs.get("input_features") is not None:
+ audio_ids_mask = input_ids == self.config.thinker_config.audio_token_index
+ audio_mask = audio_ids_mask.unsqueeze(-1).expand_as(embeds_to_talker)
+ audio_mask_tensor = torch.zeros(
+ [audio_ids_mask.sum(), embeds_to_talker.shape[-1]],
+ dtype=embeds_to_talker.dtype,
+ device=input_ids.device,
+ )
+ embeds_to_talker.masked_scatter_(audio_mask, audio_mask_tensor)
+ if thinker_kwargs.get("pixel_values") is not None:
+ image_ids_mask = input_ids == self.config.thinker_config.image_token_index
+ image_mask = image_ids_mask.unsqueeze(-1).expand_as(embeds_to_talker)
+ image_mask_tensor = torch.zeros(
+ [image_ids_mask.sum(), embeds_to_talker.shape[-1]],
+ dtype=embeds_to_talker.dtype,
+ device=input_ids.device,
+ )
+ embeds_to_talker.masked_scatter_(image_mask, image_mask_tensor)
+ if thinker_kwargs.get("pixel_values_videos") is not None:
+ video_ids_mask = input_ids == self.config.thinker_config.video_token_index
+ video_mask = video_ids_mask.unsqueeze(-1).expand_as(embeds_to_talker)
+ video_mask_tensor = torch.zeros(
+ [video_ids_mask.sum(), embeds_to_talker.shape[-1]],
+ dtype=embeds_to_talker.dtype,
+ device=input_ids.device,
+ )
+ embeds_to_talker.masked_scatter_(video_mask, video_mask_tensor)
+
+ processed_thinker_hidden = (
+ (embeds_to_talker,) + thinker_result.hidden_states[0][1:],
+ ) + thinker_result.hidden_states[1:]
+ thinker_generate_ids = thinker_result.sequences[:, input_ids.size(1) :].to(input_ids.device)
+ thinker_token_embeds = [
+ token_hidden_states[0].to(input_ids.device) for token_hidden_states in processed_thinker_hidden
+ ]
+ thinker_hidden_states = [
+ token_hidden_states[-1].to(input_ids.device) for token_hidden_states in processed_thinker_hidden
+ ]
+
+ talker_text_bos_token = speaker_params["bos_token"]
+ talker_input_text_ids = torch.cat(
+ [
+ input_ids,
+ torch.tensor([[talker_text_bos_token]], dtype=torch.long, device=input_ids.device),
+ thinker_generate_ids[:, :1],
+ ],
+ dim=-1,
+ )
+
+ talker_input_ids = torch.cat(
+ [
+ torch.full_like(input_ids, fill_value=self.talker.codec_mask_token),
+ torch.tensor([[self.talker.codec_pad_token]], dtype=torch.long, device=input_ids.device),
+ torch.tensor([[self.talker.codec_bos_token]], dtype=torch.long, device=input_ids.device),
+ ],
+ dim=1,
+ )
+
+ thinker_embed_tokens = self.thinker.get_input_embeddings()
+ thinker_reply_part = torch.cat(thinker_hidden_states[1:], dim=1) + torch.cat(thinker_token_embeds[1:], dim=1)
+ talker_inputs_embeds = thinker_hidden_states[0] + thinker_token_embeds[0]
+ talker_text_bos_token = torch.tensor([[talker_text_bos_token]], dtype=torch.long, device=input_ids.device)
+ talker_text_bos_embed = thinker_embed_tokens(talker_text_bos_token).to(input_ids.device)
+ talker_inputs_embeds = torch.cat(
+ [
+ talker_inputs_embeds,
+ talker_text_bos_embed,
+ thinker_reply_part[:, :1, :],
+ ],
+ dim=1,
+ )
+
+ eos_token = torch.tensor([[self.talker.text_eos_token]], dtype=torch.long, device=input_ids.device)
+ eos_embedding = thinker_embed_tokens(eos_token).to(input_ids.device)
+
+ pad_token = torch.tensor([[self.talker.text_pad_token]], dtype=torch.long, device=input_ids.device)
+ pad_embedding = thinker_embed_tokens(pad_token).to(input_ids.device)
+
+ thinker_reply_part = torch.cat(
+ [
+ thinker_reply_part[:, 1:, :],
+ eos_embedding,
+ pad_embedding,
+ ],
+ dim=1,
+ )
+
+ talker_attention_mask = None
+ if "attention_mask" in kwargs:
+ talker_attention_mask = torch.cat(
+ [kwargs["attention_mask"], kwargs["attention_mask"].new_ones((1, 2))], dim=1
+ ).to(input_ids.device)
+
+ talker_result = self.talker.generate(
+ input_ids=talker_input_ids,
+ input_text_ids=talker_input_text_ids,
+ thinker_reply_part=thinker_reply_part,
+ inputs_embeds=talker_inputs_embeds,
+ attention_mask=talker_attention_mask,
+ suppress_tokens=[self.talker.codec_bos_token],
+ **{k: (v.to(input_ids.device) if torch.is_tensor(v) else v) for k, v in talker_kwargs.items()},
+ )
+ talker_generate_codes = talker_result[:, talker_input_ids.shape[1] : -1]
+
+ # 3. Generate wavs from code
+ if self.token2wav.dtype != torch.float:
+ self.token2wav.float()
+
+ wav = self.token2wav(
+ talker_generate_codes.to(input_ids.device),
+ conditioning=speaker_params["cond"].to(input_ids.device).float(),
+ reference_mel=speaker_params["ref_mel"].to(input_ids.device).float(),
+ **token2wav_kwargs,
+ )
+
+ return thinker_result.sequences, wav.float()
+
+
+__all__ = [
+ "Qwen2_5OmniConfig",
+ "Qwen2_5OmniThinkerConfig",
+ "Qwen2_5OmniTalkerConfig",
+ "Qwen2_5OmniToken2WavConfig",
+ "Qwen2_5OmniAudioEncoderConfig",
+ "Qwen2_5OmniBigVGANConfig",
+ "Qwen2_5OmniDiTConfig",
+ "Qwen2_5OmniTextConfig",
+ "Qwen2_5OmniVisionEncoderConfig",
+ "Qwen2_5OmniForConditionalGeneration",
+ "Qwen2_5OmniThinkerTextModel",
+ "Qwen2_5OmniThinkerForConditionalGeneration",
+ "Qwen2_5OmniTalkerModel",
+ "Qwen2_5OmniTalkerForConditionalGeneration",
+ "Qwen2_5OmniToken2WavDiTModel",
+ "Qwen2_5OmniToken2WavBigVGANModel",
+ "Qwen2_5OmniToken2WavModel",
+ "Qwen2_5OmniPreTrainedModel",
+ "Qwen2_5OmniPreTrainedModelForConditionalGeneration",
+]