work on pip package

README.md

@@ -1,15 +1,15 @@
-<h1 align="center">ACE-Step</h1>
+# ACE-Step
-<h1 align="center">A Step Towards Music Generation Foundation Model</h1>
+
-<p align="center">
+A Step Towards Music Generation Foundation Model
-    <a href="https://ace-step.github.io/">Project</a> |
+
-    <a href="https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B">Checkpoints</a> |
+[Project Page](https://ace-step.github.io/) |
-    <a href="https://huggingface.co/spaces/ACE-Step/ACE-Step">Space Demo</a> |
+[Checkpoints](https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B) |
-     <a href="https://discord.gg/rjAZz2xBdG">Discord</a> 
+[Space Demo](https://huggingface.co/spaces/ACE-Step/ACE-Step) |
-</p>
+[Discord](https://discord.gg/rjAZz2xBdG) 
 
 ---
 <p align="center">
-    <img src="./fig/orgnization_logos.png" width="100%" alt="StepFun Logo">
+    <img src="./assets/orgnization_logos.png" width="100%" alt="StepFun Logo">
 </p>
 
 ## Table of Contents
@@ -23,7 +23,8 @@
 
 - 🚀 2025.05.06: Open source demo code and model
 
-## TODOs📋
+## 📋 Roadmap
+
 - [x] Release training code 🔥
 - [x] Release LoRA training code 🔥
 - [ ] Release RapMachine lora 🎤
@@ -34,7 +35,7 @@
 ## 🏗️ Architecture
 
 <p align="center">
-    <img src="./fig/ACE-Step_framework.png" width="100%" alt="ACE-Step Framework">
+    <img src="./assets/ACE-Step_framework.png" width="100%" alt="ACE-Step Framework">
 </p>
 
 
@@ -49,7 +50,7 @@ Rather than building yet another end-to-end text-to-music pipeline, our vision i
 ## ✨ Features
 
 <p align="center">
-    <img src="./fig/application_map.png" width="100%" alt="ACE-Step Framework">
+    <img src="./assets/application_map.png" width="100%" alt="ACE-Step Framework">
 </p>
 
 ### 🎯 Baseline Quality
@@ -201,7 +202,7 @@ It is highly recommended to use a virtual environment to manage project dependen
 
 ## 🚀 Usage
 
-![Demo Interface](fig/demo_interface.png)
+![Demo Interface](assets/demo_interface.png)
 
 ### 🔍 Basic Usage
 

acestep/__init__.py

@@ -0,0 +1,7 @@
+"""
+ACE-Step: A Step Towards Music Generation Foundation Model
+
+https://github.com/ace-step/ACE-Step
+
+Apache 2.0 License
+"""

acestep/apg_guidance.py

@@ -64,7 +64,11 @@ def cfg_double_condition_forward(
     guidance_scale_text,
     guidance_scale_lyric,
 ):
-    return (1 - guidance_scale_text) * uncond_output + (guidance_scale_text - guidance_scale_lyric) * only_text_cond_output + guidance_scale_lyric * cond_output 
+    return (
+        (1 - guidance_scale_text) * uncond_output
+        + (guidance_scale_text - guidance_scale_lyric) * only_text_cond_output
+        + guidance_scale_lyric * cond_output
+    )
 
 
 def optimized_scale(positive_flat, negative_flat):
@@ -81,14 +85,23 @@ def optimized_scale(positive_flat, negative_flat):
     return st_star
 
 
-def cfg_zero_star(noise_pred_with_cond, noise_pred_uncond, guidance_scale, i, zero_steps=1, use_zero_init=True):
+def cfg_zero_star(
+    noise_pred_with_cond,
+    noise_pred_uncond,
+    guidance_scale,
+    i,
+    zero_steps=1,
+    use_zero_init=True,
+):
     bsz = noise_pred_with_cond.shape[0]
     positive_flat = noise_pred_with_cond.view(bsz, -1)
     negative_flat = noise_pred_uncond.view(bsz, -1)
     alpha = optimized_scale(positive_flat, negative_flat)
     alpha = alpha.view(bsz, 1, 1, 1)
     if (i <= zero_steps) and use_zero_init:
-        noise_pred = noise_pred_with_cond * 0.
+        noise_pred = noise_pred_with_cond * 0.0
     else:
-        noise_pred = noise_pred_uncond * alpha + guidance_scale * (noise_pred_with_cond - noise_pred_uncond * alpha)
+        noise_pred = noise_pred_uncond * alpha + guidance_scale * (
+            noise_pred_with_cond - noise_pred_uncond * alpha
+        )
     return noise_pred

acestep/gui.py

@@ -1,30 +1,56 @@
-import argparse
+"""
-parser = argparse.ArgumentParser()
+ACE-Step: A Step Towards Music Generation Foundation Model
-parser.add_argument("--checkpoint_path", type=str, default="")
+
-parser.add_argument("--server_name", type=str, default="127.0.0.1")
+https://github.com/ace-step/ACE-Step
-parser.add_argument("--port", type=int, default=7865)
+
-parser.add_argument("--device_id", type=int, default=0)
+Apache 2.0 License
-parser.add_argument("--share", type=bool, default=False)
+"""
-parser.add_argument("--bf16", type=bool, default=True)
-parser.add_argument("--torch_compile", type=bool, default=False)
-args = parser.parse_args()
 
 import os
+import click
 
-os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device_id)
+from acestep.ui.components import create_main_demo_ui
+from acestep.pipeline_ace_step import ACEStepPipeline
+from acestep.data_sampler import DataSampler
 
 
-from ui.components import create_main_demo_ui
+@click.command()
-from pipeline_ace_step import ACEStepPipeline
+@click.option(
-from data_sampler import DataSampler
+    "--checkpoint_path", type=str, default="", help="Path to the checkpoint directory. Downloads automatically if empty."
+)
+@click.option(
+    "--server_name",
+    type=str,
+    default="127.0.0.1",
+    help="The server name to use for the Gradio app.",
+)
+@click.option(
+    "--port", type=int, default=7865, help="The port to use for the Gradio app."
+)
+@click.option("--device_id", type=int, default=0, help="The CUDA device ID to use.")
+@click.option(
+    "--share",
+    is_flag=True,
+    default=False,
+    help="Whether to create a public, shareable link for the Gradio app.",
+)
+@click.option(
+    "--bf16", is_flag=True, default=True, help="Whether to use bfloat16 precision. Turn off if using MPS."
+)
+@click.option(
+    "--torch_compile", is_flag=True, default=False, help="Whether to use torch.compile."
+)
+def main(checkpoint_path, server_name, port, device_id, share, bf16, torch_compile):
+    """
+    Main function to launch the ACE Step pipeline demo.
+    """
 
-
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
-def main(args):
 
     model_demo = ACEStepPipeline(
-        checkpoint_dir=args.checkpoint_path,
+        checkpoint_dir=checkpoint_path,
-        dtype="bfloat16" if args.bf16 else "float32",
+        dtype="bfloat16" if bf16 else "float32",
-        torch_compile=args.torch_compile
+        torch_compile=torch_compile,
     )
     data_sampler = DataSampler()
 
@@ -32,12 +58,8 @@ def main(args):
         text2music_process_func=model_demo.__call__,
         sample_data_func=data_sampler.sample,
     )
-    demo.launch(
+    demo.launch(server_name=server_name, server_port=port, share=share)
-        server_name=args.server_name,
-        server_port=args.port,
-        share=args.share
-    )
 
 
 if __name__ == "__main__":
-    main(args)
+    main()

acestep/language_segmentation/LangSegment.py

@@ -11,6 +11,41 @@ https://github.com/adbar/py3langid
 
 Projects:
 https://github.com/juntaosun/LangSegment
+
+LICENSE:
+py3langid - Language Identifier
+BSD 3-Clause License
+
+Modifications (fork): Copyright (c) 2021, Adrien Barbaresi.
+
+Original code: Copyright (c) 2011 Marco Lui <saffsd@gmail.com>.
+Based on research by Marco Lui and Tim Baldwin.
+
+All rights reserved.
+
+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.
+
+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.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER ``AS IS'' AND ANY EXPRESS 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 THE COPYRIGHT HOLDER OR
+CONTRIBUTORS 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.
 """
 
 import os
@@ -29,12 +64,7 @@ from collections import defaultdict
 # For probability normalization in library use, the user must instantiate their own . An example of such usage is as follows:
 from py3langid.langid import LanguageIdentifier, MODEL_FILE
 
-# Digital processing
+from acestep.language_segmentation.utils.num import num2str
-try:from .utils.num import num2str
-except ImportError:
-    try:from utils.num import num2str
-    except ImportError as e:
-        raise e
 
 # -----------------------------------
 # 更新日志：新版本分词更加精准。
@@ -107,16 +137,16 @@ class LangSSML:
     def __init__(self):
         # 纯数字
         self._zh_numerals_number = {
-                '0': '零',
+            "0": "零",
-                '1': '一',
+            "1": "一",
-                '2': '二',
+            "2": "二",
-                '3': '三',
+            "3": "三",
-                '4': '四',
+            "4": "四",
-                '5': '五',
+            "5": "五",
-                '6': '六',
+            "6": "六",
-                '7': '七',
+            "7": "七",
-                '8': '八',
+            "8": "八",
-                '9': '九'
+            "9": "九",
         }
 
     # 将2024/8/24, 2024-08, 08-24, 24 标准化“年月日”
@@ -124,36 +154,46 @@ class LangSSML:
     def _format_chinese_data(self, date_str: str):
         # 处理日期格式
         input_date = date_str
-        if date_str is None or date_str.strip() == "":return ""
+        if date_str is None or date_str.strip() == "":
+            return ""
         date_str = re.sub(r"[\/\._|年|月]", "-", date_str)
         date_str = re.sub(r"日", r"", date_str)
-        date_arrs = date_str.split(' ')
+        date_arrs = date_str.split(" ")
         if len(date_arrs) == 1 and ":" in date_arrs[0]:
             time_str = date_arrs[0]
             date_arrs = []
         else:
             time_str = date_arrs[1] if len(date_arrs) >= 2 else ""
+
         def nonZero(num, cn, func=None):
-            if func is not None:num=func(num)
+            if func is not None:
+                num = func(num)
             return f"{num}{cn}" if num is not None and num != "" and num != "0" else ""
+
         f_number = self.to_chinese_number
         f_currency = self.to_chinese_currency
         # year, month, day
         year_month_day = ""
         if len(date_arrs) > 0:
             year, month, day = "", "", ""
-            parts = date_arrs[0].split('-')
+            parts = date_arrs[0].split("-")
             if len(parts) == 3:  # 格式为 YYYY-MM-DD
                 year, month, day = parts
             elif len(parts) == 2:  # 格式为 MM-DD 或 YYYY-MM
                 if len(parts[0]) == 4:  # 年-月
                     year, month = parts
-                else:month, day = parts # 月-日
+                else:
+                    month, day = parts  # 月-日
             elif len(parts[0]) > 0:  # 仅有月-日或年
                 if len(parts[0]) == 4:
                     year = parts[0]
-                else:day = parts[0]
+                else:
-            year,month,day = nonZero(year,"年",f_number),nonZero(month,"月",f_currency),nonZero(day,"日",f_currency)
+                    day = parts[0]
+            year, month, day = (
+                nonZero(year, "年", f_number),
+                nonZero(month, "月", f_currency),
+                nonZero(day, "日", f_currency),
+            )
             year_month_day = re.sub(r"([年|月|日])+", r"\1", f"{year}{month}{day}")
         # hours, minutes, seconds
         time_str = re.sub(r"[\/\.\-：_]", ":", time_str)
@@ -163,24 +203,35 @@ class LangSSML:
             hours, minutes, seconds = time_arrs
         elif len(time_arrs) == 2:  # H/M
             hours, minutes = time_arrs
-        elif len(time_arrs[0]) > 0:hours = f'{time_arrs[0]}点'  # H
+        elif len(time_arrs[0]) > 0:
+            hours = f"{time_arrs[0]}点"  # H
         if len(time_arrs) > 1:
-            hours, minutes, seconds = nonZero(hours,"点",f_currency),nonZero(minutes,"分",f_currency),nonZero(seconds,"秒",f_currency)
+            hours, minutes, seconds = (
-        hours_minutes_seconds = re.sub(r"([点|分|秒])+",r"\1",f"{hours}{minutes}{seconds}")
+                nonZero(hours, "点", f_currency),
+                nonZero(minutes, "分", f_currency),
+                nonZero(seconds, "秒", f_currency),
+            )
+        hours_minutes_seconds = re.sub(
+            r"([点|分|秒])+", r"\1", f"{hours}{minutes}{seconds}"
+        )
         output_date = f"{year_month_day}{hours_minutes_seconds}"
         return output_date
 
     # 【SSML】number=中文大写数字读法（单字）
     # Chinese Numbers(single word)
     def to_chinese_number(self, num: str):
-        pattern = r'(\d+)'
+        pattern = r"(\d+)"
         zh_numerals = self._zh_numerals_number
         arrs = re.split(pattern, num)
         output = ""
         for item in arrs:
             if re.match(pattern, item):
-                output += ''.join(zh_numerals[digit] if digit in zh_numerals else "" for digit in str(item))
+                output += "".join(
-            else:output += item
+                    zh_numerals[digit] if digit in zh_numerals else ""
+                    for digit in str(item)
+                )
+            else:
+                output += item
         output = output.replace(".", "点")
         return output
 
@@ -194,13 +245,14 @@ class LangSSML:
     # 【SSML】currency=按金额发音。
     # Digital processing from GPT_SoVITS num.py （thanks）
     def to_chinese_currency(self, num: str):
-        pattern = r'(\d+)'
+        pattern = r"(\d+)"
         arrs = re.split(pattern, num)
         output = ""
         for item in arrs:
             if re.match(pattern, item):
                 output += num2str(item)
-            else:output += item
+            else:
+                output += item
         output = output.replace(".", "点")
         return output
 
@@ -225,7 +277,7 @@ class LangSegment:
         # 可自定义语言匹配标签：カスタマイズ可能な言語対応タグ:사용자 지정 가능한 언어 일치 태그:
         # Customizable language matching tags: These are supported，이 표현들은 모두 지지합니다
         # <zh>你好<zh> , <ja>佐々木</ja> , <en>OK<en> , <ko>오빠</ko> 这些写法均支持
-        self.SYMBOLS_PATTERN = r'(<([a-zA-Z|-]*)>(.*?)<\/*[a-zA-Z|-]*>)'
+        self.SYMBOLS_PATTERN = r"(<([a-zA-Z|-]*)>(.*?)<\/*[a-zA-Z|-]*>)"
 
         # 语言过滤组功能, 可以指定保留语言。不在过滤组中的语言将被清除。您可随心搭配TTS语音合成所支持的语言。
         # 언어 필터 그룹 기능을 사용하면 예약된 언어를 지정할 수 있습니다. 필터 그룹에 없는 언어는 지워집니다. TTS 텍스트에서 지원하는 언어를 원하는 대로 일치시킬 수 있습니다.
@@ -287,7 +339,7 @@ class LangSegment:
         self.keepPinyin = False
 
         # DEFINITION
-        self.PARSE_TAG = re.compile(r'(⑥\$*\d+[\d]{6,}⑥)')
+        self.PARSE_TAG = re.compile(r"(⑥\$*\d+[\d]{6,}⑥)")
 
         self.LangSSML = LangSSML()
 
@@ -300,26 +352,26 @@ class LangSegment:
         self._lang_eos = None
 
     def _is_english_word(self, word):
-        return bool(re.match(r'^[a-zA-Z]+$', word))
+        return bool(re.match(r"^[a-zA-Z]+$", word))
 
     def _is_chinese(self, word):
         for char in word:
-            if '\u4e00' <= char <= '\u9fff':
+            if "\u4e00" <= char <= "\u9fff":
                 return True
         return False
 
     def _is_japanese_kana(self, word):
-        pattern = re.compile(r'[\u3040-\u309F\u30A0-\u30FF]+')
+        pattern = re.compile(r"[\u3040-\u309F\u30A0-\u30FF]+")
         matches = pattern.findall(word)
         return len(matches) > 0
 
     def _insert_english_uppercase(self, word):
-        modified_text = re.sub(r'(?<!\b)([A-Z])', r' \1', word)
+        modified_text = re.sub(r"(?<!\b)([A-Z])", r" \1", word)
-        modified_text = modified_text.strip('-')
+        modified_text = modified_text.strip("-")
         return modified_text + " "
 
     def _split_camel_case(self, word):
-        return re.sub(r'(?<!^)(?=[A-Z])', ' ', word)
+        return re.sub(r"(?<!^)(?=[A-Z])", " ", word)
 
     def _statistics(self, language, text):
         # Language word statistics:
@@ -328,13 +380,16 @@ class LangSegment:
             self._lang_count = defaultdict(int)
         lang_count = self._lang_count
         if not "|" in language:
-            lang_count[language] += int(len(text)*2) if language == "zh" else len(text)
+            lang_count[language] += (
+                int(len(text) * 2) if language == "zh" else len(text)
+            )
         self._lang_count = lang_count
 
     def _clear_text_number(self, text):
-        if text == "\n":return text,False # Keep Line Breaks
+        if text == "\n":
-        clear_text = re.sub(r'([^\w\s]+)','',re.sub(r'\n+','',text)).strip()
+            return text, False  # Keep Line Breaks
-        is_number = len(re.sub(re.compile(r'(\d+)'),'',clear_text)) == 0
+        clear_text = re.sub(r"([^\w\s]+)", "", re.sub(r"\n+", "", text)).strip()
+        is_number = len(re.sub(re.compile(r"(\d+)"), "", clear_text)) == 0
         return clear_text, is_number
 
     def _saveData(self, words, language: str, text: str, score: float, symbol=None):
@@ -342,12 +397,15 @@ class LangSegment:
         clear_text, is_number = self._clear_text_number(text)
         # Merge the same language and save the results
         preData = words[-1] if len(words) > 0 else None
-        if symbol is not None:pass
+        if symbol is not None:
+            pass
         elif preData is not None and preData["symbol"] is None:
-            if len(clear_text) == 0:language = preData["lang"]
+            if len(clear_text) == 0:
-            elif is_number == True:language = preData["lang"]
+                language = preData["lang"]
+            elif is_number == True:
+                language = preData["lang"]
             _, pre_is_number = self._clear_text_number(preData["text"])
-            if (preData["lang"] == language):
+            if preData["lang"] == language:
                 self._statistics(preData["lang"], text)
                 text = preData["text"] + text
                 preData["text"] = text
@@ -357,70 +415,108 @@ class LangSegment:
                 words.pop()
         elif is_number == True:
             priority_language = self._get_filters_string()[:2]
-            if priority_language in "ja-zh-en-ko-fr-vi":language = priority_language
+            if priority_language in "ja-zh-en-ko-fr-vi":
+                language = priority_language
         data = {"lang": language, "text": text, "score": score, "symbol": symbol}
         filters = self.Langfilters
-        if filters is None or len(filters) == 0 or "?" in language or   \
+        if (
-            language in filters or language in filters[0] or \
+            filters is None
-            filters[0] == "*" or filters[0] in "alls-mixs-autos":
+            or len(filters) == 0
+            or "?" in language
+            or language in filters
+            or language in filters[0]
+            or filters[0] == "*"
+            or filters[0] in "alls-mixs-autos"
+        ):
             words.append(data)
             self._statistics(data["lang"], data["text"])
         return data
 
     def _addwords(self, words, language, text, score, symbol=None):
-        if text == "\n":pass # Keep Line Breaks
+        if text == "\n":
-        elif text is None or len(text.strip()) == 0:return True
+            pass  # Keep Line Breaks
-        if language is None:language = ""
+        elif text is None or len(text.strip()) == 0:
+            return True
+        if language is None:
+            language = ""
         language = language.lower()
-        if language == 'en':text = self._insert_english_uppercase(text)
+        if language == "en":
+            text = self._insert_english_uppercase(text)
         # text = re.sub(r'[(（）)]', ',' , text) # Keep it.
         text_waits = self._text_waits
         ispre_waits = len(text_waits) > 0
         preResult = text_waits.pop() if ispre_waits else None
-        if preResult is None:preResult = words[-1] if len(words) > 0 else None
+        if preResult is None:
+            preResult = words[-1] if len(words) > 0 else None
         if preResult and ("|" in preResult["lang"]):
             pre_lang = preResult["lang"]
-            if language in pre_lang:preResult["lang"] = language = language.split("|")[0]
+            if language in pre_lang:
-            else:preResult["lang"]=pre_lang.split("|")[0]
+                preResult["lang"] = language = language.split("|")[0]
-            if ispre_waits:preResult = self._saveData(words,preResult["lang"],preResult["text"],preResult["score"],preResult["symbol"])
+            else:
+                preResult["lang"] = pre_lang.split("|")[0]
+            if ispre_waits:
+                preResult = self._saveData(
+                    words,
+                    preResult["lang"],
+                    preResult["text"],
+                    preResult["score"],
+                    preResult["symbol"],
+                )
         pre_lang = preResult["lang"] if preResult else None
-        if ("|" in language) and (pre_lang and not pre_lang in language and not "…" in language):language = language.split("|")[0]
+        if ("|" in language) and (
-        if "|" in language:self._text_waits.append({"lang":language,"text": text,"score":score,"symbol":symbol})
+            pre_lang and not pre_lang in language and not "…" in language
-        else:self._saveData(words,language,text,score,symbol)
+        ):
+            language = language.split("|")[0]
+        if "|" in language:
+            self._text_waits.append(
+                {"lang": language, "text": text, "score": score, "symbol": symbol}
+            )
+        else:
+            self._saveData(words, language, text, score, symbol)
         return False
 
     def _get_prev_data(self, words):
         data = words[-1] if words and len(words) > 0 else None
-        if data:return (data["lang"] , data["text"])
+        if data:
+            return (data["lang"], data["text"])
         return (None, "")
 
     def _match_ending(self, input, index):
-        if input is None or len(input) == 0:return False,None
+        if input is None or len(input) == 0:
-        input = re.sub(r'\s+', '', input)
+            return False, None
-        if len(input) == 0 or abs(index) > len(input):return False,None
+        input = re.sub(r"\s+", "", input)
+        if len(input) == 0 or abs(index) > len(input):
+            return False, None
         ending_pattern = re.compile(r'([「」“”‘’"\':：。.！!?．？])')
         return ending_pattern.match(input[index]), input[index]
 
     def _cleans_text(self, cleans_text):
-        cleans_text = re.sub(r'(.*?)([^\w]+)', r'\1 ', cleans_text)
+        cleans_text = re.sub(r"(.*?)([^\w]+)", r"\1 ", cleans_text)
-        cleans_text = re.sub(r'(.)\1+', r'\1', cleans_text)
+        cleans_text = re.sub(r"(.)\1+", r"\1", cleans_text)
         return cleans_text.strip()
 
     def _mean_processing(self, text: str):
-        if text is None or (text.strip()) == "":return None , 0.0
+        if text is None or (text.strip()) == "":
+            return None, 0.0
         arrs = self._split_camel_case(text).split(" ")
         langs = []
         for t in arrs:
-            if len(t.strip()) <= 3:continue
+            if len(t.strip()) <= 3:
+                continue
             language, score = self.langid.classify(t)
             langs.append({"lang": language})
-        if len(langs) == 0:return None , 0.0
+        if len(langs) == 0:
-        return Counter([item['lang'] for item in langs]).most_common(1)[0][0],1.0
+            return None, 0.0
+        return Counter([item["lang"] for item in langs]).most_common(1)[0][0], 1.0
 
     def _lang_classify(self, cleans_text):
         language, score = self.langid.classify(cleans_text)
         # fix: Huggingface is np.float32
-        if score is not None and isinstance(score, np.generic) and hasattr(score,"item"):
+        if (
+            score is not None
+            and isinstance(score, np.generic)
+            and hasattr(score, "item")
+        ):
             score = score.item()
         score = round(score, 3)
         return language, score
@@ -432,49 +528,80 @@ class LangSegment:
     def _parse_language(self, words, segment):
         LANG_JA = "ja"
         LANG_ZH = "zh"
-        LANG_ZH_JA = f'{LANG_ZH}|{LANG_JA}'
+        LANG_ZH_JA = f"{LANG_ZH}|{LANG_JA}"
-        LANG_JA_ZH = f'{LANG_JA}|{LANG_ZH}'
+        LANG_JA_ZH = f"{LANG_JA}|{LANG_ZH}"
         language = LANG_ZH
-        regex_pattern = re.compile(r'([^\w\s]+)')
+        regex_pattern = re.compile(r"([^\w\s]+)")
         lines = regex_pattern.split(segment)
         lines_max = len(lines)
         LANG_EOS = self._lang_eos
         for index, text in enumerate(lines):
-            if len(text) == 0:continue
+            if len(text) == 0:
+                continue
             EOS = index >= (lines_max - 1)
             nextId = index + 1
             nextText = lines[nextId] if not EOS else ""
-            nextPunc = len(re.sub(regex_pattern,'',re.sub(r'\n+','',nextText)).strip()) == 0
+            nextPunc = (
-            textPunc = len(re.sub(regex_pattern,'',re.sub(r'\n+','',text)).strip()) == 0
+                len(re.sub(regex_pattern, "", re.sub(r"\n+", "", nextText)).strip())
-            if not EOS and (textPunc == True or ( len(nextText.strip()) >= 0 and nextPunc == True)):
+                == 0
-                lines[nextId] = f'{text}{nextText}'
+            )
+            textPunc = (
+                len(re.sub(regex_pattern, "", re.sub(r"\n+", "", text)).strip()) == 0
+            )
+            if not EOS and (
+                textPunc == True or (len(nextText.strip()) >= 0 and nextPunc == True)
+            ):
+                lines[nextId] = f"{text}{nextText}"
                 continue
-            number_tags = re.compile(r'(⑥\d{6,}⑥)')
+            number_tags = re.compile(r"(⑥\d{6,}⑥)")
-            cleans_text = re.sub(number_tags, '' ,text)
+            cleans_text = re.sub(number_tags, "", text)
-            cleans_text = re.sub(r'\d+', '' ,cleans_text)
+            cleans_text = re.sub(r"\d+", "", cleans_text)
             cleans_text = self._cleans_text(cleans_text)
             # fix:Langid's recognition of short sentences is inaccurate, and it is spliced longer.
             if not EOS and len(cleans_text) <= 2:
-                lines[nextId] = f'{text}{nextText}'
+                lines[nextId] = f"{text}{nextText}"
                 continue
             language, score = self._lang_classify(cleans_text)
             prev_language, prev_text = self._get_prev_data(words)
-            if language != LANG_ZH and all('\u4e00' <= c <= '\u9fff' for c in re.sub(r'\s','',cleans_text)):language,score = LANG_ZH,1
+            if language != LANG_ZH and all(
+                "\u4e00" <= c <= "\u9fff" for c in re.sub(r"\s", "", cleans_text)
+            ):
+                language, score = LANG_ZH, 1
             if len(cleans_text) <= 5 and self._is_chinese(cleans_text):
                 filters_string = self._get_filters_string()
                 if score < self.LangPriorityThreshold and len(filters_string) > 0:
-                    index_ja , index_zh = filters_string.find(LANG_JA) , filters_string.find(LANG_ZH)
+                    index_ja, index_zh = filters_string.find(
-                    if index_ja != -1 and index_ja < index_zh:language = LANG_JA
+                        LANG_JA
-                    elif index_zh != -1 and index_zh < index_ja:language = LANG_ZH
+                    ), filters_string.find(LANG_ZH)
-                if self._is_japanese_kana(cleans_text):language = LANG_JA
+                    if index_ja != -1 and index_ja < index_zh:
-                elif len(cleans_text) > 2 and score > 0.90:pass
+                        language = LANG_JA
-                elif EOS and LANG_EOS:language = LANG_ZH if len(cleans_text) <= 1 else language
+                    elif index_zh != -1 and index_zh < index_ja:
+                        language = LANG_ZH
+                if self._is_japanese_kana(cleans_text):
+                    language = LANG_JA
+                elif len(cleans_text) > 2 and score > 0.90:
+                    pass
+                elif EOS and LANG_EOS:
+                    language = LANG_ZH if len(cleans_text) <= 1 else language
                 else:
-                    LANG_UNKNOWN = LANG_ZH_JA if language == LANG_ZH or (len(cleans_text) <=2 and prev_language == LANG_ZH) else LANG_JA_ZH
+                    LANG_UNKNOWN = (
+                        LANG_ZH_JA
+                        if language == LANG_ZH
+                        or (len(cleans_text) <= 2 and prev_language == LANG_ZH)
+                        else LANG_JA_ZH
+                    )
                     match_end, match_char = self._match_ending(text, -1)
-                    referen = prev_language in LANG_UNKNOWN or LANG_UNKNOWN in prev_language if prev_language else False
+                    referen = (
-                    if match_char in "。.": language = prev_language if referen and len(words) > 0 else language
+                        prev_language in LANG_UNKNOWN or LANG_UNKNOWN in prev_language
-                    else:language = f"{LANG_UNKNOWN}|…"
+                        if prev_language
+                        else False
+                    )
+                    if match_char in "。.":
+                        language = (
+                            prev_language if referen and len(words) > 0 else language
+                        )
+                    else:
+                        language = f"{LANG_UNKNOWN}|…"
             text, *_ = re.subn(number_tags, self._restore_number, text)
             self._addwords(words, language, text, score)
 
@@ -518,7 +645,8 @@ class LangSegment:
         return value
 
     def _pattern_symbols(self, item, text):
-        if text is None:return text
+        if text is None:
+            return text
         tag, pattern, process = item
         matches = pattern.findall(text)
         if len(matches) == 1 and "".join(matches[0]) == text:
@@ -549,19 +677,25 @@ class LangSegment:
         enablePreview = self.EnablePreview
         if enablePreview == True:
             # Experimental: Other language support
-            regex_pattern = re.compile(r'(.*?[。.?？!！]+[\n]{,1})')
+            regex_pattern = re.compile(r"(.*?[。.?？!！]+[\n]{,1})")
             lines = regex_pattern.split(text)
             for index, text in enumerate(lines):
-                if len(text.strip()) == 0:continue
+                if len(text.strip()) == 0:
+                    continue
                 cleans_text = self._cleans_text(text)
                 language, score = self._lang_classify(cleans_text)
                 if language not in filters:
                     language, score = self._mean_processing(cleans_text)
-                if language is None or score <= 0.0:continue
+                if language is None or score <= 0.0:
-                elif language in filters:pass # pass
+                    continue
-                elif score >= 0.95:continue # High score, but not in the filter, excluded.
+                elif language in filters:
-                elif score <= 0.15 and filters[:2] == "fr":language = priority_language
+                    pass  # pass
-                else:language = "en"
+                elif score >= 0.95:
+                    continue  # High score, but not in the filter, excluded.
+                elif score <= 0.15 and filters[:2] == "fr":
+                    language = priority_language
+                else:
+                    language = "en"
                 self._addwords(words, language, text, score)
         else:
             # Default is English
@@ -627,10 +761,12 @@ class LangSegment:
         segments = re.split(self.PARSE_TAG, text)
         segments_len = len(segments) - 1
         for index, text in enumerate(segments):
-            if root_tag:self._lang_eos = index >= segments_len
+            if root_tag:
+                self._lang_eos = index >= segments_len
             if self.PARSE_TAG.match(text):
                 process, data = text_cache[text]
-                if process:process(words , data)
+                if process:
+                    process(words, data)
             else:
                 self._parse_language(words, text)
         return words
@@ -638,19 +774,31 @@ class LangSegment:
     def _merge_results(self, words):
         new_word = []
         for index, cur_data in enumerate(words):
-            if "symbol" in cur_data:del cur_data["symbol"]
+            if "symbol" in cur_data:
-            if index == 0:new_word.append(cur_data)
+                del cur_data["symbol"]
+            if index == 0:
+                new_word.append(cur_data)
             else:
                 pre_data = new_word[-1]
                 if cur_data["lang"] == pre_data["lang"]:
                     pre_data["text"] = f'{pre_data["text"]}{cur_data["text"]}'
-                else:new_word.append(cur_data)
+                else:
+                    new_word.append(cur_data)
         return new_word
 
     def _parse_symbols(self, text):
         TAG_NUM = "00"  # "00" => default channels , "$0" => testing channel
-        TAG_S1,TAG_S2,TAG_P1,TAG_P2,TAG_EN,TAG_KO,TAG_RU,TAG_TH = "$1" ,"$2" ,"$3" ,"$4" ,"$5" ,"$6" ,"$7","$8"
+        TAG_S1, TAG_S2, TAG_P1, TAG_P2, TAG_EN, TAG_KO, TAG_RU, TAG_TH = (
-        TAG_BASE = re.compile(fr'(([【《（(“‘"\']*[LANGUAGE]+[\W\s]*)+)')
+            "$1",
+            "$2",
+            "$3",
+            "$4",
+            "$5",
+            "$6",
+            "$7",
+            "$8",
+        )
+        TAG_BASE = re.compile(rf'(([【《（(“‘"\']*[LANGUAGE]+[\W\s]*)+)')
         # Get custom language filter
         filters = self.Langfilters
         filters = filters if filters is not None else ""
@@ -663,51 +811,107 @@ class LangSegment:
         # -------------------------------------------------------------------------------------------------------
         # Preview feature, other language support
         enablePreview = self.EnablePreview
-        if "fr" in filters or \
+        if "fr" in filters or "vi" in filters:
-           "vi" in filters:enablePreview = True
+            enablePreview = True
         self.EnablePreview = enablePreview
         # 实验性：法语字符支持。Prise en charge des caractères français
         RE_FR = "" if not enablePreview else "àáâãäåæçèéêëìíîïðñòóôõöùúûüýþÿ"
         # 实验性：越南语字符支持。Hỗ trợ ký tự tiếng Việt
-        RE_VI = "" if not enablePreview else "đơưăáàảãạắằẳẵặấầẩẫậéèẻẽẹếềểễệíìỉĩịóòỏõọốồổỗộớờởỡợúùủũụứừửữựôâêơưỷỹ"
+        RE_VI = (
+            ""
+            if not enablePreview
+            else "đơưăáàảãạắằẳẵặấầẩẫậéèẻẽẹếềểễệíìỉĩịóòỏõọốồổỗộớờởỡợúùủũụứừửữựôâêơưỷỹ"
+        )
         # -------------------------------------------------------------------------------------------------------
         # Basic options:
         process_list = [
-            (  TAG_S1  , re.compile(self.SYMBOLS_PATTERN) , self._process_symbol  ),               # Symbol Tag
+            (
-            (  TAG_KO  , re.compile(re.sub(r'LANGUAGE',f'\uac00-\ud7a3',TAG_BASE.pattern))    , self._process_korean  ),              # Korean words
+                TAG_S1,
-            (  TAG_TH  , re.compile(re.sub(r'LANGUAGE',f'\u0E00-\u0E7F',TAG_BASE.pattern))    , self._process_Thai ),                 # Thai words support.
+                re.compile(self.SYMBOLS_PATTERN),
-            (  TAG_RU  , re.compile(re.sub(r'LANGUAGE',f'А-Яа-яЁё',TAG_BASE.pattern))         , self._process_Russian ),              # Russian words support.
+                self._process_symbol,
-            (  TAG_NUM , re.compile(r'(\W*\d+\W+\d*\W*\d*)')        , self._process_number  ),  # Number words, Universal in all languages, Ignore it.
+            ),  # Symbol Tag
-            (  TAG_EN  , re.compile(re.sub(r'LANGUAGE',f'a-zA-Z{RE_FR}{RE_VI}',TAG_BASE.pattern))    , self._process_english ),       # English words + Other language support.
+            (
-            (  TAG_P1  , re.compile(r'(["\'])(.*?)(\1)')         , self._process_quotes  ),     # Regular quotes
+                TAG_KO,
-            (  TAG_P2  , re.compile(r'([\n]*[【《（(“‘])([^【《（(“‘’”)）》】]{3,})([’”)）》】][\W\s]*[\n]{,1})')   , self._process_quotes  ),  # Special quotes, There are left and right.
+                re.compile(re.sub(r"LANGUAGE", f"\uac00-\ud7a3", TAG_BASE.pattern)),
+                self._process_korean,
+            ),  # Korean words
+            (
+                TAG_TH,
+                re.compile(re.sub(r"LANGUAGE", f"\u0e00-\u0e7f", TAG_BASE.pattern)),
+                self._process_Thai,
+            ),  # Thai words support.
+            (
+                TAG_RU,
+                re.compile(re.sub(r"LANGUAGE", f"А-Яа-яЁё", TAG_BASE.pattern)),
+                self._process_Russian,
+            ),  # Russian words support.
+            (
+                TAG_NUM,
+                re.compile(r"(\W*\d+\W+\d*\W*\d*)"),
+                self._process_number,
+            ),  # Number words, Universal in all languages, Ignore it.
+            (
+                TAG_EN,
+                re.compile(
+                    re.sub(r"LANGUAGE", f"a-zA-Z{RE_FR}{RE_VI}", TAG_BASE.pattern)
+                ),
+                self._process_english,
+            ),  # English words + Other language support.
+            (
+                TAG_P1,
+                re.compile(r'(["\'])(.*?)(\1)'),
+                self._process_quotes,
+            ),  # Regular quotes
+            (
+                TAG_P2,
+                re.compile(
+                    r"([\n]*[【《（(“‘])([^【《（(“‘’”)）》】]{3,})([’”)）》】][\W\s]*[\n]{,1})"
+                ),
+                self._process_quotes,
+            ),  # Special quotes, There are left and right.
         ]
         # Extended options: Default False
-        if self.keepPinyin == True:process_list.insert(1 , 
+        if self.keepPinyin == True:
-            (  TAG_S2  , re.compile(r'([\(（{](?:\s*\w*\d\w*\s*)+[}）\)])') , self._process_pinyin  ),     # Chinese Pinyin Tag. 
+            process_list.insert(
+                1,
+                (
+                    TAG_S2,
+                    re.compile(r"([\(（{](?:\s*\w*\d\w*\s*)+[}）\)])"),
+                    self._process_pinyin,
+                ),  # Chinese Pinyin Tag.
             )
         # -------------------------------------------------------------------------------------------------------
         words = []
-        lines = re.findall(r'.*\n*', re.sub(self.PARSE_TAG, '' ,text))
+        lines = re.findall(r".*\n*", re.sub(self.PARSE_TAG, "", text))
         for index, text in enumerate(lines):
-            if len(text.strip()) == 0:continue
+            if len(text.strip()) == 0:
+                continue
             self._lang_eos = False
             self._text_cache = {}
             for item in process_list:
                 text = self._pattern_symbols(item, text)
             cur_word = self._process_tags([], text, True)
-            if len(cur_word) == 0:continue
+            if len(cur_word) == 0:
+                continue
             cur_data = cur_word[0] if len(cur_word) > 0 else None
             pre_data = words[-1] if len(words) > 0 else None
-            if cur_data and pre_data and cur_data["lang"] == pre_data["lang"] \
+            if (
-                and cur_data["symbol"] == False and pre_data["symbol"] :
+                cur_data
+                and pre_data
+                and cur_data["lang"] == pre_data["lang"]
+                and cur_data["symbol"] == False
+                and pre_data["symbol"]
+            ):
                 cur_data["text"] = f'{pre_data["text"]}{cur_data["text"]}'
                 words.pop()
             words += cur_word
-        if self.isLangMerge == True:words = self._merge_results(words)
+        if self.isLangMerge == True:
+            words = self._merge_results(words)
         lang_count = self._lang_count
         if lang_count and len(lang_count) > 0:
-            lang_count = dict(sorted(lang_count.items(), key=lambda x: x[1], reverse=True))
+            lang_count = dict(
+                sorted(lang_count.items(), key=lambda x: x[1], reverse=True)
+            )
             lang_count = list(lang_count.items())
             self._lang_count = lang_count
         return words
@@ -732,12 +936,19 @@ class LangSegment:
 
     def getCounts(self):
         lang_count = self._lang_count
-        if lang_count is not None:return lang_count
+        if lang_count is not None:
+            return lang_count
         text_langs = self._text_langs
-        if text_langs is None or len(text_langs) == 0:return [("zh",0)]
+        if text_langs is None or len(text_langs) == 0:
+            return [("zh", 0)]
         lang_counts = defaultdict(int)
-        for d in text_langs:lang_counts[d['lang']] += int(len(d['text'])*2) if d['lang'] == "zh" else len(d['text'])
+        for d in text_langs:
-        lang_counts = dict(sorted(lang_counts.items(), key=lambda x: x[1], reverse=True))
+            lang_counts[d["lang"]] += (
+                int(len(d["text"]) * 2) if d["lang"] == "zh" else len(d["text"])
+            )
+        lang_counts = dict(
+            sorted(lang_counts.items(), key=lambda x: x[1], reverse=True)
+        )
         lang_counts = list(lang_counts.items())
         self._lang_count = lang_counts
         return lang_counts
@@ -748,7 +959,8 @@ class LangSegment:
             return []
         # lasts
         text_langs = self._text_langs
-        if self._text_lasts == text and text_langs is not None:return text_langs 
+        if self._text_lasts == text and text_langs is not None:
+            return text_langs
         # parse
         self._text_waits = []
         self._lang_count = None
@@ -760,6 +972,7 @@ class LangSegment:
     def classify(self, text: str):
         return self.getTexts(text)
 
+
 def printList(langlist):
     """
     功能：打印数组结果
@@ -776,7 +989,6 @@ def printList(langlist):
     pass
 
 
-
 def main():
 
     # -----------------------------------
@@ -795,11 +1007,9 @@ def main():
     # 输入示例3：（包含日文，中文）Input Example 1: (including Japanese, Chinese)
     # text = "明日、私たちは海辺にバカンスに行きます。你会说日语吗：“中国語、話せますか” 你的日语真好啊！"
 
-    
     # 输入示例4：（包含日文，中文，韩语，英文）Input Example 4: (including Japanese, Chinese, Korean, English)
     # text = "你的名字叫<ja>佐々木？<ja>吗？韩语中的안녕 오빠读什么呢？あなたの体育の先生は誰ですか? 此次发布会带来了四款iPhone 15系列机型和三款Apple Watch等一系列新品，这次的iPad Air采用了LCD屏幕"
 
-    
     # 试验性支持："fr"法语 , "vi"越南语 , "ru"俄语 , "th"泰语。Experimental: Other language support.
     langsegment = LangSegment()
     langsegment.setfilters(["fr", "vi", "ja", "zh", "ko", "en", "ru", "th"])
@@ -814,13 +1024,10 @@ Tôi thích nghe nhạc vào những ngày mưa.
 ฉันชอบฟังเพลงในวันที่ฝนตก
 """
 
-
-
     # 进行分词：（接入TTS项目仅需一行代码调用）Segmentation: (Only one line of code is required to access the TTS project)
     langlist = langsegment.getTexts(text)
     printList(langlist)
 
-    
     # 语种统计:Language statistics:
     print("\n===================【语种统计】===================")
     # 获取所有语种数组结果，根据内容字数降序排列
@@ -834,7 +1041,6 @@ Tôi thích nghe nhạc vào những ngày mưa.
     print(f"输入内容的主要语言为 = {lang} ，字数 = {count}")
     print("==================================================\n")
 
-    
     # 分词输出：lang=语言，text=内容。Word output: lang = language, text = content
     # ===================【打印结果】===================
     # {'lang': 'zh', 'text': '你的名字叫'}

acestep/language_segmentation/__init__.py

@@ -1,9 +1,9 @@
-from .LangSegment import LangSegment
+from acestep.language_segmentation.LangSegment import LangSegment
 
 
 # release
-__version__ = '0.3.5'
+__version__ = "0.3.5"
 
 
 # develop
-__develop__ = 'dev-0.0.1'
+__develop__ = "dev-0.0.1"

acestep/language_segmentation/utils/__init__.py

@@ -0,0 +1 @@
+# This file intentionally left blank for Python to recognize the directory as a package.

acestep/language_segmentation/utils/num.py

@@ -21,19 +21,21 @@ import re
 from collections import OrderedDict
 from typing import List
 
-DIGITS = {str(i): tran for i, tran in enumerate('零一二三四五六七八九')}
+DIGITS = {str(i): tran for i, tran in enumerate("零一二三四五六七八九")}
-UNITS = OrderedDict({
+UNITS = OrderedDict(
-    1: '十',
+    {
-    2: '百',
+        1: "十",
-    3: '千',
+        2: "百",
-    4: '万',
+        3: "千",
-    8: '亿',
+        4: "万",
-})
+        8: "亿",
+    }
+)
 
-COM_QUANTIFIERS = '(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)'
+COM_QUANTIFIERS = "(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)"
 
 # 分数表达式
-RE_FRAC = re.compile(r'(-?)(\d+)/(\d+)')
+RE_FRAC = re.compile(r"(-?)(\d+)/(\d+)")
 
 
 def replace_frac(match) -> str:
@@ -54,7 +56,7 @@ def replace_frac(match) -> str:
 
 
 # 百分数表达式
-RE_PERCENTAGE = re.compile(r'(-?)(\d+(\.\d+)?)%')
+RE_PERCENTAGE = re.compile(r"(-?)(\d+(\.\d+)?)%")
 
 
 def replace_percentage(match) -> str:
@@ -74,7 +76,7 @@ def replace_percentage(match) -> str:
 
 # 整数表达式
 # 带负号的整数 -10
-RE_INTEGER = re.compile(r'(-)' r'(\d+)')
+RE_INTEGER = re.compile(r"(-)" r"(\d+)")
 
 
 def replace_negative_num(match) -> str:
@@ -94,7 +96,7 @@ def replace_negative_num(match) -> str:
 
 # 编号-无符号整形
 # 00078
-RE_DEFAULT_NUM = re.compile(r'\d{3}\d*')
+RE_DEFAULT_NUM = re.compile(r"\d{3}\d*")
 
 
 def replace_default_num(match):
@@ -112,15 +114,11 @@ def replace_default_num(match):
 # RE_ASMD = re.compile(
 #     r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))([\+\-\×÷=])((-?)((\d+)(\.\d+)?)|(\.(\d+)))')
 RE_ASMD = re.compile(
-    r'((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))([\+\-\×÷=])((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))')
+    r"((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))([\+\-\×÷=])((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))"
+)
+
+asmd_map = {"+": "加", "-": "减", "×": "乘", "÷": "除", "=": "等于"}
 
-asmd_map = {
-    '+': '加',
-    '-': '减',
-    '×': '乘',
-    '÷': '除',
-    '=': '等于'
-}
 
 def replace_asmd(match) -> str:
     """
@@ -134,24 +132,25 @@ def replace_asmd(match) -> str:
 
 
 # 次方专项
-RE_POWER = re.compile(r'[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]+')
+RE_POWER = re.compile(r"[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]+")
 
 power_map = {
-    '⁰': '0',
+    "⁰": "0",
-    '¹': '1',
+    "¹": "1",
-    '²': '2',
+    "²": "2",
-    '³': '3',
+    "³": "3",
-    '⁴': '4',
+    "⁴": "4",
-    '⁵': '5',
+    "⁵": "5",
-    '⁶': '6',
+    "⁶": "6",
-    '⁷': '7',
+    "⁷": "7",
-    '⁸': '8',
+    "⁸": "8",
-    '⁹': '9',
+    "⁹": "9",
-    'ˣ': 'x',
+    "ˣ": "x",
-    'ʸ': 'y',
+    "ʸ": "y",
-    'ⁿ': 'n'
+    "ⁿ": "n",
 }
 
+
 def replace_power(match) -> str:
     """
     Args:
@@ -168,10 +167,10 @@ def replace_power(match) -> str:
 
 # 数字表达式
 # 纯小数
-RE_DECIMAL_NUM = re.compile(r'(-?)((\d+)(\.\d+))' r'|(\.(\d+))')
+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+))')
+RE_NUMBER = re.compile(r"(-?)((\d+)(\.\d+)?)" r"|(\.(\d+))")
 
 
 def replace_positive_quantifier(match) -> str:
@@ -221,7 +220,9 @@ RE_RANGE = re.compile(
     [-~]                   # 匹配范围分隔符
     ((-?)((\d+)(\.\d+)?))  # 匹配范围结束的负数或正数（整数或小数）
     (?![\d\+\-\×÷=])       # 使用正向前瞻以确保数字范围之后没有其他数字和操作符
-    """, re.VERBOSE)
+    """,
+    re.VERBOSE,
+)
 
 
 def replace_range(match) -> str:
@@ -240,7 +241,9 @@ def replace_range(match) -> str:
 
 # ~至表达式
 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)')
+    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:
     """
