Spaces:
Running
on
Zero
Running
on
Zero
| from comfy import sd1_clip | |
| from .spiece_tokenizer import SPieceTokenizer | |
| import comfy.text_encoders.t5 | |
| import os | |
| import re | |
| import torch | |
| import logging | |
| from tokenizers import Tokenizer | |
| from .ace_text_cleaners import multilingual_cleaners, japanese_to_romaji | |
| SUPPORT_LANGUAGES = { | |
| "en": 259, "de": 260, "fr": 262, "es": 284, "it": 285, | |
| "pt": 286, "pl": 294, "tr": 295, "ru": 267, "cs": 293, | |
| "nl": 297, "ar": 5022, "zh": 5023, "ja": 5412, "hu": 5753, | |
| "ko": 6152, "hi": 6680 | |
| } | |
| structure_pattern = re.compile(r"\[.*?\]") | |
| DEFAULT_VOCAB_FILE = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "ace_lyrics_tokenizer"), "vocab.json") | |
| class VoiceBpeTokenizer: | |
| def __init__(self, vocab_file=DEFAULT_VOCAB_FILE): | |
| self.tokenizer = None | |
| if vocab_file is not None: | |
| self.tokenizer = Tokenizer.from_file(vocab_file) | |
| def preprocess_text(self, txt, lang): | |
| txt = multilingual_cleaners(txt, lang) | |
| return txt | |
| def encode(self, txt, lang='en'): | |
| # lang = lang.split("-")[0] # remove the region | |
| # self.check_input_length(txt, lang) | |
| txt = self.preprocess_text(txt, lang) | |
| lang = "zh-cn" if lang == "zh" else lang | |
| txt = f"[{lang}]{txt}" | |
| txt = txt.replace(" ", "[SPACE]") | |
| return self.tokenizer.encode(txt).ids | |
| def get_lang(self, line): | |
| if line.startswith("[") and line[3:4] == ']': | |
| lang = line[1:3].lower() | |
| if lang in SUPPORT_LANGUAGES: | |
| return lang, line[4:] | |
| return "en", line | |
| def __call__(self, string): | |
| lines = string.split("\n") | |
| lyric_token_idx = [261] | |
| for line in lines: | |
| line = line.strip() | |
| if not line: | |
| lyric_token_idx += [2] | |
| continue | |
| lang, line = self.get_lang(line) | |
| if lang not in SUPPORT_LANGUAGES: | |
| lang = "en" | |
| if "zh" in lang: | |
| lang = "zh" | |
| if "spa" in lang: | |
| lang = "es" | |
| try: | |
| line_out = japanese_to_romaji(line) | |
| if line_out != line: | |
| lang = "ja" | |
| line = line_out | |
| except: | |
| pass | |
| try: | |
| if structure_pattern.match(line): | |
| token_idx = self.encode(line, "en") | |
| else: | |
| token_idx = self.encode(line, lang) | |
| lyric_token_idx = lyric_token_idx + token_idx + [2] | |
| except Exception as e: | |
| logging.warning("tokenize error {} for line {} major_language {}".format(e, line, lang)) | |
| return {"input_ids": lyric_token_idx} | |
| def from_pretrained(path, **kwargs): | |
| return VoiceBpeTokenizer(path, **kwargs) | |
| def get_vocab(self): | |
| return {} | |
| class UMT5BaseModel(sd1_clip.SDClipModel): | |
| def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, model_options={}): | |
| textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "umt5_config_base.json") | |
| super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"end": 1, "pad": 0}, model_class=comfy.text_encoders.t5.T5, enable_attention_masks=True, zero_out_masked=False, model_options=model_options) | |
| class UMT5BaseTokenizer(sd1_clip.SDTokenizer): | |
| def __init__(self, embedding_directory=None, tokenizer_data={}): | |
| tokenizer = tokenizer_data.get("spiece_model", None) | |
| super().__init__(tokenizer, pad_with_end=False, embedding_size=768, embedding_key='umt5base', tokenizer_class=SPieceTokenizer, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, pad_token=0, tokenizer_data=tokenizer_data) | |
| def state_dict(self): | |
| return {"spiece_model": self.tokenizer.serialize_model()} | |
| class LyricsTokenizer(sd1_clip.SDTokenizer): | |
| def __init__(self, embedding_directory=None, tokenizer_data={}): | |
| tokenizer = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "ace_lyrics_tokenizer"), "vocab.json") | |
| super().__init__(tokenizer, pad_with_end=False, embedding_size=1024, embedding_key='lyrics', tokenizer_class=VoiceBpeTokenizer, has_start_token=True, pad_to_max_length=False, max_length=99999999, min_length=1, pad_token=2, has_end_token=False, tokenizer_data=tokenizer_data) | |
| class AceT5Tokenizer: | |
| def __init__(self, embedding_directory=None, tokenizer_data={}): | |
| self.voicebpe = LyricsTokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data) | |
| self.umt5base = UMT5BaseTokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data) | |
| def tokenize_with_weights(self, text:str, return_word_ids=False, **kwargs): | |
| out = {} | |
| out["lyrics"] = self.voicebpe.tokenize_with_weights(kwargs.get("lyrics", ""), return_word_ids, **kwargs) | |
| out["umt5base"] = self.umt5base.tokenize_with_weights(text, return_word_ids, **kwargs) | |
| return out | |
| def untokenize(self, token_weight_pair): | |
| return self.umt5base.untokenize(token_weight_pair) | |
| def state_dict(self): | |
| return self.umt5base.state_dict() | |
| class AceT5Model(torch.nn.Module): | |
| def __init__(self, device="cpu", dtype=None, model_options={}, **kwargs): | |
| super().__init__() | |
| self.umt5base = UMT5BaseModel(device=device, dtype=dtype, model_options=model_options) | |
| self.dtypes = set() | |
| if dtype is not None: | |
| self.dtypes.add(dtype) | |
| def set_clip_options(self, options): | |
| self.umt5base.set_clip_options(options) | |
| def reset_clip_options(self): | |
| self.umt5base.reset_clip_options() | |
| def encode_token_weights(self, token_weight_pairs): | |
| token_weight_pairs_umt5base = token_weight_pairs["umt5base"] | |
| token_weight_pairs_lyrics = token_weight_pairs["lyrics"] | |
| t5_out, t5_pooled = self.umt5base.encode_token_weights(token_weight_pairs_umt5base) | |
| lyrics_embeds = torch.tensor(list(map(lambda a: a[0], token_weight_pairs_lyrics[0]))).unsqueeze(0) | |
| return t5_out, None, {"conditioning_lyrics": lyrics_embeds} | |
| def load_sd(self, sd): | |
| return self.umt5base.load_sd(sd) | |