How to replace PL-BERT with XPhoneBERT? about styletts2 HOT 3 CLOSED

import os
from transformers import AutoConfig, AutoModelForMaskedLM


class CustomXPhoneBERT(AutoModelForMaskedLM):
    def forward(self, *args, **kwargs):
        # Call the original forward method
        outputs = super().forward(*args, **kwargs)

        # Only return the last_hidden_state
        return outputs.last_hidden_state


def load_xbert(model_name_or_path):
    # Load the configuration for 'xphonebert-base'
    config = AutoConfig.from_pretrained(model_name_or_path)

    # Initialize the custom XPhoneBERT model using the configuration
    xbert = CustomXPhoneBERT.from_pretrained(model_name_or_path, config=config)

    # Return the custom model
    return xbert

The inference won't be as compatible I guess, that's the current inference code which relies on the english-only bert:
    with torch.no_grad():
        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
        text_mask = length_to_mask(input_lengths).to(device)

        t_en = model.text_encoder(tokens, input_lengths, text_mask)
        bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())
        d_en = model.bert_encoder(bert_dur).transpose(-1, -2)

        s_pred = sampler(noise=torch.randn((1, 256)).unsqueeze(1).to(device),
                         embedding=bert_dur,
                         embedding_scale=embedding_scale,
                         features=ref_s,  # reference from the same speaker as the embedding
                         num_steps=diffusion_steps).squeeze(1)