@@ -249,58 +252,61 @@ def replace_to_range(match) -> str:
     Returns:
         str
     """
-    result = match.group(0).replace('~', '至')
+    result = match.group(0).replace("~", "至")
     return result
 
 
 def _get_value(value_string: str, use_zero: bool = True) -> List[str]:
-    stripped = value_string.lstrip('0')
+    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]]
+            return [DIGITS["0"], DIGITS[stripped]]
         else:
             return [DIGITS[stripped]]
     else:
         largest_unit = next(
-            power for power in reversed(UNITS.keys()) if power < len(stripped))
+            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(
+        return _get_value(first_part) + [UNITS[largest_unit]] + _get_value(second_part)
-            second_part)
 
 
 def verbalize_cardinal(value_string: str) -> str:
     if not value_string:
-        return ''
+        return ""
 
     # 000 -> '零' , 0 -> '零'
-    value_string = value_string.lstrip('0')
+    value_string = value_string.lstrip("0")
     if len(value_string) == 0:
-        return DIGITS['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[
+    if (
-            '1'] and result_symbols[1] == UNITS[1]:
+        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)
+    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)
+    result = "".join(result_symbols)
     if alt_one:
         result = result.replace("一", "幺")
     return result
 
 
 def num2str(value_string: str) -> str:
-    integer_decimal = value_string.split('.')
+    integer_decimal = value_string.split(".")
     if len(integer_decimal) == 1:
         integer = integer_decimal[0]
-        decimal = ''
+        decimal = ""
     elif len(integer_decimal) == 2:
         integer, decimal = integer_decimal
     else:
@@ -310,12 +316,12 @@ def num2str(value_string: str) -> str:
 
     result = verbalize_cardinal(integer)
 
-    decimal = decimal.rstrip('0')
+    decimal = decimal.rstrip("0")
     if decimal:
         # '.22' is verbalized as '零点二二'
         # '3.20' is verbalized as '三点二
         result = result if result else "零"
-        result += '点' + verbalize_digit(decimal)
+        result += "点" + verbalize_digit(decimal)
     return result
 
 

acestep/models/ace_step_transformer.py

@@ -31,9 +31,15 @@ from .lyrics_utils.lyric_encoder import ConformerEncoder as LyricEncoder
 
 def cross_norm(hidden_states, controlnet_input):
     # input N x T x c
-    mean_hidden_states, std_hidden_states = hidden_states.mean(dim=(1,2), keepdim=True), hidden_states.std(dim=(1,2), keepdim=True)
+    mean_hidden_states, std_hidden_states = hidden_states.mean(
-    mean_controlnet_input, std_controlnet_input = controlnet_input.mean(dim=(1,2), keepdim=True), controlnet_input.std(dim=(1,2), keepdim=True)
+        dim=(1, 2), keepdim=True
-    controlnet_input = (controlnet_input - mean_controlnet_input) * (std_hidden_states / (std_controlnet_input + 1e-12)) + mean_hidden_states
+    ), hidden_states.std(dim=(1, 2), keepdim=True)
+    mean_controlnet_input, std_controlnet_input = controlnet_input.mean(
+        dim=(1, 2), keepdim=True
+    ), controlnet_input.std(dim=(1, 2), keepdim=True)
+    controlnet_input = (controlnet_input - mean_controlnet_input) * (
+        std_hidden_states / (std_controlnet_input + 1e-12)
+    ) + mean_hidden_states
     return controlnet_input
 
 
@@ -45,17 +51,27 @@ class Qwen2RotaryEmbedding(nn.Module):
         self.dim = dim
         self.max_position_embeddings = max_position_embeddings
         self.base = base
-        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
+        inv_freq = 1.0 / (
+            self.base
+            ** (
+                torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device)
+                / self.dim
+            )
+        )
         self.register_buffer("inv_freq", inv_freq, persistent=False)
 
         # Build here to make `torch.jit.trace` work.
         self._set_cos_sin_cache(
-            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+            seq_len=max_position_embeddings,
+            device=self.inv_freq.device,
+            dtype=torch.get_default_dtype(),
         )
 
     def _set_cos_sin_cache(self, seq_len, device, dtype):
         self.max_seq_len_cached = seq_len
-        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
+        t = torch.arange(
+            self.max_seq_len_cached, device=device, dtype=torch.int64
+        ).type_as(self.inv_freq)
 
         freqs = torch.outer(t, self.inv_freq)
         # Different from paper, but it uses a different permutation in order to obtain the same calculation
@@ -82,8 +98,12 @@ class T2IFinalLayer(nn.Module):
     def __init__(self, hidden_size, patch_size=[16, 1], out_channels=256):
         super().__init__()
         self.norm_final = nn.RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.linear = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True)
+        self.linear = nn.Linear(
-        self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5)
+            hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True
+        )
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(2, hidden_size) / hidden_size**0.5
+        )
         self.out_channels = out_channels
         self.patch_size = patch_size
 
@@ -95,14 +115,28 @@ class T2IFinalLayer(nn.Module):
         # 4 unpatchify
         new_height, new_width = 1, hidden_states.size(1)
         hidden_states = hidden_states.reshape(
-            shape=(hidden_states.shape[0], new_height, new_width, self.patch_size[0], self.patch_size[1], self.out_channels)
+            shape=(
+                hidden_states.shape[0],
+                new_height,
+                new_width,
+                self.patch_size[0],
+                self.patch_size[1],
+                self.out_channels,
+            )
         ).contiguous()
         hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
         output = hidden_states.reshape(
-            shape=(hidden_states.shape[0], self.out_channels, new_height * self.patch_size[0], new_width * self.patch_size[1])
+            shape=(
+                hidden_states.shape[0],
+                self.out_channels,
+                new_height * self.patch_size[0],
+                new_width * self.patch_size[1],
+            )
         ).contiguous()
         if width > new_width:
-            output = torch.nn.functional.pad(output, (0, width - new_width, 0, 0), 'constant', 0)
+            output = torch.nn.functional.pad(
+                output, (0, width - new_width, 0, 0), "constant", 0
+            )
         elif width < new_width:
             output = output[:, :, :, :width]
         return output
@@ -131,9 +165,25 @@ class PatchEmbed(nn.Module):
         super().__init__()
         patch_size_h, patch_size_w = patch_size
         self.early_conv_layers = nn.Sequential(
-            nn.Conv2d(in_channels, in_channels*256, kernel_size=patch_size, stride=patch_size, padding=0, bias=bias),
+            nn.Conv2d(
-            torch.nn.GroupNorm(num_groups=32, num_channels=in_channels*256, eps=1e-6, affine=True),
+                in_channels,
-            nn.Conv2d(in_channels*256, embed_dim, kernel_size=1, stride=1, padding=0, bias=bias)
+                in_channels * 256,
+                kernel_size=patch_size,
+                stride=patch_size,
+                padding=0,
+                bias=bias,
+            ),
+            torch.nn.GroupNorm(
+                num_groups=32, num_channels=in_channels * 256, eps=1e-6, affine=True
+            ),
+            nn.Conv2d(
+                in_channels * 256,
+                embed_dim,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                bias=bias,
+            ),
         )
         self.patch_size = patch_size
         self.height, self.width = height // patch_size_h, width // patch_size_w
@@ -153,7 +203,9 @@ class Transformer2DModelOutput(BaseOutput):
     proj_losses: Optional[Tuple[Tuple[str, torch.Tensor]]] = None
 
 
-class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+class ACEStepTransformer2DModel(
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
+):
     _supports_gradient_checkpointing = True
 
     @register_to_config
@@ -217,9 +269,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
         )
         self.num_layers = num_layers
 
-        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.time_proj = Timesteps(
-        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim)
+            num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0
-        self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True))
+        )
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=256, time_embed_dim=self.inner_dim
+        )
+        self.t_block = nn.Sequential(
+            nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True)
+        )
 
         # speaker
         self.speaker_embedder = nn.Linear(speaker_embedding_dim, self.inner_dim)
@@ -229,25 +287,30 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
 
         # lyric
         self.lyric_embs = nn.Embedding(lyric_encoder_vocab_size, lyric_hidden_size)
-        self.lyric_encoder = LyricEncoder(input_size=lyric_hidden_size, static_chunk_size=0)
+        self.lyric_encoder = LyricEncoder(
+            input_size=lyric_hidden_size, static_chunk_size=0
+        )
         self.lyric_proj = nn.Linear(lyric_hidden_size, self.inner_dim)
 
         projector_dim = 2 * self.inner_dim
 
-        self.projectors = nn.ModuleList([
+        self.projectors = nn.ModuleList(
+            [
                 nn.Sequential(
                     nn.Linear(self.inner_dim, projector_dim),
                     nn.SiLU(),
                     nn.Linear(projector_dim, projector_dim),
                     nn.SiLU(),
                     nn.Linear(projector_dim, ssl_dim),
-            ) for ssl_dim in ssl_latent_dims
+                )
-        ])
+                for ssl_dim in ssl_latent_dims
+            ]
+        )
 
         self.ssl_latent_dims = ssl_latent_dims
         self.ssl_encoder_depths = ssl_encoder_depths
 
-        self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction='mean')
+        self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction="mean")
         self.ssl_names = ssl_names
 
         self.proj_in = PatchEmbed(
@@ -258,11 +321,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
             bias=True,
         )
 
-        self.final_layer = T2IFinalLayer(self.inner_dim, patch_size=patch_size, out_channels=out_channels)
+        self.final_layer = T2IFinalLayer(
+            self.inner_dim, patch_size=patch_size, out_channels=out_channels
+        )
         self.gradient_checkpointing = False
 
     # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
-    def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+    def enable_forward_chunking(
+        self, chunk_size: Optional[int] = None, dim: int = 0
+    ) -> None:
         """
         Sets the attention processor to use [feed forward
         chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -281,7 +348,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
         # By default chunk size is 1
         chunk_size = chunk_size or 1
 
-        def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
+        def fn_recursive_feed_forward(
+            module: torch.nn.Module, chunk_size: int, dim: int
+        ):
             if hasattr(module, "set_chunk_feed_forward"):
                 module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
 
@@ -302,7 +371,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
     ):
         # N x T x D
         lyric_embs = self.lyric_embs(lyric_token_idx)
