bark_hf_inference_workflow - Infernet ML
Description: Infernet ML Documentation

Skip to content

# Module: `bark_hf_inference_workflow`

## Bark Inference Workflow

This workflow uses huggingface's transformers library to perform inference on Suno's text-to-speech Bark model.

### Constructor Arguments

- `model_name (Optional[str])`: The source of the model. This can be either `suno/bark` or `suno/bark-small`. Default is `suno/bark`.
- `default_voice_preset (Optional[str])`: The default voice preset to be used. See list of supported presets.

### Additional Installations

Since this workflow uses some additional libraries, you'll need to install `infernet-ml[bark_inference]`. Alternatively, you can install those packages directly. The optional dependencies `"[bark_inference]"` are provided for your convenience.

uv pip

    uv pip install "infernet-ml[bark_inference]"

    pip install "infernet-ml[bark_inference]"

### Input Format

Input to the inference workflow is the following pydantic model:

    class BarkWorkflowInput(BaseModel):
        # prompt to generate audio from
        prompt: str
        # voice to be used. There is a list of supported presets here:
        # here: https://github.com/suno-ai/bark?tab=readme-ov-file#-voice-presets
        voice_preset: Optional[str]

- `"prompt"`: The text prompt to generate audio from.
- `"voice_preset"`: The voice preset to be used. See list of supported presets.

### Output Format

The output of the inference workflow is a pydantic model with the following keys:

    class AudioInferenceResult(BaseModel):
        audio_array: np.ndarray[Any, Any]

- `"audio_array"`: The audio array generated from the input prompt.

### Example

In this example, we will use the Bark Inference Workflow to generate audio from a prompt. We will then write the generated audio to a wav file.

    from scipy.io.wavfile import write as write_wav  # type: ignore
    from infernet_ml.workflows.inference.bark_hf_inference_workflow import (
        BarkHFInferenceWorkflow,
        BarkWorkflowInput,
    )

    workflow = BarkHFInferenceWorkflow(model_name="suno/bark-small", default_voice_preset="v2/en_speaker_0")

    workflow.setup()

    input = BarkWorkflowInput(
        prompt="Hello, my name is Suno. I am a text-to-speech model.",
        voice_preset="v2/en_speaker_5"
    )

    inference_result = workflow.inference(input)

    generated_audio_path = "output.wav"

    # write output to a wav file
    write_wav(
        generated_audio_path,
        BarkHFInferenceWorkflow.SAMPLE_RATE,
        inference_result.audio_array,
    )

## `BarkHFInferenceWorkflow`

Bases: `TTSInferenceWorkflow`

Implementation of Suno TTS Inference Workflow.

