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 bark model generation configuration import copy from typing import dict from generation configurationutils import generationconfig from utils import logging logger logging getloggername class barksemanticgenerationconfiggenerationconfig modeltype semantic def init self eostokenid10000 renormalizelogitstrue maxnewtokens768 outputscoresfalse returndictingeneratefalse outputhiddenstatesfalse outputattentionsfalse temperature1 0 dosamplefalse textencodingoffset10048 textpadtoken129595 semanticinfertoken129599 semanticvocabsize10000 maxinputsemanticlength256 semanticratehz49 9 mineospnone kwargs super init temperaturetemperature dosampledosample eostokenideostokenid renormalizelogitsrenormalizelogits maxnewtokensmaxnewtokens outputscoresoutputscores returndictingeneratereturndictingenerate outputhiddenstatesoutputhiddenstates outputattentionsoutputattentions kwargs self textencodingoffset textencodingoffset self textpadtoken textpadtoken self semanticpadtoken eostokenid self semanticinfertoken semanticinfertoken self semanticvocabsize semanticvocabsize self maxinputsemanticlength maxinputsemanticlength self semanticratehz semanticratehz self mineosp mineosp class barkcoarsegenerationconfiggenerationconfig modeltype coarseacoustics def init self renormalizelogitstrue outputscoresfalse returndictingeneratefalse outputhiddenstatesfalse outputattentionsfalse temperature1 0 dosamplefalse coarsesemanticpadtoken12048 coarseratehz75 ncoarsecodebooks2 coarseinfertoken12050 maxcoarseinputlength256 maxcoarsehistory int 630 slidingwindowlen int 60 kwargs super init temperaturetemperature dosampledosample renormalizelogitsrenormalizelogits outputscoresoutputscores returndictingeneratereturndictingenerate outputhiddenstatesoutputhiddenstates outputattentionsoutputattentions kwargs self coarsesemanticpadtoken coarsesemanticpadtoken self coarseratehz coarseratehz self ncoarsecodebooks ncoarsecodebooks self coarseinfertoken coarseinfertoken self maxcoarseinputlength maxcoarseinputlength self maxcoarsehistory maxcoarsehistory self slidingwindowlen slidingwindowlen class barkfinegenerationconfiggenerationconfig modeltype fineacoustics def init self temperature1 0 maxfinehistorylength512 maxfineinputlength1024 nfinecodebooks8 kwargs super inittemperaturetemperature self maxfinehistorylength maxfinehistorylength self maxfineinputlength maxfineinputlength self nfinecodebooks nfinecodebooks def validateself kwargs pass class barkgenerationconfiggenerationconfig modeltype bark iscomposition true todo joao nested fromdict def init self semanticconfig dict none coarseacousticsconfig dict none fineacousticsconfig dict none samplerate24000 codebooksize1024 kwargs if semanticconfig is none semanticconfig logger infosemanticconfig is none initializing the semantic model with default values if coarseacousticsconfig is none coarseacousticsconfig logger infocoarseacousticsconfig is none initializing the coarse model with default values if fineacousticsconfig is none fineacousticsconfig logger infofineacousticsconfig is none initializing the fine model with default values self semanticconfig barksemanticgenerationconfigsemanticconfig self coarseacousticsconfig barkcoarsegenerationconfigcoarseacousticsconfig self fineacousticsconfig barkfinegenerationconfigfineacousticsconfig self samplerate samplerate self codebooksize codebooksize classmethod def fromsubmodelconfigs cls semanticconfig barksemanticgenerationconfig coarseacousticsconfig barkcoarsegenerationconfig fineacousticsconfig barkfinegenerationconfig kwargs r instantiate a barkgenerationconfig or a derived class from bark submodels generation configuration returns barkgenerationconfig an instance of a configuration object serializes this instance to a python dictionary override the default pretrainedconfig todict returns dictstr any dictionary of all the attributes that make up this configuration instance 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 bark model generation configuration class that holds a generation configuration for barksemanticmodel this configuration inherit from generationconfig and can be used to control the model generation read the documentation from generationconfig for more information args eos_token_id int optional defaults to 10_000 the id of the end of sequence token renormalize_logits bool optional defaults to true whether to renormalize the logits after applying all the logits processors or warpers including the custom ones it s highly recommended to set this flag to true as the search algorithms suppose the score logits are normalized but some logit processors or warpers break the normalization max_new_tokens int optional defaults to 768 the maximum numbers of tokens to generate ignoring the number of tokens in the prompt output_scores bool optional defaults to false whether or not to return the prediction scores see scores under returned tensors for more details return_dict_in_generate bool optional defaults to false whether or not to return a utils modeloutput instead of a plain tuple output_hidden_states bool optional defaults to false whether or not to return the hidden states of all layers see hidden_states under returned tensors for more details output_attentions bool optional defaults to false whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more details temperature float optional defaults to 1 0 the value used to modulate the next token probabilities do_sample bool optional defaults to false whether or not to use sampling use greedy decoding otherwise text_encoding_offset int optional defaults to 10_048 text encoding offset text_pad_token int optional defaults to 129_595 text pad token semantic_infer_token int optional defaults to 129_599 semantic infer token semantic_vocab_size int optional defaults to 10_000 semantic vocab size max_input_semantic_length int optional defaults to 256 max length of semantic input vector semantic_rate_hz float optional defaults to 49 9 semantic rate in hertz min_eos_p float optional minimum threshold of the probability of the eos token for it to be sampled this is an early stopping strategy to mitigate potential unwanted generations at the end of a prompt the original implementation suggests a default value of 0 2 class that holds a generation configuration for barkcoarsemodel this configuration inherit from generationconfig and can be used to control the model generation read the documentation from generationconfig for more information args renormalize_logits bool optional defaults to true whether to renormalize the logits after applying all the logits processors or warpers including the custom ones it s highly recommended to set this flag to true as the search algorithms suppose the score logits are normalized but some logit processors or warpers break the normalization output_scores bool optional defaults to false whether or not to return the prediction scores see scores under returned tensors for more details return_dict_in_generate bool optional defaults to false whether or not to return a utils modeloutput instead of a plain tuple output_hidden_states bool optional defaults to false whether or not to return the hidden states of all layers see hidden_states under returned tensors for more details output_attentions bool optional defaults to false whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more details temperature float optional defaults to 1 0 the value used to modulate the next token probabilities do_sample bool optional defaults to false whether or not to use sampling use greedy decoding otherwise coarse_semantic_pad_token int optional defaults to 12_048 coarse semantic pad token coarse_rate_hz int optional defaults to 75 coarse rate in hertz n_coarse_codebooks int optional defaults to 2 number of coarse codebooks coarse_infer_token int optional defaults to 12_050 coarse infer token max_coarse_input_length int optional defaults to 256 max length of input coarse vector max_coarse_history int optional defaults to 630 max length of the output of the coarse acoustics model used in the fine generation step sliding_window_len int optional defaults to 60 the coarse generation step uses a sliding window to generate raw audio class that holds a generation configuration for barkfinemodel barkfinemodel is an autoencoder model so should not usually be used for generation however under the hood it uses temperature when used by barkmodel this configuration inherit from generationconfig and can be used to control the model generation read the documentation from generationconfig for more information args temperature float optional the value used to modulate the next token probabilities max_fine_history_length int optional defaults to 512 max length of the fine history vector max_fine_input_length int optional defaults to 1024 max length of fine input vector n_fine_codebooks int optional defaults to 8 number of codebooks used overrides generationconfig validate because barkfinegenerationconfig don t use any parameters outside temperature todo joao nested from_dict class that holds a generation configuration for barkmodel the barkmodel does not have a generate method but uses this class to generate speeches with a nested barkgenerationconfig which uses barksemanticgenerationconfig barkcoarsegenerationconfig barkfinegenerationconfig this configuration inherit from generationconfig and can be used to control the model generation read the documentation from generationconfig for more information args semantic_config dict optional semantic generation configuration coarse_acoustics_config dict optional coarse generation configuration fine_acoustics_config dict optional fine generation configuration sample_rate int optional defaults to 24_000 sample rate codebook_size int optional defaults to 1024 vector length for each codebook instantiate a barkgenerationconfig or a derived class from bark sub models generation configuration returns barkgenerationconfig an instance of a configuration object serializes this instance to a python dictionary override the default pretrainedconfig to_dict returns dict str any dictionary of all the attributes that make up this configuration instance
import copy
from typing import Dict