-        prompt_prenet_out, _mask = self.lyric_encoder(lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1)
+        prompt_prenet_out, _mask = self.lyric_encoder(
+            lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1
+        )
         prompt_prenet_out = self.lyric_proj(prompt_prenet_out)
         return prompt_prenet_out
 
@@ -331,8 +402,17 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
             lyric_mask=lyric_mask,
         )
 
-        encoder_hidden_states = torch.cat([encoder_spk_hidden_states, encoder_text_hidden_states, encoder_lyric_hidden_states], dim=1)
+        encoder_hidden_states = torch.cat(
-        encoder_hidden_mask = torch.cat([speaker_mask, text_attention_mask, lyric_mask], dim=1)
+            [
+                encoder_spk_hidden_states,
+                encoder_text_hidden_states,
+                encoder_lyric_hidden_states,
+            ],
+            dim=1,
+        )
+        encoder_hidden_mask = torch.cat(
+            [speaker_mask, text_attention_mask, lyric_mask], dim=1
+        )
         return encoder_hidden_states, encoder_hidden_mask
 
     def decode(
@@ -344,12 +424,16 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
         timestep: Optional[torch.Tensor],
         ssl_hidden_states: Optional[List[torch.Tensor]] = None,
         output_length: int = 0,
-        block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
+        block_controlnet_hidden_states: Optional[
+            Union[List[torch.Tensor], torch.Tensor]
+        ] = None,
         controlnet_scale: Union[float, torch.Tensor] = 1.0,
         return_dict: bool = True,
     ):
 
-        embedded_timestep = self.timestep_embedder(self.time_proj(timestep).to(dtype=hidden_states.dtype))
+        embedded_timestep = self.timestep_embedder(
+            self.time_proj(timestep).to(dtype=hidden_states.dtype)
+        )
         temb = self.t_block(embedded_timestep)
 
         hidden_states = self.proj_in(hidden_states)
@@ -361,8 +445,12 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
 
         inner_hidden_states = []
 
-        rotary_freqs_cis = self.rotary_emb(hidden_states, seq_len=hidden_states.shape[1])
+        rotary_freqs_cis = self.rotary_emb(
-        encoder_rotary_freqs_cis = self.rotary_emb(encoder_hidden_states, seq_len=encoder_hidden_states.shape[1])
+            hidden_states, seq_len=hidden_states.shape[1]
+        )
+        encoder_rotary_freqs_cis = self.rotary_emb(
+            encoder_hidden_states, seq_len=encoder_hidden_states.shape[1]
+        )
 
         for index_block, block in enumerate(self.transformer_blocks):
 
@@ -377,7 +465,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
 
                     return custom_forward
 
-                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                ckpt_kwargs: Dict[str, Any] = (
+                    {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                )
                 hidden_states = torch.utils.checkpoint.checkpoint(
                     create_custom_forward(block),
                     hidden_states=hidden_states,
@@ -406,9 +496,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
                     inner_hidden_states.append(hidden_states)
 
         proj_losses = []
-        if len(inner_hidden_states) > 0 and ssl_hidden_states is not None and len(ssl_hidden_states) > 0:
+        if (
+            len(inner_hidden_states) > 0
+            and ssl_hidden_states is not None
+            and len(ssl_hidden_states) > 0
+        ):
 
-            for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names):
+            for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(
+                inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names
+            ):
                 if ssl_hidden_state is None:
                     continue
                 # 1. N x T x D1 -> N x D x D2
@@ -416,9 +512,20 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
                 # 3. projection loss
                 bs = inner_hidden_state.shape[0]
                 proj_loss = 0.0
-                for i, (z, z_tilde) in enumerate(zip(ssl_hidden_state, est_ssl_hidden_state)):
+                for i, (z, z_tilde) in enumerate(
+                    zip(ssl_hidden_state, est_ssl_hidden_state)
+                ):
                     # 2. interpolate
-                    z_tilde = F.interpolate(z_tilde.unsqueeze(0).transpose(1, 2), size=len(z), mode='linear', align_corners=False).transpose(1, 2).squeeze(0)
+                    z_tilde = (
+                        F.interpolate(
+                            z_tilde.unsqueeze(0).transpose(1, 2),
+                            size=len(z),
+                            mode="linear",
+                            align_corners=False,
+                        )
+                        .transpose(1, 2)
+                        .squeeze(0)
+                    )
 
                     z_tilde = torch.nn.functional.normalize(z_tilde, dim=-1)
                     z = torch.nn.functional.normalize(z, dim=-1)
@@ -445,7 +552,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
         lyric_mask: Optional[torch.LongTensor] = None,
         timestep: Optional[torch.Tensor] = None,
         ssl_hidden_states: Optional[List[torch.Tensor]] = None,
-        block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
+        block_controlnet_hidden_states: Optional[
+            Union[List[torch.Tensor], torch.Tensor]
+        ] = None,
         controlnet_scale: Union[float, torch.Tensor] = 1.0,
         return_dict: bool = True,
     ):

acestep/models/attention.py

@@ -23,10 +23,18 @@ from diffusers.models.normalization import RMSNorm
 
 try:
     # from .dcformer import DCMHAttention
-    from .customer_attention_processor import Attention, CustomLiteLAProcessor2_0, CustomerAttnProcessor2_0
+    from .customer_attention_processor import (
+        Attention,
+        CustomLiteLAProcessor2_0,
+        CustomerAttnProcessor2_0,
+    )
 except ImportError:
     # from dcformer import DCMHAttention
-    from customer_attention_processor import Attention, CustomLiteLAProcessor2_0, CustomerAttnProcessor2_0
+    from customer_attention_processor import (
+        Attention,
+        CustomLiteLAProcessor2_0,
+        CustomerAttnProcessor2_0,
+    )
 
 
 logger = logging.get_logger(__name__)
@@ -55,13 +63,16 @@ def t2i_modulate(x, shift, scale):
     return x * (1 + scale) + shift
 
 
-def get_same_padding(kernel_size: Union[int, Tuple[int, ...]]) -> Union[int, Tuple[int, ...]]:
+def get_same_padding(
+    kernel_size: Union[int, Tuple[int, ...]],
+) -> Union[int, Tuple[int, ...]]:
     if isinstance(kernel_size, tuple):
         return tuple([get_same_padding(ks) for ks in kernel_size])
     else:
         assert kernel_size % 2 > 0, f"kernel size {kernel_size} should be odd number"
         return kernel_size // 2
 