Source code in `src/infernet_ml/workflows/inference/bark_hf_inference_workflow.py`

    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207
    208
    209
    210
    211
    212
    213
    214
    215
    216
    217
    218
    219
    220
    221
    222
    223
    224
    225
    226
    227
    228
    229
    230
    231
    232
    233
    234
    class BarkHFInferenceWorkflow(TTSInferenceWorkflow):
        """
        Implementation of Suno TTS Inference Workflow.
        """

        SAMPLE_RATE: int = 24_000
        model: BarkModel
        processor: BarkProcessor

        def __init__(self, *args: Any, **kwargs: Any) -> None:
            super().__init__(*args, **kwargs)
            # name of the model to be used. The allowed values are: suno/bark,
            # suno/bark-small (default: suno/bark)
            self.model_name = kwargs.get("model_name", "suno/bark")
            # default voice preset to be used. Refer to the link for the list of supported
            # presets (default: v2/en_speaker_6)
            # https://github.com/suno-ai/bark?tab=readme-ov-file#-voice-presets
            self.default_voice_preset = kwargs.get(
                "default_voice_preset", "v2/en_speaker_6"
            )
            # device to be used for inference. If cuda is available, it will be used,
            # else cpu will be used
            self.device = "cuda" if torch.cuda.is_available() else "cpu"

        def do_setup(self) -> None:
            """
            Downloads the model from huggingface.
            Returns:
                bool: True on completion of loading model
            """
            self.model = BarkModel.from_pretrained(self.model_name).to(self.device)
            self.processor = AutoProcessor.from_pretrained(self.model_name)

        def do_preprocessing(self, input_data: BarkWorkflowInput) -> BatchEncoding:
            """
            Preprocesses the input data.

            Args:
                input_data (BarkWorkflowInput): input data to be preprocessed

            Returns:
                BatchEncoding: batch encoding of the input data
            """
            text = input_data.prompt
            voice_preset = input_data.voice_preset or self.default_voice_preset
            return self.processor(text, voice_preset=voice_preset).to(self.device)

        def inference(self, input_data: BarkWorkflowInput, log_preprocessed_data: bool = True) -> AudioInferenceResult:  # type: ignore #noqa: E501
            """
            Override super class inference method to be annotated with the correct types.

            Args:
                input_data (str): prompt to generate audio from

            Returns:
                AudioInferenceResult: audio array
            """
            return cast(
                AudioInferenceResult, super().inference(input_data, log_preprocessed_data)
            )

        def do_run_model(self, preprocessed_data: BatchEncoding) -> torch.Tensor:
            """
            Run the model on the preprocessed data.

            Args:
                preprocessed_data (BatchEncoding): preprocessed data

            Returns:
                torch.Tensor: output tensor from the model
            """
            return cast(torch.Tensor, self.model.generate(**preprocessed_data))

        def do_stream(self, preprocessed_input: Any) -> Iterator[Any]:
            """
            Stream data for inference. Currently not implemented.

            Args:
                preprocessed_input (Any): preprocessed input data

            Returns:
                Iterator[Any]: iterator for streaming data

            Raises:
                NotImplementedError: if the method is not implemented
            """
            raise NotImplementedError

        def do_postprocessing(
            self, input_data: Any, output: torch.Tensor
        ) -> AudioInferenceResult:
            """
            Converts the model output to a numpy array, which then can be used to save the
            audio file.

            Args:
                input_data(Any): original input data
                output (torch.Tensor): output tensor from the model

            Returns:
                AudioInferenceResult: audio array
            """
            audio_array: numpy.ndarray[Any, Any] = output.cpu().numpy().squeeze()
            return AudioInferenceResult(audio_array=audio_array)

### `do_postprocessing(input_data, output)`

Converts the model output to a numpy array, which then can be used to save the audio file.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`input_data(Any)` |  | original input data | required
`output` | `Tensor` | output tensor from the model | required

Returns:

Name | Type | Description
--- | --- | ---
`AudioInferenceResult` | `AudioInferenceResult` | audio array

Source code in `src/infernet_ml/workflows/inference/bark_hf_inference_workflow.py`

    219
    220
    221
    222
    223
    224
    225
    226
    227
    228
    229
    230
    231
    232
    233
    234
    def do_postprocessing(
        self, input_data: Any, output: torch.Tensor
    ) -> AudioInferenceResult:
        """
        Converts the model output to a numpy array, which then can be used to save the
        audio file.

        Args:
            input_data(Any): original input data
            output (torch.Tensor): output tensor from the model

        Returns:
            AudioInferenceResult: audio array
        """
        audio_array: numpy.ndarray[Any, Any] = output.cpu().numpy().squeeze()
        return AudioInferenceResult(audio_array=audio_array)

### `do_preprocessing(input_data)`

Preprocesses the input data.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`input_data` | `BarkWorkflowInput` | input data to be preprocessed | required

Returns:

Name | Type | Description
--- | --- | ---
`BatchEncoding` | `BatchEncoding` | batch encoding of the input data

Source code in `src/infernet_ml/workflows/inference/bark_hf_inference_workflow.py`

    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    def do_preprocessing(self, input_data: BarkWorkflowInput) -> BatchEncoding:
        """
        Preprocesses the input data.

        Args:
            input_data (BarkWorkflowInput): input data to be preprocessed

        Returns:
            BatchEncoding: batch encoding of the input data
        """
        text = input_data.prompt
        voice_preset = input_data.voice_preset or self.default_voice_preset
        return self.processor(text, voice_preset=voice_preset).to(self.device)

### `do_run_model(preprocessed_data)`

Run the model on the preprocessed data.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`preprocessed_data` | `BatchEncoding` | preprocessed data | required

Returns:

Type | Description
--- | ---
`Tensor` | torch.Tensor: output tensor from the model

