codingutf8 2023 the suno ai s 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 orglicenseslicense2 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 processor class for bark tokenizerclass autotokenizer attributes tokenizer presetshape semanticprompt 1 coarseprompt 2 fineprompt 2 def initself tokenizer speakerembeddingsnone super inittokenizer self speakerembeddings speakerembeddings classmethod def frompretrained cls pretrainedprocessornameorpath speakerembeddingsdictpathspeakerembeddingspath json kwargs r instantiate a bark processor associated with a pretrained model args pretrainedmodelnameorpath str or os pathlike this can be either a string the model id of a pretrained barkprocessor hosted inside a model repo on huggingface co valid model ids can be located at the rootlevel like bertbaseuncased or namespaced under a user or organization name like dbmdzbertbasegermancased a path to a directory containing a processor saved using the barkprocessor savepretrained method e g mymodeldirectory speakerembeddingsdictpath str optional defaults to speakerembeddingspath json the name of the json file containing the speakerembeddings dictionnary located in pretrainedmodelnameorpath if none no speakerembeddings is loaded kwargs additional keyword arguments passed along to both tokenizationutilsbase pretrainedtokenizer frompretrained saves the attributes of this processor tokenizer in the specified directory so that it can be reloaded using the barkprocessor frompretrained method args savedirectory str or os pathlike directory where the tokenizer files and the speaker embeddings will be saved directory will be created if it does not exist speakerembeddingsdictpath str optional defaults to speakerembeddingspath json the name of the json file that will contains the speakerembeddings nested path dictionnary if it exists and that will be located in pretrainedmodelnameorpathspeakerembeddingsdirectory speakerembeddingsdirectory str optional defaults to speakerembeddings the name of the folder in which the speakerembeddings arrays will be saved pushtohub bool optional defaults to false whether or not to push your model to the hugging face model hub after saving it you can specify the repository you want to push to with repoid will default to the name of savedirectory in your namespace kwargs additional key word arguments passed along to the utils pushtohubmixin pushtohub method main method to prepare for the model one or several sequencess this method forwards the text and kwargs arguments to the autotokenizer s autotokenizer call to encode the text the method also proposes a voice preset which is a dictionary of arrays that conditions bark s output kwargs arguments are forwarded to the tokenizer and to cachedfile method if voicepreset is a valid filename args text str liststr listliststr the sequence or batch of sequences to be encoded each sequence can be a string or a list of strings pretokenized string if the sequences are provided as list of strings pretokenized you must set issplitintowordstrue to lift the ambiguity with a batch of sequences voicepreset str dictnp ndarray the voice preset i e the speaker embeddings it can either be a valid voicepreset name e g enspeaker1 or directly a dictionnary of np ndarray embeddings for each submodel of bark or it can be a valid file name of a local npz single voice preset returntensors str or utils tensortype optional if set will return tensors of a particular framework acceptable values are pt return pytorch torch tensor objects np return numpy np ndarray objects returns tuplebatchencoding batchfeature a tuple composed of a batchencoding i e the output of the tokenizer and a batchfeature i e the voice preset with the right tensors type coding utf 8 2023 the suno ai s 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 processor class for bark constructs a bark processor which wraps a text tokenizer and optional bark voice presets into a single processor args tokenizer pretrainedtokenizer an instance of pretrainedtokenizer speaker_embeddings dict dict str optional optional nested speaker embeddings dictionary the first level contains voice preset names e g en_speaker_4 the second level contains semantic_prompt coarse_prompt and fine_prompt embeddings the values correspond to the path of the corresponding np ndarray see here https suno ai notion site 8b8e8749ed514b0cbf3f699013548683 v bc67cff786b04b50b3ceb756fd05f68c for a list of voice_preset_names instantiate a bark processor associated with a pretrained model args pretrained_model_name_or_path str or os pathlike this can be either a string the model id of a pretrained barkprocessor hosted inside a model repo on huggingface co valid model ids can be located at the root level like bert base uncased or namespaced under a user or organization name like dbmdz bert base german cased a path to a directory containing a processor saved using the barkprocessor save_pretrained method e g my_model_directory speaker_embeddings_dict_path str optional defaults to speaker_embeddings_path json the name of the json file containing the speaker_embeddings dictionnary located in pretrained_model_name_or_path if none no speaker_embeddings is loaded kwargs additional keyword arguments passed along to both tokenization_utils_base pretrainedtokenizer from_pretrained os path join pretrained_processor_name_or_path speaker_embeddings_dict_path does not exists no preloaded speaker embeddings will be used make sure to provide a correct path to the json dictionnary if wanted otherwise set speaker_embeddings_dict_path none saves the attributes of this processor tokenizer in the specified directory so that it can be reloaded using the barkprocessor from_pretrained method args save_directory str or os pathlike directory where the tokenizer files and the speaker embeddings will be saved directory will be created if it does not exist speaker_embeddings_dict_path str optional defaults to speaker_embeddings_path json the name of the json file that will contains the speaker_embeddings nested path dictionnary if it exists and that will be located in pretrained_model_name_or_path speaker_embeddings_directory speaker_embeddings_directory str optional defaults to speaker_embeddings the name of the folder in which the speaker_embeddings arrays will be saved push_to_hub bool optional defaults to false whether or not to push your model to the hugging face model hub after saving it you can specify the repository you want to push to with repo_id will default to the name of save_directory in your namespace kwargs additional key word arguments passed along to the utils pushtohubmixin push_to_hub method os path join self speaker_embeddings get repo_or_path voice_preset_paths key does not exists no preloaded voice preset will be used make sure to provide correct paths to the voice_preset embeddings main method to prepare for the model one or several sequences s this method forwards the text and kwargs arguments to the autotokenizer s autotokenizer __call__ to encode the text the method also proposes a voice preset which is a dictionary of arrays that conditions bark s output kwargs arguments are forwarded to the tokenizer and to cached_file method if voice_preset is a valid filename args text str list str list list str the sequence or batch of sequences to be encoded each sequence can be a string or a list of strings pretokenized string if the sequences are provided as list of strings pretokenized you must set is_split_into_words true to lift the ambiguity with a batch of sequences voice_preset str dict np ndarray the voice preset i e the speaker embeddings it can either be a valid voice_preset name e g en_speaker_1 or directly a dictionnary of np ndarray embeddings for each submodel of bark or it can be a valid file name of a local npz single voice preset return_tensors str or utils tensortype optional if set will return tensors of a particular framework acceptable values are pt return pytorch torch tensor objects np return numpy np ndarray objects returns tuple batchencoding batchfeature a tuple composed of a batchencoding i e the output of the tokenizer and a batchfeature i e the voice preset with the right tensors type
import json
import os
from typing import Optional