+
 class ConvLayer(nn.Module):
     def __init__(
         self,
@@ -187,6 +198,7 @@ class LinearTransformerBlock(nn.Module):
     """
     A Sana block with global shared adaptive layer norm (adaLN-single) conditioning.
     """
+
     def __init__(
         self,
         dim,

acestep/models/customer_attention_processor.py

@@ -78,12 +78,16 @@ class CustomLiteLAProcessor2_0:
         input_ndim = hidden_states.ndim
         if input_ndim == 4:
             batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
         if encoder_hidden_states is not None:
             context_input_ndim = encoder_hidden_states.ndim
             if context_input_ndim == 4:
                 batch_size, channel, height, width = encoder_hidden_states.shape
-                encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+                encoder_hidden_states = encoder_hidden_states.view(
+                    batch_size, channel, height * width
+                ).transpose(1, 2)
 
         batch_size = hidden_states.shape[0]
 
@@ -94,7 +98,11 @@ class CustomLiteLAProcessor2_0:
         value = attn.to_v(hidden_states)
 
         # `context` projections.
-        has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
+        has_encoder_hidden_state_proj = (
+            hasattr(attn, "add_q_proj")
+            and hasattr(attn, "add_k_proj")
+            and hasattr(attn, "add_v_proj")
+        )
         if encoder_hidden_states is not None and has_encoder_hidden_state_proj:
             encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
             encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
@@ -114,7 +122,11 @@ class CustomLiteLAProcessor2_0:
         head_dim = inner_dim // attn.heads
 
         query = query.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
-        key = key.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1).transpose(-1, -2)
+        key = (
+            key.transpose(-1, -2)
+            .reshape(batch_size, attn.heads, head_dim, -1)
+            .transpose(-1, -2)
+        )
         value = value.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
 
         # RoPE需要 [B, H, S, D] 输入
@@ -140,17 +152,33 @@ class CustomLiteLAProcessor2_0:
 
         if attention_mask is not None:
             # attention_mask: [B, S] -> [B, 1, S, 1]
-            attention_mask = attention_mask[:, None, :, None].to(key.dtype)  # [B, 1, S, 1]
+            attention_mask = attention_mask[:, None, :, None].to(
-            query = query * attention_mask.permute(0, 1, 3, 2)  # [B, H, S, D] * [B, 1, S, 1]
+                key.dtype
+            )  # [B, 1, S, 1]
+            query = query * attention_mask.permute(
+                0, 1, 3, 2
+            )  # [B, H, S, D] * [B, 1, S, 1]
             if not attn.is_cross_attention:
-                key = key * attention_mask  # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘
+                key = (
-                value = value * attention_mask.permute(0, 1, 3, 2)  # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度
+                    key * attention_mask
+                )  # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘
+                value = value * attention_mask.permute(
+                    0, 1, 3, 2
+                )  # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度
 
-        if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
+        if (
-            encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(key.dtype)  # [B, 1, S_enc, 1]
+            attn.is_cross_attention
+            and encoder_attention_mask is not None
+            and has_encoder_hidden_state_proj
+        ):
+            encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(
+                key.dtype
+            )  # [B, 1, S_enc, 1]
             # 此时 key: [B, h, S_enc, D], value: [B, h, D, S_enc]
             key = key * encoder_attention_mask  # [B, h, S_enc, D] * [B, 1, S_enc, 1]
-            value = value * encoder_attention_mask.permute(0, 1, 3, 2)  # [B, h, D, S_enc] * [B, 1, 1, S_enc]
+            value = value * encoder_attention_mask.permute(
+                0, 1, 3, 2
+            )  # [B, h, D, S_enc] * [B, 1, 1, S_enc]
 
         query = self.kernel_func(query)
         key = self.kernel_func(key)
@@ -168,14 +196,20 @@ class CustomLiteLAProcessor2_0:
 
         hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps)
 
-        hidden_states = hidden_states.view(batch_size, attn.heads * head_dim, -1).permute(0, 2, 1)
+        hidden_states = hidden_states.view(
+            batch_size, attn.heads * head_dim, -1
+        ).permute(0, 2, 1)
 
         hidden_states = hidden_states.to(dtype)
         if encoder_hidden_states is not None:
             encoder_hidden_states = encoder_hidden_states.to(dtype)
 
         # Split the attention outputs.
-        if encoder_hidden_states is not None and not attn.is_cross_attention and has_encoder_hidden_state_proj:
+        if (
+            encoder_hidden_states is not None
+            and not attn.is_cross_attention
+            and has_encoder_hidden_state_proj
+        ):
             hidden_states, encoder_hidden_states = (
                 hidden_states[:, :hidden_states_len],
                 hidden_states[:, hidden_states_len:],
@@ -185,13 +219,22 @@ class CustomLiteLAProcessor2_0:
         hidden_states = attn.to_out[0](hidden_states)
         # dropout
         hidden_states = attn.to_out[1](hidden_states)
-        if encoder_hidden_states is not None and not attn.context_pre_only and not attn.is_cross_attention and hasattr(attn, "to_add_out"):
+        if (
+            encoder_hidden_states is not None
+            and not attn.context_pre_only
+            and not attn.is_cross_attention
+            and hasattr(attn, "to_add_out")
+        ):
             encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
 
         if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
         if encoder_hidden_states is not None and context_input_ndim == 4:
-            encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+            encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
 
         if torch.get_autocast_gpu_dtype() == torch.float16:
             hidden_states = hidden_states.clip(-65504, 65504)
@@ -208,7 +251,9 @@ class CustomerAttnProcessor2_0:
 
     def __init__(self):
         if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
 
     def apply_rotary_emb(
         self,
@@ -258,23 +303,35 @@ class CustomerAttnProcessor2_0:
 
         if input_ndim == 4:
             batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
 
         batch_size, sequence_length, _ = (
-            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
         )
 
-        has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
+        has_encoder_hidden_state_proj = (
+            hasattr(attn, "add_q_proj")
+            and hasattr(attn, "add_k_proj")
+            and hasattr(attn, "add_v_proj")
+        )
 
         if attn.group_norm is not None:
-            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
 
         query = attn.to_q(hidden_states)
 
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
         elif attn.norm_cross:
-            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+            encoder_hidden_states = attn.norm_encoder_hidden_states(
+                encoder_hidden_states
+            )
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
@@ -300,19 +357,33 @@ class CustomerAttnProcessor2_0:
             elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:
                 key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)
 
-        if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
+        if (
+            attn.is_cross_attention
+            and encoder_attention_mask is not None
+            and has_encoder_hidden_state_proj
+        ):
             # attention_mask: N x S1
             # encoder_attention_mask: N x S2
             # cross attention 整合attention_mask和encoder_attention_mask
-            combined_mask = attention_mask[:, :, None] * encoder_attention_mask[:, None, :]
+            combined_mask = (
+                attention_mask[:, :, None] * encoder_attention_mask[:, None, :]
+            )
             attention_mask = torch.where(combined_mask == 1, 0.0, -torch.inf)
-            attention_mask = attention_mask[:, None, :, :].expand(-1, attn.heads, -1, -1).to(query.dtype)
+            attention_mask = (
+                attention_mask[:, None, :, :]
+                .expand(-1, attn.heads, -1, -1)
+                .to(query.dtype)
+            )
 
         elif not attn.is_cross_attention and attention_mask is not None:
-            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
             # scaled_dot_product_attention expects attention_mask shape to be
             # (batch, heads, source_length, target_length)
-            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
 
         # the output of sdp = (batch, num_heads, seq_len, head_dim)
         # TODO: add support for attn.scale when we move to Torch 2.1
@@ -320,7 +391,9 @@ class CustomerAttnProcessor2_0:
             query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
         )
 
-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, attn.heads * head_dim
+        )
         hidden_states = hidden_states.to(query.dtype)
 
         # linear proj
@@ -329,7 +402,9 @@ class CustomerAttnProcessor2_0:
         hidden_states = attn.to_out[1](hidden_states)
 
         if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
 
         if attn.residual_connection:
             hidden_states = hidden_states + residual

acestep/models/lyrics_utils/lyric_encoder.py

@@ -3,16 +3,19 @@ import math
 import torch
 from torch import nn
 
+
 class ConvolutionModule(nn.Module):
     """ConvolutionModule in Conformer model."""
 
-    def __init__(self,
+    def __init__(
+        self,
         channels: int,
         kernel_size: int = 15,
         activation: nn.Module = nn.ReLU(),
         norm: str = "batch_norm",
         causal: bool = False,
-                 bias: bool = True):
+        bias: bool = True,
+    ):
         """Construct an ConvolutionModule object.
         Args:
             channels (int): The number of channels of conv layers.
@@ -51,7 +54,7 @@ class ConvolutionModule(nn.Module):
             bias=bias,
         )
 
-        assert norm in ['batch_norm', 'layer_norm']
+        assert norm in ["batch_norm", "layer_norm"]
         if norm == "batch_norm":
             self.use_layer_norm = False
             self.norm = nn.BatchNorm1d(channels)
@@ -95,12 +98,12 @@ class ConvolutionModule(nn.Module):
 
         if self.lorder > 0:
             if cache.size(2) == 0:  # cache_t == 0
-                x = nn.functional.pad(x, (self.lorder, 0), 'constant', 0.0)
+                x = nn.functional.pad(x, (self.lorder, 0), "constant", 0.0)
             else:
                 assert cache.size(0) == x.size(0)  # equal batch
                 assert cache.size(1) == x.size(1)  # equal channel
                 x = torch.cat((cache, x), dim=2)
-            assert (x.size(2) > self.lorder)
+            assert x.size(2) > self.lorder
             new_cache = x[:, :, -self.lorder :]
         else:
             # It's better we just return None if no cache is required,
@@ -126,6 +129,7 @@ class ConvolutionModule(nn.Module):
 
         return x.transpose(1, 2), new_cache
 
+
 class PositionwiseFeedForward(torch.nn.Module):
     """Positionwise feed forward layer.
 
@@ -163,6 +167,7 @@ class PositionwiseFeedForward(torch.nn.Module):
         """
         return self.w_2(self.dropout(self.activation(self.w_1(xs))))
 
+
 class Swish(torch.nn.Module):
     """Construct an Swish object."""
 
@@ -170,6 +175,7 @@ class Swish(torch.nn.Module):
         """Return Swish activation function."""
         return x * torch.sigmoid(x)
 
+
 class MultiHeadedAttention(nn.Module):
     """Multi-Head Attention layer.
 
@@ -180,11 +186,9 @@ class MultiHeadedAttention(nn.Module):
 
     """
 
-    def __init__(self,
+    def __init__(
-                 n_head: int,
+        self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True
-                 n_feat: int,
+    ):
-                 dropout_rate: float,
-                 key_bias: bool = True):
         """Construct an MultiHeadedAttention object."""
         super().__init__()
         assert n_feat % n_head == 0
@@ -229,7 +233,7 @@ class MultiHeadedAttention(nn.Module):
         self,
         value: torch.Tensor,
         scores: torch.Tensor,
-        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
+        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
     ) -> torch.Tensor:
         """Compute attention context vector.
 
@@ -252,17 +256,18 @@ class MultiHeadedAttention(nn.Module):
             mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
             # For last chunk, time2 might be larger than scores.size(-1)
             mask = mask[:, :, :, : scores.size(-1)]  # (batch, 1, *, time2)
-            scores = scores.masked_fill(mask, -float('inf'))
+            scores = scores.masked_fill(mask, -float("inf"))
             attn = torch.softmax(scores, dim=-1).masked_fill(
-                mask, 0.0)  # (batch, head, time1, time2)
+                mask, 0.0
+            )  # (batch, head, time1, time2)
 
         else:
             attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
 
         p_attn = self.dropout(attn)
         x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
-        x = (x.transpose(1, 2).contiguous().view(n_batch, -1,
+        x = (
-                                                 self.h * self.d_k)
+            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
         )  # (batch, time1, d_model)
 
         return self.linear_out(x)  # (batch, time1, d_model)
@@ -274,7 +279,7 @@ class MultiHeadedAttention(nn.Module):
         value: torch.Tensor,
         mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
         pos_emb: torch.Tensor = torch.empty(0),
-        cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+        cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Compute scaled dot product attention.
 
@@ -308,9 +313,7 @@ class MultiHeadedAttention(nn.Module):
         """
         q, k, v = self.forward_qkv(query, key, value)
         if cache.size(0) > 0:
-            key_cache, value_cache = torch.split(cache,
+            key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
-                                                 cache.size(-1) // 2,
-                                                 dim=-1)
             k = torch.cat([key_cache, k], dim=2)
             v = torch.cat([value_cache, v], dim=2)
         new_cache = torch.cat((k, v), dim=-1)
@@ -328,11 +331,9 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
         dropout_rate (float): Dropout rate.
     """
 
-    def __init__(self,
+    def __init__(
-                 n_head: int,
+        self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True
-                 n_feat: int,
+    ):
-                 dropout_rate: float,
-                 key_bias: bool = True):
         """Construct an RelPositionMultiHeadedAttention object."""
         super().__init__(n_head, n_feat, dropout_rate, key_bias)
         # linear transformation for positional encoding
@@ -355,14 +356,12 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
             torch.Tensor: Output tensor.
 
         """
-        zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1),
+        zero_pad = torch.zeros(
-                               device=x.device,
+            (x.size()[0], x.size()[1], x.size()[2], 1), device=x.device, dtype=x.dtype
-                               dtype=x.dtype)
+        )
         x_padded = torch.cat([zero_pad, x], dim=-1)
 
-        x_padded = x_padded.view(x.size()[0],
+        x_padded = x_padded.view(x.size()[0], x.size()[1], x.size(3) + 1, x.size(2))
-                                 x.size()[1],
-                                 x.size(3) + 1, x.size(2))
         x = x_padded[:, :, 1:].view_as(x)[
             :, :, :, : x.size(-1) // 2 + 1
         ]  # only keep the positions from 0 to time2
@@ -375,7 +374,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
         value: torch.Tensor,
         mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
         pos_emb: torch.Tensor = torch.empty(0),
-        cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+        cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
         Args:
@@ -399,9 +398,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
         q = q.transpose(1, 2)  # (batch, time1, head, d_k)
 
         if cache.size(0) > 0:
-            key_cache, value_cache = torch.split(cache,
+            key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
-                                                 cache.size(-1) // 2,
-                                                 dim=-1)
             k = torch.cat([key_cache, k], dim=2)
             v = torch.cat([value_cache, v], dim=2)
         # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
@@ -431,12 +428,12 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
             matrix_bd = self.rel_shift(matrix_bd)
 
         scores = (matrix_ac + matrix_bd) / math.sqrt(
-            self.d_k)  # (batch, head, time1, time2)
+            self.d_k
+        )  # (batch, head, time1, time2)
 
         return self.forward_attention(v, scores, mask), new_cache
 
 
-
 def subsequent_mask(
     size: int,
     device: torch.device = torch.device("cpu"),
@@ -510,14 +507,17 @@ def subsequent_chunk_mask(
         ret[i, start:ending] = True
     return ret
 
-def add_optional_chunk_mask(xs: torch.Tensor,
+
+def add_optional_chunk_mask(
+    xs: torch.Tensor,
     masks: torch.Tensor,
     use_dynamic_chunk: bool,
     use_dynamic_left_chunk: bool,
     decoding_chunk_size: int,
     static_chunk_size: int,
     num_decoding_left_chunks: int,
-                            enable_full_context: bool = True):
+    enable_full_context: bool = True,
+):
     """Apply optional mask for encoder.
 
     Args:
@@ -565,18 +565,17 @@ def add_optional_chunk_mask(xs: torch.Tensor,
                 chunk_size = chunk_size % 25 + 1
                 if use_dynamic_left_chunk:
                     max_left_chunks = (max_len - 1) // chunk_size
-                    num_left_chunks = torch.randint(0, max_left_chunks,
+                    num_left_chunks = torch.randint(0, max_left_chunks, (1,)).item()
-                                                    (1, )).item()
+        chunk_masks = subsequent_chunk_mask(
-        chunk_masks = subsequent_chunk_mask(xs.size(1), chunk_size,
+            xs.size(1), chunk_size, num_left_chunks, xs.device
-                                            num_left_chunks,
+        )  # (L, L)
-                                            xs.device)  # (L, L)
         chunk_masks = chunk_masks.unsqueeze(0)  # (1, L, L)
         chunk_masks = masks & chunk_masks  # (B, L, L)
     elif static_chunk_size > 0:
         num_left_chunks = num_decoding_left_chunks
-        chunk_masks = subsequent_chunk_mask(xs.size(1), static_chunk_size,
+        chunk_masks = subsequent_chunk_mask(
-                                            num_left_chunks,
+            xs.size(1), static_chunk_size, num_left_chunks, xs.device
-                                            xs.device)  # (L, L)
+        )  # (L, L)
         chunk_masks = chunk_masks.unsqueeze(0)  # (1, L, L)
         chunk_masks = masks & chunk_masks  # (B, L, L)
     else:
@@ -630,7 +629,8 @@ class ConformerEncoderLayer(nn.Module):
         if self.conv_module is not None:
             self.norm_conv = nn.LayerNorm(size, eps=1e-5)  # for the CNN module
             self.norm_final = nn.LayerNorm(
-                size, eps=1e-5)  # for the final output of the block
+                size, eps=1e-5
+            )  # for the final output of the block
         self.dropout = nn.Dropout(dropout_rate)
         self.size = size
         self.normalize_before = normalize_before
@@ -671,8 +671,7 @@ class ConformerEncoderLayer(nn.Module):
             residual = x
             if self.normalize_before:
                 x = self.norm_ff_macaron(x)
-            x = residual + self.ff_scale * self.dropout(
+            x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x))
-                self.feed_forward_macaron(x))
             if not self.normalize_before:
                 x = self.norm_ff_macaron(x)
 
@@ -680,8 +679,7 @@ class ConformerEncoderLayer(nn.Module):
         residual = x
         if self.normalize_before:
             x = self.norm_mha(x)
-        x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb,
+        x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, att_cache)
-                                              att_cache)
         x = residual + self.dropout(x_att)
         if not self.normalize_before:
             x = self.norm_mha(x)
@@ -714,7 +712,6 @@ class ConformerEncoderLayer(nn.Module):
         return x, mask, new_att_cache, new_cnn_cache
 
 
-
 class EspnetRelPositionalEncoding(torch.nn.Module):
     """Relative positional encoding module (new implementation).
 
@@ -770,8 +767,9 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
         pe = torch.cat([pe_positive, pe_negative], dim=1)
         self.pe = pe.to(device=x.device, dtype=x.dtype)
 
-    def forward(self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0) \
+    def forward(
-            -> Tuple[torch.Tensor, torch.Tensor]:
+        self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Add positional encoding.
 
         Args:
@@ -786,9 +784,9 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
         pos_emb = self.position_encoding(size=x.size(1), offset=offset)
         return self.dropout(x), self.dropout(pos_emb)
 
-    def position_encoding(self,
+    def position_encoding(
-                          offset: Union[int, torch.Tensor],
+        self, offset: Union[int, torch.Tensor], size: int
-                          size: int) -> torch.Tensor:
+    ) -> torch.Tensor:
         """For getting encoding in a streaming fashion
 
         Attention!!!!!
@@ -811,7 +809,6 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
         return pos_emb
 
 
-
 class LinearEmbed(torch.nn.Module):
     """Linear transform the input without subsampling
 
@@ -822,8 +819,9 @@ class LinearEmbed(torch.nn.Module):
 
     """
 
-    def __init__(self, idim: int, odim: int, dropout_rate: float,
+    def __init__(
-                 pos_enc_class: torch.nn.Module):
+        self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module
+    ):
         """Construct an linear object."""
         super().__init__()
         self.out = torch.nn.Sequential(
@@ -833,14 +831,13 @@ class LinearEmbed(torch.nn.Module):
         )
         self.pos_enc = pos_enc_class  # rel_pos_espnet
 
-    def position_encoding(self, offset: Union[int, torch.Tensor],
+    def position_encoding(
-                          size: int) -> torch.Tensor:
+        self, offset: Union[int, torch.Tensor], size: int
+    ) -> torch.Tensor:
         return self.pos_enc.position_encoding(offset, size)
 
     def forward(
-        self,
+        self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0
-        x: torch.Tensor,
-        offset: Union[int, torch.Tensor] = 0
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Input x.
 
@@ -894,15 +891,13 @@ def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
     """
     batch_size = lengths.size(0)
     max_len = max_len if max_len > 0 else lengths.max().item()
-    seq_range = torch.arange(0,
+    seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device)
-                             max_len,
-                             dtype=torch.int64,
-                             device=lengths.device)
     seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
     seq_length_expand = lengths.unsqueeze(-1)
     mask = seq_range_expand >= seq_length_expand
     return mask
 
+
 # https://github.com/FunAudioLLM/CosyVoice/blob/main/examples/magicdata-read/cosyvoice/conf/cosyvoice.yaml
 class ConformerEncoder(torch.nn.Module):
     """Conformer encoder module."""
@@ -917,8 +912,8 @@ class ConformerEncoder(torch.nn.Module):
         dropout_rate: float = 0.1,
         positional_dropout_rate: float = 0.1,
         attention_dropout_rate: float = 0.0,
-        input_layer: str = 'linear',
+        input_layer: str = "linear",
-        pos_enc_layer_type: str = 'rel_pos_espnet',
+        pos_enc_layer_type: str = "rel_pos_espnet",
         normalize_before: bool = True,
         static_chunk_size: int = 1,  # 1: causal_mask; 0: full_mask
         use_dynamic_chunk: bool = False,
@@ -953,8 +948,12 @@ class ConformerEncoder(torch.nn.Module):
         """
         super().__init__()
         self.output_size = output_size
-        self.embed = LinearEmbed(input_size, output_size, dropout_rate, 
+        self.embed = LinearEmbed(
-                                        EspnetRelPositionalEncoding(output_size, positional_dropout_rate))
+            input_size,
+            output_size,
+            dropout_rate,
+            EspnetRelPositionalEncoding(output_size, positional_dropout_rate),
+        )
         self.normalize_before = normalize_before
         self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5)
         self.gradient_checkpointing = gradient_checkpointing
@@ -980,40 +979,60 @@ class ConformerEncoder(torch.nn.Module):
             activation,
         )
         # convolution module definition
-        convolution_layer_args = (output_size, cnn_module_kernel, activation,
+        convolution_layer_args = (
-                                  cnn_module_norm, causal)
+            output_size,
+            cnn_module_kernel,
+            activation,
+            cnn_module_norm,
+            causal,
+        )
 