from ...generation.configuration_utils import GenerationConfig
from ...utils import logging


logger = logging.get_logger(__name__)


class BarkSemanticGenerationConfig(GenerationConfig):
    model_type = "semantic"

    def __init__(
        self,
        eos_token_id=10_000,
        renormalize_logits=True,
        max_new_tokens=768,
        output_scores=False,
        return_dict_in_generate=False,
        output_hidden_states=False,
        output_attentions=False,
        temperature=1.0,
        do_sample=False,
        text_encoding_offset=10_048,
        text_pad_token=129_595,
        semantic_infer_token=129_599,
        semantic_vocab_size=10_000,
        max_input_semantic_length=256,
        semantic_rate_hz=49.9,
        min_eos_p=None,
        **kwargs,
    ):
        
        super().__init__(
            temperature=temperature,
            do_sample=do_sample,
            eos_token_id=eos_token_id,
            renormalize_logits=renormalize_logits,
            max_new_tokens=max_new_tokens,
            output_scores=output_scores,
            return_dict_in_generate=return_dict_in_generate,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            **kwargs,
        )

        self.text_encoding_offset = text_encoding_offset
        self.text_pad_token = text_pad_token
        self.semantic_pad_token = eos_token_id
        self.semantic_infer_token = semantic_infer_token
        self.semantic_vocab_size = semantic_vocab_size
        self.max_input_semantic_length = max_input_semantic_length
        self.semantic_rate_hz = semantic_rate_hz
        self.min_eos_p = min_eos_p


