| 250 | """ |
| 251 | ) |
| 252 | class FunAudioChatAudioEncoder(FunAudioChatPreTrainedModel): |
| 253 | config_class = FunAudioChatAudioEncoderConfig |
| 254 | main_input_name = "input_features" |
| 255 | _no_split_modules = ["FunAudioChatAudioEncoderLayer"] |
| 256 | _supports_sdpa = True |
| 257 | |
| 258 | def __init__(self, config: FunAudioChatAudioEncoderConfig): |
| 259 | super().__init__(config) |
| 260 | self.dropout = config.dropout |
| 261 | |
| 262 | embed_dim = config.d_model |
| 263 | self.num_mel_bins = config.num_mel_bins |
| 264 | self.max_source_positions = config.max_source_positions |
| 265 | self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0 |
| 266 | self.n_window = config.n_window |
| 267 | self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1) |
| 268 | self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1) |
| 269 | self.layers = nn.ModuleList([FunAudioChatAudioEncoderLayer(config) for _ in range(config.encoder_layers)]) |
| 270 | self.ln_post = nn.LayerNorm(config.d_model) |
| 271 | self.avg_pooler = nn.AvgPool1d(2, stride=2) |
| 272 | self.proj = nn.Linear(config.d_model, config.output_dim) |
| 273 | self.gradient_checkpointing = False |
| 274 | self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, embed_dim) |
| 275 | self.audio_bos_eos_token = nn.Embedding(2, config.output_dim) |
| 276 | # Initialize weights and apply final processing |
| 277 | self.post_init() |
| 278 | |
| 279 | def _freeze_parameters(self): |
| 280 | for param in self.parameters(): |
| 281 | param.requires_grad = False |
| 282 | self._requires_grad = False |
| 283 | |
| 284 | def get_input_embeddings(self) -> nn.Module: |
| 285 | return self.conv1 |
| 286 | |
| 287 | def set_input_embeddings(self, value: nn.Module): |
| 288 | self.conv1 = value |
| 289 | |
| 290 | def _prepare_attention_mask(self, inputs_tensor: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: |
| 291 | # Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen` |
| 292 | if self.config._attn_implementation == "flash_attention_2": |
| 293 | return None |
| 294 | |
| 295 | seq_length = inputs_tensor.shape[0] |
| 296 | attention_mask = torch.full( |
| 297 | [1, 1, seq_length, seq_length], |
| 298 | torch.finfo(inputs_tensor.dtype).min, |
| 299 | device=inputs_tensor.device, |
| 300 | dtype=inputs_tensor.dtype, |
| 301 | ) |
| 302 | for i in range(1, len(cu_seqlens)): |
| 303 | attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0 |
| 304 | return attention_mask |
| 305 | |
| 306 | @auto_docstring |
| 307 | def forward( |
| 308 | self, |
| 309 | input_features, |