-        self.encoders = torch.nn.ModuleList([
+        self.encoders = torch.nn.ModuleList(
+            [
                 ConformerEncoderLayer(
                     output_size,
-                RelPositionMultiHeadedAttention(
+                    RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args),
-                    *encoder_selfattn_layer_args),
                     PositionwiseFeedForward(*positionwise_layer_args),
-                PositionwiseFeedForward(
+                    (
-                    *positionwise_layer_args) if macaron_style else None,
+                        PositionwiseFeedForward(*positionwise_layer_args)
-                ConvolutionModule(
+                        if macaron_style
-                    *convolution_layer_args) if use_cnn_module else None,
+                        else None
+                    ),
+                    (
+                        ConvolutionModule(*convolution_layer_args)
+                        if use_cnn_module
+                        else None
+                    ),
                     dropout_rate,
                     normalize_before,
-            ) for _ in range(num_blocks)
+                )
-        ])
+                for _ in range(num_blocks)
+            ]
+        )
 
-    def forward_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor,
+    def forward_layers(
+        self,
+        xs: torch.Tensor,
+        chunk_masks: torch.Tensor,
         pos_emb: torch.Tensor,
-        mask_pad: torch.Tensor) -> torch.Tensor:
+        mask_pad: torch.Tensor,
+    ) -> torch.Tensor:
         for layer in self.encoders:
             xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad)
         return xs
 
     @torch.jit.unused
-    def forward_layers_checkpointed(self, xs: torch.Tensor,
+    def forward_layers_checkpointed(
+        self,
+        xs: torch.Tensor,
         chunk_masks: torch.Tensor,
         pos_emb: torch.Tensor,
-                                    mask_pad: torch.Tensor) -> torch.Tensor:
+        mask_pad: torch.Tensor,
+    ) -> torch.Tensor:
         for layer in self.encoders:
-            xs, chunk_masks, _, _ = ckpt.checkpoint(layer.__call__, xs,
+            xs, chunk_masks, _, _ = ckpt.checkpoint(
-                                                    chunk_masks, pos_emb,
+                layer.__call__, xs, chunk_masks, pos_emb, mask_pad
-                                                    mask_pad)
+            )
         return xs
 
     def forward(
@@ -1050,15 +1069,17 @@ class ConformerEncoder(torch.nn.Module):
         masks = pad_mask.to(torch.bool).unsqueeze(1)  # (B, 1, T)
         xs, pos_emb = self.embed(xs)
         mask_pad = masks  # (B, 1, T/subsample_rate)
-        chunk_masks = add_optional_chunk_mask(xs, masks,
+        chunk_masks = add_optional_chunk_mask(
+            xs,
+            masks,
             self.use_dynamic_chunk,
             self.use_dynamic_left_chunk,
             decoding_chunk_size,
             self.static_chunk_size,
-                                              num_decoding_left_chunks) 
+            num_decoding_left_chunks,
+        )
         if self.gradient_checkpointing and self.training:
-            xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb,
+            xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb, mask_pad)
-                                                  mask_pad)
         else:
             xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad)
         if self.normalize_before:
@@ -1067,4 +1088,3 @@ class ConformerEncoder(torch.nn.Module):
         # return the masks before encoder layers, and the masks will be used
         # for cross attention with decoder later
         return xs, masks
-

acestep/models/lyrics_utils/lyric_normalizer.py

@@ -2,39 +2,42 @@ import re
 from opencc import OpenCC
 
 
-t2s_converter = OpenCC('t2s')
+t2s_converter = OpenCC("t2s")
-s2t_converter = OpenCC('s2t')
+s2t_converter = OpenCC("s2t")
 
 
 EMOJI_PATTERN = re.compile(
     "["
-    "\U0001F600-\U0001F64F"  # Emoticons
+    "\U0001f600-\U0001f64f"  # Emoticons
-    "]+", flags=re.UNICODE
+    "]+",
+    flags=re.UNICODE,
 )
 
 # 创建一个翻译表，用于替换和移除字符
-TRANSLATION_TABLE = str.maketrans({
+TRANSLATION_TABLE = str.maketrans(
-    '-': ' ',  # 将 '-' 替换为空格
+    {
-    ',': None,
+        "-": " ",  # 将 '-' 替换为空格
-    '.': None,
+        ",": None,
-    '，': None,
+        ".": None,
-    '。': None,
+        "，": None,
-    '!': None,
+        "。": None,
-    '！': None,
+        "!": None,
-    '?': None,
+        "！": None,
-    '？': None,
+        "?": None,
-    '…': None,
+        "？": None,
-    ';': None,
+        "…": None,
-    '；': None,
+        ";": None,
-    ':': None,
+        "；": None,
-    '：': None,
+        ":": None,
-    '\u3000': ' ',  # 将全角空格替换为空格
+        "：": None,
-})
+        "\u3000": " ",  # 将全角空格替换为空格
+    }
+)
 
 # 替换括号中的内容，包括中括号和小括号
-BACKSLASH_PATTERN = re.compile(r'\(.*?\)|\[.*?\]')
+BACKSLASH_PATTERN = re.compile(r"\(.*?\)|\[.*?\]")
 
-SPACE_PATTERN = re.compile('(?<!^)\s+(?!$)')
+SPACE_PATTERN = re.compile("(?<!^)\s+(?!$)")
 
 
 def normalize_text(text, language, strip=True):
@@ -45,10 +48,10 @@ def normalize_text(text, language, strip=True):
     text = text.translate(TRANSLATION_TABLE)
 
     # Step 2: 移除表情符号
-    text = EMOJI_PATTERN.sub('', text)
+    text = EMOJI_PATTERN.sub("", text)
 
     # Step 3: 连续空白字符替换为单个空格，首位除外
-    text = SPACE_PATTERN.sub(' ', text)
+    text = SPACE_PATTERN.sub(" ", text)
 
     # Step 4: 去除首尾空白字符（如果需要）
     if strip:

acestep/models/lyrics_utils/lyric_tokenizer.py

@@ -446,7 +446,9 @@ _ordinal_re = {
     "it": re.compile(r"([0-9]+)(º|°|ª|o|a|i|e)"),
     "pl": re.compile(r"([0-9]+)(º|ª|st|nd|rd|th)"),
     "ar": re.compile(r"([0-9]+)(ون|ين|ث|ر|ى)"),
-    "cs": re.compile(r"([0-9]+)\.(?=\s|$)"),  # In Czech, a dot is often used after the number to indicate ordinals.
+    "cs": re.compile(
+        r"([0-9]+)\.(?=\s|$)"
+    ),  # In Czech, a dot is often used after the number to indicate ordinals.
     "ru": re.compile(r"([0-9]+)(-й|-я|-е|-ое|-ье|-го)"),
     "nl": re.compile(r"([0-9]+)(de|ste|e)"),
     "tr": re.compile(r"([0-9]+)(\.|inci|nci|uncu|üncü|\.)"),
@@ -486,7 +488,9 @@ def _expand_decimal_point(m, lang="en"):
 
 def _expand_currency(m, lang="en", currency="USD"):
     amount = float((re.sub(r"[^\d.]", "", m.group(0).replace(",", "."))))
-    full_amount = num2words(amount, to="currency", currency=currency, lang=lang if lang != "cs" else "cz")
+    full_amount = num2words(
+        amount, to="currency", currency=currency, lang=lang if lang != "cs" else "cz"
+    )
 
     and_equivalents = {
         "en": ", ",
@@ -530,13 +534,21 @@ def expand_numbers_multilingual(text, lang="en"):
         else:
             text = re.sub(_dot_number_re, _remove_dots, text)
         try:
-            text = re.sub(_currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text)
+            text = re.sub(
-            text = re.sub(_currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text)
+                _currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text
-            text = re.sub(_currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text)
+            )
+            text = re.sub(
+                _currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text
+            )
+            text = re.sub(
+                _currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text
+            )
         except:
             pass
         if lang != "tr":
-            text = re.sub(_decimal_number_re, lambda m: _expand_decimal_point(m, lang), text)
+            text = re.sub(
+                _decimal_number_re, lambda m: _expand_decimal_point(m, lang), text
+            )
         text = re.sub(_ordinal_re[lang], lambda m: _expand_ordinal(m, lang), text)
         text = re.sub(_number_re, lambda m: _expand_number(m, lang), text)
     return text
@@ -582,7 +594,15 @@ def basic_cleaners(text):
 
 def chinese_transliterate(text):
     return "".join(
-        [p[0] for p in pypinyin.pinyin(text, style=pypinyin.Style.TONE3, heteronym=False, neutral_tone_with_five=True)]
+        [
+            p[0]
+            for p in pypinyin.pinyin(
+                text,
+                style=pypinyin.Style.TONE3,
+                heteronym=False,
+                neutral_tone_with_five=True,
+            )
+        ]
     )
 
 
@@ -597,7 +617,9 @@ def korean_transliterate(text):
     return r.translit(text)
 
 
-DEFAULT_VOCAB_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "vocab.json")
+DEFAULT_VOCAB_FILE = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)), "vocab.json"
+)
 
 
 class VoiceBpeTokenizer:
@@ -639,7 +661,23 @@ class VoiceBpeTokenizer:
         #     )
 
     def preprocess_text(self, txt, lang):
-        if lang in {"ar", "cs", "de", "en", "es", "fr", "hu", "it", "nl", "pl", "pt", "ru", "tr", "zh", "ko"}:
+        if lang in {
+            "ar",
+            "cs",
+            "de",
+            "en",
+            "es",
+            "fr",
+            "hu",
+            "it",
+            "nl",
+            "pl",
+            "pt",
+            "ru",
+            "tr",
+            "zh",
+            "ko",
+        }:
             txt = multilingual_cleaners(txt, lang)
             if lang == "zh":
                 txt = chinese_transliterate(txt)
@@ -672,11 +710,12 @@ class VoiceBpeTokenizer:
         # txt = txt.replace("[UNK]", "")
         return txt
 
-
     # copy from https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3936
     def batch_decode(
         self,
-        sequences: Union[List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"],
+        sequences: Union[
+            List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"
+        ],
         skip_special_tokens: bool = False,
     ) -> List[str]:
         """
@@ -693,10 +732,7 @@ class VoiceBpeTokenizer:
         Returns:
             `List[str]`: The list of decoded sentences.
         """
-        return [
+        return [self.decode(seq) for seq in sequences]
-            self.decode(seq)
-            for seq in sequences
-        ]
 
     # https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/layers/xtts/trainer/dataset.py#L202
     # def pad(self):
@@ -716,15 +752,27 @@ def test_expand_numbers_multilingual():
         ("This is a 1st test", "This is a first test", "en"),
         ("That will be $20 sir.", "That will be twenty dollars sir.", "en"),
         ("That will be 20€ sir.", "That will be twenty euro sir.", "en"),
-        ("That will be 20.15€ sir.", "That will be twenty euro, fifteen cents sir.", "en"),
+        (
+            "That will be 20.15€ sir.",
+            "That will be twenty euro, fifteen cents sir.",
+            "en",
+        ),
         ("That's 100,000.5.", "That's one hundred thousand point five.", "en"),
         # French
         ("En 12,5 secondes.", "En douze virgule cinq secondes.", "fr"),
         ("Il y avait 50 soldats.", "Il y avait cinquante soldats.", "fr"),
         ("Ceci est un 1er test", "Ceci est un premier test", "fr"),
-        ("Cela vous fera $20 monsieur.", "Cela vous fera vingt dollars monsieur.", "fr"),
+        (
+            "Cela vous fera $20 monsieur.",
+            "Cela vous fera vingt dollars monsieur.",
+            "fr",
+        ),
         ("Cela vous fera 20€ monsieur.", "Cela vous fera vingt euros monsieur.", "fr"),
-        ("Cela vous fera 20,15€ monsieur.", "Cela vous fera vingt euros et quinze centimes monsieur.", "fr"),
+        (
+            "Cela vous fera 20,15€ monsieur.",
+            "Cela vous fera vingt euros et quinze centimes monsieur.",
+            "fr",
+        ),
         ("Ce sera 100.000,5.", "Ce sera cent mille virgule cinq.", "fr"),
         # German
         ("In 12,5 Sekunden.", "In zwölf Komma fünf Sekunden.", "de"),
@@ -732,21 +780,33 @@ def test_expand_numbers_multilingual():
         ("Dies ist ein 1. Test", "Dies ist ein erste Test", "de"),  # Issue with gender
         ("Das macht $20 Herr.", "Das macht zwanzig Dollar Herr.", "de"),
         ("Das macht 20€ Herr.", "Das macht zwanzig Euro Herr.", "de"),
-        ("Das macht 20,15€ Herr.", "Das macht zwanzig Euro und fünfzehn Cent Herr.", "de"),
+        (
+            "Das macht 20,15€ Herr.",
+            "Das macht zwanzig Euro und fünfzehn Cent Herr.",
+            "de",
+        ),
         # Spanish
         ("En 12,5 segundos.", "En doce punto cinco segundos.", "es"),
         ("Había 50 soldados.", "Había cincuenta soldados.", "es"),
         ("Este es un 1er test", "Este es un primero test", "es"),
         ("Eso le costará $20 señor.", "Eso le costará veinte dólares señor.", "es"),
         ("Eso le costará 20€ señor.", "Eso le costará veinte euros señor.", "es"),
-        ("Eso le costará 20,15€ señor.", "Eso le costará veinte euros con quince céntimos señor.", "es"),
+        (
+            "Eso le costará 20,15€ señor.",
+            "Eso le costará veinte euros con quince céntimos señor.",
+            "es",
+        ),
         # Italian
         ("In 12,5 secondi.", "In dodici virgola cinque secondi.", "it"),
         ("C'erano 50 soldati.", "C'erano cinquanta soldati.", "it"),
         ("Questo è un 1° test", "Questo è un primo test", "it"),
         ("Ti costerà $20 signore.", "Ti costerà venti dollari signore.", "it"),
         ("Ti costerà 20€ signore.", "Ti costerà venti euro signore.", "it"),
-        ("Ti costerà 20,15€ signore.", "Ti costerà venti euro e quindici centesimi signore.", "it"),
+        (
+            "Ti costerà 20,15€ signore.",
+            "Ti costerà venti euro e quindici centesimi signore.",
+            "it",
+        ),
         # Portuguese
         ("Em 12,5 segundos.", "Em doze vírgula cinco segundos.", "pt"),
         ("Havia 50 soldados.", "Havia cinquenta soldados.", "pt"),
@@ -761,8 +821,16 @@ def test_expand_numbers_multilingual():
         # Polish
         ("W 12,5 sekundy.", "W dwanaście przecinek pięć sekundy.", "pl"),
         ("Było 50 żołnierzy.", "Było pięćdziesiąt żołnierzy.", "pl"),
-        ("To będzie kosztować 20€ panie.", "To będzie kosztować dwadzieścia euro panie.", "pl"),
+        (
-        ("To będzie kosztować 20,15€ panie.", "To będzie kosztować dwadzieścia euro, piętnaście centów panie.", "pl"),
+            "To będzie kosztować 20€ panie.",
+            "To będzie kosztować dwadzieścia euro panie.",
+            "pl",
+        ),
+        (
+            "To będzie kosztować 20,15€ panie.",
+            "To będzie kosztować dwadzieścia euro, piętnaście centów panie.",
+            "pl",
+        ),
         # Arabic
         ("في الـ 12,5 ثانية.", "في الـ اثنا عشر  , خمسون ثانية.", "ar"),
         ("كان هناك 50 جنديًا.", "كان هناك خمسون جنديًا.", "ar"),
@@ -776,8 +844,16 @@ def test_expand_numbers_multilingual():
         # Russian
         ("Через 12.5 секунды.", "Через двенадцать запятая пять секунды.", "ru"),
         ("Там было 50 солдат.", "Там было пятьдесят солдат.", "ru"),
-        ("Это будет 20.15€ сэр.", "Это будет двадцать евро, пятнадцать центов сэр.", "ru"),
+        (
-        ("Это будет стоить 20€ господин.", "Это будет стоить двадцать евро господин.", "ru"),
+            "Это будет 20.15€ сэр.",
+            "Это будет двадцать евро, пятнадцать центов сэр.",
+            "ru",
+        ),
+        (
+            "Это будет стоить 20€ господин.",
+            "Это будет стоить двадцать евро господин.",
+            "ru",
+        ),
         # Dutch
         ("In 12,5 seconden.", "In twaalf komma vijf seconden.", "nl"),
         ("Er waren 50 soldaten.", "Er waren vijftig soldaten.", "nl"),
@@ -817,18 +893,30 @@ def test_abbreviations_multilingual():
         ("La Dra. Martinez es muy buena.", "La doctora Martinez es muy buena.", "es"),
         # French
         ("Bonjour Mr. Dupond.", "Bonjour monsieur Dupond.", "fr"),
-        ("Mme. Moreau est absente aujourd'hui.", "madame Moreau est absente aujourd'hui.", "fr"),
+        (
+            "Mme. Moreau est absente aujourd'hui.",
+            "madame Moreau est absente aujourd'hui.",
+            "fr",
+        ),
         # German
         ("Frau Dr. Müller ist sehr klug.", "Frau doktor Müller ist sehr klug.", "de"),
         # Portuguese
         ("Olá Sr. Silva.", "Olá senhor Silva.", "pt"),
-        ("Dra. Costa, você está disponível?", "doutora Costa, você está disponível?", "pt"),
+        (
+            "Dra. Costa, você está disponível?",
+            "doutora Costa, você está disponível?",
+            "pt",
+        ),
         # Italian
         ("Buongiorno, Sig. Rossi.", "Buongiorno, signore Rossi.", "it"),
         # ("Sig.ra Bianchi, posso aiutarti?", 'signora Bianchi, posso aiutarti?', 'it'), # Issue with matching that pattern
         # Polish
         ("Dzień dobry, P. Kowalski.", "Dzień dobry, pani Kowalski.", "pl"),
-        ("M. Nowak, czy mogę zadać pytanie?", "pan Nowak, czy mogę zadać pytanie?", "pl"),
+        (
+            "M. Nowak, czy mogę zadać pytanie?",
+            "pan Nowak, czy mogę zadać pytanie?",
+            "pl",
+        ),
         # Czech
         ("P. Novák", "pan Novák", "cs"),
         ("Dr. Vojtěch", "doktor Vojtěch", "cs"),
@@ -837,7 +925,11 @@ def test_abbreviations_multilingual():
         ("Mevr. de Vries", "mevrouw de Vries", "nl"),
         # Russian
         ("Здравствуйте Г-н Иванов.", "Здравствуйте господин Иванов.", "ru"),
-        ("Д-р Смирнов здесь, чтобы увидеть вас.", "доктор Смирнов здесь, чтобы увидеть вас.", "ru"),
+        (
+            "Д-р Смирнов здесь, чтобы увидеть вас.",
+            "доктор Смирнов здесь, чтобы увидеть вас.",
+            "ru",
+        ),
         # Turkish
         ("Merhaba B. Yılmaz.", "Merhaba bay Yılmaz.", "tr"),
         ("Dr. Ayşe burada.", "doktor Ayşe burada.", "tr"),
@@ -856,8 +948,16 @@ def test_symbols_multilingual():
         ("Te veo @ la fiesta", "Te veo arroba la fiesta", "es"),
         ("J'ai 14° de fièvre", "J'ai 14 degrés de fièvre", "fr"),
         ("Die Rechnung beträgt £ 20", "Die Rechnung beträgt pfund 20", "de"),
-        ("O meu email é ana&joao@gmail.com", "O meu email é ana e joao arroba gmail.com", "pt"),
+        (
-        ("linguaggio di programmazione C#", "linguaggio di programmazione C cancelletto", "it"),
+            "O meu email é ana&joao@gmail.com",
+            "O meu email é ana e joao arroba gmail.com",
+            "pt",
+        ),
+        (
+            "linguaggio di programmazione C#",
+            "linguaggio di programmazione C cancelletto",
+            "it",
+        ),
         ("Moja temperatura to 36.6°", "Moja temperatura to 36.6 stopnie", "pl"),
         ("Mám 14% baterie", "Mám 14 procento baterie", "cs"),
         ("Těším se na tebe @ party", "Těším se na tebe na party", "cs"),
@@ -868,7 +968,11 @@ def test_symbols_multilingual():
         ("لدي 14% في البطارية", "لدي 14 في المئة في البطارية", "ar"),
         ("我的电量为 14%", "我的电量为 14 百分之", "zh"),
         ("Pilim %14 dolu.", "Pilim yüzde 14 dolu.", "tr"),
-        ("Az akkumulátorom töltöttsége 14%", "Az akkumulátorom töltöttsége 14 százalék", "hu"),
+        (
+            "Az akkumulátorom töltöttsége 14%",
+            "Az akkumulátorom töltöttsége 14 százalék",
+            "hu",
+        ),
         ("배터리 잔량이 14%입니다.", "배터리 잔량이 14 퍼센트입니다.", "ko"),
     ]
 

acestep/music_dcae/music_dcae_pipeline.py

@@ -1,3 +1,11 @@
+"""
+ACE-Step: A Step Towards Music Generation Foundation Model
+
+https://github.com/ace-step/ACE-Step
+
+Apache 2.0 License
+"""
+
 import os
 import torch
 from diffusers import AutoencoderDC
@@ -21,7 +29,12 @@ VOCODER_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_vocoder")
 
 class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
     @register_to_config
-    def __init__(self, source_sample_rate=None, dcae_checkpoint_path=DEFAULT_PRETRAINED_PATH, vocoder_checkpoint_path=VOCODER_PRETRAINED_PATH):
+    def __init__(
+        self,
+        source_sample_rate=None,
+        dcae_checkpoint_path=DEFAULT_PRETRAINED_PATH,
+        vocoder_checkpoint_path=VOCODER_PRETRAINED_PATH,
+    ):
         super(MusicDCAE, self).__init__()
 
         self.dcae = AutoencoderDC.from_pretrained(dcae_checkpoint_path)
