Function xlm_roberta_large

diffsynth/models/wan_video_image_encoder.py:150–199

XLMRobertaLarge adapted from Huggingface.

(pretrained=False,
                      return_tokenizer=False,
                      device='cpu',
                      **kwargs)

def xlm_roberta_large(pretrained=False,
                      return_tokenizer=False,
                      device='cpu',
                      **kwargs):
    """
    XLMRobertaLarge adapted from Huggingface.
    """
    # params
    cfg = dict(
        vocab_size=250002,
        max_seq_len=514,
        type_size=1,
        pad_id=1,
        dim=1024,
        num_heads=16,
        num_layers=24,
        post_norm=True,
        dropout=0.1,
        eps=1e-5)
    cfg.update(**kwargs)

    # init model
    if pretrained:
        from sora import DOWNLOAD_TO_CACHE

        # init a meta model
        with torch.device('meta'):
            model = XLMRoberta(**cfg)

        # load checkpoint
        model.load_state_dict(
            torch.load(
                DOWNLOAD_TO_CACHE('models/xlm_roberta/xlm_roberta_large.pth'),
                map_location=device),
            assign=True)
    else:
        # init a model on device
        with torch.device(device):
            model = XLMRoberta(**cfg)

    # init tokenizer
    if return_tokenizer:
        from sora.data import HuggingfaceTokenizer
        tokenizer = HuggingfaceTokenizer(
            name='xlm-roberta-large',
            seq_len=model.text_len,
            clean='whitespace')
        return model, tokenizer
    else:
        return model