init

.gitattributes

@@ -0,0 +1,38 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,215 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml
+
+__pycache__
+output
+cache
+third_party

README.md

@@ -0,0 +1,83 @@
+# [SongBloom]: *Coherent Song Generation via Interleaved Autoregressive Sketching and Diffusion Refinement*
+
+We propose **SongBloom**, a novel framework for full-length song generation that leverages an interleaved paradigm of autoregressive sketching and diffusion-based refinement. SongBloom employs an autoregressive diffusion model that combines the high fidelity of diffusion models with the scalability of language models.
+Specifically, it gradually extends a musical sketch from short to long and refines the details from coarse to fine-grained. The interleaved generation paradigm effectively integrates prior semantic and acoustic context to guide the generation process.
+Experimental results demonstrate that SongBloom outperforms existing methods across both subjective and objective metrics and achieves performance comparable to the state-of-the-art commercial music generation platforms.
+
+![img](docs/architecture.png)
+
+Demo page:  [https://cypress-yang.github.io/SongBloom_demo](https://cypress-yang.github.io/SongBloom_demo)
+
+ArXiv: [https://arxiv.org/abs/2506.07634](https://arxiv.org/abs/2506.07634)
+
+## Prepare Environments
+
+```bash
+conda create -n SongBloom python==3.8.12
+conda activate SongBloom
+
+# yum install libsndfile
+# pip install torch==2.2.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu118 # For different CUDA version
+pip install -r requirements.txt
+```
+
+## Data Preparation
+
+A  .jsonl file, where each line is a json object:
+
+```json
+{
+	"idx": "The index of each sample", 
+	"lyrics": "The lyrics to be generated",
+	"prompt_wav": "The path of the style prompt audio",
+}
+```
+
+One example can be refered to as: [example/test.jsonl](example/test.jsonl)
+
+The prompt wav should be a 10-second, 48kHz audio clip.
+
+The details about lyric format can be found in [docs/lyric_format.md](docs/lyric_format.md).
+
+## Inference
+
+```bash
+source set_env.sh
+
+python3 infer.py --input-jsonl example/test.jsonl
+
+# For GPUs with low VRAM like RTX4090, you should set the dtype as bfloat16
+python3 infer.py --input-jsonl example/test.jsonl --dtype bfloat16
+
+# SongBloom also supports flash-attn (optional). To enable it, please install flash-attn (v2.6.3 is used during training) manually and set os.environ['DISABLE_FLASH_ATTN'] = "0" in infer.py:8
+```
+
+## Models
+
+| Name                 | Size | Max Length | Prompt type | 🤗                                            |
+| -------------------- | ---- | ---------- | ----------- | -------------------------------------------- |
+| songbloom_full_150s  | 2B   | 2m30s      | 10s wav     | [link](https://huggingface.co/CypressYang/SongBloom) |
+| songbloom_mulan_150s | 2B   | 2m30s      | 10s wav / text description |           coming soon                           |
+| ... |      |            |             |                                              |
+
+
+
+## TODO List
+
+- [ ] Support Text Description
+- [ ] Full version
+
+## Citation
+
+```
+@article{yang2025songbloom,
+title={SongBloom: Coherent Song Generation via Interleaved Autoregressive Sketching and Diffusion Refinement},
+author={Yang, Chenyu and Wang, Shuai and Chen, Hangting and Tan, Wei and Yu, Jianwei and Li, Haizhou},
+journal={arXiv preprint arXiv:2506.07634},
+year={2025}
+}
+```
+
+## License
+
+SongBloom (codes and weights) is released under the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). 

SongBloom/g2p/cn_zh_g2p/__init__.py

@@ -0,0 +1,106 @@
+from . import chinese,  english # , japanese 暂时干掉看看
+from .symbols import *
+import yaml
+language_module_map = {"zh": chinese, "en": english} #, "ja": japanese
+
+def is_chinese(uchar):
+    if uchar >= u'\u4e00' and uchar <= u'\u9fa5':
+        return True
+    else:
+        return False
+
+import re
+
+# def split_text(text):
+#     chinese_pattern = r'[\u4e00-\u9fa5][\u4e00-\u9fa5\ \,\.\!\?\，\。]+'
+#     english_pattern = r'[a-zA-Z][a-zA-Z\'\ \,\.\!\?]+'
+    
+#     chinese_text = re.findall(chinese_pattern, text)
+#     print(chinese_text)
+#     english_text = re.findall(english_pattern, text)
+    
+#     return chinese_text, english_text
+
+def split_text(text):
+    pattern = re.compile("|".join(re.escape(p) for p in chinese.rep_map.keys()))
+    text = pattern.sub(lambda x: chinese.rep_map[x.group()], text)
+
+    result = []
+    lang = []
+    buffer = ""
+    chinese_pattern = r'[\u4e00-\u9fa5]'
+    special_pattern = r'[\,\.\!\?\…\-]'
+    # TODO check 一下
+    for char in text:
+        if re.match(special_pattern, char):
+            if buffer:
+                if not re.match(chinese_pattern, buffer[0]):
+                    result.append(buffer)
+                    lang.append('en')
+                else:
+                    result.append(buffer)
+                    lang.append("zh")
+            result.append(char)
+            lang.append('sp')
+            buffer = ""
+
+        
+        elif re.match(chinese_pattern, char):
+            if buffer and not re.match(chinese_pattern, buffer[-1]):
+                result.append(buffer)
+                buffer = ""
+                lang.append('en')
+            buffer += char
+        else:
+            if buffer and re.match(chinese_pattern, buffer[-1]):
+                result.append(buffer)
+                buffer = ""
+                lang.append("zh")
+            buffer += char
+
+    if buffer:
+        result.append(buffer)
+        lang.append("zh" if re.match(chinese_pattern, buffer[-1]) else 'en')
+
+    return result, lang
+
+def mixed_language_to_phoneme(text):
+    segments, lang = split_text(text)
+    # print(segments, lang)
+    result = [language_to_phoneme(s, l) for s, l in zip(segments, lang)]
+    phones, word2ph = [], []
+    for p, w, n in result:
+        phones += p
+        if w is None:
+            w = []
+        word2ph += w
+    return phones, word2ph
+
+
+def language_to_phoneme(text, language):
+    if language == 'sp':
+        return [text], None, text
+    language_module = language_module_map[language]
+    norm_text = language_module.text_normalize(text)
+    if language == "zh":
+        phones, word2ph = language_module.g2p(norm_text)
+        assert len(phones) == sum(word2ph)
+        assert len(norm_text) == len(word2ph)
+    else:
+        try:
+            phones = language_module.g2p(norm_text)
+        except:
+            phones = [norm_text]
+        word2ph = None
+
+    # for ph in phones:
+    #     assert ph in symbols, ph
+    return phones, word2ph, norm_text
+
+def gen_vocabs():
+    yaml.dump(symbols, open('./vocab.yaml', 'w'))
+
+class G2P_Mix():
+    def __call__(self, text):
+        phones, word2ph = mixed_language_to_phoneme(text)
+        return ' '.join(phones)

SongBloom/g2p/cn_zh_g2p/chinese.py

@@ -0,0 +1,173 @@
+import os
+import pdb
+import re
+
+import cn2an
+from pypinyin import lazy_pinyin, Style
+
+from .symbols import punctuation
+from .tone_sandhi import ToneSandhi
+from .zh_normalization.text_normlization import TextNormalizer
+
+normalizer = lambda x: cn2an.transform(x, "an2cn")
+
+current_file_path = os.path.dirname(__file__)
+pinyin_to_symbol_map = {
+    line.split("\t")[0]: line.strip().split("\t")[1]
+    for line in open(os.path.join(current_file_path, "opencpop-strict.txt")).readlines()
+}
+
+import jieba_fast.posseg as psg
+
+
+rep_map = {
+    "：": ",",
+    "；": ",",
+    "，": ",",
+    "。": ".",
+    "！": "!",
+    "？": "?",
+    "\n": ".",
+    "·": ",",
+    "、": ",",
+    "...": "…",
+    "$": ".",
+    "/": ",",
+    "—": "-",
+    "~": "…",
+    "～":"…",
+}
+
+tone_modifier = ToneSandhi()
+
+
+def replace_punctuation(text):
+    text = text.replace("嗯", "恩").replace("呣", "母")
+    pattern = re.compile("|".join(re.escape(p) for p in rep_map.keys()))
+
+    replaced_text = pattern.sub(lambda x: rep_map[x.group()], text)
+    replaced_text = re.sub(
+        r"[^\u4e00-\u9fa5" + "".join(punctuation) + r"]+", "", replaced_text
+    )
+
+    return replaced_text
+
+
+def g2p(text):
+    pattern = r"(?<=[{0}])\s*".format("".join(punctuation))
+    sentences = [i for i in re.split(pattern, text) if i.strip() != ""]
+    phones, word2ph = _g2p(sentences)
+    return phones, word2ph
+
+
+def _get_initials_finals(word):
+    initials = []
+    finals = []
+    orig_initials = lazy_pinyin(word, neutral_tone_with_five=True, style=Style.INITIALS)
+    orig_finals = lazy_pinyin(
+        word, neutral_tone_with_five=True, style=Style.FINALS_TONE3
+    )
+    for c, v in zip(orig_initials, orig_finals):
+        initials.append(c)
+        finals.append(v)
+    return initials, finals
+
+
+def _g2p(segments):
+    phones_list = []
+    word2ph = []
+    for seg in segments:
+        pinyins = []
+        # Replace all English words in the sentence
+        seg = re.sub("[a-zA-Z]+", "", seg)
+        seg_cut = psg.lcut(seg)
+        initials = []
+        finals = []
+        seg_cut = tone_modifier.pre_merge_for_modify(seg_cut)
+        for word, pos in seg_cut:
+            if pos == "eng":
+                continue
+            sub_initials, sub_finals = _get_initials_finals(word)
+            sub_finals = tone_modifier.modified_tone(word, pos, sub_finals)
+            initials.append(sub_initials)
+            finals.append(sub_finals)
+
+            # assert len(sub_initials) == len(sub_finals) == len(word)
+        initials = sum(initials, [])
+        finals = sum(finals, [])
+        #
+        for c, v in zip(initials, finals):
+            raw_pinyin = c + v
+            # NOTE: post process for pypinyin outputs
+            # we discriminate i, ii and iii
+            if c == v:
+                assert c in punctuation
+                phone = [c]
+                word2ph.append(1)
+            else:
+                v_without_tone = v[:-1]
+                tone = v[-1]
+
+                pinyin = c + v_without_tone
+                assert tone in "12345"
+
+                if c:
+                    # 多音节
+                    v_rep_map = {
+                        "uei": "ui",
+                        "iou": "iu",
+                        "uen": "un",
+                    }
+                    if v_without_tone in v_rep_map.keys():
+                        pinyin = c + v_rep_map[v_without_tone]
+                else:
+                    # 单音节
+                    pinyin_rep_map = {
+                        "ing": "ying",
+                        "i": "yi",
+                        "in": "yin",
+                        "u": "wu",
+                    }
+                    if pinyin in pinyin_rep_map.keys():
+                        pinyin = pinyin_rep_map[pinyin]
+                    else:
+                        single_rep_map = {
+                            "v": "yu",
+                            "e": "e",
+                            "i": "y",
+                            "u": "w",
+                        }
+                        if pinyin[0] in single_rep_map.keys():
+                            pinyin = single_rep_map[pinyin[0]] + pinyin[1:]
+
+                assert pinyin in pinyin_to_symbol_map.keys(), (pinyin, seg, raw_pinyin)
+                new_c, new_v = pinyin_to_symbol_map[pinyin].split(" ")
+                new_v = new_v + tone
+                phone = [new_c, new_v]
+                word2ph.append(len(phone))
+
+            phones_list += phone
+    return phones_list, word2ph
+
+
+def text_normalize(text):
+    # https://github.com/PaddlePaddle/PaddleSpeech/tree/develop/paddlespeech/t2s/frontend/zh_normalization
+    tx = TextNormalizer()
+    sentences = tx.normalize(text)
+    dest_text = ""
+    for sentence in sentences:
+        dest_text += replace_punctuation(sentence)
+    return dest_text
+
+
+if __name__ == "__main__":
+    text = "啊——但是《原神》是由,米哈\游自主，研发的一款全.新开放世界.冒险游戏"
+    text = "呣呣呣～就是…大人的鼹鼠党吧？"
+    text = "你好"
+    text = text_normalize(text)
+    print(g2p(text))
+
+
+# # 示例用法
+# text = "这是一个示例文本：,你好！这是一个测试..."
+# print(g2p_paddle(text))  # 输出: 这是一个示例文本你好这是一个测试

SongBloom/g2p/cn_zh_g2p/engdict-hot.rep

@@ -0,0 +1,2 @@
+CHATGPT CH AE1 T JH IY1 P IY1 T IY1
+JSON JH EY1 S AH0 N

SongBloom/g2p/cn_zh_g2p/engdict_cache.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9bff9393f4b192d873a11335efc8f124771087b6dc847d34fd240c2846889d2b
+size 5965909

SongBloom/g2p/cn_zh_g2p/english.py

@@ -0,0 +1,369 @@
+import pickle
+import os
+import re
+import wordsegment
+from g2p_en import G2p
+
+from string import punctuation
+
+from .symbols import symbols
+
+import unicodedata
+from builtins import str as unicode
+from g2p_en.expand import normalize_numbers
+
+# Set NLTK data path programmatically to avoid needing set_env.sh
+import nltk
+current_file_path = os.path.dirname(__file__)
+nltk_data_path = os.path.join(current_file_path, "nltk_data")
+if os.path.exists(nltk_data_path):
+    nltk.data.path.insert(0, nltk_data_path)
+
+from nltk.tokenize import TweetTokenizer
+word_tokenize = TweetTokenizer().tokenize
+from nltk import pos_tag
+
+CMU_DICT_PATH = os.path.join(current_file_path, "cmudict.rep")
+CMU_DICT_FAST_PATH = os.path.join(current_file_path, "cmudict-fast.rep")
+CMU_DICT_HOT_PATH = os.path.join(current_file_path, "engdict-hot.rep")
+CACHE_PATH = os.path.join(current_file_path, "engdict_cache.pickle")
+NAMECACHE_PATH = os.path.join(current_file_path, "namedict_cache.pickle")
+
+arpa = {
+    "AH0",
+    "S",
+    "AH1",
+    "EY2",
+    "AE2",
+    "EH0",
+    "OW2",
+    "UH0",
+    "NG",
+    "B",
+    "G",
+    "AY0",
+    "M",
+    "AA0",
+    "F",
+    "AO0",
+    "ER2",
+    "UH1",
+    "IY1",
+    "AH2",
+    "DH",
+    "IY0",
+    "EY1",
+    "IH0",
+    "K",
+    "N",
+    "W",
+    "IY2",
+    "T",
+    "AA1",
+    "ER1",
+    "EH2",
+    "OY0",
+    "UH2",
+    "UW1",
+    "Z",
+    "AW2",
+    "AW1",
+    "V",
+    "UW2",
+    "AA2",
+    "ER",
+    "AW0",
+    "UW0",
+    "R",
+    "OW1",
+    "EH1",
+    "ZH",
+    "AE0",
+    "IH2",
+    "IH",
+    "Y",
+    "JH",
+    "P",
+    "AY1",
+    "EY0",
+    "OY2",
+    "TH",
+    "HH",
+    "D",
+    "ER0",
+    "CH",
+    "AO1",
+    "AE1",
+    "AO2",
+    "OY1",
+    "AY2",
+    "IH1",
+    "OW0",
+    "L",
+    "SH",
+}
+
+
+def replace_phs(phs):
+    rep_map = {"'": "-"}
+    phs_new = []
+    for ph in phs:
+        if ph in symbols:
+            phs_new.append(ph)
+        elif ph in rep_map.keys():
+            phs_new.append(rep_map[ph])
+        else:
+            print("ph not in symbols: ", ph)
+    return phs_new
+
+
+def read_dict():
+    g2p_dict = {}
+    start_line = 49
+    with open(CMU_DICT_PATH) as f:
+        line = f.readline()
+        line_index = 1
+        while line:
+            if line_index >= start_line:
+                line = line.strip()
+                word_split = line.split("  ")
+                word = word_split[0].lower()
+
+                syllable_split = word_split[1].split(" - ")
+                g2p_dict[word] = []
+                for syllable in syllable_split:
+                    phone_split = syllable.split(" ")
+                    g2p_dict[word].append(phone_split)
+
+            line_index = line_index + 1
+            line = f.readline()
+
+    return g2p_dict
+
+
+def read_dict_new():
+    g2p_dict = {}
+    with open(CMU_DICT_PATH) as f:
+        line = f.readline()
+        line_index = 1
+        while line:
+            if line_index >= 57:
+                line = line.strip()
+                word_split = line.split("  ")
+                word = word_split[0].lower()
+                g2p_dict[word] = [word_split[1].split(" ")]
+
+            line_index = line_index + 1
+            line = f.readline()
+
+    with open(CMU_DICT_FAST_PATH) as f:
+        line = f.readline()
+        line_index = 1
+        while line:
+            if line_index >= 0:
+                line = line.strip()
+                word_split = line.split(" ")
+                word = word_split[0].lower()
+                if word not in g2p_dict:
+                    g2p_dict[word] = [word_split[1:]]
+
+            line_index = line_index + 1
+            line = f.readline()
+
+    return g2p_dict
+
+def hot_reload_hot(g2p_dict):
+    with open(CMU_DICT_HOT_PATH) as f:
+        line = f.readline()
+        line_index = 1
+        while line:
+            if line_index >= 0:
+                line = line.strip()
+                word_split = line.split(" ")
+                word = word_split[0].lower()
+                # 自定义发音词直接覆盖字典
+                g2p_dict[word] = [word_split[1:]]
+
+            line_index = line_index + 1
+            line = f.readline()
+
+    return g2p_dict
+
+
+def cache_dict(g2p_dict, file_path):
+    with open(file_path, "wb") as pickle_file:
+        pickle.dump(g2p_dict, pickle_file)
+
+
+def get_dict():
+    if os.path.exists(CACHE_PATH):
+        with open(CACHE_PATH, "rb") as pickle_file:
+            g2p_dict = pickle.load(pickle_file)
+    else:
+        g2p_dict = read_dict_new()
+        cache_dict(g2p_dict, CACHE_PATH)
+
+    g2p_dict = hot_reload_hot(g2p_dict)
+
+    return g2p_dict
+
+
+def get_namedict():
+    if os.path.exists(NAMECACHE_PATH):
+        with open(NAMECACHE_PATH, "rb") as pickle_file:
+            name_dict = pickle.load(pickle_file)
+    else:
+        name_dict = {}
+
+    return name_dict
+
+
+def text_normalize(text):
+    # todo: eng text normalize
+    # 适配中文及 g2p_en 标点
+    rep_map = {
+        "[;:：，；]": ",",
+        '["’]': "'",
+        "。": ".",
+        "！": "!",
+        "？": "?",
+    }
+    for p, r in rep_map.items():
+        text = re.sub(p, r, text)
+
+    # 来自 g2p_en 文本格式化处理
+    # 增加大写兼容
+    text = unicode(text)
+    text = normalize_numbers(text)
+    text = ''.join(char for char in unicodedata.normalize('NFD', text)
+                    if unicodedata.category(char) != 'Mn')  # Strip accents
+    text = re.sub("[^ A-Za-z'.,?!\-]", "", text)
+    text = re.sub(r"(?i)i\.e\.", "that is", text)
+    text = re.sub(r"(?i)e\.g\.", "for example", text)
+
+    return text
+
+
+class en_G2p(G2p):
+    def __init__(self):
+        super().__init__()
+        # 分词初始化
+        wordsegment.load()
+
+        # 扩展过时字典, 添加姓名字典
+        self.cmu = get_dict()
+        self.namedict = get_namedict()
+
+        # 剔除读音错误的几个缩写
+        for word in ["AE", "AI", "AR", "IOS", "HUD", "OS"]:
+            del self.cmu[word.lower()]
+
+        # 修正多音字
+        self.homograph2features["read"] = (['R', 'IY1', 'D'], ['R', 'EH1', 'D'], 'VBP')
+        self.homograph2features["complex"] = (['K', 'AH0', 'M', 'P', 'L', 'EH1', 'K', 'S'], ['K', 'AA1', 'M', 'P', 'L', 'EH0', 'K', 'S'], 'JJ')
+
+
+    def __call__(self, text):
+        # tokenization
+        words = word_tokenize(text)
+        tokens = pos_tag(words)  # tuples of (word, tag)
+
+        # steps
+        prons = []
+        for o_word, pos in tokens:
+            # 还原 g2p_en 小写操作逻辑
+            word = o_word.lower()
+
+            if re.search("[a-z]", word) is None:
+                pron = [word]
+            # 先把单字母推出去
+            elif len(word) == 1:
+                # 单读 A 发音修正, 这里需要原格式 o_word 判断大写
+                if o_word == "A":
+                    pron = ['EY1']
+                else:
+                    pron = self.cmu[word][0]
+            # g2p_en 原版多音字处理
+            elif word in self.homograph2features:  # Check homograph
+                pron1, pron2, pos1 = self.homograph2features[word]
+                if pos.startswith(pos1):
+                    pron = pron1
+                # pos1比pos长仅出现在read
+                elif len(pos) < len(pos1) and pos == pos1[:len(pos)]:
+                    pron = pron1
+                else:
+                    pron = pron2
+            else:
+                # 递归查找预测
+                pron = self.qryword(o_word)
+
+            prons.extend(pron)
+            prons.extend([" "])
+
+        return prons[:-1]
+
+
+    def qryword(self, o_word):
+        word = o_word.lower()
+
+        # 查字典, 单字母除外
+        if len(word) > 1 and word in self.cmu:  # lookup CMU dict
+            return self.cmu[word][0]
+
+        # 单词仅首字母大写时查找姓名字典
+        if o_word.istitle() and word in self.namedict:
+            return self.namedict[word][0]
+
+        # oov 长度小于等于 3 直接读字母
+        if len(word) <= 3:
+            phones = []
+            for w in word:
+                # 单读 A 发音修正, 此处不存在大写的情况
+                if w == "a":
+                    phones.extend(['EY1'])
+                else:
+                    phones.extend(self.cmu[w][0])
+            return phones
+
+        # 尝试分离所有格
+        if re.match(r"^([a-z]+)('s)$", word):
+            phones = self.qryword(word[:-2])[:]
+            # P T K F TH HH 无声辅音结尾 's 发 ['S']
+            if phones[-1] in ['P', 'T', 'K', 'F', 'TH', 'HH']:
+                phones.extend(['S'])
+            # S Z SH ZH CH JH 擦声结尾 's 发 ['IH1', 'Z'] 或 ['AH0', 'Z']
+            elif phones[-1] in ['S', 'Z', 'SH', 'ZH', 'CH', 'JH']:
+                phones.extend(['AH0', 'Z'])
+            # B D G DH V M N NG L R W Y 有声辅音结尾 's 发 ['Z']
+            # AH0 AH1 AH2 EY0 EY1 EY2 AE0 AE1 AE2 EH0 EH1 EH2 OW0 OW1 OW2 UH0 UH1 UH2 IY0 IY1 IY2 AA0 AA1 AA2 AO0 AO1 AO2
+            # ER ER0 ER1 ER2 UW0 UW1 UW2 AY0 AY1 AY2 AW0 AW1 AW2 OY0 OY1 OY2 IH IH0 IH1 IH2 元音结尾 's 发 ['Z']
+            else:
+                phones.extend(['Z'])
+            return phones
+
+        # 尝试进行分词，应对复合词
+        comps = wordsegment.segment(word.lower())
+
+        # 无法分词的送回去预测
+        if len(comps)==1:
+            return self.predict(word)
+
+        # 可以分词的递归处理
+        return [phone for comp in comps for phone in self.qryword(comp)]
+
+
+_g2p = en_G2p()
+
+
+def g2p(text):
+    # g2p_en 整段推理，剔除不存在的arpa返回
+    phone_list = _g2p(text)
+    phones = [ph if ph != "<unk>" else "UNK" for ph in phone_list if ph not in [" ", "<pad>", "UW", "</s>", "<s>"]]
+
+    return replace_phs(phones)
+
+
+if __name__ == "__main__":
+    print(g2p("hello"))
+    print(g2p(text_normalize("e.g. I used openai's AI tool to draw a picture.")))
+    print(g2p(text_normalize("In this; paper, we propose 1 DSPGAN, a GAN-based universal vocoder.")))

SongBloom/g2p/cn_zh_g2p/nltk_data/corpora/cmudict.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d07cca47fd72ad32ea9d8ad1219f85301eeaf4568f8b6b73747506a71fb5afd6
+size 896069

SongBloom/g2p/cn_zh_g2p/nltk_data/corpora/cmudict/README

@@ -0,0 +1,76 @@
+The Carnegie Mellon Pronouncing Dictionary [cmudict.0.7a]
+
+ftp://ftp.cs.cmu.edu/project/speech/dict/
+https://cmusphinx.svn.sourceforge.net/svnroot/cmusphinx/trunk/cmudict/cmudict.0.7a
+
+Copyright (C) 1993-2008 Carnegie Mellon University. All rights reserved.
+
+File Format: Each line consists of an uppercased word,
+a counter (for alternative pronunciations), and a transcription.
+Vowels are marked for stress (1=primary, 2=secondary, 0=no stress).
+E.g.: NATURAL 1 N AE1 CH ER0 AH0 L
+
+The dictionary contains 127069 entries.  Of these, 119400 words are assigned
+a unique pronunciation, 6830 words have two pronunciations, and 839 words have
+three or more pronunciations.  Many of these are fast-speech variants.
+
+Phonemes: There are 39 phonemes, as shown below:
+    
+    Phoneme Example Translation    Phoneme Example Translation
+    ------- ------- -----------    ------- ------- -----------
+    AA      odd     AA D           AE      at      AE T
+    AH      hut     HH AH T        AO      ought   AO T
+    AW      cow     K AW           AY      hide    HH AY D
+    B       be      B IY           CH      cheese  CH IY Z
+    D       dee     D IY           DH      thee    DH IY
+    EH      Ed      EH D           ER      hurt    HH ER T
+    EY      ate     EY T           F       fee     F IY
+    G       green   G R IY N       HH      he      HH IY
+    IH      it      IH T           IY      eat     IY T
+    JH      gee     JH IY          K       key     K IY
+    L       lee     L IY           M       me      M IY
+    N       knee    N IY           NG      ping    P IH NG
+    OW      oat     OW T           OY      toy     T OY
+    P       pee     P IY           R       read    R IY D
+    S       sea     S IY           SH      she     SH IY
+    T       tea     T IY           TH      theta   TH EY T AH
+    UH      hood    HH UH D        UW      two     T UW
+    V       vee     V IY           W       we      W IY
+    Y       yield   Y IY L D       Z       zee     Z IY
+    ZH      seizure S IY ZH ER
+
+(For NLTK, entries have been sorted so that, e.g. FIRE 1 and FIRE 2
+are contiguous, and not separated by FIRE'S 1.)
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+
+1. Redistributions of source code must retain the above copyright
+   notice, this list of conditions and the following disclaimer.
+   The contents of this file are deemed to be source code.
+
+2. Redistributions in binary form must reproduce the above copyright
+   notice, this list of conditions and the following disclaimer in
+   the documentation and/or other materials provided with the
+   distribution.
+
+This work was supported in part by funding from the Defense Advanced
+Research Projects Agency, the Office of Naval Research and the National
+Science Foundation of the United States of America, and by member
+companies of the Carnegie Mellon Sphinx Speech Consortium. We acknowledge
+the contributions of many volunteers to the expansion and improvement of
+this dictionary.
+
+THIS SOFTWARE IS PROVIDED BY CARNEGIE MELLON UNIVERSITY ``AS IS'' AND
+ANY EXPRESSED OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL CARNEGIE MELLON UNIVERSITY
+NOR ITS EMPLOYEES BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+

SongBloom/g2p/cn_zh_g2p/nltk_data/taggers/averaged_perceptron_tagger.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1f13cf2532daadfd6f3bc481a49859f0b8ea6432ccdcd83e6a49a5f19008de9
+size 2526731

SongBloom/g2p/cn_zh_g2p/nltk_data/taggers/averaged_perceptron_tagger/averaged_perceptron_tagger.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25a5a19c7ced7b2bac3831da5bc0afcc2c34e5dd01cd4f361bb799949a696238
+size 6138625

SongBloom/g2p/cn_zh_g2p/opencpop-strict.txt

@@ -0,0 +1,429 @@
+a	AA a
+ai	AA ai
+an	AA an
+ang	AA ang
+ao	AA ao
+ba	b a
+bai	b ai
+ban	b an
+bang	b ang
+bao	b ao
+bei	b ei
+ben	b en
+beng	b eng
+bi	b i
+bian	b ian
+biao	b iao
+bie	b ie
+bin	b in
+bing	b ing
+bo	b o
+bu	b u
+ca	c a
+cai	c ai
+can	c an
+cang	c ang
+cao	c ao
+ce	c e
+cei	c ei
+cen	c en
+ceng	c eng
+cha	ch a
+chai	ch ai
+chan	ch an
+chang	ch ang
+chao	ch ao
+che	ch e
+chen	ch en
+cheng	ch eng
+chi	ch ir
+chong	ch ong
+chou	ch ou
+chu	ch u
+chua	ch ua
+chuai	ch uai
+chuan	ch uan
+chuang	ch uang
+chui	ch ui
+chun	ch un
+chuo	ch uo
+ci	c i0
+cong	c ong
+cou	c ou
+cu	c u
+cuan	c uan
+cui	c ui
+cun	c un
+cuo	c uo
+da	d a
+dai	d ai
+dan	d an
+dang	d ang
+dao	d ao
+de	d e
+dei	d ei
+den	d en
+deng	d eng
+di	d i
+dia	d ia
+dian	d ian
+diao	d iao
+die	d ie
+ding	d ing
+diu	d iu
+dong	d ong
+dou	d ou
+du	d u
+duan	d uan
+dui	d ui
+dun	d un
+duo	d uo
+e	EE e
+ei	EE ei
+en	EE en
+eng	EE eng
+er	EE er
+fa	f a
+fan	f an
+fang	f ang
+fei	f ei
+fen	f en
+feng	f eng
+fo	f o
+fou	f ou
+fu	f u
+ga	g a
+gai	g ai
+gan	g an
+gang	g ang
+gao	g ao
+ge	g e
+gei	g ei
+gen	g en
+geng	g eng
+gong	g ong
+gou	g ou
+gu	g u
+gua	g ua
+guai	g uai
+guan	g uan
+guang	g uang
+gui	g ui
+gun	g un
+guo	g uo
+ha	h a
+hai	h ai
+han	h an
+hang	h ang
+hao	h ao
+he	h e
+hei	h ei
+hen	h en
+heng	h eng
+hong	h ong
+hou	h ou
+hu	h u
+hua	h ua
+huai	h uai
+huan	h uan
+huang	h uang
+hui	h ui
+hun	h un
+huo	h uo
+ji	j i
+jia	j ia
+jian	j ian
+jiang	j iang
+jiao	j iao
+jie	j ie
+jin	j in
+jing	j ing
+jiong	j iong
+jiu	j iu
+ju	j v
+jv	j v
+juan	j van
+jvan	j van
+jue	j ve
+jve	j ve
+jun	j vn
+jvn	j vn
+ka	k a
+kai	k ai
+kan	k an
+kang	k ang
+kao	k ao
+ke	k e
+kei	k ei
+ken	k en
+keng	k eng
+kong	k ong
+kou	k ou
+ku	k u
+kua	k ua
+kuai	k uai
+kuan	k uan
+kuang	k uang
+kui	k ui
+kun	k un
+kuo	k uo
+la	l a
+lai	l ai
+lan	l an
+lang	l ang
+lao	l ao
+le	l e
+lei	l ei
+leng	l eng
+li	l i
+lia	l ia
+lian	l ian
+liang	l iang
+liao	l iao
+lie	l ie
+lin	l in
+ling	l ing
+liu	l iu
+lo	l o
+long	l ong
+lou	l ou
+lu	l u
+luan	l uan
+lun	l un
+luo	l uo
+lv	l v
+lve	l ve
+ma	m a
+mai	m ai
+man	m an
+mang	m ang
+mao	m ao
+me	m e
+mei	m ei
+men	m en
+meng	m eng
+mi	m i
+mian	m ian
+miao	m iao
+mie	m ie
+min	m in
+ming	m ing
+miu	m iu
+mo	m o
+mou	m ou
+mu	m u
+na	n a
+nai	n ai
+nan	n an
+nang	n ang
+nao	n ao
+ne	n e
+nei	n ei
+nen	n en
+neng	n eng
+ni	n i
+nian	n ian
+niang	n iang
+niao	n iao
+nie	n ie
+nin	n in
+ning	n ing
+niu	n iu
+nong	n ong
+nou	n ou
+nu	n u
+nuan	n uan
+nun	n un
+nuo	n uo
+nv	n v
+nve	n ve
+o	OO o
+ou	OO ou
+pa	p a
+pai	p ai
+pan	p an
+pang	p ang
+pao	p ao
+pei	p ei
+pen	p en
+peng	p eng
+pi	p i
+pian	p ian
+piao	p iao
+pie	p ie
+pin	p in
+ping	p ing
+po	p o
+pou	p ou
+pu	p u
+qi	q i
+qia	q ia
+qian	q ian
+qiang	q iang
+qiao	q iao
+qie	q ie
+qin	q in
+qing	q ing
+qiong	q iong
+qiu	q iu
+qu	q v
+qv	q v
+quan	q van
+qvan	q van
+que	q ve
+qve	q ve
+qun	q vn
+qvn	q vn
+ran	r an
+rang	r ang
+rao	r ao
+re	r e
+ren	r en
+reng	r eng
+ri	r ir
+rong	r ong
+rou	r ou
+ru	r u
+rua	r ua
+ruan	r uan
+rui	r ui
+run	r un
+ruo	r uo
+sa	s a
+sai	s ai
+san	s an
+sang	s ang
+sao	s ao
+se	s e
+sen	s en
+seng	s eng
+sha	sh a
+shai	sh ai
+shan	sh an
+shang	sh ang
+shao	sh ao
+she	sh e
+shei	sh ei
+shen	sh en
+sheng	sh eng
+shi	sh ir
+shou	sh ou
+shu	sh u
+shua	sh ua
+shuai	sh uai
+shuan	sh uan
+shuang	sh uang
+shui	sh ui
+shun	sh un
+shuo	sh uo
+si	s i0
+song	s ong
+sou	s ou
+su	s u
+suan	s uan
+sui	s ui
+sun	s un
+suo	s uo
+ta	t a
+tai	t ai
+tan	t an
+tang	t ang
+tao	t ao
+te	t e
+tei	t ei
+teng	t eng
+ti	t i
+tian	t ian
+tiao	t iao
+tie	t ie
+ting	t ing
+tong	t ong
+tou	t ou
+tu	t u
+tuan	t uan
+tui	t ui
+tun	t un
+tuo	t uo
+wa	w a
+wai	w ai
+wan	w an
+wang	w ang
+wei	w ei
+wen	w en
+weng	w eng
+wo	w o
+wu	w u
+xi	x i
+xia	x ia
+xian	x ian
+xiang	x iang
+xiao	x iao
+xie	x ie
+xin	x in
+xing	x ing
+xiong	x iong
+xiu	x iu
+xu	x v
+xv	x v
+xuan	x van
+xvan	x van
+xue	x ve
+xve	x ve
+xun	x vn
+xvn	x vn
+ya	y a
+yan	y En
+yang	y ang
+yao	y ao
+ye	y E
+yi	y i
+yin	y in
+ying	y ing
+yo	y o
+yong	y ong
+you	y ou
+yu	y v
+yv	y v
+yuan	y van
+yvan	y van
+yue	y ve
+yve	y ve
+yun	y vn
+yvn	y vn
+za	z a
+zai	z ai
+zan	z an
+zang	z ang
+zao	z ao
+ze	z e
+zei	z ei
+zen	z en
+zeng	z eng
+zha	zh a
+zhai	zh ai
+zhan	zh an
+zhang	zh ang
+zhao	zh ao
+zhe	zh e
+zhei	zh ei
+zhen	zh en
+zheng	zh eng
+zhi	zh ir
+zhong	zh ong
+zhou	zh ou
+zhu	zh u
+zhua	zh ua
+zhuai	zh uai
+zhuan	zh uan
+zhuang	zh uang
+zhui	zh ui
+zhun	zh un
+zhuo	zh uo
+zi	z i0
+zong	z ong
+zou	z ou
+zu	z u
+zuan	z uan
+zui	z ui
+zun	z un
+zuo	z uo

SongBloom/g2p/cn_zh_g2p/symbols.py

@@ -0,0 +1,401 @@
+import os
+
+# punctuation = ['!', '?', '…', ",", ".","@"]#@是SP停顿
+punctuation = ["!", "?", "…", ",", "."]  # @是SP停顿
+punctuation.append("-")
+pu_symbols = punctuation + ["SP", "SP2", "SP3", "UNK"]
+# pu_symbols = punctuation + ["SP", 'SP2', 'SP3','SP4', "UNK"]
+pad = "_"
+
+c = [
+    "AA",
+    "EE",
+    "OO",
+    "b",
+    "c",
+    "ch",
+    "d",
+    "f",
+    "g",
+    "h",
+    "j",
+    "k",
+    "l",
+    "m",
+    "n",
+    "p",
+    "q",
+    "r",
+    "s",
+    "sh",
+    "t",
+    "w",
+    "x",
+    "y",
+    "z",
+    "zh",
+]
+v = [
+    "E1",
+    "En1",
+    "a1",
+    "ai1",
+    "an1",
+    "ang1",
+    "ao1",
+    "e1",
+    "ei1",
+    "en1",
+    "eng1",
+    "er1",
+    "i1",
+    "i01",
+    "ia1",
+    "ian1",
+    "iang1",
+    "iao1",
+    "ie1",
+    "in1",
+    "ing1",
+    "iong1",
+    "ir1",
+    "iu1",
+    "o1",
+    "ong1",
+    "ou1",
+    "u1",
+    "ua1",
+    "uai1",
+    "uan1",
+    "uang1",
+    "ui1",
+    "un1",
+    "uo1",
+    "v1",
+    "van1",
+    "ve1",
+    "vn1",
+    "E2",
+    "En2",
+    "a2",
+    "ai2",
+    "an2",
+    "ang2",
+    "ao2",
+    "e2",
+    "ei2",
+    "en2",
+    "eng2",
+    "er2",
+    "i2",
+    "i02",
+    "ia2",
+    "ian2",
+    "iang2",
+    "iao2",
+    "ie2",
+    "in2",
+    "ing2",
+    "iong2",
+    "ir2",
+    "iu2",
+    "o2",
+    "ong2",
+    "ou2",
+    "u2",
+    "ua2",
+    "uai2",
+    "uan2",
+    "uang2",
+    "ui2",
+    "un2",
+    "uo2",
+    "v2",
+    "van2",
+    "ve2",
+    "vn2",
+    "E3",
+    "En3",
+    "a3",
+    "ai3",
+    "an3",
+    "ang3",
+    "ao3",
+    "e3",
+    "ei3",
+    "en3",
+    "eng3",
+    "er3",
+    "i3",
+    "i03",
+    "ia3",
+    "ian3",
+    "iang3",
+    "iao3",
+    "ie3",
+    "in3",
+    "ing3",
+    "iong3",
+    "ir3",
+    "iu3",
+    "o3",
+    "ong3",
+    "ou3",
+    "u3",
+    "ua3",
+    "uai3",
+    "uan3",
+    "uang3",
+    "ui3",
+    "un3",
+    "uo3",
+    "v3",
+    "van3",
+    "ve3",
+    "vn3",
+    "E4",
+    "En4",
+    "a4",
+    "ai4",
+    "an4",
+    "ang4",
+    "ao4",
+    "e4",
+    "ei4",
+    "en4",
+    "eng4",
+    "er4",
+    "i4",
+    "i04",
+    "ia4",
+    "ian4",
+    "iang4",
+    "iao4",
+    "ie4",
+    "in4",
+    "ing4",
+    "iong4",
+    "ir4",
+    "iu4",
+    "o4",
+    "ong4",
+    "ou4",
+    "u4",
+    "ua4",
+    "uai4",
+    "uan4",
+    "uang4",
+    "ui4",
+    "un4",
+    "uo4",
+    "v4",
+    "van4",
+    "ve4",
+    "vn4",
+    "E5",
+    "En5",
+    "a5",
+    "ai5",
+    "an5",
+    "ang5",
+    "ao5",
+    "e5",
+    "ei5",
+    "en5",
+    "eng5",
+    "er5",
+    "i5",
+    "i05",
+    "ia5",
+    "ian5",
+    "iang5",
+    "iao5",
+    "ie5",
+    "in5",
+    "ing5",
+    "iong5",
+    "ir5",
+    "iu5",
+    "o5",
+    "ong5",
+    "ou5",
+    "u5",
+    "ua5",
+    "uai5",
+    "uan5",
+    "uang5",
+    "ui5",
+    "un5",
+    "uo5",
+    "v5",
+    "van5",
+    "ve5",
+    "vn5",
+]
+
+v_without_tone = [
+    "E",
+    "En",
+    "a",
+    "ai",
+    "an",
+    "ang",
+    "ao",
+    "e",
+    "ei",
+    "en",
+    "eng",
+    "er",
+    "i",
+    "i0",
+    "ia",
+    "ian",
+    "iang",
+    "iao",
+    "ie",
+    "in",
+    "ing",
+    "iong",
+    "ir",
+    "iu",
+    "o",
+    "ong",
+    "ou",
+    "u",
+    "ua",
+    "uai",
+    "uan",
+    "uang",
+    "ui",
+    "un",
+    "uo",
+    "v",
+    "van",
+    "ve",
+    "vn",
+]
+
+# japanese
+ja_symbols = [
+    "I",
+    "N",
+    "U",
+    "a",
+    "b",
+    "by",
+    "ch",
+    "cl",
+    "d",
+    "dy",
+    "e",
+    "f",
+    "g",
+    "gy",
+    "h",
+    "hy",
+    "i",
+    "j",
+    "k",
+    "ky",
+    "m",
+    "my",
+    "n",
+    "ny",
+    "o",
+    "p",
+    "py",
+    "r",
+    "ry",
+    "s",
+    "sh",
+    "t",
+    "ts",
+    "u",
+    "v",
+    "w",
+    "y",
+    "z",
+    # "[", #上升调型
+    # "]", #下降调型
+    # "$", #结束符
+    # "^", #开始符
+]
+
+arpa = {
+    "AH0",
+    "S",
+    "AH1",
+    "EY2",
+    "AE2",
+    "EH0",
+    "OW2",
+    "UH0",
+    "NG",
+    "B",
+    "G",
+    "AY0",
+    "M",
+    "AA0",
+    "F",
+    "AO0",
+    "ER2",
+    "UH1",
+    "IY1",
+    "AH2",
+    "DH",
+    "IY0",
+    "EY1",
+    "IH0",
+    "K",
+    "N",
+    "W",
+    "IY2",
+    "T",
+    "AA1",
+    "ER1",
+    "EH2",
+    "OY0",
+    "UH2",
+    "UW1",
+    "Z",
+    "AW2",
+    "AW1",
+    "V",
+    "UW2",
+    "AA2",
+    "ER",
+    "AW0",
+    "UW0",
+    "R",
+    "OW1",
+    "EH1",
+    "ZH",
+    "AE0",
+    "IH2",
+    "IH",
+    "Y",
+    "JH",
+    "P",
+    "AY1",
+    "EY0",
+    "OY2",
+    "TH",
+    "HH",
+    "D",
+    "ER0",
+    "CH",
+    "AO1",
+    "AE1",
+    "AO2",
+    "OY1",
+    "AY2",
+    "IH1",
+    "OW0",
+    "L",
+    "SH",
+}
+
+symbols = [pad] + c + v + ja_symbols + pu_symbols + list(arpa)
+symbols = sorted(set(symbols))
+if __name__ == "__main__":
+    print(len(symbols))

SongBloom/g2p/cn_zh_g2p/tone_sandhi.py

@@ -0,0 +1,806 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+from typing import List
+from typing import Tuple
+
+import jieba_fast as jieba
+from pypinyin import lazy_pinyin
+from pypinyin import Style
+
+
+class ToneSandhi:
+    def __init__(self):
+        self.must_neural_tone_words = {
+            "麻烦",
+            "麻利",
+            "鸳鸯",
+            "高粱",
+            "骨头",
+            "骆驼",
+            "马虎",
+            "首饰",
+            "馒头",
+            "馄饨",
+            "风筝",
+            "难为",
+            "队伍",
+            "阔气",
+            "闺女",
+            "门道",
+            "锄头",
+            "铺盖",
+            "铃铛",
+            "铁匠",
+            "钥匙",
+            "里脊",
+            "里头",
+            "部分",
+            "那么",
+            "道士",
+            "造化",
+            "迷糊",
+            "连累",
+            "这么",
+            "这个",
+            "运气",
+            "过去",
+            "软和",
+            "转悠",
+            "踏实",
+            "跳蚤",
+            "跟头",
+            "趔趄",
+            "财主",
+            "豆腐",
+            "讲究",
+            "记性",
+            "记号",
+            "认识",
+            "规矩",
+            "见识",
+            "裁缝",
+            "补丁",
+            "衣裳",
+            "衣服",
+            "衙门",
+            "街坊",
+            "行李",
+            "行当",
+            "蛤蟆",
+            "蘑菇",
+            "薄荷",
+            "葫芦",
+            "葡萄",
+            "萝卜",
+            "荸荠",
+            "苗条",
+            "苗头",
+            "苍蝇",
+            "芝麻",
+            "舒服",
+            "舒坦",
+            "舌头",
+            "自在",
+            "膏药",
+            "脾气",
+            "脑袋",
+            "脊梁",
+            "能耐",
+            "胳膊",
+            "胭脂",
+            "胡萝",
+            "胡琴",
+            "胡同",
+            "聪明",
+            "耽误",
+            "耽搁",
+            "耷拉",
+            "耳朵",
+            "老爷",
+            "老实",
+            "老婆",
+            "老头",
+            "老太",
+            "翻腾",
+            "罗嗦",
+            "罐头",
+            "编辑",
+            "结实",
+            "红火",
+            "累赘",
+            "糨糊",
+            "糊涂",
+            "精神",
+            "粮食",
+            "簸箕",
+            "篱笆",
+            "算计",
+            "算盘",
+            "答应",
+            "笤帚",
+            "笑语",
+            "笑话",
+            "窟窿",
+            "窝囊",
+            "窗户",
+            "稳当",
+            "稀罕",
+            "称呼",
+            "秧歌",
+            "秀气",
+            "秀才",
+            "福气",
+            "祖宗",
+            "砚台",
+            "码头",
+            "石榴",
+            "石头",
+            "石匠",
+            "知识",
+            "眼睛",
+            "眯缝",
+            "眨巴",
+            "眉毛",
+            "相声",
+            "盘算",
+            "白净",
+            "痢疾",
+            "痛快",
+            "疟疾",
+            "疙瘩",
+            "疏忽",
+            "畜生",
+            "生意",
+            "甘蔗",
+            "琵琶",
+            "琢磨",
+            "琉璃",
+            "玻璃",
+            "玫瑰",
+            "玄乎",
+            "狐狸",
+            "状元",
+            "特务",
+            "牲口",
+            "牙碜",
+            "牌楼",
+            "爽快",
+            "爱人",
+            "热闹",
+            "烧饼",
+            "烟筒",
+            "烂糊",
+            "点心",
+            "炊帚",
+            "灯笼",
+            "火候",
+            "漂亮",
+            "滑溜",
+            "溜达",
+            "温和",
+            "清楚",
+            "消息",
+            "浪头",
+            "活泼",
+            "比方",
+            "正经",
+            "欺负",
+            "模糊",
+            "槟榔",
+            "棺材",
+            "棒槌",
+            "棉花",
+            "核桃",
+            "栅栏",
+            "柴火",
+            "架势",
+            "枕头",
+            "���杷",
+            "机灵",
+            "本事",
+            "木头",
+            "木匠",
+            "朋友",
+            "月饼",
+            "月亮",
+            "暖和",
+            "明白",
+            "时候",
+            "新鲜",
+            "故事",
+            "收拾",
+            "收成",
+            "提防",
+            "挖苦",
+            "挑剔",
+            "指甲",
+            "指头",
+            "拾掇",
+            "拳头",
+            "拨弄",
+            "招牌",
+            "招呼",
+            "抬举",
+            "护士",
+            "折腾",
+            "扫帚",
+            "打量",
+            "打算",
+            "打点",
+            "打扮",
+            "打听",
+            "打发",
+            "扎实",
+            "扁担",
+            "戒指",
+            "懒得",
+            "意识",
+            "意思",
+            "情形",
+            "悟性",
+            "怪物",
+            "思量",
+            "怎么",
+            "念头",
+            "念叨",
+            "快活",
+            "忙活",
+            "志气",
+            "心思",
+            "得罪",
+            "张罗",
+            "弟兄",
+            "开通",
+            "应酬",
+            "庄稼",
+            "干事",
+            "帮手",
+            "帐篷",
+            "希罕",
+            "师父",
+            "师傅",
+            "巴结",
+            "巴掌",
+            "差事",
+            "工夫",
+            "岁数",
+            "屁股",
+            "尾巴",
+            "少爷",
+            "小气",
+            "小伙",
+            "将就",
+            "对头",
+            "对付",
+            "寡妇",
+            "家伙",
+            "客气",
+            "实在",
+            "官司",
+            "学问",
+            "学生",
+            "字号",
+            "嫁妆",
+            "媳妇",
+            "媒人",
+            "婆家",
+            "娘家",
+            "委屈",
+            "姑娘",
+            "姐夫",
+            "妯娌",
+            "妥当",
+            "妖精",
+            "奴才",
+            "女婿",
+            "头发",
+            "太阳",
+            "大爷",
+            "大方",
+            "大意",
+            "大夫",
+            "多少",
+            "多么",
+            "外甥",
+            "壮实",
+            "地道",
+            "地方",
+            "在乎",
+            "困难",
+            "嘴巴",
+            "嘱咐",
+            "嘟囔",
+            "嘀咕",
+            "喜欢",
+            "喇嘛",
+            "喇叭",
+            "商量",
+            "唾沫",
+            "哑巴",
+            "哈欠",
+            "哆嗦",
+            "咳嗽",
+            "和尚",
+            "告诉",
+            "告示",
+            "含糊",
+            "吓唬",
+            "后头",
+            "名字",
+            "名堂",
+            "合同",
+            "吆喝",
+            "叫唤",
+            "口袋",
+            "厚道",
+            "厉害",
+            "千斤",
+            "包袱",
+            "包涵",
+            "匀称",
+            "勤快",
+            "动静",
+            "动弹",
+            "功夫",
+            "力气",
+            "前头",
+            "刺猬",
+            "刺激",
+            "别扭",
+            "利落",
+            "利索",
+            "利害",
+            "分析",
+            "出息",
+            "凑合",
+            "凉快",
+            "冷战",
+            "冤枉",
+            "冒失",
+            "养活",
+            "关系",
+            "先生",
+            "兄弟",
+            "便宜",
+            "使唤",
+            "佩服",
+            "作坊",
+            "体面",
+            "位置",
+            "似的",
+            "伙计",
+            "休息",
+            "什么",
+            "人家",
+            "亲戚",
+            "亲家",
+            "交情",
+            "云彩",
+            "事情",
+            "买卖",
+            "主意",
+            "丫头",
+            "丧气",
+            "两口",
+            "东西",
+            "东家",
+            "世故",
+            "不由",
+            "不在",
+            "下水",
+            "下巴",
+            "上头",
+            "上司",
+            "丈夫",
+            "丈人",
+            "一辈",
+            "那个",
+            "菩萨",
+            "父亲",
+            "母亲",
+            "咕噜",
+            "邋遢",
+            "费用",
+            "冤家",
+            "甜头",
+            "介绍",
+            "荒唐",
+            "大人",
+            "泥鳅",
+            "幸福",
+            "熟悉",
+            "计划",
+            "扑腾",
+            "蜡烛",
+            "姥爷",
+            "照顾",
+            "喉咙",
+            "吉他",
+            "弄堂",
+            "蚂蚱",
+            "凤凰",
+            "拖沓",
+            "寒碜",
+            "糟蹋",
+            "倒腾",
+            "报复",
+            "逻辑",
+            "盘缠",
+            "喽啰",
+            "牢骚",
+            "咖喱",
+            "扫把",
+            "惦记",
+        }
+        self.must_not_neural_tone_words = {
+            "男子",
+            "女子",
+            "分子",
+            "原子",
+            "量子",
+            "莲子",
+            "石子",
+            "瓜子",
+            "电子",
+            "人人",
+            "虎虎",
+            "幺幺",
+            "干嘛",
+            "学子",
+            "哈哈",
+            "数数",
+            "袅袅",
+            "局地",
+            "以下",
+            "娃哈哈",
+            "花花草草",
+            "留得",
+            "耕地",
+            "想想",
+            "熙熙",
+            "攘攘",
+            "卵子",
+            "死死",
+            "冉冉",
+            "恳恳",
+            "佼佼",
+            "吵吵",
+            "打打",
+            "考考",
+            "整整",
+            "莘莘",
+            "落地",
+            "算子",
+            "家家户户",
+            "青青",
+        }
+        self.punc = "：，；。？！“”‘’':,;.?!"
+
+    # the meaning of jieba pos tag: https://blog.csdn.net/weixin_44174352/article/details/113731041
+    # e.g.
+    # word: "家里"
+    # pos: "s"
+    # finals: ['ia1', 'i3']
+    def _neural_sandhi(self, word: str, pos: str, finals: List[str]) -> List[str]:
+        # reduplication words for n. and v. e.g. 奶奶, 试试, 旺旺
+        for j, item in enumerate(word):
+            if (
+                j - 1 >= 0
+                and item == word[j - 1]
+                and pos[0] in {"n", "v", "a"}
+                and word not in self.must_not_neural_tone_words
+            ):
+                finals[j] = finals[j][:-1] + "5"
+        ge_idx = word.find("个")
+        if len(word) >= 1 and word[-1] in "吧呢哈啊呐噻嘛吖嗨呐哦哒额滴哩哟喽啰耶喔诶":
+            finals[-1] = finals[-1][:-1] + "5"
+        elif len(word) >= 1 and word[-1] in "的地得":
+            finals[-1] = finals[-1][:-1] + "5"
+        # e.g. 走了, 看着, 去过
+        elif len(word) == 1 and word in "了着过" and pos in {"ul", "uz", "ug"}:
+            finals[-1] = finals[-1][:-1] + "5"
+        elif (
+            len(word) > 1
+            and word[-1] in "们子"
+            and pos in {"r", "n"}
+            and word not in self.must_not_neural_tone_words
+        ):
+            finals[-1] = finals[-1][:-1] + "5"
+        # e.g. 桌上, 地下, 家里
+        elif len(word) > 1 and word[-1] in "上下里" and pos in {"s", "l", "f"}:
+            finals[-1] = finals[-1][:-1] + "5"
+        # e.g. 上来, 下去
+        elif len(word) > 1 and word[-1] in "来去" and word[-2] in "上下进出回过起开":
+            finals[-1] = finals[-1][:-1] + "5"
+        # 个做量词
+        elif (
+            ge_idx >= 1
+            and (word[ge_idx - 1].isnumeric() or word[ge_idx - 1] in "几有两半多各整每做是")
+        ) or word == "个":
+            finals[ge_idx] = finals[ge_idx][:-1] + "5"
+        else:
+            if (
+                word in self.must_neural_tone_words
+                or word[-2:] in self.must_neural_tone_words
+            ):
+                finals[-1] = finals[-1][:-1] + "5"
+
+        word_list = self._split_word(word)
+        finals_list = [finals[: len(word_list[0])], finals[len(word_list[0]) :]]
+        for i, word in enumerate(word_list):
+            # conventional neural in Chinese
+            if (
+                word in self.must_neural_tone_words
+                or word[-2:] in self.must_neural_tone_words
+            ):
+                finals_list[i][-1] = finals_list[i][-1][:-1] + "5"
+        finals = sum(finals_list, [])
+        return finals
+
+    def _bu_sandhi(self, word: str, finals: List[str]) -> List[str]:
+        # e.g. 看不懂
+        if len(word) == 3 and word[1] == "不":
+            finals[1] = finals[1][:-1] + "5"
+        else:
+            for i, char in enumerate(word):
+                # "不" before tone4 should be bu2, e.g. 不怕
+                if char == "不" and i + 1 < len(word) and finals[i + 1][-1] == "4":
+                    finals[i] = finals[i][:-1] + "2"
+        return finals
+
+    def _yi_sandhi(self, word: str, finals: List[str]) -> List[str]:
+        # "一" in number sequences, e.g. 一零零, 二一零
+        if word.find("一") != -1 and all(
+            [item.isnumeric() for item in word if item != "一"]
+        ):
+            return finals
+        # "一" between reduplication words shold be yi5, e.g. 看一看
+        elif len(word) == 3 and word[1] == "一" and word[0] == word[-1]:
+            finals[1] = finals[1][:-1] + "5"
+        # when "一" is ordinal word, it should be yi1
+        elif word.startswith("第一"):
+            finals[1] = finals[1][:-1] + "1"
+        else:
+            for i, char in enumerate(word):
+                if char == "一" and i + 1 < len(word):
+                    # "一" before tone4 should be yi2, e.g. 一段
+                    if finals[i + 1][-1] == "4":
+                        finals[i] = finals[i][:-1] + "2"
+                    # "一" before non-tone4 should be yi4, e.g. 一天
+                    else:
+                        # "一" 后面如果是标点，还读一声
+                        if word[i + 1] not in self.punc:
+                            finals[i] = finals[i][:-1] + "4"
+        return finals
+
+    def _split_word(self, word: str) -> List[str]:
+        word_list = jieba.cut_for_search(word)
+        word_list = sorted(word_list, key=lambda i: len(i), reverse=False)
+        first_subword = word_list[0]
+        first_begin_idx = word.find(first_subword)
+        if first_begin_idx == 0:
+            second_subword = word[len(first_subword) :]
+            new_word_list = [first_subword, second_subword]
+        else:
+            second_subword = word[: -len(first_subword)]
+            new_word_list = [second_subword, first_subword]
+        return new_word_list
+
+    def _three_sandhi(self, word: str, finals: List[str]) -> List[str]:
+        if len(word) == 2 and self._all_tone_three(finals):
+            finals[0] = finals[0][:-1] + "2"
+        elif len(word) == 3:
+            word_list = self._split_word(word)
+            if self._all_tone_three(finals):
+                #  disyllabic + monosyllabic, e.g. 蒙古/包
+                if len(word_list[0]) == 2:
+                    finals[0] = finals[0][:-1] + "2"
+                    finals[1] = finals[1][:-1] + "2"
+                #  monosyllabic + disyllabic, e.g. 纸/老虎
+                elif len(word_list[0]) == 1:
+                    finals[1] = finals[1][:-1] + "2"
+            else:
+                finals_list = [finals[: len(word_list[0])], finals[len(word_list[0]) :]]
+                if len(finals_list) == 2:
+                    for i, sub in enumerate(finals_list):
+                        # e.g. 所有/人
+                        if self._all_tone_three(sub) and len(sub) == 2:
+                            finals_list[i][0] = finals_list[i][0][:-1] + "2"
+                        # e.g. 好/喜欢
+                        elif (
+                            i == 1
+                            and not self._all_tone_three(sub)
+                            and finals_list[i][0][-1] == "3"
+                            and finals_list[0][-1][-1] == "3"
+                        ):
+                            finals_list[0][-1] = finals_list[0][-1][:-1] + "2"
+                        finals = sum(finals_list, [])
+        # split idiom into two words who's length is 2
+        elif len(word) == 4:
+            finals_list = [finals[:2], finals[2:]]
+            finals = []
+            for sub in finals_list:
+                if self._all_tone_three(sub):
+                    sub[0] = sub[0][:-1] + "2"
+                finals += sub
+
+        return finals
+
+    def _all_tone_three(self, finals: List[str]) -> bool:
+        return all(x[-1] == "3" for x in finals)
+
+    # merge "不" and the word behind it
+    # if don't merge, "不" sometimes appears alone according to jieba, which may occur sandhi error
+    def _merge_bu(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+        new_seg = []
+        last_word = ""
+        for word, pos in seg:
+            if last_word == "不":
+                word = last_word + word
+            if word != "不":
+                new_seg.append((word, pos))
+            last_word = word[:]
+        if last_word == "不":
+            new_seg.append((last_word, "d"))
+            last_word = ""
+        return new_seg
+
+    # function 1: merge "一" and reduplication words in it's left and right, e.g. "听","一","听" ->"听一听"
+    # function 2: merge single  "一" and the word behind it
+    # if don't merge, "一" sometimes appears alone according to jieba, which may occur sandhi error
+    # e.g.
+    # input seg: [('听', 'v'), ('一', 'm'), ('听', 'v')]
+    # output seg: [['听一听', 'v']]
+    def _merge_yi(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+        new_seg = []
+        # function 1
+        for i, (word, pos) in enumerate(seg):
+            if (
+                i - 1 >= 0
+                and word == "一"
+                and i + 1 < len(seg)
+                and seg[i - 1][0] == seg[i + 1][0]
+                and seg[i - 1][1] == "v"
+                and seg[i + 1][1] == "v"
+            ):
+                new_seg[i - 1][0] = new_seg[i - 1][0] + "一" + new_seg[i - 1][0]
+            else:
+                if (
+                    i - 2 >= 0
+                    and seg[i - 1][0] == "一"
+                    and seg[i - 2][0] == word
+                    and pos == "v"
+                ):
+                    continue
+                else:
+                    new_seg.append([word, pos])
+        seg = new_seg
+        new_seg = []
+        # function 2
+        for i, (word, pos) in enumerate(seg):
+            if new_seg and new_seg[-1][0] == "一":
+                new_seg[-1][0] = new_seg[-1][0] + word
+            else:
+                new_seg.append([word, pos])
+        return new_seg
+
+    # the first and the second words are all_tone_three
+    def _merge_continuous_three_tones(
+        self, seg: List[Tuple[str, str]]
+    ) -> List[Tuple[str, str]]:
+        new_seg = []
+        sub_finals_list = [
+            lazy_pinyin(word, neutral_tone_with_five=True, style=Style.FINALS_TONE3)
+            for (word, pos) in seg
+        ]
+        assert len(sub_finals_list) == len(seg)
+        merge_last = [False] * len(seg)
+        for i, (word, pos) in enumerate(seg):
+            if (
+                i - 1 >= 0
+                and self._all_tone_three(sub_finals_list[i - 1])
+                and self._all_tone_three(sub_finals_list[i])
+                and not merge_last[i - 1]
+            ):
+                # if the last word is reduplication, not merge, because reduplication need to be _neural_sandhi
+                if (
+                    not self._is_reduplication(seg[i - 1][0])
+                    and len(seg[i - 1][0]) + len(seg[i][0]) <= 3
+                ):
+                    new_seg[-1][0] = new_seg[-1][0] + seg[i][0]
+                    merge_last[i] = True
+                else:
+                    new_seg.append([word, pos])
+            else:
+                new_seg.append([word, pos])
+
+        return new_seg
+
+    def _is_reduplication(self, word: str) -> bool:
+        return len(word) == 2 and word[0] == word[1]
+
+    # the last char of first word and the first char of second word is tone_three
+    def _merge_continuous_three_tones_2(
+        self, seg: List[Tuple[str, str]]
+    ) -> List[Tuple[str, str]]:
+        new_seg = []
+        sub_finals_list = [
+            lazy_pinyin(word, neutral_tone_with_five=True, style=Style.FINALS_TONE3)
+            for (word, pos) in seg
+        ]
+        assert len(sub_finals_list) == len(seg)
+        merge_last = [False] * len(seg)
+        for i, (word, pos) in enumerate(seg):
+            if (
+                i - 1 >= 0
+                and sub_finals_list[i - 1][-1][-1] == "3"
+                and sub_finals_list[i][0][-1] == "3"
+                and not merge_last[i - 1]
+            ):
+                # if the last word is reduplication, not merge, because reduplication need to be _neural_sandhi
+                if (
+                    not self._is_reduplication(seg[i - 1][0])
+                    and len(seg[i - 1][0]) + len(seg[i][0]) <= 3
+                ):
+                    new_seg[-1][0] = new_seg[-1][0] + seg[i][0]
+                    merge_last[i] = True
+                else:
+                    new_seg.append([word, pos])
+            else:
+                new_seg.append([word, pos])
+        return new_seg
+
+    def _merge_er(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+        new_seg = []
+        for i, (word, pos) in enumerate(seg):
+            if i - 1 >= 0 and word == "儿" and seg[i - 1][0] != "#":
+                new_seg[-1][0] = new_seg[-1][0] + seg[i][0]
+            else:
+                new_seg.append([word, pos])
+        return new_seg
+
+    def _merge_reduplication(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+        new_seg = []
+        for i, (word, pos) in enumerate(seg):
+            if new_seg and word == new_seg[-1][0]:
+                new_seg[-1][0] = new_seg[-1][0] + seg[i][0]
+            else:
+                new_seg.append([word, pos])
+        return new_seg
+
+    def pre_merge_for_modify(self, seg: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+        seg = self._merge_bu(seg)
+        try:
+            seg = self._merge_yi(seg)
+        except:
+            print("_merge_yi failed")
+        seg = self._merge_reduplication(seg)
+        try:
+            seg = self._merge_continuous_three_tones(seg)
+        except:
+            print("_merge_continuous_three_tones failed")
+        try:
+            seg = self._merge_continuous_three_tones_2(seg)
+        except:
+            print("_merge_continuous_three_tones_2 failed")
+
+        seg = self._merge_er(seg)
+        return seg
+
+    def modified_tone(self, word: str, pos: str, finals: List[str]) -> List[str]:
+        finals = self._bu_sandhi(word, finals)
+        finals = self._yi_sandhi(word, finals)
+        finals = self._neural_sandhi(word, pos, finals)
+        finals = self._three_sandhi(word, finals)
+        return finals

SongBloom/g2p/cn_zh_g2p/zh_normalization/README.md

@@ -0,0 +1,16 @@
+## Supported NSW (Non-Standard-Word) Normalization
+
+|NSW type|raw|normalized|
+|:--|:-|:-|
+|serial number|电影中梁朝伟扮演的陈永仁的编号27149|电影中梁朝伟扮演的陈永仁的编号二七一四九|
+|cardinal|这块黄金重达324.75克<br>我们班的最高总分为583分|这块黄金重达三百二十四点七五克<br>我们班的最高总分为五百八十三分|
+|numeric range |12\~23<br>-1.5\~2|十二到二十三<br>负一点五到二|
+|date|她出生于86年8月18日，她弟弟出生于1995年3月1日|她出生于八六年八月十八日， 她弟弟出生于一九九五年三月一日|
+|time|等会请在12:05请通知我|等会请在十二点零五分请通知我
+|temperature|今天的最低气温达到-10°C|今天的最低气温达到零下十度
+|fraction|现场有7/12的观众投出了赞成票|现场有十二分之七的观众投出了赞成票|
+|percentage|明天有62％的概率降雨|明天有百分之六十二的概率降雨|
+|money|随便来几个价格12块5，34.5元，20.1万|随便来几个价格十二块五，三十四点五元，二十点一万|
+|telephone|这是固话0421-33441122<br>这是手机+86 18544139121|这是固话零四二一三三四四一一二二<br>这是手机八六一八五四四一三九一二一|
+## References
+[Pull requests #658 of DeepSpeech](https://github.com/PaddlePaddle/DeepSpeech/pull/658/files)

SongBloom/g2p/cn_zh_g2p/zh_normalization/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+from .text_normlization import *

SongBloom/g2p/cn_zh_g2p/zh_normalization/char_convert.py

@@ -0,0 +1,46 @@
+# coding=utf-8
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+"""Traditional and simplified Chinese conversion, a simplified character may correspond to multiple traditional characters.
+"""
+simplified_charcters = '制咖片型超声盘鉴定仔点他命书歌粉巾字帐恤手指记忆棒形转弯沟光○〇㐄㐅㐆㐌㐖毒㐜㐡㐤㐰㐺㑇㑳㒳㒸㔾㗂㗎㝵㞎㞙㞞以㢲㢴㤅㥁㥯㨗㫺㬎㮎㮚㮸㲋㲱㲾㳮涧㵪㶸㷖㷭㹢㹴犬㺢狓㺵碗㽮㿝䍃䔢䖟䖸䗈䗥䗪䝓射䥯䦉䯝鲃鱼䲔䳗鹅䵹鼄䶑一对应映射丁不识下儿子做二休世丘之貉并中台原则串为甚谓干净了百事无成八变五十些人得道鸡升天代如并来去个国政策劲幽灵在欧洲游荡接样萝卜坑侧化传价元论醇共再准刀两断切分耕耘收获钱货物向看旧就绪险刻千金动劳永逸匙零夜半卡通回复返影踪反常态口咬气句话同吐快吹周味呼诺呜品红锅哄而散起唱和问三知生熟团漆黑火糟堆场空块面塌糊涂尘染壁厢夔已足多情露水大早到晚夫妻当关万莫开失古恨套所料既往孔见提师要家主审寸阴难买斗牛小撮部阵局展身层巴掌帆风顺席地带过年计于春头载四季期被蛇怕井绳度愿式份弹顷深前律径心意念差愁孤行俱全房厅交遮打技长把抓死拿眼泪鼻涕钥锁折段抿拍即合扫排掬挥拨拥上入击洞掷揽改故辙败文值名斑方面旁族日秋餐隔雅里终父旦时晌会霎间晃暴寒曝更月望垠际朝夕本正经利杯羹东西板枝独秀根筋杆进条龙服务概模次函数又性程总付步脚印趋登毛拔呵氧氮碳决雌雄波未平派谎言流清楚白准溜烟潭有获闻是处降琴鹤甲病发可拾沙目然了直以相眨穿睹瞥瞬矢的解石鸟神教秉虔诚秘种窝蜂穷窍笑置笔苟勾销抹杀煞等奖箍节吃箭仇双雕诗筹箩筐系列纸级士官统丝毫挂维网尽线微吭响股脑胎脉承腔臂力致效资源址器举功投般说讲规贸易叶障着慎满皆输号木电池衣倾钟高低视仁觉醒览遗角银币触溃九鼎蔽抄出驷马追重语破贫洗贯走路安蹴至几蹶振跃役胆汗较辈轮辞赞退六连遍递边针血锤音错门思闪真倒项栽雾类保护川先惊乍体哄鳞爪鸣滴泡邻域党专鼓作齐炒丑烯亥克内酯冬加奴卯肝炎基尺梁街裤镐客宠庭巳汝昌烷玲磊糖肇酉醛啷青县韪良香骨鲷丂七集河市弦喜嘴张舌堵区工业姊妹星架构巧彩扭歪拼凑余热曜武州爷浮屠美乡老阶树荤素碎落能魄鳃鳗珠丄丅丆万俟丈尚摸母娘量管群亚虎必我堂令申件装伏位博侠义界表女墟台戏臭皮匠胜诸葛亮赛顶倍催请运算包立叉戟离疫苗土史志演围揭瓦晒夷姑婆帝村宝烂尖杉碱屉桌山岔岛由纪峡坝库镇废从德后拗汤治旬食明昧曹朋友框栏极权幂曲归依猫民氟硼氯磷铁江侗自旅法司洋浦梅园温暖湾焦班幸用田略番叠皇炮捶硝苯酸腺苷棱草镜穗跳远索锦纲聚氰胺联店胚膲爱色堇紫罗兰芝茶饭菱云虫藏藩乱叛苏亲债凳学座恐恋柱测肌腹衩锥系貂企乌跪叩军车农题迭都甘油屯奏键短阿姨陪姐只顾茅庐槽驾魂鲜鹿页其菜单乘任供势午齿汉组织吊调泻唇坡城报坟外夸将尉建筑岸岗公床扬新剑升杭林栗校楼标款汽社浣海商馆剧院钢华港机械广媒环球融第医科证券综财乐育游涨犹岭疏瘾睑确兵领导缴肢膛船艾瑟尔苍蔡虞效衫覆访诉课谕议轨述野钩限敌鞋颌颔颚饶首龈站例修凡划垂届属崽颏厨拜挫摆放旋削棋榻槛礼沉注滑营狱画确仪聘花葬诏员跌辖周达酒锚闸陷陆雨雪飞威丌于丹久乏予理评产亢卑亦乎舞己悲矩圆词害志但住佞佳便俗信票案幅翁倦伦假偏倚斜亏鬼敲停备伤脾胃仅此像俭匮免宜穴焉戴兼容许冻伯仲负彼昼皂轩轾实刊划颠卫战哥比省非好黄饰别拘束掩奶睬选择摇扰烦苦枚写协厌及格受欢迎约只估侵犯割状告或缺抗拒挽撤救药喻磨灭端倪少逆逾越避靠适吉誉吝玉含延咎歹听啻渊善谋均匀堪忍够太惹妙妥妨孕症孝术室完纳推冠积宣疑辩栗碴称屈挠屑干涉衡待很忙恶忿怎么怠急耻恭息悦惑惜惟想愉愧怍慌愤启懂懈怀材才紧招认扣抵拉舍也罢插揣冒搭撞南墙扩核支攻敢雷攀敬里吗需景智暇曾罪遇朽枉止况竞争辱求愈渝溶济左右袒困补爽特寂寞示弱找谢畏强疾徐痛痒冤符眠睦瞅董何厚云措活疲羞者轻玻璃祥兆禁���稂莠稳佛换答简结果盟绝缕途给谈否羁翼耐肖胫毋宁兴舒若菲莱痕迹窠臼虚衰脸兔撒鹰棺范该详讳抬泰让须眉象众赀账费灰赖奇虑训辍辨菽麦辛近送透逞徒速续逮捕遂遑违逊斧钺艰醉锈随观弃显饱脂肪使丏丐帮丒且慢末丕替桃宗王尊凉爵各图屋脊粮署录坛吾禄职胄袭君厦丗北壑桐疹损逢陵鹬丙寅戌氨腈唑纶辰酮脱氢酶醚丞丢现掉纱帽弄扯炮碗丠両丣坐存激肩臻蒂莲悖序驱丨丩丫挺杈髻鬟细介俄伊犁京尼布订普渡央委监察检查剂圈设警队斯督剩震境航舶革防托播促质版蝾螈锋研艺历残消频谱精密制造陲邮候埔坚压坜凹汇执府究邦俘摄寮彬狼岳肺肿庸英讯诊埋粒胞括控码韩暑枪枢砥澳哇牟寿甸钻探篇签缀缝继耳肯照妇埃悬璧轴柜台辣搁浅邪跑纤阮阳私囊魔丮丰姿采丱烧丳丵丶丷丸参寨朗桂瑞砂衷霞貌凤仆舰因嫌宰峰干络牌持旨祭祷簿编罚宾办丼丿乀乂乃乄仰慕盛旷留考验阔乆乇么丑麽乊湖燃乑乒乓乕乖僻忤戾离谬迕乗危肥劫除隙浪婿乙炔肠酰吡咯盐乚乛乜嘢卿玄宫尾狐龟塔嶷兄弟泉章霄钉耙乞扎哀怜恕讨乢乣乤乥乧乨乩童乪乫乭乳晕汁液瑶浆牙癌突窦罩腐胶猪酪蛋糕菌瘤乴乵乶乷乸乹乺乼乾俸冰嘉哕嚎坤妈尸垒旱枯涸俐渴潮涩煸豆燥爹瘦瘪癣瞪袋脆姜贝隆馏乿亀亁叫咕攘扔搞男砸窜蓬麻亃亄亅却亇迟典今临繁累卵奉婚聪躬巨与迁添裂副宿岁怪恶尕仑愣杆硅硫钛铀锰芑杂异钠砷胂磺琥珀舱棍簧胡茬盗浩盆贩郎腿亍洪亐互欠助勉惠操斥诿系户译亓墓碑刑铃卅渠缤纷斗米旗宪钒灯徽瘟祖拳福谷丰脏腑绑肉腌苓蕴桥铺霸颜闹判喷冈底蛙陉矿亖亘亜罕们娜桑那努哈喀弗烈曼松森杜氏杯奥琛敦戊穆圣裔汇薛孙亟亡佚虏羊牢奋释卷卸契媾感额睫缠谊趾塞挤纽阻还配驰庄亨洛祚亪享津沪畿郊慈菴枇杷膏亭阁锃丽亳亶亹诛初责翻疯偶杰丛稠妖拖寰居吸授慧蜗吞壮魅狗矛盾益渣患忧稀描猿梦暂涯畜祸缘沸搜引擎臣横纭谁混援蒸兽狮税剖亻亼亽亡什献刹邡么仂仃仄仆富怨仈仉毕昔晨壳绍仍仏仒仕宦仗欺恃腰叹叹炬梓讫施仙后琼逝仚仝仞仟悔仡佬偿填泊拓扑簇羔购顿钦佩发棻阃驭养亿儆尤借帧赈凌叙帖李柔刚沃眦睚戒讹取飨读仨仫仮著泳卧躺韶夏裁仳仵唯贤凭钓诞仿似宋佛讽伀硕盼鹅伄儅伈伉俪柯始娃迈戈坦堡帕茨萨庙玛莉莎藤霍姆伋伍奢胥廷芳豪伎俩侍汛勒希羲雏伐憩整谟闲闲伕伙伴颐伜伝伢叔恒兹恩翰伱伲侣伶俜悧鼬伸懒缩喇叭伹伺伻伽倻辐伾似佃伫布乔妮墨佉卢佌贷劣廉昂档浓矮伞洼缓耗胸谷迷挡率龋宅沫舍疗佐贰佑占优据铧尝呢须鲁晓佗佘余坪寺瓜铳僧蒙芒陀龛哼呕坊奸孽弊揖祟茧缚誓贼佝偻瞀佟你夺赶佡佢佣佤佧贾佪佫佯佰佱洁绩酿肴佴卷佶佷佸佹佺佻佼佽佾具唤窘坏娱怒慨硬习惯聋膨胀蔓骇贵痹侀侁侂侃侄侅鸿燕侇侈糜靡侉侌妾侏儒仓鼠侐侑侔仑侘侚链侜偎傍钴循柳葫芦附価侮骂蔑侯岩截蚀局贴壶嬛宴捷携桶笺酌俣狭膝狄俅俉俊俏俎俑俓俔谚俚俛黎健呈固墒增守康箱湿祐镖镳杠盒靖膜龄俞豹猎噪孚封札筒托衍鸽剪撰稿炼厂禊练缮葺俯瞰撑冲效俳俴俵俶俷俺备俾伥倂倅储卒惶敷猝逃颉蓄崇隐倌倏忽刺蜡烛噍嚼坍扁抽毙葱楣灌灶粪背薮卖赔闭霉腾倓倔幸倘倜傥倝借箸挹浇阅倡狂倢倣値倥偬倨傲倩匡嗣冲柝珍倬倭寇猩倮倶倷倹勤赞偁偃充伪吏嗓寐惺扮拱芫茜藉虢钞偈伟晶偌宕距析滤殿疼瘫注颇偓偕鸭歇滞偝偟偢忘怡旺偨偩逼偫偭偯偰偱偲侦缉蹄偷减惰漏窥窃偸偺迹傀儡傅傈僳骂篱傎奎琳迪叟芭傒傔傕伧悉荒傜傞傢傣芽逼佣婢傮睨寄檄诵谣颂伛担辜弓惨蒿悼疤傺傻屄臆巢泄箧羡盖轧颓傿㑩僄僇佥僊働僎侨僔僖僚僝伪僣僤侥僦猴偾僩僬僭僮僯僰雇僵殖签静僾僿征陇儁侬儃儇侩朴薄儊儋儌儍傧儓俦侪拟尽儜儞儤儦儩汰哉寡渥裕酷儭儱罐儳儵儹傩俨儽兀臬臲鹫允勋勋宙宵帅憝彝谐嫂阋畅沛溢盈饥赫凶悍狠猛顽愚妣斩秦遣鞭耀敏荣槃泽爆碟磁秃缆辉霁卤朵娄孜烽酱勃汀箕裘钳耶蒙蕾彻兑软遭黜兎児韵媳爸兕觥兖兙兛兜售鍪肚兝兞兟兡兢兣樽殓涅睡禀籍赘泌啡肽奸幕涵涝熵疚眷稃衬讧赴焕椒歼植跏没试误猜栖窗肋袖颊兪卦撇胡岐廓轿疸枫茴珑厕秩募勺吨寓斤历亩迫筷厘最淫螺韬兮宽匪筛襄赢轭复兲诈刃堰戎痞蚁饷它冀铸冂冃円冇冉册嫁厉砺竭醮冏牧冑冓冔冕冖冗冘冞冢窄抑诬冥冫烘菇蛰冷凝坨橇淇淋炭饼砖碛窖醋雕雹霜冱冶炉艳嘲峻滩淡漠煖飕饮冼冽凃凄怆梗凅凇净凊凋敝蒙凔凛遵汞脢凞几凢処凰凯凵凶焰凸折刷纹预丧喽奔巡榜殡芙蓉租笼辑鞘萃凼锯镬刁蛮刂娩崩批拆摊掰蘖骤歧颗秒袂赃勿嘱忌磋琢肤刈羽刎讼戮舂桨艇刓刖霹雳刜创犊刡恙墅帜筵致劫劫刨昏默攸尿欲熏润薰圭删刮痧铲刱刲刳刴刵踏磅戳柏槐绣芹苋猬舟铭鹄鹜劫剁剃辫刭锉履铅克剌姻咽哨廊掠桅沿召瞻翅赵卜渺茫郭剒剔剕沥剚愎毅讷才剜剥啄采剞剟剡剣剤䌽剐肾驶黏剰袍剀紊铲剸剺剽剿劁劂札劈啪柴扳啦刘奭姥夼昫涓熙禅禹锡翔雁鹗刽刿弩柄蜻蛉劒劓劖劘劙澜篑赏矶釜晋甜薪逐劦熔纣虐赤囚劬劭労劵效劻劼劾峭艮勅勇励勍勐腊脖庞漫饲荡粥辄勖勗勘骄馁碌泮雇捐竹骑殊阱绩朴恳谨剿勧勩勯勰劢勋勷劝惩慰诫谏勹芡践阑匁庇拯粟扎袱裹饺匆遽匈匉匊匋匍匐茎匏匕妆痰脓蛹斋苑烤蹈塘羌熊阀螳螂疆碚竿纬荷茵邙魏匚匜匝匟扶稷匣匦拢匸匹耦匽匾匿卂叮疮禧轸堤棚迢钧炼卄卆遐卉瓷盲瓶当胱腱裸卋卌卍卐怯污贱鄙龌龊陋卓溪唐梯渔陈枣泥漳浔涧梨芬谯赡辕迦郑単驴弈洽鳌卛占筮卝卞卟吩啉屎翠厄卣卨卪卬卮榫袄玺绶钮蚤惧殆笃耸卲帘帙绕恤卼卽厂厎厓厔厖厗奚厘厍厜厝谅厕厤厥厪腻孢厮厰厳厣厹厺粕垢芜菁厼厾叁悟茸薯叄吵笄悌哺讥坫垄弧芯杠潜婴刍袁诘贪谍煽馈驳収岳缔灾贿骗叚叡吻拦蘑蜜诀燧玩砚筝椎蔺铜逗骊另觅叨唠谒杵姓喊嚷嚣咚咛塑寻恼憎擦只泣渗蝠叱吒咄咤喝籀黛舵舷叵叶铎懿昭穰苴辽叻叼吁堑嫖赌瞧爬众抒吅吆夥卺橡涤抱纵摩郡唁坠扇篮膀袜颈吋忾谘酬哭妓媛暗表缰迩妃羿絮蕃浑拐葵暮隅吔吖啶嗪戚吜啬噬咽吟哦咏吠吧唧嗒咐吪隽咀征燐苞茹钙哧吮吰吱嘎吲哚吴栋娇窟孟箫忠晗淞阖闾趼宇呐睛嘘拂捧疵熄竽笛糠吼吽呀吕韦蒙呃呆笨呇贡呉罄呋喃呎呏呔呠呡痴呣呤呦呧瑛眩扒晬淑姬瑜璇鹃呪呫哔嚅嗫呬呯呰呱呲咧噌钝呴呶呷呸呺呻哱咻啸噜吁坎坷逻呿咁咂咆哮咇咈咋蟹煦珅蔼咍咑咒诅咔哒嚓咾哝哩喱咗咠咡咢咣咥咦咨嗟询咩咪咫啮啮咭咮咱咲咳呛嗽咴啕咸咹咺呙喉咿婉恸悯赋矜绿茗蓝哂抢瞒哆嗦啰噻啾滨彗哋哌哎唷哟哏哐哞哢哤哪里哫啼喘哰哲萎蚌哳咩哽哿呗唅唆唈唉唎唏哗尧棣殇璜睿肃唔睇唕吣唞唣喳唪唬唰喏唲唳唵嘛唶唸唹唻唼唾唿啁啃鹦鹉啅埠栈榷祺铺鞅飙啊啍啎啐啓啕啖啗啜哑祈啢衔啤啥啫啱啲啵啺饥啽噶昆沁喁喂喆裙喈咙喋喌喎喑喒喓喔粗喙幛庆滋鹊喟喣喤喥喦喧骚喨喩梆吃葡萄喭驼挑吓碰枞瓣纯疱藻趟铬喵営喹喺喼喿嗀嗃嗄嗅嗈嗉嗊嗍嗐嗑嗔诟嗕嗖嗙嗛嗜痂癖嗝嗡嗤嗥嗨唢嗬嗯嗰嗲嗵叽嗷嗹嗾嗿嘀嘁嘂嘅惋嘈峪禾荫啀嘌嘏嘐嘒啯啧嘚唛嘞嘟囔嘣嘥嘦嘧嘬嘭这谑严敞馋松哓嘶嗥呒虾嘹嘻啴嘿噀噂噅噇噉噎噏噔噗噘噙噚咝噞噢噤蝉皿噩噫噭嗳噱哙噳嚏涌洒欲巫霏噷噼嚃嚄嚆抖哜尝嚔苏嚚嚜嚞嚟呖嚬嚭嚮嚯亸喾饬按竣苛嚵嘤啭冁呓膪谦囍囒囓囗囘萧酚飘溅谛囝溯眸纥銮鹘囟殉囡団囤囥囧囨囱囫囵囬囮囯囲図囶囷囸囹圄圉拟囻囿圀圂圃圊粹蠹赦圌垦圏滚鲱凿枘圕圛圜圞坯埂壤骸炕祠窑豚绅魠鲮鳖圧握圩圪垯圬圮圯炸岬幔毯祇窨菩溉圳圴圻圾坂坆沾坋坌舛壈昆垫墩椅坒坓坩埚坭坰坱坳坴坵坻坼杨挣涎帘垃垈垌垍垓垔垕垗垚垛垝垣垞垟垤垧垮垵垺垾垿埀畔埄埆埇埈埌殃隍埏埒埕埗埜垭埤埦埧埭埯埰埲埳埴埵埶绋埸培怖桩础辅埼埽堀诃侄庑堃堄摧磐贞韧砌堈堉垩堋堌堍堎垴堙堞堠礁堧堨舆堭堮蜓摘堲堳堽堿塁塄塈煤茔棵塍垲埘塓绸塕鸦沽虱塙冢塝缪塡坞埙塥塩塬塱场螨塼塽塾塿墀墁墈墉墐夯増毁墝墠墦渍钵墫墬堕墰墺墙橱壅壆壊壌壎壒榨蒜壔壕壖圹垆壜壝垅壡壬壭壱売壴壹壻壸寝壿夂夅夆変夊夌漱邑夓腕泄甥御骼夗夘夙衮瑙妊娠醣枭珊莺鹭戗幻魇夤蹀秘擂鸫姚宛闺屿庾挞拇賛蛤裨菠氅漓捞湄蚊霆鲨箐篆篷荆肆舅荔鲆巷惭骰辟邱镕镰阪漂烩鲵鲽鳄鸨胪鹏妒峨谭枰晏玑癸祝秤竺牡籁恢罡蝼蝎赐绒御梭夬夭砣榆怙枕夶夹馅奄崛葩谲奈贺祀赠奌奂奓奕䜣詝奘奜奠奡奣陶奨奁魁奫奬奰娲孩贬隶酥宄狡猾她姹嫣妁毡荼皋膻蝇嫔妄妍嫉媚娆妗趣妚妞妤碍妬娅妯娌妲妳妵妺姁姅姉姗姒姘姙姜姝姞姣姤姧姫姮娥姱姸姺姽婀娀诱慑胁娉婷娑娓娟娣娭娯娵娶娸娼婊婐婕婞婤婥溪孺婧婪婬婹婺婼婽媁媄媊媕媞媟媠媢媬媮妫媲媵媸媺媻媪眯媿嫄嫈袅嫏嫕妪嫘嫚嫜嫠嫡嫦嫩嫪毐嫫嫬嫰妩嫺娴嫽嫿妫嬃嬅嬉耍婵痴艳嬔嬖嬗嫱袅嫒嬢嬷嬦嬬嬭幼嬲嬴婶嬹嬾嬿孀娘孅娈孏曰癫屏孑孓雀孖斟篓谜摺孛矻鸠崮轲祜鸾孥邈毓棠膑孬孭孰孱孳孵泛罔衔孻孪宀宁冗拙株薇掣抚琪瓿榴谧弥宊濂祁瑕宍宏碁宓邸谳実潢町宥宧宨宬徵崎骏掖阙臊煮禽蚕宸豫寀寁寥寃檐庶寎暄碜寔寖寘寙寛寠苫寤肘洱滥蒗陕核寪弘绰螽宝擅疙瘩晷対檐専尃尅赎绌缭畴衅尌峙醌襟痲碧屁昊槌淘恵瀑牝畑莓缸羚觑蔻脏躁尔尓锐尗尙尜尟尢��尨尪尬尭尰擒尲尶尴尸尹潽蠖蛾尻扣梢蚴鳍脬蹲屇屌蚵屐屃挪屖屘屙屛屝屡屣峦嶂岩舄屧屦屩屪屃屮戍驻钾崖嵛巅旮旯楂榄榉芋茱萸靛麓屴屹屺屼岀岊岌岍阜岑彭巩岒岝岢岚岣岧岨岫岱岵岷峁峇峋峒峓峞峠嵋峨峰峱岘峹峿崀崁崆祯崋崌崃岖昆崒崔嵬巍萤颢崚崞崟崠峥巆崤崦崧殂岽崱崳崴崶崿嵂嵇嵊泗嵌嵎嵒嵓岁嵙嵞嵡嵩嵫嵯嵴嵼嵾嵝崭崭晴嶋嶌嶒嶓嵚崂嶙嶝嶞峤嶡嶢峄嶨嶭嶮嶰嶲岙嵘巂巃巇巉岿巌巓巘巛滇芎巟巠弋回巣巤炊擘蜥蟒蛊觋巰蜀彦淖杏茂甫楞巻巽帼巿帛斐鲫蕊帑帔帗帚琉汶帟帡帣帨裙帯帰帷帹暆帏幄帮幋幌幏帻幙帮幞幠幡幢幦幨幩幪帱幭幯幰遥蹉跎馀庚鉴幵幷稚邃庀庁広庄庈庉笠庋跋庖牺庠庤庥鲸庬庱庳庴庵馨衢庹庿廃厩廆廋廌廎廏廐廑廒荫廖廛厮搏锣廞弛袤廥廧廨廪廱绵踵髓廸迫瓯邺廻廼廾廿躔弁皱弇弌弍弎弐弑吊诡憾荐弝弢弣弤弨弭弮弰弪霖繇焘斌旭溥骞弶弸弼弾彀彄别累纠强彔彖彘彟彟陌彤贻彧绘虹彪炳雕蔚鸥彰瘅彲彳彴仿彷徉徨彸彽踩敛旆徂徇徊渭畲铉裼従筌徘徙徜徕膳苏萌渐徬徭醺徯徳徴潘徻徼忀瘁胖燎怦悸颤扉犀澎湃砰恍惚绞隘忉惮挨饿忐忑忒忖応忝忞耿忡忪忭忮忱忸怩忻悠懑怏遏怔怗怚怛怞怼黍讶怫怭懦怱怲恍怵惕怸怹恁恂恇恉恌恏恒恓恔恘恚恛恝恞恟恠恣恧眄恪恫恬澹恰恿悀悁悃悄悆悊悐悒晦悚悛悜悝悤您悩悪悮悰悱凄恻德悴怅惘闷悻悾惄愫钟蒐惆惇惌惎惏惓惔惙惛耄惝疟浊恿惦德恽惴蠢惸拈愀愃愆愈愊愍愐愑愒愓愔愕恪氓蠢騃昵惬赧悫愬愮愯恺愼慁恿慅慆慇霭慉慊愠慝慥怄怂慬慱悭慴慵慷戚焚憀灼郁憃惫憋憍眺捏轼愦憔憖憙憧憬憨憪憭怃憯憷憸憹憺懃懅懆邀懊懋怿懔懐懞懠懤懥恹懫懮懰懱毖懵遁梁雍忏懽戁戄戆戉戋戕戛戝戛戠戡戢戣戤戥戦戬戭戯轰戱披菊牖戸戹戺戻卯戽锹扂楔扃扆扈扊杖牵绢铐镯赉扐搂搅烊盹瞌跟趸镲靶鼾払扗玫腮扛扞扠扡扢盔押扤扦扱罾揄绥鞍郤窾扻扼扽抃抆抈抉抌抏瞎抔缳缢擞抜拗択抨摔歉蹿牾抶抻搐泵菸拃拄拊髀抛拌脯拎拏拑擢秧沓曳挛迂拚拝拠拡拫拭拮踢拴拶拷攒拽掇芥橐簪摹疔挈瓢骥捺蹻挌挍挎挐拣挓挖掘浚挙揍聩挲挶挟挿捂捃捄捅捆捉捋胳膊揎捌捍捎躯蛛捗捘捙捜捥捩扪捭据捱捻捼捽掀掂抡臀膘掊掎掏掐笙掔掗掞棉芍掤搪阐掫掮掯揉掱掲掽掾揃揅揆搓揌诨揕揗揘揜揝揞揠揥揩揪揫橥遒麈揰揲揵揶揸背揺搆搉搊搋搌搎搔搕撼橹捣搘搠搡搢搣搤搥搦搧搨搬楦裢讪赸掏搰搲搳搴揾搷搽搾搿摀摁摂摃摎掴摒摓跤摙摛掼摞摠摦喉羯摭摮挚摰摲抠摴抟摷掺摽撂撃撅稻撊撋挦锏泼撕撙撚㧑挢撢掸撦撅撩撬撱朔揿蚍蜉挝捡擀掳闯擉缶觚擐擕擖擗擡擣擤澡腚擧擨擩擫擭摈拧撷擸撸擽擿攃摅撵攉攥攐攓撄搀撺每攩攫辔澄攮攰攲攴轶攷砭讦攽碘敁敃敇敉叙敎筏敔敕敖闰诲敜煌敧敪敳敹敺敻敿斁衽斄牒绉诌斉斎斓鹑谰驳鳢斒筲斛斝斞斠斡斢斨斫斮晾沂潟颖绛邵斲斸釳於琅斾斿旀旗旃旄涡旌旎旐旒旓旖旛旝旟旡旣浴旰獭魃旴时旻旼旽昀昃昄昇昉晰躲澈熹皎皓矾昑昕昜昝昞昡昤晖笋昦昨是昱昳昴昶昺昻晁蹇隧蔬髦晄晅晒晛晜晞晟晡晢晤晥曦晩萘莹顗晿暁暋暌暍暐暔暕煅旸暝暠暡曚暦暨暪朦胧昵暲殄冯暵暸暹暻暾曀晔昙曈曌曏曐暧曘曙曛叠昽曩骆曱甴肱曷牍禺锟曽沧耽朁朅朆杪栓夸竟粘绦朊膺朏朐朓朕朘朙瞄觐溘饔飧朠朢朣栅椆淀虱朩朮朰朱炆璋钰炽鹮朳槿朵朾朿杅杇杌陧欣钊湛漼楷瀍煜玟缨翱肇舜贽适逵杓杕杗杙荀蘅杝杞脩珓筊杰榔狍閦颦缅莞杲杳眇杴杶杸杻杼枋枌枒枓衾葄翘纾逋枙狸桠枟槁枲枳枴枵枷枸橼枹枻柁柂柃柅柈柊柎某柑橘柒柘柙柚柜柞栎柟柢柣柤柩柬柮柰柲橙柶柷柸柺査柿栃栄栒栔栘栝栟柏栩栫栭栱栲栳栴檀栵栻桀骜桁镁桄桉桋桎梏椹葚桓桔桕桜桟桫椤桭杯桯桲桴桷桹湘溟梃梊梍梐潼栀枧梜梠梡梣梧梩梱梲梳梴梵梹棁棃樱棐棑棕榈簑绷蓑枨棘棜棨棩棪棫棬棯棰棱棳棸棹椁棼碗椄苕椈椊椋椌椐椑椓椗検椤椪椰椳椴椵椷椸椽椿楀匾楅篪楋楍楎楗楘楙楛楝楟楠楢楥桢楩楪楫楬楮楯楰梅楸楹楻楽榀榃榊榎槺榕榖榘榛狉莽搒笞榠榡榤榥榦榧杩榭榰榱梿霰榼榾桤槊闩槎槑槔槖様槜槢槥椠槪槭椮槱槲槻槼槾樆樊樏樑樕樗樘樛樟樠樧樨権樲樴樵猢狲桦樻罍樾樿橁橄橆桡笥龠橕橚橛辆椭橤橧竖膈跨橾橿檩檃檇柽檍檎檑檖檗桧槚檠樯檨檫檬梼槟檴檵柠棹櫆櫌栉櫜椟櫡槠栌枥榇栊櫹棂茄櫽欀欂欃欐欑栾欙棂溴欨欬欱欵欶欷歔欸欹欻欼欿歁歃歆艎歈歊莳蝶歓歕歘歙歛歜欤歠蹦诠镶蹒跚升陟歩歮歯歰歳歴璞歺瞑歾殁夭殈殍殑殗殜殙殛殒殢殣殥殪殚僵殰殳荃殷殸殹蛟殻肴谤殴毈毉喂毎���蕈毗毘毚茛邓毧毬毳毷毹毽毾毵牦氄氆靴氉氊氇氍氐聊氕氖気氘氙氚氛氜氝氡汹焊痉氤氲氥氦铝锌氪烃氩铵痤汪浒漉痘盂碾菖蒲蕹蛭螅氵冰氹氺氽烫氾氿渚汆汊汋汍汎汏汐汔汕褟汙汚汜蓠沼秽蔑汧汨汩汭汲汳汴堤汾沄沅沆瀣沇沈葆浸沦湎溺痼疴沌沍沏沐沔沕沘浜畹砾沚沢沬沭沮沰沱灢沴沷籽沺烹濡洄泂肛泅泆涌肓泐泑泒泓泔泖泙泚泜泝泠漩馍涛粼泞藓鳅泩泫泭泯铢泱泲洇洊泾琵琶荽蓟箔洌洎洏洑潄濯洙洚洟洢洣洧洨洩痢滔洫洮洳洴洵洸洹洺洼洿淌蜚浄浉浙赣渫浠浡浤浥淼瀚浬浭翩萍浯浰蜃淀苔蛞蝓蜇螵蛸煲鲤浃浼浽溦涂涊涐涑涒涔滂莅涘涙涪涫涬涮涴涶涷涿淄淅淆淊凄黯淓淙涟淜淝淟淠淢淤渌淦淩猥藿亵淬淮淯淰淳诣涞纺淸淹炖癯绮渇済渉渋渓渕涣渟渢滓渤澥渧渨渮渰渲渶渼湅湉湋湍湑湓湔黔湜湝浈湟湢湣湩湫湮麟湱湲湴涅満沩溍溎溏溛舐漭溠溤溧驯溮溱溲溳溵溷溻溼溽溾滁滃滉滊荥滏稽滕滘汇滝滫滮羼耷卤滹浐煎漈漊漎绎漕漖漘漙沤漜漪漾漥漦漯漰溆漶漷濞潀颍潎潏潕潗潚潝潞潠潦祉疡潲潵滗潸潺潾涠澁澂澃澉澌澍澐澒澔澙渑澣澦澧澨澫澬浍澰澴澶澼熏郁濆濇濈濉濊貊濔疣濜濠濩觞浚濮盥潍濲泺瀁滢渎渖瀌浏瀒瀔濒泸瀛潇潆瀡潴泷濑瀬弥潋瀳瀵瀹瀺瀼沣滠灉灋灒漓灖灏灞灠滦灥灨滟灪蜴灮烬獴灴灸灺炁炅鱿炗炘炙炤炫疽烙钎炯炰炱炲炴炷毁炻烀烋瘴鲳烓烔焙烜烝烳饪烺焃焄耆焌焐焓焗焜焞焠焢焮焯焱焼煁煃煆煇煊熠煍熬煐炜煕暖熏硷霾煚煝煟煠茕矸煨琐炀萁煳煺煻熀熅熇熉罴荧穹炝熘熛熜稔谙烁熤熨熯熰眶蚂颎熳熸熿燀烨燂燄盏燊燋燏燔隼燖焖燠燡灿燨燮燹燻燽燿爇爊爓爚爝爟爨蟾爯爰为爻丬爿牀牁牂牄牋窗牏牓窗釉牚腩蒡虻牠虽蛎牣牤牮牯牲牳牴牷牸牼绊牿靬犂犄犆犇犉犍犎犒荦犗犛犟犠犨犩犪犮犰狳犴犵犺狁甩狃狆狎狒獾狘狙黠狨狩狫狴狷狺狻豕狈蜘猁猇猈猊猋猓猖獗猗猘狰狞犸猞猟獕猭猱猲猳猷猸猹猺玃獀獃獉獍獏獐獒毙獙獚獜獝獞獠獢獣獧鼇蹊狯猃獬豸狝獯鬻獳犷猕猡玁菟玅玆玈珉糁禛郅玍玎玓瓅玔玕玖玗玘玞玠玡玢玤玥玦珏瑰玭玳瑁玶玷玹玼珂珇珈瑚珌馐馔珔珖珙珛珞珡珣珥珧珩珪佩珶珷珺珽琀琁陨玡琇琖琚琠琤琦琨琫琬琭琮琯琰琱琲琅琴珐珲瑀瑂瑄瑉玮瑑瑔瑗瑢瑭瑱瑲瑳瑽瑾瑿璀璨璁璅璆璈琏璊璐璘璚璝璟璠璡璥瑷璩璪璫璯璲玙璸璺璿瓀璎瓖瓘瓒瓛脐瓞瓠瓤瓧瓩瓮瓰瓱瓴瓸瓻瓼甀甁甃甄甇甋甍甎甏甑甒甓甔瓮甖甗饴蔗甙诧钜粱盎锈团甡褥産甪甬甭甮宁铠甹甽甾甿畀畁畇畈畊畋畎畓畚畛畟鄂畤畦畧荻畯畳畵畷畸畽畾疃叠疋疍疎箪疐疒疕疘疝疢疥疧疳疶疿痁痄痊痌痍痏痐痒痔痗瘢痚痠痡痣痦痩痭痯痱痳痵痻痿瘀痖瘃瘈瘉瘊瘌瘏瘐痪瘕瘖瘙瘚瘛疭瘜瘝瘗瘠瘥瘨瘭瘆瘯瘰疬瘳疠瘵瘸瘺瘘瘼癃痨痫癈癎癐癔癙癜癠疖症癞蟆癪瘿痈発踔绀蔫酵皙砬砒翎翳蔹钨镴皑鹎驹暨粤褶皀皁荚皃镈皈皌皋皒朱皕皖皘皜皝皞皤皦皨皪皫皭糙绽皴皲皻皽盅盋碗盍盚盝踞盦盩秋千盬盭眦睁瞤盯盱眙裰盵盻睐眂眅眈眊県眑眕眚眛眞眢眣眭眳眴眵眹瞓眽郛睃睅睆睊睍睎困睒睖睙睟睠睢睥睪睾睯睽睾眯瞈瞋瞍逛瞏瞕瞖眍䁖瞟瞠瞢瞫瞭瞳瞵瞷瞹瞽阇瞿眬矉矍铄矔矗矙瞩矞矟矠矣矧矬矫矰矱硪碇磙罅舫阡、矼矽礓砃砅砆砉砍砑砕砝砟砠砢砦砧砩砫砮砳艏砵砹砼硇硌硍硎硏硐硒硜硖砗磲茚钡硭硻硾碃碉碏碣碓碔碞碡碪碫碬砀碯碲砜碻礴磈磉磎硙磔磕磖磛磟磠磡磤磥蹭磪磬磴磵磹磻硗礀硚礅礌礐礚礜礞礤礧礮砻礲礵礽礿祂祄祅祆禳祊祍祏祓祔祕祗祘祛祧祫祲祻祼饵脔锢禂禇禋祦禔祎隋禖禘禚禜禝禠祃禢禤禥禨禫祢禴禸秆秈秊闱飒秋秏秕笈蘵赁秠秣秪秫秬秭秷秸稊稌稍稑稗稙稛稞稬秸稲稹稼颡稿穂穄穇穈穉穋稣贮穏穜穟秾穑穣穤穧穨穭穮穵穸窿阒窀窂窅窆窈窕窊窋窌窒窗窔窞窣窬黩蹙窑窳窴窵窭窸窗竁竃竈竑竜并竦竖篦篾笆鲛竾笉笊笎笏笐靥笓笤箓笪笫笭笮笰笱笲笳笵笸笻筀筅筇筈筎筑筘筠筤筥筦笕筒筭箸筰筱筳筴宴筸箂个箊箎箑箒箘箙箛箜篌箝箠箬镞箯箴箾篁筼筜篘篙篚篛篜篝篟篠篡篢篥篧篨篭篰篲筚篴篶篹篼箦簁簃簆簉簋簌簏簜簟簠簥簦簨簬簰簸簻籊藤籒籓籔签籚篯箨籣籥籧笾簖籫籯芾麴籵籸籹籼粁秕粋粑粔粝粛粞粢粧粨粲粳稗粻粽辟粿糅糆糈糌糍糒糔萼糗蛆蹋糢糨糬粽糯糱籴粜糸糺紃蹼鲣霉纡纨绔纫闽襻紑纰纮锭鸢鹞纴紞紟扎紩紬绂绁纻紽紾绐絁絃絅経絍绗絏缡褵絓絖絘絜绚絣螯絪絫聒絰絵绝絺絻絿綀绡綅绠绨绣綌綍綎捆綖綘継続缎绻綦綪线綮綯绾罟蝽綷縩绺绫緁绲緅緆缁绯緌緎総緑绱緖缃缄缂绵缗緤褓缌纂緪緰缑缈缏缇縁縃縄萦缙缒縏缣縕缞縚缜缟缛縠縡縢縦绦縯縰骋缧縳纤缦絷缥縻衙縿繄缫繈繊繋繐缯繖繘繙繠缋繣繨缰缲繸繻缱纁纆纇缬缵纩纑纕缵纙纚纛缾罃罆坛罋罂罎罏罖罘罛罝罠罣罥罦罨罫罭锾罳罶罹罻罽罿羂羃羇芈蕉５１鸵羑羖羌羜羝羢羣羟羧羭羮羰羱羵羶羸藜鲐翀翃翅翊翌翏翕翛翟翡翣翥翦跹翪翫翚翮翯翱翽翾翿板饕鸹锨耋耇耎耏专耒耜耔耞耡耤耨耩耪耧耰鬓耵聍聃聆聎聝聡聦聱聴聂聼阈聿肄肏肐肕腋肙肜肟肧胛肫肬肭肰肴肵肸肼胊胍胏胑胔胗胙胝胠铨胤胦胩胬胭胯胰胲胴胹胻胼胾脇脘脝脞脡脣脤脥脧脰脲脳腆腊腌臜腍腒腓胨腜腠脶腥腧腬腯踝蹬镣腴腶蠕诽膂腽嗉膇膋膔腘膗膙膟黐膣膦膫膰膴膵膷脍臃臄臇臈臌臐臑臓膘臖臙臛臝臞臧蓐诩臽臾臿舀舁鳑鲏舋舎舔舗馆舝舠舡舢舨舭舲舳舴舸舺艁艄艅艉艋艑艕艖艗艘艚艜艟艣舣艨艩舻艬艭荏艴艳艸艹艻艿芃芄芊萰陂藭芏芔芘芚蕙芟芣芤茉芧芨芩芪芮芰鲢芴芷芸荛豢芼芿苄苒苘苙苜蓿苠苡苣荬苤苎苪镑苶苹苺苻苾茀茁范蠡萣茆茇茈茌茍茖茞茠茢茥茦菰茭茯茳藨茷藘茼荁荄荅荇荈菅蜢鸮荍荑荘豆荵荸荠莆莒莔莕莘莙莚莛莜莝莦莨菪莩莪莭莰莿菀菆菉菎菏菐菑菓菔芲菘菝菡菢菣菥蓂菧菫毂蓥菶菷菹醢菺菻菼菾萅萆苌萋萏萐萑萜萩萱萴莴扁萻葇葍葎葑荭葖葙葠葥苇葧葭药葳葴葶葸葹葽蒄蒎莼茏薹莅蒟蒻蒢蒦蒨蒭藁蒯蒱鉾蒴蒹蒺蒽荪蓁蓆蓇蓊蓌蓍蓏蓓蓖蓧蓪蓫荜跣藕苁蓰蓱莼蓷蓺蓼蔀蔂蔃蔆蔇蔉蔊蔋蔌蔎蔕蔘蔙蒌蔟锷蒋雯茑蔯蔳麻蔵蔸蔾荨蒇蕋蕍荞蕐蕑芸莸蕖蕗蕝蕞蕠蕡蒉蕣蕤蕨蕳蓣蕸蕺蕻薀薁薃薅薆荟薉芗薏薐蔷薖薘剃谔钗薜薠薢薤薧薨薫薬薳薶薷薸薽薾薿藄藇藋荩藐藙藚藟藦藳藴苈藷藾蘀蘁蕲苹蘗蘘蘝蘤蘧蘩蘸蘼虀虆虍蟠虒虓虖虡虣虥虩虬虰蛵蛇虷鳟虺虼蚆蚈蚋蚓蚔蚖蚘蚜蚡蚣蚧蚨蚩蚪蚯蚰蜒蚱蚳蚶蚹蚺蚻蚿蛀蛁蛄蛅蝮蛌蛍蛐蟮蛑蛓蛔蛘蛚蛜蛡蛣蜊蛩蛱蜕螫蜅蚬蜈蝣蜋蜍蜎蜑蠊蜛饯蜞蜣蜨蜩蜮蜱蜷蜺蜾蜿蝀蝃蝋蝌蝍蝎蝏蝗蝘蝙蝝鲼蝡蝤蝥猿蝰虻蝲蝴蝻螃蠏蛳螉螋螒螓螗螘螙螚蟥螟螣螥螬螭䗖螾螀蟀蟅蝈蟊蟋蟑蟓蟛蟜蟟蟢虮蟨蟪蟭蛲蟳蛏蟷蟺蟿蠁蠂蠃虿蠋蛴蠓蚝蠗蠙蠚蠛蠜蠧蟏蠩蜂蠮蠰蠲蠵蠸蠼蠽衁衄衄衇衈衉衋衎衒同衖胡衞裳钩衭衲衵衹衺衿袈裟袗袚袟袢袪袮袲袴袷袺袼褙袽裀裉袅裋夹裍裎裒裛裯裱裲裴裾褀褂褉褊裈褎褐褒褓褔褕袆褚褡褢褦褧褪褫袅褯褰褱裆褛褽褾襁褒襆裥襉襋襌襏襚襛襜裣襞襡襢褴襦襫襬襭襮襕襶襼襽襾覂覃覅霸覉覊覌覗觇覚覜觍觎覧覩觊觏覰観觌觔觕觖觜觽觝觡酲觩觫觭觱觳觯觷觼觾觿言赅讣訇訏訑訒诂讬訧訬訳訹证訾詀詅诋毁詈詊讵詑诒诐詗诎察詨诜詶詸詹詻诙诖誂誃诔锄诓誋诳诶悖誙诮诰誧説読誯谇訚谄谆諆諌诤诹诼諕谂谀諝谝諟喧谥諴諵谌谖誊謆謇歌謍謏謑谡谥謡謦謪谪讴謷謼谩哗譅譆譈譊讹譒撰谮鑫譞噪譩谵譬譱譲谴譸譹谫讅讆詟䜩雠讐谗谶讙谠讟谽豁豉豇岂豊豋豌豏豔豞豖豗豜豝豣豦豨豭豱豳豵豶豷豺豻貅貆狸猊貔貘䝙貜貤餍贳餸贶贲赂賏赊赇赒賝赓赕賨赍斗賮賵賸赚赙赜赟贉赆赑贕赝赬赭赱赳迄趁趂趄趐趑趒趔趡趦趫趮趯趱趴趵趷趹趺趿跁跂跅跆踬跄跐跕跖跗跙跛跦跧跩跫跬跮跱跲跴跺跼跽踅踆踈踉踊踒踖踘踜踟躇蹰踠踡踣踤踥踦踧跷踫踮逾踱踊踶踹踺踼踽躞蹁蹂躏蹎蹐蹓蹔跸蹚蹜蹝迹蹠蹡蹢跶蹧蹩蹪蹯鞠蹽躃躄躅踌跻躐踯跞躘躙躗躝躠蹑躜躧躩躭躰躬躶軃軆辊軏轫軘軜軝腭転軥軨軭軱轱辘軷轵轺軽軿輀輂辇辂辁輈挽輗辄辎辋輠輤輬輭輮辏輴輵輶輹輼辗辒轇轏轑轒辚轕轖轗轘轙轝轞轹轳罪辣辞辵辶辺込辿迅迋迍麿迓迣迤逦迥迨迮迸迺迻迿逄逅逌逍逑逓迳逖逡逭逯逴逶逹遄遅侦遘遛遝遢遨遫遯遰遴绕遹遻邂邅邉邋邎邕邗邘邛邠邢邧邨邯郸邰邲邳邴邶邷邽邾邿郃郄郇郈郔郕郗郙郚郜郝郞郏郠郢郪郫郯郰郲郳郴郷郹郾郿鄀鄄郓鄇鄈鄋鄍鄎鄏鄐鄑邹邬鄕郧鄗鄘鄚鄜鄞鄠鄢鄣鄤鄦鄩鄫鄬鄮鄯鄱郐鄷鄹邝鄻鄾鄿酃酅酆酇郦酊酋酎酏酐酣酔酕醄酖酗酞酡酢酤酩酴酹酺醁醅醆醊醍醐醑醓醖醝酝醡醤醨醪醭醯醰酦醲醴醵醸醹醼醽醾釂酾酽釆釈鲈镏阊钆钇钌钯钋鼢鼹钐钏釪釬釭釱钍釸钕钫鈃钭鈆鈇钚鈊鈌钤钣鈒鈤钬钪鈬铌铈钶铛钹铍钸钿鉄鉆铊铇鉌铋鉏铂钷铆钵鉥钲鉨钼钽鉱鉲鉶铰铒鉼铪銍銎铣銕镂铫铦铑铷銤铱铟銧铥铕铯銭銰焊銶锑锉汞鋂锒鋆鋈鋊铤鋍铗鋐鋑鋕鋘鋙锊锓锔锇铓鋭铖锆锂铽鋳鋹鋺鉴镚钎錀锞锖锫锩錍铔锕錔锱铮锛錞锬锜錤錩錬録铼錼锝钔锴鍉镀鍏鍐铡鍚锻锽锸锲锘鍫鍭鍱鍴锶鍹锗针锺锿镅鎉鎋鎌鎍鎏鎒鎓鎗镉鎚鎞镃鎤铩锼鎭鎯镒镍鎴镓��鎹镎镟鏊镆镠镝鏖铿锵鏚镗镘镛鏠鏦錾镤鏸镪鏻鏽鏾铙鐄鐇鐏铹镦镡鐗馗镫镢镨鐡锎镄鐩镌鐬鐱镭鐶鐻鐽镱鑀鑅镔鑐鑕鑚鑛鑢鑤镥鑪镧鑯鑱鑴鑵镊镢钃镻闫闬闶闳閒闵閗閟阂関合閤哄阆閲阉閺阎阏阍阌暗闉阕阗闑闒闿闘闚阚闟闠闤闼阞阢阤阨阬阯阹阼阽陁陑陔陛陜陡陥陬骘陴険陼陾阴隃隈隒隗隞隠隣隤隩隮隰颧隳隷隹雂雈雉雊雎雑雒雗雘雚雝雟雩雰雱驿霂霅霈霊沾霒霓霙霝霢霣霤霨霩霪霫霮靁叇叆靑靓靣腼靪靮靰靳靷靸靺靼靿鞀鞃鞄鞍鞗鞙鞚鞝鞞鞡鞣鞨鞫鞬鞮鞶鞹鞾鞑韅鞯驮韍韎韔韖韘韝韫韡韣韭韭韱韹韺頀刮頄顸顼頍颀颃颁頖頞頠頫頬颅頯頲颕頼悴顋顑颙颛颜顕顚顜颟顣颥颞飐飑台飓颸飏飖颽颾颿飀飂飚飌翻飡飣饲飥饨饫飮飧飶餀餂饸饹餇餈饽哺馂餖餗餚馄馃餟餠餤餧餩餪餫糊餮糇餲饧馎糕饩馈馊馌馒饇馑馓膳饎饐饘饟馕馘馥馝馡馣骝骡馵馹駃駄駅駆駉駋驽駓驵駗骀驸駜骂骈駪駬骃駴骎駹駽駾騂騄骓騆騉騋骒骐麟騑騒験騕骛騠騢騣騤騧骧騵驺骟騺蓦骖骠骢驆驈骅驌骁驎骣驒驔驖驙驦驩驫骺鲠骫骭肮骱骴骶骷髅骾髁髂髄髆膀髇髑髌髋髙髝髞髟髡髣髧髪髫髭髯髲髳髹髺髽髾鬁鬃鬅鬈鬋鬎鬏鬐鬑鬒鬖鬗鬘鬙鬠鬣斗鬫鬬阄鬯鬰鬲鬵鬷魆魈魊魋魍魉魑魖鳔魛魟魣魦魨魬鲂魵魸鮀鲅鮆鲧鲇鲍鲋鮓鲒鲕鮟鱇鮠鮦鮨鲔鲑鮶鮸鮿鲧鯄鯆鲩鯈鲻鯕鲭鲞鯙鯠鲲鯥鲰鲶鳀鯸鳊鲗䲠鹣鳇鰋鳄鳆鰕鰛鰜鲥鰤鳏鰦鳎鳐鳁鳓鰶鲦鲡鰼鰽鱀鱄鳙鱆鳕鱎鱐鳝鳝鳜鲟鲎鱠鳣鱨鲚鱮鱲鱵鱻鲅鳦凫鳯鳲鳷鳻鴂鴃鴄鸩鴈鴎鸰鴔鴗鸳鸯鸲鹆鸱鴠鴢鸪鴥鸸鹋鴳鸻鴷鴽鵀鵁鸺鹁鵖鵙鹈鹕鹅鵟鵩鹌鵫鵵鵷鵻鹍鶂鶊鶏鶒鹙鶗鶡鶤鶦鶬鶱鹟鶵鶸鶹鹡鶿鹚鷁鷃鷄鷇䴘䴘鷊鷏鹧鷕鹥鸷鷞鷟鸶鹪鹩鷩鷫鷭鹇鹇鸴鷾䴙鸂鸇䴙鸏鸑鸒鸓鸬鹳鸜鹂鹸咸鹾麀麂麃麄麇麋麌麐麑麒麚麛麝麤麸面麫麮麯麰麺麾黁黈黉黢黒黓黕黙黝黟黥黦黧黮黰黱黪黶黹黻黼黾鼋鼂鼃鼅鼈鼍鼏鼐鼒冬鼖鼙鼚鼛鼡鼩鼱鼪鼫鼯鼷鼽齁齆齇齈齉齌赍齑龀齕齗龅齚龇齞龃龉龆齢出齧齩齮齯齰齱齵齾厐龑龒龚龖龘龝龡龢龤'
+
+traditional_characters = '制咖片型超聲盤鑒定仔點他命書歌粉巾字帳恤手指記憶棒形轉彎溝光○〇㐄㐅㐆㐌㐖毒㐜㐡㐤㐰㐺㑇㑳㒳㒸㔾㗂㗎㝵㞎㞙㞞㠯㢲㢴㤅㥁㥯㨗㫺㬎㮎㮚㮸㲋㲱㲾㳮㵎㵪㶸㷖㷭㹢㹴犬㺢狓㺵㼝㽮㿝䍃䔢䖟䖸䗈䗥䗪䝓䠶䥯䦉䯝䰾魚䲔䳗䳘䵹鼄䶑一對應映射丁不識下兒子做二休世丘之貉並中台原則串為甚謂乾淨了百事無成八變五十些人得道雞升天代如併來去個國政策勁幽靈在歐洲遊蕩接樣蘿蔔坑側化傳價元論醇共再准刀兩斷切分耕耘收穫錢貨物向看舊就緒險刻千金動勞永逸匙零夜半卡通回復返影蹤反常態口咬氣句話同吐快吹周味呼諾嗚品紅鍋哄而散起唱和問三知生熟團漆黑火糟堆場空塊麵塌糊塗塵染壁廂夔已足多情露水大早到晚夫妻當關萬莫開失古恨套所料既往孔見提師要家主審寸陰難買鬥牛小撮部陣局展身層巴掌帆風順席地帶過年計於春頭載四季期被蛇怕井繩度願式份彈頃深前律徑心意念差愁孤行俱全房廳交遮打技長把抓死拿眼淚鼻涕鑰鎖折段抿拍即合掃排掬揮撥擁上入擊洞擲攬改故轍敗文值名斑方面旁族日秋餐隔雅里終父旦時晌會霎間晃暴寒曝更月望垠際朝夕本正經利杯羹東西板枝獨秀根筋桿進條龍服務概模次函數又性程總付步腳印趨登毛拔呵氧氮碳決雌雄波未平派謊言流清楚白準溜煙潭有獲聞是處降琴鶴甲病發可拾沙目然瞭直以相眨穿睹瞥瞬矢的解石鳥神教秉虔誠秘種窩蜂窮竅笑置筆苟勾銷抹殺煞等獎箍節吃箭仇雙鵰詩籌籮筐系列紙級士官統絲毫掛維網盡線微吭響股腦胎脈承腔臂力致效資源址器舉功投般說講規貿易葉障著慎滿皆輸號木電池衣傾鐘高低視仁覺醒覽遺角銀幣觸潰九鼎蔽抄出駟馬追重語破貧洗貫走路安蹴至幾蹶振躍役膽汗較輩輪辭贊退六連遍遞邊針血錘音錯門思閃真倒項栽霧類保護川先驚乍體鬨鱗爪鳴滴泡鄰域黨專鼓作齊炒丑烯亥克內酯冬加奴卯肝炎基尺梁街褲鎬客寵庭巳汝昌烷玲磊糖肇酉醛啷青縣韙良香骨鯛丂七集河市弦喜嘴張舌堵區工業姊妹星架構巧彩扭歪拼湊餘熱曜武州爺浮屠美鄉老階樹葷素碎落能魄鰓鰻珠丄丅丆万俟丈尚摸母娘量管群亞虎必我堂令申件裝伏位博俠義界表女墟臺戲臭皮匠勝諸葛亮賽頂倍催請運算包立叉戟離疫苗土史志演圍揭瓦曬夷姑婆帝村寶爛尖杉鹼屜桌山岔島由紀峽壩庫鎮廢從德後拗湯治旬食明昧曹朋友框欄極權冪曲歸依貓民氟硼氯磷鐵江侗自旅法司洋浦梅園溫暖灣焦班幸用田略番疊皇炮捶硝苯酸腺苷稜草鏡穗跳遠索錦綱聚氰胺聯店胚膲愛色堇紫羅蘭芝茶飯菱雲蟲藏藩亂叛蘇親債凳學座恐戀柱測肌腹衩錐係貂企烏跪叩軍車農題迭都甘油屯奏鍵短阿姨陪姐隻顧茅廬槽駕魂鮮鹿頁其菜單乘任供勢午齒漢組織吊調瀉唇坡城報墳外夸將尉建築岸崗公床揚新劍昇杭林栗校樓標款汽社浣海商館劇院鋼華港機械廣媒環球融第醫科證券綜財樂育游漲猶嶺疏癮瞼確兵領導繳肢膛船艾瑟爾蒼蔡虞傚衫覆訪訴課諭議軌述野鉤限敵鞋頜頷顎饒首齦站例修凡劃垂屆屬崽頦廚拜挫擺放旋削棋榻檻禮沉注滑營獄畫确儀聘花葬詔員跌轄週達酒錨閘陷陸雨雪飛威丌于丹久乏予理評產亢卑亦乎舞己悲矩圓詞害誌但住佞佳便俗信票案幅翁倦倫假偏倚斜虧鬼敲停備傷脾胃僅此像儉匱免宜穴焉戴兼容許凍伯仲負彼晝皂軒輊實刊划顛衛戰哥比省非好黃飾別拘束掩奶睬選擇搖擾煩苦枚寫協厭及格受歡迎約只估侵犯割狀告或缺抗拒挽撤救藥喻磨滅端倪少逆逾越避靠適吉譽吝玉含延咎歹聽啻淵善謀均勻堪忍夠太惹妙妥妨孕症孝術室完納推冠積宣疑辯慄碴稱屈撓屑干涉衡待很忙惡忿怎麼怠急恥恭息悅惑惜惟想愉愧怍慌憤啟懂懈懷材才緊招認扣抵拉捨也罷插揣冒搭撞南牆擴核支攻敢雷攀敬裡嗎需景智暇曾罪遇朽枉止況競爭辱求癒渝溶濟左右袒困補爽特寂寞示弱找謝畏強疾徐痛癢冤符眠睦瞅董何厚云措活疲羞者輕玻璃祥兆禁移稂莠穩佛換答簡結果盟絕縷途給談否羈翼耐肖脛毋寧興舒若菲萊痕跡窠臼虛衰臉兔撒鷹棺範該詳諱抬泰讓鬚眉象眾貲賬費灰賴奇慮訓輟辨菽麥辛近送透逞徒速續逮捕遂遑違遜斧鉞艱醉鏽隨觀棄顯飽脂肪使丏丐幫丒且慢末丕替桃宗王尊涼爵各圖屋脊糧署錄壇吾祿職胄襲君廈丗北壑桐疹損逢陵鷸丙寅戌氨腈唑綸辰酮脫氫酶醚丞丟現掉紗帽弄扯砲碗丠両丣坐存激肩臻蒂蓮悖序驅丨丩丫挺杈髻鬟細介俄伊犁京尼布訂普渡央委監察檢查劑圈設警隊斯督剩震境航舶革防托播促質版蠑螈鋒研藝歷殘消頻譜精密製造陲郵候埔堅壓壢凹匯執府究邦俘攝寮彬狼嶽肺腫庸英訊診埋粒胞括控碼韓暑槍樞砥澳哇牟壽甸鑽探篇簽綴縫繼耳肯照婦埃懸璧軸櫃檯辣擱淺邪跑纖阮陽私囊魔丮丰姿采丱燒丳丵丶丷丸參寨朗桂瑞砂衷霞貌鳳僕艦因嫌宰峰幹絡牌持旨祭禱簿編罰賓辦丼丿乀乂乃乄仰慕盛曠留考驗闊乆乇么醜麼乊湖燃乑乒乓乕乖僻忤戾离謬迕乗危肥劫除隙浪婿乙炔腸酰吡咯鹽乚乛乜嘢卿玄宮尾狐龜塔嶷兄弟泉章霄釘耙乞扎哀憐恕討乢乣乤乥乧乨乩童乪乫乭乳暈汁液瑤漿牙癌突竇罩腐膠豬酪蛋糕菌瘤乴乵乶乷乸乹乺乼乾俸冰嘉噦嚎坤媽屍壘旱枯涸俐渴潮澀煸豆燥爹瘦癟癬瞪袋脆薑貝隆餾乿亀亁叫咕攘扔搞男砸竄蓬麻亃亄亅卻亇遲典今臨繁累卵奉婚聰躬巨與遷添裂副宿歲怪噁尕崙愣杆硅硫鈦鈾錳芑雜異鈉砷胂磺琥珀艙棍簧胡茬盜浩盆販郎腿亍洪亐互欠助勉惠操斥諉繫戶譯亓墓碑刑鈴卅渠繽紛斗米旗憲釩燈徽瘟祖拳福穀豐臟腑綁肉醃苓蘊橋鋪霸顏鬧判噴岡底蛙陘礦亖亙亜罕們娜桑那努哈喀弗烈曼松森杜氏盃奧琛敦戊穆聖裔彙薛孫亟亡佚虜羊牢奮釋卷卸契媾感額睫纏誼趾塞擠紐阻還配馳莊亨洛祚亪享津滬畿郊慈菴枇杷膏亭閣鋥麗亳亶亹誅初責翻瘋偶傑叢稠妖拖寰居吸授慧蝸吞壯魅狗矛盾益渣患憂稀描猿夢暫涯畜禍緣沸搜引擎臣橫紜誰混援蒸獸獅稅剖亻亼亽亾什獻剎邡麽仂仃仄仆富怨仈仉畢昔晨殼紹仍仏仒仕宦仗欺恃腰嘆歎炬梓訖施仙后瓊逝仚仝仞仟悔仡佬償填泊拓撲簇羔購頓欽佩髮棻閫馭養億儆尤藉幀賑凌敘帖李柔剛沃眥睚戒訛取饗讀仨仫仮著泳臥躺韶夏裁仳仵唯賢憑釣誕仿似宋彿諷伀碩盼鵝伄儅伈伉儷柯始娃邁戈坦堡帕茨薩廟瑪莉莎藤霍姆伋伍奢胥廷芳豪伎倆侍汛勒希羲雛伐憩整謨閑閒伕伙伴頤伜伝伢叔恆茲恩翰伱伲侶伶俜悧鼬伸懶縮喇叭伹伺伻伽倻輻伾佀佃佇佈喬妮墨佉盧佌貸劣廉昂檔濃矮傘窪緩耗胸谷迷擋率齲宅沫舍療佐貳佑佔優據鏵嘗呢須魯曉佗佘余坪寺瓜銃僧蒙芒陀龕哼嘔坊姦孽弊揖祟繭縛誓賊佝僂瞀佟你奪趕佡佢佣佤佧賈佪佫佯佰佱潔績釀餚佴捲佶佷佸佹佺佻佼佽佾具喚窘壞娛怒慨硬習慣聾膨脹蔓駭貴痺侀侁侂侃侄侅鴻燕侇侈糜靡侉侌妾侏儒倉鼠侐侑侔侖侘侚鏈侜偎傍鈷循柳葫蘆附価侮罵蔑侯岩截蝕侷貼壺嬛宴捷攜桶箋酌俁狹膝狄俅俉俊俏俎俑俓俔諺俚俛黎健呈固墒增守康箱濕祐鏢鑣槓盒靖膜齡俞豹獵噪孚封札筒託衍鴿剪撰稿煉廠禊練繕葺俯瞰撐衝俲俳俴俵俶俷俺俻俾倀倂倅儲卒惶敷猝逃頡蓄崇隱倌倏忽刺蠟燭噍嚼坍扁抽斃蔥楣灌灶糞背藪賣賠閉霉騰倓倔倖倘倜儻倝借箸挹澆閱倡狂倢倣値倥傯倨��倩匡嗣沖柝珍倬倭寇猩倮倶倷倹勤讚偁偃充偽吏嗓寐惺扮拱芫茜藉虢鈔偈偉晶偌宕距析濾殿疼癱註頗偓偕鴨歇滯偝偟偢忘怡旺偨偩偪偫偭偯偰偱偲偵緝蹄偷減惰漏窺竊偸偺迹傀儡傅傈僳傌籬傎奎琳迪叟芭傒傔傕傖悉荒傜傞傢傣芽逼傭婢傮睨寄檄誦謠頌傴擔辜弓慘蒿悼疤傺傻屄臆巢洩篋羨蓋軋頹傿儸僄僇僉僊働僎僑僔僖僚僝僞僣僤僥僦猴僨僩僬僭僮僯僰僱僵殖籤靜僾僿征隴儁儂儃儇儈朴薄儊儋儌儍儐儓儔儕儗儘儜儞儤儦儩汰哉寡渥裕酷儭儱罐儳儵儹儺儼儽兀臬臲鷲允勛勳宙宵帥憝彞諧嫂鬩暢沛溢盈飢赫兇悍狠猛頑愚妣斬秦遣鞭耀敏榮槃澤爆碟磁禿纜輝霽鹵朵婁孜烽醬勃汀箕裘鉗耶懞蕾徹兌軟遭黜兎児韻媳爸兕觥兗兙兛兜售鍪肚兝兞兟兡兢兣樽殮涅睡稟籍贅泌啡肽奸幕涵澇熵疚眷稃襯訌赴煥椒殲植跏沒試誤猜棲窗肋袖頰兪卦撇鬍岐廓轎疸楓茴瓏廁秩募勺噸寓斤曆畝迫筷釐最淫螺韜兮寬匪篩襄贏軛複兲詐刃堰戎痞蟻餉它冀鑄冂冃円冇冉冊嫁厲礪竭醮冏牧冑冓冔冕冖冗冘冞冢窄抑誣冥冫烘菇蟄冷凝坨橇淇淋炭餅磚磧窖醋雕雹霜冱冶爐艷嘲峻灘淡漠煖颼飲冼冽凃凄愴梗凅凇凈凊凋敝濛凔凜遵汞脢凞几凢処凰凱凵凶焰凸摺刷紋預喪嘍奔巡榜殯芙蓉租籠輯鞘萃凼鋸鑊刁蠻刂娩崩批拆攤掰櫱驟歧顆秒袂贓勿囑忌磋琢膚刈羽刎訟戮舂槳艇刓刖霹靂刜創犢刡恙墅幟筵緻刦刧刨昏默攸尿慾薰潤薰圭刪刮痧鏟刱刲刳刴刵踏磅戳柏槐繡芹莧蝟舟銘鵠鶩刼剁剃辮剄剉履鉛剋剌姻咽哨廊掠桅沿召瞻翅趙卜渺茫郭剒剔剕瀝剚愎毅訥纔剜剝啄採剞剟剡剣剤綵剮腎駛黏剰袍剴紊剷剸剺剽剿劁劂劄劈啪柴扳啦劉奭姥夼昫涓熙禪禹錫翔雁鶚劊劌弩柄蜻蛉劒劓劖劘劙瀾簣賞磯釜晉甜薪逐劦熔紂虐赤囚劬劭労劵効劻劼劾峭艮勅勇勵勍勐臘脖龐漫飼盪粥輒勖勗勘驕餒碌泮雇捐竹騎殊阱勣樸懇謹勦勧勩勯勰勱勲勷勸懲慰誡諫勹芡踐闌匁庇拯粟紮袱裹餃匆遽匈匉匊匋匍匐莖匏匕妝痰膿蛹齋苑烤蹈塘羌熊閥螳螂疆碚竿緯荷茵邙魏匚匜匝匟扶稷匣匭攏匸匹耦匽匾匿卂叮瘡禧軫堤棚迢鈞鍊卄卆遐卉瓷盲瓶噹胱腱裸卋卌卍卐怯污賤鄙齷齪陋卓溪唐梯漁陳棗泥漳潯澗梨芬譙贍轅迦鄭単驢弈洽鰲卛占筮卝卞卟吩啉屎翠厄卣卨卪卬卮榫襖璽綬鈕蚤懼殆篤聳卲帘帙繞卹卼卽厂厎厓厔厖厗奚厘厙厜厝諒厠厤厥厪膩孢厮厰厳厴厹厺粕垢蕪菁厼厾叁悟茸薯叄吵笄悌哺譏坫壟弧芯杠潛嬰芻袁詰貪諜煽饋駁収岳締災賄騙叚叡吻攔蘑蜜訣燧玩硯箏椎藺銅逗驪另覓叨嘮謁杵姓喊嚷囂咚嚀塑尋惱憎擦祇泣滲蝠叱吒咄咤喝籀黛舵舷叵叶鐸懿昭穰苴遼叻叼吁塹嫖賭瞧爬衆抒吅吆夥巹橡滌抱縱摩郡唁墜扇籃膀襪頸吋愾諮酬哭妓媛暗錶韁邇妃羿絮蕃渾拐葵暮隅吔吖啶嗪戚吜嗇噬嚥吟哦詠吠吧唧嗒咐吪雋咀徵燐苞茹鈣哧吮吰吱嘎吲哚吳棟嬌窟孟簫忠晗淞闔閭趼宇吶睛噓拂捧疵熄竽笛糠吼吽呀呂韋矇呃呆笨呇貢呉罄呋喃呎呏呔呠呡癡呣呤呦呧瑛眩扒晬淑姬瑜璇鵑呪呫嗶嚅囁呬呯呰呱呲咧噌鈍呴呶呷呸呺呻哱咻嘯嚕籲坎坷邏呿咁咂咆哮咇咈咋蟹煦珅藹咍咑咒詛咔噠嚓咾噥哩喱咗咠咡咢咣咥咦咨嗟詢咩咪咫嚙齧咭咮咱咲咳嗆嗽咴咷咸咹咺咼喉咿婉慟憫賦矜綠茗藍哂搶瞞哆嗦囉噻啾濱彗哋哌哎唷喲哏哐哞哢哤哪裏哫啼喘哰哲萎蚌哳哶哽哿唄唅唆唈唉唎唏嘩堯棣殤璜睿肅唔睇唕唚唞唣喳唪唬唰喏唲唳唵嘛唶唸唹唻唼唾唿啁啃鸚鵡啅埠棧榷祺舖鞅飆啊啍啎啐啓啕啖啗啜啞祈啢啣啤啥啫啱啲啵啺饑啽噶崑沁喁喂喆裙喈嚨喋喌喎喑喒喓喔粗喙幛慶滋鵲喟喣喤喥喦喧騷喨喩梆喫葡萄喭駝挑嚇碰樅瓣純皰藻趟鉻喵営喹喺喼喿嗀嗃嗄嗅嗈嗉嗊嗍嗐嗑嗔詬嗕嗖嗙嗛嗜痂癖嗝嗡嗤嗥嗨嗩嗬嗯嗰嗲嗵嘰嗷嗹嗾嗿嘀嘁嘂嘅惋嘈峪禾蔭嘊嘌嘏嘐嘒嘓嘖嘚嘜嘞嘟囔嘣嘥嘦嘧嘬嘭這謔嚴敞饞鬆嘵嘶嘷嘸蝦嘹嘻嘽嘿噀噂噅噇噉噎噏噔噗噘噙噚噝噞噢噤蟬皿噩噫噭噯噱噲噳嚏涌灑欲巫霏噷噼嚃嚄嚆抖嚌嚐嚔囌嚚嚜嚞嚟嚦嚬嚭嚮嚯嚲嚳飭按竣苛嚵嚶囀囅囈膪謙囍囒囓囗囘蕭酚飄濺諦囝溯眸紇鑾鶻囟殉囡団囤囥囧囨囪囫圇囬囮囯囲図囶囷囸囹圄圉擬囻囿圀圂圃圊粹蠹赦圌墾圏滾鯡鑿枘圕圛圜圞坯埂壤骸炕祠窯豚紳魠鯪鱉圧握圩圪垯圬圮圯炸岬幔毯祇窨菩溉圳圴圻圾坂坆沾坋坌舛壈昆墊墩椅坒坓坩堝坭坰坱坳坴坵坻坼楊掙涎簾垃垈垌垍垓垔垕垗垚垛垝垣垞垟垤垧垮垵垺垾垿埀畔埄埆埇埈埌殃隍埏埒埕埗埜埡埤埦埧埭埯埰埲埳埴埵埶紼埸培怖樁礎輔埼埽堀訶姪廡堃堄摧磐貞韌砌堈堉堊堋堌堍堎堖堙堞堠礁堧堨輿堭堮蜓摘堲堳堽堿塁塄塈煤塋棵塍塏塒塓綢���鴉沽虱塙塚塝繆塡塢塤塥塩塬塱塲蟎塼塽塾塿墀墁墈墉墐夯増毀墝墠墦漬缽墫墬墮墰墺墻櫥壅壆壊壌壎壒榨蒜壔壕壖壙壚壜壝壠壡壬壭壱売壴壹壻壼寢壿夂夅夆変夊夌漱邑夓腕泄甥禦骼夗夘夙袞瑙妊娠醣梟珊鶯鷺戧幻魘夤蹀祕擂鶇姚宛閨嶼庾撻拇賛蛤裨菠氅漓撈湄蚊霆鯊箐篆篷荊肆舅荔鮃巷慚骰辟邱鎔鐮阪漂燴鯢鰈鱷鴇臚鵬妒峨譚枰晏璣癸祝秤竺牡籟恢罡螻蠍賜絨御梭夬夭砣榆怙枕夶夾餡奄崛葩譎奈賀祀贈奌奐奓奕訢詝奘奜奠奡奣陶奨奩魁奫奬奰媧孩貶隸酥宄狡猾她奼嫣妁氈荼皋膻蠅嬪妄妍嫉媚嬈妗趣妚妞妤礙妬婭妯娌妲妳妵妺姁姅姉姍姒姘姙姜姝姞姣姤姧姫姮娥姱姸姺姽婀娀誘懾脅娉婷娑娓娟娣娭娯娵娶娸娼婊婐婕婞婤婥谿孺婧婪婬婹婺婼婽媁媄媊媕媞媟媠媢媬媮媯媲媵媸媺媻媼眯媿嫄嫈嫋嫏嫕嫗嫘嫚嫜嫠嫡嫦嫩嫪毐嫫嫬嫰嫵嫺嫻嫽嫿嬀嬃嬅嬉耍嬋痴豔嬔嬖嬗嬙嬝嬡嬢嬤嬦嬬嬭幼嬲嬴嬸嬹嬾嬿孀孃孅孌孏曰癲屏孑孓雀孖斟簍謎摺孛矻鳩崮軻祜鸞孥邈毓棠臏孬孭孰孱孳孵泛罔銜孻孿宀宁宂拙株薇掣撫琪瓿榴謐彌宊濂祁瑕宍宏碁宓邸讞実潢町宥宧宨宬徵崎駿掖闕臊煮禽蠶宸豫寀寁寥寃簷庶寎暄磣寔寖寘寙寛寠苫寤肘洱濫蒗陝覈寪弘綽螽寳擅疙瘩晷対檐専尃尅贖絀繚疇釁尌峙醌襟痲碧屁昊槌淘恵瀑牝畑莓缸羚覷蔻髒躁尒尓銳尗尙尜尟尢尥尨尪尬尭尰擒尲尶尷尸尹潽蠖蛾尻釦梢蚴鰭脬蹲屇屌蚵屐屓挪屖屘屙屛屝屢屣巒嶂巖舄屧屨屩屪屭屮戍駐鉀崖嵛巔旮旯楂欖櫸芋茱萸靛麓屴屹屺屼岀岊岌岍阜岑彭鞏岒岝岢嵐岣岧岨岫岱岵岷峁峇峋峒峓峞峠嵋峩峯峱峴峹峿崀崁崆禎崋崌崍嶇崐崒崔嵬巍螢顥崚崞崟崠崢巆崤崦崧殂崬崱崳崴崶崿嵂嵇嵊泗嵌嵎嵒嵓嵗嵙嵞嵡嵩嵫嵯嵴嵼嵾嶁嶃嶄晴嶋嶌嶒嶓嶔嶗嶙嶝嶞嶠嶡嶢嶧嶨嶭嶮嶰嶲嶴嶸巂巃巇巉巋巌巓巘巛滇芎巟巠弋迴巣巤炊擘蜥蟒蠱覡巰蜀彥淖杏茂甫楞巻巽幗巿帛斐鯽蕊帑帔帗帚琉汶帟帡帣帨帬帯帰帷帹暆幃幄幇幋幌幏幘幙幚幞幠幡幢幦幨幩幪幬幭幯幰遙蹉跎餘庚鑑幵幷稚邃庀庁広庄庈庉笠庋跋庖犧庠庤庥鯨庬庱庳庴庵馨衢庹庿廃廄廆廋廌廎廏廐廑廒廕廖廛廝搏鑼廞弛袤廥廧廨廩廱綿踵髓廸廹甌鄴廻廼廾廿躔弁皺弇弌弍弎弐弒弔詭憾薦弝弢弣弤弨弭弮弰弳霖繇燾斌旭溥騫弶弸弼弾彀彄彆纍糾彊彔彖彘彟彠陌彤貽彧繪虹彪炳彫蔚鷗彰癉彲彳彴彷彷徉徨彸彽踩斂旆徂徇徊渭畬鉉裼従筌徘徙徜徠膳甦萌漸徬徭醺徯徳徴潘徻徼忀瘁胖燎怦悸顫扉犀澎湃砰恍惚絞隘忉憚挨餓忐忑忒忖応忝忞耿忡忪忭忮忱忸怩忻悠懣怏遏怔怗怚怛怞懟黍訝怫怭懦怱怲怳怵惕怸怹恁恂恇恉恌恏恒恓恔恘恚恛恝恞恟恠恣恧眄恪恫恬澹恰恿悀悁悃悄悆悊悐悒晦悚悛悜悝悤您悩悪悮悰悱悽惻悳悴悵惘悶悻悾惄愫鍾蒐惆惇惌惎惏惓惔惙惛耄惝瘧濁惥惦惪惲惴惷惸拈愀愃愆愈愊愍愐愑愒愓愔愕愙氓蠢騃昵愜赧愨愬愮愯愷愼慁慂慅慆慇靄慉慊慍慝慥慪慫慬慱慳慴慵慷慼焚憀灼鬱憃憊憋憍眺捏軾憒憔憖憙憧憬憨憪憭憮憯憷憸憹憺懃懅懆邀懊懋懌懍懐懞懠懤懥懨懫懮懰懱毖懵遁樑雍懺懽戁戄戇戉戔戕戛戝戞戠戡戢戣戤戥戦戩戭戯轟戱披菊牖戸戹戺戻戼戽鍬扂楔扃扆扈扊杖牽絹銬鐲賚扐摟攪烊盹瞌跟躉鑔靶鼾払扗玫腮扛扞扠扡扢盔押扤扦扱罾揄綏鞍郤窾扻扼扽抃抆抈抉抌抏瞎抔繯縊擻抜抝択抨摔歉躥牾抶抻搐泵菸拃拄拊髀拋拌脯拎拏拑擢秧沓曳攣迂拚拝拠拡拫拭拮踢拴拶拷攢拽掇芥橐簪摹疔挈瓢驥捺蹻挌挍挎挐揀挓挖掘浚挙揍聵挲挶挾挿捂捃捄捅捆捉捋胳膊揎捌捍捎軀蛛捗捘捙捜捥捩捫捭据捱捻捼捽掀掂掄臀膘掊掎掏掐笙掔掗掞棉芍掤搪闡掫掮掯揉掱掲掽掾揃揅揆搓揌諢揕揗揘揜揝揞揠揥揩揪揫櫫遒麈揰揲揵揶揸揹揺搆搉搊搋搌搎搔搕撼櫓搗搘搠搡搢搣搤搥搦搧搨搬楦褳訕赸搯搰搲搳搴搵搷搽搾搿摀摁摂摃摎摑摒摓跤摙摛摜摞摠摦睺羯摭摮摯摰摲摳摴摶摷摻摽撂撃撅稻撊撋撏鐧潑撕撙撚撝撟撢撣撦撧撩撬撱朔撳蚍蜉撾撿擀擄闖擉缶觚擐擕擖擗擡擣擤澡腚擧擨擩擫擭擯擰擷擸擼擽擿攃攄攆攉攥攐攓攖攙攛每攩攫轡澄攮攰攲攴軼攷砭訐攽碘敁敃敇敉敍敎筏敔敕敖閏誨敜煌敧敪敱敹敺敻敿斁衽斄牒縐謅斉斎斕鶉讕駮鱧斒筲斛斝斞斠斡斢斨斫斮晾沂潟穎絳邵斲斸釳於琅斾斿旀旂旃旄渦旌旎旐旒旓旖旛旝旟旡旣浴旰獺魃旴旹旻旼旽昀昃昄昇昉晰躲澈熹皎皓礬昑昕昜昝昞昡昤暉筍昦昨昰昱昳昴昶昺昻晁蹇隧蔬髦晄晅晒晛晜晞晟晡晢晤晥曦晩萘瑩顗晿暁暋暌暍暐暔暕煅暘暝暠暡曚暦暨暪朦朧暱暲殄馮暵暸暹暻暾曀曄曇曈曌曏曐曖曘曙曛曡曨曩駱曱甴肱曷牘禺錕曽滄耽朁朅朆杪栓誇竟粘絛朊膺朏朐朓朕朘朙瞄覲溘饔飧朠朢朣柵椆澱蝨朩朮朰朱炆璋鈺熾鹮朳槿朶朾朿杅杇杌隉欣釗湛漼楷瀍煜玟纓翱肈舜贄适逵杓杕杗杙荀蘅杝杞脩珓筊杰榔狍閦顰緬莞杲杳眇杴杶杸杻杼枋枌枒枓衾葄翹紓逋枙狸椏枟槁枲枳枴枵枷枸櫞枹枻柁柂柃柅柈柊柎某柑橘柒柘柙柚柜柞櫟柟柢柣柤柩柬柮柰柲橙柶柷柸柺査柿栃栄栒栔栘栝栟栢栩栫栭栱栲栳栴檀栵栻桀驁桁鎂桄桉桋桎梏椹葚桓桔桕桜桟桫欏桭桮桯桲桴桷桹湘溟梃梊梍梐潼梔梘梜梠梡梣梧梩梱梲梳梴梵梹棁棃櫻棐棑棕櫚簑繃蓑棖棘棜棨棩棪棫棬棯棰棱棳棸棹槨棼椀椄苕椈椊椋椌椐椑椓椗検椤椪椰椳椴椵椷椸椽椿楀楄楅篪楋楍楎楗楘楙楛楝楟楠楢楥楨楩楪楫楬楮楯楰楳楸楹楻楽榀榃榊榎槺榕榖榘榛狉莽榜笞榠榡榤榥榦榧榪榭榰榱槤霰榼榾榿槊閂槎槑槔槖様槜槢槥槧槪槭槮槱槲槻槼槾樆樊樏樑樕樗樘樛樟樠樧樨権樲樴樵猢猻樺樻罍樾樿橁橄橆橈笥龠橕橚橛輛橢橤橧豎膈跨橾橿檁檃檇檉檍檎檑檖檗檜檟檠檣檨檫檬檮檳檴檵檸櫂櫆櫌櫛櫜櫝櫡櫧櫨櫪櫬櫳櫹櫺茄櫽欀欂欃欐欑欒欙欞溴欨欬欱欵欶欷歔欸欹欻欼欿歁歃歆艎歈歊蒔蝶歓歕歘歙歛歜歟歠蹦詮鑲蹣跚陞陟歩歮歯歰歳歴璞歺瞑歾歿殀殈殍殑殗殜殙殛殞殢殣殥殪殫殭殰殳荃殷殸殹蛟殻殽謗毆毈毉餵毎毑蕈毗毘毚茛鄧毧毬毳毷毹毽毾毿氂氄氆靴氉氊氌氍氐聊氕氖気氘氙氚氛氜氝氡洶焊痙氤氳氥氦鋁鋅氪烴氬銨痤汪滸漉痘盂碾菖蒲蕹蛭螅氵氷氹氺氽燙氾氿渚汆汊汋汍汎汏汐汔汕褟汙汚汜蘺沼穢衊汧汨汩汭汲汳汴隄汾沄沅沆瀣沇沈葆浸淪湎溺痼痾沌沍沏沐沔沕沘浜畹礫沚沢沬沭沮沰沱灢沴沷籽沺烹濡洄泂肛泅泆湧肓泐泑泒泓泔泖泙泚泜泝泠漩饃濤粼濘蘚鰍泩泫泭泯銖泱泲洇洊涇琵琶荽薊箔洌洎洏洑潄濯洙洚洟洢洣洧洨洩痢滔洫洮洳洴洵洸洹洺洼洿淌蜚浄浉浙贛渫浠浡浤浥淼瀚浬浭翩萍浯浰蜃淀苔蛞蝓蜇螵蛸煲鯉浹浼浽溦涂涊涐涑涒涔滂涖涘涙涪涫涬涮涴涶涷涿淄淅淆淊淒黯淓淙漣淜淝淟淠淢淤淥淦淩猥藿褻淬淮淯淰淳詣淶紡淸淹燉癯綺渇済渉渋渓渕渙渟渢滓渤澥渧渨渮渰渲渶渼湅湉湋湍湑湓湔黔湜湝湞湟湢湣湩湫湮麟湱湲湴湼満溈溍溎溏溛舐漭溠溤溧馴溮溱溲溳溵溷溻溼溽溾滁滃滉滊滎滏稽滕滘滙滝滫滮羼耷滷滹滻煎漈漊漎繹漕漖漘漙漚漜漪漾漥漦漯漰漵漶漷濞潀潁潎潏潕潗潚潝潞潠潦祉瘍潲潵潷潸潺潾潿澁澂澃澉澌澍澐澒澔澙澠澣澦澧澨澫澬澮澰澴澶澼熏郁濆濇濈濉濊貊濔疣濜濠濩觴濬濮盥濰濲濼瀁瀅瀆瀋瀌瀏瀒瀔瀕瀘瀛瀟瀠瀡瀦瀧瀨瀬瀰瀲瀳瀵瀹瀺瀼灃灄灉灋灒灕灖灝灞灠灤灥灨灩灪蜴灮燼獴灴灸灺炁炅魷炗炘炙炤炫疽烙釺炯炰炱炲炴炷燬炻烀烋瘴鯧烓烔焙烜烝烳飪烺焃焄耆焌焐焓焗焜焞焠焢焮焯焱焼煁煃煆煇煊熠煍熬煐煒煕煗燻礆霾煚煝煟煠煢矸煨瑣煬萁煳煺煻熀熅熇熉羆熒穹熗熘熛熜稔諳爍熤熨熯熰眶螞熲熳熸熿燀燁燂燄盞燊燋燏燔隼燖燜燠燡燦燨燮燹燻燽燿爇爊爓爚爝爟爨蟾爯爰爲爻爿爿牀牁牂牄牋牎牏牓牕釉牚腩蒡虻牠雖蠣牣牤牮牯牲牳牴牷牸牼絆牿靬犂犄犆犇犉犍犎犒犖犗犛犟犠犨犩犪犮犰狳犴犵犺狁甩狃狆狎狒獾狘狙黠狨狩狫狴狷狺狻豕狽蜘猁猇猈猊猋猓猖獗猗猘猙獰獁猞猟獕猭猱猲猳猷猸猹猺玃獀獃獉獍獏獐獒獘獙獚獜獝獞獠獢獣獧鼇蹊獪獫獬豸獮獯鬻獳獷獼玀玁菟玅玆玈珉糝禛郅玍玎玓瓅玔玕玖玗玘玞玠玡玢玤玥玦玨瑰玭玳瑁玶玷玹玼珂珇珈瑚珌饈饌珔珖珙珛珞珡珣珥珧珩珪珮珶珷珺珽琀琁隕琊琇琖琚琠琤琦琨琫琬琭琮琯琰琱琲瑯琹琺琿瑀瑂瑄瑉瑋瑑瑔瑗瑢瑭瑱瑲瑳瑽瑾瑿璀璨璁璅璆璈璉璊璐璘璚璝璟璠璡璥璦璩璪璫璯璲璵璸璺璿瓀瓔瓖瓘瓚瓛臍瓞瓠瓤瓧瓩瓮瓰瓱瓴瓸瓻瓼甀甁甃甄甇甋甍甎甏甑甒甓甔甕甖甗飴蔗甙詫鉅粱盎銹糰甡褥産甪甬甭甮甯鎧甹甽甾甿畀畁畇畈畊畋畎畓畚畛畟鄂畤畦畧荻畯畳畵畷畸畽畾疃疉疋疍疎簞疐疒疕疘疝疢疥疧疳疶疿痁痄痊痌痍痏痐痒痔痗瘢痚痠痡痣痦痩痭痯痱痳痵痻痿瘀瘂瘃瘈瘉瘊瘌瘏瘐瘓瘕瘖瘙瘚瘛瘲瘜瘝瘞瘠瘥瘨瘭瘮瘯瘰癧瘳癘瘵瘸瘺瘻瘼癃癆癇癈癎癐癔癙癜癠癤癥癩蟆癪癭癰発踔紺蔫酵皙砬砒翎翳蘞鎢鑞皚鵯駒鱀粵褶皀皁莢皃鎛皈皌皐皒硃皕皖皘皜皝皞皤皦皨皪皫皭糙綻皴皸皻皽盅盋盌盍盚盝踞盦盩鞦韆盬盭眦睜瞤盯盱眙裰盵盻睞眂眅眈眊県眑眕眚眛眞眢眣眭眳眴眵眹瞓眽郛睃睅睆睊睍睎睏睒睖睙睟睠睢睥睪睪睯睽睾瞇瞈瞋瞍逛瞏瞕瞖瞘瞜瞟瞠瞢瞫瞭瞳瞵瞷瞹瞽闍瞿矓矉矍鑠矔矗矙矚矞矟矠矣矧矬矯矰矱硪碇磙��舫阡、矼矽礓砃砅砆砉砍砑砕砝砟砠砢砦砧砩砫砮砳艏砵砹砼硇硌硍硎硏硐硒硜硤硨磲茚鋇硭硻硾碃碉碏碣碓碔碞碡碪碫碬碭碯碲碸碻礡磈磉磎磑磔磕磖磛磟磠磡磤磥蹭磪磬磴磵磹磻磽礀礄礅礌礐礚礜礞礤礧礮礱礲礵礽礿祂祄祅祆禳祊祍祏祓祔祕祗祘祛祧祫祲祻祼餌臠錮禂禇禋禑禔禕隋禖禘禚禜禝禠禡禢禤禥禨禫禰禴禸稈秈秊闈颯秌秏秕笈蘵賃秠秣秪秫秬秭秷秸稊稌稍稑稗稙稛稞稬稭稲稹稼顙稾穂穄穇穈穉穋穌貯穏穜穟穠穡穣穤穧穨穭穮穵穸窿闃窀窂窅窆窈窕窊窋窌窒窓窔窞窣窬黷蹙窰窳窴窵窶窸窻竁竃竈竑竜竝竦竪篦篾笆鮫竾笉笊笎笏笐靨笓笤籙笪笫笭笮笰笱笲笳笵笸笻筀筅筇筈筎筑筘筠筤筥筦筧筩筭筯筰筱筳筴讌筸箂箇箊箎箑箒箘箙箛箜篌箝箠箬鏃箯箴箾篁篔簹篘篙篚篛篜篝篟篠篡篢篥篧篨篭篰篲篳篴篶篹篼簀簁簃簆簉簋簌簏簜簟簠簥簦簨簬簰簸簻籊籐籒籓籔籖籚籛籜籣籥籧籩籪籫籯芾麴籵籸籹籼粁粃粋粑粔糲粛粞粢粧粨粲粳粺粻粽闢粿糅糆糈糌糍糒糔萼糗蛆蹋糢糨糬糭糯糱糴糶糸糺紃蹼鰹黴紆紈絝紉閩襻紑紕紘錠鳶鷂紝紞紟紥紩紬紱紲紵紽紾紿絁絃絅経絍絎絏縭褵絓絖絘絜絢絣螯絪絫聒絰絵絶絺絻絿綀綃綅綆綈綉綌綍綎綑綖綘継続緞綣綦綪綫綮綯綰罟蝽綷縩綹綾緁緄緅緆緇緋緌緎総緑緔緖緗緘緙緜緡緤緥緦纂緪緰緱緲緶緹縁縃縄縈縉縋縏縑縕縗縚縝縞縟縠縡縢縦縧縯縰騁縲縳縴縵縶縹縻衙縿繄繅繈繊繋繐繒繖繘繙繠繢繣繨繮繰繸繻繾纁纆纇纈纉纊纑纕纘纙纚纛缾罃罆罈罋罌罎罏罖罘罛罝罠罣罥罦罨罫罭鍰罳罶罹罻罽罿羂羃羇羋蕉５１鴕羑羖羗羜羝羢羣羥羧羭羮羰羱羵羶羸藜鮐翀翃翄翊翌翏翕翛翟翡翣翥翦躚翪翫翬翮翯翺翽翾翿闆饕鴰鍁耋耇耎耏耑耒耜耔耞耡耤耨耩耪耬耰鬢耵聹聃聆聎聝聡聦聱聴聶聼閾聿肄肏肐肕腋肙肜肟肧胛肫肬肭肰肴肵肸肼胊胍胏胑胔胗胙胝胠銓胤胦胩胬胭胯胰胲胴胹胻胼胾脇脘脝脞脡脣脤脥脧脰脲脳腆腊腌臢腍腒腓腖腜腠腡腥腧腬腯踝蹬鐐腴腶蠕誹膂膃膆膇膋膔膕膗膙膟黐膣膦膫膰膴膵膷膾臃臄臇臈臌臐臑臓臕臖臙臛臝臞臧蓐詡臽臾臿舀舁鰟鮍舋舎舔舗舘舝舠舡舢舨舭舲舳舴舸舺艁艄艅艉艋艑艕艖艗艘艚艜艟艣艤艨艩艫艬艭荏艴艶艸艹艻艿芃芄芊萰陂藭芏芔芘芚蕙芟芣芤茉芧芨芩芪芮芰鰱芴芷芸蕘豢芼芿苄苒苘苙苜蓿苠苡苣蕒苤苧苪鎊苶苹苺苻苾茀茁范蠡萣茆茇茈茌茍茖茞茠茢茥茦菰茭茯茳藨茷藘茼荁荄荅荇荈菅蜢鴞荍荑荘荳荵荸薺莆莒莔莕莘莙莚莛莜莝莦莨菪莩莪莭莰莿菀菆菉菎菏菐菑菓菔菕菘菝菡菢菣菥蓂菧菫轂鎣菶菷菹醢菺菻菼菾萅萆萇萋萏萐萑萜萩萱萴萵萹萻葇葍葎葑葒葖葙葠葥葦葧葭葯葳葴葶葸葹葽蒄蒎蒓蘢薹蒞蒟蒻蒢蒦蒨蒭藁蒯蒱鉾蒴蒹蒺蒽蓀蓁蓆蓇蓊蓌蓍蓏蓓蓖蓧蓪蓫蓽跣藕蓯蓰蓱蓴蓷蓺蓼蔀蔂蔃蔆蔇蔉蔊蔋蔌蔎蔕蔘蔙蔞蔟鍔蔣雯蔦蔯蔳蔴蔵蔸蔾蕁蕆蕋蕍蕎蕐蕑蕓蕕蕖蕗蕝蕞蕠蕡蕢蕣蕤蕨蕳蕷蕸蕺蕻薀薁薃薅薆薈薉薌薏薐薔薖薘薙諤釵薜薠薢薤薧薨薫薬薳薶薷薸薽薾薿藄藇藋藎藐藙藚藟藦藳藴藶藷藾蘀蘁蘄蘋蘗蘘蘝蘤蘧蘩蘸蘼虀虆虍蟠虒虓虖虡虣虥虩虯虰蛵虵虷鱒虺虼蚆蚈蚋蚓蚔蚖蚘蚜蚡蚣蚧蚨蚩蚪蚯蚰蜒蚱蚳蚶蚹蚺蚻蚿蛀蛁蛄蛅蝮蛌蛍蛐蟮蛑蛓蛔蛘蛚蛜蛡蛣蜊蛩蛺蛻螫蜅蜆蜈蝣蜋蜍蜎蜑蠊蜛餞蜞蜣蜨蜩蜮蜱蜷蜺蜾蜿蝀蝃蝋蝌蝍蝎蝏蝗蝘蝙蝝鱝蝡蝤蝥蝯蝰蝱蝲蝴蝻螃蠏螄螉螋螒螓螗螘螙螚蟥螟螣螥螬螭螮螾螿蟀蟅蟈蟊蟋蟑蟓蟛蟜蟟蟢蟣蟨蟪蟭蟯蟳蟶蟷蟺蟿蠁蠂蠃蠆蠋蠐蠓蠔蠗蠙蠚蠛蠜蠧蠨蠩蠭蠮蠰蠲蠵蠸蠼蠽衁衂衄衇衈衉衋衎衒衕衖衚衞裳鈎衭衲衵衹衺衿袈裟袗袚袟袢袪袮袲袴袷袺袼褙袽裀裉裊裋裌裍裎裒裛裯裱裲裴裾褀褂褉褊褌褎褐褒褓褔褕褘褚褡褢褦褧褪褫褭褯褰褱襠褸褽褾襁襃襆襇襉襋襌襏襚襛襜襝襞襡襢襤襦襫襬襭襮襴襶襼襽襾覂覃覅覇覉覊覌覗覘覚覜覥覦覧覩覬覯覰観覿觔觕觖觜觽觝觡酲觩觫觭觱觳觶觷觼觾觿言賅訃訇訏訑訒詁託訧訬訳訹証訾詀詅詆譭詈詊詎詑詒詖詗詘詧詨詵詶詸詹詻詼詿誂誃誄鋤誆誋誑誒誖誙誚誥誧説読誯誶誾諂諄諆諌諍諏諑諕諗諛諝諞諟諠諡諴諵諶諼謄謆謇謌謍謏謑謖謚謡謦謪謫謳謷謼謾譁譅譆譈譊譌譒譔譖鑫譞譟譩譫譬譱譲譴譸譹譾讅讆讋讌讎讐讒讖讙讜讟谽豁豉豇豈豊豋豌豏豔豞豖豗豜豝豣豦豨豭豱豳豵豶豷豺豻貅貆貍貎貔貘貙貜貤饜貰餸貺賁賂賏賒賕賙賝賡賧賨賫鬭賮賵賸賺賻賾贇贉贐贔贕贗赬赭赱赳迄趁趂趄趐趑趒趔趡趦趫趮趯趲趴趵趷趹趺趿跁跂跅跆躓蹌跐跕跖跗跙跛跦跧跩跫跬跮跱跲跴跺跼跽踅踆踈踉踊踒���踘踜踟躇躕踠踡踣踤踥踦踧蹺踫踮踰踱踴踶踹踺踼踽躞蹁蹂躪蹎蹐蹓蹔蹕蹚蹜蹝蹟蹠蹡蹢躂蹧蹩蹪蹯鞠蹽躃躄躅躊躋躐躑躒躘躙躛躝躠躡躦躧躩躭躰躳躶軃軆輥軏軔軘軜軝齶転軥軨軭軱軲轆軷軹軺軽軿輀輂輦輅輇輈輓輗輙輜輞輠輤輬輭輮輳輴輵輶輹輼輾轀轇轏轑轒轔轕轖轗轘轙轝轞轢轤辠辢辤辵辶辺込辿迅迋迍麿迓迣迤邐迥迨迮迸迺迻迿逄逅逌逍逑逓逕逖逡逭逯逴逶逹遄遅遉遘遛遝遢遨遫遯遰遴遶遹遻邂邅邉邋邎邕邗邘邛邠邢邧邨邯鄲邰邲邳邴邶邷邽邾邿郃郄郇郈郔郕郗郙郚郜郝郞郟郠郢郪郫郯郰郲郳郴郷郹郾郿鄀鄄鄆鄇鄈鄋鄍鄎鄏鄐鄑鄒鄔鄕鄖鄗鄘鄚鄜鄞鄠鄢鄣鄤鄦鄩鄫鄬鄮鄯鄱鄶鄷鄹鄺鄻鄾鄿酃酅酆酇酈酊酋酎酏酐酣酔酕醄酖酗酞酡酢酤酩酴酹酺醁醅醆醊醍醐醑醓醖醝醞醡醤醨醪醭醯醰醱醲醴醵醸醹醼醽醾釂釃釅釆釈鱸鎦閶釓釔釕鈀釙鼢鼴釤釧釪釬釭釱釷釸釹鈁鈃鈄鈆鈇鈈鈊鈌鈐鈑鈒鈤鈥鈧鈬鈮鈰鈳鐺鈸鈹鈽鈿鉄鉆鉈鉋鉌鉍鉏鉑鉕鉚鉢鉥鉦鉨鉬鉭鉱鉲鉶鉸鉺鉼鉿銍銎銑銕鏤銚銛銠銣銤銥銦銧銩銪銫銭銰銲銶銻銼銾鋂鋃鋆鋈鋊鋌鋍鋏鋐鋑鋕鋘鋙鋝鋟鋦鋨鋩鋭鋮鋯鋰鋱鋳鋹鋺鋻鏰鐱錀錁錆錇錈錍錏錒錔錙錚錛錞錟錡錤錩錬録錸錼鍀鍆鍇鍉鍍鍏鍐鍘鍚鍛鍠鍤鍥鍩鍫鍭鍱鍴鍶鍹鍺鍼鍾鎄鎇鎉鎋鎌鎍鎏鎒鎓鎗鎘鎚鎞鎡鎤鎩鎪鎭鎯鎰鎳鎴鎵鎸鎹鎿鏇鏊鏌鏐鏑鏖鏗鏘鏚鏜鏝鏞鏠鏦鏨鏷鏸鏹鏻鏽鏾鐃鐄鐇鐏鐒鐓鐔鐗馗鐙鐝鐠鐡鐦鐨鐩鐫鐬鐱鐳鐶鐻鐽鐿鑀鑅鑌鑐鑕鑚鑛鑢鑤鑥鑪鑭鑯鑱鑴鑵鑷钁钃镻閆閈閌閎閒閔閗閟閡関閤閤閧閬閲閹閺閻閼閽閿闇闉闋闐闑闒闓闘闚闞闟闠闤闥阞阢阤阨阬阯阹阼阽陁陑陔陛陜陡陥陬騭陴険陼陾隂隃隈隒隗隞隠隣隤隩隮隰顴隳隷隹雂雈雉雊雎雑雒雗雘雚雝雟雩雰雱驛霂霅霈霊霑霒霓霙霝霢霣霤霨霩霪霫霮靁靆靉靑靚靣靦靪靮靰靳靷靸靺靼靿鞀鞃鞄鞌鞗鞙鞚鞝鞞鞡鞣鞨鞫鞬鞮鞶鞹鞾韃韅韉馱韍韎韔韖韘韝韞韡韣韭韮韱韹韺頀颳頄頇頊頍頎頏頒頖頞頠頫頬顱頯頲頴頼顇顋顑顒顓顔顕顚顜顢顣顬顳颭颮颱颶颸颺颻颽颾颿飀飂飈飌飜飡飣飤飥飩飫飮飱飶餀餂餄餎餇餈餑餔餕餖餗餚餛餜餟餠餤餧餩餪餫餬餮餱餲餳餺餻餼餽餿饁饅饇饉饊饍饎饐饘饟饢馘馥馝馡馣騮騾馵馹駃駄駅駆駉駋駑駓駔駗駘駙駜駡駢駪駬駰駴駸駹駽駾騂騄騅騆騉騋騍騏驎騑騒験騕騖騠騢騣騤騧驤騵騶騸騺驀驂驃驄驆驈驊驌驍驎驏驒驔驖驙驦驩驫骺鯁骫骭骯骱骴骶骷髏骾髁髂髄髆髈髐髑髕髖髙髝髞髟髡髣髧髪髫髭髯髲髳髹髺髽髾鬁鬃鬅鬈鬋鬎鬏鬐鬑鬒鬖鬗鬘鬙鬠鬣鬪鬫鬬鬮鬯鬰鬲鬵鬷魆魈魊魋魍魎魑魖鰾魛魟魣魦魨魬魴魵魸鮀鮁鮆鮌鮎鮑鮒鮓鮚鮞鮟鱇鮠鮦鮨鮪鮭鮶鮸鮿鯀鯄鯆鯇鯈鯔鯕鯖鯗鯙鯠鯤鯥鯫鯰鯷鯸鯿鰂鰆鶼鰉鰋鰐鰒鰕鰛鰜鰣鰤鰥鰦鰨鰩鰮鰳鰶鰷鱺鰼鰽鱀鱄鱅鱆鱈鱎鱐鱓鱔鱖鱘鱟鱠鱣鱨鱭鱮鱲鱵鱻鲅鳦鳧鳯鳲鳷鳻鴂鴃鴄鴆鴈鴎鴒鴔鴗鴛鴦鴝鵒鴟鴠鴢鴣鴥鴯鶓鴳鴴鴷鴽鵀鵁鵂鵓鵖鵙鵜鶘鵞鵟鵩鵪鵫鵵鵷鵻鵾鶂鶊鶏鶒鶖鶗鶡鶤鶦鶬鶱鶲鶵鶸鶹鶺鶿鷀鷁鷃鷄鷇鷈鷉鷊鷏鷓鷕鷖鷙鷞鷟鷥鷦鷯鷩鷫鷭鷳鷴鷽鷾鷿鸂鸇鸊鸏鸑鸒鸓鸕鸛鸜鸝鹸鹹鹺麀麂麃麄麇麋麌麐麑麒麚麛麝麤麩麪麫麮麯麰麺麾黁黈黌黢黒黓黕黙黝黟黥黦黧黮黰黱黲黶黹黻黼黽黿鼂鼃鼅鼈鼉鼏鼐鼒鼕鼖鼙鼚鼛鼡鼩鼱鼪鼫鼯鼷鼽齁齆齇齈齉齌齎齏齔齕齗齙齚齜齞齟齬齠齢齣齧齩齮齯齰齱齵齾龎龑龒龔龖龘龝龡龢龤'
+
+assert len(simplified_charcters) == len(simplified_charcters)
+
+s2t_dict = {}
+t2s_dict = {}
+for i, item in enumerate(simplified_charcters):
+    s2t_dict[item] = traditional_characters[i]
+    t2s_dict[traditional_characters[i]] = item
+
+
+def tranditional_to_simplified(text: str) -> str:
+    return "".join(
+        [t2s_dict[item] if item in t2s_dict else item for item in text])
+
+
+def simplified_to_traditional(text: str) -> str:
+    return "".join(
+        [s2t_dict[item] if item in s2t_dict else item for item in text])
+
+
+if __name__ == "__main__":
+    text = "一般是指存取一個應用程式啟動時始終顯示在網站或網頁瀏覽器中的一個或多個初始網頁等畫面存在的站點"
+    print(text)
+    text_simple = tranditional_to_simplified(text)
+    print(text_simple)
+    text_traditional = simplified_to_traditional(text_simple)
+    print(text_traditional)

SongBloom/g2p/cn_zh_g2p/zh_normalization/chronology.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+import re
+
+from .num import DIGITS
+from .num import num2str
+from .num import verbalize_cardinal
+from .num import verbalize_digit
+
+
+def _time_num2str(num_string: str) -> str:
+    """A special case for verbalizing number in time."""
+    result = num2str(num_string.lstrip('0'))
+    if num_string.startswith('0'):
+        result = DIGITS['0'] + result
+    return result
+
+
+# 时刻表达式
+RE_TIME = re.compile(r'([0-1]?[0-9]|2[0-3])'
+                     r':([0-5][0-9])'
+                     r'(:([0-5][0-9]))?')
+
+# 时间范围，如8:30-12:30
+RE_TIME_RANGE = re.compile(r'([0-1]?[0-9]|2[0-3])'
+                           r':([0-5][0-9])'
+                           r'(:([0-5][0-9]))?'
+                           r'(~|-)'
+                           r'([0-1]?[0-9]|2[0-3])'
+                           r':([0-5][0-9])'
+                           r'(:([0-5][0-9]))?')
+
+
+def replace_time(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+
+    is_range = len(match.groups()) > 5
+
+    hour = match.group(1)
+    minute = match.group(2)
+    second = match.group(4)
+
+    if is_range:
+        hour_2 = match.group(6)
+        minute_2 = match.group(7)
+        second_2 = match.group(9)
+
+    result = f"{num2str(hour)}点"
+    if minute.lstrip('0'):
+        if int(minute) == 30:
+            result += "半"
+        else:
+            result += f"{_time_num2str(minute)}分"
+    if second and second.lstrip('0'):
+        result += f"{_time_num2str(second)}秒"
+
+    if is_range:
+        result += "至"
+        result += f"{num2str(hour_2)}点"
+        if minute_2.lstrip('0'):
+            if int(minute) == 30:
+                result += "半"
+            else:
+                result += f"{_time_num2str(minute_2)}分"
+        if second_2 and second_2.lstrip('0'):
+            result += f"{_time_num2str(second_2)}秒"
+
+    return result
+
+
+RE_DATE = re.compile(r'(\d{4}|\d{2})年'
+                     r'((0?[1-9]|1[0-2])月)?'
+                     r'(((0?[1-9])|((1|2)[0-9])|30|31)([日号]))?')
+
+
+def replace_date(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    year = match.group(1)
+    month = match.group(3)
+    day = match.group(5)
+    result = ""
+    if year:
+        result += f"{verbalize_digit(year)}年"
+    if month:
+        result += f"{verbalize_cardinal(month)}月"
+    if day:
+        result += f"{verbalize_cardinal(day)}{match.group(9)}"
+    return result
+
+
+# 用 / 或者 - 分隔的 YY/MM/DD 或者 YY-MM-DD 日期
+RE_DATE2 = re.compile(
+    r'(\d{4})([- /.])(0[1-9]|1[012])\2(0[1-9]|[12][0-9]|3[01])')
+
+
+def replace_date2(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    year = match.group(1)
+    month = match.group(3)
+    day = match.group(4)
+    result = ""
+    if year:
+        result += f"{verbalize_digit(year)}年"
+    if month:
+        result += f"{verbalize_cardinal(month)}月"
+    if day:
+        result += f"{verbalize_cardinal(day)}日"
+    return result

SongBloom/g2p/cn_zh_g2p/zh_normalization/constants.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+import re
+import string
+
+from pypinyin.constants import SUPPORT_UCS4
+
+# 全角半角转换
+# 英文字符全角 -> 半角映射表 (num: 52)
+F2H_ASCII_LETTERS = {
+    ord(char) + 65248: ord(char)
+    for char in string.ascii_letters
+}
+
+# 英文字符半角 -> 全角映射表
+H2F_ASCII_LETTERS = {value: key for key, value in F2H_ASCII_LETTERS.items()}
+
+# 数字字符全角 -> 半角映射表 (num: 10)
+F2H_DIGITS = {ord(char) + 65248: ord(char) for char in string.digits}
+# 数字字符半角 -> 全角映射表
+H2F_DIGITS = {value: key for key, value in F2H_DIGITS.items()}
+
+# 标点符号全角 -> 半角映射表 (num: 32)
+F2H_PUNCTUATIONS = {ord(char) + 65248: ord(char) for char in string.punctuation}
+# 标点符号半角 -> 全角映射表
+H2F_PUNCTUATIONS = {value: key for key, value in F2H_PUNCTUATIONS.items()}
+
+# 空格 (num: 1)
+F2H_SPACE = {'\u3000': ' '}
+H2F_SPACE = {' ': '\u3000'}
+
+# 非"有拼音的汉字"的字符串，可用于NSW提取
+if SUPPORT_UCS4:
+    RE_NSW = re.compile(r'(?:[^'
+                        r'\u3007'  # 〇
+                        r'\u3400-\u4dbf'  # CJK扩展A:[3400-4DBF]
+                        r'\u4e00-\u9fff'  # CJK基本:[4E00-9FFF]
+                        r'\uf900-\ufaff'  # CJK兼容:[F900-FAFF]
+                        r'\U00020000-\U0002A6DF'  # CJK扩展B:[20000-2A6DF]
+                        r'\U0002A703-\U0002B73F'  # CJK扩展C:[2A700-2B73F]
+                        r'\U0002B740-\U0002B81D'  # CJK扩展D:[2B740-2B81D]
+                        r'\U0002F80A-\U0002FA1F'  # CJK兼容扩展:[2F800-2FA1F]
+                        r'])+')
+else:
+    RE_NSW = re.compile(  # pragma: no cover
+        r'(?:[^'
+        r'\u3007'  # 〇
+        r'\u3400-\u4dbf'  # CJK扩展A:[3400-4DBF]
+        r'\u4e00-\u9fff'  # CJK基本:[4E00-9FFF]
+        r'\uf900-\ufaff'  # CJK兼容:[F900-FAFF]
+        r'])+')

SongBloom/g2p/cn_zh_g2p/zh_normalization/num.py

@@ -0,0 +1,282 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+"""
+Rules to verbalize numbers into Chinese characters.
+https://zh.wikipedia.org/wiki/中文数字#現代中文
+"""
+import re
+from collections import OrderedDict
+from typing import List
+
+DIGITS = {str(i): tran for i, tran in enumerate('零一二三四五六七八九')}
+UNITS = OrderedDict({
+    1: '十',
+    2: '百',
+    3: '千',
+    4: '万',
+    8: '亿',
+})
+
+COM_QUANTIFIERS = '(封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)'
+
+# 分数表达式
+RE_FRAC = re.compile(r'(-?)(\d+)/(\d+)')
+
+
+def replace_frac(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    sign = match.group(1)
+    nominator = match.group(2)
+    denominator = match.group(3)
+    sign: str = "负" if sign else ""
+    nominator: str = num2str(nominator)
+    denominator: str = num2str(denominator)
+    result = f"{sign}{denominator}分之{nominator}"
+    return result
+
+
+# 百分数表达式
+RE_PERCENTAGE = re.compile(r'(-?)(\d+(\.\d+)?)%')
+
+
+def replace_percentage(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    sign = match.group(1)
+    percent = match.group(2)
+    sign: str = "负" if sign else ""
+    percent: str = num2str(percent)
+    result = f"{sign}百分之{percent}"
+    return result
+
+
+# 整数表达式
+# 带负号的整数 -10
+RE_INTEGER = re.compile(r'(-)' r'(\d+)')
+
+
+def replace_negative_num(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    sign = match.group(1)
+    number = match.group(2)
+    sign: str = "负" if sign else ""
+    number: str = num2str(number)
+    result = f"{sign}{number}"
+    return result
+
+
+# 编号-无符号整形
+# 00078
+RE_DEFAULT_NUM = re.compile(r'\d{3}\d*')
+
+
+def replace_default_num(match):
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    number = match.group(0)
+    return verbalize_digit(number, alt_one=True)
+
+
+# 加减乘除
+RE_ASMD = re.compile(
+    r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))([\+\-\×÷=])((-?)((\d+)(\.\d+)?)|(\.(\d+)))')
+asmd_map = {
+    '+': '加',
+    '-': '减',
+    '×': '乘',
+    '÷': '除',
+    '=': '等于'
+}
+
+
+def replace_asmd(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    result = match.group(1) + asmd_map[match.group(8)] + match.group(9)
+    return result
+
+
+# 数字表达式
+# 纯小数
+RE_DECIMAL_NUM = re.compile(r'(-?)((\d+)(\.\d+))' r'|(\.(\d+))')
+# 正整数 + 量词
+RE_POSITIVE_QUANTIFIERS = re.compile(r"(\d+)([多余几\+])?" + COM_QUANTIFIERS)
+RE_NUMBER = re.compile(r'(-?)((\d+)(\.\d+)?)' r'|(\.(\d+))')
+
+
+def replace_positive_quantifier(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    number = match.group(1)
+    match_2 = match.group(2)
+    if match_2 == "+":
+        match_2 = "多"
+    match_2: str = match_2 if match_2 else ""
+    quantifiers: str = match.group(3)
+    number: str = num2str(number)
+    result = f"{number}{match_2}{quantifiers}"
+    return result
+
+
+def replace_number(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    sign = match.group(1)
+    number = match.group(2)
+    pure_decimal = match.group(5)
+    if pure_decimal:
+        result = num2str(pure_decimal)
+    else:
+        sign: str = "负" if sign else ""
+        number: str = num2str(number)
+        result = f"{sign}{number}"
+    return result
+
+
+# 范围表达式
+# match.group(1) and match.group(8) are copy from RE_NUMBER
+
+RE_RANGE = re.compile(
+    r"""
+    (?<![\d\+\-\×÷=])      # 使用反向前瞻以确保数字范围之前没有其他数字和操作符
+    ((-?)((\d+)(\.\d+)?))  # 匹配范围起始的负数或正数（整数或小数）
+    [-~]                   # 匹配范围分隔符
+    ((-?)((\d+)(\.\d+)?))  # 匹配范围结束的负数或正数（整数或小数）
+    (?![\d\+\-\×÷=])       # 使用正向前瞻以确保数字范围之后没有其他数字和操作符
+    """, re.VERBOSE)
+
+
+def replace_range(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    first, second = match.group(1), match.group(6)
+    first = RE_NUMBER.sub(replace_number, first)
+    second = RE_NUMBER.sub(replace_number, second)
+    result = f"{first}到{second}"
+    return result
+
+
+# ~至表达式
+RE_TO_RANGE = re.compile(
+    r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)[~]((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)')
+
+def replace_to_range(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    result = match.group(0).replace('~', '至')
+    return result
+
+
+def _get_value(value_string: str, use_zero: bool=True) -> List[str]:
+    stripped = value_string.lstrip('0')
+    if len(stripped) == 0:
+        return []
+    elif len(stripped) == 1:
+        if use_zero and len(stripped) < len(value_string):
+            return [DIGITS['0'], DIGITS[stripped]]
+        else:
+            return [DIGITS[stripped]]
+    else:
+        largest_unit = next(
+            power for power in reversed(UNITS.keys()) if power < len(stripped))
+        first_part = value_string[:-largest_unit]
+        second_part = value_string[-largest_unit:]
+        return _get_value(first_part) + [UNITS[largest_unit]] + _get_value(
+            second_part)
+
+
+def verbalize_cardinal(value_string: str) -> str:
+    if not value_string:
+        return ''
+
+    # 000 -> '零' , 0 -> '零'
+    value_string = value_string.lstrip('0')
+    if len(value_string) == 0:
+        return DIGITS['0']
+
+    result_symbols = _get_value(value_string)
+    # verbalized number starting with '一十*' is abbreviated as `十*`
+    if len(result_symbols) >= 2 and result_symbols[0] == DIGITS[
+            '1'] and result_symbols[1] == UNITS[1]:
+        result_symbols = result_symbols[1:]
+    return ''.join(result_symbols)
+
+
+def verbalize_digit(value_string: str, alt_one=False) -> str:
+    result_symbols = [DIGITS[digit] for digit in value_string]
+    result = ''.join(result_symbols)
+    if alt_one:
+        result = result.replace("一", "幺")
+    return result
+
+
+def num2str(value_string: str) -> str:
+    integer_decimal = value_string.split('.')
+    if len(integer_decimal) == 1:
+        integer = integer_decimal[0]
+        decimal = ''
+    elif len(integer_decimal) == 2:
+        integer, decimal = integer_decimal
+    else:
+        raise ValueError(
+            f"The value string: '${value_string}' has more than one point in it."
+        )
+
+    result = verbalize_cardinal(integer)
+
+    decimal = decimal.rstrip('0')
+    if decimal:
+        # '.22' is verbalized as '零点二二'
+        # '3.20' is verbalized as '三点二
+        result = result if result else "零"
+        result += '点' + verbalize_digit(decimal)
+    return result

SongBloom/g2p/cn_zh_g2p/zh_normalization/phonecode.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+import re
+
+from .num import verbalize_digit
+
+# 规范化固话/手机号码
+# 手机
+# http://www.jihaoba.com/news/show/13680
+# 移动：139、138、137、136、135、134、159、158、157、150、151、152、188、187、182、183、184、178、198
+# 联通：130、131、132、156、155、186、185、176
+# 电信：133、153、189、180、181、177
+RE_MOBILE_PHONE = re.compile(
+    r"(?<!\d)((\+?86 ?)?1([38]\d|5[0-35-9]|7[678]|9[89])\d{8})(?!\d)")
+RE_TELEPHONE = re.compile(
+    r"(?<!\d)((0(10|2[1-3]|[3-9]\d{2})-?)?[1-9]\d{6,7})(?!\d)")
+
+# 全国统一的号码400开头
+RE_NATIONAL_UNIFORM_NUMBER = re.compile(r"(400)(-)?\d{3}(-)?\d{4}")
+
+
+def phone2str(phone_string: str, mobile=True) -> str:
+    if mobile:
+        sp_parts = phone_string.strip('+').split()
+        result = '，'.join(
+            [verbalize_digit(part, alt_one=True) for part in sp_parts])
+        return result
+    else:
+        sil_parts = phone_string.split('-')
+        result = '，'.join(
+            [verbalize_digit(part, alt_one=True) for part in sil_parts])
+        return result
+
+
+def replace_phone(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    return phone2str(match.group(0), mobile=False)
+
+
+def replace_mobile(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    return phone2str(match.group(0))

SongBloom/g2p/cn_zh_g2p/zh_normalization/quantifier.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+import re
+
+from .num import num2str
+
+# 温度表达式，温度会影响负号的读法
+# -3°C 零下三度
+RE_TEMPERATURE = re.compile(r'(-?)(\d+(\.\d+)?)(°C|℃|度|摄氏度)')
+measure_dict = {
+    "cm2": "平方厘米",
+    "cm²": "平方厘米",
+    "cm3": "立方厘米",
+    "cm³": "立方厘米",
+    "cm": "厘米",
+    "db": "分贝",
+    "ds": "毫秒",
+    "kg": "千克",
+    "km": "千米",
+    "m2": "平方米",
+    "m²": "平方米",
+    "m³": "立方米",
+    "m3": "立方米",
+    "ml": "毫升",
+    "m": "米",
+    "mm": "毫米",
+    "s": "秒"
+}
+
+
+def replace_temperature(match) -> str:
+    """
+    Args:
+        match (re.Match)
+    Returns:
+        str
+    """
+    sign = match.group(1)
+    temperature = match.group(2)
+    unit = match.group(3)
+    sign: str = "零下" if sign else ""
+    temperature: str = num2str(temperature)
+    unit: str = "摄氏度" if unit == "摄氏度" else "度"
+    result = f"{sign}{temperature}{unit}"
+    return result
+
+
+def replace_measure(sentence) -> str:
+    for q_notation in measure_dict:
+        if q_notation in sentence:
+            sentence = sentence.replace(q_notation, measure_dict[q_notation])
+    return sentence

SongBloom/g2p/cn_zh_g2p/zh_normalization/text_normlization.py

@@ -0,0 +1,165 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+import re
+from typing import List
+
+from .char_convert import tranditional_to_simplified
+from .chronology import RE_DATE
+from .chronology import RE_DATE2
+from .chronology import RE_TIME
+from .chronology import RE_TIME_RANGE
+from .chronology import replace_date
+from .chronology import replace_date2
+from .chronology import replace_time
+from .constants import F2H_ASCII_LETTERS
+from .constants import F2H_DIGITS
+from .constants import F2H_SPACE
+from .num import RE_DECIMAL_NUM
+from .num import RE_DEFAULT_NUM
+from .num import RE_FRAC
+from .num import RE_INTEGER
+from .num import RE_NUMBER
+from .num import RE_PERCENTAGE
+from .num import RE_POSITIVE_QUANTIFIERS
+from .num import RE_RANGE
+from .num import RE_TO_RANGE
+from .num import RE_ASMD
+from .num import replace_default_num
+from .num import replace_frac
+from .num import replace_negative_num
+from .num import replace_number
+from .num import replace_percentage
+from .num import replace_positive_quantifier
+from .num import replace_range
+from .num import replace_to_range
+from .num import replace_asmd
+from .phonecode import RE_MOBILE_PHONE
+from .phonecode import RE_NATIONAL_UNIFORM_NUMBER
+from .phonecode import RE_TELEPHONE
+from .phonecode import replace_mobile
+from .phonecode import replace_phone
+from .quantifier import RE_TEMPERATURE
+from .quantifier import replace_measure
+from .quantifier import replace_temperature
+
+
+class TextNormalizer():
+    def __init__(self):
+        self.SENTENCE_SPLITOR = re.compile(r'([：、，；。？！,;?!][”’]?)')
+
+    def _split(self, text: str, lang="zh") -> List[str]:
+        """Split long text into sentences with sentence-splitting punctuations.
+        Args:
+            text (str): The input text.
+        Returns:
+            List[str]: Sentences.
+        """
+        # Only for pure Chinese here
+        if lang == "zh":
+            text = text.replace(" ", "")
+            # 过滤掉特殊字符
+            text = re.sub(r'[——《》【】<>{}()（）#&@“”^_|\\]', '', text)
+        text = self.SENTENCE_SPLITOR.sub(r'\1\n', text)
+        text = text.strip()
+        sentences = [sentence.strip() for sentence in re.split(r'\n+', text)]
+        return sentences
+
+    def _post_replace(self, sentence: str) -> str:
+        sentence = sentence.replace('/', '每')
+        # sentence = sentence.replace('~', '至')
+        # sentence = sentence.replace('～', '至')
+        sentence = sentence.replace('①', '一')
+        sentence = sentence.replace('②', '二')
+        sentence = sentence.replace('③', '三')
+        sentence = sentence.replace('④', '四')
+        sentence = sentence.replace('⑤', '五')
+        sentence = sentence.replace('⑥', '六')
+        sentence = sentence.replace('⑦', '七')
+        sentence = sentence.replace('⑧', '八')
+        sentence = sentence.replace('⑨', '九')
+        sentence = sentence.replace('⑩', '十')
+        sentence = sentence.replace('α', '阿尔法')
+        sentence = sentence.replace('β', '贝塔')
+        sentence = sentence.replace('γ', '伽玛').replace('Γ', '伽玛')
+        sentence = sentence.replace('δ', '德尔塔').replace('Δ', '德尔塔')
+        sentence = sentence.replace('ε', '艾普西龙')
+        sentence = sentence.replace('ζ', '捷塔')
+        sentence = sentence.replace('η', '依塔')
+        sentence = sentence.replace('θ', '西塔').replace('Θ', '西塔')
+        sentence = sentence.replace('ι', '艾欧塔')
+        sentence = sentence.replace('κ', '喀帕')
+        sentence = sentence.replace('λ', '拉姆达').replace('Λ', '拉姆达')
+        sentence = sentence.replace('μ', '缪')
+        sentence = sentence.replace('ν', '拗')
+        sentence = sentence.replace('ξ', '克西').replace('Ξ', '克西')
+        sentence = sentence.replace('ο', '欧米克伦')
+        sentence = sentence.replace('π', '派').replace('Π', '派')
+        sentence = sentence.replace('ρ', '肉')
+        sentence = sentence.replace('ς', '西格玛').replace('Σ', '西格玛').replace(
+            'σ', '西格玛')
+        sentence = sentence.replace('τ', '套')
+        sentence = sentence.replace('υ', '宇普西龙')
+        sentence = sentence.replace('φ', '服艾').replace('Φ', '服艾')
+        sentence = sentence.replace('χ', '器')
+        sentence = sentence.replace('ψ', '普赛').replace('Ψ', '普赛')
+        sentence = sentence.replace('ω', '欧米伽').replace('Ω', '欧米伽')
+        # re filter special characters, have one more character "-" than line 68
+        sentence = re.sub(r'[-——《》【】<=>{}()（）#&@“”^_|\\]', '', sentence)
+        return sentence
+
+    def normalize_sentence(self, sentence: str) -> str:
+        # basic character conversions
+        sentence = tranditional_to_simplified(sentence)
+        sentence = sentence.translate(F2H_ASCII_LETTERS).translate(
+            F2H_DIGITS).translate(F2H_SPACE)
+
+        # number related NSW verbalization
+        sentence = RE_DATE.sub(replace_date, sentence)
+        sentence = RE_DATE2.sub(replace_date2, sentence)
+
+        # range first
+        sentence = RE_TIME_RANGE.sub(replace_time, sentence)
+        sentence = RE_TIME.sub(replace_time, sentence)
+
+        # 处理~波浪号作为至的替换
+        sentence = RE_TO_RANGE.sub(replace_to_range, sentence)
+        sentence = RE_TEMPERATURE.sub(replace_temperature, sentence)
+        sentence = replace_measure(sentence)
+        sentence = RE_FRAC.sub(replace_frac, sentence)
+        sentence = RE_PERCENTAGE.sub(replace_percentage, sentence)
+        sentence = RE_MOBILE_PHONE.sub(replace_mobile, sentence)
+
+        sentence = RE_TELEPHONE.sub(replace_phone, sentence)
+        sentence = RE_NATIONAL_UNIFORM_NUMBER.sub(replace_phone, sentence)
+
+        sentence = RE_RANGE.sub(replace_range, sentence)
+
+        # 处理加减乘除
+        while RE_ASMD.search(sentence):
+            sentence = RE_ASMD.sub(replace_asmd, sentence)
+
+        sentence = RE_INTEGER.sub(replace_negative_num, sentence)
+        sentence = RE_DECIMAL_NUM.sub(replace_number, sentence)
+        sentence = RE_POSITIVE_QUANTIFIERS.sub(replace_positive_quantifier,
+                                               sentence)
+        sentence = RE_DEFAULT_NUM.sub(replace_default_num, sentence)
+        sentence = RE_NUMBER.sub(replace_number, sentence)
+        sentence = self._post_replace(sentence)
+
+        return sentence
+
+    def normalize(self, text: str) -> List[str]:
+        sentences = self._split(text)
+        sentences = [self.normalize_sentence(sent) for sent in sentences]
+        return sentences

SongBloom/g2p/lyric_common.py

@@ -0,0 +1,81 @@
+import os, sys
+
+sys.path.insert(0, os.path.dirname(__file__))
+from pinyin.pinyin import G2P_PinYin
+from cn_zh_g2p import G2P_Mix, symbols
+
+key2processor = {
+    'pinyin': G2P_PinYin(), 
+    'phoneme': G2P_Mix(),
+}
+
+valid_struct_type = ['[chorus]', '[verse]', '[bridge]']
+start_struct_type = ['[intro]', '[start]']
+end_struct_type = ['[outro]', '[end]']
+conn_struct_type = ['[inst]', '[solo]', '[break]']
+
+LABELS = {
+    '[intro]': 0,
+    '[outro]': 1,
+    '[bridge]': 2,
+    '[inst]': 3,
+    '[verse]': 4,
+    '[chorus]': 5,
+    '[silence]': 6,
+}
+
+NUMBERS = {
+    '0': ['零', 'zero'],
+    '1': ['一', 'one'],
+    '2': ['二', 'two'],
+    '3': ['三', 'three'],
+    '4': ['四', 'four'],
+    '5': ['五', 'five'],
+    '6': ['六', 'six'],
+    '7': ['七', 'seven'],
+    '8': ['八', 'eight'],
+    '9': ['九', 'nine']
+}
+
+def detect_structure(structure):
+    valid_start = ['start', 'intro']
+    valid_end = ['outro', 'end']
+    valid_instru = ['solo', 'inst', 'break']
+    valid_bridge = ['bridge']
+
+    if structure in ['verse', 'chorus', 'silence']:
+        return structure
+    
+    if structure in valid_start:
+        return 'intro'
+    if structure in valid_end:
+        return 'outro'
+    if structure in valid_instru:
+        return 'inst'
+    if structure in valid_bridge:
+        return 'bridge'
+
+def merge_structure(start_time, end_time, structure, lyric):
+    cnt = 1
+    while cnt < len(start_time):
+        if structure[cnt] == structure[cnt-1]:
+            end_time[cnt-1] = end_time[cnt]
+            if structure[cnt] not in ["verse", "chorus", "bridge"]:
+                del start_time[cnt]
+                del end_time[cnt]
+                del structure[cnt]
+                del lyric[cnt]
+            else:
+                cnt += 1
+        else:
+            cnt += 1
+    
+    return start_time, end_time, structure, lyric
+
+
+def is_struct_legal(struct, text):
+    if struct in valid_struct_type and text != "":
+        return True
+    elif struct not in valid_struct_type and text == "":
+        return True
+    return False

SongBloom/g2p/pinyin/__init__.py

@@ -0,0 +1,430 @@
+from .symbols import symbols
+
+
+
+
+pinyin_dict = {
+    "a": ("^", "a"),
+    "ai": ("^", "ai"),
+    "an": ("^", "an"),
+    "ang": ("^", "ang"),
+    "ao": ("^", "ao"),
+    "ba": ("b", "a"),
+    "bai": ("b", "ai"),
+    "ban": ("b", "an"),
+    "bang": ("b", "ang"),
+    "bao": ("b", "ao"),
+    "be": ("b", "e"),
+    "bei": ("b", "ei"),
+    "ben": ("b", "en"),
+    "beng": ("b", "eng"),
+    "bi": ("b", "i"),
+    "bian": ("b", "ian"),
+    "biao": ("b", "iao"),
+    "bie": ("b", "ie"),
+    "bin": ("b", "in"),
+    "bing": ("b", "ing"),
+    "bo": ("b", "o"),
+    "bu": ("b", "u"),
+    "ca": ("c", "a"),
+    "cai": ("c", "ai"),
+    "can": ("c", "an"),
+    "cang": ("c", "ang"),
+    "cao": ("c", "ao"),
+    "ce": ("c", "e"),
+    "cen": ("c", "en"),
+    "ceng": ("c", "eng"),
+    "cha": ("ch", "a"),
+    "chai": ("ch", "ai"),
+    "chan": ("ch", "an"),
+    "chang": ("ch", "ang"),
+    "chao": ("ch", "ao"),
+    "che": ("ch", "e"),
+    "chen": ("ch", "en"),
+    "cheng": ("ch", "eng"),
+    "chi": ("ch", "iii"),
+    "chong": ("ch", "ong"),
+    "chou": ("ch", "ou"),
+    "chu": ("ch", "u"),
+    "chua": ("ch", "ua"),
+    "chuai": ("ch", "uai"),
+    "chuan": ("ch", "uan"),
+    "chuang": ("ch", "uang"),
+    "chui": ("ch", "uei"),
+    "chun": ("ch", "uen"),
+    "chuo": ("ch", "uo"),
+    "ci": ("c", "ii"),
+    "cong": ("c", "ong"),
+    "cou": ("c", "ou"),
+    "cu": ("c", "u"),
+    "cuan": ("c", "uan"),
+    "cui": ("c", "uei"),
+    "cun": ("c", "uen"),
+    "cuo": ("c", "uo"),
+    "da": ("d", "a"),
+    "dai": ("d", "ai"),
+    "dan": ("d", "an"),
+    "dang": ("d", "ang"),
+    "dao": ("d", "ao"),
+    "de": ("d", "e"),
+    "dei": ("d", "ei"),
+    "den": ("d", "en"),
+    "deng": ("d", "eng"),
+    "di": ("d", "i"),
+    "dia": ("d", "ia"),
+    "dian": ("d", "ian"),
+    "diao": ("d", "iao"),
+    "die": ("d", "ie"),
+    "ding": ("d", "ing"),
+    "diu": ("d", "iou"),
+    "dong": ("d", "ong"),
+    "dou": ("d", "ou"),
+    "du": ("d", "u"),
+    "duan": ("d", "uan"),
+    "dui": ("d", "uei"),
+    "dun": ("d", "uen"),
+    "duo": ("d", "uo"),
+    "e": ("^", "e"),
+    "ei": ("^", "ei"),
+    "en": ("^", "en"),
+    "ng": ("^", "en"),
+    "eng": ("^", "eng"),
+    "er": ("^", "er"),
+    "fa": ("f", "a"),
+    "fan": ("f", "an"),
+    "fang": ("f", "ang"),
+    "fei": ("f", "ei"),
+    "fen": ("f", "en"),
+    "feng": ("f", "eng"),
+    "fo": ("f", "o"),
+    "fou": ("f", "ou"),
+    "fu": ("f", "u"),
+    "ga": ("g", "a"),
+    "gai": ("g", "ai"),
+    "gan": ("g", "an"),
+    "gang": ("g", "ang"),
+    "gao": ("g", "ao"),
+    "ge": ("g", "e"),
+    "gei": ("g", "ei"),
+    "gen": ("g", "en"),
+    "geng": ("g", "eng"),
+    "gong": ("g", "ong"),
+    "gou": ("g", "ou"),
+    "gu": ("g", "u"),
+    "gua": ("g", "ua"),
+    "guai": ("g", "uai"),
+    "guan": ("g", "uan"),
+    "guang": ("g", "uang"),
+    "gui": ("g", "uei"),
+    "gun": ("g", "uen"),
+    "guo": ("g", "uo"),
+    "ha": ("h", "a"),
+    "hai": ("h", "ai"),
+    "han": ("h", "an"),
+    "hang": ("h", "ang"),
+    "hao": ("h", "ao"),
+    "he": ("h", "e"),
+    "hei": ("h", "ei"),
+    "hen": ("h", "en"),
+    "heng": ("h", "eng"),
+    "hong": ("h", "ong"),
+    "hou": ("h", "ou"),
+    "hu": ("h", "u"),
+    "hua": ("h", "ua"),
+    "huai": ("h", "uai"),
+    "huan": ("h", "uan"),
+    "huang": ("h", "uang"),
+    "hui": ("h", "uei"),
+    "hun": ("h", "uen"),
+    "huo": ("h", "uo"),
+    "ji": ("j", "i"),
+    "jia": ("j", "ia"),
+    "jian": ("j", "ian"),
+    "jiang": ("j", "iang"),
+    "jiao": ("j", "iao"),
+    "jie": ("j", "ie"),
+    "jin": ("j", "in"),
+    "jing": ("j", "ing"),
+    "jiong": ("j", "iong"),
+    "jiu": ("j", "iou"),
+    "ju": ("j", "v"),
+    "juan": ("j", "van"),
+    "jue": ("j", "ve"),
+    "jun": ("j", "vn"),
+    "ka": ("k", "a"),
+    "kai": ("k", "ai"),
+    "kan": ("k", "an"),
+    "kang": ("k", "ang"),
+    "kao": ("k", "ao"),
+    "ke": ("k", "e"),
+    "kei": ("k", "ei"),
+    "ken": ("k", "en"),
+    "keng": ("k", "eng"),
+    "kong": ("k", "ong"),
+    "kou": ("k", "ou"),
+    "ku": ("k", "u"),
+    "kua": ("k", "ua"),
+    "kuai": ("k", "uai"),
+    "kuan": ("k", "uan"),
+    "kuang": ("k", "uang"),
+    "kui": ("k", "uei"),
+    "kun": ("k", "uen"),
+    "kuo": ("k", "uo"),
+    "la": ("l", "a"),
+    "lai": ("l", "ai"),
+    "lan": ("l", "an"),
+    "lang": ("l", "ang"),
+    "lao": ("l", "ao"),
+    "le": ("l", "e"),
+    "lei": ("l", "ei"),
+    "leng": ("l", "eng"),
+    "li": ("l", "i"),
+    "lia": ("l", "ia"),
+    "lian": ("l", "ian"),
+    "liang": ("l", "iang"),
+    "liao": ("l", "iao"),
+    "lie": ("l", "ie"),
+    "lin": ("l", "in"),
+    "ling": ("l", "ing"),
+    "liu": ("l", "iou"),
+    "lo": ("l", "o"),
+    "long": ("l", "ong"),
+    "lou": ("l", "ou"),
+    "lu": ("l", "u"),
+    "lv": ("l", "v"),
+    "luan": ("l", "uan"),
+    "lve": ("l", "ve"),
+    "lue": ("l", "ve"),
+    "lun": ("l", "uen"),
+    "luo": ("l", "uo"),
+    "ma": ("m", "a"),
+    "mai": ("m", "ai"),
+    "man": ("m", "an"),
+    "mang": ("m", "ang"),
+    "mao": ("m", "ao"),
+    "me": ("m", "e"),
+    "mei": ("m", "ei"),
+    "men": ("m", "en"),
+    "meng": ("m", "eng"),
+    "mi": ("m", "i"),
+    "mian": ("m", "ian"),
+    "miao": ("m", "iao"),
+    "mie": ("m", "ie"),
+    "min": ("m", "in"),
+    "ming": ("m", "ing"),
+    "miu": ("m", "iou"),
+    "mo": ("m", "o"),
+    "mou": ("m", "ou"),
+    "mu": ("m", "u"),
+    "na": ("n", "a"),
+    "nai": ("n", "ai"),
+    "nan": ("n", "an"),
+    "nang": ("n", "ang"),
+    "nao": ("n", "ao"),
+    "ne": ("n", "e"),
+    "nei": ("n", "ei"),
+    "nen": ("n", "en"),
+    "neng": ("n", "eng"),
+    "ni": ("n", "i"),
+    "nia": ("n", "ia"),
+    "nian": ("n", "ian"),
+    "niang": ("n", "iang"),
+    "niao": ("n", "iao"),
+    "nie": ("n", "ie"),
+    "nin": ("n", "in"),
+    "ning": ("n", "ing"),
+    "niu": ("n", "iou"),
+    "nong": ("n", "ong"),
+    "nou": ("n", "ou"),
+    "nu": ("n", "u"),
+    "nv": ("n", "v"),
+    "nuan": ("n", "uan"),
+    "nve": ("n", "ve"),
+    "nue": ("n", "ve"),
+    "nuo": ("n", "uo"),
+    "o": ("^", "o"),
+    "ou": ("^", "ou"),
+    "pa": ("p", "a"),
+    "pai": ("p", "ai"),
+    "pan": ("p", "an"),
+    "pang": ("p", "ang"),
+    "pao": ("p", "ao"),
+    "pe": ("p", "e"),
+    "pei": ("p", "ei"),
+    "pen": ("p", "en"),
+    "peng": ("p", "eng"),
+    "pi": ("p", "i"),
+    "pian": ("p", "ian"),
+    "piao": ("p", "iao"),
+    "pie": ("p", "ie"),
+    "pin": ("p", "in"),
+    "ping": ("p", "ing"),
+    "po": ("p", "o"),
+    "pou": ("p", "ou"),
+    "pu": ("p", "u"),
+    "qi": ("q", "i"),
+    "qia": ("q", "ia"),
+    "qian": ("q", "ian"),
+    "qiang": ("q", "iang"),
+    "qiao": ("q", "iao"),
+    "qie": ("q", "ie"),
+    "qin": ("q", "in"),
+    "qing": ("q", "ing"),
+    "qiong": ("q", "iong"),
+    "qiu": ("q", "iou"),
+    "qu": ("q", "v"),
+    "quan": ("q", "van"),
+    "que": ("q", "ve"),
+    "qun": ("q", "vn"),
+    "ran": ("r", "an"),
+    "rang": ("r", "ang"),
+    "rao": ("r", "ao"),
+    "re": ("r", "e"),
+    "ren": ("r", "en"),
+    "reng": ("r", "eng"),
+    "ri": ("r", "iii"),
+    "rong": ("r", "ong"),
+    "rou": ("r", "ou"),
+    "ru": ("r", "u"),
+    "rua": ("r", "ua"),
+    "ruan": ("r", "uan"),
+    "rui": ("r", "uei"),
+    "run": ("r", "uen"),
+    "ruo": ("r", "uo"),
+    "sa": ("s", "a"),
+    "sai": ("s", "ai"),
+    "san": ("s", "an"),
+    "sang": ("s", "ang"),
+    "sao": ("s", "ao"),
+    "se": ("s", "e"),
+    "sen": ("s", "en"),
+    "seng": ("s", "eng"),
+    "sha": ("sh", "a"),
+    "shai": ("sh", "ai"),
+    "shan": ("sh", "an"),
+    "shang": ("sh", "ang"),
+    "shao": ("sh", "ao"),
+    "she": ("sh", "e"),
+    "shei": ("sh", "ei"),
+    "shen": ("sh", "en"),
+    "sheng": ("sh", "eng"),
+    "shi": ("sh", "iii"),
+    "shou": ("sh", "ou"),
+    "shu": ("sh", "u"),
+    "shua": ("sh", "ua"),
+    "shuai": ("sh", "uai"),
+    "shuan": ("sh", "uan"),
+    "shuang": ("sh", "uang"),
+    "shui": ("sh", "uei"),
+    "shun": ("sh", "uen"),
+    "shuo": ("sh", "uo"),
+    "si": ("s", "ii"),
+    "song": ("s", "ong"),
+    "sou": ("s", "ou"),
+    "su": ("s", "u"),
+    "suan": ("s", "uan"),
+    "sui": ("s", "uei"),
+    "sun": ("s", "uen"),
+    "suo": ("s", "uo"),
+    "ta": ("t", "a"),
+    "tai": ("t", "ai"),
+    "tan": ("t", "an"),
+    "tang": ("t", "ang"),
+    "tao": ("t", "ao"),
+    "te": ("t", "e"),
+    "tei": ("t", "ei"),
+    "teng": ("t", "eng"),
+    "ti": ("t", "i"),
+    "tian": ("t", "ian"),
+    "tiao": ("t", "iao"),
+    "tie": ("t", "ie"),
+    "ting": ("t", "ing"),
+    "tong": ("t", "ong"),
+    "tou": ("t", "ou"),
+    "tu": ("t", "u"),
+    "tuan": ("t", "uan"),
+    "tui": ("t", "uei"),
+    "tun": ("t", "uen"),
+    "tuo": ("t", "uo"),
+    "wa": ("^", "ua"),
+    "wai": ("^", "uai"),
+    "wan": ("^", "uan"),
+    "wang": ("^", "uang"),
+    "wei": ("^", "uei"),
+    "wen": ("^", "uen"),
+    "weng": ("^", "ueng"),
+    "wo": ("^", "uo"),
+    "wu": ("^", "u"),
+    "xi": ("x", "i"),
+    "xia": ("x", "ia"),
+    "xian": ("x", "ian"),
+    "xiang": ("x", "iang"),
+    "xiao": ("x", "iao"),
+    "xie": ("x", "ie"),
+    "xin": ("x", "in"),
+    "xing": ("x", "ing"),
+    "xiong": ("x", "iong"),
+    "xiu": ("x", "iou"),
+    "xu": ("x", "v"),
+    "xuan": ("x", "van"),
+    "xue": ("x", "ve"),
+    "xun": ("x", "vn"),
+    "ya": ("^", "ia"),
+    "yan": ("^", "ian"),
+    "yang": ("^", "iang"),
+    "yao": ("^", "iao"),
+    "ye": ("^", "ie"),
+    "yi": ("^", "i"),
+    "yin": ("^", "in"),
+    "ying": ("^", "ing"),
+    "yo": ("^", "iou"),
+    "yong": ("^", "iong"),
+    "you": ("^", "iou"),
+    "yu": ("^", "v"),
+    "yuan": ("^", "van"),
+    "yue": ("^", "ve"),
+    "yun": ("^", "vn"),
+    "za": ("z", "a"),
+    "zai": ("z", "ai"),
+    "zan": ("z", "an"),
+    "zang": ("z", "ang"),
+    "zao": ("z", "ao"),
+    "ze": ("z", "e"),
+    "zei": ("z", "ei"),
+    "zen": ("z", "en"),
+    "zeng": ("z", "eng"),
+    "zha": ("zh", "a"),
+    "zhai": ("zh", "ai"),
+    "zhan": ("zh", "an"),
+    "zhang": ("zh", "ang"),
+    "zhao": ("zh", "ao"),
+    "zhe": ("zh", "e"),
+    "zhei": ("zh", "ei"),
+    "zhen": ("zh", "en"),
+    "zheng": ("zh", "eng"),
+    "zhi": ("zh", "iii"),
+    "zhong": ("zh", "ong"),
+    "zhou": ("zh", "ou"),
+    "zhu": ("zh", "u"),
+    "zhua": ("zh", "ua"),
+    "zhuai": ("zh", "uai"),
+    "zhuan": ("zh", "uan"),
+    "zhuang": ("zh", "uang"),
+    "zhui": ("zh", "uei"),
+    "zhun": ("zh", "uen"),
+    "zhuo": ("zh", "uo"),
+    "zi": ("z", "ii"),
+    "zong": ("z", "ong"),
+    "zou": ("z", "ou"),
+    "zu": ("z", "u"),
+    "zuan": ("z", "uan"),
+    "zui": ("z", "uei"),
+    "zun": ("z", "uen"),
+    "zuo": ("z", "uo"),
+}
+
+
+def gen_vocabs():
+  import yaml
+  vocab = [f"<{c}{i}>" for c in list(pinyin_dict.keys()) for i in range(1,6)]
+  yaml.dump(vocab, open('./vocab.yaml', 'w'))

SongBloom/g2p/pinyin/pinyin.py

@@ -0,0 +1,137 @@
+import re
+
+from pypinyin import Style
+from pypinyin.contrib.neutral_tone import NeutralToneWith5Mixin
+from pypinyin.converter import DefaultConverter
+from pypinyin.core import Pinyin
+
+from . import pinyin_dict
+import torch
+
+
+class MyConverter(NeutralToneWith5Mixin, DefaultConverter):
+    pass
+
+
+def is_chinese(uchar):
+    if uchar >= u'\u4e00' and uchar <= u'\u9fa5':
+        return True
+    else:
+        return False
+
+
+def clean_chinese(text: str):
+    text = text.strip()
+    text_clean = []
+    for char in text:
+        if (is_chinese(char)):
+            text_clean.append(char)
+        else:
+            if len(text_clean) > 1 and is_chinese(text_clean[-1]):
+                text_clean.append(',')
+    text_clean = ''.join(text_clean).strip(',')
+    return text_clean
+
+
+class G2P_PinYin():
+
+    def __init__(self):
+        super(G2P_PinYin, self).__init__()
+        self.pinyin_parser = Pinyin(MyConverter())
+
+    def get_phoneme4pinyin(self, pinyins):
+        result = []
+        count_phone = []
+        for pinyin in pinyins:
+            if pinyin[:-1] in pinyin_dict:
+                tone = pinyin[-1]
+                a = pinyin[:-1]
+                a1, a2 = pinyin_dict[a]
+                result += [a1, a2 + tone]
+                count_phone.append(2)
+        return result, count_phone
+
+    # def chinese_to_phonemes(self, text):
+    #     text = clean_chinese(text)
+    #     phonemes = ["sil"]
+    #     chars = ['[PAD]']
+    #     all_pinyins = []
+    #     count_phone = []
+    #     count_phone.append(1)
+    #     for subtext in text.split(","):
+    #         if (len(subtext) == 0):
+    #             continue
+    #         pinyins = self.correct_pinyin_tone3(subtext)
+    #         all_pinyins.append(' '.join(pinyins))
+    #         sub_p, sub_c = self.get_phoneme4pinyin(pinyins)
+    #         phonemes.extend(sub_p)
+    #         phonemes.append(",")
+    #         count_phone.extend(sub_c)
+    #         count_phone.append(1)
+    #         chars.append(subtext)
+    #         chars.append(',')
+    #     phonemes.append("sil")
+    #     count_phone.append(1)
+    #     chars.append('[PAD]')
+    #     # char_embeds = self.prosody.expand_for_phone(char_embeds, count_phone)
+    #     return " ".join(phonemes), " ".join(chars), ' , '.join(all_pinyins)
+
+    def chinese_to_phonemes(self, text):
+        all_pinyins = []
+        subtext = []
+        for chr in text:
+            if is_chinese(chr):
+                subtext.append(chr)
+            else:
+                if subtext != []:
+                    subtext = ''.join(subtext)
+                    pinyins = self.correct_pinyin_tone3(subtext) 
+                    pinyins = [f"<{i}>" for i in pinyins]   
+                    all_pinyins.append(' '+ ' '.join(pinyins)+ ' ')
+                all_pinyins.append(chr)
+                subtext = []
+        if subtext != []:
+            subtext = ''.join(subtext)
+            pinyins = self.correct_pinyin_tone3(subtext)
+            pinyins = [f"<{i}>" for i in pinyins]     
+            all_pinyins.append(' '+ ' '.join(pinyins)+ ' ')
+        # char_embeds = self.prosody.expand_for_phone(char_embeds, count_phone)
+        return  ''.join(all_pinyins)
+    
+    def correct_pinyin_tone3(self, text):
+        pinyin_list = [
+            p[0]
+            for p in self.pinyin_parser.pinyin(text,
+                                               style=Style.TONE3,
+                                               strict=False,
+                                               neutral_tone_with_five=True)
+        ]
+        if len(pinyin_list) >= 2:
+            for i in range(1, len(pinyin_list)):
+                try:
+                    if re.findall(r'\d',
+                                  pinyin_list[i - 1])[0] == '3' and re.findall(
+                                      r'\d', pinyin_list[i])[0] == '3':
+                        pinyin_list[i - 1] = pinyin_list[i - 1].replace(
+                            '3', '2')
+                except IndexError:
+                    pass
+        return pinyin_list
+
+    # def expand_for_phone(self, char_embeds, length):  # length of phones for char
+    #     if(char_embeds.size(0) > len(length)):
+    #         print(char_embeds.shape, len(length))
+    #         char_embeds = char_embeds[0:len(length),:]
+    #     elif(char_embeds.size(0) < len(length)):
+    #         print(char_embeds.shape, len(length))
+    #         length = length[0:char_embeds.size(0)]
+    #     expand_vecs = list()
+    #     for vec, leng in zip(char_embeds, length):
+    #         vec = vec.expand(leng, -1)
+    #         expand_vecs.append(vec)
+    #     expand_embeds = torch.cat(expand_vecs, 0)
+    #     assert expand_embeds.size(0) == sum(length)
+    #     return expand_embeds.numpy()
+
+    def __call__(self, text):
+        return self.chinese_to_phonemes(text)

SongBloom/g2p/pinyin/symbols.py

@@ -0,0 +1,71 @@
+_pause = ["sil", "eos", "sp", "#0", "#1", "#2", "#3"]
+
+_initials = [
+    "^",
+    "b",
+    "c",
+    "ch",
+    "d",
+    "f",
+    "g",
+    "h",
+    "j",
+    "k",
+    "l",
+    "m",
+    "n",
+    "p",
+    "q",
+    "r",
+    "s",
+    "sh",
+    "t",
+    "x",
+    "z",
+    "zh",
+]
+
+_tones = ["1", "2", "3", "4", "5"]
+
+_finals = [
+    "a",
+    "ai",
+    "an",
+    "ang",
+    "ao",
+    "e",
+    "ei",
+    "en",
+    "eng",
+    "er",
+    "i",
+    "ia",
+    "ian",
+    "iang",
+    "iao",
+    "ie",
+    "ii",
+    "iii",
+    "in",
+    "ing",
+    "iong",
+    "iou",
+    "o",
+    "ong",
+    "ou",
+    "u",
+    "ua",
+    "uai",
+    "uan",
+    "uang",
+    "uei",
+    "uen",
+    "ueng",
+    "uo",
+    "v",
+    "van",
+    "ve",
+    "vn",
+]
+
+symbols = _pause + _initials + [i + j for i in _finals for j in _tones]

SongBloom/models/base/sample.py

@@ -0,0 +1,57 @@
+import torch
+
+
+def multinomial(input: torch.Tensor, num_samples: int, replacement=False, *, generator=None):
+    """torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension.
+
+    Args:
+        input (torch.Tensor): The input tensor containing probabilities.
+        num_samples (int): Number of samples to draw.
+        replacement (bool): Whether to draw with replacement or not.
+    Keywords args:
+        generator (torch.Generator): A pseudorandom number generator for sampling.
+    Returns:
+        torch.Tensor: Last dimension contains num_samples indices
+            sampled from the multinomial probability distribution
+            located in the last dimension of tensor input.
+    """
+    input_ = input.reshape(-1, input.shape[-1])
+    output_ = torch.multinomial(input_, num_samples=num_samples, replacement=replacement, generator=generator)
+    output = output_.reshape(*list(input.shape[:-1]), -1)
+    return output
+
+
+def sample_top_k(probs: torch.Tensor, k: int) -> torch.Tensor:
+    """Sample next token from top K values along the last dimension of the input probs tensor.
+
+    Args:
+        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
+        k (int): The k in “top-k”.
+    Returns:
+        torch.Tensor: Sampled tokens.
+    """
+    top_k_value, _ = torch.topk(probs, k, dim=-1)
+    min_value_top_k = top_k_value[..., [-1]]
+    probs *= (probs >= min_value_top_k).float()
+    probs.div_(probs.sum(dim=-1, keepdim=True))
+    next_token = multinomial(probs, num_samples=1)
+    return next_token
+
+
+def sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
+    """Sample next token from top P probabilities along the last dimension of the input probs tensor.
+
+    Args:
+        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
+        p (int): The p in “top-p”.
+    Returns:
+        torch.Tensor: Sampled tokens.
+    """
+    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    mask = probs_sum - probs_sort > p
+    probs_sort *= (~mask).float()
+    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+    next_token = multinomial(probs_sort, num_samples=1)
+    next_token = torch.gather(probs_idx, -1, next_token)
+    return next_token

SongBloom/models/base/utils.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+import typing as tp
+
+def length_to_mask(lengths: torch.Tensor, max_len: tp.Optional[int] = None) -> torch.Tensor:
+    """Utility function to convert a tensor of sequence lengths to a mask (useful when working on padded sequences).
+    For example: [3, 5] => [[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]]
+
+    Args:
+        lengths (torch.Tensor): tensor with lengths
+        max_len (int): can set the max length manually. Defaults to None.
+    Returns:
+        torch.Tensor: mask with 0s where there is pad tokens else 1s
+    """
+    assert len(lengths.shape) == 1, "Length shape should be 1 dimensional."
+    final_length = lengths.max().item() if not max_len else max_len
+    final_length = max(final_length, 1)  # if all seqs are of len zero we don't want a zero-size tensor
+    return torch.arange(final_length)[None, :].to(lengths.device) < lengths[:, None]
+
+
+def create_sin_embedding(positions: torch.Tensor, dim: int, max_period: float = 10000,
+                         dtype: torch.dtype = torch.float32) -> torch.Tensor:
+    """Create sinusoidal positional embedding, with shape `[B, T, C]`.
+
+    Args:
+        positions (torch.Tensor): LongTensor of positions.
+        dim (int): Dimension of the embedding.
+        max_period (float): Maximum period of the cosine/sine functions.
+        dtype (torch.dtype or str): dtype to use to generate the embedding.
+    Returns:
+        torch.Tensor: Sinusoidal positional embedding.
+    """
+    # We aim for BTC format
+    assert dim % 2 == 0
+    half_dim = dim // 2
+    positions = positions.to(dtype)
+    adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1)
+    max_period_tensor = torch.full([], max_period, device=positions.device, dtype=dtype)  # avoid sync point
+    phase = positions / (max_period_tensor ** (adim / (half_dim - 1)))
+    # phase = phase.to(torch.bfloat16)
+    return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
+
+
+def create_norm_fn(norm_type: str, dim: int, **kwargs) -> nn.Module:
+    """Create normalization module for transformer encoder layer.
+
+    Args:
+        norm_type (str): Normalization method.
+        dim (int): Dimension of the normalized layer.
+        **kwargs (dict): Additional parameters for normalization layer.
+    Returns:
+        nn.Module: Normalization module.
+    """
+    if norm_type == 'layer_norm':
+        return nn.LayerNorm(dim, eps=1e-5, **kwargs)
+    else:
+        raise ValueError(f"Unknown norm type: {norm_type}")

SongBloom/models/musicgen/conditioners/__init__.py

@@ -0,0 +1,37 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import omegaconf
+from .base import *
+from .text import *
+from .wav import *
+
+KLASS = {
+    'phoneme_tokenizer': PhonemeTokenizerConditioner,
+    'audio_tokenizer_wrapper': AudioTokenizerConditioner,
+}
+
+def get_condition_fuser(fuser_cfgs) -> ConditionFuser:
+    """Instantiate a condition fuser object."""
+    fuser_methods = ['sum', 'cross', 'prepend', 'input_interpolate']
+    fuse2cond = {k: fuser_cfgs[k] for k in fuser_methods}
+    kwargs = {k: v for k, v in fuser_cfgs.items() if k not in fuser_methods}
+    fuser = ConditionFuser(fuse2cond=fuse2cond, **kwargs)
+    return fuser
+
+def get_conditioner_provider(cfg) -> ConditioningProvider:
+    """Instantiate a conditioning model."""
+
+    dict_cfg = {} if cfg is None else dict(cfg)
+    conditioners: tp.Dict[str, BaseConditioner] = {}
+
+    # import pdb; pdb.set_trace()
+    for cond, cond_cfg in dict_cfg.items():
+        model_args = cond_cfg.copy()
+        model_type = model_args.pop('type')
+        conditioners[str(cond)] = KLASS[model_type](**model_args)
+    conditioner = ConditioningProvider(conditioners)
+    return conditioner

SongBloom/models/musicgen/conditioners/base.py

@@ -0,0 +1,872 @@
+from collections import defaultdict
+from copy import deepcopy
+from dataclasses import dataclass, field
+from itertools import chain
+import logging
+import typing as tp
+import einops
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+
+from dataclasses import dataclass, field, fields, replace
+
+from ..modules.streaming import StreamingModule
+from ...base.utils import length_to_mask, create_sin_embedding
+
+
+def collate(tensors: tp.List[torch.Tensor], dim: int = 0) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+    """Get a list of tensors and collate them to a single tensor. according to the following logic:
+    - `dim` specifies the time dimension which will be stacked and padded.
+    - The output will contain 1 new dimension (dimension index 0) which will be the size of
+    of the original list.
+
+    Args:
+        tensors (tp.List[torch.Tensor]): List of tensors to collate.
+        dim (int): Dimension which will be stacked and padded.
+    Returns:
+        tp.Tuple[torch.Tensor, torch.Tensor]:
+            torch.Tensor: Stacked and padded tensor. The output will contain 1 new dimension
+                (dimension index 0) which will be the size of the original list.
+            torch.Tensor: Tensor containing length of original tensor sizes (without padding).
+    """
+    tensors = [x.transpose(0, dim) for x in tensors]
+    lens = torch.LongTensor([len(x) for x in tensors])
+    padded_tensors = pad_sequence(tensors)
+    padded_tensors = padded_tensors.transpose(0, 1)
+    padded_tensors = padded_tensors.transpose(1, dim + 1)
+    return padded_tensors, lens
+
+
+
+@dataclass(order=True)
+class PathInZip:
+    """Hold a path of file within a zip file.
+
+    Args:
+        path (str): The convention is <path_to_zip>:<relative_path_inside_zip>.
+            Let's assume there is a zip file /some/location/foo.zip
+            and inside of it is a json file located at /data/file1.json,
+            Then we expect path = "/some/location/foo.zip:/data/file1.json".
+    """
+
+    INFO_PATH_SEP = ':'
+    zip_path: str
+    file_path: str
+
+    def __init__(self, path: str) -> None:
+        split_path = path.split(self.INFO_PATH_SEP)
+        assert len(split_path) == 2
+        self.zip_path, self.file_path = split_path
+
+    @classmethod
+    def from_paths(cls, zip_path: str, file_path: str):
+        return cls(zip_path + cls.INFO_PATH_SEP + file_path)
+
+    def __str__(self) -> str:
+        return self.zip_path + self.INFO_PATH_SEP + self.file_path
+
+
+@dataclass(order=True)
+class BaseInfo:
+
+    @classmethod
+    def _dict2fields(cls, dictionary: dict):
+        return {
+                field.name: dictionary[field.name]
+                for field in fields(cls) if field.name in dictionary
+            }
+        # try:
+        #     return {
+        #         field.name: dictionary[field.name]
+        #         for field in fields(cls) if field.name in dictionary
+        #     }
+        # except:
+        #     print(dictionary)
+
+    @classmethod
+    def from_dict(cls, dictionary: dict):
+        _dictionary = cls._dict2fields(dictionary)
+        return cls(**_dictionary)
+
+    def to_dict(self):
+        return {
+            field.name: self.__getattribute__(field.name)
+            for field in fields(self)
+            }
+
+
+@dataclass(order=True)
+class AudioMeta(BaseInfo):
+    path: str
+    duration: float
+    sample_rate: int
+    amplitude: tp.Optional[float] = None
+    weight: tp.Optional[float] = None
+    # info_path is used to load additional information about the audio file that is stored in zip files.
+    info_path: tp.Optional[PathInZip] = None
+
+    @classmethod
+    def from_dict(cls, dictionary: dict):
+        base = cls._dict2fields(dictionary)
+        if 'info_path' in base and base['info_path'] is not None:
+            base['info_path'] = PathInZip(base['info_path'])
+        return cls(**base)
+
+    def to_dict(self):
+        d = super().to_dict()
+        if d['info_path'] is not None:
+            d['info_path'] = str(d['info_path'])
+        return d
+
+
+@dataclass(order=True)
+class SegmentInfo(BaseInfo):
+    meta: AudioMeta
+    seek_time: float
+    # The following values are given once the audio is processed, e.g.
+    # at the target sample rate and target number of channels.
+    n_frames: int      # actual number of frames without padding
+    total_frames: int  # total number of frames, padding included
+    sample_rate: int   # actual sample rate
+    channels: int      # number of audio channels.
+
+
+logger = logging.getLogger(__name__)
+TextCondition = tp.Optional[str]  # a text condition can be a string or None (if doesn't exist)
+ConditionType = tp.Tuple[torch.Tensor, torch.Tensor]  # condition, mask
+
+
+class WavCondition(tp.NamedTuple):
+    wav: torch.Tensor
+    length: torch.Tensor
+    sample_rate: tp.List[int]
+    path: tp.List[tp.Optional[str]] = []
+    seek_time: tp.List[tp.Optional[float]] = []
+
+
+class JointEmbedCondition(tp.NamedTuple):
+    wav: torch.Tensor
+    text: tp.List[tp.Optional[str]]
+    length: torch.Tensor
+    sample_rate: tp.List[int]
+    path: tp.List[tp.Optional[str]] = []
+    seek_time: tp.List[tp.Optional[float]] = []
+
+
+@dataclass
+class ConditioningAttributes:
+    text: tp.Dict[str, tp.Optional[str]] = field(default_factory=dict)
+    wav: tp.Dict[str, WavCondition] = field(default_factory=dict)
+    joint_embed: tp.Dict[str, JointEmbedCondition] = field(default_factory=dict)
+
+    def __getitem__(self, item):
+        return getattr(self, item)
+
+    @property
+    def text_attributes(self):
+        return self.text.keys()
+
+    @property
+    def wav_attributes(self):
+        return self.wav.keys()
+
+    @property
+    def joint_embed_attributes(self):
+        return self.joint_embed.keys()
+
+    @property
+    def attributes(self):
+        return {
+            "text": self.text_attributes,
+            "wav": self.wav_attributes,
+            "joint_embed": self.joint_embed_attributes,
+        }
+
+    def to_flat_dict(self):
+        return {
+            **{f"text.{k}": v for k, v in self.text.items()},
+            **{f"wav.{k}": v for k, v in self.wav.items()},
+            **{f"joint_embed.{k}": v for k, v in self.joint_embed.items()}
+        }
+
+    @classmethod
+    def from_flat_dict(cls, x):
+        out = cls()
+        for k, v in x.items():
+            kind, att = k.split(".")
+            out[kind][att] = v
+        return out
+
+
+
+# class SegmentWithAttributes(SegmentInfo):
+#     """Base class for all dataclasses that are used for conditioning.
+#     All child classes should implement `to_condition_attributes` that converts
+#     the existing attributes to a dataclass of type ConditioningAttributes.
+#     """
+#     def to_condition_attributes(self) -> ConditioningAttributes:
+#         raise NotImplementedError()
+
+
+
+def nullify_condition(condition: ConditionType, dim: int = 1):
+    """Transform an input condition to a null condition.
+    The way it is done by converting it to a single zero vector similarly
+    to how it is done inside WhiteSpaceTokenizer and NoopTokenizer.
+
+    Args:
+        condition (ConditionType): A tuple of condition and mask (tuple[torch.Tensor, torch.Tensor])
+        dim (int): The dimension that will be truncated (should be the time dimension)
+        WARNING!: dim should not be the batch dimension!
+    Returns:
+        ConditionType: A tuple of null condition and mask
+    """
+    assert dim != 0, "dim cannot be the batch dimension!"
+    assert isinstance(condition, tuple) and \
+        isinstance(condition[0], torch.Tensor) and \
+        isinstance(condition[1], torch.Tensor), "'nullify_condition' got an unexpected input type!"
+    cond, mask = condition
+    B = cond.shape[0]
+    last_dim = cond.dim() - 1
+    out = cond.transpose(dim, last_dim)
+    out = 0. * out[..., :1]
+    out = out.transpose(dim, last_dim)
+    mask = torch.zeros((B, 1), device=out.device).int()
+    assert cond.dim() == out.dim()
+    return out, mask
+
+
+def nullify_wav(cond: WavCondition) -> WavCondition:
+    """Transform a WavCondition to a nullified WavCondition.
+    It replaces the wav by a null tensor, forces its length to 0, and replaces metadata by dummy attributes.
+
+    Args:
+        cond (WavCondition): Wav condition with wav, tensor of shape [B, T].
+    Returns:
+        WavCondition: Nullified wav condition.
+    """
+    #TODO by YCY, fix this to support zero-length input (as None)
+    null_wav, _ = nullify_condition((cond.wav, torch.zeros_like(cond.wav)), dim=cond.wav.dim() - 1) # B,1 all-zero
+    return WavCondition(
+        wav=null_wav,
+        length=torch.tensor([0] * cond.wav.shape[0], device=cond.wav.device),
+        sample_rate=cond.sample_rate,
+        path=[None] * cond.wav.shape[0],
+        seek_time=[None] * cond.wav.shape[0],
+    )
+
+
+def nullify_joint_embed(embed: JointEmbedCondition) -> JointEmbedCondition:
+    """Nullify the joint embedding condition by replacing it by a null tensor, forcing its length to 0,
+    and replacing metadata by dummy attributes.
+
+    Args:
+        cond (JointEmbedCondition): Joint embedding condition with wav and text, wav tensor of shape [B, C, T].
+    """
+    null_wav, _ = nullify_condition((embed.wav, torch.zeros_like(embed.wav)), dim=embed.wav.dim() - 1)
+    return JointEmbedCondition(
+        wav=null_wav, text=[None] * len(embed.text),
+        length=torch.LongTensor([0]).to(embed.wav.device),
+        sample_rate=embed.sample_rate,
+        path=[None] * embed.wav.shape[0],
+        seek_time=[0] * embed.wav.shape[0],
+    )
+
+
+
+class BaseConditioner(nn.Module):
+    """Base model for all conditioner modules.
+    We allow the output dim to be different than the hidden dim for two reasons:
+    1) keep our LUTs small when the vocab is large;
+    2) make all condition dims consistent.
+
+    Args:
+        dim (int): Hidden dim of the model.
+        output_dim (int): Output dim of the conditioner.
+    """
+    def __init__(self, dim: int, output_dim: int, input_token = False, padding_idx=None):
+        super().__init__()
+        self.dim = dim
+        self.output_dim = output_dim
+        if input_token:
+            self.output_proj = nn.Embedding(dim, output_dim, padding_idx)
+        else:
+            self.output_proj = nn.Linear(dim, output_dim)
+
+    def tokenize(self, *args, **kwargs) -> tp.Any:
+        """Should be any part of the processing that will lead to a synchronization
+        point, e.g. BPE tokenization with transfer to the GPU.
+
+        The returned value will be saved and return later when calling forward().
+        """
+        raise NotImplementedError()
+
+    def forward(self, inputs: tp.Any) -> ConditionType:
+        """Gets input that should be used as conditioning (e.g, genre, description or a waveform).
+        Outputs a ConditionType, after the input data was embedded as a dense vector.
+
+        Returns:
+            ConditionType:
+                - A tensor of size [B, T, D] where B is the batch size, T is the length of the
+                  output embedding and D is the dimension of the embedding.
+                - And a mask indicating where the padding tokens.
+        """
+        raise NotImplementedError()
+
+
+
+def dropout_condition(sample: ConditioningAttributes, condition_type: str, condition: str) -> ConditioningAttributes:
+    """Utility function for nullifying an attribute inside an ConditioningAttributes object.
+    If the condition is of type "wav", then nullify it using `nullify_condition` function.
+    If the condition is of any other type, set its value to None.
+    Works in-place.
+    """
+    if condition_type not in ['text', 'wav', 'joint_embed']:
+        raise ValueError(
+            "dropout_condition got an unexpected condition type!"
+            f" expected 'text', 'wav' or 'joint_embed' but got '{condition_type}'"
+        )
+
+    if condition not in getattr(sample, condition_type):
+        raise ValueError(
+            "dropout_condition received an unexpected condition!"
+            f" expected wav={sample.wav.keys()} and text={sample.text.keys()}"
+            f" but got '{condition}' of type '{condition_type}'!"
+        )
+
+    if condition_type == 'wav':
+        wav_cond = sample.wav[condition]
+        sample.wav[condition] = nullify_wav(wav_cond)
+    elif condition_type == 'joint_embed':
+        embed = sample.joint_embed[condition]
+        sample.joint_embed[condition] = nullify_joint_embed(embed)
+    else:
+        sample.text[condition] = None
+
+    return sample
+
+
+class DropoutModule(nn.Module):
+    """Base module for all dropout modules."""
+    def __init__(self, seed: int = 1234):
+        super().__init__()
+        self.rng = torch.Generator()
+        self.rng.manual_seed(seed)
+
+
+class AttributeDropout(DropoutModule):
+    """Dropout with a given probability per attribute.
+    This is different from the behavior of ClassifierFreeGuidanceDropout as this allows for attributes
+    to be dropped out separately. For example, "artist" can be dropped while "genre" remains.
+    This is in contrast to ClassifierFreeGuidanceDropout where if "artist" is dropped "genre"
+    must also be dropped.
+
+    Args:
+        p (tp.Dict[str, float]): A dict mapping between attributes and dropout probability. For example:
+            ...
+            "genre": 0.1,
+            "artist": 0.5,
+            "wav": 0.25,
+            ...
+        active_on_eval (bool, optional): Whether the dropout is active at eval. Default to False.
+        seed (int, optional): Random seed.
+    """
+    def __init__(self, p: tp.Dict[str, tp.Dict[str, float]], active_on_eval: bool = False, seed: int = 1234):
+        super().__init__(seed=seed)
+        self.active_on_eval = active_on_eval
+        # construct dict that return the values from p otherwise 0
+        self.p = {}
+        for condition_type, probs in p.items():
+            self.p[condition_type] = defaultdict(lambda: 0, probs)
+
+    def forward(self, samples: tp.List[ConditioningAttributes]) -> tp.List[ConditioningAttributes]:
+        """
+        Args:
+            samples (list[ConditioningAttributes]): List of conditions.
+        Returns:
+            list[ConditioningAttributes]: List of conditions after certain attributes were set to None.
+        """
+        if not self.training and not self.active_on_eval:
+            return samples
+
+        samples = deepcopy(samples)
+        for condition_type, ps in self.p.items():  # for condition types [text, wav]
+            for condition, p in ps.items():  # for attributes of each type (e.g., [artist, genre])
+                # import pdb; pdb.set_trace()
+                # print(condition, p)
+                if torch.rand(1, generator=self.rng).item() < p:
+                    for sample in samples:
+                        dropout_condition(sample, condition_type, condition)
+        return samples
+
+    def __repr__(self):
+        return f"AttributeDropout({dict(self.p)})"
+
+
+class ClassifierFreeGuidanceDropout(DropoutModule):
+    """Classifier Free Guidance dropout.
+    All attributes are dropped with the same probability.
+
+    Args:
+        p (float): Probability to apply condition dropout during training.
+        seed (int): Random seed.
+    """
+    def __init__(self, p: float, seed: int = 1234):
+        super().__init__(seed=seed)
+        self.p = p
+
+    def forward(self, samples: tp.List[ConditioningAttributes]) -> tp.List[ConditioningAttributes]:
+        """
+        Args:
+            samples (list[ConditioningAttributes]): List of conditions.
+        Returns:
+            list[ConditioningAttributes]: List of conditions after all attributes were set to None.
+        """
+
+        if not self.training:
+            return samples
+        # import pdb; pdb.set_trace()
+        # decide on which attributes to drop in a batched fashion
+        drop = torch.rand(1, generator=self.rng).item() < self.p
+        if not drop:
+            return samples
+
+        # nullify conditions of all attributes
+        samples = deepcopy(samples)
+        for condition_type in ["text", "wav","joint_embed"]:
+            for sample in samples:
+                for condition in sample.attributes[condition_type]:
+                    dropout_condition(sample, condition_type, condition)
+        return samples
+
+    def __repr__(self):
+        return f"ClassifierFreeGuidanceDropout(p={self.p})"
+
+
+class TextConditioner(BaseConditioner):
+    ...
+
+
+class WaveformConditioner(BaseConditioner):
+    """Base class for all conditioners that take a waveform as input.
+    Classes that inherit must implement `_get_wav_embedding` that outputs
+    a continuous tensor, and `_downsampling_factor` that returns the down-sampling
+    factor of the embedding model.
+
+    Args:
+        dim (int): The internal representation dimension.
+        output_dim (int): Output dimension.
+    """
+    def __init__(self, dim: int, output_dim: int, input_token = False, padding_idx=None):
+        super().__init__(dim, output_dim, input_token, padding_idx)
+
+    def tokenize(self, x: WavCondition) -> WavCondition:
+        wav, length, sample_rate, path, seek_time = x
+        assert length is not None
+        return WavCondition(wav, length, sample_rate, path, seek_time)
+
+    def _get_wav_embedding(self, x: WavCondition) -> torch.Tensor:
+        """Gets as input a WavCondition and returns a dense embedding."""
+        raise NotImplementedError()
+
+    def _downsampling_factor(self):
+        """Returns the downsampling factor of the embedding model."""
+        raise NotImplementedError()
+
+    def forward(self, x: WavCondition) -> ConditionType:
+        """Extract condition embedding and mask from a waveform and its metadata.
+        Args:
+            x (WavCondition): Waveform condition containing raw waveform and metadata.
+        Returns:
+            ConditionType: a dense vector representing the conditioning along with its mask
+        """
+
+        wav, lengths, *_ = x
+        # import pdb; pdb.set_trace()
+        with torch.no_grad():
+            embeds = self._get_wav_embedding(x)
+        embeds = embeds.to(self.output_proj.weight)
+        embeds = self.output_proj(embeds)
+        # import pdb; pdb.set_trace()
+        if lengths is not None:
+            lengths = lengths / self._downsampling_factor()
+            mask = length_to_mask(lengths, max_len=embeds.shape[1]).int()  # type: ignore
+        else:
+            mask = torch.ones_like(embeds)
+        embeds = (embeds * mask.unsqueeze(2))
+
+        return embeds, mask
+
+
+class JointEmbeddingConditioner(BaseConditioner):
+    """Joint embedding conditioning supporting both audio or text conditioning.
+
+    Args:
+        dim (int): Dimension.
+        output_dim (int): Output dimension.
+        autocast_dtype (str): Autocast for the conditioner.
+        quantize (bool): Whether to quantize the CLAP embedding.
+        n_q (int): Number of residual quantizers (used if quantize is true).
+        bins (int): Quantizers' codebooks size (used if quantize is true).
+        kwargs: Additional parameters for residual vector quantizer.
+    """
+    def __init__(self, dim: int, output_dim: int, 
+                 autocast_dtype: tp.Optional[str] = 'float32', #quantize: bool = False,
+                 **kwargs):
+        super().__init__(dim=dim, output_dim=output_dim)
+        self.autocast_dtype = getattr(torch, autocast_dtype) if autocast_dtype is not None \
+                                else None
+        if self.autocast_dtype is None:
+            logger.warning("JointEmbeddingConditioner has no autocast, this might lead to NaN.")
+
+        # # residual vector quantizer to discretize the conditioned embedding
+        # self.quantizer  = None
+        # if quantize:
+        #     from ..modules.quantization import ResidualVectorQuantizer
+        #     self.quantizer = ResidualVectorQuantizer(dim, n_q=n_q, bins=bins, **kwargs)
+
+    def _get_embed(self, x: JointEmbedCondition) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+        """Get joint embedding in latent space from the inputs.
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor]: Tensor for the latent embedding
+                and corresponding empty indexes.
+        """
+        raise NotImplementedError()
+
+    def forward(self, x: JointEmbedCondition) -> ConditionType:
+        with torch.cuda.amp.autocast(dtype=self.autocast_dtype):
+            embed, empty_idx = self._get_embed(x)
+            if self.quantizer is not None:
+                embed = embed.view(-1, self.dim, 1)
+                q_res = self.quantizer(embed, frame_rate=1)
+                out_embed = q_res.x.view(-1, self.dim)
+            else:
+                out_embed = embed
+            out_embed = self.output_proj(out_embed).view(-1, 1, self.output_dim)
+            mask = torch.ones(*out_embed.shape[:2], device=out_embed.device)
+            mask[empty_idx, :] = 0  # zero-out index where the input is non-existant
+            out_embed = (out_embed * mask.unsqueeze(-1))
+            return out_embed, mask
+
+    def tokenize(self, x: JointEmbedCondition) -> JointEmbedCondition:
+        return x
+
+
+class ConditioningProvider(nn.Module):
+    """Prepare and provide conditions given all the supported conditioners.
+
+    Args:
+        conditioners (dict): Dictionary of conditioners.
+    """
+    def __init__(self, conditioners: tp.Dict[str, BaseConditioner]):
+        super().__init__()
+        self.conditioners = nn.ModuleDict(conditioners)
+        def _check_conditioner_type(c):
+            if isinstance(c, WaveformConditioner):
+                return "wav"
+            elif isinstance(c, TextConditioner):
+                return "text"
+            elif isinstance(c, JointEmbeddingConditioner):
+                return "joint_embed"
+            else:
+                raise NotImplementedError(f"{type(c)} are not Implemented!")
+        self.conditioner_type = {k: _check_conditioner_type(self.conditioners[k]) for k in self.conditioners}
+
+
+    @property
+    def joint_embed_conditions(self):
+        return [k for k, v in self.conditioners.items() if isinstance(v, JointEmbeddingConditioner)]
+
+    @property
+    def has_joint_embed_conditions(self):
+        return len(self.joint_embed_conditions) > 0
+
+    @property
+    def text_conditions(self):
+        return [k for k, v in self.conditioners.items() if isinstance(v, TextConditioner)]
+
+    @property
+    def wav_conditions(self):
+        return [k for k, v in self.conditioners.items() if isinstance(v, WaveformConditioner)]
+
+    @property
+    def has_wav_condition(self):
+        return len(self.wav_conditions) > 0
+
+    def tokenize(self, inputs: tp.List[ConditioningAttributes]) -> tp.Dict[str, tp.Any]:
+        """Match attributes/wavs with existing conditioners in self, and compute tokenize them accordingly.
+        This should be called before starting any real GPU work to avoid synchronization points.
+        This will return a dict matching conditioner names to their arbitrary tokenized representations.
+
+        Args:
+            inputs (list[ConditioningAttributes]): List of ConditioningAttributes objects containing
+                text and wav conditions.
+        """
+        assert all([isinstance(x, ConditioningAttributes) for x in inputs]), (
+            "Got unexpected types input for conditioner! should be tp.List[ConditioningAttributes]",
+            f" but types were {set([type(x) for x in inputs])}"
+        )
+
+        # import pdb; pdb.set_trace()
+        output = {}
+        text = self._collate_text(inputs)
+        wavs = self._collate_wavs(inputs)
+        joint_embeds = self._collate_joint_embeds(inputs)
+
+        assert set(text.keys() | wavs.keys() | joint_embeds.keys()).issubset(set(self.conditioners.keys())), (
+            f"Got an unexpected attribute! Expected {self.conditioners.keys()}, ",
+            f"got {text.keys(), wavs.keys(), joint_embeds.keys()}"
+        )
+
+        for attribute, batch in chain(text.items(), wavs.items(), joint_embeds.items()):
+            output[attribute] = self.conditioners[attribute].tokenize(batch)
+        return output
+
+    def forward(self, tokenized: tp.Dict[str, tp.Any], texts = None) -> tp.Dict[str, ConditionType]:
+        """Compute pairs of `(embedding, mask)` using the configured conditioners and the tokenized representations.
+        The output is for example:
+        {
+            "genre": (torch.Tensor([B, 1, D_genre]), torch.Tensor([B, 1])),
+            "description": (torch.Tensor([B, T_desc, D_desc]), torch.Tensor([B, T_desc])),
+            ...
+        }
+
+        Args:
+            tokenized (dict): Dict of tokenized representations as returned by `tokenize()`.
+        """
+        # import pdb; pdb.set_trace()
+        output = {}
+        for attribute, inputs in tokenized.items():
+            if attribute == 'self_wav' and texts is not None:
+                condition, mask = self.conditioners[attribute](inputs, texts = texts)
+            else:
+                condition, mask = self.conditioners[attribute](inputs)
+            output[attribute] = (condition, mask)
+        return output
+
+    def _collate_text(self, samples: tp.List[ConditioningAttributes]) -> tp.Dict[str, tp.List[tp.Optional[str]]]:
+        """Given a list of ConditioningAttributes objects, compile a dictionary where the keys
+        are the attributes and the values are the aggregated input per attribute.
+        For example:
+        Input:
+        [
+            ConditioningAttributes(text={"genre": "Rock", "description": "A rock song with a guitar solo"}, wav=...),
+            ConditioningAttributes(text={"genre": "Hip-hop", "description": "A hip-hop verse"}, wav=...),
+        ]
+        Output:
+        {
+            "genre": ["Rock", "Hip-hop"],
+            "description": ["A rock song with a guitar solo", "A hip-hop verse"]
+        }
+
+        Args:
+            samples (list of ConditioningAttributes): List of ConditioningAttributes samples.
+        Returns:
+            dict[str, list[str, optional]]: A dictionary mapping an attribute name to text batch.
+        """
+        out: tp.Dict[str, tp.List[tp.Optional[str]]] = defaultdict(list)
+        texts = [x.text for x in samples]
+        for text in texts:
+            for condition in self.text_conditions:
+                out[condition].append(text[condition])
+        return out
+
+    def _collate_wavs(self, samples: tp.List[ConditioningAttributes]) -> tp.Dict[str, WavCondition]:
+        """Generate a dict where the keys are attributes by which we fetch similar wavs,
+        and the values are Tensors of wavs according to said attributes.
+
+        *Note*: by the time the samples reach this function, each sample should have some waveform
+        inside the "wav" attribute. It should be either:
+        1. A real waveform
+        2. A null waveform due to the sample having no similar waveforms (nullified by the dataset)
+        3. A null waveform due to it being dropped in a dropout module (nullified by dropout)
+
+        Args:
+            samples (list of ConditioningAttributes): List of ConditioningAttributes samples.
+        Returns:
+            dict[str, WavCondition]: A dictionary mapping an attribute name to wavs.
+        """
+        # import pdb; pdb.set_trace()
+        wavs = defaultdict(list)
+        lengths = defaultdict(list)
+        sample_rates = defaultdict(list)
+        paths = defaultdict(list)
+        seek_times = defaultdict(list)
+        out: tp.Dict[str, WavCondition] = {}
+
+        for sample in samples:
+            for attribute in self.wav_conditions:
+                wav, length, sample_rate, path, seek_time = sample.wav[attribute]
+                assert wav.dim() == 3, f"Got wav with dim={wav.dim()}, but expected 3 [1, C, T]"
+                assert wav.size(0) == 1, f"Got wav [B, C, T] with shape={wav.shape}, but expected B == 1"
+                # mono-channel conditioning
+                # wav = wav.mean(1, keepdim=True)  # [1, 1, T] # by cyy, 为了实现后续功能注释掉了，请手动确保channel=1，or 输入channel 符合预期
+                wavs[attribute].append(wav.flatten())  # [C*T]
+                lengths[attribute].append(length)
+                sample_rates[attribute].extend(sample_rate)
+                paths[attribute].extend(path)
+                seek_times[attribute].extend(seek_time)
+
+        # stack all wavs to a single tensor
+        for attribute in self.wav_conditions:
+            stacked_wav, _ = collate(wavs[attribute], dim=0)
+            out[attribute] = WavCondition(
+                stacked_wav.unsqueeze(1), torch.cat(lengths[attribute]), sample_rates[attribute],
+                paths[attribute], seek_times[attribute])
+
+        return out
+
+    def _collate_joint_embeds(self, samples: tp.List[ConditioningAttributes]) -> tp.Dict[str, JointEmbedCondition]:
+        """Generate a dict where the keys are attributes by which we compute joint embeddings,
+        and the values are Tensors of pre-computed embeddings and the corresponding text attributes.
+
+        Args:
+            samples (list[ConditioningAttributes]): List of ConditioningAttributes samples.
+        Returns:
+            A dictionary mapping an attribute name to joint embeddings.
+        """
+        texts = defaultdict(list)
+        wavs = defaultdict(list)
+        lengths = defaultdict(list)
+        sample_rates = defaultdict(list)
+        paths = defaultdict(list)
+        seek_times = defaultdict(list)
+        channels: int = 0
+        
+        out = {}
+        for sample in samples:
+            for attribute in self.joint_embed_conditions:
+                wav, text, length, sample_rate, path, seek_time = sample.joint_embed[attribute]
+                assert wav.dim() == 3
+                if channels == 0:
+                    channels = wav.size(1)
+                else:
+                    assert channels == wav.size(1), "not all audio has same number of channels in batch"
+                assert wav.size(0) == 1, "Expecting single-wav batch in the collate method"
+                wav = einops.rearrange(wav, "b c t -> (b c t)")  # [1, C, T] => [C * T] 
+                wavs[attribute].append(wav)
+                texts[attribute].extend(text)
+                lengths[attribute].append(length)
+                sample_rates[attribute].extend(sample_rate)
+                paths[attribute].extend(path)
+                seek_times[attribute].extend(seek_time)
+
+        for attribute in self.joint_embed_conditions:
+            stacked_texts = texts[attribute]
+            stacked_paths = paths[attribute]
+            stacked_seek_times = seek_times[attribute]
+            stacked_wavs = pad_sequence(wavs[attribute])
+            stacked_wavs = einops.rearrange(stacked_wavs, "(c t) b -> b c t", c=channels)
+            stacked_sample_rates = sample_rates[attribute]
+            stacked_lengths = torch.cat(lengths[attribute])
+
+            assert stacked_lengths.size(0) == stacked_wavs.size(0)
+            assert len(stacked_sample_rates) == stacked_wavs.size(0)
+            assert len(stacked_texts) == stacked_wavs.size(0)
+            out[attribute] = JointEmbedCondition(
+                text=stacked_texts, wav=stacked_wavs,
+                length=stacked_lengths, sample_rate=stacked_sample_rates,
+                path=stacked_paths, seek_time=stacked_seek_times)
+
+        return out
+
+
+class ConditionFuser(StreamingModule):
+    """Condition fuser handles the logic to combine the different conditions
+    to the actual model input.
+
+    Args:
+        fuse2cond (tp.Dict[str, str]): A dictionary that says how to fuse
+            each condition. For example:
+            {
+                "prepend": ["description"],
+                "sum": ["genre", "bpm"],
+                "cross": ["description"],
+            }
+        cross_attention_pos_emb (bool, optional): Use positional embeddings in cross attention.
+        cross_attention_pos_emb_scale (int): Scale for positional embeddings in cross attention if used.
+    """
+    FUSING_METHODS = ["sum", "prepend", "cross", "input_interpolate"]
+
+    def __init__(self, fuse2cond: tp.Dict[str, tp.List[str]], cross_attention_pos_emb: bool = False,
+                 cross_attention_pos_emb_scale: float = 1.0):
+        super().__init__()
+        assert all(
+            [k in self.FUSING_METHODS for k in fuse2cond.keys()]
+        ), f"Got invalid fuse method, allowed methods: {self.FUSING_METHODS}"
+        self.cross_attention_pos_emb = cross_attention_pos_emb
+        self.cross_attention_pos_emb_scale = cross_attention_pos_emb_scale
+        self.fuse2cond: tp.Dict[str, tp.List[str]] = fuse2cond
+        self.cond2fuse: tp.Dict[str, str] = {}
+        for fuse_method, conditions in fuse2cond.items():
+            for condition in conditions:
+                self.cond2fuse[condition] = fuse_method
+
+    def forward(
+        self,
+        input: torch.Tensor,
+        conditions: tp.Dict[str, ConditionType]
+    ) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        """Fuse the conditions to the provided model input.
+
+        Args:
+            input (torch.Tensor): Transformer input.
+            conditions (dict[str, ConditionType]): Dict of conditions.
+        Returns:
+            tuple[torch.Tensor, torch.Tensor]: The first tensor is the transformer input
+                after the conditions have been fused. The second output tensor is the tensor
+                used for cross-attention or None if no cross attention inputs exist.
+        """
+        # import pdb; pdb.set_trace()
+        B, T, _ = input.shape
+
+        if 'offsets' in self._streaming_state:
+            first_step = False
+            offsets = self._streaming_state['offsets']
+        else:
+            first_step = True
+            offsets = torch.zeros(input.shape[0], dtype=torch.long, device=input.device)
+
+        assert set(conditions.keys()).issubset(set(self.cond2fuse.keys())), \
+            f"given conditions contain unknown attributes for fuser, " \
+            f"expected {self.cond2fuse.keys()}, got {conditions.keys()}"
+        cross_attention_output = None
+        prepend_input = input[:, :0]
+        for cond_type, (cond, cond_mask) in conditions.items():
+            op = self.cond2fuse[cond_type]
+            if op == 'sum': 
+                input += cond
+            elif op == 'input_interpolate':
+                cond = einops.rearrange(cond, "b t d -> b d t")
+                cond = F.interpolate(cond, size=input.shape[1])
+                input += einops.rearrange(cond, "b d t -> b t d")
+            elif op == 'prepend':
+                prepend_input = torch.cat([cond.to(input.dtype), prepend_input], dim=1) 
+                # NOTE 这里cond应该在后,这样顺序才符合配置文件,否则为逆序
+                # 但是之前实验是这样的为了保持一致就没改
+            elif op == 'cross':
+                if cross_attention_output is not None:
+                    cross_attention_output = torch.cat([cross_attention_output, cond], dim=1)
+                else:
+                    cross_attention_output = cond
+            else:
+                raise ValueError(f"unknown op ({op})")
+
+        if self.cross_attention_pos_emb and cross_attention_output is not None:
+            positions = torch.arange(
+                cross_attention_output.shape[1],
+                device=cross_attention_output.device
+            ).view(1, -1, 1)
+            pos_emb = create_sin_embedding(positions, cross_attention_output.shape[-1])
+            cross_attention_output = cross_attention_output + self.cross_attention_pos_emb_scale * pos_emb
+
+        if first_step:
+            input = torch.cat([prepend_input, input], dim=1)
+        if self._is_streaming:
+            self._streaming_state['offsets'] = offsets + T
+
+        return input, cross_attention_output

SongBloom/models/musicgen/conditioners/text.py

@@ -0,0 +1,254 @@
+from .base import *
+
+import spacy
+import warnings
+import random
+import hashlib
+from transformers import RobertaTokenizer, T5EncoderModel, T5Tokenizer, AutoTokenizer, XLMRobertaModel, XLMRobertaTokenizer  # type: ignore
+from num2words import num2words
+
+def hash_trick(word: str, vocab_size: int) -> int:
+    """Hash trick to pair each word with an index
+
+    Args:
+        word (str): word we wish to convert to an index
+        vocab_size (int): size of the vocabulary
+    Returns:
+        int: index of the word in the embedding LUT
+    """
+
+    hash = int(hashlib.sha256(word.encode("utf-8")).hexdigest(), 16)
+    return hash % vocab_size
+
+
+
+class PhonemeTokenizerConditioner(TextConditioner):
+    def __init__(self, 
+                 output_dim: int, 
+                 vocab_list,
+                 max_len = 600, 
+                 max_sentence_per_structure = 50,
+                 structure_tokens=None,
+                 structure_split_tokens=[','],
+                 sentence_split_tokens=['.'],
+                 mode='sum',
+                 structure_output_dim = 64,
+                 sentence_output_dim = 64,
+                 max_duration = 120,
+                 interpolate = False,
+                 ): 
+        
+        self.vocab_list = vocab_list
+        self.max_len = max_len
+        self.mode = mode
+        self.max_sentence_per_structure = max_sentence_per_structure
+        voc_size = len(self.vocab_list)
+        self.interpolate = interpolate
+        
+        if structure_tokens is None:
+            structure_tokens = [i for i in vocab_list if len(i) > 1 and i[0] == '[' and i[-1] == ']']
+        self.structure_token_ids = [vocab_list.index(i) for i in structure_tokens if i in vocab_list]
+        self.structure_split_token_ids = [vocab_list.index(i) for i in structure_split_tokens]
+        self.sentence_split_token_ids = [vocab_list.index(i) for i in sentence_split_tokens]
+
+        # here initialize a output_proj (nn.Embedding) layer
+        # By default the first vocab is "" (null)
+        if mode == 'sum':
+            content_output_dim = output_dim
+            sentence_output_dim = output_dim
+            structure_output_dim = output_dim
+        else:   # concat
+            content_output_dim = output_dim - sentence_output_dim - structure_output_dim   # by default
+            
+        super().__init__(voc_size, content_output_dim, input_token=True, padding_idx=0)
+        if self.mode != 'sum':
+            self.special_emb = nn.Embedding(len(self.structure_token_ids)+len(self.structure_split_token_ids)+len(self.sentence_split_token_ids)+1, 
+                                             structure_output_dim, padding_idx=0)
+            
+        self.blank_emb = nn.Parameter(torch.zeros(1, output_dim), requires_grad=False)
+
+        # the first index is "empty structure" token
+        self.sentence_idx_in_structure_emb = nn.Embedding(max_sentence_per_structure, sentence_output_dim, padding_idx=0) 
+
+        # print("max_len", self.max_len)
+        print(self.structure_token_ids)
+        
+        self.resolution = max_duration / max_len    # e.g., 120 / 600 = 0.2s 
+        print(self.__class__, f"resolution = {self.resolution}")
+    
+    def tokenize(self, x: tp.List[tp.Optional[str]]) -> tp.Dict[str, torch.Tensor]:
+        inputs = []
+        for xx in x:
+            xx = '' if xx is None else xx
+            vocab_id = [self.vocab_list.index(item) for item in xx.split(" ") if item in self.vocab_list]
+            inputs.append(torch.tensor(vocab_id).long()) # [T]        
+        return inputs
+    
+    
+    def interpolate_with_structure_duration(self, special_tokens, embeds, structure_dur):
+        # embeds: [T, N]
+        def sec2idx(sec):   # convert duration sec to token index
+            return int(sec / self.resolution)
+        
+        def target_token_types2list(tokens, target_token_types):
+
+            is_target_list = torch.any(torch.stack([tokens == i for i in target_token_types], dim=-1), dim=-1)
+            is_target_list = torch.where(is_target_list)[0].tolist()
+            return is_target_list
+        
+        structure_ids = []
+        for (structure, st, et) in structure_dur:
+            structure_ids.append([structure, sec2idx(st), sec2idx(et)])
+            
+        """
+        interpolate embeddings of each structure according to its duration 
+        """
+        is_structure_list = target_token_types2list(special_tokens, self.structure_token_ids)
+        is_structure_list.append(special_tokens.shape[-1])
+        
+        split_tokens = deepcopy(self.structure_split_token_ids)
+        split_tokens.extend(self.sentence_split_token_ids)
+        # is_split_list = target_token_types2list(special_tokens, split_tokens)
+                
+        
+        interpolated_embeds = embeds[:is_structure_list[0]]
+        for i, st in enumerate(is_structure_list[:-1]):
+            # (lorry) Explain "-tmp": 
+            # All structures are connected with " , " token,
+            # " ," is also the final token of each structure except the final one,
+            # but here we dont want to interpolate " , " token
+            tmp = 1
+            if i == len(is_structure_list[:-1]) - 1:  # the final structure, no need for "-1"
+                tmp = 0
+            
+       #     print(st, is_structure_list[i+1]-tmp)
+            to_interpolate = embeds[st: is_structure_list[i+1] - tmp]
+            interpolate_size = structure_ids[i][2] - structure_ids[i][1] - tmp
+       #     print(interpolate_size)
+            
+            #import pdb; pdb.set_trace()
+            # print(interpolated_embeds.shape, to_interpolate.shape, interpolate_size, )
+            if to_interpolate.shape[0] == 0:
+                import pdb; pdb.set_trace()
+            this_interpolated_embeds = F.interpolate(to_interpolate.unsqueeze(0).transpose(2, 1), 
+                                        size=interpolate_size, 
+                                        mode='nearest-exact').squeeze(0).transpose(1, 0)
+            
+            if tmp == 1:
+                interpolated_embeds = torch.cat((interpolated_embeds, this_interpolated_embeds, 
+                                                embeds[is_structure_list[i+1]].unsqueeze(0)), 0)
+            else:
+                interpolated_embeds = torch.cat((interpolated_embeds, this_interpolated_embeds), 0)
+        return interpolated_embeds
+            
+            
+    def forward(self, batch_tokens: tp.List, structure_dur = None) -> ConditionType:
+        """
+        Encode token_id into three types of embeddings:
+        1) content embedding: phoneme only (or meaningful contents to be sung out) 
+        2) structure embedding: structure / separation embeddings, including structures (verse/chorus/...), separators (. / ,)
+        The two above share the same embedding layer, can be changed to separate embedding layers.
+        3) sentence_idx embedding (per structure): 
+        """
+        embeds_batch = []
+        # print(batch_tokens)
+        for b in range(len(batch_tokens)):
+            tokens = batch_tokens[b]  
+
+            content_tokens = torch.zeros_like(tokens)
+            special_tokens = torch.zeros_like(tokens)
+            sentence_idx_in_structure_tokens = torch.zeros_like(tokens) 
+
+            current_structure_idx = 1
+            current_sentence_in_structure_idx = 1
+            current_structure = 0
+
+            for i in range(tokens.shape[-1]):
+                token = tokens[i]
+                if token in self.structure_token_ids:       # structure token
+                    # only update structure token, leave content and sentence index token null (default 0)
+                    if self.mode == 'sum':
+                        special_tokens[i] = token
+                    else:
+                        special_tokens[i] = self.structure_token_ids.index(token) + 1
+                    current_structure = token
+                    current_structure_idx += 1
+                    current_sentence_in_structure_idx = 1
+
+                elif token in self.sentence_split_token_ids:    # utterance split token
+                    # only update structure token, leave content and sentence index token null (default 0)
+                    # add up sentence index
+                    if self.mode == 'sum':
+                        special_tokens[i] = token
+                    else:
+                        special_tokens[i] = self.sentence_split_token_ids.index(token) + 1 + len(self.structure_token_ids)
+                    current_sentence_in_structure_idx += 1
+
+                elif token in self.structure_split_token_ids:    # structure split token
+                    # update structure token (current structure), content token (current token), 
+                    # blank index token 
+                    if self.mode == 'sum':
+                        special_tokens[i] = token
+                    else:
+                        special_tokens[i] = self.structure_split_token_ids.index(token) + 1 + len(self.structure_token_ids) + len(self.sentence_split_token_ids)
+
+                else:       # content tokens
+                    content_tokens[i] = token
+                    special_tokens[i] = current_structure
+                    sentence_idx_in_structure_tokens[i] = min(current_sentence_in_structure_idx, self.max_sentence_per_structure - 1)
+
+            # print("tokens", tokens.max(), tokens.min())
+            # print("special tokens", special_tokens.max(), special_tokens.min())
+            # print("sentence idx in structure", sentence_idx_in_structure_tokens.max(), sentence_idx_in_structure_tokens.min())
+            device = self.output_proj.weight.device
+            
+            # import pdb; pdb.set_trace()
+            content_embeds = self.output_proj(tokens.to(device))    # [T, N]
+            if self.mode == 'sum':
+                structure_embeds = self.output_proj(special_tokens.to(device))
+            else:
+                structure_embeds = self.special_emb(special_tokens.to(device))
+            sentence_idx_embeds = self.sentence_idx_in_structure_emb(sentence_idx_in_structure_tokens.to(device))
+
+            if self.mode == 'sum':
+                embeds = content_embeds + structure_embeds + sentence_idx_embeds
+            else:
+                embeds = torch.cat((content_embeds, structure_embeds, sentence_idx_embeds), -1) # [T, N]
+                
+            if self.interpolate:
+                embeds = self.interpolate_with_structure_duration(tokens, embeds, structure_dur[b])
+            embeds_batch.append(embeds)
+
+        # set batch_size = 1, [B, T, N]
+        if self.max_len is not None:
+            max_len = self.max_len
+        else:
+            max_len = max([e.shape[0] for e in embeds_batch])
+        embeds, mask = self.pad_2d_tensor(embeds_batch, max_len)
+        
+        return embeds, mask
+    
+    
+    def pad_2d_tensor(self, xs, max_len):
+        new_tensor = []
+        new_mask = []
+        for x in xs:
+            seq_len, dim = x.size()
+            pad_len = max_len - seq_len
+
+            if pad_len > 0:
+                pad_tensor = self.blank_emb.repeat(pad_len, 1).to(x.device)  # T, D
+                padded_tensor = torch.cat([x, pad_tensor], dim=0)
+                mask = torch.cat((torch.ones_like(x[:, 0]), 
+                                  torch.zeros_like(pad_tensor[:, 0])), 0)   # T
+            elif pad_len < 0:
+                padded_tensor = x[:max_len]
+                mask = torch.ones_like(padded_tensor[:, 0])
+            else:
+                padded_tensor = x
+                mask = torch.ones_like(x[:, 0])
+
+            new_tensor.append(padded_tensor)
+            new_mask.append(mask)
+        # [B, T, D] & [B, T]
+        return torch.stack(new_tensor, 0), torch.stack(new_mask, 0)   

SongBloom/models/musicgen/conditioners/wav.py

@@ -0,0 +1,74 @@
+
+from .base import *
+import omegaconf
+from ...vae_frontend import AbstractVAE
+
+def pad_to_fix_length(x, max_len, pad_value=0.):
+    bsz, seq_len = x.shape[:2]
+    if seq_len >= max_len:
+        return x[:, :max_len]
+    else:
+        pad_len = max_len - seq_len
+        pad_tensor = torch.full((bsz, pad_len, *x.shape[2:]), pad_value, dtype=x.dtype, device=x.device)
+        padded_tensor = torch.cat([x, pad_tensor], dim=1)
+        return padded_tensor
+
+class AudioTokenizerConditioner(WaveformConditioner):
+    def __init__(self, output_dim, audio_tokenizer, cache=False, max_len=None):
+        super().__init__(output_dim, output_dim)
+        self.max_len = max_len
+        self.use_cache = cache
+        
+        self.tokenizer = audio_tokenizer
+        # breakpoint()
+        
+        # TODO if cached and not load vae, receive a dict instead
+        if isinstance(self.tokenizer, dict):
+            self.tokenizer = omegaconf.DictConfig(self.tokenizer)
+            self.code_depth = self.tokenizer.channel_dim
+            
+            
+        elif isinstance(self.tokenizer, AbstractVAE):
+            self.tokenizer_tp = "vae"
+            if self.use_cache:
+                self.code_depth = self.tokenizer.channel_dim
+            else:
+                self.code_depth = 1 # TODO 强制把输入channel设成1了 self.tokenizer.input_channel
+            self.output_proj = nn.Identity() if self.output_dim == self.tokenizer.channel_dim \
+                            else nn.Linear(self.tokenizer.channel_dim, self.output_dim, bias=False)
+                            
+        else:
+            raise NotImplementedError
+        
+        
+    def forward(self, x: WavCondition):
+        wav, lengths, *_ = x
+        B = wav.shape[0]
+        wav = wav.reshape(B, self.code_depth, -1)
+        # print(wav.shape)
+        # import torchaudio
+        # torchaudio.save("/apdcephfs_cq7/share_1297902/common/erichtchen/shixisheng/cyy/project/music_generation_repo/core/models/musicgen/conditioners/111.wav", wav[0].cpu(), 48000)
+        if self.tokenizer_tp == "vae":
+            if self.use_cache:
+                audio_latents = wav.transpose(-1,-2)
+            else:
+                with torch.no_grad():
+                    audio_latents = self.tokenizer.encode(wav).transpose(-1,-2)
+                    # print('transform wav to vae')
+            audio_latents = self.output_proj(audio_latents)
+
+        # print(audio_latents.shape)
+        if self.max_len is not None:
+            audio_latents = pad_to_fix_length(audio_latents, self.max_len, 0.)
+                    
+        if lengths is not None:
+            lengths = torch.round(lengths.float() * audio_latents.shape[1] / wav.shape[-1])
+            mask = length_to_mask(lengths, max_len=audio_latents.shape[1]).int()  # type: ignore
+        else:
+            mask = torch.ones((B, audio_latents.shape[1]), device=audio_latents.device,dtype=torch.int)
+
+        audio_latents = audio_latents * mask[..., None]
+        
+        return audio_latents, mask
+     
+

SongBloom/models/musicgen/get_backend.py

@@ -0,0 +1,76 @@
+import torch
+import os,sys
+from transformers.utils import is_flash_attn_2_available
+from transformers.models.llama import LlamaModel, LlamaConfig
+from transformers.models.bart.modeling_bart import BartEncoder, BartDecoder, BartConfig
+import warnings
+# from transformers.models.musicgen.modeling_musicgen import MusicgenModel, MusicgenDecoder, MusicgenDecoderConfig # 用的就是BartDecoder，但是没有cross-attn
+
+try:
+    assert is_flash_attn_2_available()
+    assert torch.cuda.get_device_capability(torch.device("cuda")) >= (8, 0)
+    assert os.environ.get("DISABLE_FLASH_ATTN",'0') != "1"
+    _enable_flash_attention = True
+except:
+    _enable_flash_attention = False
+
+if not _enable_flash_attention:
+    warnings.warn("Not support flash-attn!")
+
+def get_backend(name, dim, num_heads, num_layers, hidden_scale, init_std=0.02, rope_theta=10000,):
+    # SA (causal) - FF
+    if name == 'llama':
+        model_cfg = LlamaConfig(
+            hidden_size=dim,
+            intermediate_size=dim * hidden_scale,
+            num_attention_heads=num_heads,
+            num_hidden_layers=num_layers,
+            num_key_value_heads=num_heads,
+            vocab_size=dim,
+            use_cache=False,
+            max_position_embeddings=4096,
+            hidden_act="silu",
+            initializer_range=init_std,
+            rope_theta=rope_theta,
+            _attn_implementation="flash_attention_2" if _enable_flash_attention else "eager",
+        )
+        model = LlamaModel(model_cfg)
+        
+    # SA -FF
+    elif name == 'bart_enc':
+        model_cfg = BartConfig(
+            d_model=dim,
+            max_position_embeddings=4096,
+            dropout=0., 
+            use_cache=False,
+            _attn_implementation="flash_attention_2" if _enable_flash_attention else "eager",
+            activation_function='gelu',
+            # for BartEncoder
+            encoder_layers=num_layers, 
+            encoder_attention_heads=num_heads,
+            init_std=init_std,
+            encoder_ffn_dim=dim * hidden_scale,
+        )
+        model = BartEncoder(model_cfg)
+        
+    # SA - CA - FF
+    elif name == 'bart_dec':
+        model_cfg = BartConfig(
+            d_model=dim,
+            max_position_embeddings=4096,
+            dropout=0., 
+            use_cache=False,
+            _attn_implementation="flash_attention_2" if _enable_flash_attention else "eager",
+            activation_function='gelu',
+            # for BartDecoder
+            decoder_layers=num_layers,
+            decoder_attention_heads=num_heads,
+            decoder_ffn_dim=dim * hidden_scale,
+        )
+        model = BartDecoder(model_cfg)
+
+    else:
+        raise NotImplementedError
+
+    delattr(model, "embed_tokens")
+    return model

SongBloom/models/musicgen/modules/streaming.py

@@ -0,0 +1,125 @@
+"""
+Streaming module API that should be implemented by all Streaming components,
+"""
+
+from contextlib import contextmanager
+import typing as tp
+from torch import nn
+import torch
+
+
+State = tp.Dict[str, torch.Tensor]
+
+
+class StreamingModule(nn.Module):
+    """Common API for streaming components.
+
+    Each streaming component has a streaming state, which is just a dict[str, Tensor].
+    By convention, the first dim of each tensor must be the batch size.
+    Don't use dots in the key names, as this would clash with submodules
+    (like in state_dict).
+
+    If `self._is_streaming` is True, the component should use and remember
+    the proper state inside `self._streaming_state`.
+
+    To set a streaming component in streaming state, use
+
+        with module.streaming():
+            ...
+
+    This will automatically reset the streaming state when exiting the context manager.
+    This also automatically propagates to all streaming children module.
+
+    Some module might also implement the `StreamingModule.flush` method, although
+    this one is trickier, as all parents module must be StreamingModule and implement
+    it as well for it to work properly. See `StreamingSequential` after.
+    """
+    def __init__(self) -> None:
+        super().__init__()
+        self._streaming_state: State = {}
+        self._is_streaming = False
+
+    def _apply_named_streaming(self, fn: tp.Any):
+        for name, module in self.named_modules():
+            if isinstance(module, StreamingModule):
+                fn(name, module)
+
+    def _set_streaming(self, streaming: bool):
+        def _set_streaming(name, module):
+            module._is_streaming = streaming
+        self._apply_named_streaming(_set_streaming)
+
+    @contextmanager
+    def streaming(self):
+        """Context manager to enter streaming mode. Reset streaming state on exit."""
+        self._set_streaming(True)
+        try:
+            yield
+        finally:
+            self._set_streaming(False)
+            self.reset_streaming()
+
+    def reset_streaming(self):
+        """Reset the streaming state."""
+        def _reset(name: str, module: StreamingModule):
+            module._streaming_state.clear()
+
+        self._apply_named_streaming(_reset)
+
+    def get_streaming_state(self) -> State:
+        """Return the streaming state, including that of sub-modules."""
+        state: State = {}
+
+        def _add(name: str, module: StreamingModule):
+            if name:
+                name += "."
+            for key, value in module._streaming_state.items():
+                state[name + key] = value
+
+        self._apply_named_streaming(_add)
+        return state
+
+    def set_streaming_state(self, state: State):
+        """Set the streaming state, including that of sub-modules."""
+        state = dict(state)
+
+        def _set(name: str, module: StreamingModule):
+            if name:
+                name += "."
+            module._streaming_state.clear()
+            for key, value in list(state.items()):
+                # complexity is not ideal here, but probably fine.
+                if key.startswith(name):
+                    local_key = key[len(name):]
+                    if '.' not in local_key:
+                        module._streaming_state[local_key] = value
+                        del state[key]
+
+        self._apply_named_streaming(_set)
+        assert len(state) == 0, list(state.keys())
+
+    def flush(self, x: tp.Optional[torch.Tensor] = None):
+        """Flush any remaining outputs that were waiting for completion.
+        Typically, for convolutions, this will add the final padding
+        and process the last buffer.
+
+        This should take an optional argument `x`, which will be provided
+        if a module before this one in the streaming pipeline has already
+        spitted out a flushed out buffer.
+        """
+        if x is None:
+            return None
+        else:
+            return self(x)
+
+
+class StreamingSequential(StreamingModule, nn.Sequential):
+    """A streaming compatible alternative of `nn.Sequential`.
+    """
+    def flush(self, x: tp.Optional[torch.Tensor] = None):
+        for module in self:
+            if isinstance(module, StreamingModule):
+                x = module.flush(x)
+            elif x is not None:
+                x = module(x)
+        return x

SongBloom/models/musicldm/inference/sampling.py

@@ -0,0 +1,271 @@
+import torch
+import math
+from tqdm import trange, tqdm
+
+# import k_diffusion as K
+
+# Define the noise schedule and sampling loop
+def get_alphas_sigmas(t):
+    """Returns the scaling factors for the clean image (alpha) and for the
+    noise (sigma), given a timestep."""
+    return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2)
+
+def alpha_sigma_to_t(alpha, sigma):
+    """Returns a timestep, given the scaling factors for the clean image and for
+    the noise."""
+    return torch.atan2(sigma, alpha) / math.pi * 2
+
+def t_to_alpha_sigma(t):
+    """Returns the scaling factors for the clean image and for the noise, given
+    a timestep."""
+    return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2)
+
+
+@torch.no_grad()
+def sample_discrete_euler(model, x, steps, sigma_max=1.0, prog_bar=False, **extra_args):
+    """Draws samples from a model given starting noise. Euler method"""
+
+    # Make tensor of ones to broadcast the single t values
+    ts = x.new_ones([x.shape[0]])
+
+    # Create the noise schedule
+    t = torch.linspace(sigma_max, 0, steps + 1)
+    # all = {}
+
+    #alphas, sigmas = 1-t, t
+    iterator = tqdm(zip(t[:-1], t[1:]), total=steps) if prog_bar else zip(t[:-1], t[1:])
+    for t_curr, t_prev in iterator:
+            # Broadcast the current timestep to the correct shape
+            t_curr_tensor = t_curr * torch.ones(
+                (x.shape[0],), dtype=x.dtype, device=x.device
+            )
+            dt = t_prev - t_curr  # we solve backwards in our formulation
+            v = model(x, t_curr_tensor, **extra_args)
+            # all[t_curr.item()] = x-t_curr*v
+            x = x + dt * v #.denoise(x, denoiser, t_curr_tensor, cond, uc)
+            
+    # If we are on the last timestep, output the denoised image
+    return x #, all
+
+@torch.no_grad()
+def sample_discrete_euler_with_temperature(model, x, steps, temperature=1.0, sigma_max=1.0, prog_bar=False, **extra_args):
+    """Draws samples from a model given starting noise. Euler method"""
+
+    # Make tensor of ones to broadcast the single t values
+    ts = x.new_ones([x.shape[0]])
+    noise = x
+
+    # Create the noise schedule
+    t = torch.linspace(sigma_max, 0, steps + 1)
+    # all = {}
+    x = torch.zeros_like(noise)
+    if temperature >= sigma_max:
+        x = noise
+        
+    #alphas, sigmas = 1-t, t
+    iterator = tqdm(zip(t[:-1], t[1:]), total=steps) if prog_bar else zip(t[:-1], t[1:])
+    for t_curr, t_prev in iterator:
+            # Broadcast the current timestep to the correct shape
+                
+            t_curr_tensor = t_curr * torch.ones(
+                (x.shape[0],), dtype=x.dtype, device=x.device
+            )
+            dt = t_prev - t_curr  # we solve backwards in our formulation
+            v = model(x, t_curr_tensor, **extra_args)
+            # all[t_curr.item()] = x-t_curr*v
+            if  t_curr > temperature and t_prev <= temperature:
+                x_0 = x - v
+                x = (1-t_prev) * x_0 + t_prev * noise
+            else:
+                x = x + dt * v #.denoise(x, denoiser, t_curr_tensor, cond, uc)
+            
+    # If we are on the last timestep, output the denoised image
+    return x #, all
+
+
+@torch.no_grad()
+def sample(model, x, steps, eta, prog_bar=False, **extra_args):
+    """Draws samples from a model given starting noise. v-diffusion"""
+    ts = x.new_ones([x.shape[0]])
+    origin_dtype = x.dtype
+    # Create the noise schedule
+    t = torch.linspace(1, 0, steps + 1)[:-1]
+
+    alphas, sigmas = get_alphas_sigmas(t)
+
+    # The sampling loop
+    bar = trange if prog_bar else range
+    for i in bar(steps):
+
+        # Get the model output (v, the predicted velocity)
+        with torch.cuda.amp.autocast():
+            v = model(x, ts * t[i], **extra_args).float()
+
+        # Predict the noise and the denoised image
+        pred = x * alphas[i] - v * sigmas[i]
+        eps = x * sigmas[i] + v * alphas[i]
+
+        # If we are not on the last timestep, compute the noisy image for the
+        # next timestep.
+        if i < steps - 1:
+            # If eta > 0, adjust the scaling factor for the predicted noise
+            # downward according to the amount of additional noise to add
+            ddim_sigma = eta * (sigmas[i + 1]**2 / sigmas[i]**2).sqrt() * \
+                (1 - alphas[i]**2 / alphas[i + 1]**2).sqrt()
+            adjusted_sigma = (sigmas[i + 1]**2 - ddim_sigma**2).sqrt()
+
+            # Recombine the predicted noise and predicted denoised image in the
+            # correct proportions for the next step
+            x = pred * alphas[i + 1] + eps * adjusted_sigma
+            # Add the correct amount of fresh noise
+            if eta:
+                x += torch.randn_like(x) * ddim_sigma
+
+    # If we are on the last timestep, output the denoised image
+    return pred.to(origin_dtype)
+
+# Soft mask inpainting is just shrinking hard (binary) mask inpainting
+# Given a float-valued soft mask (values between 0 and 1), get the binary mask for this particular step
+def get_bmask(i, steps, mask):
+    strength = (i+1)/(steps)
+    # convert to binary mask
+    bmask = torch.where(mask<=strength,1,0)
+    return bmask
+
+def make_cond_model_fn(model, cond_fn):
+    def cond_model_fn(x, sigma, **kwargs):
+        with torch.enable_grad():
+            x = x.detach().requires_grad_()
+            denoised = model(x, sigma, **kwargs)
+            cond_grad = cond_fn(x, sigma, denoised=denoised, **kwargs).detach()
+            cond_denoised = denoised.detach() + cond_grad * K.utils.append_dims(sigma**2, x.ndim)
+        return cond_denoised
+    return cond_model_fn
+
+# Uses k-diffusion from https://github.com/crowsonkb/k-diffusion
+# init_data is init_audio as latents (if this is latent diffusion)
+# For sampling, set both init_data and mask to None
+# For variations, set init_data 
+# For inpainting, set both init_data & mask 
+def sample_k(
+        model_fn, 
+        noise, 
+        init_data=None,
+        mask=None,
+        steps=100, 
+        sampler_type="dpmpp-2m-sde", 
+        sigma_min=0.5, 
+        sigma_max=50, 
+        rho=1.0, device="cuda", 
+        callback=None, 
+        cond_fn=None,
+        **extra_args
+    ):
+
+    denoiser = K.external.VDenoiser(model_fn)
+
+    if cond_fn is not None:
+        denoiser = make_cond_model_fn(denoiser, cond_fn)
+
+    # Make the list of sigmas. Sigma values are scalars related to the amount of noise each denoising step has
+    sigmas = K.sampling.get_sigmas_polyexponential(steps, sigma_min, sigma_max, rho, device=device)
+    # Scale the initial noise by sigma 
+    noise = noise * sigmas[0]
+
+    wrapped_callback = callback
+
+    if mask is None and init_data is not None:
+        # VARIATION (no inpainting)
+        # set the initial latent to the init_data, and noise it with initial sigma
+        x = init_data + noise 
+    elif mask is not None and init_data is not None:
+        # INPAINTING
+        bmask = get_bmask(0, steps, mask)
+        # initial noising
+        input_noised = init_data + noise
+        # set the initial latent to a mix of init_data and noise, based on step 0's binary mask
+        x = input_noised * bmask + noise * (1-bmask)
+        # define the inpainting callback function (Note: side effects, it mutates x)
+        # See https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py#L596C13-L596C105
+        # callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        # This is called immediately after `denoised = model(x, sigmas[i] * s_in, **extra_args)`
+        def inpainting_callback(args):
+            i = args["i"]
+            x = args["x"]
+            sigma = args["sigma"]
+            #denoised = args["denoised"]
+            # noise the init_data input with this step's appropriate amount of noise
+            input_noised = init_data + torch.randn_like(init_data) * sigma
+            # shrinking hard mask
+            bmask = get_bmask(i, steps, mask)
+            # mix input_noise with x, using binary mask
+            new_x = input_noised * bmask + x * (1-bmask)
+            # mutate x
+            x[:,:,:] = new_x[:,:,:]
+        # wrap together the inpainting callback and the user-submitted callback. 
+        if callback is None: 
+            wrapped_callback = inpainting_callback
+        else:
+            wrapped_callback = lambda args: (inpainting_callback(args), callback(args))
+    else:
+        # SAMPLING
+        # set the initial latent to noise
+        x = noise
+
+
+    with torch.cuda.amp.autocast():
+        if sampler_type == "k-heun":
+            return K.sampling.sample_heun(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "k-lms":
+            return K.sampling.sample_lms(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "k-dpmpp-2s-ancestral":
+            return K.sampling.sample_dpmpp_2s_ancestral(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "k-dpm-2":
+            return K.sampling.sample_dpm_2(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "k-dpm-fast":
+            return K.sampling.sample_dpm_fast(denoiser, x, sigma_min, sigma_max, steps, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "k-dpm-adaptive":
+            return K.sampling.sample_dpm_adaptive(denoiser, x, sigma_min, sigma_max, rtol=0.01, atol=0.01, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "dpmpp-2m-sde":
+            return K.sampling.sample_dpmpp_2m_sde(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "dpmpp-3m-sde":
+            return K.sampling.sample_dpmpp_3m_sde(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+
+# Uses discrete Euler sampling for rectified flow models
+# init_data is init_audio as latents (if this is latent diffusion)
+# For sampling, set both init_data and mask to None
+# For variations, set init_data 
+# For inpainting, set both init_data & mask 
+def sample_rf(
+        model_fn, 
+        noise, 
+        init_data=None,
+        steps=100, 
+        sigma_max=1,
+        device="cuda", 
+        callback=None, 
+        cond_fn=None,
+        **extra_args
+    ):
+
+    if sigma_max > 1:
+        sigma_max = 1
+
+    if cond_fn is not None:
+        denoiser = make_cond_model_fn(denoiser, cond_fn)
+
+    wrapped_callback = callback
+
+    if init_data is not None:
+        # VARIATION (no inpainting)
+        # Interpolate the init data and the noise for init audio
+        x = init_data * (1 - sigma_max) + noise * sigma_max
+    else:
+        # SAMPLING
+        # set the initial latent to noise
+        x = noise
+
+    with torch.cuda.amp.autocast():
+        # TODO: Add callback support
+        #return sample_discrete_euler(model_fn, x, steps, sigma_max, callback=wrapped_callback, **extra_args)
+        return sample_discrete_euler(model_fn, x, steps, sigma_max, **extra_args)

SongBloom/models/musicldm/musicldm_dit.py

@@ -0,0 +1,24 @@
+from contextlib import contextmanager
+from dataclasses import dataclass
+from functools import partial
+import logging
+import math
+
+import torch
+from torch import nn
+
+
+
+class FourierFeatures(nn.Module):
+    def __init__(self, in_features, out_features, std=1.):
+        super().__init__()
+        assert out_features % 2 == 0
+        self.weight = nn.Parameter(torch.randn(
+            [out_features // 2, in_features]) * std)
+
+    def forward(self, input):
+        f = 2 * math.pi * input @ self.weight.T
+        return torch.cat([f.cos(), f.sin()], dim=-1)
+
+
+

SongBloom/models/songbloom/songbloom_mvsa.py

@@ -0,0 +1,572 @@
+from contextlib import contextmanager
+from dataclasses import dataclass
+from functools import partial
+import logging
+import math
+import typing as tp
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from einops.layers.torch import Rearrange
+from einops import rearrange
+import tqdm
+
+from ..base.utils import create_norm_fn
+from ..base.sample import sample_top_k, sample_top_p, multinomial
+from ..musicgen.modules.streaming import StreamingModule
+from ..musicgen.conditioners import (
+    get_condition_fuser,
+    get_conditioner_provider,
+     ConditionType,
+     ConditioningProvider,
+     ConditionFuser,
+     AttributeDropout,
+     ClassifierFreeGuidanceDropout,
+     ConditioningAttributes,
+     WavCondition,
+    JointEmbedCondition
+)
+
+from ..musicgen.get_backend import get_backend
+
+from ..transformer import ContinuousTransformer as DiT_block
+from ..musicldm.musicldm_dit import FourierFeatures
+from ..musicldm.inference.sampling import get_alphas_sigmas, sample, sample_discrete_euler, sample_discrete_euler_with_temperature
+
+ConditionTensors = tp.Dict[str, ConditionType]
+CFGConditions = tp.Union[ConditionTensors, tp.Tuple[ConditionTensors, ConditionTensors]]
+
+@dataclass
+class DiTAROutput:
+    ar_logit: torch.Tensor  
+    ar_target: torch.Tensor
+    nar_pred: torch.Tensor
+    nar_target: torch.Tensor
+    nar_t: torch.Tensor
+
+
+
+class MVSA_DiTAR(StreamingModule):
+    """
+    Multiple skeleton embedding, single compressed vae latent 
+    eg. V1 V2 V3 A1-3 V4 V5 V6 A4-6
+        V -> cross entropy (skeleton)
+        A -> local-DiT uncompress -> (A1-3 -> E1 E2 E3)
+
+    Args:
+        StreamingModule (_type_): _description_
+    """
+
+    def __init__(self, condition_provider_cfg, fuser_cfg, 
+                 block_size: int = 32, dim: int = 1024, num_heads: int = 8,
+                 num_pitch: int = 128, hidden_scale: int = 4, lm_layers: int = 16,
+                 norm: str = 'layer_norm', pre_norm: bool = False,
+                 backend='llama',init_std: float=0.02, 
+                 # ======================
+                 latent_dim: int = 64, diff_layers: int = 8,
+                 time_cond_type: tp.Literal['adaLM', "prepend"] = "prepend", 
+                 timestep_features_dim: int = 256,
+                 diffusion_objective: tp.Literal["v", "rectified_flow"] = "v",
+                 timestep_sampler: tp.Literal["uniform", "logit_normal"] = "uniform",
+                 rotary_base_val=10000, h_dropout: float = None, 
+                 # ======================
+                 cfg_dropout: float = 0, cfg_coef: float = 1.0,
+                 attribute_dropout: tp.Dict[str, tp.Dict[str, float]] = {}
+                 ):
+        super().__init__()
+
+        self.condition_provider = get_conditioner_provider(condition_provider_cfg)
+        self.fuser = get_condition_fuser(fuser_cfg)
+        
+        self.dim = dim  
+        self.latent_dim = latent_dim
+        self.block_size = block_size
+            
+        self.cfg_coef = cfg_coef
+        self.h_dropout = h_dropout if h_dropout is not None else 0.
+        self.cfg_dropout = ClassifierFreeGuidanceDropout(p=cfg_dropout)
+        self.att_dropout = AttributeDropout(p=attribute_dropout)
+
+
+        # Build AR lm
+        self.num_pitch = num_pitch + 1 # self.num_pitch = <EOS>, self.num_pitch+1 = special
+        self.skeleton_emb = nn.Embedding(self.num_pitch + 1, dim)
+        self.bos_token = nn.Parameter(torch.empty(dim).normal_(mean=0.0, std=init_std), requires_grad=True)
+
+        
+        # self.lm_type = lm_type
+        self.backend = backend
+        
+        if self.backend == 'llama':
+            self.ar_transformer = get_backend('llama', 
+                                        dim, num_heads, lm_layers, hidden_scale,init_std=init_std, rope_theta=rotary_base_val)
+            self.ar_transformer.gradient_checkpointing_enable()
+        elif self.backend == 'bart':
+            self.cross_encoder = get_backend('bart_enc',
+                                        dim, num_heads, lm_layers // 4, hidden_scale,init_std=init_std)
+            self.ar_transformer = get_backend('bart_dec',
+                                        dim, num_heads, lm_layers, hidden_scale,init_std=init_std)
+        else:
+            raise NotImplementedError(f"Illegal backend: {self.backend}!")
+        
+
+        self.skeleton_classifier =  nn.Sequential(nn.Linear(dim, dim, bias=False), 
+                                                    nn.SiLU(),
+                                                    nn.Linear(dim, self.num_pitch),)
+        
+        self.pre_norm: tp.Optional[nn.Module] = None
+        if pre_norm:
+            self.pre_norm = create_norm_fn(norm, dim)
+        self.reset_streaming()
+        
+        # Build NAR DiT
+        self.block_conv = nn.Sequential(
+            Rearrange("b d (n s) -> b n (s d)", s=self.block_size),
+            nn.Linear(self.block_size * latent_dim, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim)
+        )
+        self.project_in = nn.Linear(latent_dim, dim) if latent_dim != dim else nn.Identity()
+        self.project_out = nn.Linear(dim, latent_dim) if latent_dim != dim else nn.Identity()
+        
+        self.timestep_features_dim = timestep_features_dim
+        self.time_cond_type = time_cond_type 
+        assert self.time_cond_type in ['adaLN', "prepend"]
+        self.timestep_features = FourierFeatures(1, timestep_features_dim)
+        self.to_timestep_embed = nn.Sequential(
+            nn.Linear(timestep_features_dim, dim, bias=False),
+            nn.SiLU(),
+            nn.Linear(dim, dim),
+        )
+        
+        self.time_cond_type = time_cond_type
+        self.nar_dit = DiT_block(
+            dim=dim,
+            depth=diff_layers,
+            dim_heads= dim // num_heads,
+            rotary_pos_emb=True,
+            cross_attend=False,
+            causal=False,
+            ff_kwargs={"dim_ff": dim * hidden_scale, "no_bias": True},
+            global_cond_dim=self.dim if self.time_cond_type=="adaLN" else None,
+            rotary_base_val = rotary_base_val,
+            # init_std=init_std
+        )
+        self.nar_dit.gradient_checkpointing_enable()
+        
+        self.diffusion_objective = diffusion_objective
+        self.timestep_sampler = timestep_sampler
+        self.rng = torch.quasirandom.SobolEngine(1, scramble=True)
+        
+        self.init_weights(init_std=init_std)
+            
+        
+
+
+    @property
+    def special_token_id(self) -> int:
+        return self.num_pitch
+
+
+    @property
+    def eos_token_id(self) -> int:
+        return self.num_pitch-1
+
+    def forward(self, x_sketch, x_latent, x_len, condition_tensors) -> DiTAROutput:
+        '''
+        only for train: lm_forward + diffusion_forward (random_t)
+        x_sketch: (B,T) # T % block_sz == 0 (no <eos> token) padded with <eos>
+        x_latent: (B, D_{in}, T)
+        '''
+        # AR 
+        assert torch.all(x_len % self.block_size == 0), f"{x_len}"
+        block_num = x_len // self.block_size
+        
+        sketch_emb = self.skeleton_emb(x_sketch)
+        latent_emb = self.block_conv(x_latent)
+
+        B, T, D = sketch_emb.shape
+        
+        lm_input = rearrange(torch.cat([rearrange(sketch_emb, "b (n s) d -> b n s d", s=self.block_size), 
+                            latent_emb.unsqueeze(dim=2)], dim=2), "b n s d -> b (n s) d")
+        lm_input = torch.cat([self.bos_token.reshape(1,1,-1).expand(B,-1,-1),
+                            lm_input], dim=1) #add <sos>
+        
+        new_seq_len = x_len + block_num + 1
+        
+        ar_target = F.pad(x_sketch, (0,1), value=self.eos_token_id)
+        for b,l in enumerate(x_len):
+            ar_target[b, l+1:] = self.special_token_id # 用来mask掉多余的eos
+        
+
+    
+        lm_out = self.lm_forward(lm_input, condition_tensors)
+        
+        
+        indices = torch.arange(lm_out.shape[1])
+        h_ind = indices[(indices+1) % (self.block_size+1) == 0]
+        not_h_ind = indices[(indices+1) % (self.block_size+1) != 0]
+        
+        x_sketch_logit = self.skeleton_classifier(lm_out[:, not_h_ind])
+        
+        # NAR (h + prev_block)
+        h_pad = lm_out[:, h_ind] # B, N, D        
+        h = torch.cat([hh[:hl] for hh, hl in zip(h_pad, block_num)], dim=0) 
+        block_semantic = rearrange(sketch_emb, "b (n s) d -> b n s d", s=self.block_size) # B, N, 32, D 
+        current_block_semantic = torch.cat([bb[:bl] for bb, bl in zip(block_semantic, block_num)], dim=0) 
+        
+        if self.training: # for CFG
+            drop_h_idx = torch.rand((h.shape[0], 1), device=h.device) < self.h_dropout
+            h = torch.masked_fill(h, drop_h_idx, 0)
+            # current_block_semantic = torch.masked_fill(current_block_semantic, drop_h_idx.unsqueeze(-1), 0)
+            
+            drop_s_idx = torch.rand((current_block_semantic.shape[0], 1), device=current_block_semantic.device) < self.h_dropout
+            current_block_semantic = torch.masked_fill(current_block_semantic, drop_s_idx.unsqueeze(-1), 0)
+        
+        with torch.no_grad():
+            block_latent = rearrange(x_latent, "b d (n s) -> b n s d", s=self.block_size) # B, N, 32, D
+            current_block = torch.cat([bb[:bl] for bb, bl in zip(block_latent, block_num)], dim=0) 
+            prev_block = torch.cat([bb[:bl] for bb, bl in zip(F.pad(block_latent, (0,0,0,0,1,0)), block_num)], dim=0)
+       
+            # b_indices = torch.randperm(block_latent.shape[0])[:B*16]
+            # h, current_block, prev_block = h[b_indices], current_block[b_indices], prev_block[b_indices]
+            
+            orig_type = x_latent.dtype
+            with torch.cuda.amp.autocast(enabled=False):
+                if self.timestep_sampler == "uniform":
+                    # Draw uniformly distributed continuous timesteps
+                    t = self.rng.draw(h.shape[0])[:, 0].to(device=h.device, dtype=h.dtype)
+                elif self.timestep_sampler == "logit_normal":
+                    t = torch.sigmoid(torch.randn(h.shape[0], device=h.device, dtype=h.dtype))
+                elif self.timestep_sampler == "trunc_logit_normal":
+                    # Draw from logistic truncated normal distribution
+                    from ..musicldm.musicldm_pl import truncated_logistic_normal_rescaled
+                    t = truncated_logistic_normal_rescaled(h.shape[0]).to(h.device)
+                    # Flip the distribution
+                    t = 1 - t
+                    
+                # Calculate the noise schedule parameters for those timesteps
+                if self.diffusion_objective == "v":
+                    alphas, sigmas = get_alphas_sigmas(t)
+                elif self.diffusion_objective == "rectified_flow":
+                    alphas, sigmas = 1-t, t
+                # Combine the ground truth data and the noise
+                alphas = alphas[:, None, None]
+                sigmas = sigmas[:, None, None]
+                noise = torch.randn_like(current_block)
+                noised_inputs = current_block * alphas + noise * sigmas
+                if self.diffusion_objective == "v": # (a_t - a_{t-1})x_0 + (b_t-b_{t-1}) e = -b x_0 + a e
+                    targets = noise * alphas - current_block * sigmas
+                elif self.diffusion_objective == "rectified_flow": #||(XT-X0) - p(x_t, t)||      
+                    targets = noise - current_block
+    
+        nar_output = self.diffusion_forward(noised_inputs.to(orig_type), t.to(orig_type), h, current_block_semantic, prev_block)
+
+        return DiTAROutput(
+            ar_logit=x_sketch_logit,
+            ar_target=ar_target, 
+            nar_pred=nar_output,
+            nar_target=targets.to(orig_type),
+            nar_t=t
+        )
+
+
+
+    def lm_forward(self, sequence, condition_tensors: tp.Optional[ConditionTensors] = None) -> torch.Tensor:
+        # import pdb; pdb.set_trace()
+        B, T, D = sequence.shape
+        if self.pre_norm:
+            sequence = self.pre_norm(sequence.to(self.pre_norm.weight.data.dtype))
+            
+        input_, cross_attention_input = self.fuser(sequence, condition_tensors)
+
+        transformer_input = {
+            "inputs_embeds":input_,
+            "use_cache": self._is_streaming, 
+            "past_key_values": self._streaming_state.get('past_key_values', None),
+        }
+        if self.backend == 'bart': # TODO infer 的时候这个玩意不用重复算
+            # TODO attention_mask
+            cross_attention_input = self.cross_encoder(inputs_embeds=cross_attention_input) 
+            transformer_input["encoder_hidden_states"] = cross_attention_input.last_hidden_state
+
+        output = self.ar_transformer(**transformer_input)
+        if self._is_streaming:
+            self._streaming_state['past_key_values'] = output.past_key_values
+        out = output.last_hidden_state
+             
+
+            
+        if len(self.fuser.fuse2cond['prepend']) > 0:
+            out = out[:, -T:, :]
+
+        return out
+
+
+
+
+    def diffusion_forward(self, 
+                x: torch.Tensor,
+                t: torch.Tensor, # B,
+                h: torch.Tensor,
+                s: torch.Tensor, # B, self.block_size, D
+                history_x: torch.Tensor,
+                cfg_coef: float = None) -> torch.Tensor:
+
+        if cfg_coef is not None:
+            # only for infer
+            assert not self.training # only for inference
+            x = torch.cat([x,x], dim=0)
+            t = torch.cat([t,t], dim=0)
+            h = torch.cat([h,torch.zeros_like(h)], dim=0)
+            s = torch.cat([s,torch.zeros_like(s)], dim=0)
+            history_x = torch.cat([history_x,history_x], dim=0)
+            
+        B, T, _ = x.shape
+        
+        input_ = self.project_in(torch.cat([history_x, x], dim=1))
+        # print(h.shape, s.shape, input_.shape)
+        input_ = torch.cat([h.unsqueeze(1), s, input_], dim=1)
+        # Get the batch of timestep embeddings
+        timestep_embed = self.to_timestep_embed(self.timestep_features(t[:, None]))# (b, embed_dim)
+        # breakpoint()
+        if self.time_cond_type == "prepend":
+            input_ = torch.cat([timestep_embed.unsqueeze(1), input_], dim=1)
+
+        transformer_input = {
+            "x": input_,
+            "global_cond": timestep_embed if self.time_cond_type == "adaLN" else None}
+        
+        output = self.nar_dit(**transformer_input)
+
+        # remove the prefix from the model outputs
+        output = output[:, -T:, :]
+        output = self.project_out(output)
+
+        if cfg_coef is not None:
+            cond_output, uncond_output = torch.chunk(output, 2, dim=0)
+            output = uncond_output + (cond_output - uncond_output) * cfg_coef
+        
+        return output  # [B, T, D]
+
+
+
+    def _sample_next_block(self,
+                           sequence: torch.Tensor,
+                           prev_latents: torch.Tensor, 
+                           condition_tensors: tp.Optional[ConditionTensors] = None,
+                           cfg_coef: tp.Optional[tp.Union[float, tp.List[float]]] = None,
+                           steps: int = 50,
+                           dit_cfg_type: str = 'h',
+                           use_sampling: bool = False,
+                           temp: float = 1.0,
+                           diff_temp: float = 1.0, 
+                           top_k: int = 0,
+                           top_p: float = 0.0,
+                           penalty_token_pool: tp.Optional[list] = None) -> torch.Tensor:
+        # infer: lm next_token -> (if % block_sz == 0) infer diff
+        # 1. sample sketch (lm) -> 2. sample latent (lm+diff)
+        sequence = sequence.clone()
+        
+        if isinstance(cfg_coef, tp.Iterable):
+            assert len(cfg_coef) == 2
+            cfg_coef_lm, cfg_coef_diff = cfg_coef
+        else:
+            cfg_coef_lm, cfg_coef_diff = cfg_coef, cfg_coef
+        
+        B = sequence.shape[0]
+        # import pdb; pdb.set_trace()
+
+        if condition_tensors:
+            # Preparing for CFG, predicting both conditional and unconditional logits.
+            sequence = torch.cat([sequence, sequence], dim=0)
+        
+        
+        # ############### decode sketch #########################
+        next_tokens = []
+        next_token_embs = []
+            
+        for k in range(self.block_size):
+            if self._is_streaming and k > 0:
+                lm_inp = sequence[:,-1:]
+            else:
+                lm_inp = sequence
+                
+            lm_out = self.lm_forward(
+                lm_inp,
+                condition_tensors=condition_tensors)
+            next_pitch_logit = self.skeleton_classifier(lm_out[:, -1:]) # B, 1, card
+
+            if condition_tensors:
+                cond_logit, uncond_logit = next_pitch_logit.split(B, dim=0)  
+                next_pitch_logit = uncond_logit + (cond_logit - uncond_logit) * cfg_coef_lm
+
+            # add penalty to pre-sampled tokens
+            if penalty_token_pool is not None and len(penalty_token_pool) > 0: # B, T
+                for b in range(B):
+                    # q_count = torch.bincount(penalty_token_pool)
+                    q_count = torch.bincount(torch.unique(penalty_token_pool[b]))
+                    tmp = min(q_count.shape[-1], self.num_pitch - 1) 
+                    next_pitch_logit[b, -1, :tmp] /= (1.1 ** q_count[:tmp])
+                    
+            # sample k
+            if use_sampling and temp > 0.0:
+                probs = torch.softmax(next_pitch_logit  / temp, dim=-1)
+                if top_p > 0.0:
+                    next_token = sample_top_p(probs, p=top_p)
+                elif top_k > 0:
+                    next_token = sample_top_k(probs, k=top_k)
+                else:
+                    next_token = multinomial(probs, num_samples=1)
+                next_token = next_token.squeeze(-1)
+            else:
+                next_token = torch.argmax(next_pitch_logit, dim=-1) # B, 1
+            if penalty_token_pool is not None and len(penalty_token_pool) > 0: # B, T
+                penalty_token_pool = torch.cat([penalty_token_pool, next_token], dim=-1)[:,1:]
+            next_token_emb = self.skeleton_emb(next_token) #B, 1, d
+            next_tokens.append(next_token)
+            next_token_embs.append(next_token_emb)
+            
+            if condition_tensors:
+                doubled_next_emb = torch.cat([next_token_emb, next_token_emb], dim=0)
+                sequence = torch.cat([sequence, doubled_next_emb], dim=1)
+            else:
+                sequence = torch.cat([sequence, next_token_emb], dim=1)
+            
+        next_tokens = torch.cat(next_tokens, dim=1)
+        next_token_embs = torch.cat(next_token_embs, dim=1)
+        
+        # ############### decode latent ###########################
+        # 这里求h虽然double了 但是没用classifier-free guidance
+        if self._is_streaming:
+            lm_inp = sequence[:,-1:]
+        else:
+            lm_inp = sequence
+                
+        lm_out = self.lm_forward(
+            lm_inp,
+            condition_tensors=condition_tensors)
+        
+        h = lm_out[:,-1] 
+
+        noise = torch.randn((B, self.block_size, self.latent_dim), device=h.device, dtype=h.dtype)
+        
+        assert dit_cfg_type in ['h', 'global', 'none']
+        """
+        global: same cfg setting as next-token-prediction
+        none: no cfg
+        h: no cfg during ar-stage and apply cfg via ar output
+        """
+        if condition_tensors:
+            if dit_cfg_type == 'global':
+                noise = torch.cat([noise, noise], dim=0)
+                prev_latents = torch.cat([prev_latents, prev_latents], dim=0)
+                semantic_embs = torch.cat([next_token_embs, next_token_embs], dim=0)
+            else:
+                h, _ = h.chunk(2, dim=0)      
+                semantic_embs = next_token_embs
+
+        
+        if self.diffusion_objective == "v":
+            next_latent = sample(self.diffusion_forward, noise, steps=steps, eta=0, h=h, s=semantic_embs, history_x=prev_latents, 
+                                 cfg_coef=(cfg_coef_diff if dit_cfg_type=='h' else None))
+        elif self.diffusion_objective == "rectified_flow":
+            # next_latent = sample_discrete_euler(self.diffusion_forward, noise, steps=steps, h=h, s=semantic_embs, history_x=prev_latents, 
+            #                                     cfg_coef=(cfg_coef_diff if dit_cfg_type=='h' else None))
+            next_latent = sample_discrete_euler_with_temperature(self.diffusion_forward, noise, steps=steps, temperature=diff_temp, h=h, s=semantic_embs, history_x=prev_latents, 
+                                                cfg_coef=(cfg_coef_diff if dit_cfg_type=='h' else None))
+        if condition_tensors and dit_cfg_type == 'global':
+            cond_next_latent, uncond_next_latent = torch.chunk(next_latent, 2, dim=0)
+            next_latent = uncond_next_latent + (cond_next_latent - uncond_next_latent) * cfg_coef_diff
+            
+        latent_emb = self.block_conv(next_latent.transpose(1,2))
+
+        next_block_seq = torch.cat([next_token_embs, latent_emb], dim=1) # B, self.block_size+1, d
+        
+        return next_tokens, next_latent, next_block_seq
+        
+            
+        
+    @torch.no_grad()
+    def generate(self,
+                 prompt: tp.Optional[torch.Tensor] = None,
+                 conditions: tp.List[ConditioningAttributes] = [],
+                 cfg_coef: tp.Optional[tp.Union[float, tp.List[float]]] = None,
+                 steps=50,
+                 dit_cfg_type: str = 'h',
+                 max_frames: int = 1500, # 60 * 25
+                 use_sampling: bool = True,
+                 temp: float = 1.0,
+                 diff_temp: float = 1.0,
+                 top_k: int = 0,
+                 top_p: float = 0.0,
+                 penalty_repeat: bool = False,
+                 penalty_window: int = 50) -> torch.Tensor: 
+        assert not self.training, "generation shouldn't be used in training mode."
+
+        B = len(conditions)
+        assert B==1, "currently  do not support batch decoding"
+        null_conditions = ClassifierFreeGuidanceDropout(p=1.0)(conditions)
+        conditions = conditions + null_conditions
+        tokenized = self.condition_provider.tokenize(conditions)
+        condition_tensors = self.condition_provider(tokenized)          
+        
+        
+        sequence = self.bos_token.reshape(1,1,-1).expand(B, 1, -1)
+        if prompt is not None:
+            # TODO 
+            raise NotImplementedError
+            # sequence = torch.cat([sequence, prompt])
+
+            
+        prev_blocks = torch.zeros((B, self.block_size, self.latent_dim), device=sequence.device, dtype=sequence.dtype)
+        latent_seq, token_seq = None, None
+            
+        with self.streaming():
+            prog_bar = tqdm.tqdm()
+            while True:
+                if token_seq is None or not penalty_repeat:
+                    penalty_token_pool = None
+                else:
+                    penalty_token_pool = token_seq[: ,-penalty_window:]
+                    if penalty_token_pool.shape[-1] < penalty_window:
+                        penalty_token_pool = F.pad(penalty_token_pool, (penalty_window - penalty_token_pool.shape[-1], 0), value=self.eos_token_id)
+                next_tokens, next_latent, next_block_seq = self._sample_next_block(sequence[:, -1: ], prev_blocks, condition_tensors, 
+                                                                                   cfg_coef=cfg_coef, steps=steps, dit_cfg_type=dit_cfg_type,
+                                                                                    use_sampling=use_sampling, temp=temp, diff_temp=diff_temp,
+                                                                                    top_k=top_k, top_p=top_p,
+                                                                                    penalty_token_pool=penalty_token_pool)
+                
+                if (next_tokens == self.eos_token_id).any() or sequence.shape[1] > max_frames  / self.block_size * (self.block_size+1):
+                    break
+                
+                latent_seq = next_latent if latent_seq is None else torch.cat([latent_seq, next_latent], dim=1) # B,T, D
+                token_seq = next_tokens if token_seq is None else torch.cat([token_seq, next_tokens], dim=1) # B,T
+                sequence = torch.cat([sequence, next_block_seq], dim=1)
+                prev_blocks = next_latent
+                
+                prog_bar.update(self.block_size)
+                
+                
+        if latent_seq is None:
+            latent_seq = prev_blocks
+        return latent_seq.transpose(1,2), token_seq    
+        
+        
+    
+    
+    def init_weights(self, init_std=0.02):
+        
+        def _init_weights(module, init_std=0.02):
+            if isinstance(module, nn.Linear):
+                module.weight.data.normal_(mean=0.0, std=init_std)
+                # torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+            elif isinstance(module, nn.Embedding):
+                module.weight.data.normal_(mean=0.0, std=init_std)
+                if module.padding_idx is not None:
+                    module.weight.data[module.padding_idx].zero_()
+                    
+        self.apply(partial(_init_weights, init_std=init_std))

SongBloom/models/songbloom/songbloom_pl.py

@@ -0,0 +1,224 @@
+
+from functools import partial
+import typing as tp
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import torchaudio
+import numpy as np
+import random
+from omegaconf import OmegaConf
+import copy
+import lightning as pl
+
+import os, sys
+
+from ..musicgen.conditioners import WavCondition, JointEmbedCondition, ConditioningAttributes
+from ..vae_frontend import StableVAE
+from .songbloom_mvsa import MVSA_DiTAR
+from ...g2p.lyric_common import key2processor, symbols, LABELS
+
+
+os.environ['TOKENIZERS_PARALLELISM'] = "false"
+
+
+class SongBloom_PL(pl.LightningModule):
+    def __init__(self, cfg):
+        super().__init__()
+        # 关闭自动优化
+        # self.automatic_optimization = False
+
+        self.cfg = cfg
+
+        # Build VAE
+        self.vae = StableVAE(**cfg.vae).eval()
+        assert self.cfg.model['latent_dim'] == self.vae.channel_dim
+
+            
+        self.save_hyperparameters(cfg)
+        if self.vae is not None:
+            for param in self.vae.parameters():
+                param.requires_grad = False
+                
+        # Build DiT
+        model_cfg = OmegaConf.to_container(copy.deepcopy(cfg.model), resolve=True)
+        for cond_name in model_cfg["condition_provider_cfg"]:
+            if model_cfg["condition_provider_cfg"][cond_name]['type'] == 'audio_tokenizer_wrapper':
+                model_cfg["condition_provider_cfg"][cond_name]["audio_tokenizer"] = self.vae
+                model_cfg["condition_provider_cfg"][cond_name]["cache"] = False
+        
+        
+        self.model = MVSA_DiTAR(**model_cfg)
+        # print(self.model)
+        
+
+
+
+
+
+####################################
+
+class SongBloom_Sampler:    
+    
+    def __init__(self, compression_model: StableVAE, diffusion: MVSA_DiTAR, lyric_processor_key,
+                 max_duration: float, prompt_duration: tp.Optional[float] = None):
+        self.compression_model = compression_model
+        self.diffusion = diffusion
+        self.lyric_processor_key = lyric_processor_key
+        self.lyric_processor = key2processor.get(lyric_processor_key) if lyric_processor_key is not None else lambda x: x
+        # import pdb; pdb.set_trace()
+
+        assert max_duration is not None
+        self.max_duration: float = max_duration
+        self.prompt_duration = prompt_duration
+        
+        
+        self.device = next(iter(diffusion.parameters())).device
+        self.generation_params: dict = {}
+        # self.set_generation_params(duration=15)  # 15 seconds by default
+        self.set_generation_params(cfg_coef=1.5, steps=50, dit_cfg_type='h',
+                                   use_sampling=True, top_k=200, max_frames=self.max_duration * 25)
+        self._progress_callback: tp.Optional[tp.Callable[[int, int], None]] = None
+
+    @classmethod
+    def build_from_trainer(cls, cfg, strict=True, dtype=torch.float32):
+        model_light = SongBloom_PL(cfg)
+        incompatible = model_light.load_state_dict(torch.load(cfg.pretrained_path, map_location='cpu'), strict=strict)
+        
+        lyric_processor_key = cfg.train_dataset.lyric_processor
+    
+        print(incompatible)
+        
+        model_light = model_light.eval().cuda().to(dtype=dtype)  
+        model = cls(
+            compression_model = model_light.vae,
+            diffusion = model_light.model,
+            lyric_processor_key = lyric_processor_key,
+            max_duration = cfg.max_dur,
+            prompt_duration = cfg.sr * cfg.train_dataset.prompt_len
+            
+        )
+        model.set_generation_params(**cfg.inference)
+        return model
+        
+    @property
+    def frame_rate(self) -> float:
+        """Roughly the number of AR steps per seconds."""
+        return self.compression_model.frame_rate
+
+    @property
+    def sample_rate(self) -> int:
+        """Sample rate of the generated audio."""
+        return self.compression_model.sample_rate
+
+
+    def set_generation_params(self, **kwargs):
+        """Set the generation parameters."""
+        self.generation_params.update(kwargs)
+
+    # Mulan Inference
+    @torch.no_grad()
+    def generate(self, lyrics, prompt_wav) -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]:
+        """ Generate samples conditioned on text and melody.
+        """
+        # breakpoint()
+        assert prompt_wav.ndim == 2
+        if self.prompt_duration is not None:
+            prompt_wav = prompt_wav[..., :self.prompt_duration]
+            
+        attributes, _ = self._prepare_tokens_and_attributes(conditions={"lyrics": [self._process_lyric(lyrics)], "prompt_wav": [prompt_wav]}, 
+                                                                        prompt=None, prompt_tokens=None)
+
+        # breakpoint()
+        print(self.generation_params)
+        latent_seq, token_seq = self.diffusion.generate(None, attributes, **self.generation_params)
+        # print(token_seq)
+        audio_recon = self.compression_model.decode(latent_seq).float()
+        
+        return audio_recon
+    
+
+    def _process_lyric(self, input_lyric):
+        if self.lyric_processor_key == 'pinyin':
+            processed_lyric = self.lyric_processor(input_lyric)
+        else:
+            processed_lyric = []
+            check_lyric = input_lyric.split(" ")
+            for ii in range(len(check_lyric)):
+                if check_lyric[ii] not in symbols and check_lyric[ii] not in LABELS.keys() and len(check_lyric[ii]) > 0:
+                    new = self.lyric_processor(check_lyric[ii])
+                    check_lyric[ii] = new
+            processed_lyric = " ".join(check_lyric)
+        
+        return processed_lyric
+    
+    @torch.no_grad()
+    def _prepare_tokens_and_attributes(
+            self,
+            conditions: tp.Dict[str, tp.List[tp.Union[str, torch.Tensor]]],
+            prompt: tp.Optional[torch.Tensor],
+            prompt_tokens: tp.Optional[torch.Tensor] = None,
+    ) -> tp.Tuple[tp.List[ConditioningAttributes], tp.Optional[torch.Tensor]]:
+        """Prepare model inputs.
+
+        Args:
+            descriptions (list of str): A list of strings used as text conditioning.
+            prompt (torch.Tensor): A batch of waveforms used for continuation.
+            melody_wavs (torch.Tensor, optional): A batch of waveforms
+                used as melody conditioning. Defaults to None.
+        """
+        batch_size = len(list(conditions.values())[0])
+        assert batch_size == 1
+        # breakpoint()
+        attributes = [ConditioningAttributes() for _ in range(batch_size)]
+        for k in self.diffusion.condition_provider.conditioners:
+            conds = conditions.pop(k, [None for _ in attributes])
+            for attr, cond in zip(attributes, conds):
+                if self.diffusion.condition_provider.conditioner_type[k] == 'wav':
+                    if cond is None:
+                        attr.wav[k] = WavCondition(
+                            torch.zeros((1, 1, 1), device=self.device),
+                            torch.tensor([0], device=self.device).long(),
+                            sample_rate=[self.sample_rate],
+                            path=[None])  
+                    else:
+                        attr.wav[k] = WavCondition(
+                            cond.to(device=self.device).unsqueeze(0), # 1,C,T .mean(dim=0, keepdim=True)
+                            torch.tensor([cond.shape[-1]], device=self.device).long(),
+                            sample_rate=[self.sample_rate],
+                            path=[None])  
+                elif self.diffusion.condition_provider.conditioner_type[k] == 'text':
+                    attr.text[k] = cond
+                elif self.diffusion.condition_provider.conditioner_type[k] == 'joint_embed':
+                    if cond is None or isinstance(cond, str):
+                        attr.joint_embed[k] = JointEmbedCondition(
+                            torch.zeros((1, 1, 1), device=self.device),
+                            [cond],
+                            torch.tensor([0], device=self.device).long(),
+                            sample_rate=[self.sample_rate],
+                            path=[None])  
+                    elif isinstance(cond, torch.Tensor):
+                        attr.joint_embed[k] = JointEmbedCondition(
+                            cond.to(device=self.device).mean(dim=0, keepdim=True).unsqueeze(0),
+                            [None], 
+                            torch.tensor([cond.shape[-1]], device=self.device).long(),
+                            sample_rate=[self.sample_rate],
+                            path=[None])  
+                    else:
+                        raise NotImplementedError
+        assert conditions == {}, f"Find illegal conditions: {conditions}, support keys: {self.lm.condition_provider.conditioners}"
+        # breakpoint()
+        print(attributes)
+        
+        if prompt_tokens is not None:
+            prompt_tokens = prompt_tokens.to(self.device)
+            assert prompt is None
+        elif prompt is not None:
+            assert len(attributes) == len(prompt), "Prompt and nb. attributes doesn't match"
+            prompt = prompt.to(self.device)
+            prompt_tokens = self.compression_model.encode(prompt)
+        else:
+            prompt_tokens = None
+
+        return attributes, prompt_tokens
+

SongBloom/models/transformer.py

@@ -0,0 +1,937 @@
+from functools import reduce, partial
+from packaging import version
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+from torch.cuda.amp import autocast
+from typing import Callable, Literal
+import os, sys
+import warnings
+from torch.utils import checkpoint
+from transformers.utils import is_flash_attn_2_available
+
+try:
+    assert is_flash_attn_2_available()
+    assert torch.cuda.get_device_capability(torch.device("cuda")) >= (8, 0)
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, unpad_input, pad_input
+    assert os.environ.get("DISABLE_FLASH_ATTN",'0') != "1"
+except Exception as e:
+    flash_attn_kvpacked_func = None
+    flash_attn_func = None
+    warnings.warn("Not support flash-attn!")
+
+try:
+    import natten
+except ImportError:
+    natten = None
+
+def checkpoint(function, *args, **kwargs):
+    kwargs.setdefault("use_reentrant", False)
+    return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
+
+
+# Copied and modified from https://github.com/lucidrains/x-transformers/blob/main/x_transformers/attend.py under MIT License
+# License can be found in LICENSES/LICENSE_XTRANSFORMERS.txt
+
+def create_causal_mask(i, j, device):
+    return torch.ones((i, j), device = device, dtype = torch.bool).triu(j - i + 1)
+
+def or_reduce(masks):
+    head, *body = masks
+    for rest in body:
+        head = head | rest
+    return head
+
+# positional embeddings
+
+class AbsolutePositionalEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.max_seq_len = max_seq_len
+        self.emb = nn.Embedding(max_seq_len, dim)
+
+    def forward(self, x, pos = None, seq_start_pos = None):
+        seq_len, device = x.shape[1], x.device
+        assert seq_len <= self.max_seq_len, f'you are passing in a sequence length of {seq_len} but your absolute positional embedding has a max sequence length of {self.max_seq_len}'
+
+        if pos is None:
+            pos = torch.arange(seq_len, device = device)
+
+        if seq_start_pos is not None:
+            pos = (pos - seq_start_pos[..., None]).clamp(min = 0)
+
+        pos_emb = self.emb(pos)
+        pos_emb = pos_emb * self.scale
+        return pos_emb
+
+class ScaledSinusoidalEmbedding(nn.Module):
+    def __init__(self, dim, theta = 10000):
+        super().__init__()
+        assert (dim % 2) == 0, 'dimension must be divisible by 2'
+        self.scale = nn.Parameter(torch.ones(1) * dim ** -0.5)
+
+        half_dim = dim // 2
+        freq_seq = torch.arange(half_dim).float() / half_dim
+        inv_freq = theta ** -freq_seq
+        self.register_buffer('inv_freq', inv_freq, persistent = False)
+
+    def forward(self, x, pos = None, seq_start_pos = None):
+        seq_len, device = x.shape[1], x.device
+
+        if pos is None:
+            pos = torch.arange(seq_len, device = device)
+
+        if seq_start_pos is not None:
+            pos = pos - seq_start_pos[..., None]
+
+        emb = einsum('i, j -> i j', pos, self.inv_freq)
+        emb = torch.cat((emb.sin(), emb.cos()), dim = -1)
+        return emb * self.scale
+    
+class RotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim,
+        use_xpos = False,
+        scale_base = 512,
+        interpolation_factor = 1.,
+        base = 10000,
+        base_rescale_factor = 1.
+    ):
+        super().__init__()
+        # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+        # has some connection to NTK literature
+        # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+        base *= base_rescale_factor ** (dim / (dim - 2))
+
+        inv_freq = 1. / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq, persistent=False)
+
+        assert interpolation_factor >= 1.
+        self.interpolation_factor = interpolation_factor
+
+        if not use_xpos:
+            self.register_buffer('scale', None)
+        else:
+            scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
+
+            self.scale_base = scale_base
+            self.register_buffer('scale', scale)
+
+    def forward_from_seq_len(self, seq_len):
+        device = self.inv_freq.device
+
+        t = torch.arange(seq_len, device = device)
+        return self.forward(t)
+
+    @autocast(enabled = False)
+    def forward(self, t):
+        device = self.inv_freq.device
+
+        t = t.to(torch.float32)
+        seq_len = t.shape[0]
+
+        t = t / self.interpolation_factor
+
+        freqs = torch.einsum('i , j -> i j', t, self.inv_freq)
+        freqs = torch.cat((freqs, freqs), dim = -1)
+
+        if self.scale is None:
+            return freqs, 1.
+
+        power = (torch.arange(seq_len, device = device) - (seq_len // 2)) / self.scale_base
+        scale = self.scale ** rearrange(power, 'n -> n 1')
+        scale = torch.cat((scale, scale), dim = -1)
+
+        return freqs, scale
+
+class RotaryEmbedding2D(RotaryEmbedding):
+    def __init__(self, dim, w, **kwargs):
+        super().__init__(dim // 2, **kwargs)
+        self.w = w
+        
+    
+    def forward_from_seq_len(self, seq_len):
+        device = self.inv_freq.device
+        assert seq_len % self.w == 0 , f"{seq_len} % {self.w} != 0"
+        h_len = seq_len // self.w
+        
+        t_h = torch.arange(h_len, device = device)
+        t_w = torch.arange(self.w, device = device)
+        
+        return self.forward(t_h, t_w)
+
+    @autocast(enabled = False)
+    def forward(self, t_h: torch.Tensor, t_w: torch.Tensor):
+        repeat_t_h = t_h.repeat_interleave(t_w.shape[0], dim=0)
+        repeat_t_w = t_w.repeat(t_h.shape[0])
+        freq_h, scale_h = super().forward(repeat_t_h)
+        freq_w, scale_w = super().forward(repeat_t_w)
+        freq = torch.stack([freq_h, freq_w], dim=-1) #h*w, D//2, 2
+        freq = torch.cat(torch.unbind(freq, dim=-2), dim=-1)
+        
+        if self.scale is None:
+            scale = 1.
+        else:
+            scale = torch.stack([scale_h, scale_w], dim=-1)
+            scale = torch.cat(torch.unbind(scale, dim=-2), dim=-1)
+            
+        return freq, scale
+        
+        
+    
+
+def rotate_half(x):
+    x = rearrange(x, '... (j d) -> ... j d', j = 2)
+    x1, x2 = x.unbind(dim = -2)
+    return torch.cat((-x2, x1), dim = -1)
+
+@autocast(enabled = False)
+def apply_rotary_pos_emb(t, freqs, scale = 1):
+    out_dtype = t.dtype
+
+    # cast to float32 if necessary for numerical stability
+    dtype = reduce(torch.promote_types, (t.dtype, freqs.dtype, torch.float32))
+    rot_dim, seq_len = freqs.shape[-1], t.shape[-2]
+    freqs, t = freqs.to(dtype), t.to(dtype)
+    freqs = freqs[-seq_len:, :]
+
+    if t.ndim == 4 and freqs.ndim == 3:
+        freqs = rearrange(freqs, 'b n d -> b 1 n d')
+
+    # partial rotary embeddings, Wang et al. GPT-J
+    t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:]
+    t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
+
+    t, t_unrotated = t.to(out_dtype), t_unrotated.to(out_dtype)
+    return torch.cat((t, t_unrotated), dim = -1)
+
+# norms
+class LayerNorm(nn.Module):
+    def __init__(self, dim, bias=False, fix_scale=False):
+        """
+        bias-less layernorm has been shown to be more stable. most newer models have moved towards rmsnorm, also bias-less
+        """
+        super().__init__()
+
+        if fix_scale:
+            self.register_buffer("gamma", torch.ones(dim))
+        else:
+            self.gamma = nn.Parameter(torch.ones(dim))
+
+        if bias:
+            self.beta = nn.Parameter(torch.zeros(dim))
+        else:
+            self.register_buffer("beta", torch.zeros(dim))
+
+
+    def forward(self, x):
+        return F.layer_norm(x, x.shape[-1:], weight=self.gamma, bias=self.beta)
+
+# feedforward
+
+class GLU(nn.Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        activation: Callable,
+        use_conv = False,
+        conv_kernel_size = 3,
+        bias = False,
+    ):
+        super().__init__()
+        self.act = activation
+        self.up_proj = nn.Linear(dim_in, dim_out, bias=bias) if not use_conv else nn.Conv1d(dim_in, dim_out, conv_kernel_size, padding = (conv_kernel_size // 2))
+        self.gate_proj = nn.Linear(dim_in, dim_out, bias=bias) if not use_conv else nn.Conv1d(dim_in, dim_out, conv_kernel_size, padding = (conv_kernel_size // 2))
+        self.use_conv = use_conv
+
+    def forward(self, x):
+        if self.use_conv:
+            x = rearrange(x, 'b n d -> b d n')
+            gate = self.gate_proj(x)
+            x = self.up_proj(x)
+            x = rearrange(x, 'b d n -> b n d')
+            gate = rearrange(gate, 'b d n -> b n d')
+        else:
+            gate = self.gate_proj(x)
+            x = self.up_proj(x)
+
+        return x * self.act(gate)
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out = None,
+        dim_ff = None,
+        no_bias = False,
+        glu = True,
+        use_conv = False,
+        conv_kernel_size = 3,
+        zero_init_output = True,
+    ):
+        super().__init__()
+        inner_dim = dim_ff if dim_ff is not None else 4 * dim
+
+        # Default to SwiGLU
+
+        activation = nn.SiLU()
+
+        dim_out = dim if dim_out is None else dim_out
+
+        if glu:
+            linear_in = GLU(dim, inner_dim, activation, bias=not no_bias)
+        else:
+            linear_in = nn.Sequential(
+                Rearrange('b n d -> b d n') if use_conv else nn.Identity(),
+                nn.Linear(dim, inner_dim, bias = not no_bias) if not use_conv else nn.Conv1d(dim, inner_dim, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias),
+                Rearrange('b n d -> b d n') if use_conv else nn.Identity(),
+                activation
+            )
+
+        linear_out = nn.Linear(inner_dim, dim_out, bias = not no_bias) if not use_conv else nn.Conv1d(inner_dim, dim_out, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias)
+
+        # init last linear layer to 0
+        if zero_init_output:
+            nn.init.zeros_(linear_out.weight)
+            if not no_bias:
+                nn.init.zeros_(linear_out.bias)
+
+
+        self.ff = nn.Sequential(
+            linear_in,
+            Rearrange('b d n -> b n d') if use_conv else nn.Identity(),
+            linear_out,
+            Rearrange('b n d -> b d n') if use_conv else nn.Identity(),
+        )
+
+    def forward(self, x):
+        return self.ff(x)
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_heads = 64,
+        dim_context = None,
+        causal = False,
+        zero_init_output=True,
+        qk_norm: Literal['l2', 'ln', 'none'] = 'none',
+        natten_kernel_size = None
+    ):
+        super().__init__()
+        self.dim = dim
+        self.dim_heads = dim_heads
+        self.causal = causal
+
+        dim_kv = dim_context if dim_context is not None else dim
+        
+        self.num_heads = dim // dim_heads
+        self.kv_heads = dim_kv // dim_heads
+
+        if dim_context is not None:
+            self.to_q = nn.Linear(dim, dim, bias=False)
+            self.to_kv = nn.Linear(dim_kv, dim_kv * 2, bias=False)
+        else:
+            self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
+
+        self.to_out = nn.Linear(dim, dim, bias=False)
+
+        if zero_init_output:
+            nn.init.zeros_(self.to_out.weight)
+
+        self.qk_norm = qk_norm
+
+        if self.qk_norm == "ln":
+            self.q_norm = nn.LayerNorm(dim_heads, elementwise_affine=True, eps=1.0e-6)
+            self.k_norm = nn.LayerNorm(dim_heads, elementwise_affine=True, eps=1.0e-6)
+
+        # Using 1d neighborhood attention
+        self.natten_kernel_size = natten_kernel_size
+        if natten_kernel_size is not None:
+            return
+
+        self.use_pt_flash = torch.cuda.is_available() and version.parse(torch.__version__) >= version.parse('2.0.0')
+
+        self.use_fa_flash = torch.cuda.is_available() and flash_attn_func is not None
+
+        self.sdp_kwargs = dict(
+            enable_flash = True,
+            enable_math = True,
+            enable_mem_efficient = True
+        )
+
+    def flash_attn(
+            self,
+            q, 
+            k, 
+            v,
+            mask = None,
+            causal = None
+    ):
+        batch, heads, q_len, _, k_len, device = *q.shape, k.shape[-2], q.device
+        kv_heads = k.shape[1]
+        # Recommended for multi-query single-key-value attention by Tri Dao
+        # kv shape torch.Size([1, 512, 64]) -> torch.Size([1, 8, 512, 64])
+
+        if heads != kv_heads:
+            # Repeat interleave kv_heads to match q_heads
+            heads_per_kv_head = heads // kv_heads
+            k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
+
+        if k.ndim == 3:
+            k = rearrange(k, 'b ... -> b 1 ...').expand_as(q)
+
+        if v.ndim == 3:
+            v = rearrange(v, 'b ... -> b 1 ...').expand_as(q)
+
+        causal = self.causal if causal is None else causal
+
+        if q_len == 1 and causal:
+            causal = False
+        
+        if mask is not None:
+            assert mask.ndim == 4
+            mask = mask.expand(batch, heads, q_len, k_len)
+
+        # handle kv cache - this should be bypassable in updated flash attention 2
+
+        if k_len > q_len and causal:
+            causal_mask = self.create_causal_mask(q_len, k_len, device = device)
+            if mask is None:
+                mask = ~causal_mask
+            else:
+                mask = mask & ~causal_mask
+            causal = False
+
+        # manually handle causal mask, if another mask was given
+
+        row_is_entirely_masked = None
+
+        if mask is not None and causal:
+            causal_mask = self.create_causal_mask(q_len, k_len, device = device)
+            mask = mask & ~causal_mask
+
+            # protect against an entire row being masked out
+
+            row_is_entirely_masked = ~mask.any(dim = -1)
+            mask[..., 0] = mask[..., 0] | row_is_entirely_masked
+
+            causal = False
+        
+        with torch.backends.cuda.sdp_kernel(**self.sdp_kwargs):
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask = mask,
+                is_causal = causal
+            )
+
+        # for a row that is entirely masked out, should zero out the output of that row token
+
+        if row_is_entirely_masked is not None:
+            out = out.masked_fill(row_is_entirely_masked[..., None], 0.)
+
+        return out
+
+    def forward(
+        self,
+        x,
+        context = None,
+        mask = None,
+        context_mask = None,
+        rotary_pos_emb = None,
+        causal = None
+    ):
+        h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None
+
+        kv_input = context if has_context else x
+
+        if hasattr(self, 'to_q'):
+            # Use separate linear projections for q and k/v
+            q = self.to_q(x)
+            q = rearrange(q, 'b n (h d) -> b h n d', h = h)
+
+            k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+            k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = kv_h), (k, v))
+        else:
+            # Use fused linear projection
+            q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+            q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
+        
+        # Normalize q and k for cosine sim attention
+        if self.qk_norm == "l2":
+            q = F.normalize(q, dim=-1)
+            k = F.normalize(k, dim=-1)
+        elif self.qk_norm == "ln":
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if rotary_pos_emb is not None and not has_context:
+            freqs, _ = rotary_pos_emb
+
+            q_dtype = q.dtype
+            k_dtype = k.dtype
+
+            q = q.to(torch.float32)
+            k = k.to(torch.float32)
+            freqs = freqs.to(torch.float32)
+
+            q = apply_rotary_pos_emb(q, freqs)
+            k = apply_rotary_pos_emb(k, freqs)
+
+            q = q.to(q_dtype)
+            k = k.to(k_dtype)
+        
+        # TODO 这里这俩都是 [B, k/Q_len]这样的格式
+        # context mask也许应该改成 [B, Q_len, K_len]
+        # 并且下面flash_attn 默认假设attn靠左部分全为1
+        input_mask = context_mask # cross-attn
+        if input_mask is None and not has_context: # self-attn
+            input_mask = mask
+
+        # determine masking
+        masks = []
+        final_attn_mask = None # The mask that will be applied to the attention matrix, taking all masks into account
+
+        if input_mask is not None:
+            input_mask = rearrange(input_mask, 'b j -> b 1 1 j')
+            masks.append(~input_mask)
+
+        # Other masks will be added here later
+
+        if len(masks) > 0:
+            final_attn_mask = ~or_reduce(masks)
+
+        n, device = q.shape[-2], q.device
+
+        causal = self.causal if causal is None else causal
+        if n == 1 and causal:
+            causal = False
+        if self.natten_kernel_size is not None:
+            if natten is None:
+                raise ImportError('natten not installed, please install natten to use neighborhood attention')
+            
+            dtype_in = q.dtype
+            q, k, v = map(lambda t: t.to(torch.float32), (q, k, v))
+
+            attn = natten.functional.natten1dqk(q, k, kernel_size = self.natten_kernel_size, dilation=1)
+
+            if final_attn_mask is not None:
+                attn = attn.masked_fill(final_attn_mask, -torch.finfo(attn.dtype).max)
+
+            attn = F.softmax(attn, dim=-1, dtype=torch.float32)
+
+            out = natten.functional.natten1dav(attn, v, kernel_size = self.natten_kernel_size, dilation=1).to(dtype_in)
+
+        # Prioritize Flash Attention 2
+        elif self.use_fa_flash:
+            fa_dtype_in = q.dtype
+            if q.dtype in [torch.float, torch.float32]:
+                target_dtype = self.to_out.weight.dtype if self.to_out.weight.dtype not in [torch.float, torch.float32] else torch.float16
+                warnings.warn(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+                )
+                q, k, v = map(lambda t: t.to(target_dtype), (q, k, v))
+            q, k, v = map(lambda t: rearrange(t, 'b h n d -> b n h d'), (q, k, v))
+            # out = flash_attn_func(q, k, v, causal = causal) 
+            if final_attn_mask is not None:
+                # Check if the mask meets the requirement of FlashAttn
+                kv_seq_mask = final_attn_mask.squeeze(dim=[1,2])
+                kv_reallens = kv_seq_mask.sum(dim=-1, dtype=torch.int32)
+                first_zero_indices = torch.argmax((kv_seq_mask == 0).int(), dim=1).masked_fill(kv_seq_mask[:,-1] != 0, kv_seq_mask.shape[1])
+                assert (kv_reallens == first_zero_indices).all(), f'{kv_reallens} , {first_zero_indices}'
+                
+                batch_size, kv_seq_len, num_key_value_heads, head_dim = k.shape
+                unpad_k, indices_k, cu_seqlens_k, max_seqlen_in_batch_k = unpad_input(k, kv_seq_mask)
+                unpad_v = index_first_axis(
+                    v.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+                )
+                q_seq_len = q.shape[1]
+                unpad_q, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(q, torch.ones((batch_size, q_seq_len), device=q.device, dtype=torch.bool))
+                # print(q.shape, k.shape)
+                # print(cu_seqlens_q, cu_seqlens_k)
+                # breakpoint()
+                out_unpad = flash_attn_varlen_func(
+                    unpad_q,
+                    unpad_k,
+                    unpad_v,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_in_batch_q,
+                    max_seqlen_k=max_seqlen_in_batch_k,
+                    causal=causal,
+                )
+                out = pad_input(out_unpad, indices_q, batch_size, q_seq_len)
+            else:
+                out = flash_attn_func(q, k, v, causal = causal)
+                
+                
+            out = rearrange(out.to(fa_dtype_in), 'b n h d -> b h n d')
+        # Fall back to PyTorch implementation
+        elif self.use_pt_flash:
+            out = self.flash_attn(q, k, v, causal = causal, mask = final_attn_mask)
+
+        else:
+            # Fall back to custom implementation
+
+            if h != kv_h:
+                # Repeat interleave kv_heads to match q_heads
+                heads_per_kv_head = h // kv_h
+                k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
+
+            scale = 1. / (q.shape[-1] ** 0.5)
+
+            kv_einsum_eq = 'b j d' if k.ndim == 3 else 'b h j d'
+
+            dots = einsum(f'b h i d, {kv_einsum_eq} -> b h i j', q, k) * scale
+            
+            i, j, dtype = *dots.shape[-2:], dots.dtype
+
+            mask_value = -torch.finfo(dots.dtype).max
+
+            if final_attn_mask is not None:
+                dots = dots.masked_fill(~final_attn_mask, mask_value)
+
+            if causal:
+                causal_mask = self.create_causal_mask(i, j, device = device)
+                dots = dots.masked_fill(causal_mask, mask_value)
+
+            attn = F.softmax(dots, dim=-1, dtype=torch.float32)
+            attn = attn.type(dtype)
+
+            out = einsum(f'b h i j, {kv_einsum_eq} -> b h i d', attn, v)
+
+        # merge heads
+        out = rearrange(out, ' b h n d -> b n (h d)')
+
+        # Communicate between heads
+        
+        # with autocast(enabled = False):
+        #     out_dtype = out.dtype
+        #     out = out.to(torch.float32)
+        #     out = self.to_out(out).to(out_dtype)
+        out = self.to_out(out)
+
+        if mask is not None:
+            mask = rearrange(mask, 'b n -> b n 1')
+            out = out.masked_fill(~mask, 0.)
+
+        return out
+
+class ConformerModule(nn.Module):
+    def __init__(
+        self,
+        dim,
+        norm_kwargs = {},
+    ):     
+
+        super().__init__()
+
+        self.dim = dim
+        
+        self.in_norm = LayerNorm(dim, **norm_kwargs)
+        self.pointwise_conv = nn.Conv1d(dim, dim, kernel_size=1, bias=False)
+        self.glu = GLU(dim, dim, nn.SiLU())
+        self.depthwise_conv = nn.Conv1d(dim, dim, kernel_size=17, groups=dim, padding=8, bias=False)
+        self.mid_norm = LayerNorm(dim, **norm_kwargs) # This is a batch norm in the original but I don't like batch norm
+        self.swish = nn.SiLU()
+        self.pointwise_conv_2 = nn.Conv1d(dim, dim, kernel_size=1, bias=False)
+
+    def forward(self, x):
+        x = self.in_norm(x)
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.pointwise_conv(x)
+        x = rearrange(x, 'b d n -> b n d')
+        x = self.glu(x)
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.depthwise_conv(x)
+        x = rearrange(x, 'b d n -> b n d')
+        x = self.mid_norm(x)
+        x = self.swish(x)
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.pointwise_conv_2(x)
+        x = rearrange(x, 'b d n -> b n d')
+
+        return x
+
+class TransformerBlock(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_heads = 64,
+            cross_attend = False,
+            dim_context = None,
+            global_cond_dim = None,
+            causal = False,
+            zero_init_branch_outputs = True,
+            conformer = False,
+            layer_ix = -1,
+            remove_norms = False,
+            attn_kwargs = {},
+            ff_kwargs = {},
+            norm_kwargs = {}
+    ):
+        
+        super().__init__()
+        self.dim = dim
+        self.dim_heads = dim_heads
+        self.cross_attend = cross_attend
+        self.dim_context = dim_context
+        self.causal = causal
+
+        self.pre_norm = LayerNorm(dim, **norm_kwargs) if not remove_norms else nn.Identity()
+
+        self.self_attn = Attention(
+            dim,
+            dim_heads = dim_heads,
+            causal = causal,
+            zero_init_output=zero_init_branch_outputs,
+            **attn_kwargs
+        )
+
+        if cross_attend:
+            self.cross_attend_norm = LayerNorm(dim, **norm_kwargs) if not remove_norms else nn.Identity()
+            self.cross_attn = Attention(
+                dim,
+                dim_heads = dim_heads,
+                dim_context=dim_context,
+                causal = causal,
+                zero_init_output=zero_init_branch_outputs,
+                **attn_kwargs
+            )
+        
+        self.ff_norm = LayerNorm(dim, **norm_kwargs) if not remove_norms else nn.Identity()
+        self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs, **ff_kwargs)
+
+        self.layer_ix = layer_ix
+
+        self.conformer = ConformerModule(dim, norm_kwargs=norm_kwargs) if conformer else None
+
+        self.global_cond_dim = global_cond_dim
+
+        if global_cond_dim is not None:
+            self.to_scale_shift_gate = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(global_cond_dim, dim * 6, bias=False)
+            )
+
+            nn.init.zeros_(self.to_scale_shift_gate[1].weight)
+            #nn.init.zeros_(self.to_scale_shift_gate_self[1].bias)
+
+    def forward(
+        self,
+        x,
+        mask = None,
+        global_cond=None,
+        context = None,
+        context_mask = None,
+        rotary_pos_emb = None
+    ):
+        if self.global_cond_dim is not None:
+            assert global_cond is not None
+            # scale_self, shift_self, gate_self, scale_ff, shift_ff, gate_ff = checkpoint(self.to_scale_shift_gate, global_cond).unsqueeze(1).chunk(6, dim = -1)
+            scale_self, shift_self, gate_self, scale_ff, shift_ff, gate_ff = self.to_scale_shift_gate(global_cond).unsqueeze(1).chunk(6, dim = -1)
+
+            # self-attention with adaLN
+            residual = x
+            x = self.pre_norm(x)
+            x = x * (1 + scale_self) + shift_self
+            x = self.self_attn(x, mask = mask, rotary_pos_emb = rotary_pos_emb)
+            x = x * torch.sigmoid(1 - gate_self)
+            x = x + residual
+
+            if context is not None:
+                x = x + self.cross_attn(self.cross_attend_norm(x), context = context, context_mask = context_mask)
+
+            if self.conformer is not None:
+                x = x + self.conformer(x)
+
+            # feedforward with adaLN
+            residual = x
+            x = self.ff_norm(x)
+            x = x * (1 + scale_ff) + shift_ff
+            x = self.ff(x)
+            x = x * torch.sigmoid(1 - gate_ff)
+            x = x + residual
+
+        else:
+            assert global_cond is None
+            x = x + self.self_attn(self.pre_norm(x), mask = mask, rotary_pos_emb = rotary_pos_emb)
+
+            if context is not None:
+                x = x + self.cross_attn(self.cross_attend_norm(x), context = context, context_mask = context_mask)
+
+            if self.conformer is not None:
+                x = x + self.conformer(x)
+
+            x = x + self.ff(self.ff_norm(x))
+
+        return x
+        
+class ContinuousTransformer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        *,
+        dim_in = None,
+        dim_out = None,
+        dim_heads = 64,
+        cross_attend=False,
+        cross_atten_layer_idx=None,
+        cond_token_dim=None,
+        global_cond_dim=None,
+        causal=False,
+        rotary_pos_emb=True,
+        zero_init_branch_outputs=True,
+        conformer=False,
+        use_sinusoidal_emb=False,
+        use_abs_pos_emb=False,
+        abs_pos_emb_max_length=10000,
+        pos_emb_2d_size=1,
+        rotary_base_val=10000,
+        init_std=0.02,
+        **kwargs
+        ):
+
+        super().__init__()
+
+        self.dim = dim
+        self.depth = depth
+        self.causal = causal
+        self.layers = nn.ModuleList([])
+
+        self.project_in = nn.Linear(dim_in, dim, bias=False) if dim_in is not None else nn.Identity()
+        self.project_out = nn.Linear(dim, dim_out, bias=False) if dim_out is not None else nn.Identity()
+
+        if rotary_pos_emb:
+            if pos_emb_2d_size == 1:
+                self.rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32), base=rotary_base_val)
+            else:
+                self.rotary_pos_emb = RotaryEmbedding2D(max(dim_heads // 2, 32), pos_emb_2d_size, base=rotary_base_val)
+        else:
+            self.rotary_pos_emb = None
+
+        self.use_sinusoidal_emb = use_sinusoidal_emb
+        if use_sinusoidal_emb:
+            if pos_emb_2d_size != 1:
+                raise NotImplementedError
+            self.pos_emb = ScaledSinusoidalEmbedding(dim)
+
+        self.use_abs_pos_emb = use_abs_pos_emb
+        if use_abs_pos_emb:
+            if pos_emb_2d_size != 1:
+                raise NotImplementedError
+            self.pos_emb = AbsolutePositionalEmbedding(dim, abs_pos_emb_max_length)
+
+        
+        if cross_atten_layer_idx is None:
+            cross_atten_layer_idx = list(range(depth))
+        for i in range(depth):
+            self.layers.append(
+                TransformerBlock(
+                    dim,
+                    dim_heads = dim_heads,
+                    cross_attend = cross_attend and (i in cross_atten_layer_idx),
+                    dim_context = cond_token_dim,
+                    global_cond_dim = global_cond_dim,
+                    causal = causal,
+                    zero_init_branch_outputs = zero_init_branch_outputs,
+                    conformer=conformer,
+                    layer_ix=i,
+                    **kwargs
+                )
+            )
+        self.gradient_checkpointing = False
+        
+        self.apply(partial(self._init_weights,init_std=init_std))
+        
+    def forward(
+        self,
+        x,
+        mask = None,
+        prepend_embeds = None,
+        prepend_mask = None,
+        global_cond = None,
+        return_info = False,
+        **kwargs
+    ):
+        batch, seq, device = *x.shape[:2], x.device
+
+        info = {
+            "hidden_states": [],
+        }
+
+        x = self.project_in(x)
+
+        if prepend_embeds is not None:
+            prepend_length, prepend_dim = prepend_embeds.shape[1:]
+
+            assert prepend_dim == x.shape[-1], 'prepend dimension must match sequence dimension'
+
+            x = torch.cat((prepend_embeds, x), dim = -2)
+
+            if prepend_mask is not None or mask is not None:
+                mask = mask if mask is not None else torch.ones((batch, seq), device = device, dtype = torch.bool)
+                prepend_mask = prepend_mask if prepend_mask is not None else torch.ones((batch, prepend_length), device = device, dtype = torch.bool)
+
+                mask = torch.cat((prepend_mask, mask), dim = -1)
+
+        # Attention layers 
+
+        if self.rotary_pos_emb is not None:
+            rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1])
+        else:
+            rotary_pos_emb = None
+
+        if self.use_sinusoidal_emb or self.use_abs_pos_emb:
+            x = x + self.pos_emb(x)
+
+        # Iterate over the transformer layers
+        context, context_mask = kwargs.pop('context', None), kwargs.pop("context_mask", None)
+
+        for layer_idx, layer in enumerate(self.layers):
+            if layer.cross_attend:  
+                # x = layer(x, mask, global_cond=global_cond, rotary_pos_emb=rotary_pos_emb, context=context, context_mask=context_mask,**kwargs)
+                if self.gradient_checkpointing:
+                    x = checkpoint(layer, x, mask, global_cond, context, context_mask, rotary_pos_emb=rotary_pos_emb, **kwargs)
+                else:
+                    x = layer(x, mask, global_cond, context, context_mask, rotary_pos_emb=rotary_pos_emb, **kwargs)
+            else:
+                # x = layer(x, mask, global_cond=global_cond, rotary_pos_emb=rotary_pos_emb, **kwargs)
+                if self.gradient_checkpointing:
+                    x = checkpoint(layer, x, mask, global_cond, rotary_pos_emb=rotary_pos_emb, **kwargs)
+                else:
+                    x = layer(x, mask, global_cond, rotary_pos_emb=rotary_pos_emb, **kwargs)
+            if return_info:
+                info["hidden_states"].append(x)
+
+        x = self.project_out(x)
+
+        if return_info:
+            return x, info
+        
+        return x
+
+    def gradient_checkpointing_enable(self):
+        self.gradient_checkpointing = True
+        
+        
+    def _init_weights(self, module, init_std=0.02):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=init_std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=init_std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()

SongBloom/models/vae_frontend/__init__.py

@@ -0,0 +1,96 @@
+import torch
+from torch import nn
+import typing as tp
+import torchaudio
+import einops
+from abc import ABC, abstractmethod
+
+
+class AbstractVAE(ABC, nn.Module):
+
+    @property
+    @abstractmethod
+    def frame_rate(self) -> float:
+        ...
+
+    @property
+    @abstractmethod
+    def orig_sample_rate(self) -> int:
+        ...
+    
+
+    @property
+    @abstractmethod
+    def channel_dim(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def split_bands(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def input_channel(self) -> int:
+        ...
+
+
+    def encode(self, wav) -> torch.Tensor:
+        ...
+        
+    def decode(self, latents) -> torch.Tensor:
+        ...
+
+
+from .autoencoders import create_autoencoder_from_config, AudioAutoencoder
+class StableVAE(AbstractVAE):
+    def __init__(self, vae_ckpt, vae_cfg, sr=48000) -> None:
+        super().__init__()
+        import json
+        with open(vae_cfg) as f:
+            config = json.load(f)
+        self.vae: AudioAutoencoder = create_autoencoder_from_config(config)
+        self.vae.load_state_dict(torch.load(vae_ckpt)['state_dict'])
+        self.sample_rate = sr
+        self.rsp48k = torchaudio.transforms.Resample(sr, self.orig_sample_rate) if sr != self.orig_sample_rate else nn.Identity()
+       
+    @torch.no_grad()
+    def encode(self, wav: torch.Tensor, sample=True) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        wav = self.rsp48k(wav)
+        if wav.shape[-1] < 2048:
+            return torch.zeros((wav.shape[0], self.channel_dim, 0), device=wav.device, dtype=wav.dtype)
+        if wav.ndim == 2:
+            wav = wav.unsqueeze(1)
+        if wav.shape[1] == 1:
+            wav = wav.repeat(1, self.vae.in_channels, 1)
+        latent = self.vae.encode_audio(wav) # B, 64, T
+        return latent
+            
+
+        
+    def decode(self, latents: torch.Tensor, **kwargs):
+        # B, 64, T
+        with torch.no_grad():
+            audio_recon = self.vae.decode_audio(latents, **kwargs)
+            
+        return audio_recon
+        
+    @property
+    def frame_rate(self) -> float:
+        return float(self.vae.sample_rate) / self.vae.downsampling_ratio
+
+    @property
+    def orig_sample_rate(self) -> int:
+        return self.vae.sample_rate
+
+    @property
+    def channel_dim(self) -> int:
+        return self.vae.latent_dim
+
+    @property
+    def split_bands(self) -> int:
+        return 1
+    
+    @property
+    def input_channel(self) -> int:
+        return self.vae.in_channels

SongBloom/models/vae_frontend/autoencoders.py

@@ -0,0 +1,657 @@
+# https://github.com/Stability-AI/stable-audio-tools/tree/main/stable_audio_tools/models
+
+import torch
+import math
+import numpy as np
+
+from torch import nn
+from torch.nn import functional as F
+from torchaudio import transforms as T
+from dac.nn.layers import WNConv1d, WNConvTranspose1d
+from typing import Literal, Dict, Any
+import os,sys
+
+sys.path.insert(0, os.path.dirname(__file__))
+from bottleneck import create_bottleneck_from_config
+
+class Bottleneck(nn.Module):
+    def __init__(self, is_discrete: bool = False):
+        super().__init__()
+
+        self.is_discrete = is_discrete
+
+    def encode(self, x, return_info=False, **kwargs):
+        raise NotImplementedError
+
+    def decode(self, x):
+        raise NotImplementedError
+
+class DiscreteBottleneck(Bottleneck):
+    def __init__(self, num_quantizers, codebook_size, tokens_id):
+        super().__init__(is_discrete=True)
+
+        self.num_quantizers = num_quantizers
+        self.codebook_size = codebook_size
+        self.tokens_id = tokens_id
+
+    def decode_tokens(self, codes, **kwargs):
+        raise NotImplementedError
+    
+
+def checkpoint(function, *args, **kwargs):
+    kwargs.setdefault("use_reentrant", False)
+    return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
+
+
+
+# Adapted from https://github.com/NVIDIA/BigVGAN/blob/main/activations.py under MIT license
+# License available in LICENSES/LICENSE_NVIDIA.txt
+def snake_beta(x, alpha, beta):
+    return x + (1.0 / (beta + 1e-9)) * pow(torch.sin(x * alpha), 2)
+
+class SnakeBeta(nn.Module):
+
+    def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=True):
+        super(SnakeBeta, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale: # log scale alphas initialized to zeros
+            self.alpha = nn.Parameter(torch.zeros(in_features) * alpha)
+            self.beta = nn.Parameter(torch.zeros(in_features) * alpha)
+        else: # linear scale alphas initialized to ones
+            self.alpha = nn.Parameter(torch.ones(in_features) * alpha)
+            self.beta = nn.Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 1e-9
+
+    def forward(self, x):
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = snake_beta(x, alpha, beta)
+
+        return x
+    
+def get_activation(activation: Literal["elu", "snake", "none"], antialias=False, channels=None) -> nn.Module:
+    if activation == "elu":
+        act = nn.ELU()
+    elif activation == "snake":
+        act = SnakeBeta(channels)
+    elif activation == "none":
+        act = nn.Identity()
+    else:
+        raise ValueError(f"Unknown activation {activation}")
+    
+    if antialias:
+        from alias_free_torch import Activation1d
+        act = Activation1d(act)
+    
+    return act
+
+class ResidualUnit(nn.Module):
+    def __init__(self, in_channels, out_channels, dilation, use_snake=False, antialias_activation=False):
+        super().__init__()
+        
+        self.dilation = dilation
+
+        padding = (dilation * (7-1)) // 2
+
+        self.layers = nn.Sequential(
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels),
+            WNConv1d(in_channels=in_channels, out_channels=out_channels,
+                      kernel_size=7, dilation=dilation, padding=padding),
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels),
+            WNConv1d(in_channels=out_channels, out_channels=out_channels,
+                      kernel_size=1)
+        )
+
+    def forward(self, x):
+        res = x
+        
+        #x = checkpoint(self.layers, x)
+        x = self.layers(x)
+
+        return x + res
+
+class EncoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False):
+        super().__init__()
+
+        self.layers = nn.Sequential(
+            ResidualUnit(in_channels=in_channels,
+                         out_channels=in_channels, dilation=1, use_snake=use_snake),
+            ResidualUnit(in_channels=in_channels,
+                         out_channels=in_channels, dilation=3, use_snake=use_snake),
+            ResidualUnit(in_channels=in_channels,
+                         out_channels=in_channels, dilation=9, use_snake=use_snake),
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels),
+            WNConv1d(in_channels=in_channels, out_channels=out_channels,
+                      kernel_size=2*stride, stride=stride, padding=math.ceil(stride/2)),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+class DecoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False, use_nearest_upsample=False):
+        super().__init__()
+
+        if use_nearest_upsample:
+            upsample_layer = nn.Sequential(
+                nn.Upsample(scale_factor=stride, mode="nearest"),
+                WNConv1d(in_channels=in_channels,
+                        out_channels=out_channels, 
+                        kernel_size=2*stride,
+                        stride=1,
+                        bias=False,
+                        padding='same')
+            )
+        else:
+            upsample_layer = WNConvTranspose1d(in_channels=in_channels,
+                               out_channels=out_channels,
+                               kernel_size=2*stride, stride=stride, padding=math.ceil(stride/2))
+
+        self.layers = nn.Sequential(
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels),
+            upsample_layer,
+            ResidualUnit(in_channels=out_channels, out_channels=out_channels,
+                         dilation=1, use_snake=use_snake),
+            ResidualUnit(in_channels=out_channels, out_channels=out_channels,
+                         dilation=3, use_snake=use_snake),
+            ResidualUnit(in_channels=out_channels, out_channels=out_channels,
+                         dilation=9, use_snake=use_snake),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+class OobleckEncoder(nn.Module):
+    def __init__(self, 
+                 in_channels=2, 
+                 channels=128, 
+                 latent_dim=32, 
+                 c_mults = [1, 2, 4, 8], 
+                 strides = [2, 4, 8, 8],
+                 use_snake=False,
+                 antialias_activation=False
+        ):
+        super().__init__()
+          
+        c_mults = [1] + c_mults
+
+        self.depth = len(c_mults)
+
+        layers = [
+            WNConv1d(in_channels=in_channels, out_channels=c_mults[0] * channels, kernel_size=7, padding=3)
+        ]
+        
+        for i in range(self.depth-1):
+            layers += [EncoderBlock(in_channels=c_mults[i]*channels, out_channels=c_mults[i+1]*channels, stride=strides[i], use_snake=use_snake)]
+
+        layers += [
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[-1] * channels),
+            WNConv1d(in_channels=c_mults[-1]*channels, out_channels=latent_dim, kernel_size=3, padding=1)
+        ]
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class OobleckDecoder(nn.Module):
+    def __init__(self, 
+                 out_channels=2, 
+                 channels=128, 
+                 latent_dim=32, 
+                 c_mults = [1, 2, 4, 8], 
+                 strides = [2, 4, 8, 8],
+                 use_snake=False,
+                 antialias_activation=False,
+                 use_nearest_upsample=False,
+                 final_tanh=True):
+        super().__init__()
+
+        c_mults = [1] + c_mults
+        
+        self.depth = len(c_mults)
+
+        layers = [
+            WNConv1d(in_channels=latent_dim, out_channels=c_mults[-1]*channels, kernel_size=7, padding=3),
+        ]
+        
+        for i in range(self.depth-1, 0, -1):
+            layers += [DecoderBlock(
+                in_channels=c_mults[i]*channels, 
+                out_channels=c_mults[i-1]*channels, 
+                stride=strides[i-1], 
+                use_snake=use_snake, 
+                antialias_activation=antialias_activation,
+                use_nearest_upsample=use_nearest_upsample
+                )
+            ]
+
+        layers += [
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[0] * channels),
+            WNConv1d(in_channels=c_mults[0] * channels, out_channels=out_channels, kernel_size=7, padding=3, bias=False),
+            nn.Tanh() if final_tanh else nn.Identity()
+        ]
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+
+class AudioAutoencoder(nn.Module):
+    def __init__(
+        self,
+        encoder,
+        decoder,
+        latent_dim,
+        downsampling_ratio,
+        sample_rate,
+        io_channels=2,
+        bottleneck = None,
+        pretransform = None,
+        in_channels = None,
+        out_channels = None,
+        soft_clip = False
+    ):
+        super().__init__()
+
+        self.downsampling_ratio = downsampling_ratio
+        self.sample_rate = sample_rate
+
+        self.latent_dim = latent_dim
+        self.io_channels = io_channels
+        self.in_channels = io_channels
+        self.out_channels = io_channels
+
+        self.min_length = self.downsampling_ratio
+
+        if in_channels is not None:
+            self.in_channels = in_channels
+
+        if out_channels is not None:
+            self.out_channels = out_channels
+
+        self.bottleneck = bottleneck
+
+        self.encoder = encoder
+
+        self.decoder = decoder
+
+        self.pretransform = pretransform
+
+        self.soft_clip = soft_clip
+ 
+        self.is_discrete = self.bottleneck is not None and self.bottleneck.is_discrete
+
+    def encode(self, audio, return_info=False, skip_pretransform=False, iterate_batch=False, **kwargs):
+
+        info = {}
+
+        if self.pretransform is not None and not skip_pretransform:
+            if self.pretransform.enable_grad:
+                if iterate_batch:
+                    audios = []
+                    for i in range(audio.shape[0]):
+                        audios.append(self.pretransform.encode(audio[i:i+1]))
+                    audio = torch.cat(audios, dim=0)
+                else:
+                    audio = self.pretransform.encode(audio)
+            else:
+                with torch.no_grad():
+                    if iterate_batch:
+                        audios = []
+                        for i in range(audio.shape[0]):
+                            audios.append(self.pretransform.encode(audio[i:i+1]))
+                        audio = torch.cat(audios, dim=0)
+                    else:
+                        audio = self.pretransform.encode(audio)
+
+        if self.encoder is not None:
+            if iterate_batch:
+                latents = []
+                for i in range(audio.shape[0]):
+                    latents.append(self.encoder(audio[i:i+1]))
+                latents = torch.cat(latents, dim=0)
+            else:
+                latents = self.encoder(audio)
+        else:
+            latents = audio
+
+        if self.bottleneck is not None:
+            # TODO: Add iterate batch logic, needs to merge the info dicts
+            latents, bottleneck_info = self.bottleneck.encode(latents, return_info=True, **kwargs)
+
+            info.update(bottleneck_info)
+        
+        if return_info:
+            return latents, info
+
+        return latents
+
+    def decode(self, latents, iterate_batch=False, **kwargs):
+
+        if self.bottleneck is not None:
+            if iterate_batch:
+                decoded = []
+                for i in range(latents.shape[0]):
+                    decoded.append(self.bottleneck.decode(latents[i:i+1]))
+                latents = torch.cat(decoded, dim=0)
+            else:
+                latents = self.bottleneck.decode(latents)
+
+        if iterate_batch:
+            decoded = []
+            for i in range(latents.shape[0]):
+                decoded.append(self.decoder(latents[i:i+1]))
+            decoded = torch.cat(decoded, dim=0)
+        else:
+            decoded = self.decoder(latents, **kwargs)
+
+        if self.pretransform is not None:
+            if self.pretransform.enable_grad:
+                if iterate_batch:
+                    decodeds = []
+                    for i in range(decoded.shape[0]):
+                        decodeds.append(self.pretransform.decode(decoded[i:i+1]))
+                    decoded = torch.cat(decodeds, dim=0)
+                else:
+                    decoded = self.pretransform.decode(decoded)
+            else:
+                with torch.no_grad():
+                    if iterate_batch:
+                        decodeds = []
+                        for i in range(latents.shape[0]):
+                            decodeds.append(self.pretransform.decode(decoded[i:i+1]))
+                        decoded = torch.cat(decodeds, dim=0)
+                    else:
+                        decoded = self.pretransform.decode(decoded)
+
+        if self.soft_clip:
+            decoded = torch.tanh(decoded)
+        
+        return decoded
+          
+    def decode_tokens(self, tokens, **kwargs):
+        '''
+        Decode discrete tokens to audio
+        Only works with discrete autoencoders
+        '''
+
+        assert isinstance(self.bottleneck, DiscreteBottleneck), "decode_tokens only works with discrete autoencoders"
+
+        latents = self.bottleneck.decode_tokens(tokens, **kwargs)
+
+        return self.decode(latents, **kwargs)
+        
+    
+
+    def encode_audio(self, audio, chunked=False, overlap=32, chunk_size=128, **kwargs):
+        '''
+        Encode audios into latents. Audios should already be preprocesed by preprocess_audio_for_encoder.
+        If chunked is True, split the audio into chunks of a given maximum size chunk_size, with given overlap.
+        Overlap and chunk_size params are both measured in number of latents (not audio samples) 
+        # and therefore you likely could use the same values with decode_audio. 
+        A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size. 
+        Every autoencoder will have a different receptive field size, and thus ideal overlap.
+        You can determine it empirically by diffing unchunked vs chunked output and looking at maximum diff.
+        The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks.
+        Smaller chunk_size uses less memory, but more compute.
+        The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version
+        For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks
+        '''
+        if not chunked:
+            # default behavior. Encode the entire audio in parallel
+            return self.encode(audio, **kwargs)
+        else:
+            # CHUNKED ENCODING
+            # samples_per_latent is just the downsampling ratio (which is also the upsampling ratio)
+            samples_per_latent = self.downsampling_ratio
+            total_size = audio.shape[2] # in samples
+            batch_size = audio.shape[0]
+            chunk_size *= samples_per_latent # converting metric in latents to samples
+            overlap *= samples_per_latent # converting metric in latents to samples
+            hop_size = chunk_size - overlap
+            chunks = []
+            for i in range(0, total_size - chunk_size + 1, hop_size):
+                chunk = audio[:,:,i:i+chunk_size]
+                chunks.append(chunk)
+            if i+chunk_size != total_size:
+                # Final chunk
+                chunk = audio[:,:,-chunk_size:]
+                chunks.append(chunk)
+            chunks = torch.stack(chunks)
+            num_chunks = chunks.shape[0]
+            # Note: y_size might be a different value from the latent length used in diffusion training
+            # because we can encode audio of varying lengths
+            # However, the audio should've been padded to a multiple of samples_per_latent by now.
+            y_size = total_size // samples_per_latent
+            # Create an empty latent, we will populate it with chunks as we encode them
+            y_final = torch.zeros((batch_size,self.latent_dim,y_size)).to(audio.device)
+            for i in range(num_chunks):
+                x_chunk = chunks[i,:]
+                # encode the chunk
+                y_chunk = self.encode(x_chunk)
+                # figure out where to put the audio along the time domain
+                if i == num_chunks-1:
+                    # final chunk always goes at the end
+                    t_end = y_size
+                    t_start = t_end - y_chunk.shape[2]
+                else:
+                    t_start = i * hop_size // samples_per_latent
+                    t_end = t_start + chunk_size // samples_per_latent
+                #  remove the edges of the overlaps
+                ol = overlap//samples_per_latent//2
+                chunk_start = 0
+                chunk_end = y_chunk.shape[2]
+                if i > 0:
+                    # no overlap for the start of the first chunk
+                    t_start += ol
+                    chunk_start += ol
+                if i < num_chunks-1:
+                    # no overlap for the end of the last chunk
+                    t_end -= ol
+                    chunk_end -= ol
+                # paste the chunked audio into our y_final output audio
+                y_final[:,:,t_start:t_end] = y_chunk[:,:,chunk_start:chunk_end]
+            return y_final
+    
+    def decode_audio(self, latents, chunked=False, overlap=32, chunk_size=128, **kwargs):
+        '''
+        Decode latents to audio. 
+        If chunked is True, split the latents into chunks of a given maximum size chunk_size, with given overlap, both of which are measured in number of latents. 
+        A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size. 
+        Every autoencoder will have a different receptive field size, and thus ideal overlap.
+        You can determine it empirically by diffing unchunked vs chunked audio and looking at maximum diff.
+        The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks.
+        Smaller chunk_size uses less memory, but more compute.
+        The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version
+        For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks
+        '''
+        if not chunked:
+            # default behavior. Decode the entire latent in parallel
+            return self.decode(latents, **kwargs)
+        else:
+            # chunked decoding
+            hop_size = chunk_size - overlap
+            total_size = latents.shape[2]
+            batch_size = latents.shape[0]
+            chunks = []
+            for i in range(0, total_size - chunk_size + 1, hop_size):
+                chunk = latents[:,:,i:i+chunk_size]
+                chunks.append(chunk)
+            if i+chunk_size != total_size:
+                # Final chunk
+                chunk = latents[:,:,-chunk_size:]
+                chunks.append(chunk)
+            chunks = torch.stack(chunks)
+            num_chunks = chunks.shape[0]
+            # samples_per_latent is just the downsampling ratio
+            samples_per_latent = self.downsampling_ratio
+            # Create an empty waveform, we will populate it with chunks as decode them
+            y_size = total_size * samples_per_latent
+            y_final = torch.zeros((batch_size,self.out_channels,y_size)).to(latents.device)
+            for i in range(num_chunks):
+                x_chunk = chunks[i,:]
+                # decode the chunk
+                y_chunk = self.decode(x_chunk)
+                # figure out where to put the audio along the time domain
+                if i == num_chunks-1:
+                    # final chunk always goes at the end
+                    t_end = y_size
+                    t_start = t_end - y_chunk.shape[2]
+                else:
+                    t_start = i * hop_size * samples_per_latent
+                    t_end = t_start + chunk_size * samples_per_latent
+                #  remove the edges of the overlaps
+                ol = (overlap//2) * samples_per_latent
+                chunk_start = 0
+                chunk_end = y_chunk.shape[2]
+                if i > 0:
+                    # no overlap for the start of the first chunk
+                    t_start += ol
+                    chunk_start += ol
+                if i < num_chunks-1:
+                    # no overlap for the end of the last chunk
+                    t_end -= ol
+                    chunk_end -= ol
+                # paste the chunked audio into our y_final output audio
+                y_final[:,:,t_start:t_end] = y_chunk[:,:,chunk_start:chunk_end]
+            return y_final
+
+        
+# AE factories
+
+def create_encoder_from_config(encoder_config: Dict[str, Any]):
+    encoder_type = encoder_config.get("type", None)
+    assert encoder_type is not None, "Encoder type must be specified"
+
+    if encoder_type == "oobleck":
+        encoder = OobleckEncoder(
+            **encoder_config["config"]
+        )
+    
+    elif encoder_type == "seanet":
+        from encodec.modules import SEANetEncoder
+        seanet_encoder_config = encoder_config["config"]
+
+        #SEANet encoder expects strides in reverse order
+        seanet_encoder_config["ratios"] = list(reversed(seanet_encoder_config.get("ratios", [2, 2, 2, 2, 2])))
+        encoder = SEANetEncoder(
+            **seanet_encoder_config
+        )
+    else:
+        raise ValueError(f"Unknown encoder type {encoder_type}")
+    
+    requires_grad = encoder_config.get("requires_grad", True)
+    if not requires_grad:
+        for param in encoder.parameters():
+            param.requires_grad = False
+
+    return encoder
+
+def create_decoder_from_config(decoder_config: Dict[str, Any]):
+    decoder_type = decoder_config.get("type", None)
+    assert decoder_type is not None, "Decoder type must be specified"
+
+    if decoder_type == "oobleck":
+        decoder = OobleckDecoder(
+            **decoder_config["config"]
+        )
+    elif decoder_type == "seanet":
+        from encodec.modules import SEANetDecoder
+
+        decoder = SEANetDecoder(
+            **decoder_config["config"]
+        )
+    else:
+        raise ValueError(f"Unknown decoder type {decoder_type}")
+    
+    requires_grad = decoder_config.get("requires_grad", True)
+    if not requires_grad:
+        for param in decoder.parameters():
+            param.requires_grad = False
+
+    return decoder
+
+def create_autoencoder_from_config(config: Dict[str, Any]):
+    
+    # print(config)
+    ae_config = config["model"]
+
+    encoder = create_encoder_from_config(ae_config["encoder"])
+    decoder = create_decoder_from_config(ae_config["decoder"])
+
+    bottleneck = ae_config.get("bottleneck", None)
+
+    latent_dim = ae_config.get("latent_dim", None)
+    assert latent_dim is not None, "latent_dim must be specified in model config"
+    downsampling_ratio = ae_config.get("downsampling_ratio", None)
+    assert downsampling_ratio is not None, "downsampling_ratio must be specified in model config"
+    io_channels = ae_config.get("io_channels", None)
+    assert io_channels is not None, "io_channels must be specified in model config"
+    sample_rate = config.get("sample_rate", None)
+    assert sample_rate is not None, "sample_rate must be specified in model config"
+
+    in_channels = ae_config.get("in_channels", None)
+    out_channels = ae_config.get("out_channels", None)
+
+    pretransform = ae_config.get("pretransform", None)
+
+    if pretransform is not None:
+        from stable_audio_tools.models.factory import create_pretransform_from_config
+        pretransform = create_pretransform_from_config(pretransform, sample_rate)
+
+    if bottleneck is not None:
+        bottleneck = create_bottleneck_from_config(bottleneck)
+
+    soft_clip = ae_config["decoder"].get("soft_clip", False)
+
+    return AudioAutoencoder(
+        encoder,
+        decoder,
+        io_channels=io_channels,
+        latent_dim=latent_dim,
+        downsampling_ratio=downsampling_ratio,
+        sample_rate=sample_rate,
+        bottleneck=bottleneck,
+        pretransform=pretransform,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        soft_clip=soft_clip
+    )
+
+
+
+if __name__ == "__main__":
+    import json
+    import torchaudio
+    config_path = 'modelzoo/stable_audio_vae/stable_audio_2_0_vae.json'
+    with open(config_path) as f:
+        config = json.load(f)
+    with torch.no_grad():
+        vae_model = create_autoencoder_from_config(config).cuda()
+        model_ckpt_path = 'modelzoo/stable_audio_vae/autoencoder.ckpt'
+        vae_model.load_state_dict(torch.load(model_ckpt_path)['state_dict'])
+        
+        
+        input_audios, sr = torchaudio.load("music_example/加勒比海盗 主题.wav")
+        input_audios = torchaudio.functional.resample(input_audios, sr, 48000)[...,:2048]
+        input_audios = input_audios.unsqueeze(1).repeat(1, 2, 1).cuda()
+        latents = vae_model.encode_audio(input_audios)
+        recover_audio = vae_model.decode_audio(latents)
+        print(recover_audio)
+    
+    breakpoint()