@@ -32,9 +45,11 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
 
         self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
 
-        self.transform = transforms.Compose([
+        self.transform = transforms.Compose(
+            [
                 transforms.Normalize(0.5, 0.5),
-        ])
+            ]
+        )
         self.min_mel_value = -11.0
         self.max_mel_value = 3.0
         self.audio_chunk_size = int(round((1024 * 512 / 44100 * 48000)))
@@ -76,7 +91,9 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
 
         max_audio_len = audio.shape[-1]
         if max_audio_len % (8 * 512) != 0:
-            audio = torch.nn.functional.pad(audio, (0, 8 * 512 - max_audio_len % (8 * 512)))
+            audio = torch.nn.functional.pad(
+                audio, (0, 8 * 512 - max_audio_len % (8 * 512))
+            )
 
         mels = self.forward_mel(audio)
         mels = (mels - self.min_mel_value) / (self.max_mel_value - self.min_mel_value)
@@ -86,7 +103,9 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
             latent = self.dcae.encoder(mel.unsqueeze(0))
             latents.append(latent)
         latents = torch.cat(latents, dim=0)
-        latent_lengths = (audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple).long()
+        latent_lengths = (
+            audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple
+        ).long()
         latents = (latents - self.shift_factor) * self.scale_factor
         return latents, latent_lengths
 
@@ -103,18 +122,26 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
             wav = self.vocoder.decode(mels[0]).squeeze(1)
 
             if sr is not None:
-                resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
+                resampler = (
+                    torchaudio.transforms.Resample(44100, sr)
+                    .to(latents.device)
+                    .to(latents.dtype)
+                )
                 wav = resampler(wav)
             else:
                 sr = 44100
             pred_wavs.append(wav)
 
         if audio_lengths is not None:
-            pred_wavs = [wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)]
+            pred_wavs = [
+                wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)
+            ]
         return sr, pred_wavs
 
     def forward(self, audios, audio_lengths=None, sr=None):
-        latents, latent_lengths = self.encode(audios=audios, audio_lengths=audio_lengths, sr=sr)
+        latents, latent_lengths = self.encode(
+            audios=audios, audio_lengths=audio_lengths, sr=sr
+        )
         sr, pred_wavs = self.decode(latents=latents, audio_lengths=audio_lengths, sr=sr)
         return sr, pred_wavs, latents, latent_lengths
 

acestep/music_dcae/music_log_mel.py

@@ -1,3 +1,11 @@
+"""
+ACE-Step: A Step Towards Music Generation Foundation Model
+
+https://github.com/ace-step/ACE-Step
+
+Apache 2.0 License
+"""
+
 import torch
 import torch.nn as nn
 from torch import Tensor

acestep/music_dcae/music_vocoder.py

@@ -1,3 +1,11 @@
+"""
+ACE-Step: A Step Towards Music Generation Foundation Model
+
+https://github.com/ace-step/ACE-Step
+
+Apache 2.0 License
+"""
+
 import librosa
 import torch
 from torch import nn
@@ -132,13 +140,11 @@ class ConvNeXtBlock(nn.Module):
         self.act = nn.GELU()
         self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)
         self.gamma = (
-            nn.Parameter(layer_scale_init_value *
+            nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
-                         torch.ones((dim)), requires_grad=True)
             if layer_scale_init_value > 0
             else None
         )
-        self.drop_path = DropPath(
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
-            drop_path) if drop_path > 0.0 else nn.Identity()
 
     def forward(self, x, apply_residual: bool = True):
         input = x
@@ -367,8 +373,7 @@ class HiFiGANGenerator(nn.Module):
         upsample_rates: Tuple[int] = (8, 8, 2, 2, 2),
         upsample_kernel_sizes: Tuple[int] = (16, 16, 8, 2, 2),
         resblock_kernel_sizes: Tuple[int] = (3, 7, 11),
-        resblock_dilation_sizes: Tuple[Tuple[int]] = (
+        resblock_dilation_sizes: Tuple[Tuple[int]] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)),
-            (1, 3, 5), (1, 3, 5), (1, 3, 5)),
         num_mels: int = 128,
         upsample_initial_channel: int = 512,
         use_template: bool = True,
@@ -498,7 +503,11 @@ class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
         upsample_kernel_sizes: Tuple[int] = (8, 8, 4, 4, 4, 4, 4),
         resblock_kernel_sizes: Tuple[int] = (3, 7, 11, 13),
         resblock_dilation_sizes: Tuple[Tuple[int]] = (
-            (1, 3, 5), (1, 3, 5), (1, 3, 5), (1, 3, 5)),
+            (1, 3, 5),
+            (1, 3, 5),
+            (1, 3, 5),
+            (1, 3, 5),
+        ),
         num_mels: int = 512,
         upsample_initial_channel: int = 1024,
         use_template: bool = False,
@@ -566,7 +575,9 @@ if __name__ == "__main__":
     import soundfile as sf
 
     x = "test_audio.flac"
-    model = ADaMoSHiFiGANV1.from_pretrained("./checkpoints/music_vocoder", local_files_only=True)
+    model = ADaMoSHiFiGANV1.from_pretrained(
+        "./checkpoints/music_vocoder", local_files_only=True
+    )
 
     wav, sr = librosa.load(x, sr=44100, mono=True)
     wav = torch.from_numpy(wav).float()[None]

acestep/schedulers/scheduling_flow_match_euler_discrete.py

@@ -72,7 +72,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
         base_image_seq_len: Optional[int] = 256,
         max_image_seq_len: Optional[int] = 4096,
     ):
-        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        timesteps = np.linspace(
+            1, num_train_timesteps, num_train_timesteps, dtype=np.float32
+        )[::-1].copy()
         timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
 
         sigmas = timesteps / num_train_timesteps
@@ -146,7 +148,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
 
         # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
         if self.begin_index is None:
-            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps) for t in timestep
+            ]
         elif self.step_index is not None:
             # add_noise is called after first denoising step (for inpainting)
             step_indices = [self.step_index] * timestep.shape[0]
@@ -186,12 +190,16 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
         """
 
         if self.config.use_dynamic_shifting and mu is None:
-            raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`")
+            raise ValueError(
+                " you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+            )
 
         if sigmas is None:
             self.num_inference_steps = num_inference_steps
             timesteps = np.linspace(
-                self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+                self._sigma_to_t(self.sigma_max),
+                self._sigma_to_t(self.sigma_min),
+                num_inference_steps,
             )
 
             sigmas = timesteps / self.config.num_train_timesteps
@@ -243,7 +251,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
         s_noise: float = 1.0,
         generator: Optional[torch.Generator] = None,
         return_dict: bool = True,
-        omega: Union[float, np.array] = 0.0
+        omega: Union[float, np.array] = 0.0,
     ) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion

acestep/schedulers/scheduling_flow_match_heun_discrete.py

@@ -67,7 +67,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
         num_train_timesteps: int = 1000,
         shift: float = 1.0,
     ):
-        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        timesteps = np.linspace(
+            1, num_train_timesteps, num_train_timesteps, dtype=np.float32
+        )[::-1].copy()
         timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
 
         sigmas = timesteps / num_train_timesteps
@@ -137,7 +139,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
     def _sigma_to_t(self, sigma):
         return sigma * self.config.num_train_timesteps
 
-    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+    def set_timesteps(
+        self, num_inference_steps: int, device: Union[str, torch.device] = None
+    ):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
 
@@ -150,7 +154,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
         self.num_inference_steps = num_inference_steps
 
         timesteps = np.linspace(
-            self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+            self._sigma_to_t(self.sigma_max),
+            self._sigma_to_t(self.sigma_min),
+            num_inference_steps,
         )
 
         sigmas = timesteps / self.config.num_train_timesteps
@@ -162,7 +168,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
         self.timesteps = timesteps.to(device=device)
 
         sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
-        self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
+        self.sigmas = torch.cat(
+            [sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]]
+        )
 
         # empty dt and derivative
         self.prev_derivative = None
@@ -208,7 +216,7 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
         s_noise: float = 1.0,
         generator: Optional[torch.Generator] = None,
         return_dict: bool = True,
-        omega: Union[float, np.array] = 0.0
+        omega: Union[float, np.array] = 0.0,
     ) -> Union[FlowMatchHeunDiscreteSchedulerOutput, Tuple]:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
@@ -285,13 +293,20 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
             sigma = self.sigmas[self.step_index - 1]
             sigma_next = self.sigmas[self.step_index]
 
-        gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0
+        gamma = (
+            min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1)
+            if s_tmin <= sigma <= s_tmax
+            else 0.0
+        )
 
         sigma_hat = sigma * (gamma + 1)
 
         if gamma > 0:
             noise = randn_tensor(
-                model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
+                model_output.shape,
+                dtype=model_output.dtype,
+                device=model_output.device,
+                generator=generator,
             )
             eps = noise * s_noise
             sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5

acestep/text2music_dataset.py

@@ -17,6 +17,7 @@ warnings.simplefilter("ignore", category=FutureWarning)
 
 DEFAULT_TRAIN_PATH = "./data/example_dataset"
 
+
 def is_silent_audio(audio_tensor, silence_threshold=0.95):
     """
     Determine if an audio is silent and should be discarded
@@ -39,10 +40,23 @@ def is_silent_audio(audio_tensor, silence_threshold=0.95):
 
 # Supported languages for tokenization
 SUPPORT_LANGUAGES = {
-    "en": 259, "de": 260, "fr": 262, "es": 284, "it": 285, 
+    "en": 259,
-    "pt": 286, "pl": 294, "tr": 295, "ru": 267, "cs": 293, 
+    "de": 260,
-    "nl": 297, "ar": 5022, "zh": 5023, "ja": 5412, "hu": 5753,
+    "fr": 262,
-    "ko": 6152, "hi": 6680
+    "es": 284,
+    "it": 285,
+    "pt": 286,
+    "pl": 294,
+    "tr": 295,
+    "ru": 267,
+    "cs": 293,
+    "nl": 297,
+    "ar": 5022,
+    "zh": 5023,
+    "ja": 5412,
+    "hu": 5753,
+    "ko": 6152,
+    "hi": 6680,
 }
 
 # Regex pattern for structure markers like [Verse], [Chorus], etc.
@@ -54,9 +68,15 @@ class Text2MusicDataset(Dataset):
     Dataset for text-to-music generation that processes lyrics and audio files
     """
 
-    def __init__(self, train=True, train_dataset_path=DEFAULT_TRAIN_PATH, 
+    def __init__(
-                 max_duration=240.0, sample_size=None, shuffle=True, 
+        self,
-                 minibatch_size=1):
+        train=True,
+        train_dataset_path=DEFAULT_TRAIN_PATH,
+        max_duration=240.0,
+        sample_size=None,
+        shuffle=True,
+        minibatch_size=1,
+    ):
         """
         Initialize the Text2Music dataset
 
@@ -75,14 +95,107 @@ class Text2MusicDataset(Dataset):
 
         # Initialize language segmentation
         self.lang_segment = LangSegment()
-        self.lang_segment.setfilters([
+        self.lang_segment.setfilters(
-            'af', 'am', 'an', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'br', 'bs', 'ca', 'cs', 'cy', 'da', 'de', 'dz', 'el',
+            [
-            'en', 'eo', 'es', 'et', 'eu', 'fa', 'fi', 'fo', 'fr', 'ga', 'gl', 'gu', 'he', 'hi', 'hr', 'ht', 'hu', 'hy',
+                "af",
-            'id', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'ku', 'ky', 'la', 'lb', 'lo', 'lt', 'lv', 'mg',
+                "am",
-            'mk', 'ml', 'mn', 'mr', 'ms', 'mt', 'nb', 'ne', 'nl', 'nn', 'no', 'oc', 'or', 'pa', 'pl', 'ps', 'pt', 'qu',
+                "an",
-            'ro', 'ru', 'rw', 'se', 'si', 'sk', 'sl', 'sq', 'sr', 'sv', 'sw', 'ta', 'te', 'th', 'tl', 'tr', 'ug', 'uk',
+                "ar",
-            'ur', 'vi', 'vo', 'wa', 'xh', 'zh', 'zu'
+                "as",
-        ])
+                "az",
+                "be",
+                "bg",
+                "bn",
+                "br",
+                "bs",
+                "ca",
+                "cs",
+                "cy",
+                "da",
+                "de",
+                "dz",
+                "el",
+                "en",
+                "eo",
+                "es",
+                "et",
+                "eu",
+                "fa",
+                "fi",
+                "fo",
+                "fr",
+                "ga",
+                "gl",
+                "gu",
+                "he",
+                "hi",
+                "hr",
+                "ht",
+                "hu",
+                "hy",
+                "id",
+                "is",
+                "it",
+                "ja",
+                "jv",
+                "ka",
+                "kk",
+                "km",
+                "kn",
+                "ko",
+                "ku",
+                "ky",
+                "la",
+                "lb",
+                "lo",
+                "lt",
+                "lv",
+                "mg",
+                "mk",
+                "ml",
+                "mn",
+                "mr",
+                "ms",
+                "mt",
+                "nb",
+                "ne",
+                "nl",
+                "nn",
+                "no",
+                "oc",
+                "or",
+                "pa",
+                "pl",
+                "ps",
+                "pt",
+                "qu",
+                "ro",
+                "ru",
+                "rw",
+                "se",
+                "si",
+                "sk",
+                "sl",
+                "sq",
+                "sr",
+                "sv",
+                "sw",
+                "ta",
+                "te",
+                "th",
+                "tl",
+                "tr",
+                "ug",
+                "uk",
+                "ur",
+                "vi",
+                "vo",
+                "wa",
+                "xh",
+                "zh",
+                "zu",
+            ]
+        )
 
         # Initialize lyric tokenizer
         self.lyric_tokenizer = VoiceBpeTokenizer()
@@ -196,22 +309,30 @@ class Text2MusicDataset(Dataset):
 
                         # If debug mode, show tokenization results
                         if debug:
-                            toks = self.lyric_tokenizer.batch_decode([[tok_id] for tok_id in token_idx])
+                            toks = self.lyric_tokenizer.batch_decode(
-                            logger.info(f"debug using most_common_lang {line} --> {most_common_lang} --> {toks}")
+                                [[tok_id] for tok_id in token_idx]
+                            )
+                            logger.info(
+                                f"debug using most_common_lang {line} --> {most_common_lang} --> {toks}"
+                            )
 
                         # If tokenization contains unknown token (1), try with segment language
                         if 1 in token_idx:
                             token_idx = self.lyric_tokenizer.encode(line, lang)
 
                     if debug:
-                        toks = self.lyric_tokenizer.batch_decode([[tok_id] for tok_id in token_idx])
+                        toks = self.lyric_tokenizer.batch_decode(
+                            [[tok_id] for tok_id in token_idx]
+                        )
                         logger.info(f"debug {line} --> {lang} --> {toks}")
 
                     # Add tokens and line break
                     lyric_token_idx = lyric_token_idx + token_idx + [2]
 
                 except Exception as e:
-                    logger.error(f"Tokenize error: {e} for line: {line}, major_language: {lang}")
+                    logger.error(
+                        f"Tokenize error: {e} for line: {line}, major_language: {lang}"
+                    )
 
         return lyric_token_idx
 
@@ -302,7 +423,9 @@ class Text2MusicDataset(Dataset):
 
         # Pad to minimum 3 seconds if needed
         if audio.shape[-1] < 48000 * 3:
-            audio = torch.nn.functional.pad(audio, (0, 48000 * 3 - audio.shape[-1]), 'constant', 0)
+            audio = torch.nn.functional.pad(
+                audio, (0, 48000 * 3 - audio.shape[-1]), "constant", 0
+            )
 
         # Check if audio is silent
         if is_silent_audio(audio):
@@ -368,9 +491,11 @@ class Text2MusicDataset(Dataset):
         lyrics = item["norm_lyrics"]
         lyrics_lines = lyrics.split("\n")
         for lyric_line in lyrics_lines:
-            candidate_lyric_chunk.append({
+            candidate_lyric_chunk.append(
+                {
                     "lyric": lyric_line,
-            })
+                }
+            )
 
         # Limit audio length
         longest_length = 24 * 10 * 48000  # 240 seconds
@@ -482,28 +607,43 @@ class Text2MusicDataset(Dataset):
             elif k in ["src_wavs", "target_wavs", "vocal_wavs"]:
                 # Pad audio to max length
                 max_length = max(seq.shape[1] for seq in v)
-                padded_input_list = torch.stack([
+                padded_input_list = torch.stack(
-                    torch.nn.functional.pad(seq, (0, max_length - seq.shape[1]), 'constant', 0) 
+                    [
+                        torch.nn.functional.pad(
+                            seq, (0, max_length - seq.shape[1]), "constant", 0
+                        )
                         for seq in v
-                ])
+                    ]
+                )
             elif k in ["clap_conditions"]:
                 # Pad time dimension of embeddings
                 max_length = max(seq.shape[0] for seq in v)
                 v = [
-                    torch.nn.functional.pad(seq, (0, 0, 0, max_length - seq.shape[0]), 'constant', 0) 
+                    torch.nn.functional.pad(
+                        seq, (0, 0, 0, max_length - seq.shape[0]), "constant", 0
+                    )
                     for seq in v
                 ]
                 padded_input_list = torch.stack(v)
             elif k == "speaker_embs":
                 # Stack speaker embeddings
                 padded_input_list = torch.stack(v)
-            elif k in ["chunk_masks", "clap_attention_masks", "lyric_token_ids", "lyric_masks"]:
+            elif k in [
+                "chunk_masks",
+                "clap_attention_masks",
+                "lyric_token_ids",
+                "lyric_masks",
+            ]:
                 # Pad sequence tensors
                 max_length = max(len(seq) for seq in v)
-                padded_input_list = torch.stack([
+                padded_input_list = torch.stack(
-                    torch.nn.functional.pad(seq, (0, max_length - len(seq)), 'constant', 0) 
+                    [
+                        torch.nn.functional.pad(
+                            seq, (0, max_length - len(seq)), "constant", 0
+                        )
                         for seq in v
-                ])
+                    ]
+                )
 
             output[k] = padded_input_list
 

acestep/ui/components.py

@@ -1,3 +1,11 @@
+"""
+ACE-Step: A Step Towards Music Generation Foundation Model
+
+https://github.com/ace-step/ACE-Step
+
+Apache 2.0 License
+"""
+
 import gradio as gr
 import librosa
 
