Encodes the input audio waveform into discrete codes. Args: input_values (mx.array): The input audio waveform with shape ``(batch_size, channels, sequence_length)``. padding_mask (mx.array): Padding mask used to pad the ``input_values``.
(
self,
input_values: mx.array,
padding_mask: mx.array = None,
bandwidth: Optional[float] = None,
)
| 508 | return codes, scale |
| 509 | |
| 510 | def encode( |
| 511 | self, |
| 512 | input_values: mx.array, |
| 513 | padding_mask: mx.array = None, |
| 514 | bandwidth: Optional[float] = None, |
| 515 | ) -> Tuple[mx.array, Optional[mx.array]]: |
| 516 | """ |
| 517 | Encodes the input audio waveform into discrete codes. |
| 518 | |
| 519 | Args: |
| 520 | input_values (mx.array): The input audio waveform with shape |
| 521 | ``(batch_size, channels, sequence_length)``. |
| 522 | padding_mask (mx.array): Padding mask used to pad the ``input_values``. |
| 523 | bandwidth (float, optional): The target bandwidth. Must be one of |
| 524 | ``config.target_bandwidths``. If ``None``, uses the smallest |
| 525 | possible bandwidth. bandwidth is represented as a thousandth of |
| 526 | what it is, e.g. 6kbps bandwidth is represented as bandwidth == 6.0 |
| 527 | |
| 528 | Returns: |
| 529 | A list of frames containing the discrete encoded codes for the |
| 530 | input audio waveform, along with rescaling factors for each chunk |
| 531 | when ``config.normalize==True``. Each frame is a tuple ``(codebook, |
| 532 | scale)``, with ``codebook`` of shape ``(batch_size, num_codebooks, |
| 533 | frames)``. |
| 534 | """ |
| 535 | |
| 536 | if bandwidth is None: |
| 537 | bandwidth = self.config.target_bandwidths[0] |
| 538 | if bandwidth not in self.config.target_bandwidths: |
| 539 | raise ValueError( |
| 540 | f"This model doesn't support the bandwidth {bandwidth}. " |
| 541 | f"Select one of {self.config.target_bandwidths}." |
| 542 | ) |
| 543 | |
| 544 | _, input_length, channels = input_values.shape |
| 545 | |
| 546 | if channels < 1 or channels > 2: |
| 547 | raise ValueError( |
| 548 | f"Number of audio channels must be 1 or 2, but got {channels}" |
| 549 | ) |
| 550 | |
| 551 | chunk_length = self.chunk_length |
| 552 | if chunk_length is None: |
| 553 | chunk_length = input_length |
| 554 | stride = input_length |
| 555 | else: |
| 556 | stride = self.chunk_stride |
| 557 | |
| 558 | if padding_mask is None: |
| 559 | padding_mask = mx.ones(input_values.shape[:2], dtype=mx.bool_) |
| 560 | encoded_frames = [] |
| 561 | scales = [] |
| 562 | |
| 563 | step = chunk_length - stride |
| 564 | if (input_length % stride) != step: |
| 565 | raise ValueError( |
| 566 | "The input length is not properly padded for batched chunked " |
| 567 | "encoding. Make sure to pad the input correctly." |