Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2025 The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Processor class for Dia"""16 17import math18from pathlib import Path19from typing import Optional, Union20 21from ...audio_utils import AudioInput, make_list_of_audio22from ...feature_extraction_utils import BatchFeature23from ...processing_utils import AudioKwargs, ProcessingKwargs, ProcessorMixin, Unpack24from ...utils import is_soundfile_available, is_torch_available25 26 27if is_torch_available():28    import torch29 30if is_soundfile_available():31    import soundfile as sf32 33 34class DiaAudioKwargs(AudioKwargs, total=False):35    bos_token_id: int36    eos_token_id: int37    pad_token_id: int38    delay_pattern: list[int]39    generation: bool40 41 42class DiaProcessorKwargs(ProcessingKwargs, total=False):43    audio_kwargs: DiaAudioKwargs44    _defaults = {45        "text_kwargs": {46            "padding": True,47            "padding_side": "right",48            "add_special_tokens": False,49        },50        "audio_kwargs": {51            "eos_token_id": 1024,52            "pad_token_id": 1025,53            "bos_token_id": 1026,54            "delay_pattern": [0, 8, 9, 10, 11, 12, 13, 14, 15],55            "generation": True,56            "sampling_rate": 44100,57        },58        "common_kwargs": {"return_tensors": "pt"},59    }60 61 62class DiaProcessor(ProcessorMixin):63    r"""64    Constructs a Dia processor which wraps a [`DiaFeatureExtractor`], [`DiaTokenizer`], and a [`DacModel`] into65    a single processor. It inherits, the audio feature extraction, tokenizer, and audio encode/decode functio-66    nalities. See [`~DiaProcessor.__call__`], [`~DiaProcessor.encode`], and [`~DiaProcessor.decode`] for more67    information.68 69    Args:70        feature_extractor (`DiaFeatureExtractor`):71            An instance of [`DiaFeatureExtractor`]. The feature extractor is a required input.72        tokenizer (`DiaTokenizer`):73            An instance of [`DiaTokenizer`]. The tokenizer is a required input.74        audio_tokenizer (`DacModel`):75            An instance of [`DacModel`] used to encode/decode audio into/from codebooks. It is is a required input.76    """77 78    feature_extractor_class = "DiaFeatureExtractor"79    tokenizer_class = "DiaTokenizer"80    audio_tokenizer_class = "DacModel"81 82    def __init__(self, feature_extractor, tokenizer, audio_tokenizer):83        super().__init__(feature_extractor, tokenizer, audio_tokenizer=audio_tokenizer)84 85    def __call__(86        self,87        text: Union[str, list[str]],88        audio: Optional[AudioInput] = None,89        output_labels: Optional[bool] = False,90        **kwargs: Unpack[DiaProcessorKwargs],91    ):92        """93        Main method to prepare text(s) and audio to be fed as input to the model. The `audio` argument is94        forwarded to the DiaFeatureExtractor's [`~DiaFeatureExtractor.__call__`] and subsequently to the95        DacModel's [`~DacModel.encode`]. The `text` argument to [`~DiaTokenizer.__call__`]. Please refer96        to the docstring of the above methods for more information.97        """98        if not is_torch_available():99            raise ValueError(100                "The `DiaProcessor` relies on the `audio_tokenizer` which requires `torch` but we couldn't "101                "find it in your environment. You can install torch via `pip install torch`."102            )103 104        if text is None:105            raise ValueError("You need to specify the `text` input to process.")106 107        output_kwargs = self._merge_kwargs(108            DiaProcessorKwargs,109            **kwargs,110        )111 112        text_kwargs = output_kwargs["text_kwargs"]113        audio_kwargs = output_kwargs["audio_kwargs"]114        common_kwargs = output_kwargs["common_kwargs"]115 116        return_tensors = common_kwargs.pop("return_tensors", None)117        if return_tensors != "pt":118            raise ValueError(f"{self.__class__.__name__} only supports `return_tensors='pt'`.")119 120        data = {}121 122        # Text123        if isinstance(text, str):124            text = [text]125        elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)):126            raise ValueError("Invalid input text. Please provide a string, or a list of strings")127 128        encodings = self.tokenizer(text, **text_kwargs)129        data.update(encodings)130 131        # Audio132        delay_pattern = audio_kwargs.pop("delay_pattern", None)133        audio_bos_token_id = audio_kwargs.pop("bos_token_id", None)134        audio_eos_token_id = audio_kwargs.pop("eos_token_id", None)135        audio_pad_token_id = audio_kwargs.pop("pad_token_id", None)136        generation = audio_kwargs.pop("generation", True)137        if (138            audio_bos_token_id is None139            or audio_eos_token_id is None140            or audio_pad_token_id is None141            or delay_pattern is None142        ):143            raise ValueError(144                "To enable processing for Dia, we need the `bos_token_id`, `eos_token_id`, "145                "`pad_token_id`, and `delay_pattern`. You may have accidentally overwritten one of those."146            )147 148        if generation and output_labels:149            raise ValueError(150                f"Labels with `generation` is incompatible, got generation={generation}, output_labels={output_labels}."151            )152 153        batch_size = data["input_ids"].shape[0]154        num_channels = len(delay_pattern)155        max_delay = max(delay_pattern)156 157        # Voice cloning generation / general training158        if audio is not None:159            audio = make_list_of_audio(audio)160            input_audios = self.feature_extractor(audio, **audio_kwargs)161 162            compression_rate = math.prod(self.audio_tokenizer.config.downsampling_ratios)163            max_encoded_sequence_len = input_audios["padding_mask"][0].shape[-1] // compression_rate164 165            decoder_input_ids = []166            decoder_attention_mask = []167            # TODO: dac with batching is currently broken, but non-batch is working168            # refer to https://gist.github.com/vasqu/643a45b680cf39fd7467271ee2eb6f80 for a validation script169            for padding_mask, audio in zip(input_audios["padding_mask"], input_audios["input_values"]):170                # get current length with hop length in mind (as if it were sampled as a single audio)171                base_pad_len = self.feature_extractor.hop_length172                current_audio_len = math.ceil(padding_mask.sum(dim=-1) / base_pad_len) * base_pad_len173 174                encoded_sequence_len = current_audio_len // compression_rate175                padding_len = max_encoded_sequence_len - encoded_sequence_len176 177                # compute non-padded forward pass; one extra bos (and eos if training) is added178                with torch.no_grad():179                    audio = audio[None, ..., :current_audio_len].to(self.audio_tokenizer.device)180                    input_ids = self.audio_tokenizer.encode(audio).audio_codes.transpose(1, 2)181 182                if not generation:183                    input_ids = torch.nn.functional.pad(184                        input_ids, pad=(0, 0, 0, 1, 0, 0), mode="constant", value=audio_eos_token_id185                    )186 187                # apply padding188                # +1 for the bos within the real sequence189                input_ids = torch.nn.functional.pad(190                    input_ids, pad=(0, 0, padding_len + 1, 0, 0, 0), mode="constant", value=audio_bos_token_id191                )192                num_valid_inputs = encoded_sequence_len + 1 + max_delay  # sequence + bos + delay193                num_valid_inputs += 0 if generation else 1  # eos if training194                attention_mask = torch.tensor([0] * padding_len + [1] * num_valid_inputs, dtype=torch.long)[None, :]195 196                decoder_input_ids.append(input_ids)197                decoder_attention_mask.append(attention_mask)198 199            decoder_input_ids = torch.cat(decoder_input_ids, dim=0)200            decoder_attention_mask = torch.cat(decoder_attention_mask, dim=0)201        # TTS generation202        elif generation:203            # all bos to start with TTS204            decoder_input_ids = torch.full((batch_size, 1, num_channels), audio_bos_token_id, dtype=torch.long)205 206            # we preemptively add the delay207            decoder_attention_mask = torch.ones(size=(batch_size, 1 + max_delay), dtype=torch.long)208        else:209            raise ValueError("If you try to train, you should provide audio data as well.")210 211        if batch_size != decoder_input_ids.shape[0]:212            raise ValueError(213                f"Need the same amount of samples for both text and audio, but got text samples={batch_size} and "214                f"audio samples = {decoder_input_ids.shape[0]} instead."215            )216 217        # prepare shift indices per delay218        max_seq_len = decoder_attention_mask.shape[-1]219        max_audio_len = max_seq_len - max_delay220        precomputed_idx = self.build_indices(221            bsz=batch_size,222            seq_len=max_seq_len,223            num_channels=num_channels,224            delay_pattern=delay_pattern,225            revert=False,226        )227 228        # create delay pattern input229        # the pad token will be used for masking which input is valid for prediction during generation230        prefill = torch.full(231            (batch_size, max_seq_len, num_channels),232            fill_value=audio_pad_token_id,233            dtype=torch.int,234        )235        prefill[:, :max_audio_len] = decoder_input_ids236 237        delayed_decoder_input_ids = self.apply_audio_delay(238            audio=prefill,239            pad_token_id=audio_pad_token_id,240            bos_token_id=audio_bos_token_id,241            precomputed_idx=precomputed_idx,242        )243 244        data.update({"decoder_input_ids": delayed_decoder_input_ids, "decoder_attention_mask": decoder_attention_mask})245 246        if output_labels:247            # Base idea is to shift on the sequence dim248            labels = data["decoder_input_ids"].clone()[:, 1:]249            labels[labels == audio_pad_token_id] = -100250            labels[labels == audio_bos_token_id] = -100251 252            data["labels"] = labels.transpose(1, 2).reshape(batch_size * num_channels, -1).contiguous().long()253            data["decoder_input_ids"] = data["decoder_input_ids"][:, :-1]254            data["decoder_attention_mask"] = data["decoder_attention_mask"][:, :-1]255 256        return BatchFeature(data=data, tensor_type=return_tensors)257 258    def batch_decode(259        self,260        decoder_input_ids: "torch.Tensor",261        audio_prompt_len: Optional[int] = None,262        **kwargs: Unpack[DiaProcessorKwargs],263    ) -> list["torch.Tensor"]:264        """265        Decodes a batch of audio codebook sequences into their respective audio waveforms via the266        `audio_tokenizer`. See [`~DacModel.decode`] for more information.267 268        Args:269            decoder_input_ids (`torch.Tensor`): The complete output sequence of the decoder.270            audio_prompt_len (`int`): The audio prefix length (e.g. when using voice cloning).271        """272        output_kwargs = self._merge_kwargs(273            DiaProcessorKwargs,274            **kwargs,275        )276        audio_kwargs = output_kwargs["audio_kwargs"]277 278        delay_pattern = audio_kwargs.pop("delay_pattern", None)279        audio_bos_token_id = audio_kwargs.pop("bos_token_id", None)280        audio_pad_token_id = audio_kwargs.pop("pad_token_id", None)281        if audio_bos_token_id is None or audio_pad_token_id is None or delay_pattern is None:282            raise ValueError(283                "To enable decoding for Dia, we need the `bos_token_id`, `pad_token_id`, "284                "and `delay_pattern`. You may have accidentally overwritten one of those."285            )286 287        # either decode the whole audio sequence or only the generated parts288        if audio_prompt_len is not None:289            audio_prompt_len = torch.tensor(audio_prompt_len, device=decoder_input_ids.device, dtype=torch.long)290            start_of_generation_idx = audio_prompt_len[None].expand(decoder_input_ids.shape[0])291        else:292            start_of_generation_idx = (decoder_input_ids[:, :, 0] == audio_bos_token_id).sum(dim=-1)293        # -1 for the eos token294        end_of_generation_idx = (295            decoder_input_ids.shape[1] - (decoder_input_ids[:, :, 0] == audio_pad_token_id).sum(dim=-1) - 1296        )297 298        # revert delay299        bsz, seq_len, num_channels = decoder_input_ids.shape300        precomputed_idx = self.build_indices(301            bsz=bsz,302            seq_len=seq_len,303            num_channels=num_channels,304            delay_pattern=delay_pattern,305            revert=True,306        )307 308        output_sequences = self.apply_audio_delay(309            audio=decoder_input_ids,310            # We do not care about these values as we cut them out311            # with `start_of_generation_idx` and `end_of_generation_idx`312            pad_token_id=-1,313            bos_token_id=-1,314            precomputed_idx=precomputed_idx,315        ).transpose(1, 2)316 317        # retrieve the correct sequences each318        audios = []319        # TODO: see above, dac doesn't work in batches yet320        with torch.no_grad():321            for i in range(start_of_generation_idx.shape[0]):322                output_i = output_sequences[i, :, start_of_generation_idx[i] : end_of_generation_idx[i]][None, ...]323                output_i = output_i.to(self.audio_tokenizer.device)324                audio_i = self.audio_tokenizer.decode(audio_codes=output_i).audio_values.cpu().squeeze()325                audios.append(audio_i)326 327        return audios328 329    def decode(330        self,331        decoder_input_ids: "torch.Tensor",332        audio_prompt_len: Optional[int] = None,333        **kwargs: Unpack[DiaProcessorKwargs],334    ) -> "torch.Tensor":335        """336        Decodes a single sequence of audio codebooks into the respective audio waveform via the337        `audio_tokenizer`. See [`~DacModel.decode`] and [`~DiaProcessor.batch_decode`] for more information.338        """339        if decoder_input_ids.shape[0] != 1:340            raise ValueError(341                f"Expecting a single output to be decoded but received {decoder_input_ids.shape[0]} samples instead."342            )343 344        return self.batch_decode(decoder_input_ids, audio_prompt_len, **kwargs)[0]345 346    def get_audio_prompt_len(347        self,348        decoder_attention_mask: "torch.Tensor",349        **kwargs: Unpack[DiaProcessorKwargs],350    ) -> int:351        """Utility function to get the audio prompt length."""352        output_kwargs = self._merge_kwargs(353            DiaProcessorKwargs,354            **kwargs,355        )356        audio_kwargs = output_kwargs["audio_kwargs"]357 358        delay_pattern = audio_kwargs.pop("delay_pattern", None)359        if delay_pattern is None:360            raise ValueError(361                "To enable the utility of retrieving the prompt length for Dia, we need the "362                "`delay_pattern`. You may have accidentally overwritten this."363            )364        return decoder_attention_mask.shape[1] - max(delay_pattern)365 366    # Copied from transformers.models.csm.processing_csm.CsmProcessor.save_audio with Csm->Dia367    def save_audio(368        self,369        audio: AudioInput,370        saving_path: Union[str, Path, list[Union[str, Path]]],371        **kwargs: Unpack[DiaProcessorKwargs],372    ):373        # TODO: @eustlb, this should be in AudioProcessor374        if not is_soundfile_available():375            raise ImportError("Please install `soundfile` to save audio files.")376 377        # ensure correct audio input378        audio = make_list_of_audio(audio)379 380        # ensure correct saving path381        if isinstance(saving_path, (str, Path)):382            saving_path = [saving_path]383        elif not (isinstance(saving_path, (list, tuple)) and all(isinstance(p, (str, Path)) for p in saving_path)):384            raise ValueError("Invalid input path. Please provide a string, or a list of strings")385 386        if len(audio) != len(saving_path):387            raise ValueError("The number of audio and saving paths must be the same")388 389        output_kwargs = self._merge_kwargs(390            DiaProcessorKwargs,391            **kwargs,392        )393        audio_kwargs = output_kwargs["audio_kwargs"]394        sampling_rate = audio_kwargs["sampling_rate"]395 396        for audio_value, p in zip(audio, saving_path):397            if isinstance(audio_value, torch.Tensor):398                audio_value = audio_value.cpu().float().numpy()399            sf.write(p, audio_value, sampling_rate)400 401    @staticmethod402    def build_indices(403        bsz: int,404        seq_len: int,405        num_channels: int,406        delay_pattern: list[int],407        revert: bool = False,408    ) -> tuple["torch.Tensor", "torch.Tensor"]:409        """410        Precompute (sequence_idx, all_idx) so that out[seq, channel] = in[seq - delay[channel], channel]411        or in[seq, channel] = out[seq + delay[channel], channel] if `revert`.412        Negative sequence_idx => BOS; sequence_idx >= seq_len => PAD.413        """414        delay_array = torch.tensor(delay_pattern, dtype=torch.int32)415 416        # (0..seq_len-1)417        sequence_idx = torch.arange(seq_len, dtype=torch.int32)[None, :].expand(bsz, seq_len)[..., None]418        # + or - delay depending if we delay or revert the delay419        if not revert:420            sequence_idx = sequence_idx - delay_array[None, None, :]421        else:422            sequence_idx = sequence_idx + delay_array[None, None, :]423        # if delay goes over the range we clamp back to valid values424        valid_sequence_idx = torch.clamp(sequence_idx, 0, seq_len - 1)425 426        batch_idx = torch.arange(bsz, dtype=torch.int32)[:, None, None].expand(bsz, seq_len, num_channels)427        channel_idx = torch.arange(num_channels, dtype=torch.int32)[None, None, :].expand(bsz, seq_len, num_channels)428 429        all_idx = torch.stack(430            [batch_idx.reshape(-1), valid_sequence_idx.reshape(-1), channel_idx.reshape(-1)],431            dim=1,432        ).long()433 434        return sequence_idx, all_idx435 436    @staticmethod437    def apply_audio_delay(438        audio: "torch.Tensor",439        pad_token_id: int,440        bos_token_id: int,441        precomputed_idx: tuple["torch.Tensor", "torch.Tensor"],442    ) -> "torch.Tensor":443        """444        Applies or reverts the delay pattern to batched audio tokens using precomputed indices,445        inserting BOS where sequence_idx < 0 and PAD where sequence_idx >= seq_len.446 447        Args:448            audio: audio tokens of shape [bsz, seq_len, num_channels]449            pad_token_id: the PAD token450            bos_token_id: the BOS token451            precomputed_idx: from `build_indices`452 453        Returns:454            final_audio: delayed or reverted audio tokens of shape [bsz, seq_len, num_channels]455        """456        # Move everything to the same device457        device = audio.device458        sequence_idx, all_idx = precomputed_idx459        sequence_idx = sequence_idx.to(device)460        all_idx = all_idx.to(device)461 462        # Gather per precomputed indices463        batch_idx, valid_sequence_idx, channel_idx = torch.unbind(all_idx, dim=-1)464        gathered_audio = audio[batch_idx, valid_sequence_idx, channel_idx].view(audio.size())465 466        # Mask according to negative sequence_idx => BOS; sequence_idx >= seq_len => PAD467        mask_bos = sequence_idx < 0468        mask_pad = sequence_idx >= audio.shape[1]469        final_audio = torch.where(mask_bos, bos_token_id, torch.where(mask_pad, pad_token_id, gathered_audio))470 471        return final_audio472 473 474__all__ = ["DiaProcessor"]475