@@ -41,8 +49,6 @@ In this moment we take flight
 """
 
 
-
-
 def create_output_ui(task_name="Text2Music"):
     # For many consumer-grade GPU devices, only one batch can be run
     output_audio1 = gr.Audio(type="filepath", label=f"{task_name} Generated Audio 1")
@@ -68,41 +74,162 @@ def create_text2music_ui(
         with gr.Column():
             with gr.Row(equal_height=True):
                 # add markdown, tags and lyrics examples are from ai music generation community
-                audio_duration = gr.Slider(-1, 240.0, step=0.00001, value=-1, label="Audio Duration", interactive=True, info="-1 means random duration (30 ~ 240).", scale=9)
+                audio_duration = gr.Slider(
+                    -1,
+                    240.0,
+                    step=0.00001,
+                    value=-1,
+                    label="Audio Duration",
+                    interactive=True,
+                    info="-1 means random duration (30 ~ 240).",
+                    scale=9,
+                )
                 sample_bnt = gr.Button("Sample", variant="primary", scale=1)
 
-            prompt = gr.Textbox(lines=2, label="Tags", max_lines=4, value=TAG_DEFAULT, info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community")
+            prompt = gr.Textbox(
-            lyrics = gr.Textbox(lines=9, label="Lyrics", max_lines=13, value=LYRIC_DEFAULT, info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics")
+                lines=2,
+                label="Tags",
+                max_lines=4,
+                value=TAG_DEFAULT,
+                info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community",
+            )
+            lyrics = gr.Textbox(
+                lines=9,
+                label="Lyrics",
+                max_lines=13,
+                value=LYRIC_DEFAULT,
+                info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics",
+            )
 
             with gr.Accordion("Basic Settings", open=False):
-                infer_step = gr.Slider(minimum=1, maximum=1000, step=1, value=27, label="Infer Steps", interactive=True)
+                infer_step = gr.Slider(
-                guidance_scale = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=15.0, label="Guidance Scale", interactive=True, info="When guidance_scale_lyric > 1 and guidance_scale_text > 1, the guidance scale will not be applied.")
+                    minimum=1,
-                guidance_scale_text = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Text", interactive=True, info="Guidance scale for text condition. It can only apply to cfg. set guidance_scale_text=5.0, guidance_scale_lyric=1.5 for start")
+                    maximum=1000,
-                guidance_scale_lyric = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Lyric", interactive=True)
+                    step=1,
+                    value=27,
+                    label="Infer Steps",
+                    interactive=True,
+                )
+                guidance_scale = gr.Slider(
+                    minimum=0.0,
+                    maximum=200.0,
+                    step=0.1,
+                    value=15.0,
+                    label="Guidance Scale",
+                    interactive=True,
+                    info="When guidance_scale_lyric > 1 and guidance_scale_text > 1, the guidance scale will not be applied.",
+                )
+                guidance_scale_text = gr.Slider(
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=0.0,
+                    label="Guidance Scale Text",
+                    interactive=True,
+                    info="Guidance scale for text condition. It can only apply to cfg. set guidance_scale_text=5.0, guidance_scale_lyric=1.5 for start",
+                )
+                guidance_scale_lyric = gr.Slider(
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=0.0,
+                    label="Guidance Scale Lyric",
+                    interactive=True,
+                )
 
-                manual_seeds = gr.Textbox(label="manual seeds (default None)", placeholder="1,2,3,4", value=None, info="Seed for the generation")
+                manual_seeds = gr.Textbox(
+                    label="manual seeds (default None)",
+                    placeholder="1,2,3,4",
+                    value=None,
+                    info="Seed for the generation",
+                )
 
             with gr.Accordion("Advanced Settings", open=False):
-                scheduler_type = gr.Radio(["euler", "heun"], value="euler", label="Scheduler Type", elem_id="scheduler_type", info="Scheduler type for the generation. euler is recommended. heun will take more time.")
+                scheduler_type = gr.Radio(
-                cfg_type = gr.Radio(["cfg", "apg", "cfg_star"], value="apg", label="CFG Type", elem_id="cfg_type", info="CFG type for the generation. apg is recommended. cfg and cfg_star are almost the same.")
+                    ["euler", "heun"],
-                use_erg_tag = gr.Checkbox(label="use ERG for tag", value=True, info="Use Entropy Rectifying Guidance for tag. It will multiple a temperature to the attention to make a weaker tag condition and make better diversity.")
+                    value="euler",
-                use_erg_lyric = gr.Checkbox(label="use ERG for lyric", value=True, info="The same but apply to lyric encoder's attention.")
+                    label="Scheduler Type",
-                use_erg_diffusion = gr.Checkbox(label="use ERG for diffusion", value=True, info="The same but apply to diffusion model's attention.")
+                    elem_id="scheduler_type",
+                    info="Scheduler type for the generation. euler is recommended. heun will take more time.",
+                )
+                cfg_type = gr.Radio(
+                    ["cfg", "apg", "cfg_star"],
+                    value="apg",
+                    label="CFG Type",
+                    elem_id="cfg_type",
+                    info="CFG type for the generation. apg is recommended. cfg and cfg_star are almost the same.",
+                )
+                use_erg_tag = gr.Checkbox(
+                    label="use ERG for tag",
+                    value=True,
+                    info="Use Entropy Rectifying Guidance for tag. It will multiple a temperature to the attention to make a weaker tag condition and make better diversity.",
+                )
+                use_erg_lyric = gr.Checkbox(
+                    label="use ERG for lyric",
+                    value=True,
+                    info="The same but apply to lyric encoder's attention.",
+                )
+                use_erg_diffusion = gr.Checkbox(
+                    label="use ERG for diffusion",
+                    value=True,
+                    info="The same but apply to diffusion model's attention.",
+                )
 
-                omega_scale = gr.Slider(minimum=-100.0, maximum=100.0, step=0.1, value=10.0, label="Granularity Scale", interactive=True, info="Granularity scale for the generation. Higher values can reduce artifacts")
+                omega_scale = gr.Slider(
+                    minimum=-100.0,
+                    maximum=100.0,
+                    step=0.1,
+                    value=10.0,
+                    label="Granularity Scale",
+                    interactive=True,
+                    info="Granularity scale for the generation. Higher values can reduce artifacts",
+                )
 
-                guidance_interval = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Guidance Interval", interactive=True, info="Guidance interval for the generation. 0.5 means only apply guidance in the middle steps (0.25 * infer_steps to 0.75 * infer_steps)")
+                guidance_interval = gr.Slider(
-                guidance_interval_decay = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.0, label="Guidance Interval Decay", interactive=True, info="Guidance interval decay for the generation. Guidance scale will decay from guidance_scale to min_guidance_scale in the interval. 0.0 means no decay.")
+                    minimum=0.0,
-                min_guidance_scale = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=3.0, label="Min Guidance Scale", interactive=True, info="Min guidance scale for guidance interval decay's end scale")
+                    maximum=1.0,
-                oss_steps = gr.Textbox(label="OSS Steps", placeholder="16, 29, 52, 96, 129, 158, 172, 183, 189, 200", value=None, info="Optimal Steps for the generation. But not test well")
+                    step=0.01,
+                    value=0.5,
+                    label="Guidance Interval",
+                    interactive=True,
+                    info="Guidance interval for the generation. 0.5 means only apply guidance in the middle steps (0.25 * infer_steps to 0.75 * infer_steps)",
+                )
+                guidance_interval_decay = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.0,
+                    label="Guidance Interval Decay",
+                    interactive=True,
+                    info="Guidance interval decay for the generation. Guidance scale will decay from guidance_scale to min_guidance_scale in the interval. 0.0 means no decay.",
+                )
+                min_guidance_scale = gr.Slider(
+                    minimum=0.0,
+                    maximum=200.0,
+                    step=0.1,
+                    value=3.0,
+                    label="Min Guidance Scale",
+                    interactive=True,
+                    info="Min guidance scale for guidance interval decay's end scale",
+                )
+                oss_steps = gr.Textbox(
+                    label="OSS Steps",
+                    placeholder="16, 29, 52, 96, 129, 158, 172, 183, 189, 200",
+                    value=None,
+                    info="Optimal Steps for the generation. But not test well",
+                )
 
             text2music_bnt = gr.Button("Generate", variant="primary")
 
         with gr.Column():
             outputs, input_params_json = create_output_ui()
             with gr.Tab("retake"):
-                retake_variance = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance")
+                retake_variance = gr.Slider(
-                retake_seeds = gr.Textbox(label="retake seeds (default None)", placeholder="", value=None)
+                    minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance"
+                )
+                retake_seeds = gr.Textbox(
+                    label="retake seeds (default None)", placeholder="", value=None
+                )
                 retake_bnt = gr.Button("Retake", variant="primary")
                 retake_outputs, retake_input_params_json = create_output_ui("Retake")
 
@@ -124,8 +251,16 @@ def create_text2music_ui(
                         json_data["use_erg_lyric"],
                         json_data["use_erg_diffusion"],
                         ", ".join(map(str, json_data["oss_steps"])),
-                        json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
+                        (
-                        json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
+                            json_data["guidance_scale_text"]
+                            if "guidance_scale_text" in json_data
+                            else 0.0
+                        ),
+                        (
+                            json_data["guidance_scale_lyric"]
+                            if "guidance_scale_lyric" in json_data
+                            else 0.0
+                        ),
                         retake_seeds=retake_seeds,
                         retake_variance=retake_variance,
                         task="retake",
@@ -141,15 +276,45 @@ def create_text2music_ui(
                     outputs=retake_outputs + [retake_input_params_json],
                 )
             with gr.Tab("repainting"):
-                retake_variance = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance")
+                retake_variance = gr.Slider(
-                retake_seeds = gr.Textbox(label="repaint seeds (default None)", placeholder="", value=None)
+                    minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance"
-                repaint_start = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=0.0, label="Repaint Start Time", interactive=True)
+                )
-                repaint_end = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=30.0, label="Repaint End Time", interactive=True)
+                retake_seeds = gr.Textbox(
-                repaint_source = gr.Radio(["text2music", "last_repaint", "upload"], value="text2music", label="Repaint Source", elem_id="repaint_source")
+                    label="repaint seeds (default None)", placeholder="", value=None
+                )
+                repaint_start = gr.Slider(
+                    minimum=0.0,
+                    maximum=240.0,
+                    step=0.01,
+                    value=0.0,
+                    label="Repaint Start Time",
+                    interactive=True,
+                )
+                repaint_end = gr.Slider(
+                    minimum=0.0,
+                    maximum=240.0,
+                    step=0.01,
+                    value=30.0,
+                    label="Repaint End Time",
+                    interactive=True,
+                )
+                repaint_source = gr.Radio(
+                    ["text2music", "last_repaint", "upload"],
+                    value="text2music",
+                    label="Repaint Source",
+                    elem_id="repaint_source",
+                )
 
-                repaint_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="repaint_source_audio_upload")
+                repaint_source_audio_upload = gr.Audio(
+                    label="Upload Audio",
+                    type="filepath",
+                    visible=False,
+                    elem_id="repaint_source_audio_upload",
+                )
                 repaint_source.change(
-                    fn=lambda x: gr.update(visible=x == "upload", elem_id="repaint_source_audio_upload"),
+                    fn=lambda x: gr.update(
+                        visible=x == "upload", elem_id="repaint_source_audio_upload"
+                    ),
                     inputs=[repaint_source],
                     outputs=[repaint_source_audio_upload],
                 )
@@ -187,9 +352,7 @@ def create_text2music_ui(
                     if repaint_source == "upload":
                         src_audio_path = repaint_source_audio_upload
                         audio_duration = librosa.get_duration(filename=src_audio_path)
-                        json_data = {
+                        json_data = {"audio_duration": audio_duration}
-                            "audio_duration": audio_duration
-                        }
                     elif repaint_source == "text2music":
                         json_data = text2music_json_data
                         src_audio_path = json_data["audio_path"]
@@ -258,11 +421,33 @@ def create_text2music_ui(
             with gr.Tab("edit"):
                 edit_prompt = gr.Textbox(lines=2, label="Edit Tags", max_lines=4)
                 edit_lyrics = gr.Textbox(lines=9, label="Edit Lyrics", max_lines=13)
-                retake_seeds = gr.Textbox(label="edit seeds (default None)", placeholder="", value=None)
+                retake_seeds = gr.Textbox(
+                    label="edit seeds (default None)", placeholder="", value=None
+                )
 
-                edit_type = gr.Radio(["only_lyrics", "remix"], value="only_lyrics", label="Edit Type", elem_id="edit_type", info="`only_lyrics` will keep the whole song the same except lyrics difference. Make your diffrence smaller, e.g. one lyrc line change.\nremix can change the song melody and genre")
+                edit_type = gr.Radio(
-                edit_n_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.6, label="edit_n_min", interactive=True)
+                    ["only_lyrics", "remix"],
-                edit_n_max = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="edit_n_max", interactive=True)
+                    value="only_lyrics",
+                    label="Edit Type",
+                    elem_id="edit_type",
+                    info="`only_lyrics` will keep the whole song the same except lyrics difference. Make your diffrence smaller, e.g. one lyrc line change.\nremix can change the song melody and genre",
+                )
+                edit_n_min = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.6,
+                    label="edit_n_min",
+                    interactive=True,
+                )
+                edit_n_max = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=1.0,
+                    label="edit_n_max",
+                    interactive=True,
+                )
 
                 def edit_type_change_func(edit_type):
                     if edit_type == "only_lyrics":
@@ -276,13 +461,25 @@ def create_text2music_ui(
                 edit_type.change(
                     edit_type_change_func,
                     inputs=[edit_type],
-                    outputs=[edit_n_min, edit_n_max]
+                    outputs=[edit_n_min, edit_n_max],
                 )
 
-                edit_source = gr.Radio(["text2music", "last_edit", "upload"], value="text2music", label="Edit Source", elem_id="edit_source")
+                edit_source = gr.Radio(
-                edit_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="edit_source_audio_upload")
+                    ["text2music", "last_edit", "upload"],
+                    value="text2music",
+                    label="Edit Source",
+                    elem_id="edit_source",
+                )
+                edit_source_audio_upload = gr.Audio(
+                    label="Upload Audio",
+                    type="filepath",
+                    visible=False,
+                    elem_id="edit_source_audio_upload",
+                )
                 edit_source.change(
-                    fn=lambda x: gr.update(visible=x == "upload", elem_id="edit_source_audio_upload"),
+                    fn=lambda x: gr.update(
+                        visible=x == "upload", elem_id="edit_source_audio_upload"
+                    ),
                     inputs=[edit_source],
                     outputs=[edit_source_audio_upload],
                 )
@@ -321,9 +518,7 @@ def create_text2music_ui(
                     if edit_source == "upload":
                         src_audio_path = edit_source_audio_upload
                         audio_duration = librosa.get_duration(filename=src_audio_path)
-                        json_data = {
+                        json_data = {"audio_duration": audio_duration}
-                            "audio_duration": audio_duration
-                        }
                     elif edit_source == "text2music":
                         json_data = text2music_json_data
                         src_audio_path = json_data["audio_path"]
@@ -397,14 +592,42 @@ def create_text2music_ui(
                     outputs=edit_outputs + [edit_input_params_json],
                 )
             with gr.Tab("extend"):
-                extend_seeds = gr.Textbox(label="extend seeds (default None)", placeholder="", value=None)
+                extend_seeds = gr.Textbox(
-                left_extend_length = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=0.0, label="Left Extend Length", interactive=True)
+                    label="extend seeds (default None)", placeholder="", value=None
-                right_extend_length = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=30.0, label="Right Extend Length", interactive=True)
+                )
-                extend_source = gr.Radio(["text2music", "last_extend", "upload"], value="text2music", label="Extend Source", elem_id="extend_source")
+                left_extend_length = gr.Slider(
+                    minimum=0.0,
+                    maximum=240.0,
+                    step=0.01,
+                    value=0.0,
+                    label="Left Extend Length",
+                    interactive=True,
+                )
+                right_extend_length = gr.Slider(
+                    minimum=0.0,
+                    maximum=240.0,
+                    step=0.01,
+                    value=30.0,
+                    label="Right Extend Length",
+                    interactive=True,
+                )
+                extend_source = gr.Radio(
+                    ["text2music", "last_extend", "upload"],
+                    value="text2music",
+                    label="Extend Source",
+                    elem_id="extend_source",
+                )
 
-                extend_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="extend_source_audio_upload")
+                extend_source_audio_upload = gr.Audio(
+                    label="Upload Audio",
+                    type="filepath",
+                    visible=False,
+                    elem_id="extend_source_audio_upload",
+                )
                 extend_source.change(
-                    fn=lambda x: gr.update(visible=x == "upload", elem_id="extend_source_audio_upload"),
+                    fn=lambda x: gr.update(
+                        visible=x == "upload", elem_id="extend_source_audio_upload"
+                    ),
                     inputs=[extend_source],
                     outputs=[extend_source_audio_upload],
                 )
@@ -442,9 +665,7 @@ def create_text2music_ui(
                         src_audio_path = extend_source_audio_upload
                         # get audio duration
                         audio_duration = librosa.get_duration(filename=src_audio_path)
-                        json_data = {
+                        json_data = {"audio_duration": audio_duration}
-                            "audio_duration": audio_duration
-                        }
                     elif extend_source == "text2music":
                         json_data = text2music_json_data
                         src_audio_path = json_data["audio_path"]
@@ -531,8 +752,16 @@ def create_text2music_ui(
                 json_data["use_erg_lyric"],
                 json_data["use_erg_diffusion"],
                 ", ".join(map(str, json_data["oss_steps"])),
-                json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
+                (
-                json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
+                    json_data["guidance_scale_text"]
+                    if "guidance_scale_text" in json_data
+                    else 0.0
+                ),
+                (
+                    json_data["guidance_scale_lyric"]
+                    if "guidance_scale_lyric" in json_data
+                    else 0.0
+                ),
             )
 
         sample_bnt.click(
@@ -580,7 +809,8 @@ def create_text2music_ui(
             oss_steps,
             guidance_scale_text,
             guidance_scale_lyric,
-        ], outputs=outputs + [input_params_json]
+        ],
+        outputs=outputs + [input_params_json],
     )
 
 
@@ -594,7 +824,8 @@ def create_main_demo_ui(
         gr.Markdown(
             """
             <h1 style="text-align: center;">ACE-Step: A Step Towards Music Generation Foundation Model</h1>
-        """)
+        """
+        )
 
         with gr.Tab("text2music"):
             create_text2music_ui(

infer.py

@@ -1,4 +1,5 @@
 import argparse
+
 parser = argparse.ArgumentParser()
 parser.add_argument("--checkpoint_path", type=str, default="")
 parser.add_argument("--bf16", type=bool, default=True)
@@ -11,8 +12,8 @@ import os
 
 os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device_id)
 
-from pipeline_ace_step import ACEStepPipeline
+from acestep.pipeline_ace_step import ACEStepPipeline
-from data_sampler import DataSampler
+from acestep.data_sampler import DataSampler
 
 
 def sample_data(json_data):
@@ -34,15 +35,20 @@ def sample_data(json_data):
         json_data["use_erg_diffusion"],
         ", ".join(map(str, json_data["oss_steps"])),
         json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
-            json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
+        (
+            json_data["guidance_scale_lyric"]
+            if "guidance_scale_lyric" in json_data
+            else 0.0
+        ),
     )
 
+
 def main(args):
 
     model_demo = ACEStepPipeline(
         checkpoint_dir=args.checkpoint_path,
         dtype="bfloat16" if args.bf16 else "float32",
-        torch_compile=args.torch_compile
+        torch_compile=args.torch_compile,
     )
     print(model_demo)
 
@@ -52,46 +58,49 @@ def main(args):
     json_data = sample_data(json_data)
     print(json_data)
 
-
+    (
-    audio_duration,\
+        audio_duration,
-    prompt, \
+        prompt,
-    lyrics,\
+        lyrics,
-    infer_step, \
+        infer_step,
-    guidance_scale,\
+        guidance_scale,
-    scheduler_type, \
+        scheduler_type,
-    cfg_type, \
+        cfg_type,
-    omega_scale, \
+        omega_scale,
-    manual_seeds, \
+        manual_seeds,
-    guidance_interval, \
+        guidance_interval,
-    guidance_interval_decay, \
+        guidance_interval_decay,
-    min_guidance_scale, \
+        min_guidance_scale,
-    use_erg_tag, \
+        use_erg_tag,
-    use_erg_lyric, \
+        use_erg_lyric,
-    use_erg_diffusion, \
+        use_erg_diffusion,
-    oss_steps, \
+        oss_steps,
-    guidance_scale_text, \
+        guidance_scale_text,
-    guidance_scale_lyric = json_data
-
-
-    model_demo(audio_duration, \
-            prompt, \
-            lyrics,\
-            infer_step, \
-            guidance_scale,\
-            scheduler_type, \
-            cfg_type, \
-            omega_scale, \
-            manual_seeds, \
-            guidance_interval, \
-            guidance_interval_decay, \
-            min_guidance_scale, \
-            use_erg_tag, \
-            use_erg_lyric, \
-            use_erg_diffusion, \
-            oss_steps, \
-            guidance_scale_text, \
         guidance_scale_lyric,
-            save_path=args.output_path)
+    ) = json_data
+
+    model_demo(
+        audio_duration,
+        prompt,
+        lyrics,
+        infer_step,
+        guidance_scale,
+        scheduler_type,
+        cfg_type,
+        omega_scale,
+        manual_seeds,
+        guidance_interval,
+        guidance_interval_decay,
+        min_guidance_scale,
+        use_erg_tag,
+        use_erg_lyric,
+        use_erg_diffusion,
+        oss_steps,
+        guidance_scale_text,
+        guidance_scale_lyric,
+        save_path=args.output_path,
+    )
+
 
 if __name__ == "__main__":
     main(args)

requirements.txt

@@ -20,3 +20,4 @@ spacy==3.8.4
 accelerate==1.6.0
 cutlet
 fugashi[unidic-lite]
+click

setup.py

@@ -0,0 +1,23 @@
+from setuptools import setup
+
+setup(
+    name="ace_step",
+    description="ACE Step: A Step Towards Music Generation Foundation Model",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    version="0.1.0",
+    packages=["acestep"],
+    install_requires=open("requirements.txt").read().splitlines(),
+    author="ACE Studio, StepFun AI",
+    license="Apache 2.0",
+    classifiers=[
+        "Development Status :: 3 - Alpha",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: Apache Software License",
+    ],
+    entry_points={
+        "console_scripts": [
+            "acestep=acestep.gui:main",
+        ],
+    },
+)

trainer.py

@@ -10,23 +10,27 @@ import torch.nn.functional as F
 import torch.utils.data
 from pytorch_lightning.core import LightningModule
 from torch.utils.data import DataLoader
-from schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from acestep.schedulers.scheduling_flow_match_euler_discrete import (
-from text2music_dataset import Text2MusicDataset
+    FlowMatchEulerDiscreteScheduler,
+)
+from acestep.text2music_dataset import Text2MusicDataset
 from loguru import logger
 from transformers import AutoModel, Wav2Vec2FeatureExtractor
 import torchaudio
-from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import retrieve_timesteps
+from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import (
+    retrieve_timesteps,
+)
 from diffusers.utils.torch_utils import randn_tensor
-from apg_guidance import apg_forward, MomentumBuffer
+from acestep.apg_guidance import apg_forward, MomentumBuffer
 from tqdm import tqdm
 import random
 import os
-from pipeline_ace_step import ACEStepPipeline
+from acestep.pipeline_ace_step import ACEStepPipeline
 
 
 matplotlib.use("Agg")
 torch.backends.cudnn.benchmark = False
-torch.set_float32_matmul_precision('high')
+torch.set_float32_matmul_precision("high")
 
 
 class Pipeline(LightningModule):
@@ -47,7 +51,7 @@ class Pipeline(LightningModule):
         max_steps: int = 200000,
         warmup_steps: int = 4000,
         dataset_path: str = "./data/your_dataset_path",
-        lora_config_path: str = None
+        lora_config_path: str = None,
     ):
         super().__init__()
 
@@ -91,15 +95,31 @@ class Pipeline(LightningModule):
         if self.is_train:
             self.transformers.train()
 
-            self.mert_model = AutoModel.from_pretrained("m-a-p/MERT-v1-330M", trust_remote_code=True, cache_dir=checkpoint_dir).eval()
+            self.mert_model = AutoModel.from_pretrained(
+                "m-a-p/MERT-v1-330M", trust_remote_code=True, cache_dir=checkpoint_dir
+            ).eval()
             self.mert_model.requires_grad_(False)
-            self.resampler_mert = torchaudio.transforms.Resample(orig_freq=48000, new_freq=24000)
+            self.resampler_mert = torchaudio.transforms.Resample(
-            self.processor_mert = Wav2Vec2FeatureExtractor.from_pretrained("m-a-p/MERT-v1-330M", trust_remote_code=True)
+                orig_freq=48000, new_freq=24000
+            )
+            self.processor_mert = Wav2Vec2FeatureExtractor.from_pretrained(
+                "m-a-p/MERT-v1-330M", trust_remote_code=True
+            )
 
-            self.hubert_model = AutoModel.from_pretrained("utter-project/mHuBERT-147", local_files_only=True, cache_dir=checkpoint_dir).eval()
+            self.hubert_model = AutoModel.from_pretrained(
+                "utter-project/mHuBERT-147",
+                local_files_only=True,
+                cache_dir=checkpoint_dir,
+            ).eval()
             self.hubert_model.requires_grad_(False)