class BarkCoarseGenerationConfig(GenerationConfig):
    model_type = "coarse_acoustics"

    def __init__(
        self,
        renormalize_logits=True,
        output_scores=False,
        return_dict_in_generate=False,
        output_hidden_states=False,
        output_attentions=False,
        temperature=1.0,
        do_sample=False,
        coarse_semantic_pad_token=12_048,
        coarse_rate_hz=75,
        n_coarse_codebooks=2,
        coarse_infer_token=12_050,
        max_coarse_input_length=256,
        max_coarse_history: int = 630,
        sliding_window_len: int = 60,
        **kwargs,
    ):
        
        super().__init__(
            temperature=temperature,
            do_sample=do_sample,
            renormalize_logits=renormalize_logits,
            output_scores=output_scores,
            return_dict_in_generate=return_dict_in_generate,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            **kwargs,
        )

        self.coarse_semantic_pad_token = coarse_semantic_pad_token
        self.coarse_rate_hz = coarse_rate_hz
        self.n_coarse_codebooks = n_coarse_codebooks
        self.coarse_infer_token = coarse_infer_token
        self.max_coarse_input_length = max_coarse_input_length
        self.max_coarse_history = max_coarse_history
        self.sliding_window_len = sliding_window_len


class BarkFineGenerationConfig(GenerationConfig):
    model_type = "fine_acoustics"

    def __init__(
        self,
        temperature=1.0,
        max_fine_history_length=512,
        max_fine_input_length=1024,
        n_fine_codebooks=8,
        **kwargs,
    ):
        
        super().__init__(temperature=temperature)

        self.max_fine_history_length = max_fine_history_length
        self.max_fine_input_length = max_fine_input_length
        self.n_fine_codebooks = n_fine_codebooks

    def validate(self, **kwargs):
        
        pass


class BarkGenerationConfig(GenerationConfig):
    model_type = "bark"
    is_composition = True

    

    def __init__(
        self,
        semantic_config: Dict = None,
        coarse_acoustics_config: Dict = None,
        fine_acoustics_config: Dict = None,
        sample_rate=24_000,
        codebook_size=1024,
        **kwargs,
    ):
        
        if semantic_config is None:
            semantic_config = {}
            logger.info("semantic_config is None. initializing the semantic model with default values.")

        if coarse_acoustics_config is None:
            coarse_acoustics_config = {}
            logger.info("coarse_acoustics_config is None. initializing the coarse model with default values.")

        if fine_acoustics_config is None:
            fine_acoustics_config = {}
            logger.info("fine_acoustics_config is None. initializing the fine model with default values.")

        self.semantic_config = BarkSemanticGenerationConfig(**semantic_config)
        self.coarse_acoustics_config = BarkCoarseGenerationConfig(**coarse_acoustics_config)
        self.fine_acoustics_config = BarkFineGenerationConfig(**fine_acoustics_config)

        self.sample_rate = sample_rate
        self.codebook_size = codebook_size

    @classmethod
    def from_sub_model_configs(
        cls,
        semantic_config: BarkSemanticGenerationConfig,
        coarse_acoustics_config: BarkCoarseGenerationConfig,
        fine_acoustics_config: BarkFineGenerationConfig,
        **kwargs,
    ):
        r
        return cls(
            semantic_config=semantic_config.to_dict(),
            coarse_acoustics_config=coarse_acoustics_config.to_dict(),
            fine_acoustics_config=fine_acoustics_config.to_dict(),
            **kwargs,
        )

    def to_dict(self):
        
        output = copy.deepcopy(self.__dict__)

        output["semantic_config"] = self.semantic_config.to_dict()
        output["coarse_acoustics_config"] = self.coarse_acoustics_config.to_dict()
        output["fine_acoustics_config"] = self.fine_acoustics_config.to_dict()

        output["model_type"] = self.__class__.model_type
        return output