Source code in `src/infernet_ml/workflows/inference/bark_hf_inference_workflow.py`

    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    def do_run_model(self, preprocessed_data: BatchEncoding) -> torch.Tensor:
        """
        Run the model on the preprocessed data.

        Args:
            preprocessed_data (BatchEncoding): preprocessed data

        Returns:
            torch.Tensor: output tensor from the model
        """
        return cast(torch.Tensor, self.model.generate(**preprocessed_data))

### `do_setup()`

Downloads the model from huggingface. Returns: bool: True on completion of loading model

Source code in `src/infernet_ml/workflows/inference/bark_hf_inference_workflow.py`

    155
    156
    157
    158
    159
    160
    161
    162
    def do_setup(self) -> None:
        """
        Downloads the model from huggingface.
        Returns:
            bool: True on completion of loading model
        """
        self.model = BarkModel.from_pretrained(self.model_name).to(self.device)
        self.processor = AutoProcessor.from_pretrained(self.model_name)

### `do_stream(preprocessed_input)`

Stream data for inference. Currently not implemented.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`preprocessed_input` | `Any` | preprocessed input data | required

Returns:

Type | Description
--- | ---
`Iterator[Any]` | Iterator[Any]: iterator for streaming data

Raises:

Type | Description
--- | ---
`NotImplementedError` | if the method is not implemented

Source code in `src/infernet_ml/workflows/inference/bark_hf_inference_workflow.py`

    204
    205
    206
    207
    208
    209
    210
    211
    212
    213
    214
    215
    216
    217
    def do_stream(self, preprocessed_input: Any) -> Iterator[Any]:
        """
        Stream data for inference. Currently not implemented.

        Args:
            preprocessed_input (Any): preprocessed input data

        Returns:
            Iterator[Any]: iterator for streaming data

        Raises:
            NotImplementedError: if the method is not implemented
        """
        raise NotImplementedError

### `inference(input_data, log_preprocessed_data=True)`

Override super class inference method to be annotated with the correct types.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`input_data` | `str` | prompt to generate audio from | required

Returns:

Name | Type | Description
--- | --- | ---
`AudioInferenceResult` | `AudioInferenceResult` | audio array

Source code in `src/infernet_ml/workflows/inference/bark_hf_inference_workflow.py`

    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    def inference(self, input_data: BarkWorkflowInput, log_preprocessed_data: bool = True) -> AudioInferenceResult:  # type: ignore #noqa: E501
        """
        Override super class inference method to be annotated with the correct types.

        Args:
            input_data (str): prompt to generate audio from

        Returns:
            AudioInferenceResult: audio array
        """
        return cast(
            AudioInferenceResult, super().inference(input_data, log_preprocessed_data)
        )

## `BarkProcessor`

Bases: `Protocol`

Type for the Suno Processor function. Used for type-safety.

Source code in `src/infernet_ml/workflows/inference/bark_hf_inference_workflow.py`

    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    class BarkProcessor(Protocol):
        """
        Type for the Suno Processor function. Used for type-safety.
        """

        def __call__(self, input_data: str, voice_preset: str) -> BatchEncoding:
            """
            Args:
                input_data (str): prompt to generate audio from
                voice_preset (str): voice to be used. There is a list of supported presets
                here: https://github.com/suno-ai/bark?tab=readme-ov-file#-voice-presets

            Returns:
                BatchEncoding: batch encoding of the input data
            """
            ...

### `__call__(input_data, voice_preset)`

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`input_data` | `str` | prompt to generate audio from | required
`voice_preset` | `str` | voice to be used. There is a list of supported presets | required
`here` |  |  | required

Returns:

Name | Type | Description
--- | --- | ---
`BatchEncoding` | `BatchEncoding` | batch encoding of the input data

Source code in `src/infernet_ml/workflows/inference/bark_hf_inference_workflow.py`

    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    def __call__(self, input_data: str, voice_preset: str) -> BatchEncoding:
        """
        Args:
            input_data (str): prompt to generate audio from
            voice_preset (str): voice to be used. There is a list of supported presets
            here: https://github.com/suno-ai/bark?tab=readme-ov-file#-voice-presets

        Returns:
            BatchEncoding: batch encoding of the input data
        """
        ...