-            self.resampler_mhubert = torchaudio.transforms.Resample(orig_freq=48000, new_freq=16000)
+            self.resampler_mhubert = torchaudio.transforms.Resample(
-            self.processor_mhubert = Wav2Vec2FeatureExtractor.from_pretrained("utter-project/mHuBERT-147", local_files_only=True, cache_dir=checkpoint_dir)
+                orig_freq=48000, new_freq=16000
+            )
+            self.processor_mhubert = Wav2Vec2FeatureExtractor.from_pretrained(
+                "utter-project/mHuBERT-147",
+                local_files_only=True,
+                cache_dir=checkpoint_dir,
+            )
 
             self.ssl_coeff = ssl_coeff
 
@@ -110,9 +130,21 @@ class Pipeline(LightningModule):
         actual_lengths_24k = wav_lengths // 2  # 48kHz -> 24kHz
 
         # Normalize the actual audio part
-        means = torch.stack([mert_input_wavs_mono_24k[i, :actual_lengths_24k[i]].mean() for i in range(bsz)])
+        means = torch.stack(
-        vars = torch.stack([mert_input_wavs_mono_24k[i, :actual_lengths_24k[i]].var() for i in range(bsz)])
+            [
-        mert_input_wavs_mono_24k = (mert_input_wavs_mono_24k - means.view(-1, 1)) / torch.sqrt(vars.view(-1, 1) + 1e-7)
+                mert_input_wavs_mono_24k[i, : actual_lengths_24k[i]].mean()
+                for i in range(bsz)
+            ]
+        )
+        vars = torch.stack(
+            [
+                mert_input_wavs_mono_24k[i, : actual_lengths_24k[i]].var()
+                for i in range(bsz)
+            ]
+        )
+        mert_input_wavs_mono_24k = (
+            mert_input_wavs_mono_24k - means.view(-1, 1)
+        ) / torch.sqrt(vars.view(-1, 1) + 1e-7)
 
         # MERT SSL constraint
         # Define the length of each chunk (5 seconds of samples)
@@ -131,7 +163,9 @@ class Pipeline(LightningModule):
                 end = min(start + chunk_size, actual_length)
                 chunk = audio[start:end]
                 if len(chunk) < chunk_size:
-                    chunk = F.pad(chunk, (0, chunk_size - len(chunk)))  # Pad insufficient parts with zeros
+                    chunk = F.pad(
+                        chunk, (0, chunk_size - len(chunk))
+                    )  # Pad insufficient parts with zeros
                 all_chunks.append(chunk)
                 chunk_actual_lengths.append(end - start)
 
@@ -147,14 +181,21 @@ class Pipeline(LightningModule):
         chunk_num_features = [(length + 319) // 320 for length in chunk_actual_lengths]
 
         # Trim the hidden states of each chunk
-        chunk_hidden_states = [mert_ssl_hidden_states[i, :chunk_num_features[i], :] for i in range(len(all_chunks))]
+        chunk_hidden_states = [
+            mert_ssl_hidden_states[i, : chunk_num_features[i], :]
+            for i in range(len(all_chunks))
+        ]
 
         # Organize hidden states by audio
         mert_ssl_hidden_states_list = []
         chunk_idx = 0
         for i in range(bsz):
-            audio_chunks = chunk_hidden_states[chunk_idx:chunk_idx + num_chunks_per_audio[i]]
+            audio_chunks = chunk_hidden_states[
-            audio_hidden = torch.cat(audio_chunks, dim=0)  # Concatenate chunks of the same audio
+                chunk_idx : chunk_idx + num_chunks_per_audio[i]
+            ]
+            audio_hidden = torch.cat(
+                audio_chunks, dim=0
+            )  # Concatenate chunks of the same audio
             mert_ssl_hidden_states_list.append(audio_hidden)
             chunk_idx += num_chunks_per_audio[i]
 
@@ -168,18 +209,29 @@ class Pipeline(LightningModule):
         actual_lengths_16k = wav_lengths // 3  # Convert lengths from 48kHz to 16kHz
 
         # Step 2: Zero-mean unit-variance normalization (only on actual audio)
-        means = torch.stack([mhubert_input_wavs_mono_16k[i, :actual_lengths_16k[i]].mean() 
+        means = torch.stack(
-                            for i in range(bsz)])
+            [
-        vars = torch.stack([mhubert_input_wavs_mono_16k[i, :actual_lengths_16k[i]].var() 
+                mhubert_input_wavs_mono_16k[i, : actual_lengths_16k[i]].mean()
-                            for i in range(bsz)])
+                for i in range(bsz)
-        mhubert_input_wavs_mono_16k = (mhubert_input_wavs_mono_16k - means.view(-1, 1)) / \
+            ]
-                                    torch.sqrt(vars.view(-1, 1) + 1e-7)
+        )
+        vars = torch.stack(
+            [
+                mhubert_input_wavs_mono_16k[i, : actual_lengths_16k[i]].var()
+                for i in range(bsz)
+            ]
+        )
+        mhubert_input_wavs_mono_16k = (
+            mhubert_input_wavs_mono_16k - means.view(-1, 1)
+        ) / torch.sqrt(vars.view(-1, 1) + 1e-7)
 
         # Step 3: Define chunk size for MHubert (30 seconds at 16kHz)
         chunk_size = 16000 * 30  # 30 seconds = 480,000 samples
 
         # Step 4: Split audio into chunks
-        num_chunks_per_audio = (actual_lengths_16k + chunk_size - 1) // chunk_size  # Ceiling division
+        num_chunks_per_audio = (
+            actual_lengths_16k + chunk_size - 1
+        ) // chunk_size  # Ceiling division
         all_chunks = []
         chunk_actual_lengths = []
 
@@ -206,20 +258,33 @@ class Pipeline(LightningModule):
         chunk_num_features = [(length + 319) // 320 for length in chunk_actual_lengths]
 
         # Step 8: Trim hidden states to remove padding effects
-        chunk_hidden_states = [mhubert_ssl_hidden_states[i, :chunk_num_features[i], :] for i in range(len(all_chunks))]
+        chunk_hidden_states = [
+            mhubert_ssl_hidden_states[i, : chunk_num_features[i], :]
+            for i in range(len(all_chunks))
+        ]
 
         # Step 9: Reorganize hidden states by original audio
         mhubert_ssl_hidden_states_list = []
         chunk_idx = 0
         for i in range(bsz):
-            audio_chunks = chunk_hidden_states[chunk_idx:chunk_idx + num_chunks_per_audio[i]]
+            audio_chunks = chunk_hidden_states[
-            audio_hidden = torch.cat(audio_chunks, dim=0)  # Concatenate chunks for this audio
+                chunk_idx : chunk_idx + num_chunks_per_audio[i]
+            ]
+            audio_hidden = torch.cat(
+                audio_chunks, dim=0
+            )  # Concatenate chunks for this audio
             mhubert_ssl_hidden_states_list.append(audio_hidden)
             chunk_idx += num_chunks_per_audio[i]
         return mhubert_ssl_hidden_states_list
 
     def get_text_embeddings(self, texts, device, text_max_length=256):
-        inputs = self.text_tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=text_max_length)
+        inputs = self.text_tokenizer(
+            texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=text_max_length,
+        )
         inputs = {key: value.to(device) for key, value in inputs.items()}
         if self.text_encoder_model.device != device:
             self.text_encoder_model.to(device)
@@ -243,15 +308,21 @@ class Pipeline(LightningModule):
         if train:
             with torch.amp.autocast(device_type="cuda", dtype=dtype):
                 mert_ssl_hidden_states = self.infer_mert_ssl(target_wavs, wav_lengths)
-                mhubert_ssl_hidden_states = self.infer_mhubert_ssl(target_wavs, wav_lengths)
+                mhubert_ssl_hidden_states = self.infer_mhubert_ssl(
+                    target_wavs, wav_lengths
+                )
 
         # 1: text embedding
         texts = batch["prompts"]
-        encoder_text_hidden_states, text_attention_mask = self.get_text_embeddings(texts, device)
+        encoder_text_hidden_states, text_attention_mask = self.get_text_embeddings(
+            texts, device
+        )
         encoder_text_hidden_states = encoder_text_hidden_states.to(dtype)
 
         target_latents, _ = self.dcae.encode(target_wavs, wav_lengths)
-        attention_mask = torch.ones(bs, target_latents.shape[-1], device=device, dtype=dtype)
+        attention_mask = torch.ones(
+            bs, target_latents.shape[-1], device=device, dtype=dtype
+        )
 
         speaker_embds = batch["speaker_embs"].to(dtype)
         keys = batch["keys"]
@@ -263,27 +334,43 @@ class Pipeline(LightningModule):
             full_cfg_condition_mask = torch.where(
                 (torch.rand(size=(bs,), device=device) < 0.15),
                 torch.zeros(size=(bs,), device=device),
-                torch.ones(size=(bs,), device=device)
+                torch.ones(size=(bs,), device=device),
             ).long()
             # N x T x 768
-            encoder_text_hidden_states = torch.where(full_cfg_condition_mask.unsqueeze(1).unsqueeze(1).bool(), encoder_text_hidden_states, torch.zeros_like(encoder_text_hidden_states))
+            encoder_text_hidden_states = torch.where(
+                full_cfg_condition_mask.unsqueeze(1).unsqueeze(1).bool(),
+                encoder_text_hidden_states,
+                torch.zeros_like(encoder_text_hidden_states),
+            )
 
             full_cfg_condition_mask = torch.where(
                 (torch.rand(size=(bs,), device=device) < 0.50),
                 torch.zeros(size=(bs,), device=device),
-                torch.ones(size=(bs,), device=device)
+                torch.ones(size=(bs,), device=device),
             ).long()
             # N x 512
-            speaker_embds = torch.where(full_cfg_condition_mask.unsqueeze(1).bool(), speaker_embds, torch.zeros_like(speaker_embds))
+            speaker_embds = torch.where(
+                full_cfg_condition_mask.unsqueeze(1).bool(),
+                speaker_embds,
+                torch.zeros_like(speaker_embds),
+            )
 
             # Lyrics
             full_cfg_condition_mask = torch.where(
                 (torch.rand(size=(bs,), device=device) < 0.15),
                 torch.zeros(size=(bs,), device=device),
-                torch.ones(size=(bs,), device=device)
+                torch.ones(size=(bs,), device=device),
             ).long()
-            lyric_token_ids = torch.where(full_cfg_condition_mask.unsqueeze(1).bool(), lyric_token_ids, torch.zeros_like(lyric_token_ids))
+            lyric_token_ids = torch.where(
-            lyric_mask = torch.where(full_cfg_condition_mask.unsqueeze(1).bool(), lyric_mask, torch.zeros_like(lyric_mask))
+                full_cfg_condition_mask.unsqueeze(1).bool(),
+                lyric_token_ids,
+                torch.zeros_like(lyric_token_ids),
+            )
+            lyric_mask = torch.where(
+                full_cfg_condition_mask.unsqueeze(1).bool(),
+                lyric_mask,
+                torch.zeros_like(lyric_mask),
+            )
 
         return (
             keys,
@@ -305,10 +392,12 @@ class Pipeline(LightningModule):
         )
 
     def configure_optimizers(self):
-        trainable_params = [p for name, p in self.transformers.named_parameters() if p.requires_grad]
+        trainable_params = [
+            p for name, p in self.transformers.named_parameters() if p.requires_grad
+        ]
         optimizer = torch.optim.AdamW(
             params=[
-                {'params': trainable_params},
+                {"params": trainable_params},
             ],
             lr=self.hparams.learning_rate,
             weight_decay=self.hparams.weight_decay,
@@ -324,13 +413,13 @@ class Pipeline(LightningModule):
                 return float(current_step) / float(max(1, warmup_steps))
             else:
                 # Linear decay from learning_rate to 0
-                progress = float(current_step - warmup_steps) / float(max(1, max_steps - warmup_steps))
+                progress = float(current_step - warmup_steps) / float(
+                    max(1, max_steps - warmup_steps)
+                )
                 return max(0.0, 1.0 - progress)
 
         lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
-            optimizer,
+            optimizer, lr_lambda, last_epoch=-1
-            lr_lambda,
-            last_epoch=-1
         )
         return [optimizer], lr_scheduler
 
@@ -362,10 +451,17 @@ class Pipeline(LightningModule):
             # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
             # In practice, we sample the random variable u from a normal distribution u ∼ N (u; m, s)
             # and map it through the standard logistic function
-            u = torch.normal(mean=self.hparams.logit_mean, std=self.hparams.logit_std, size=(bsz, ), device="cpu")
+            u = torch.normal(
+                mean=self.hparams.logit_mean,
+                std=self.hparams.logit_std,
+                size=(bsz,),
+                device="cpu",
+            )
             u = torch.nn.functional.sigmoid(u)
             indices = (u * self.scheduler.config.num_train_timesteps).long()
-            indices = torch.clamp(indices, 0, self.scheduler.config.num_train_timesteps - 1)
+            indices = torch.clamp(
+                indices, 0, self.scheduler.config.num_train_timesteps - 1
+            )
             timesteps = self.scheduler.timesteps[indices].to(device)
 
         return timesteps
@@ -394,7 +490,9 @@ class Pipeline(LightningModule):
         timesteps = self.get_timestep(bsz, device)
 
         # Add noise according to flow matching.
-        sigmas = self.get_sd3_sigmas(timesteps=timesteps, device=device, n_dim=target_image.ndim, dtype=dtype)
+        sigmas = self.get_sd3_sigmas(
+            timesteps=timesteps, device=device, n_dim=target_image.ndim, dtype=dtype
+        )
         noisy_image = sigmas * noise + (1.0 - sigmas) * target_image
 
         # This is the flow-matching target for vanilla SD3.
@@ -431,7 +529,11 @@ class Pipeline(LightningModule):
         # Compute loss. Only calculate loss where chunk_mask is 1 and there is no padding
         # N x T x 64
         # N x T -> N x c x W x T
-        mask = attention_mask.unsqueeze(1).unsqueeze(1).expand(-1, target_image.shape[1], target_image.shape[2], -1)
+        mask = (
+            attention_mask.unsqueeze(1)
+            .unsqueeze(1)
+            .expand(-1, target_image.shape[1], target_image.shape[2], -1)
+        )
 
         selected_model_pred = (model_pred * mask).reshape(bsz, -1).contiguous()
         selected_target = (target * mask).reshape(bsz, -1).contiguous()
@@ -443,11 +545,19 @@ class Pipeline(LightningModule):
 
         prefix = "train"
 
-        self.log(f"{prefix}/denoising_loss", loss, on_step=True, on_epoch=False, prog_bar=True)
+        self.log(
+            f"{prefix}/denoising_loss",
+            loss,
+            on_step=True,
+            on_epoch=False,
+            prog_bar=True,
+        )
 
         total_proj_loss = 0.0
         for k, v in proj_losses:
-            self.log(f"{prefix}/{k}_loss", v, on_step=True, on_epoch=False, prog_bar=True)
+            self.log(
+                f"{prefix}/{k}_loss", v, on_step=True, on_epoch=False, prog_bar=True
+            )
             total_proj_loss += v
 
         if len(proj_losses) > 0:
@@ -459,7 +569,13 @@ class Pipeline(LightningModule):
         # Log learning rate if scheduler exists
         if self.lr_schedulers() is not None:
             learning_rate = self.lr_schedulers().get_last_lr()[0]
-            self.log(f"{prefix}/learning_rate", learning_rate, on_step=True, on_epoch=False, prog_bar=True)
+            self.log(
+                f"{prefix}/learning_rate",
+                learning_rate,
+                on_step=True,
+                on_epoch=False,
+                prog_bar=True,
+            )
         # with torch.autograd.detect_anomaly():
         #     self.manual_backward(loss)
         return loss
@@ -496,18 +612,35 @@ class Pipeline(LightningModule):
         )
 
         frame_length = int(duration * 44100 / 512 / 8)
-        timesteps, num_inference_steps = retrieve_timesteps(scheduler, num_inference_steps=infer_steps, device=device, timesteps=None)
+        timesteps, num_inference_steps = retrieve_timesteps(
+            scheduler, num_inference_steps=infer_steps, device=device, timesteps=None
+        )
 
-        target_latents = randn_tensor(shape=(bsz, 8, 16, frame_length), generator=random_generators, device=device, dtype=dtype)
+        target_latents = randn_tensor(
+            shape=(bsz, 8, 16, frame_length),
+            generator=random_generators,
+            device=device,
+            dtype=dtype,
+        )
         attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
         if do_classifier_free_guidance:
             attention_mask = torch.cat([attention_mask] * 2, dim=0)
-            encoder_text_hidden_states = torch.cat([encoder_text_hidden_states, torch.zeros_like(encoder_text_hidden_states)], 0)
+            encoder_text_hidden_states = torch.cat(
+                [
+                    encoder_text_hidden_states,
+                    torch.zeros_like(encoder_text_hidden_states),
+                ],
+                0,
+            )
             text_attention_mask = torch.cat([text_attention_mask] * 2, dim=0)
 
-            speaker_embds = torch.cat([speaker_embds, torch.zeros_like(speaker_embds)], 0)
+            speaker_embds = torch.cat(
+                [speaker_embds, torch.zeros_like(speaker_embds)], 0
+            )
 
-            lyric_token_ids = torch.cat([lyric_token_ids, torch.zeros_like(lyric_token_ids)], 0)
+            lyric_token_ids = torch.cat(
+                [lyric_token_ids, torch.zeros_like(lyric_token_ids)], 0
+            )
             lyric_mask = torch.cat([lyric_mask, torch.zeros_like(lyric_mask)], 0)
 
         momentum_buffer = MomentumBuffer()
@@ -515,7 +648,9 @@ class Pipeline(LightningModule):
         for i, t in tqdm(enumerate(timesteps), total=num_inference_steps):
             # expand the latents if we are doing classifier free guidance
             latents = target_latents
-            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            latent_model_input = (
+                torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            )
             timestep = t.expand(latent_model_input.shape[0])
             noise_pred = self.transformers(
                 hidden_states=latent_model_input,
@@ -537,7 +672,13 @@ class Pipeline(LightningModule):
                     momentum_buffer=momentum_buffer,
                 )
 
-            target_latents = scheduler.step(model_output=noise_pred, timestep=t, sample=target_latents, return_dict=False, omega=omega_scale)[0]
+            target_latents = scheduler.step(
+                model_output=noise_pred,
+                timestep=t,
+                sample=target_latents,
+                return_dict=False,
+                omega=omega_scale,
+            )[0]
 
         return target_latents
 
@@ -582,7 +723,9 @@ class Pipeline(LightningModule):
         )
 
         audio_lengths = batch["wav_lengths"]
-        sr, pred_wavs = self.dcae.decode(pred_latents, audio_lengths=audio_lengths, sr=48000)
+        sr, pred_wavs = self.dcae.decode(
+            pred_latents, audio_lengths=audio_lengths, sr=48000
+        )
         return {
             "target_wavs": batch["target_wavs"],
             "pred_wavs": pred_wavs,
@@ -603,7 +746,12 @@ class Pipeline(LightningModule):
 
     def plot_step(self, batch, batch_idx):
         global_step = self.global_step
-        if global_step % self.hparams.every_plot_step != 0 or self.local_rank != 0 or torch.distributed.get_rank() != 0 or torch.cuda.current_device() != 0:
+        if (
+            global_step % self.hparams.every_plot_step != 0
+            or self.local_rank != 0
+            or torch.distributed.get_rank() != 0
+            or torch.cuda.current_device() != 0
+        ):
             return
         results = self.predict_step(batch)
 
@@ -615,7 +763,9 @@ class Pipeline(LightningModule):
         sr = results["sr"]
         seeds = results["seeds"]
         i = 0
-        for key, target_wav, pred_wav, prompt, candidate_lyric_chunk, seed in zip(keys, target_wavs, pred_wavs, prompts, candidate_lyric_chunks, seeds):
+        for key, target_wav, pred_wav, prompt, candidate_lyric_chunk, seed in zip(
+            keys, target_wavs, pred_wavs, prompts, candidate_lyric_chunks, seeds
+        ):
             key = key
             prompt = prompt
             lyric = self.construct_lyrics(candidate_lyric_chunk)
@@ -624,9 +774,15 @@ class Pipeline(LightningModule):
             save_dir = f"{log_dir}/eval_results/step_{self.global_step}"
             if not os.path.exists(save_dir):
                 os.makedirs(save_dir, exist_ok=True)
-            torchaudio.save(f"{save_dir}/target_wav_{key}_{i}.flac", target_wav.float().cpu(), sr)
+            torchaudio.save(
-            torchaudio.save(f"{save_dir}/pred_wav_{key}_{i}.flac", pred_wav.float().cpu(), sr)
+                f"{save_dir}/target_wav_{key}_{i}.flac", target_wav.float().cpu(), sr
-            with open(f"{save_dir}/key_prompt_lyric_{key}_{i}.txt", "w", encoding="utf-8") as f:
+            )
+            torchaudio.save(
+                f"{save_dir}/pred_wav_{key}_{i}.flac", pred_wav.float().cpu(), sr
+            )
+            with open(
+                f"{save_dir}/key_prompt_lyric_{key}_{i}.txt", "w", encoding="utf-8"
+            ) as f:
                 f.write(key_prompt_lyric)
             i += 1
 
@@ -642,11 +798,14 @@ def main(args):
         checkpoint_dir=args.checkpoint_dir,
     )
     checkpoint_callback = ModelCheckpoint(
-        monitor=None, every_n_train_steps=args.every_n_train_steps, save_top_k=-1,
+        monitor=None,
+        every_n_train_steps=args.every_n_train_steps,
+        save_top_k=-1,
     )
     # add datetime str to version
     logger_callback = TensorBoardLogger(
-        version=datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + args.exp_name, save_dir=args.logger_dir
+        version=datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + args.exp_name,
+        save_dir=args.logger_dir,
     )
     trainer = Trainer(
         accelerator="gpu",