import numpy as np

from ...feature_extraction_utils import BatchFeature
from ...processing_utils import ProcessorMixin
from ...utils import logging
from ...utils.hub import get_file_from_repo
from ..auto import AutoTokenizer


logger = logging.get_logger(__name__)


class BarkProcessor(ProcessorMixin):
    r

    tokenizer_class = "AutoTokenizer"
    attributes = ["tokenizer"]

    preset_shape = {
        "semantic_prompt": 1,
        "coarse_prompt": 2,
        "fine_prompt": 2,
    }

    def __init__(self, tokenizer, speaker_embeddings=None):
        super().__init__(tokenizer)

        self.speaker_embeddings = speaker_embeddings

    @classmethod
    def from_pretrained(
        cls, pretrained_processor_name_or_path, speaker_embeddings_dict_path="speaker_embeddings_path.json", **kwargs
    ):
        r

        if speaker_embeddings_dict_path is not None:
            speaker_embeddings_path = get_file_from_repo(
                pretrained_processor_name_or_path,
                speaker_embeddings_dict_path,
                subfolder=kwargs.pop("subfolder", None),
                cache_dir=kwargs.pop("cache_dir", None),
                force_download=kwargs.pop("force_download", False),
                proxies=kwargs.pop("proxies", None),
                resume_download=kwargs.pop("resume_download", False),
                local_files_only=kwargs.pop("local_files_only", False),
                token=kwargs.pop("use_auth_token", None),
                revision=kwargs.pop("revision", None),
            )
            if speaker_embeddings_path is None:
                logger.warning(
                    f
                )
                speaker_embeddings = None
            else:
                with open(speaker_embeddings_path) as speaker_embeddings_json:
                    speaker_embeddings = json.load(speaker_embeddings_json)
        else:
            speaker_embeddings = None

        tokenizer = AutoTokenizer.from_pretrained(pretrained_processor_name_or_path, **kwargs)

        return cls(tokenizer=tokenizer, speaker_embeddings=speaker_embeddings)

    def save_pretrained(
        self,
        save_directory,
        speaker_embeddings_dict_path="speaker_embeddings_path.json",
        speaker_embeddings_directory="speaker_embeddings",
        push_to_hub: bool = False,
        **kwargs,
    ):
        
        if self.speaker_embeddings is not None:
            os.makedirs(os.path.join(save_directory, speaker_embeddings_directory, "v2"), exist_ok=True)

            embeddings_dict = {}

            embeddings_dict["repo_or_path"] = save_directory

            for prompt_key in self.speaker_embeddings:
                if prompt_key != "repo_or_path":
                    voice_preset = self._load_voice_preset(prompt_key)

                    tmp_dict = {}
                    for key in self.speaker_embeddings[prompt_key]:
                        np.save(
                            os.path.join(
                                embeddings_dict["repo_or_path"], speaker_embeddings_directory, f"{prompt_key}_{key}"
                            ),
                            voice_preset[key],
                            allow_pickle=False,
                        )
                        tmp_dict[key] = os.path.join(speaker_embeddings_directory, f"{prompt_key}_{key}.npy")

                    embeddings_dict[prompt_key] = tmp_dict

            with open(os.path.join(save_directory, speaker_embeddings_dict_path), "w") as fp:
                json.dump(embeddings_dict, fp)

        super().save_pretrained(save_directory, push_to_hub, **kwargs)

    def _load_voice_preset(self, voice_preset: str = None, **kwargs):
        voice_preset_paths = self.speaker_embeddings[voice_preset]

        voice_preset_dict = {}
        for key in ["semantic_prompt", "coarse_prompt", "fine_prompt"]:
            if key not in voice_preset_paths:
                raise ValueError(
                    f"Voice preset unrecognized, missing {key} as a key in self.speaker_embeddings[{voice_preset}]."
                )

            path = get_file_from_repo(
                self.speaker_embeddings.get("repo_or_path", "/"),
                voice_preset_paths[key],
                subfolder=kwargs.pop("subfolder", None),
                cache_dir=kwargs.pop("cache_dir", None),
                force_download=kwargs.pop("force_download", False),
                proxies=kwargs.pop("proxies", None),
                resume_download=kwargs.pop("resume_download", False),
                local_files_only=kwargs.pop("local_files_only", False),
                token=kwargs.pop("use_auth_token", None),
                revision=kwargs.pop("revision", None),
            )
            if path is None:
                raise ValueError(
                    f
                )

            voice_preset_dict[key] = np.load(path)

        return voice_preset_dict

    def _validate_voice_preset_dict(self, voice_preset: Optional[dict] = None):
        for key in ["semantic_prompt", "coarse_prompt", "fine_prompt"]:
            if key not in voice_preset:
                raise ValueError(f"Voice preset unrecognized, missing {key} as a key.")

            if not isinstance(voice_preset[key], np.ndarray):
                raise ValueError(f"{key} voice preset must be a {str(self.preset_shape[key])}D ndarray.")

            if len(voice_preset[key].shape) != self.preset_shape[key]:
                raise ValueError(f"{key} voice preset must be a {str(self.preset_shape[key])}D ndarray.")

    def __call__(
        self,
        text=None,
        voice_preset=None,
        return_tensors="pt",
        max_length=256,
        add_special_tokens=False,
        return_attention_mask=True,
        return_token_type_ids=False,
        **kwargs,
    ):
        
        if voice_preset is not None and not isinstance(voice_preset, dict):
            if (
                isinstance(voice_preset, str)
                and self.speaker_embeddings is not None
                and voice_preset in self.speaker_embeddings
            ):
                voice_preset = self._load_voice_preset(voice_preset)

            else:
                if isinstance(voice_preset, str) and not voice_preset.endswith(".npz"):
                    voice_preset = voice_preset + ".npz"

                voice_preset = np.load(voice_preset)

        if voice_preset is not None:
            self._validate_voice_preset_dict(voice_preset, **kwargs)
            voice_preset = BatchFeature(data=voice_preset, tensor_type=return_tensors)

        encoded_text = self.tokenizer(
            text,
            return_tensors=return_tensors,
            padding="max_length",
            max_length=max_length,
            return_attention_mask=return_attention_mask,
            return_token_type_ids=return_token_type_ids,
            add_special_tokens=add_special_tokens,
            **kwargs,
        )

        if voice_preset is not None:
            encoded_text["history_prompt"] = voice_preset

        return encoded_text