nodes/billwuhao-ComfyUI_ACE-Step
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

v1.0.0

Browse Source
This commit is contained in:
billwuhao
2025-05-07 20:38:55 +08:00
parent 968f27eb42
commit e6e7908259
57 changed files with 25290 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/publish_action.yml
+22
View File
@@ -0,0 +1,22 @@
name: Publish to Comfy registry
on:
workflow_dispatch:
push:
branches:
- master
- main
paths:
- "pyproject.toml"
jobs:
publish-node:
name: Publish Custom Node to registry
runs-on: ubuntu-latest
steps:
- name: Check out code
uses: actions/checkout@v4
- name: Publish Custom Node
uses: Comfy-Org/publish-node-action@main
with:
## Add your own personal access token to your Github Repository secrets and reference it here.
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
LICENSE
+202
View File
@@ -0,0 +1,202 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [2025] Timedomain Inc. and stepfun
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
README-CN.md
+72
View File
@@ -0,0 +1,72 @@
[中文](README-CN.md)|[English](README.md)
# ACE-Step 的 ComfyUI 节点
快速, 高质量音乐生成, "重绘", Remix, 编辑, 扩展等, Windows, Linux, Mac 应该都支持(未做完整测试).
示例:
- 生成:
![](https://github.com/billwuhao/ComfyUI_ACE-Step/blob/main/images/2025-05-07_19-53-51.png)
- "重绘":
![](https://github.com/billwuhao/ComfyUI_ACE-Step/blob/main/images/2025-05-07_19-59-22.png)
- 扩展:
![](https://github.com/billwuhao/ComfyUI_ACE-Step/blob/main/images/2025-05-07_20-04-02.png)
- 编辑:
![](https://github.com/billwuhao/ComfyUI_ACE-Step/blob/main/images/2025-05-07_20-09-52.png)
## 📣 更新
[2025-05-07]⚒️: 发布版本 v1.0.0.
## 安装
```
cd ComfyUI/custom_nodes
git clone https://github.com/billwuhao/ComfyUI_ACE-Step.git
cd ComfyUI_ACE-Step
pip install -r requirements.txt
# python_embeded
./python_embeded/python.exe -m pip install -r requirements.txt
```
## 模型下载
模型会自动下载到 `models\TTS\ACE-Step-v1-3.5B` 目录下, 也可以手动下载放到该目录下, 手动下载结构如下:
https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B
```
ACE-Step-v1-3.5B
│
├─ace_step_transformer
│ config.json
│ diffusion_pytorch_model.safetensors
│
├─music_dcae_f8c8
│ config.json
│ diffusion_pytorch_model.safetensors
│
├─music_vocoder
│ config.json
│ diffusion_pytorch_model.safetensors
│
└─umt5-base
config.json
model.safetensors
special_tokens_map.json
tokenizer.json
tokenizer_config.json
```
## 鸣谢
[ACE-Step](https://github.com/ace-step/ACE-Step)
README.md
+72 -1
View File
@@ -1 +1,72 @@
# ComfyUI_ACE-Step
[中文](README-CN.md)|[English](README.md)
# ACE-Step Nodes for ComfyUI
Fast, high-quality music generation, "redrawing", remixing, editing, extending, and more. Windows, Linux, and Mac should all be supported (not fully tested).
Examples:
- Generation:
![](https://github.com/billwuhao/ComfyUI_ACE-Step/blob/main/images/2025-05-07_19-53-51.png)
- "Redrawing":
![](https://github.com/billwuhao/ComfyUI_ACE-Step/blob/main/images/2025-05-07_19-59-22.png)
- Extending:
![](https://github.com/billwuhao/ComfyUI_ACE-Step/blob/main/images/2025-05-07_20-04-02.png)
- Editing:
![](https://github.com/billwuhao/ComfyUI_ACE-Step/blob/main/images/2025-05-07_20-09-52.png)
## 📣 Updates
[2025-05-07]⚒️: Released version v1.0.0.
## Installation
```
cd ComfyUI/custom_nodes
git clone https://github.com/billwuhao/ComfyUI_ACE-Step.git
cd ComfyUI_ACE-Step
pip install -r requirements.txt
# python_embeded
./python_embeded/python.exe -m pip install -r requirements.txt
```
## Model Download
The model will be automatically downloaded to the `models\TTS\ACE-Step-v1-3.5B` directory. You can also manually download it and place it in this directory. The manual download structure is as follows:
https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B
```
ACE-Step-v1-3.5B
│
├─ace_step_transformer
│ config.json
│ diffusion_pytorch_model.safetensors
│
├─music_dcae_f8c8
│ config.json
│ diffusion_pytorch_model.safetensors
│
├─music_vocoder
│ config.json
│ diffusion_pytorch_model.safetensors
│
└─umt5-base
config.json
model.safetensors
special_tokens_map.json
tokenizer.json
tokenizer_config.json
```
## Acknowledgements
[ACE-Step](https://github.com/ace-step/ACE-Step)
__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .ace_step_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
ace_step_nodes.py
+485
View File
@@ -0,0 +1,485 @@
import torchaudio
import tempfile
from typing import Optional, List
import torch
import os
import ast
import sys
import librosa
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
sys.path.append(current_dir)
from pipeline_ace_step import ACEStepPipeline as AP
import folder_paths
cache_dir = folder_paths.get_temp_directory()
models_dir = folder_paths.models_dir
model_path = os.path.join(models_dir, "TTS", "ACE-Step-v1-3.5B")
class AudioCacher:
"""
一个用于缓存音频张量到临时文件，并在之后清理这些文件的类。
支持作为上下文管理器使用，以便自动清理。
"""
def __init__(self, cache_dir: Optional[str] = None, default_format: str = "wav"):
"""
初始化 AudioCacher。
Args:
cache_dir (Optional[str]): 缓存文件存放的目录。
如果为 None，则使用系统默认的临时目录。
default_format (str): 默认的音频文件格式后缀 (例如 "wav", "mp3", "flac")。
"""
if cache_dir is None:
self.cache_dir = tempfile.gettempdir()
else:
self.cache_dir = cache_dir
# 确保缓存目录存在
if not os.path.exists(self.cache_dir):
try:
os.makedirs(self.cache_dir, exist_ok=True)
except OSError as e:
raise # 重新抛出异常，因为这是一个关键的初始化步骤
self.default_format = default_format.lstrip('.') # 确保没有前导点
self._files_to_cleanup_in_context: List[str] = [] # 用于上下文管理器
def cache_audio_tensor(
self,
audio_tensor: torch.Tensor,
sample_rate: int,
filename_prefix: str = "cached_audio_",
audio_format: Optional[str] = None
) -> str:
"""
将音频张量保存到缓存文件，并返回文件路径。
Args:
audio_tensor (torch.Tensor): 要保存的音频张量。
sample_rate (int): 音频的采样率。
filename_prefix (str): 缓存文件名的前缀。
audio_format (Optional[str]): 要使用的音频格式 (例如 "wav", "mp3")。
如果为 None，则使用初始化时设置的 default_format。
Returns:
str: 保存的缓存文件的绝对路径。
Raises:
RuntimeError: 如果保存音频失败。
"""
current_format = (audio_format or self.default_format).lstrip('.')
# 创建一个带特定后缀的临时文件，但不立即删除
# NamedTemporaryFile 会在创建时打开文件，我们需要先关闭它才能让 torchaudio.save 使用
try:
with tempfile.NamedTemporaryFile(
prefix=filename_prefix,
suffix=f".{current_format}",
dir=self.cache_dir,
delete=False # 这是关键，我们手动管理删除
) as tmp_file:
temp_filepath = tmp_file.name
# 此时 tmp_file 已经关闭，但文件因 delete=False 而保留
torchaudio.save(temp_filepath, audio_tensor, sample_rate)
# 如果在上下文管理器中使用，则记录此文件以备自动清理
self._files_to_cleanup_in_context.append(temp_filepath)
return temp_filepath
except Exception as e:
# 如果 temp_filepath 已定义且文件存在，尝试删除，因为它可能不完整或损坏
if 'temp_filepath' in locals() and os.path.exists(temp_filepath):
try:
os.remove(temp_filepath)
except OSError as e_clean:
logger.warning(f"Error cleaning temporary file {temp_filepath}: {e_clean}")
raise RuntimeError(f"Failed to save audio: {e}") from e
def cleanup_file(self, filepath: str) -> bool:
"""
清理指定的缓存文件。
Args:
filepath (str): 要删除的文件的路径。
Returns:
bool: 如果文件成功删除或文件不存在，则返回 True；如果删除失败，则返回 False。
"""
if not filepath:
return True # 没有文件可以删除，所以认为是“成功”
if os.path.exists(filepath):
try:
os.remove(filepath)
# 如果文件在上下文中被跟踪，也从中移除
if filepath in self._files_to_cleanup_in_context:
self._files_to_cleanup_in_context.remove(filepath)
return True
except OSError as e:
return False
else:
# 如果文件在上下文中被跟踪，也从中移除
if filepath in self._files_to_cleanup_in_context:
self._files_to_cleanup_in_context.remove(filepath)
return True # 文件不存在，也视为清理“成功”
def cleanup_all_tracked_files(self) -> None:
"""
清理所有由当前上下文管理器实例跟踪的缓存文件。
"""
# 迭代列表的副本，因为 cleanup_file 可能会修改列表
for f_path in list(self._files_to_cleanup_in_context):
self.cleanup_file(f_path)
self._files_to_cleanup_in_context.clear() # 确保列表被清空
def __enter__(self):
"""进入上下文管理器时调用。"""
# 重置跟踪文件列表，以防同一个实例被多次用于 'with' 语句
self._files_to_cleanup_in_context = []
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""退出上下文管理器时调用，负责清理。"""
self.cleanup_all_tracked_files()
# 返回 False 以便在发生异常时重新抛出异常
return False
from data_sampler import DataSampler
def sample_data(json_data):
return (
json_data["audio_duration"],
json_data["prompt"],
json_data["lyrics"],
json_data["infer_step"],
json_data["guidance_scale"],
json_data["scheduler_type"],
json_data["cfg_type"],
json_data["omega_scale"],
json_data["actual_seeds"][0],
json_data["guidance_interval"],
json_data["guidance_interval_decay"],
json_data["min_guidance_scale"],
json_data["use_erg_tag"],
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,
)
data_sampler = DataSampler()
json_data = data_sampler.sample()
json_data = sample_data(json_data)
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 = json_data
class GenerationParameters:
@classmethod
def INPUT_TYPES(s):
return {"required":
{ "audio_duration": ("FLOAT", {"default": audio_duration, "min": 0.0, "max": 240.0, "step": 1.0, "tooltip": "0 is a random length"}),
"infer_step": ("INT", {"default": infer_step, "min": 1, "max": 60, "step": 1}),
"guidance_scale": ("FLOAT", {"default": guidance_scale, "min": 0.0, "max": 200.0, "step": 0.1, "tooltip": "When guidance_scale_lyric > 1 and guidance_scale_text > 1, the guidance scale will not be applied."}),
"scheduler_type": (["euler", "heun"], {"default": scheduler_type, "tooltip": "euler is recommended. heun will take more time."}),
"cfg_type": (["cfg", "apg", "cfg_star"], {"default": cfg_type, "tooltip": "apg is recommended. cfg and cfg_star are almost the same."}),
"omega_scale": ("FLOAT", {"default": omega_scale, "min": -100.0, "max": 100.0, "step": 0.1, "tooltip": "Higher values can reduce artifacts"}),
"seed": ("INT", {"default":manual_seeds, "min": 0, "max": 4294967295, "step": 1}),
"guidance_interval": ("FLOAT", {"default": guidance_interval, "min": 0, "max": 1, "step": 0.01, "tooltip": "0.5 means only apply guidance in the middle steps"}),
"guidance_interval_decay": ("FLOAT", {"default": guidance_interval_decay, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Guidance scale will decay from guidance_scale to min_guidance_scale in the interval. 0.0 means no decay."}),
"min_guidance_scale": ("INT", {"default": min_guidance_scale, "min": 0, "max": 200, "step": 1}),
"use_erg_tag": ("BOOLEAN", {"default": use_erg_tag}),
"use_erg_lyric": ("BOOLEAN", {"default": use_erg_lyric}),
"use_erg_diffusion": ("BOOLEAN", {"default": use_erg_diffusion}),
"oss_steps": ("STRING", {"default": oss_steps}),
"guidance_scale_text": ("FLOAT", {"default": guidance_scale_text, "min": 0.0, "max": 10.0, "step": 0.1}),
"guidance_scale_lyric": ("FLOAT", {"default": guidance_scale_lyric, "min": 0.0, "max": 10.0, "step": 0.1}),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("parameters",)
FUNCTION = "generate"
CATEGORY = "🎤MW/MW-ACE-Step"
def generate(self, **kwargs):
kwargs["manual_seeds"] = kwargs.pop("seed")
return (str(kwargs),)
class MultiLinePromptACES:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"multi_line_prompt": ("STRING", {
"multiline": True,
"default": prompt}),
},
}
CATEGORY = "🎤MW/MW-ACE-Step"
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("prompt",)
FUNCTION = "promptgen"
def promptgen(self, multi_line_prompt: str):
return (multi_line_prompt.strip(),)
class MultiLineLyrics:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"multi_line_prompt": ("STRING", {
"multiline": True,
"default": lyrics}),
},
}
CATEGORY = "🎤MW/MW-ACE-Step"
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("lyrics",)
FUNCTION = "lyricsgen"
def lyricsgen(self, multi_line_prompt: str):
return (multi_line_prompt.strip(),)
class ACEStepGen:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"prompt": ("STRING", {"forceInput": True}),
"lyrics": ("STRING", {"forceInput": True}),
"parameters": ("STRING", {"forceInput": True}),
# "unload_model": ("BOOLEAN", {"default": False}),
},
}
CATEGORY = "🎤MW/MW-ACE-Step"
RETURN_TYPES = ("AUDIO",)
RETURN_NAMES = ("music",)
FUNCTION = "acestepgen"
def acestepgen(self, prompt: str, lyrics: str, parameters: str, unload_model=True):
parameters = ast.literal_eval(parameters)
ap = AP(model_path)
audio_output = ap(prompt=prompt, lyrics=lyrics, task="text2music", **parameters)
audio, sr = audio_output[0][0].unsqueeze(0), audio_output[0][1]
if unload_model:
ap.cleanup()
return ({"waveform": audio, "sample_rate": sr},)
class ACEStepRepainting:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"src_audio": ("AUDIO",),
"prompt": ("STRING", {"forceInput": True}),
"lyrics": ("STRING", {"forceInput": True}),
"parameters": ("STRING", {"forceInput": True}),
"repaint_start": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
"repaint_end": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
"repaint_variance": ("FLOAT", {"default": 0.01, "min": 0.01, "max": 1.0, "step": 0.01}),
"seed": ("INT", {"default":0, "min": 0, "max": 4294967295, "step": 1}),
# "unload_model": ("BOOLEAN", {"default": False}),
},
}
CATEGORY = "🎤MW/MW-ACE-Step"
RETURN_TYPES = ("AUDIO",)
RETURN_NAMES = ("music",)
FUNCTION = "acesteprepainting"
def acesteprepainting(self, src_audio, prompt: str, lyrics: str, parameters: str, repaint_start, repaint_end, repaint_variance, seed, unload_model=True):
retake_seeds = [str(seed)]
ac = AudioCacher(cache_dir=cache_dir)
src_audio_path = ac.cache_audio_tensor(src_audio["waveform"].squeeze(0), src_audio["sample_rate"], filename_prefix="src_audio_")
audio_duration = librosa.get_duration(filename=src_audio_path)
if repaint_end > audio_duration:
repaint_end = audio_duration
parameters = ast.literal_eval(parameters)
parameters["audio_duration"] = audio_duration
ap = AP(model_path)
audio_output = ap(
prompt=prompt,
lyrics=lyrics,
task="repaint",
retake_seeds=retake_seeds,
src_audio_path=src_audio_path,
repaint_start=repaint_start,
repaint_end=repaint_end,
retake_variance=repaint_variance,
**parameters)
audio, sr = audio_output[0][0].unsqueeze(0), audio_output[0][1]
ac.cleanup_file(src_audio_path)
if unload_model:
ap.cleanup()
return ({"waveform": audio, "sample_rate": sr},)
class ACEStepEdit:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"src_audio": ("AUDIO",),
"prompt": ("STRING", {"forceInput": True}),
"lyrics": ("STRING", {"forceInput": True}),
"parameters": ("STRING", {"forceInput": True}),
"edit_prompt": ("STRING", {"forceInput": True}),
"edit_lyrics": ("STRING", {"forceInput": True}),
"edit_n_min": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.01}),
"edit_n_max": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"seed": ("INT", {"default":0, "min": 0, "max": 4294967295, "step": 1}),
# "unload_model": ("BOOLEAN", {"default": False}),
},
}
CATEGORY = "🎤MW/MW-ACE-Step"
RETURN_TYPES = ("AUDIO",)
RETURN_NAMES = ("music",)
FUNCTION = "acestepedit"
def acestepedit(self, src_audio, prompt: str, lyrics: str, parameters: str, edit_prompt, edit_lyrics, edit_n_min, edit_n_max, seed, unload_model=True):
retake_seeds = [str(seed)]
ac = AudioCacher(cache_dir=cache_dir)
src_audio_path = ac.cache_audio_tensor(src_audio["waveform"].squeeze(0), src_audio["sample_rate"], filename_prefix="src_audio_")
audio_duration = librosa.get_duration(filename=src_audio_path)
parameters = ast.literal_eval(parameters)
parameters["audio_duration"] = audio_duration
ap = AP(model_path)
audio_output = ap(
prompt=prompt,
lyrics=lyrics,
task="edit",
retake_seeds=retake_seeds,
src_audio_path=src_audio_path,
edit_target_prompt = edit_prompt,
edit_target_lyrics = edit_lyrics,
edit_n_min = edit_n_min,
edit_n_max = edit_n_max,
**parameters)
audio, sr = audio_output[0][0].unsqueeze(0), audio_output[0][1]
ac.cleanup_file(src_audio_path)
if unload_model:
ap.cleanup()
return ({"waveform": audio, "sample_rate": sr},)
class ACEStepExtend:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"src_audio": ("AUDIO",),
"prompt": ("STRING", {"forceInput": True}),
"lyrics": ("STRING", {"forceInput": True}),
"parameters": ("STRING", {"forceInput": True}),
"left_extend_length": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
"right_extend_length": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
# "repaint_variance": ("FLOAT", {"default": 0.01, "min": 0.01, "max": 1.0, "step": 0.01}),
"seed": ("INT", {"default":0, "min": 0, "max": 4294967295, "step": 1}),
# "unload_model": ("BOOLEAN", {"default": False}),
},
}
CATEGORY = "🎤MW/MW-ACE-Step"
RETURN_TYPES = ("AUDIO",)
RETURN_NAMES = ("music",)
FUNCTION = "acestepextend"
def acestepextend(self, src_audio, prompt: str, lyrics: str, parameters: str, left_extend_length, right_extend_length, seed, unload_model=True):
retake_seeds = [str(seed)]
ac = AudioCacher(cache_dir=cache_dir)
src_audio_path = ac.cache_audio_tensor(src_audio["waveform"].squeeze(0), src_audio["sample_rate"], filename_prefix="src_audio_")
audio_duration = librosa.get_duration(filename=src_audio_path)
repaint_start = -left_extend_length
repaint_end = audio_duration + right_extend_length
parameters = ast.literal_eval(parameters)
parameters["audio_duration"] = audio_duration
ap = AP(model_path)
audio_output = ap(
prompt=prompt,
lyrics=lyrics,
task="extend",
retake_seeds=retake_seeds,
src_audio_path=src_audio_path,
repaint_start=repaint_start,
repaint_end=repaint_end,
retake_variance=1.0,
**parameters)
audio, sr = audio_output[0][0].unsqueeze(0), audio_output[0][1]
ac.cleanup_file(src_audio_path)
if unload_model:
ap.cleanup()
return ({"waveform": audio, "sample_rate": sr},)
NODE_CLASS_MAPPINGS = {
"ACEStepGen": ACEStepGen,
"GenerationParameters": GenerationParameters,
"MultiLinePromptACES": MultiLinePromptACES,
"MultiLineLyrics": MultiLineLyrics,
"ACEStepRepainting": ACEStepRepainting,
"ACEStepEdit": ACEStepEdit,
"ACEStepExtend": ACEStepExtend,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ACEStepGen": "ACE-Step",
"GenerationParameters": "ACE-Step Parameters",
"MultiLinePromptACES": "ACE-Step Prompt",
"MultiLineLyrics": "ACE-Step Lyrics",
"ACEStepRepainting": "ACE-Step Repainting",
"ACEStepEdit": "ACE-Step Edit",
"ACEStepExtend": "ACE-Step Extend",
}
apg_guidance.py
+94
View File
@@ -0,0 +1,94 @@
import torch
class MomentumBuffer:
def __init__(self, momentum: float = -0.75):
self.momentum = momentum
self.running_average = 0
def update(self, update_value: torch.Tensor):
new_average = self.momentum * self.running_average
self.running_average = update_value + new_average
def project(
v0: torch.Tensor, # [B, C, H, W]
v1: torch.Tensor, # [B, C, H, W]
dims=[-1, -2],
):
dtype = v0.dtype
if v0.device.type == "mps":
v0, v1 = v0.float(), v1.float()
else:
v0, v1 = v0.double(), v1.double()
v1 = torch.nn.functional.normalize(v1, dim=dims)
v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
v0_orthogonal = v0 - v0_parallel
return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
def apg_forward(
pred_cond: torch.Tensor, # [B, C, H, W]
pred_uncond: torch.Tensor, # [B, C, H, W]
guidance_scale: float,
momentum_buffer: MomentumBuffer = None,
eta: float = 0.0,
norm_threshold: float = 2.5,
dims=[-1, -2],
):
diff = pred_cond - pred_uncond
if momentum_buffer is not None:
momentum_buffer.update(diff)
diff = momentum_buffer.running_average
if norm_threshold > 0:
ones = torch.ones_like(diff)
diff_norm = diff.norm(p=2, dim=dims, keepdim=True)
scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
diff = diff * scale_factor
diff_parallel, diff_orthogonal = project(diff, pred_cond, dims)
normalized_update = diff_orthogonal + eta * diff_parallel
pred_guided = pred_cond + (guidance_scale - 1) * normalized_update
return pred_guided
def cfg_forward(cond_output, uncond_output, cfg_strength):
return uncond_output + cfg_strength * (cond_output - uncond_output)
def cfg_double_condition_forward(
cond_output,
uncond_output,
only_text_cond_output,
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
def optimized_scale(positive_flat, negative_flat):
# Calculate dot production
dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
# Squared norm of uncondition
squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
# st_star = v_cond^T * v_uncond / ||v_uncond||^2
st_star = dot_product / squared_norm
return st_star
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.
else:
noise_pred = noise_pred_uncond * alpha + guidance_scale * (noise_pred_with_cond - noise_pred_uncond * alpha)
return noise_pred
data_sampler.py
+22
View File
@@ -0,0 +1,22 @@
import json
from pathlib import Path
import random
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
DEFAULT_ROOT_DIR = f"{current_dir}/examples/input_params"
class DataSampler:
def __init__(self, root_dir=DEFAULT_ROOT_DIR):
self.root_dir = root_dir
self.input_params_files = list(Path(self.root_dir).glob("*.json"))
def load_json(self, file_path):
with open(file_path, "r", encoding="utf-8") as f:
return json.load(f)
def sample(self):
json_path = random.choice(self.input_params_files)
json_data = self.load_json(json_path)
return json_data
examples/input_params/output_20250426071706_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "pop, rap, electronic, blues, hip-house, rhythm and blues",
"lyrics": "[verse]\n我走过深夜的街道\n冷风吹乱思念的漂亮外套\n你的微笑像星光很炫耀\n照亮了我孤独的每分每秒\n\n[chorus]\n愿你是风吹过我的脸\n带我飞过最远最遥远的山间\n愿你是风轻触我的梦\n停在心头不再飘散无迹无踪\n\n[verse]\n一起在喧哗避开世俗的骚动\n独自在天台探望月色的朦胧\n你说爱像音乐带点重节奏\n一拍一跳让我忘了心的温度多空洞\n\n[bridge]\n唱起对你的想念不隐藏\n像诗又像画写满藏不了的渴望\n你的影子挥不掉像风的倔强\n追着你飞扬穿越云海一样泛光\n\n[chorus]\n愿你是风吹过我的手\n暖暖的触碰像春日细雨温柔\n愿你是风盘绕我的身\n深情万万重不会有一天走远走\n\n[verse]\n深夜的钢琴弹起动人的旋律\n低音鼓砸进心底的每一次呼吸\n要是能将爱化作歌声传递\n你是否会听见我心里的真心实意",
"audio_duration": 170.63997916666668,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 3.191075086593628,
"diffusion": 17.459356784820557,
"latent2audio": 1.7095518112182617
},
"actual_seeds": [
3299954530
]
}
examples/input_params/output_20250426071812_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "country rock, folk rock, southern rock, bluegrass, country pop",
"lyrics": "[verse]\nWoke up to the sunrise glow\nTook my heart and hit the road\nWheels hummin' the only tune I know\nStraight to where the wildflowers grow\n\n[verse]\nGot that old map all wrinkled and torn\nDestination unknown but I'm reborn\nWith a smile that the wind has worn\nChasin' dreams that can't be sworn\n\n[chorus]\nRidin' on a highway to sunshine\nGot my shades and my radio on fine\nLeave the shadows in the rearview rhyme\nHeart's racing as we chase the time\n\n[verse]\nMet a girl with a heart of gold\nTold stories that never get old\nHer laugh like a tale that's been told\nA melody so bold yet uncontrolled\n\n[bridge]\nClouds roll by like silent ghosts\nAs we drive along the coast\nWe toast to the days we love the most\nFreedom's song is what we post\n\n[chorus]\nRidin' on a highway to sunshine\nGot my shades and my radio on fine\nLeave the shadows in the rearview rhyme\nHeart's racing as we chase the time",
"audio_duration": 224.23997916666667,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 4.262240648269653,
"diffusion": 15.380569219589233,
"latent2audio": 2.3227272033691406
},
"actual_seeds": [
401640
]
}
examples/input_params/output_20250426072346_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "hip-house, funk",
"lyrics": "[verse]\n哎呀跳起来，脚尖踩节拍 (oo-yeah!)\n灯光闪烁像星星盛开 (uh-huh!)\n人人都醒来，把烦恼踹开 (get it!)\n热血沸腾，汗水自己安排\n\n[chorus]\n嘿，你还等啥？快抓住节拍 (come on!)\n光芒指引，让心都不存在 (whoa!)\n点燃热火，我们一起飙high (let’s go!)\n跳入午夜的狂欢时代\n\n[bridge]\n咚咚鼓声啊，让你的灵魂起飞 (woo!)\n手心拍一拍，能量翻倍 (ah-hah!)\n键盘响起来，如宇宙的交汇 (oh yeah!)\n就是这感觉，兄弟姐妹都陶醉\n\n[verse]\n灵魂从不睡，只想继续燃烧 (woo!)\n节奏像热浪，席卷这街道 (ow!)\n大伙儿涌上楼台，满面微笑 (yeah!)\n这一刻属于我们，无可替代\n\n[chorus]\n嘿，你还等啥？快抓住节拍 (come on!)\n光芒指引，让心都不存在 (whoa!)\n点燃热火，我们一起飙high (let’s go!)\n跳入午夜的狂欢时代\n\n[verse]\n世界多精彩，握紧把它打开 (alright!)\n每一步都像星球在摇摆 (uh-huh!)\n无边无际的律动像大海 (oo-yeah!)\n跟着光芒之舞，一起澎湃",
"audio_duration": 204.19997916666668,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.05196118354797363,
"diffusion": 15.530808210372925,
"latent2audio": 2.5604095458984375
},
"actual_seeds": [
401640
]
}
examples/input_params/output_20250426072508_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "funk, pop, soul, rock, melodic, guitar, drums, bass, keyboard, percussion, 105 BPM, energetic, upbeat, groovy, vibrant, dynamic",
"lyrics": "[verse]\nNeon lights they flicker bright\nCity hums in dead of night\nRhythms pulse through concrete veins\nLost in echoes of refrains\n\n[verse]\nBassline groovin' in my chest\nHeartbeats match the city's zest\nElectric whispers fill the air\nSynthesized dreams everywhere\n\n[chorus]\nTurn it up and let it flow\nFeel the fire let it grow\nIn this rhythm we belong\nHear the night sing out our song\n\n[verse]\nGuitar strings they start to weep\nWake the soul from silent sleep\nEvery note a story told\nIn this night we’re bold and gold\n\n[bridge]\nVoices blend in harmony\nLost in pure cacophony\nTimeless echoes timeless cries\nSoulful shouts beneath the skies\n\n[verse]\nKeyboard dances on the keys\nMelodies on evening breeze\nCatch the tune and hold it tight\nIn this moment we take flight",
"audio_duration": 178.87997916666666,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.02882218360900879,
"diffusion": 16.91233205795288,
"latent2audio": 1.7794082164764404
},
"actual_seeds": [
401640
]
}
examples/input_params/output_20250426073829_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "electronic rap",
"lyrics": "[verse]\nWaves on the bass, pulsing in the speakers,\nTurn the dial up, we chasing six-figure features,\nGrinding on the beats, codes in the creases,\nDigital hustler, midnight in sneakers.\n\n[chorus]\nElectro vibes, hearts beat with the hum,\nUrban legends ride, we ain't ever numb,\nCircuits sparking live, tapping on the drum,\nLiving on the edge, never succumb.\n\n[verse]\nSynthesizers blaze, city lights a glow,\nRhythm in the haze, moving with the flow,\nSwagger on stage, energy to blow,\nFrom the blocks to the booth, you already know.\n\n[bridge]\nNight's electric, streets full of dreams,\nBass hits collective, bursting at seams,\nHustle perspective, all in the schemes,\nRise and reflective, ain't no in-betweens.\n\n[verse]\nVibin' with the crew, sync in the wire,\nGot the dance moves, fire in the attire,\nRhythm and blues, soul's our supplier,\nRun the digital zoo, higher and higher.\n\n[chorus]\nElectro vibes, hearts beat with the hum,\nUrban legends ride, we ain't ever numb,\nCircuits sparking live, tapping on the drum,\nLiving on the edge, never succumb.",
"audio_duration": 221.42547916666666,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.024875164031982422,
"diffusion": 20.566852569580078,
"latent2audio": 2.2281734943389893
},
"actual_seeds": [
401640
]
}
examples/input_params/output_20250426074037_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "electronic, house, electro house, synthesizer, drums, bass, percussion, fast, energetic, uplifting, exciting",
"lyrics": "[verse]\n霓虹灯下我们追逐\n人群跃动像潮水满布\n热浪袭来吹散孤独\n跳进节奏不如停下脚步\n\n[pre-chorus]\n脚尖触电快点感受\n迎着风声释放自由\n心跳节拍配合节奏\n一切烦恼请靠边游\n\n[chorus]\n夏夜狂奔没有尽头\n星光闪烁舞池不朽\n尽情挥洒所有节奏\n无边热情把你包裹哦\n\n[verse]\n天空翻滚黑云入夜\n每颗星星像音乐律贴\n耳边回响那低音线\n环绕耳际如梦境般甜\n\n[pre-chorus]\n脚尖触电快点感受\n迎着风声释放自由\n心跳节拍配合节奏\n一切烦恼请靠边游\n\n[chorus]\n夏夜狂奔没有尽头\n星光闪烁舞池不朽\n尽情挥洒所有节奏\n无边热情把你包裹哦",
"audio_duration": 221.47997916666668,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.028400182723999023,
"diffusion": 13.195815324783325,
"latent2audio": 2.1679723262786865
},
"actual_seeds": [
3440445703
]
}
examples/input_params/output_20250426074214_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "synth-pop, electronic, pop, synthesizer, drums, bass, piano, 128 BPM, energetic, uplifting, modern",
"lyrics": "[verse]\nWoke up in a city that's always alive\nNeon lights they shimmer they thrive\nElectric pulses beat they drive\nMy heart races just to survive\n\n[chorus]\nOh electric dreams they keep me high\nThrough the wires I soar and fly\nMidnight rhythms in the sky\nElectric dreams together we’ll defy\n\n[verse]\nLost in the labyrinth of screens\nVirtual love or so it seems\nIn the night the city gleams\nDigital faces haunted by memes\n\n[chorus]\nOh electric dreams they keep me high\nThrough the wires I soar and fly\nMidnight rhythms in the sky\nElectric dreams together we’ll defy\n\n[bridge]\nSilent whispers in my ear\nPixelated love serene and clear\nThrough the chaos find you near\nIn electric dreams no fear\n\n[verse]\nBound by circuits intertwined\nLove like ours is hard to find\nIn this world we’re truly blind\nBut electric dreams free the mind",
"audio_duration": 221.27997916666666,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.025463581085205078,
"diffusion": 15.243804454803467,
"latent2audio": 2.170398473739624
},
"actual_seeds": [
3400270027
]
}
examples/input_params/output_20250426074413_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "Cuban music, salsa, son, Afro-Cuban, traditional Cuban",
"lyrics": "[verse]\nSun dips low the night ignites\nBassline hums with gleaming lights\nElectric guitar singing tales so fine\nIn the rhythm we all intertwine\n\n[verse]\nDrums beat steady calling out\nPercussion guides no room for doubt\nElectric pulse through every vein\nDance away every ounce of pain\n\n[chorus]\nFeel the rhythm feel the flow\nLet the music take control\nBassline deep electric hum\nIn this night we're never numb\n\n[bridge]\nStars above they start to glow\nEchoes of the night's soft glow\nElectric strings weave through the air\nIn this moment none compare\n\n[verse]\nHeartbeats sync with every tone\nLost in music never alone\nElectric tales of love and peace\nIn this groove we find release\n\n[chorus]\nFeel the rhythm feel the flow\nLet the music take control\nBassline deep electric hum\nIn this night we're never numb",
"audio_duration": 208.27997916666666,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.026132583618164062,
"diffusion": 15.139378070831299,
"latent2audio": 2.2071540355682373
},
"actual_seeds": [
3358899399
]
}
examples/input_params/output_20250426075107_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "pop, piano, rap, dark, atmospheric",
"lyrics": "[verse]\n月光爬上窗 染白冷的床\n心跳的方向 带我入迷惘\n黑夜吞噬光 命运的纸张\n爱是血色霜 邪恶又芬芳\n\n[chorus]\n你是猎人的欲望 我是迷途的小羊\n深陷你眼眸的荒 唐突献出心脏\n我在夜里回荡 是谁给我希望\n黑暗风中飘荡 假装不再受伤\n\n[verse]\n心锁在门外 谁会解开关怀\n温柔的手拍 藏着冷酷杀害\n思绪如尘埃 撞击爱的霹雳\n灵魂的独白 为你沾满血迹\n\n[bridge]\n你是噩梦的歌唱 是灵魂的捆绑\n绝望中带着光 悬崖边的渴望\n心跳被你鼓掌 恶魔也痴痴想\n渐渐没了抵抗 古老诡计流淌\n\n[chorus]\n你是猎人的欲望 我是迷途的小羊\n深陷你眼眸的荒 唐突献出心脏\n我在夜里回荡 是谁给我希望\n黑暗风中飘荡 假装不再受伤\n\n[outro]\n爱如月黑无光 渗进梦的战场\n逃入无声的场 放手或心嚷嚷\n隐秘的极端 爱是极致风浪\n灵魂彻底交偿 你是终极虚妄",
"audio_duration": 146.91997916666668,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.03876018524169922,
"diffusion": 15.962624549865723,
"latent2audio": 1.4594337940216064
},
"actual_seeds": [
2065110378
]
}
examples/input_params/output_20250426075537_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "surf music",
"lyrics": "[verse]\nSunshine on the boulevard the beach is calling loud\nWaves are dancing golden sand under a cotton cloud\nElectric heartbeat pounding fast the tide is on our side\nCatch a wave and feel alive we’ll take it for a ride\n\n[verse]\nPalm trees swaying left to right they know where we belong\nFeel the rhythm of the night it keeps us moving strong\nSea spray kisses salty air we’re flying with the breeze\nChampagne states of mind we ride we do just as we please\n\n[chorus]\nWe’re riding waves of life together hand in hand\nWith every beat we chase the beat it’s our own wonderland\nFeel the music take you higher as the shorelines blur\nThis is our world our endless summer as we live and learn\n\n[bridge]\nMoonlight paints the ocean blue reflections in our eyes\nStars align to light our path we’re surfing through the skies\nEvery moment like a song we sing it loud and clear\nEvery day’s a new adventure with you always near\n\n[verse]\nNeon lights and city sounds they blend with ocean views\nWe’re unstoppable tonight no way that we can lose\nDreams are written in the sand they sparkle in the sun\nTogether we’re a masterpiece our story’s just begun\n\n[chorus]\nWe’re riding waves of life together hand in hand\nWith every beat we chase the beat it’s our own wonderland\nFeel the music take you higher as the shorelines blur\nThis is our world our endless summer as we live and learn",
"audio_duration": 236.55997916666666,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.033666133880615234,
"diffusion": 16.291455507278442,
"latent2audio": 2.3726775646209717
},
"actual_seeds": [
508630535
]
}
examples/input_params/output_20250426075843_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "alternative rock, pop, rock",
"lyrics": "[verse]\nBright lights flashing in the city sky\nRunning fast and we don't know why\nElectric nights got our hearts on fire\nChasing dreams we'll never tire\n\n[verse]\nGrit in our eyes wind in our hair\nBreaking rules we don't even care\nShouting loud above the crowd\nLiving life like we're unbowed\n\n[chorus]\nRunning wild in the night so free\nFeel the beat pumping endlessly\nHearts collide in the midnight air\nWe belong we don't have a care\n\n[verse]\nPiercing through like a lightning strike\nEvery moment feels like a hike\nDaring bold never backing down\nKings and queens without a crown\n\n[chorus]\nRunning wild in the night so free\nFeel the beat pumping endlessly\nHearts collide in the midnight air\nWe belong we don't have a care\n\n[bridge]\nClose your eyes let your spirit soar\nWe are the ones who wanted more\nBreaking chains of the mundane\nIn this world we'll make our claim",
"audio_duration": 202.19997916666668,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.02512216567993164,
"diffusion": 18.860822677612305,
"latent2audio": 2.0361969470977783
},
"actual_seeds": [
1255121549
]
}
examples/input_params/output_20250426080234_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "rock, hip - hop, orchestral, bass, drums, electric guitar, piano, synthesizer, violin, viola, cello, fast, energetic, motivational, inspirational, empowering",
"lyrics": "### **[Intro – Spoken]** \n*\"The streets whisper, their echoes never fade. \nEvery step I take leaves a mark—this ain't just a game.\"* \n\n### **[Hook/Chorus]** \nBorn in the chaos, I weather the storm, \nRising from ashes where warriors are born. \nChains couldn't hold me, the system’s a maze, \nI rewrite the rules, set the city ablaze! \n\n### **[Verse 1]** \nCold nights, empty pockets, dreams laced with fight, \nEvery loss made me sharper, cut deep like a knife. \nThey said I wouldn’t make it, now they watch in despair, \nFrom the curb to the throne, took the pain, made it rare. \nEvery siren’s a melody, every alley holds a tale, \nRose from the shadows, left my name on the trail. \nStreetlights flicker like warnings in the haze, \nBut I move like a phantom, unfazed by the blaze. \n\n### **[Hook/Chorus]** \nBorn in the chaos, I weather the storm, \nRising from ashes where warriors are born. \nChains couldn't hold me, the system’s a maze, \nI rewrite the rules, set the city ablaze! \n\n### **[Verse 2]** \nBarbed wire fences couldn't lock in my mind, \nEvery cage they designed, I left broken behind. \nThey want control, but I’m destined to roam, \nWhere the lost find their voice, where the heart sets the tone. \nSteel and concrete, where the lessons run deep, \nEvery crack in the pavement tells a story of heat. \nBut I rise, undefeated, like a king with no throne, \nWriting scripts in the struggle, my legacy’s stone. \n\n### **[Bridge]** \nFeel the rhythm of the underground roar, \nEvery wound tells a story of the battles before. \nBlood, sweat, and echoes fill the cold midnight, \nBut we move with the fire—unshaken, upright. \n\n### **[Verse 3]** \nNo regrets, no retreat, this game has no pause, \nEvery step that I take is a win for the lost. \nI took lessons from hustlers, wisdom from pain, \nNow the echoes of struggle carve power in my name. \nThey built walls, but I walk through the cracks, \nTurned dirt into gold, never looked back. \nThrough the struggle we rise, through the fire we claim, \nThis is more than just music—it's life in the frame. \n\n### **[Hook/Chorus – Reprise]** \nBorn in the chaos, I weather the storm, \nRising from ashes where warriors are born. \nChains couldn't hold me, the system’s a maze, \nI rewrite the rules, set the city ablaze! \n\n### **[Outro – Spoken]** \n*\"The scars, the struggle, the grind—it’s all part of the rhythm. \nWe never break, we never fold. We rise.\"*",
"audio_duration": 153.95997916666667,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.04368758201599121,
"diffusion": 17.16369390487671,
"latent2audio": 1.5405471324920654
},
"actual_seeds": [
2659225017
]
}
examples/input_params/output_20250426080407_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "tango finlandés, campanas, disco, dark pop, electro, guitarra clásica, corridos tumba",
"lyrics": "[inst]",
"audio_duration": 162.79997916666667,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.011058568954467773,
"diffusion": 9.924944400787354,
"latent2audio": 1.6034839153289795
},
"actual_seeds": [
780297686
]
}
examples/input_params/output_20250426080601_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "Nightclubs, dance parties, workout playlists, radio broadcasts",
"lyrics": "Burning in motion, set me alight!\nEvery heartbeat turns into a fight!\nCaged in rhythm, chained in time!\nLove’s a battle— You're Mine! You're Mine!",
"audio_duration": 221.83997916666667,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.012485980987548828,
"diffusion": 14.345409154891968,
"latent2audio": 2.174558639526367
},
"actual_seeds": [
1318394052
]
}
examples/input_params/output_20250426081134_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "melancholic, world, sad, medieval, soulful",
"lyrics": "[Verse]\nIn a world so grand he roams the skies alone\nHis heart a heavy stone a tale untold\nWhispers of his past echo through the night\nA lonely dragon searching for the light\n\n[Verse 2]\nOnce a mighty force now he drifts in pain\nHis scales once shimmered now they're dark with shame\nCast out by his kin in shadows he does hide\nA haunting sorrow burns deep inside\n\n[Chorus]\nRoaming endless fields with no friend in sight\nHis roar a mournful cry beneath the moon's pale light\nTears fall like stars as he flies on his way\nA lonely dragon yearning for the break of day\n\n[Bridge]\nThe world turns cold the nights grow long\nIn his heart he carries an ancient song\nOf battles fought and love long gone\nA legend now but his soul is torn\n\n[Verse 3]\nHoping for a day he'll find a kindred soul\nTo share his pain and make him whole\nTill then he drifts a shadow in the sky\nA lonely dragon with tears in his eye\n\n[Chorus]\nRoaming endless fields with no friend in sight\nHis roar a mournful cry beneath the moon's pale light\nTears fall like stars as he flies on his way\nA lonely dragon yearning for the break of day",
"audio_duration": 239.99997916666666,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.029100656509399414,
"diffusion": 22.503791570663452,
"latent2audio": 2.3603708744049072
},
"actual_seeds": [
2166832218
]
}
examples/input_params/output_20250426091716_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "anime, cute female vocals, kawaii pop, j-pop, childish, piano, guitar, synthesizer, fast, happy, cheerful, lighthearted",
"lyrics": "[Chorus]\nねぇ、顔が赤いよ？\nどうしたの？ 熱があるの？\nそれとも怒ってるの？\nねぇ、言ってよ！\n\nどうしてそんな目で見るの？\n私、悪いことした？\n何か間違えたの？\nお願い、やめて… 怖いから…\nだから、やめてよ…\n\n[Bridge]\n目を閉じて、くるっと背を向けて、\n何も見なかったフリするから、\n怒らないで… 許してよ…\n\n[Chorus]\nねぇ、顔が赤いよ？\nどうしたの？ 熱があるの？\nそれとも怒ってるの？\nねぇ、言ってよ！\n\nどうしてそんな目で見るの？\n私、悪いことした？\n何か間違えたの？\nお願い、やめて… 怖いから…\nだから、やめてよ…\n\n[Bridge 2]\n待って、もし私が悪いなら、\nごめんなさいって言うから、\nアイスクリームあげるから、\nもう怒らないで？\n\nOoooh… 言ってよ！",
"audio_duration": 160,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.0282442569732666,
"diffusion": 12.104875326156616,
"latent2audio": 1.587641954421997
},
"actual_seeds": [
4028738662
]
}
examples/input_params/output_20250426092025_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "dark, death rock, metal, hardcore, electric guitar, powerful, bass, drums, 110 bpm, G major",
"lyrics": "[Verse]\nMy lovers betray me\nThe snake in my garden is hissing\nIn the air is the sweetness of roses\nAnd under my skin\nThere's a thorn\n\n[Verse 2]\nI should have known\nThat God sends his angel in shadows\nWith blood in his veins\nI watch the enemy\nGivin' me the hand of my savior\n\n[Chorus]\nAnd I can't love again\nWith the echo of your name in my head\nWith the demons in my bed\nWith the memories\nYour ghost\nI see it\n'Cause it comes to haunt me\nJust to taunt me\nIt comes to haunt me\nJust to taunt me\n\n[Verse 3]\nWith sugar and spice\nIt's hard to ignore the nostalgia\nWith the men on their knees\nAt the gates of my heart\nHow they beg me\n\n[Verse 4]\nThey say\n\"No one will ever love you\nThe way that I do\nNo one will ever touch you\nThe way that I do\"\n\n[Chorus]\nAnd I can't love again\nWith the echo of your name in my head\nWith the demons in my bed\nWith the memories\nYour ghost\nI see it\n'Cause it comes to haunt me\nJust to taunt me\nIt comes to haunt me\nJust to taunt me",
"audio_duration": 174.27997916666666,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 3.8372838497161865,
"diffusion": 13.039669275283813,
"latent2audio": 1.7923030853271484
},
"actual_seeds": [
4064916393
]
}
examples/input_params/output_20250426093007_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "aggressive, Heavy Riffs, Blast Beats, Satanic Black Metal",
"lyrics": "[verse]\nFloating through the galaxy on a midnight ride\nStars are dancing all around in cosmic tides\nFeel the pulse of space and time beneath our feet\nEvery beat a heartbeat in this endless suite\n\n[chorus]\nGalactic dreams under neon lights\nSailing through the velvet nights\nWe are echoes in a cosmic sea\nIn a universe where we are free\n\n[verse]\nPlanetary whispers in the sky tonight\nEvery constellation's got a secret sight\nDistant worlds and moons we have yet to see\nIn the void of space where we can just be\n\n[bridge]\nAsteroids and comets in a ballet they spin\nLost in the rhythm of where our dreams begin\nClose your eyes and let the synths take flight\nWe're voyagers on an electric night\n\n[verse]\nLet the piano keys unlock the stars above\nEvery chord a memory every note is love\nIn this synth symphony we find our grace\nDrifting forever in this boundless space\n\n[chorus]\nGalactic dreams under neon lights\nSailing through the velvet nights\nWe are echoes in a cosmic sea\nIn a universe where we are free",
"audio_duration": 181.99997916666666,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.025065898895263672,
"diffusion": 17.176705837249756,
"latent2audio": 1.8225171566009521
},
"actual_seeds": [
1132623236
]
}
examples/input_params/output_20250426093146_0_input_params.json
+25
View File
@@ -0,0 +1,25 @@
{
"prompt": "r&b, soul, funk/soul",
"lyrics": "[verse]\nDancing through electric fires\nHeart is buzzing like live wires\nIn your arms I find desire\nFeel the beat as we get higher\n\n[chorus]\nElectric love in the night sky\nWe’re gonna soar baby you and I\nDrop the bass let the rhythm fly\nFeel the heat and don't ask why\n\n[verse]\nWhisper secrets that make me blush\nUnder the neon city hush\nYour touch gives me such a rush\nTurn it up we're feeling lush\n\n[chorus]\nElectric love in the night sky\nWe’re gonna soar baby you and I\nDrop the bass let the rhythm fly\nFeel the heat and don't ask why\n\n[bridge]\nThrough the lights and the smoky haze\nI see you in a thousand ways\nLove's a script and we’re the play\nTurn the page stay till we sway\n\n[chorus]\nElectric love in the night sky\nWe’re gonna soar baby you and I\nDrop the bass let the rhythm fly\nFeel the heat and don't ask why",
"audio_duration": 195.15997916666666,
"infer_step": 60,
"guidance_scale": 15,
"scheduler_type": "euler",
"cfg_type": "apg",
"omega_scale": 10,
"guidance_interval": 0.5,
"guidance_interval_decay": 0,
"min_guidance_scale": 3,
"use_erg_tag": true,
"use_erg_lyric": true,
"use_erg_diffusion": true,
"oss_steps": [],
"timecosts": {
"preprocess": 0.025553464889526367,
"diffusion": 18.250118494033813,
"latent2audio": 1.9400627613067627
},
"actual_seeds": [
2853131993
]
}
images/2025-05-07_19-53-51.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 117 KiB

images/2025-05-07_19-59-22.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 154 KiB

images/2025-05-07_20-04-02.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 145 KiB

images/2025-05-07_20-09-52.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 142 KiB

language_segmentation/LangSegment.py
+866
View File
@@ -0,0 +1,866 @@
"""
This file bundles language identification functions.
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.
See LICENSE file for more info.
https://github.com/adbar/py3langid
Projects:
https://github.com/juntaosun/LangSegment
"""
import os
import re
import sys
import numpy as np
from collections import Counter
from collections import defaultdict
# import langid
# import py3langid as langid
# pip install py3langid==0.2.2
# 启用语言预测概率归一化，概率预测的分数。因此，实现重新规范化 产生 0-1 范围内的输出。
# langid disables probability normalization by default. For command-line usages of , it can be enabled by passing the flag.
# 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
try:from .utils.num import num2str
except ImportError:
try:from utils.num import num2str
except ImportError as e:
raise e
# -----------------------------------
# 更新日志：新版本分词更加精准。
# Changelog: The new version of the word segmentation is more accurate.
# チェンジログ:新しいバージョンの単語セグメンテーションはより正確です。
# Changelog: 분할이라는 단어의 새로운 버전이 더 정확합니다.
# -----------------------------------
# Word segmentation function:
# automatically identify and split the words (Chinese/English/Japanese/Korean) in the article or sentence according to different languages,
# making it more suitable for TTS processing.
# This code is designed for front-end text multi-lingual mixed annotation distinction, multi-language mixed training and inference of various TTS projects.
# This processing result is mainly for (Chinese = zh, Japanese = ja, English = en, Korean = ko), and can actually support up to 97 different language mixing processing.
#===========================================================================================================
#分かち書き機能:文章や文章の中の例えば（中国語/英語/日本語/韓国語）を、異なる言語で自動的に認識して分割し、TTS処理により適したものにします。
#このコードは、さまざまなTTSプロジェクトのフロントエンドテキストの多言語混合注釈区別、多言語混合トレーニング、および推論のために特別に作成されています。
#===========================================================================================================
#(1)自動分詞:「韓国語では何を読むのですかあなたの体育の先生は誰ですか?今回の発表会では、iPhone 15シリーズの4機種が登場しました」
#（2）手动分词:“あなたの名前は<ja>佐々木ですか?<ja>ですか?”
#この処理結果は主に（中国語=ja、日本語=ja、英語=en、韓国語=ko）を対象としており、実際には最大97の異なる言語の混合処理をサポートできます。
#===========================================================================================================
#===========================================================================================================
# 단어 분할 기능: 기사 또는 문장에서 단어(중국어/영어/일본어/한국어)를 다른 언어에 따라 자동으로 식별하고 분할하여 TTS 처리에 더 적합합니다.
# 이 코드는 프런트 엔드 텍스트 다국어 혼합 주석 분화, 다국어 혼합 교육 및 다양한 TTS 프로젝트의 추론을 위해 설계되었습니다.
#===========================================================================================================
# (1) 자동 단어 분할: "한국어로 무엇을 읽습니까? 스포츠 씨? 이 컨퍼런스는 4개의 iPhone 15 시리즈 모델을 제공합니다."
# (2) 수동 참여: "이름이 <ja>Saki입니까? <ja>?"
# 이 처리 결과는 주로 (중국어 = zh, 일본어 = ja, 영어 = en, 한국어 = ko)를 위한 것이며 실제로 혼합 처리를 위해 최대 97개의 언어를 지원합니다.
#===========================================================================================================
# ===========================================================================================================
# 分词功能：将文章或句子里的例如（中/英/日/韩），按不同语言自动识别并拆分，让它更适合TTS处理。
# 本代码专为各种 TTS 项目的前端文本多语种混合标注区分，多语言混合训练和推理而编写。
# ===========================================================================================================
# （1）自动分词：“韩语中的오빠读什么呢？あなたの体育の先生は誰ですか? 此次发布会带来了四款iPhone 15系列机型”
# （2）手动分词：“你的名字叫<ja>佐々木？<ja>吗？”
# 本处理结果主要针对（中文=zh , 日文=ja , 英文=en , 韩语=ko）, 实际上可支持多达 97 种不同的语言混合处理。
# ===========================================================================================================
# 手动分词标签规范：<语言标签>文本内容</语言标签>
# 수동 단어 분할 태그 사양: <언어 태그> 텍스트 내용</언어 태그>
# Manual word segmentation tag specification: <language tags> text content </language tags>
# 手動分詞タグ仕様:<言語タグ>テキスト内容</言語タグ>
# ===========================================================================================================
# For manual word segmentation, labels need to appear in pairs, such as:
# 如需手动分词，标签需要成对出现，例如：“<ja>佐々木<ja>” 或者 “<ja>佐々木</ja>”
# 错误示范：“你的名字叫<ja>佐々木。” 此句子中出现的单个<ja>标签将被忽略，不会处理。
# Error demonstration: "Your name is <ja>佐々木。" Single <ja> tags that appear in this sentence will be ignored and will not be processed.
# ===========================================================================================================
# ===========================================================================================================
# 语音合成标记语言 SSML , 这里只支持它的标签（非 XML）Speech Synthesis Markup Language SSML, only its tags are supported here (not XML)
# 想支持更多的 SSML 标签？欢迎 PR！ Want to support more SSML tags? PRs are welcome!
# 说明：除了中文以外，它也可改造成支持多语种 SSML ，不仅仅是中文。
# Note: In addition to Chinese, it can also be modified to support multi-language SSML, not just Chinese.
# ===========================================================================================================
# 中文实现：Chinese implementation:
# 【SSML】<number>=中文大写数字读法（单字）
# 【SSML】<telephone>=数字转成中文电话号码大写汉字（单字）
# 【SSML】<currency>=按金额发音。
# 【SSML】<date>=按日期发音。支持 2024年08月24, 2024/8/24, 2024-08, 08-24, 24 等输入。
# ===========================================================================================================
class LangSSML:
def __init__(self):
# 纯数字
self._zh_numerals_number = {
'0': '零',
'1': '一',
'2': '二',
'3': '三',
'4': '四',
'5': '五',
'6': '六',
'7': '七',
'8': '八',
'9': '九'
}
# 将2024/8/24, 2024-08, 08-24, 24 标准化“年月日”
# Standardize 2024/8/24, 2024-08, 08-24, 24 to "year-month-day"
def _format_chinese_data(self, date_str:str):
# 处理日期格式
input_date = date_str
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(' ')
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)
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('-')
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 # 月-日
elif len(parts[0]) > 0: # 仅有月-日或年
if len(parts[0]) == 4:
year = parts[0]
else: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)
time_arrs = time_str.split(":")
hours, minutes, seconds = "","",""
if len(time_arrs) == 3: # H/M/S
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
if len(time_arrs) > 1:
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+)'
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))
else:output += item
output = output.replace(".","点")
return output
# 【SSML】telephone=数字转成中文电话号码大写汉字（单字）
# Convert numbers to Chinese phone numbers in uppercase Chinese characters(single word)
def to_chinese_telephone(self, num:str):
output = self.to_chinese_number(num.replace("+86","")) # zh +86
output = output.replace("一","幺")
return output
# 【SSML】currency=按金额发音。
# Digital processing from GPT_SoVITS num.py （thanks）
def to_chinese_currency(self, num:str):
pattern = r'(\d+)'
arrs = re.split(pattern, num)
output = ""
for item in arrs:
if re.match(pattern,item):
output += num2str(item)
else:output += item
output = output.replace(".","点")
return output
# 【SSML】date=按日期发音。支持 2024年08月24, 2024/8/24, 2024-08, 08-24, 24 等输入。
def to_chinese_date(self, num:str):
chinese_date = self._format_chinese_data(num)
return chinese_date
class LangSegment:
def __init__(self):
self.langid = LanguageIdentifier.from_pickled_model(MODEL_FILE, norm_probs=True)
self._text_cache = None
self._text_lasts = None
self._text_langs = None
self._lang_count = None
self._lang_eos = None
# 可自定义语言匹配标签：カスタマイズ可能な言語対応タグ:사용자 지정 가능한 언어 일치 태그:
# 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|-]*>)'
# 语言过滤组功能, 可以指定保留语言。不在过滤组中的语言将被清除。您可随心搭配TTS语音合成所支持的语言。
# 언어 필터 그룹 기능을 사용하면 예약된 언어를 지정할 수 있습니다. 필터 그룹에 없는 언어는 지워집니다. TTS 텍스트에서 지원하는 언어를 원하는 대로 일치시킬 수 있습니다.
# 言語フィルターグループ機能では、予約言語を指定できます。フィルターグループに含まれていない言語はクリアされます。TTS音声合成がサポートする言語を自由に組み合わせることができます。
# The language filter group function allows you to specify reserved languages.
# Languages not in the filter group will be cleared. You can match the languages supported by TTS Text To Speech as you like.
# 排名越前，优先级越高，The higher the ranking, the higher the priority，ランキングが上位になるほど、優先度が高くなります。
# 系统默认过滤器。System default filter。(ISO 639-1 codes given)
# ----------------------------------------------------------------------------------------------------------------------------------
# "zh"中文=Chinese ,"en"英语=English ,"ja"日语=Japanese ,"ko"韩语=Korean ,"fr"法语=French ,"vi"越南语=Vietnamese , "ru"俄语=Russian
# "th"泰语=Thai
# ----------------------------------------------------------------------------------------------------------------------------------
self.DEFAULT_FILTERS = ["zh", "ja", "ko", "en"]
# 用户可自定义过滤器。User-defined filters
self.Langfilters = self.DEFAULT_FILTERS[:] # 创建副本
# 合并文本
self.isLangMerge = True
# 试验性支持：您可自定义添加："fr"法语 , "vi"越南语。Experimental: You can customize to add: "fr" French, "vi" Vietnamese.
# 请使用API启用：self.setfilters(["zh", "en", "ja", "ko", "fr", "vi" , "ru" , "th"]) # 您可自定义添加，如："fr"法语 , "vi"越南语。
# 预览版功能，自动启用或禁用，无需设置
# Preview feature, automatically enabled or disabled, no settings required
self.EnablePreview = False
# 除此以外，它支持简写过滤器，只需按不同语种任意组合即可。
# In addition to that, it supports abbreviation filters, allowing for any combination of different languages.
# 示例：您可以任意指定多种组合，进行过滤
# Example: You can specify any combination to filter
# 中/日语言优先级阀值（评分范围为 0 ~ 1）:评分低于设定阀值 <0.89 时，启用 filters 中的优先级。\n
# 중/일본어 우선 순위 임계값(점수 범위 0-1): 점수가 설정된 임계값 <0.89보다 낮을 때 필터에서 우선 순위를 활성화합니다.
# 中国語/日本語の優先度しきい値（スコア範囲0〜1）:スコアが設定されたしきい値<0.89未満の場合、フィルターの優先度が有効になります。\n
# Chinese and Japanese language priority threshold (score range is 0 ~ 1): The default threshold is 0.89. \n
# Only the common characters between Chinese and Japanese are processed with confidence and priority. \n
self.LangPriorityThreshold = 0.89
# Langfilters = ["zh"] # 按中文识别
# Langfilters = ["en"] # 按英文识别
# Langfilters = ["ja"] # 按日文识别
# Langfilters = ["ko"] # 按韩文识别
# Langfilters = ["zh_ja"] # 中日混合识别
# Langfilters = ["zh_en"] # 中英混合识别
# Langfilters = ["ja_en"] # 日英混合识别
# Langfilters = ["zh_ko"] # 中韩混合识别
# Langfilters = ["ja_ko"] # 日韩混合识别
# Langfilters = ["en_ko"] # 英韩混合识别
# Langfilters = ["zh_ja_en"] # 中日英混合识别
# Langfilters = ["zh_ja_en_ko"] # 中日英韩混合识别
# 更多过滤组合，请您随意。。。For more filter combinations, please feel free to......
# より多くのフィルターの組み合わせ、お気軽に。。。더 많은 필터 조합을 원하시면 자유롭게 해주세요. .....
# 可选保留：支持中文数字拼音格式，更方便前端实现拼音音素修改和推理，默认关闭 False 。
# 开启后 True ，括号内的数字拼音格式均保留，并识别输出为："zh"中文。
self.keepPinyin = False
# DEFINITION
self.PARSE_TAG = re.compile(r'(⑥\$*\d+[\d]{6,}⑥)')
self.LangSSML = LangSSML()
def _clears(self):
self._text_cache = None
self._text_lasts = None
self._text_langs = None
self._text_waits = None
self._lang_count = None
self._lang_eos = None
def _is_english_word(self, word):
return bool(re.match(r'^[a-zA-Z]+$', word))
def _is_chinese(self, word):
for char in word:
if '\u4e00' <= char <= '\u9fff':
return True
return False
def _is_japanese_kana(self, word):
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 = modified_text.strip('-')
return modified_text + " "
def _split_camel_case(self, word):
return re.sub(r'(?<!^)(?=[A-Z])', ' ', word)
def _statistics(self, language, text):
# Language word statistics:
# Chinese characters usually occupy double bytes
if self._lang_count is None or not isinstance(self._lang_count, defaultdict):
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)
self._lang_count = lang_count
def _clear_text_number(self, text):
if text == "\n":return text,False # Keep Line Breaks
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):
# Pre-detection
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
elif preData is not None and preData["symbol"] is None:
if len(clear_text) == 0:language = preData["lang"]
elif is_number == True:language = preData["lang"]
_ , pre_is_number = self._clear_text_number(preData["text"])
if (preData["lang"] == language):
self._statistics(preData["lang"],text)
text = preData["text"] + text
preData["text"] = text
return preData
elif pre_is_number == True:
text = f'{preData["text"]}{text}'
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
data = {"lang":language,"text": text,"score":score,"symbol":symbol}
filters = self.Langfilters
if filters is None 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
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)
# 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 and ("|" in preResult["lang"]):
pre_lang = preResult["lang"]
if language in pre_lang:preResult["lang"] = language = language.split("|")[0]
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: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"])
return (None,"")
def _match_ending(self, input , index):
if input is None or len(input) == 0: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'(.)\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
arrs = self._split_camel_case(text).split(" ")
langs = []
for t in arrs:
if len(t.strip()) <= 3:continue
language, score = self.langid.classify(t)
langs.append({"lang":language})
if len(langs) == 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"):
score = score.item()
score = round(score , 3)
return language, score
def _get_filters_string(self):
filters = self.Langfilters
return "-".join(filters).lower().strip() if filters is not None else ""
def _parse_language(self, words , segment):
LANG_JA = "ja"
LANG_ZH = "zh"
LANG_ZH_JA = f'{LANG_ZH}|{LANG_JA}'
LANG_JA_ZH = f'{LANG_JA}|{LANG_ZH}'
language = LANG_ZH
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
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
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,}⑥)')
cleans_text = re.sub(number_tags, '' ,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}'
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 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)
if index_ja != -1 and index_ja < index_zh:language = LANG_JA
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
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
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)
# ----------------------------------------------------------
# 【SSML】中文数字处理：Chinese Number Processing (SSML support)
# 这里默认都是中文，用于处理 SSML 中文标签。当然可以支持任意语言，例如：
# The default here is Chinese, which is used to process SSML Chinese tags. Of course, any language can be supported, for example:
# 中文电话号码：<telephone>1234567</telephone>
# 中文数字号码：<number>1234567</number>
def _process_symbol_SSML(self, words,data):
tag , match = data
language = SSML = match[1]
text = match[2]
score = 1.0
if SSML == "telephone":
# 中文-电话号码
language = "zh"
text = self.LangSSML.to_chinese_telephone(text)
elif SSML == "number":
# 中文-数字读法
language = "zh"
text = self.LangSSML.to_chinese_number(text)
elif SSML == "currency":
# 中文-按金额发音
language = "zh"
text = self.LangSSML.to_chinese_currency(text)
elif SSML == "date":
# 中文-按金额发音
language = "zh"
text = self.LangSSML.to_chinese_date(text)
self._addwords(words,language,text,score,SSML)
# ----------------------------------------------------------
def _restore_number(self, matche):
value = matche.group(0)
text_cache = self._text_cache
if value in text_cache:
process , data = text_cache[value]
tag , match = data
value = match
return value
def _pattern_symbols(self, item , text):
if text is None:return text
tag , pattern , process = item
matches = pattern.findall(text)
if len(matches) == 1 and "".join(matches[0]) == text:
return text
for i , match in enumerate(matches):
key = f"⑥{tag}{i:06d}⑥"
text = re.sub(pattern , key , text , count=1)
self._text_cache[key] = (process , (tag , match))
return text
def _process_symbol(self, words,data):
tag , match = data
language = match[1]
text = match[2]
score = 1.0
filters = self._get_filters_string()
if language not in filters:
self._process_symbol_SSML(words,data)
else:
self._addwords(words,language,text,score,True)
def _process_english(self, words,data):
tag , match = data
text = match[0]
filters = self._get_filters_string()
priority_language = filters[:2]
# Preview feature, other language segmentation processing
enablePreview = self.EnablePreview
if enablePreview == True:
# Experimental: Other language support
regex_pattern = re.compile(r'(.*?[。.?？!！]+[\n]{,1})')
lines = regex_pattern.split(text)
for index , text in enumerate(lines):
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
elif language in filters:pass # pass
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
language, score = "en", 1.0
self._addwords(words,language,text,score)
def _process_Russian(self, words,data):
tag , match = data
text = match[0]
language = "ru"
score = 1.0
self._addwords(words,language,text,score)
def _process_Thai(self, words,data):
tag , match = data
text = match[0]
language = "th"
score = 1.0
self._addwords(words,language,text,score)
def _process_korean(self, words,data):
tag , match = data
text = match[0]
language = "ko"
score = 1.0
self._addwords(words,language,text,score)
def _process_quotes(self, words,data):
tag , match = data
text = "".join(match)
childs = self.PARSE_TAG.findall(text)
if len(childs) > 0:
self._process_tags(words , text , False)
else:
cleans_text = self._cleans_text(match[1])
if len(cleans_text) <= 5:
self._parse_language(words,text)
else:
language,score = self._lang_classify(cleans_text)
self._addwords(words,language,text,score)
def _process_pinyin(self, words,data):
tag , match = data
text = match
language = "zh"
score = 1.0
self._addwords(words,language,text,score)
def _process_number(self, words,data): # "$0" process only
"""
Numbers alone cannot accurately identify language.
Because numbers are universal in all languages.
So it won't be executed here, just for testing.
"""
tag , match = data
language = words[0]["lang"] if len(words) > 0 else "zh"
text = match
score = 0.0
self._addwords(words,language,text,score)
def _process_tags(self, words , text , root_tag):
text_cache = self._text_cache
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 self.PARSE_TAG.match(text):
process , data = text_cache[text]
if process:process(words , data)
else:
self._parse_language(words , text)
return words
def _merge_results(self, words):
new_word = []
for index , cur_data in enumerate(words):
if "symbol" in 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)
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_BASE = re.compile(fr'(([【《（(“‘"\']*[LANGUAGE]+[\W\s]*)+)')
# Get custom language filter
filters = self.Langfilters
filters = filters if filters is not None else ""
# =======================================================================================================
# Experimental: Other language support.Thử nghiệm: Hỗ trợ ngôn ngữ khác.Expérimental : prise en charge d’autres langues.
# 相关语言字符如有缺失，熟悉相关语言的朋友，可以提交把缺失的发音符号补全。
# If relevant language characters are missing, friends who are familiar with the relevant languages can submit a submission to complete the missing pronunciation symbols.
# S'il manque des caractères linguistiques pertinents, les amis qui connaissent les langues concernées peuvent soumettre une soumission pour compléter les symboles de prononciation manquants.
# Nếu thiếu ký tự ngôn ngữ liên quan, những người bạn quen thuộc với ngôn ngữ liên quan có thể gửi bài để hoàn thành các ký hiệu phát âm còn thiếu.
# -------------------------------------------------------------------------------------------------------
# Preview feature, other language support
enablePreview = self.EnablePreview
if "fr" in filters or \
"vi" in filters: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 "đơưăáàảãạắằẳẵặấầẩẫậéèẻẽẹếềểễệíìỉĩịóòỏõọốồổỗộớờởỡợúùủũụứừửữựôâêơưỷỹ"
# -------------------------------------------------------------------------------------------------------
# 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_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 ,
( 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))
for index , text in enumerate(lines):
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
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"] \
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)
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 = list(lang_count.items())
self._lang_count = lang_count
return words
def setfilters(self, filters):
# 当过滤器更改时，清除缓存
# 필터가 변경되면 캐시를 지웁니다.
# フィルタが変更されると、キャッシュがクリアされます
# When the filter changes, clear the cache
if self.Langfilters != filters:
self._clears()
self.Langfilters = filters
def getfilters(self):
return self.Langfilters
def setPriorityThreshold(self, threshold:float):
self.LangPriorityThreshold = threshold
def getPriorityThreshold(self):
return self.LangPriorityThreshold
def getCounts(self):
lang_count = self._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)]
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'])
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
def getTexts(self, text:str):
if text is None or len(text.strip()) == 0:
self._clears()
return []
# lasts
text_langs = self._text_langs
if self._text_lasts == text and text_langs is not None:return text_langs
# parse
self._text_waits = []
self._lang_count = None
self._text_lasts = text
text = self._parse_symbols(text)
self._text_langs = text
return text
def classify(self, text:str):
return self.getTexts(text)
def printList(langlist):
"""
功能：打印数组结果
기능: 어레이 결과 인쇄
機能:配列結果を印刷
Function: Print array results
"""
print("\n===================【打印结果】===================")
if langlist is None or len(langlist) == 0:
print("无内容结果,No content result")
return
for line in langlist:
print(line)
pass
def main():
# -----------------------------------
# 更新日志：新版本分词更加精准。
# Changelog: The new version of the word segmentation is more accurate.
# チェンジログ:新しいバージョンの単語セグメンテーションはより正確です。
# Changelog: 분할이라는 단어의 새로운 버전이 더 정확합니다.
# -----------------------------------
# 输入示例1：（包含日文，中文）Input Example 1: (including Japanese, Chinese)
# text = "“昨日は雨が降った，音楽、映画。。。”你今天学习日语了吗？春は桜の季節です。语种分词是语音合成必不可少的环节。言語分詞は音声合成に欠かせない環節である！"
# 输入示例2：（包含日文，中文）Input Example 1: (including Japanese, Chinese)
# text = "欢迎来玩。東京，は日本の首都です。欢迎来玩. 太好了!"
# 输入示例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"])
text = """
我喜欢在雨天里听音乐。
I enjoy listening to music on rainy days.
雨の日に音楽を聴くのが好きです。
비 오는 날에 음악을 듣는 것을 즐깁니다。
J'aime écouter de la musique les jours de pluie.
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===================【语种统计】===================")
# 获取所有语种数组结果，根据内容字数降序排列
# Get the array results in all languages, sorted in descending order according to the number of content words
langCounts = langsegment.getCounts()
print(langCounts , "\n")
# 根据结果获取内容的主要语种 (语言，字数含标点)
# Get the main language of content based on the results (language, word count including punctuation)
lang , count = langCounts[0]
print(f"输入内容的主要语言为 = {lang} ，字数 = {count}")
print("==================================================\n")
# 分词输出：lang=语言，text=内容。Word output: lang = language, text = content
# ===================【打印结果】===================
# {'lang': 'zh', 'text': '你的名字叫'}
# {'lang': 'ja', 'text': '佐々木？'}
# {'lang': 'zh', 'text': '吗？韩语中的'}
# {'lang': 'ko', 'text': '안녕 오빠'}
# {'lang': 'zh', 'text': '读什么呢？'}
# {'lang': 'ja', 'text': 'あなたの体育の先生は誰ですか?'}
# {'lang': 'zh', 'text': ' 此次发布会带来了四款'}
# {'lang': 'en', 'text': 'i Phone '}
# {'lang': 'zh', 'text': '15系列机型和三款'}
# {'lang': 'en', 'text': 'Apple Watch '}
# {'lang': 'zh', 'text': '等一系列新品，这次的'}
# {'lang': 'en', 'text': 'i Pad Air '}
# {'lang': 'zh', 'text': '采用了'}
# {'lang': 'en', 'text': 'L C D '}
# {'lang': 'zh', 'text': '屏幕'}
# ===================【语种统计】===================
# ===================【语种统计】===================
# [('zh', 51), ('ja', 19), ('en', 18), ('ko', 5)]
# 输入内容的主要语言为 = zh ，字数 = 51
# ==================================================
# The main language of the input content is = zh, word count = 51
if __name__ == "__main__":
main()
language_segmentation/__init__.py
+9
View File
@@ -0,0 +1,9 @@
from .LangSegment import LangSegment
# release
__version__ = '0.3.5'
# develop
__develop__ = 'dev-0.0.1'
language_segmentation/utils/__init__.py
View File
language_segmentation/utils/num.py
+327
View File
@@ -0,0 +1,327 @@
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Digital processing from GPT_SoVITS num.py （thanks）
"""
Rules to verbalize numbers into Chinese characters.
https://zh.wikipedia.org/wiki/中文数字#現代中文
"""
import re
from collections import OrderedDict
from typing import List
DIGITS = {str(i): tran for i, tran in enumerate('零一二三四五六七八九')}
UNITS = OrderedDict({
1: '十',
2: '百',
3: '千',
4: '万',
8: '亿',
})
COM_QUANTIFIERS = '(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)'
# 分数表达式
RE_FRAC = re.compile(r'(-?)(\d+)/(\d+)')
def replace_frac(match) -> str:
"""
Args:
match (re.Match)
Returns:
str
"""
sign = match.group(1)
nominator = match.group(2)
denominator = match.group(3)
sign: str = "负" if sign else ""
nominator: str = num2str(nominator)
denominator: str = num2str(denominator)
result = f"{sign}{denominator}分之{nominator}"
return result
# 百分数表达式
RE_PERCENTAGE = re.compile(r'(-?)(\d+(\.\d+)?)%')
def replace_percentage(match) -> str:
"""
Args:
match (re.Match)
Returns:
str
"""
sign = match.group(1)
percent = match.group(2)
sign: str = "负" if sign else ""
percent: str = num2str(percent)
result = f"{sign}百分之{percent}"
return result
# 整数表达式
# 带负号的整数 -10
RE_INTEGER = re.compile(r'(-)' r'(\d+)')
def replace_negative_num(match) -> str:
"""
Args:
match (re.Match)
Returns:
str
"""
sign = match.group(1)
number = match.group(2)
sign: str = "负" if sign else ""
number: str = num2str(number)
result = f"{sign}{number}"
return result
# 编号-无符号整形
# 00078
RE_DEFAULT_NUM = re.compile(r'\d{3}\d*')
def replace_default_num(match):
"""
Args:
match (re.Match)
Returns:
str
"""
number = match.group(0)
return verbalize_digit(number, alt_one=True)
# 加减乘除
# RE_ASMD = re.compile(
# r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))([\+\-\×÷=])((-?)((\d+)(\.\d+)?)|(\.(\d+)))')
RE_ASMD = re.compile(
r'((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))([\+\-\×÷=])((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))')
asmd_map = {
'+': '加',
'-': '减',
'×': '乘',
'÷': '除',
'=': '等于'
}
def replace_asmd(match) -> str:
"""
Args:
match (re.Match)
Returns:
str
"""
result = match.group(1) + asmd_map[match.group(8)] + match.group(9)
return result
# 次方专项
RE_POWER = re.compile(r'[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]+')
power_map = {
'⁰': '0',
'¹': '1',
'²': '2',
'³': '3',
'⁴': '4',
'⁵': '5',
'⁶': '6',
'⁷': '7',
'⁸': '8',
'⁹': '9',
'ˣ': 'x',
'ʸ': 'y',
'ⁿ': 'n'
}
def replace_power(match) -> str:
"""
Args:
match (re.Match)
Returns:
str
"""
power_num = ""
for m in match.group(0):
power_num += power_map[m]
result = "的" + power_num + "次方"
return result
# 数字表达式
# 纯小数
RE_DECIMAL_NUM = re.compile(r'(-?)((\d+)(\.\d+))' r'|(\.(\d+))')
# 正整数 + 量词
RE_POSITIVE_QUANTIFIERS = re.compile(r"(\d+)([多余几\+])?" + COM_QUANTIFIERS)
RE_NUMBER = re.compile(r'(-?)((\d+)(\.\d+)?)' r'|(\.(\d+))')
def replace_positive_quantifier(match) -> str:
"""
Args:
match (re.Match)
Returns:
str
"""
number = match.group(1)
match_2 = match.group(2)
if match_2 == "+":
match_2 = "多"
match_2: str = match_2 if match_2 else ""
quantifiers: str = match.group(3)
number: str = num2str(number)
result = f"{number}{match_2}{quantifiers}"
return result
def replace_number(match) -> str:
"""
Args:
match (re.Match)
Returns:
str
"""
sign = match.group(1)
number = match.group(2)
pure_decimal = match.group(5)
if pure_decimal:
result = num2str(pure_decimal)
else:
sign: str = "负" if sign else ""
number: str = num2str(number)
result = f"{sign}{number}"
return result
# 范围表达式
# match.group(1) and match.group(8) are copy from RE_NUMBER
RE_RANGE = re.compile(
r"""
(?<![\d\+\-\×÷=]) # 使用反向前瞻以确保数字范围之前没有其他数字和操作符
((-?)((\d+)(\.\d+)?)) # 匹配范围起始的负数或正数（整数或小数）
[-~] # 匹配范围分隔符
((-?)((\d+)(\.\d+)?)) # 匹配范围结束的负数或正数（整数或小数）
(?![\d\+\-\×÷=]) # 使用正向前瞻以确保数字范围之后没有其他数字和操作符
""", re.VERBOSE)
def replace_range(match) -> str:
"""
Args:
match (re.Match)
Returns:
str
"""
first, second = match.group(1), match.group(6)
first = RE_NUMBER.sub(replace_number, first)
second = RE_NUMBER.sub(replace_number, second)
result = f"{first}到{second}"
return result
# ~至表达式
RE_TO_RANGE = re.compile(
r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)[~]((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)')
def replace_to_range(match) -> str:
"""
Args:
match (re.Match)
Returns:
str
"""
result = match.group(0).replace('~', '至')
return result
def _get_value(value_string: str, use_zero: bool=True) -> List[str]:
stripped = value_string.lstrip('0')
if len(stripped) == 0:
return []
elif len(stripped) == 1:
if use_zero and len(stripped) < len(value_string):
return [DIGITS['0'], DIGITS[stripped]]
else:
return [DIGITS[stripped]]
else:
largest_unit = next(
power for power in reversed(UNITS.keys()) if power < len(stripped))
first_part = value_string[:-largest_unit]
second_part = value_string[-largest_unit:]
return _get_value(first_part) + [UNITS[largest_unit]] + _get_value(
second_part)
def verbalize_cardinal(value_string: str) -> str:
if not value_string:
return ''
# 000 -> '零' , 0 -> '零'
value_string = value_string.lstrip('0')
if len(value_string) == 0:
return DIGITS['0']
result_symbols = _get_value(value_string)
# verbalized number starting with '一十*' is abbreviated as `十*`
if len(result_symbols) >= 2 and result_symbols[0] == DIGITS[
'1'] and result_symbols[1] == UNITS[1]:
result_symbols = result_symbols[1:]
return ''.join(result_symbols)
def verbalize_digit(value_string: str, alt_one=False) -> str:
result_symbols = [DIGITS[digit] for digit in value_string]
result = ''.join(result_symbols)
if alt_one:
result = result.replace("一", "幺")
return result
def num2str(value_string: str) -> str:
integer_decimal = value_string.split('.')
if len(integer_decimal) == 1:
integer = integer_decimal[0]
decimal = ''
elif len(integer_decimal) == 2:
integer, decimal = integer_decimal
else:
raise ValueError(
f"The value string: '${value_string}' has more than one point in it."
)
result = verbalize_cardinal(integer)
decimal = decimal.rstrip('0')
if decimal:
# '.22' is verbalized as '零点二二'
# '3.20' is verbalized as '三点二
result = result if result else "零"
result += '点' + verbalize_digit(decimal)
return result
if __name__ == "__main__":
text = ""
text = num2str(text)
print(text)
pass
models/ace_step_transformer.py
+475
View File
@@ -0,0 +1,475 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, List, Union
import torch
import torch.nn.functional as F
from torch import nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.utils import BaseOutput, is_torch_version
from diffusers.models.modeling_utils import ModelMixin
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
from .attention import LinearTransformerBlock, t2i_modulate
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_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
# Copied from transformers.models.mixtral.modeling_mixtral.MixtralRotaryEmbedding with Mixtral->Qwen2
class Qwen2RotaryEmbedding(nn.Module):
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
super().__init__()
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))
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()
)
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)
freqs = torch.outer(t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim=-1)
self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
def forward(self, x, seq_len=None):
# x: [bs, num_attention_heads, seq_len, head_size]
if seq_len > self.max_seq_len_cached:
self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
return (
self.cos_cached[:seq_len].to(dtype=x.dtype),
self.sin_cached[:seq_len].to(dtype=x.dtype),
)
class T2IFinalLayer(nn.Module):
"""
The final layer of Sana.
"""
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.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5)
self.out_channels = out_channels
self.patch_size = patch_size
def unpatchfy(
self,
hidden_states: torch.Tensor,
width: int,
):
# 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)
).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])
).contiguous()
if width > new_width:
output = torch.nn.functional.pad(output, (0, width - new_width, 0, 0), 'constant', 0)
elif width < new_width:
output = output[:, :, :, :width]
return output
def forward(self, x, t, output_length):
shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2, dim=1)
x = t2i_modulate(self.norm_final(x), shift, scale)
x = self.linear(x)
# unpatchify
output = self.unpatchfy(x, output_length)
return output
class PatchEmbed(nn.Module):
"""2D Image to Patch Embedding"""
def __init__(
self,
height=16,
width=4096,
patch_size=(16, 1),
in_channels=8,
embed_dim=1152,
bias=True,
):
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),
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
self.base_size = self.width
def forward(self, latent):
# early convolutions, N x C x H x W -> N x 256 * sqrt(patch_size) x H/patch_size x W/patch_size
latent = self.early_conv_layers(latent)
latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC
return latent
@dataclass
class Transformer2DModelOutput(BaseOutput):
sample: torch.FloatTensor
proj_losses: Optional[Tuple[Tuple[str, torch.Tensor]]] = None
class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
_supports_gradient_checkpointing = True
@register_to_config
def __init__(
self,
in_channels: Optional[int] = 8,
num_layers: int = 28,
inner_dim: int = 1536,
attention_head_dim: int = 64,
num_attention_heads: int = 24,
mlp_ratio: float = 4.0,
out_channels: int = 8,
max_position: int = 32768,
rope_theta: float = 1000000.0,
speaker_embedding_dim: int = 512,
text_embedding_dim: int = 768,
ssl_encoder_depths: List[int] = [9, 9],
ssl_names: List[str] = ["mert", "m-hubert"],
ssl_latent_dims: List[int] = [1024, 768],
lyric_encoder_vocab_size: int = 6681,
lyric_hidden_size: int = 1024,
patch_size: List[int] = [16, 1],
max_height: int = 16,
max_width: int = 4096,
**kwargs,
):
super().__init__()
self.num_attention_heads = num_attention_heads
self.attention_head_dim = attention_head_dim
inner_dim = num_attention_heads * attention_head_dim
self.inner_dim = inner_dim
self.out_channels = out_channels
self.max_position = max_position
self.patch_size = patch_size
self.rope_theta = rope_theta
self.rotary_emb = Qwen2RotaryEmbedding(
dim=self.attention_head_dim,
max_position_embeddings=self.max_position,
base=self.rope_theta,
)
# 2. Define input layers
self.in_channels = in_channels
# 3. Define transformers blocks
self.transformer_blocks = nn.ModuleList(
[
LinearTransformerBlock(
dim=self.inner_dim,
num_attention_heads=self.num_attention_heads,
attention_head_dim=attention_head_dim,
mlp_ratio=mlp_ratio,
add_cross_attention=True,
add_cross_attention_dim=self.inner_dim,
)
for i in range(self.config.num_layers)
]
)
self.num_layers = num_layers
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
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)
# genre
self.genre_embedder = nn.Linear(text_embedding_dim, self.inner_dim)
# 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_proj = nn.Linear(lyric_hidden_size, self.inner_dim)
projector_dim = 2 * self.inner_dim
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
])
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.ssl_names = ssl_names
self.proj_in = PatchEmbed(
height=max_height,
width=max_width,
patch_size=patch_size,
embed_dim=self.inner_dim,
bias=True,
)
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:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
Parameters:
chunk_size (`int`, *optional*):
The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
over each tensor of dim=`dim`.
dim (`int`, *optional*, defaults to `0`):
The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
or dim=1 (sequence length).
"""
if dim not in [0, 1]:
raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
# 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):
if hasattr(module, "set_chunk_feed_forward"):
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
for child in module.children():
fn_recursive_feed_forward(child, chunk_size, dim)
for module in self.children():
fn_recursive_feed_forward(module, chunk_size, dim)
def _set_gradient_checkpointing(self, module, value=False):
if hasattr(module, "gradient_checkpointing"):
module.gradient_checkpointing = value
def forward_lyric_encoder(
self,
lyric_token_idx: Optional[torch.LongTensor] = None,
lyric_mask: Optional[torch.LongTensor] = None,
):
# 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 = self.lyric_proj(prompt_prenet_out)
return prompt_prenet_out
def encode(
self,
encoder_text_hidden_states: Optional[torch.Tensor] = None,
text_attention_mask: Optional[torch.LongTensor] = None,
speaker_embeds: Optional[torch.FloatTensor] = None,
lyric_token_idx: Optional[torch.LongTensor] = None,
lyric_mask: Optional[torch.LongTensor] = None,
):
bs = encoder_text_hidden_states.shape[0]
device = encoder_text_hidden_states.device
# speaker embedding
encoder_spk_hidden_states = self.speaker_embedder(speaker_embeds).unsqueeze(1)
speaker_mask = torch.ones(bs, 1, device=device)
# genre embedding
encoder_text_hidden_states = self.genre_embedder(encoder_text_hidden_states)
# lyric
encoder_lyric_hidden_states = self.forward_lyric_encoder(
lyric_token_idx=lyric_token_idx,
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_mask = torch.cat([speaker_mask, text_attention_mask, lyric_mask], dim=1)
return encoder_hidden_states, encoder_hidden_mask
def decode(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
encoder_hidden_states: torch.Tensor,
encoder_hidden_mask: torch.Tensor,
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,
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))
temb = self.t_block(embedded_timestep)
hidden_states = self.proj_in(hidden_states)
# controlnet logic
if block_controlnet_hidden_states is not None:
control_condi = cross_norm(hidden_states, block_controlnet_hidden_states)
hidden_states = hidden_states + control_condi * controlnet_scale
inner_hidden_states = []
rotary_freqs_cis = self.rotary_emb(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):
if self.training and self.gradient_checkpointing:
def create_custom_forward(module, return_dict=None):
def custom_forward(*inputs):
if return_dict is not None:
return module(*inputs, return_dict=return_dict)
else:
return module(*inputs)
return custom_forward
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,
attention_mask=attention_mask,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_hidden_mask,
rotary_freqs_cis=rotary_freqs_cis,
rotary_freqs_cis_cross=encoder_rotary_freqs_cis,
temb=temb,
**ckpt_kwargs,
)
else:
hidden_states = block(
hidden_states=hidden_states,
attention_mask=attention_mask,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_hidden_mask,
rotary_freqs_cis=rotary_freqs_cis,
rotary_freqs_cis_cross=encoder_rotary_freqs_cis,
temb=temb,
)
for ssl_encoder_depth in self.ssl_encoder_depths:
if index_block == ssl_encoder_depth:
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:
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
est_ssl_hidden_state = projector(inner_hidden_state)
# 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)):
# 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 = torch.nn.functional.normalize(z_tilde, dim=-1)
z = torch.nn.functional.normalize(z, dim=-1)
# T x d -> T x 1 -> 1
target = torch.ones(z.shape[0], device=z.device)
proj_loss += self.cosine_loss(z, z_tilde, target)
proj_losses.append((ssl_name, proj_loss / bs))
output = self.final_layer(hidden_states, embedded_timestep, output_length)
if not return_dict:
return (output, proj_losses)
return Transformer2DModelOutput(sample=output, proj_losses=proj_losses)
# @torch.compile
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
encoder_text_hidden_states: Optional[torch.Tensor] = None,
text_attention_mask: Optional[torch.LongTensor] = None,
speaker_embeds: Optional[torch.FloatTensor] = None,
lyric_token_idx: Optional[torch.LongTensor] = None,
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,
controlnet_scale: Union[float, torch.Tensor] = 1.0,
return_dict: bool = True,
):
encoder_hidden_states, encoder_hidden_mask = self.encode(
encoder_text_hidden_states=encoder_text_hidden_states,
text_attention_mask=text_attention_mask,
speaker_embeds=speaker_embeds,
lyric_token_idx=lyric_token_idx,
lyric_mask=lyric_mask,
)
output_length = hidden_states.shape[-1]
output = self.decode(
hidden_states=hidden_states,
attention_mask=attention_mask,
encoder_hidden_states=encoder_hidden_states,
encoder_hidden_mask=encoder_hidden_mask,
timestep=timestep,
ssl_hidden_states=ssl_hidden_states,
output_length=output_length,
block_controlnet_hidden_states=block_controlnet_hidden_states,
controlnet_scale=controlnet_scale,
return_dict=return_dict,
)
return output
models/attention.py
+319
View File
@@ -0,0 +1,319 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Tuple, Union
import torch
import torch.nn.functional as F
from torch import nn
from diffusers.utils import logging
from diffusers.models.normalization import RMSNorm
try:
# from .dcformer import DCMHAttention
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
logger = logging.get_logger(__name__)
def val2list(x: list or tuple or any, repeat_time=1) -> list: # type: ignore
"""Repeat `val` for `repeat_time` times and return the list or val if list/tuple."""
if isinstance(x, (list, tuple)):
return list(x)
return [x for _ in range(repeat_time)]
def val2tuple(x: list or tuple or any, min_len: int = 1, idx_repeat: int = -1) -> tuple: # type: ignore
"""Return tuple with min_len by repeating element at idx_repeat."""
# convert to list first
x = val2list(x)
# repeat elements if necessary
if len(x) > 0:
x[idx_repeat:idx_repeat] = [x[idx_repeat] for _ in range(min_len - len(x))]
return tuple(x)
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, ...]]:
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,
in_dim: int,
out_dim: int,
kernel_size=3,
stride=1,
dilation=1,
groups=1,
padding: Union[int, None] = None,
use_bias=False,
norm=None,
act=None,
):
super().__init__()
if padding is None:
padding = get_same_padding(kernel_size)
padding *= dilation
self.in_dim = in_dim
self.out_dim = out_dim
self.kernel_size = kernel_size
self.stride = stride
self.dilation = dilation
self.groups = groups
self.padding = padding
self.use_bias = use_bias
self.conv = nn.Conv1d(
in_dim,
out_dim,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=use_bias,
)
if norm is not None:
self.norm = RMSNorm(out_dim, elementwise_affine=False)
else:
self.norm = None
if act is not None:
self.act = nn.SiLU(inplace=True)
else:
self.act = None
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.conv(x)
if self.norm:
x = self.norm(x)
if self.act:
x = self.act(x)
return x
class GLUMBConv(nn.Module):
def __init__(
self,
in_features: int,
hidden_features: int,
out_feature=None,
kernel_size=3,
stride=1,
padding: Union[int, None] = None,
use_bias=False,
norm=(None, None, None),
act=("silu", "silu", None),
dilation=1,
):
out_feature = out_feature or in_features
super().__init__()
use_bias = val2tuple(use_bias, 3)
norm = val2tuple(norm, 3)
act = val2tuple(act, 3)
self.glu_act = nn.SiLU(inplace=False)
self.inverted_conv = ConvLayer(
in_features,
hidden_features * 2,
1,
use_bias=use_bias[0],
norm=norm[0],
act=act[0],
)
self.depth_conv = ConvLayer(
hidden_features * 2,
hidden_features * 2,
kernel_size,
stride=stride,
groups=hidden_features * 2,
padding=padding,
use_bias=use_bias[1],
norm=norm[1],
act=None,
dilation=dilation,
)
self.point_conv = ConvLayer(
hidden_features,
out_feature,
1,
use_bias=use_bias[2],
norm=norm[2],
act=act[2],
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x.transpose(1, 2)
x = self.inverted_conv(x)
x = self.depth_conv(x)
x, gate = torch.chunk(x, 2, dim=1)
gate = self.glu_act(gate)
x = x * gate
x = self.point_conv(x)
x = x.transpose(1, 2)
return x
class LinearTransformerBlock(nn.Module):
"""
A Sana block with global shared adaptive layer norm (adaLN-single) conditioning.
"""
def __init__(
self,
dim,
num_attention_heads,
attention_head_dim,
use_adaln_single=True,
cross_attention_dim=None,
added_kv_proj_dim=None,
context_pre_only=False,
mlp_ratio=4.0,
add_cross_attention=False,
add_cross_attention_dim=None,
qk_norm=None,
):
super().__init__()
self.norm1 = RMSNorm(dim, elementwise_affine=False, eps=1e-6)
self.attn = Attention(
query_dim=dim,
cross_attention_dim=cross_attention_dim,
added_kv_proj_dim=added_kv_proj_dim,
dim_head=attention_head_dim,
heads=num_attention_heads,
out_dim=dim,
bias=True,
qk_norm=qk_norm,
processor=CustomLiteLAProcessor2_0(),
)
self.add_cross_attention = add_cross_attention
self.context_pre_only = context_pre_only
if add_cross_attention and add_cross_attention_dim is not None:
self.cross_attn = Attention(
query_dim=dim,
cross_attention_dim=add_cross_attention_dim,
added_kv_proj_dim=add_cross_attention_dim,
dim_head=attention_head_dim,
heads=num_attention_heads,
out_dim=dim,
context_pre_only=context_pre_only,
bias=True,
qk_norm=qk_norm,
processor=CustomerAttnProcessor2_0(),
)
self.norm2 = RMSNorm(dim, 1e-06, elementwise_affine=False)
self.ff = GLUMBConv(
in_features=dim,
hidden_features=int(dim * mlp_ratio),
use_bias=(True, True, False),
norm=(None, None, None),
act=("silu", "silu", None),
)
self.use_adaln_single = use_adaln_single
if use_adaln_single:
self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
def forward(
self,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
attention_mask: torch.FloatTensor = None,
encoder_attention_mask: torch.FloatTensor = None,
rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
temb: torch.FloatTensor = None,
):
N = hidden_states.shape[0]
# step 1: AdaLN single
if self.use_adaln_single:
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
self.scale_shift_table[None] + temb.reshape(N, 6, -1)
).chunk(6, dim=1)
norm_hidden_states = self.norm1(hidden_states)
if self.use_adaln_single:
norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
# step 2: attention
if not self.add_cross_attention:
attn_output, encoder_hidden_states = self.attn(
hidden_states=norm_hidden_states,
attention_mask=attention_mask,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_attention_mask,
rotary_freqs_cis=rotary_freqs_cis,
rotary_freqs_cis_cross=rotary_freqs_cis_cross,
)
else:
attn_output, _ = self.attn(
hidden_states=norm_hidden_states,
attention_mask=attention_mask,
encoder_hidden_states=None,
encoder_attention_mask=None,
rotary_freqs_cis=rotary_freqs_cis,
rotary_freqs_cis_cross=None,
)
if self.use_adaln_single:
attn_output = gate_msa * attn_output
hidden_states = attn_output + hidden_states
if self.add_cross_attention:
attn_output = self.cross_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_attention_mask,
rotary_freqs_cis=rotary_freqs_cis,
rotary_freqs_cis_cross=rotary_freqs_cis_cross,
)
hidden_states = attn_output + hidden_states
# step 3: add norm
norm_hidden_states = self.norm2(hidden_states)
if self.use_adaln_single:
norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
# step 4: feed forward
ff_output = self.ff(norm_hidden_states)
if self.use_adaln_single:
ff_output = gate_mlp * ff_output
hidden_states = hidden_states + ff_output
return hidden_states
models/config.json
+23
View File
@@ -0,0 +1,23 @@
{
"_class_name": "Transformer2DModel",
"_diffusers_version": "0.27.2",
"in_channels": 8,
"num_layers": 24,
"inner_dim": 2560,
"attention_head_dim": 128,
"num_attention_heads": 20,
"mlp_ratio": 2.5,
"out_channels": 8,
"max_position": 32768,
"rope_theta": 1000000.0,
"speaker_embedding_dim": 512,
"text_embedding_dim": 768,
"ssl_encoder_depths": [8, 8],
"ssl_names": ["mert", "m-hubert"],
"ssl_latent_dims": [1024, 768],
"patch_size": [16, 1],
"max_height": 16,
"max_width": 32768,
"lyric_encoder_vocab_size": 6693,
"lyric_hidden_size": 1024
}
models/customer_attention_processor.py
+339
View File
@@ -0,0 +1,339 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Optional, Union, Tuple
import torch
import torch.nn.functional as F
from torch import nn
from diffusers.utils import logging
from diffusers.models.attention_processor import Attention
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class CustomLiteLAProcessor2_0:
"""Attention processor used typically in processing the SD3-like self-attention projections. add rms norm for query and key and apply RoPE"""
def __init__(self):
self.kernel_func = nn.ReLU(inplace=False)
self.eps = 1e-15
self.pad_val = 1.0
def apply_rotary_emb(
self,
x: torch.Tensor,
freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
tensors contain rotary embeddings and are returned as real tensors.
Args:
x (`torch.Tensor`):
Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
Returns:
Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
"""
cos, sin = freqs_cis # [S, D]
cos = cos[None, None]
sin = sin[None, None]
cos, sin = cos.to(x.device), sin.to(x.device)
x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
return out
def __call__(
self,
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
attention_mask: Optional[torch.FloatTensor] = None,
encoder_attention_mask: Optional[torch.FloatTensor] = None,
rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
*args,
**kwargs,
) -> torch.FloatTensor:
hidden_states_len = hidden_states.shape[1]
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)
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)
batch_size = hidden_states.shape[0]
# `sample` projections.
dtype = hidden_states.dtype
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
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")
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)
encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
# attention
if not attn.is_cross_attention:
query = torch.cat([query, encoder_hidden_states_query_proj], dim=1)
key = torch.cat([key, encoder_hidden_states_key_proj], dim=1)
value = torch.cat([value, encoder_hidden_states_value_proj], dim=1)
else:
query = hidden_states
key = encoder_hidden_states
value = encoder_hidden_states
inner_dim = key.shape[-1]
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)
value = value.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
# RoPE需要 [B, H, S, D] 输入
# 此时 query是 [B, H, D, S], 需要转成 [B, H, S, D] 才能应用RoPE
query = query.permute(0, 1, 3, 2) # [B, H, S, D] (从 [B, H, D, S])
# Apply query and key normalization if needed
if attn.norm_q is not None:
query = attn.norm_q(query)
if attn.norm_k is not None:
key = attn.norm_k(key)
# Apply RoPE if needed
if rotary_freqs_cis is not None:
query = self.apply_rotary_emb(query, rotary_freqs_cis)
if not attn.is_cross_attention:
key = self.apply_rotary_emb(key, rotary_freqs_cis)
elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:
key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)
# 此时 query是 [B, H, S, D]，需要还原成 [B, H, D, S]
query = query.permute(0, 1, 3, 2) # [B, H, D, S]
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]
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] 相乘
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:
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]
query = self.kernel_func(query)
key = self.kernel_func(key)
query, key, value = query.float(), key.float(), value.float()
value = F.pad(value, (0, 0, 0, 1), mode="constant", value=self.pad_val)
vk = torch.matmul(value, key)
hidden_states = torch.matmul(vk, query)
if hidden_states.dtype in [torch.float16, torch.bfloat16]:
hidden_states = hidden_states.float()
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.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:
hidden_states, encoder_hidden_states = (
hidden_states[:, : hidden_states_len],
hidden_states[:, hidden_states_len:],
)
# linear proj
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"):
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)
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)
if torch.get_autocast_gpu_dtype() == torch.float16:
hidden_states = hidden_states.clip(-65504, 65504)
if encoder_hidden_states is not None:
encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
return hidden_states, encoder_hidden_states
class CustomerAttnProcessor2_0:
r"""
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.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.")
def apply_rotary_emb(
self,
x: torch.Tensor,
freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
tensors contain rotary embeddings and are returned as real tensors.
Args:
x (`torch.Tensor`):
Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
Returns:
Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
"""
cos, sin = freqs_cis # [S, D]
cos = cos[None, None]
sin = sin[None, None]
cos, sin = cos.to(x.device), sin.to(x.device)
x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
return out
def __call__(
self,
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
attention_mask: Optional[torch.FloatTensor] = None,
encoder_attention_mask: Optional[torch.FloatTensor] = None,
rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
*args,
**kwargs,
) -> torch.Tensor:
residual = hidden_states
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)
batch_size, sequence_length, _ = (
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")
if attn.group_norm is not None:
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)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
if attn.norm_q is not None:
query = attn.norm_q(query)
if attn.norm_k is not None:
key = attn.norm_k(key)
# Apply RoPE if needed
if rotary_freqs_cis is not None:
query = self.apply_rotary_emb(query, rotary_freqs_cis)
if not attn.is_cross_attention:
key = self.apply_rotary_emb(key, rotary_freqs_cis)
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:
# 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, :]
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)
elif not attn.is_cross_attention and attention_mask is not None:
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])
# the output of sdp = (batch, num_heads, seq_len, head_dim)
# TODO: add support for attn.scale when we move to Torch 2.1
hidden_states = F.scaled_dot_product_attention(
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.to(query.dtype)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
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)
if attn.residual_connection:
hidden_states = hidden_states + residual
hidden_states = hidden_states / attn.rescale_output_factor
return hidden_states
models/lyrics_utils/lyric_encoder.py
+1070
View File
File diff suppressed because it is too large Load Diff
models/lyrics_utils/lyric_normalizer.py
+66
View File
@@ -0,0 +1,66 @@
import re
from opencc import OpenCC
t2s_converter = OpenCC('t2s')
s2t_converter = OpenCC('s2t')
EMOJI_PATTERN = re.compile(
"["
"\U0001F600-\U0001F64F" # Emoticons
"]+", flags=re.UNICODE
)
# 创建一个翻译表，用于替换和移除字符
TRANSLATION_TABLE = str.maketrans({
'-': ' ', # 将 '-' 替换为空格
',': None,
'.': None,
'，': None,
'。': None,
'!': None,
'！': None,
'?': None,
'？': None,
'…': None,
';': None,
'；': None,
':': None,
'：': None,
'\u3000': ' ', # 将全角空格替换为空格
})
# 替换括号中的内容，包括中括号和小括号
BACKSLASH_PATTERN = re.compile(r'\(.*?\)|\[.*?\]')
SPACE_PATTERN = re.compile('(?<!^)\s+(?!$)')
def normalize_text(text, language, strip=True):
"""
对文本进行标准化处理，去除标点符号，转为小写（如果适用）
"""
# Step 1: 替换 '-' 为 ' ' 并移除标点符号
text = text.translate(TRANSLATION_TABLE)
# Step 2: 移除表情符号
text = EMOJI_PATTERN.sub('', text)
# Step 3: 连续空白字符替换为单个空格，首位除外
text = SPACE_PATTERN.sub(' ', text)
# Step 4: 去除首尾空白字符（如果需要）
if strip:
text = text.strip()
# Step 5: 转为小写
text = text.lower()
# Step 6: 多语言转换
if language == "zh":
text = t2s_converter.convert(text)
if language == "yue":
text = s2t_converter.convert(text)
# 其他语言根据需要添加
return text
models/lyrics_utils/lyric_tokenizer.py
+883
View File
@@ -0,0 +1,883 @@
import os
import re
import textwrap
from functools import cached_property
import pypinyin
import torch
from hangul_romanize import Transliter
from hangul_romanize.rule import academic
from num2words import num2words
from spacy.lang.ar import Arabic
from spacy.lang.en import English
from spacy.lang.es import Spanish
from spacy.lang.ja import Japanese
from spacy.lang.zh import Chinese
from tokenizers import Tokenizer
from models.lyrics_utils.zh_num2words import TextNorm as zh_num2words
from typing import Dict, List, Optional, Set, Union
#copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/tokenizer.py
def get_spacy_lang(lang):
if lang == "zh":
return Chinese()
elif lang == "ja":
return Japanese()
elif lang == "ar":
return Arabic()
elif lang == "es":
return Spanish()
else:
# For most languages, Enlish does the job
return English()
def split_sentence(text, lang, text_split_length=250):
"""Preprocess the input text"""
text_splits = []
if text_split_length is not None and len(text) >= text_split_length:
text_splits.append("")
nlp = get_spacy_lang(lang)
nlp.add_pipe("sentencizer")
doc = nlp(text)
for sentence in doc.sents:
if len(text_splits[-1]) + len(str(sentence)) <= text_split_length:
# if the last sentence + the current sentence is less than the text_split_length
# then add the current sentence to the last sentence
text_splits[-1] += " " + str(sentence)
text_splits[-1] = text_splits[-1].lstrip()
elif len(str(sentence)) > text_split_length:
# if the current sentence is greater than the text_split_length
for line in textwrap.wrap(
str(sentence),
width=text_split_length,
drop_whitespace=True,
break_on_hyphens=False,
tabsize=1,
):
text_splits.append(str(line))
else:
text_splits.append(str(sentence))
if len(text_splits) > 1:
if text_splits[0] == "":
del text_splits[0]
else:
text_splits = [text.lstrip()]
return text_splits
_whitespace_re = re.compile(r"\s+")
# List of (regular expression, replacement) pairs for abbreviations:
_abbreviations = {
"en": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
("mrs", "misess"),
("mr", "mister"),
("dr", "doctor"),
("st", "saint"),
("co", "company"),
("jr", "junior"),
("maj", "major"),
("gen", "general"),
("drs", "doctors"),
("rev", "reverend"),
("lt", "lieutenant"),
("hon", "honorable"),
("sgt", "sergeant"),
("capt", "captain"),
("esq", "esquire"),
("ltd", "limited"),
("col", "colonel"),
("ft", "fort"),
]
],
"es": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
("sra", "señora"),
("sr", "señor"),
("dr", "doctor"),
("dra", "doctora"),
("st", "santo"),
("co", "compañía"),
("jr", "junior"),
("ltd", "limitada"),
]
],
"fr": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
("mme", "madame"),
("mr", "monsieur"),
("dr", "docteur"),
("st", "saint"),
("co", "compagnie"),
("jr", "junior"),
("ltd", "limitée"),
]
],
"de": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
("fr", "frau"),
("dr", "doktor"),
("st", "sankt"),
("co", "firma"),
("jr", "junior"),
]
],
"pt": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
("sra", "senhora"),
("sr", "senhor"),
("dr", "doutor"),
("dra", "doutora"),
("st", "santo"),
("co", "companhia"),
("jr", "júnior"),
("ltd", "limitada"),
]
],
"it": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
# ("sig.ra", "signora"),
("sig", "signore"),
("dr", "dottore"),
("st", "santo"),
("co", "compagnia"),
("jr", "junior"),
("ltd", "limitata"),
]
],
"pl": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
("p", "pani"),
("m", "pan"),
("dr", "doktor"),
("sw", "święty"),
("jr", "junior"),
]
],
"ar": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
# There are not many common abbreviations in Arabic as in English.
]
],
"zh": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
# Chinese doesn't typically use abbreviations in the same way as Latin-based scripts.
]
],
"cs": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
("dr", "doktor"), # doctor
("ing", "inženýr"), # engineer
("p", "pan"), # Could also map to pani for woman but no easy way to do it
# Other abbreviations would be specialized and not as common.
]
],
"ru": [
(re.compile("\\b%s\\b" % x[0], re.IGNORECASE), x[1])
for x in [
("г-жа", "госпожа"), # Mrs.
("г-н", "господин"), # Mr.
("д-р", "доктор"), # doctor
# Other abbreviations are less common or specialized.
]
],
"nl": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
("dhr", "de heer"), # Mr.
("mevr", "mevrouw"), # Mrs.
("dr", "dokter"), # doctor
("jhr", "jonkheer"), # young lord or nobleman
# Dutch uses more abbreviations, but these are the most common ones.
]
],
"tr": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
("b", "bay"), # Mr.
("byk", "büyük"), # büyük
("dr", "doktor"), # doctor
# Add other Turkish abbreviations here if needed.
]
],
"hu": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
("dr", "doktor"), # doctor
("b", "bácsi"), # Mr.
("nőv", "nővér"), # nurse
# Add other Hungarian abbreviations here if needed.
]
],
"ko": [
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
for x in [
# Korean doesn't typically use abbreviations in the same way as Latin-based scripts.
]
],
}
def expand_abbreviations_multilingual(text, lang="en"):
for regex, replacement in _abbreviations[lang]:
text = re.sub(regex, replacement, text)
return text
_symbols_multilingual = {
"en": [
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " and "),
("@", " at "),
("%", " percent "),
("#", " hash "),
("$", " dollar "),
("£", " pound "),
("°", " degree "),
]
],
"es": [
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " y "),
("@", " arroba "),
("%", " por ciento "),
("#", " numeral "),
("$", " dolar "),
("£", " libra "),
("°", " grados "),
]
],
"fr": [
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " et "),
("@", " arobase "),
("%", " pour cent "),
("#", " dièse "),
("$", " dollar "),
("£", " livre "),
("°", " degrés "),
]
],
"de": [
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " und "),
("@", " at "),
("%", " prozent "),
("#", " raute "),
("$", " dollar "),
("£", " pfund "),
("°", " grad "),
]
],
"pt": [
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " e "),
("@", " arroba "),
("%", " por cento "),
("#", " cardinal "),
("$", " dólar "),
("£", " libra "),
("°", " graus "),
]
],
"it": [
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " e "),
("@", " chiocciola "),
("%", " per cento "),
("#", " cancelletto "),
("$", " dollaro "),
("£", " sterlina "),
("°", " gradi "),
]
],
"pl": [
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " i "),
("@", " małpa "),
("%", " procent "),
("#", " krzyżyk "),
("$", " dolar "),
("£", " funt "),
("°", " stopnie "),
]
],
"ar": [
# Arabic
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " و "),
("@", " على "),
("%", " في المئة "),
("#", " رقم "),
("$", " دولار "),
("£", " جنيه "),
("°", " درجة "),
]
],
"zh": [
# Chinese
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " 和 "),
("@", " 在 "),
("%", " 百分之 "),
("#", " 号 "),
("$", " 美元 "),
("£", " 英镑 "),
("°", " 度 "),
]
],
"cs": [
# Czech
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " a "),
("@", " na "),
("%", " procento "),
("#", " křížek "),
("$", " dolar "),
("£", " libra "),
("°", " stupně "),
]
],
"ru": [
# Russian
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " и "),
("@", " собака "),
("%", " процентов "),
("#", " номер "),
("$", " доллар "),
("£", " фунт "),
("°", " градус "),
]
],
"nl": [
# Dutch
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " en "),
("@", " bij "),
("%", " procent "),
("#", " hekje "),
("$", " dollar "),
("£", " pond "),
("°", " graden "),
]
],
"tr": [
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " ve "),
("@", " at "),
("%", " yüzde "),
("#", " diyez "),
("$", " dolar "),
("£", " sterlin "),
("°", " derece "),
]
],
"hu": [
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " és "),
("@", " kukac "),
("%", " százalék "),
("#", " kettőskereszt "),
("$", " dollár "),
("£", " font "),
("°", " fok "),
]
],
"ko": [
# Korean
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
for x in [
("&", " 그리고 "),
("@", " 에 "),
("%", " 퍼센트 "),
("#", " 번호 "),
("$", " 달러 "),
("£", " 파운드 "),
("°", " 도 "),
]
],
}
def expand_symbols_multilingual(text, lang="en"):
for regex, replacement in _symbols_multilingual[lang]:
text = re.sub(regex, replacement, text)
text = text.replace(" ", " ") # Ensure there are no double spaces
return text.strip()
_ordinal_re = {
"en": re.compile(r"([0-9]+)(st|nd|rd|th)"),
"es": re.compile(r"([0-9]+)(º|ª|er|o|a|os|as)"),
"fr": re.compile(r"([0-9]+)(º|ª|er|re|e|ème)"),
"de": re.compile(r"([0-9]+)(st|nd|rd|th|º|ª|\.(?=\s|$))"),
"pt": re.compile(r"([0-9]+)(º|ª|o|a|os|as)"),
"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.
"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ü|\.)"),
"hu": re.compile(r"([0-9]+)(\.|adik|edik|odik|edik|ödik|ödike|ik)"),
"ko": re.compile(r"([0-9]+)(번째|번|차|째)"),
}
_number_re = re.compile(r"[0-9]+")
_currency_re = {
"USD": re.compile(r"((\$[0-9\.\,]*[0-9]+)|([0-9\.\,]*[0-9]+\$))"),
"GBP": re.compile(r"((£[0-9\.\,]*[0-9]+)|([0-9\.\,]*[0-9]+£))"),
"EUR": re.compile(r"(([0-9\.\,]*[0-9]+€)|((€[0-9\.\,]*[0-9]+)))"),
}
_comma_number_re = re.compile(r"\b\d{1,3}(,\d{3})*(\.\d+)?\b")
_dot_number_re = re.compile(r"\b\d{1,3}(.\d{3})*(\,\d+)?\b")
_decimal_number_re = re.compile(r"([0-9]+[.,][0-9]+)")
def _remove_commas(m):
text = m.group(0)
if "," in text:
text = text.replace(",", "")
return text
def _remove_dots(m):
text = m.group(0)
if "." in text:
text = text.replace(".", "")
return text
def _expand_decimal_point(m, lang="en"):
amount = m.group(1).replace(",", ".")
return num2words(float(amount), lang=lang if lang != "cs" else "cz")
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")
and_equivalents = {
"en": ", ",
"es": " con ",
"fr": " et ",
"de": " und ",
"pt": " e ",
"it": " e ",
"pl": ", ",
"cs": ", ",
"ru": ", ",
"nl": ", ",
"ar": ", ",
"tr": ", ",
"hu": ", ",
"ko": ", ",
}
if amount.is_integer():
last_and = full_amount.rfind(and_equivalents[lang])
if last_and != -1:
full_amount = full_amount[:last_and]
return full_amount
def _expand_ordinal(m, lang="en"):
return num2words(int(m.group(1)), ordinal=True, lang=lang if lang != "cs" else "cz")
def _expand_number(m, lang="en"):
return num2words(int(m.group(0)), lang=lang if lang != "cs" else "cz")
def expand_numbers_multilingual(text, lang="en"):
if lang == "zh":
text = zh_num2words()(text)
else:
if lang in ["en", "ru"]:
text = re.sub(_comma_number_re, _remove_commas, text)
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(_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(_ordinal_re[lang], lambda m: _expand_ordinal(m, lang), text)
text = re.sub(_number_re, lambda m: _expand_number(m, lang), text)
return text
def lowercase(text):
return text.lower()
def collapse_whitespace(text):
return re.sub(_whitespace_re, " ", text)
def multilingual_cleaners(text, lang):
text = text.replace('"', "")
if lang == "tr":
text = text.replace("İ", "i")
text = text.replace("Ö", "ö")
text = text.replace("Ü", "ü")
text = lowercase(text)
try:
text = expand_numbers_multilingual(text, lang)
except:
pass
try:
text = expand_abbreviations_multilingual(text, lang)
except:
pass
try:
text = expand_symbols_multilingual(text, lang=lang)
except:
pass
text = collapse_whitespace(text)
return text
def basic_cleaners(text):
"""Basic pipeline that lowercases and collapses whitespace without transliteration."""
text = lowercase(text)
text = collapse_whitespace(text)
return 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)]
)
def japanese_cleaners(text, katsu):
text = katsu.romaji(text)
text = lowercase(text)
return text
def korean_transliterate(text):
r = Transliter(academic)
return r.translit(text)
DEFAULT_VOCAB_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "vocab.json")
class VoiceBpeTokenizer:
def __init__(self, vocab_file=DEFAULT_VOCAB_FILE):
self.tokenizer = None
if vocab_file is not None:
self.tokenizer = Tokenizer.from_file(vocab_file)
self.char_limits = {
"en": 10000,
"de": 253,
"fr": 273,
"es": 239,
"it": 213,
"pt": 203,
"pl": 224,
"zh": 82,
"ar": 166,
"cs": 186,
"ru": 182,
"nl": 251,
"tr": 226,
"ja": 71,
"hu": 224,
"ko": 95,
}
@cached_property
def katsu(self):
import cutlet
return cutlet.Cutlet()
def check_input_length(self, txt, lang):
lang = lang.split("-")[0] # remove the region
limit = self.char_limits.get(lang, 250)
# if len(txt) > limit:
# print(
# f"[!] Warning: The text length exceeds the character limit of {limit} for language '{lang}', this might cause truncated audio."
# )
def preprocess_text(self, txt, lang):
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)
if lang == "ko":
txt = korean_transliterate(txt)
elif lang == "ja":
txt = japanese_cleaners(txt, self.katsu)
elif lang == "hi":
# @manmay will implement this
txt = basic_cleaners(txt)
else:
raise NotImplementedError(f"Language '{lang}' is not supported.")
return txt
def encode(self, txt, lang):
lang = lang.split("-")[0] # remove the region
self.check_input_length(txt, lang)
txt = self.preprocess_text(txt, lang)
lang = "zh-cn" if lang == "zh" else lang
txt = f"[{lang}]{txt}"
txt = txt.replace(" ", "[SPACE]")
return self.tokenizer.encode(txt).ids
def decode(self, seq, skip_special_tokens=False):
if isinstance(seq, torch.Tensor):
seq = seq.cpu().numpy()
txt = self.tokenizer.decode(seq, skip_special_tokens=False).replace(" ", "")
txt = txt.replace("[SPACE]", " ")
txt = txt.replace("[STOP]", "")
# 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"],
skip_special_tokens: bool = False,
) -> List[str]:
"""
Convert a list of lists of token ids into a list of strings by calling decode.
Args:
sequences (`Union[List[int], List[List[int]], np.ndarray, torch.Tensor, tf.Tensor]`):
List of tokenized input ids. Can be obtained using the `__call__` method.
skip_special_tokens (`bool`, *optional*, defaults to `False`):
Whether or not to remove special tokens in the decoding.
kwargs (additional keyword arguments, *optional*):
Will be passed to the underlying model specific decode method.
Returns:
`List[str]`: The list of decoded sentences.
"""
return [
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):
def __len__(self):
return self.tokenizer.get_vocab_size()
def get_number_tokens(self):
return max(self.tokenizer.get_vocab().values()) + 1
def test_expand_numbers_multilingual():
test_cases = [
# English
("In 12.5 seconds.", "In twelve point five seconds.", "en"),
("There were 50 soldiers.", "There were fifty soldiers.", "en"),
("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'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 euros 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"),
("Es gab 50 Soldaten.", "Es gab fünfzig Soldaten.", "de"),
("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"),
# 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"),
# 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"),
# Portuguese
("Em 12,5 segundos.", "Em doze vírgula cinco segundos.", "pt"),
("Havia 50 soldados.", "Havia cinquenta soldados.", "pt"),
("Este é um 1º teste", "Este é um primeiro teste", "pt"),
("Isso custará $20 senhor.", "Isso custará vinte dólares senhor.", "pt"),
("Isso custará 20€ senhor.", "Isso custará vinte euros senhor.", "pt"),
(
"Isso custará 20,15€ senhor.",
"Isso custará vinte euros e quinze cêntimos senhor.",
"pt",
), # "cêntimos" should be "centavos" num2words issue
# 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"),
# Arabic
("في الـ 12,5 ثانية.", "في الـ اثنا عشر , خمسون ثانية.", "ar"),
("كان هناك 50 جنديًا.", "كان هناك خمسون جنديًا.", "ar"),
# ("ستكون النتيجة $20 يا سيد.", 'ستكون النتيجة عشرون دولار يا سيد.', 'ar'), # $ and € are mising from num2words
# ("ستكون النتيجة 20€ يا سيد.", 'ستكون النتيجة عشرون يورو يا سيد.', 'ar'),
# Czech
("Za 12,5 vteřiny.", "Za dvanáct celá pět vteřiny.", "cs"),
("Bylo tam 50 vojáků.", "Bylo tam padesát vojáků.", "cs"),
("To bude stát 20€ pane.", "To bude stát dvacet euro pane.", "cs"),
("To bude 20.15€ pane.", "To bude dvacet euro, patnáct centů pane.", "cs"),
# Russian
("Через 12.5 секунды.", "Через двенадцать запятая пять секунды.", "ru"),
("Там было 50 солдат.", "Там было пятьдесят солдат.", "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"),
("Dat wordt dan $20 meneer.", "Dat wordt dan twintig dollar meneer.", "nl"),
("Dat wordt dan 20€ meneer.", "Dat wordt dan twintig euro meneer.", "nl"),
# Chinese (Simplified)
("在12.5秒内", "在十二点五秒内", "zh"),
("有50名士兵", "有五十名士兵", "zh"),
# ("那将是$20先生", '那将是二十美元先生', 'zh'), currency doesn't work
# ("那将是20€先生", '那将是二十欧元先生', 'zh'),
# Turkish
# ("12,5 saniye içinde.", 'On iki virgül beş saniye içinde.', 'tr'), # decimal doesn't work for TR
("50 asker vardı.", "elli asker vardı.", "tr"),
("Bu 1. test", "Bu birinci test", "tr"),
# ("Bu 100.000,5.", 'Bu yüz bin virgül beş.', 'tr'),
# Hungarian
("12,5 másodperc alatt.", "tizenkettő egész öt tized másodperc alatt.", "hu"),
("50 katona volt.", "ötven katona volt.", "hu"),
("Ez az 1. teszt", "Ez az első teszt", "hu"),
# Korean
("12.5 초 안에.", "십이 점 다섯 초 안에.", "ko"),
("50 명의 병사가 있었다.", "오십 명의 병사가 있었다.", "ko"),
("이것은 1 번째 테스트입니다", "이것은 첫 번째 테스트입니다", "ko"),
]
for a, b, lang in test_cases:
out = expand_numbers_multilingual(a, lang=lang)
assert out == b, f"'{out}' vs '{b}'"
def test_abbreviations_multilingual():
test_cases = [
# English
("Hello Mr. Smith.", "Hello mister Smith.", "en"),
("Dr. Jones is here.", "doctor Jones is here.", "en"),
# Spanish
("Hola Sr. Garcia.", "Hola señor Garcia.", "es"),
("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"),
# 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"),
# 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"),
# Czech
("P. Novák", "pan Novák", "cs"),
("Dr. Vojtěch", "doktor Vojtěch", "cs"),
# Dutch
("Dhr. Jansen", "de heer Jansen", "nl"),
("Mevr. de Vries", "mevrouw de Vries", "nl"),
# Russian
("Здравствуйте Г-н Иванов.", "Здравствуйте господин Иванов.", "ru"),
("Д-р Смирнов здесь, чтобы увидеть вас.", "доктор Смирнов здесь, чтобы увидеть вас.", "ru"),
# Turkish
("Merhaba B. Yılmaz.", "Merhaba bay Yılmaz.", "tr"),
("Dr. Ayşe burada.", "doktor Ayşe burada.", "tr"),
# Hungarian
("Dr. Szabó itt van.", "doktor Szabó itt van.", "hu"),
]
for a, b, lang in test_cases:
out = expand_abbreviations_multilingual(a, lang=lang)
assert out == b, f"'{out}' vs '{b}'"
def test_symbols_multilingual():
test_cases = [
("I have 14% battery", "I have 14 percent battery", "en"),
("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"),
("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"),
("У меня 14% заряда", "У меня 14 процентов заряда", "ru"),
("Я буду @ дома", "Я буду собака дома", "ru"),
("Ik heb 14% batterij", "Ik heb 14 procent batterij", "nl"),
("Ik zie je @ het feest", "Ik zie je bij het feest", "nl"),
("لدي 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"),
("배터리 잔량이 14%입니다.", "배터리 잔량이 14 퍼센트입니다.", "ko"),
]
for a, b, lang in test_cases:
out = expand_symbols_multilingual(a, lang=lang)
assert out == b, f"'{out}' vs '{b}'"
if __name__ == "__main__":
test_expand_numbers_multilingual()
test_abbreviations_multilingual()
test_symbols_multilingual()
models/lyrics_utils/vocab.json
+15535
View File
File diff suppressed because it is too large Load Diff
models/lyrics_utils/zh_num2words.py
+1209
View File
File diff suppressed because it is too large Load Diff
music_dcae/__init__.py
View File
music_dcae/music_dcae_pipeline.py
+141
View File
@@ -0,0 +1,141 @@
import os
import torch
from diffusers import AutoencoderDC
import torchaudio
import torchvision.transforms as transforms
from diffusers.models.modeling_utils import ModelMixin
from diffusers.loaders import FromOriginalModelMixin
from diffusers.configuration_utils import ConfigMixin, register_to_config
try:
from .music_vocoder import ADaMoSHiFiGANV1
except ImportError:
from music_vocoder import ADaMoSHiFiGANV1
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_dcae_f8c8")
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):
super(MusicDCAE, self).__init__()
self.dcae = AutoencoderDC.from_pretrained(dcae_checkpoint_path)
self.vocoder = ADaMoSHiFiGANV1.from_pretrained(vocoder_checkpoint_path)
if source_sample_rate is None:
source_sample_rate = 48000
self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
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)))
self.mel_chunk_size = 1024
self.time_dimention_multiple = 8
self.latent_chunk_size = self.mel_chunk_size // self.time_dimention_multiple
self.scale_factor = 0.1786
self.shift_factor = -1.9091
def load_audio(self, audio_path):
audio, sr = torchaudio.load(audio_path)
return audio, sr
def forward_mel(self, audios):
mels = []
for i in range(len(audios)):
image = self.vocoder.mel_transform(audios[i])
mels.append(image)
mels = torch.stack(mels)
return mels
@torch.no_grad()
def encode(self, audios, audio_lengths=None, sr=None):
if audio_lengths is None:
audio_lengths = torch.tensor([audios.shape[2]] * audios.shape[0])
audio_lengths = audio_lengths.to(audios.device)
# audios: N x 2 x T, 48kHz
device = audios.device
dtype = audios.dtype
if sr is None:
sr = 48000
resampler = self.resampler
else:
resampler = torchaudio.transforms.Resample(sr, 44100).to(device).to(dtype)
audio = resampler(audios)
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)))
mels = self.forward_mel(audio)
mels = (mels - self.min_mel_value) / (self.max_mel_value - self.min_mel_value)
mels = self.transform(mels)
latents = []
for mel in mels:
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()
latents = (latents - self.shift_factor) * self.scale_factor
return latents, latent_lengths
@torch.no_grad()
def decode(self, latents, audio_lengths=None, sr=None):
latents = latents / self.scale_factor + self.shift_factor
pred_wavs = []
for latent in latents:
mels = self.dcae.decoder(latent.unsqueeze(0))
mels = mels * 0.5 + 0.5
mels = mels * (self.max_mel_value - self.min_mel_value) + self.min_mel_value
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)
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)]
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)
sr, pred_wavs = self.decode(latents=latents, audio_lengths=audio_lengths, sr=sr)
return sr, pred_wavs, latents, latent_lengths
if __name__ == "__main__":
audio, sr = torchaudio.load("test.wav")
audio_lengths = torch.tensor([audio.shape[1]])
audios = audio.unsqueeze(0)
# test encode only
model = MusicDCAE()
# latents, latent_lengths = model.encode(audios, audio_lengths)
# print("latents shape: ", latents.shape)
# print("latent_lengths: ", latent_lengths)
# test encode and decode
sr, pred_wavs, latents, latent_lengths = model(audios, audio_lengths, sr)
print("reconstructed wavs: ", pred_wavs[0].shape)
print("latents shape: ", latents.shape)
print("latent_lengths: ", latent_lengths)
print("sr: ", sr)
torchaudio.save("test_reconstructed.flac", pred_wavs[0], sr)
print("test_reconstructed.flac")
music_dcae/music_log_mel.py
+107
View File
@@ -0,0 +1,107 @@
import torch
import torch.nn as nn
from torch import Tensor
from torchaudio.transforms import MelScale
class LinearSpectrogram(nn.Module):
def __init__(
self,
n_fft=2048,
win_length=2048,
hop_length=512,
center=False,
mode="pow2_sqrt",
):
super().__init__()
self.n_fft = n_fft
self.win_length = win_length
self.hop_length = hop_length
self.center = center
self.mode = mode
self.register_buffer("window", torch.hann_window(win_length))
def forward(self, y: Tensor) -> Tensor:
if y.ndim == 3:
y = y.squeeze(1)
y = torch.nn.functional.pad(
y.unsqueeze(1),
(
(self.win_length - self.hop_length) // 2,
(self.win_length - self.hop_length + 1) // 2,
),
mode="reflect",
).squeeze(1)
dtype = y.dtype
spec = torch.stft(
y.float(),
self.n_fft,
hop_length=self.hop_length,
win_length=self.win_length,
window=self.window,
center=self.center,
pad_mode="reflect",
normalized=False,
onesided=True,
return_complex=True,
)
spec = torch.view_as_real(spec)
if self.mode == "pow2_sqrt":
spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
spec = spec.to(dtype)
return spec
class LogMelSpectrogram(nn.Module):
def __init__(
self,
sample_rate=44100,
n_fft=2048,
win_length=2048,
hop_length=512,
n_mels=128,
center=False,
f_min=0.0,
f_max=None,
):
super().__init__()
self.sample_rate = sample_rate
self.n_fft = n_fft
self.win_length = win_length
self.hop_length = hop_length
self.center = center
self.n_mels = n_mels
self.f_min = f_min
self.f_max = f_max or sample_rate // 2
self.spectrogram = LinearSpectrogram(n_fft, win_length, hop_length, center)
self.mel_scale = MelScale(
self.n_mels,
self.sample_rate,
self.f_min,
self.f_max,
self.n_fft // 2 + 1,
"slaney",
"slaney",
)
def compress(self, x: Tensor) -> Tensor:
return torch.log(torch.clamp(x, min=1e-5))
def decompress(self, x: Tensor) -> Tensor:
return torch.exp(x)
def forward(self, x: Tensor, return_linear: bool = False) -> Tensor:
linear = self.spectrogram(x)
x = self.mel_scale(linear)
x = self.compress(x)
# print(x.shape)
if return_linear:
return x, self.compress(linear)
return x
music_dcae/music_vocoder.py
+576
View File
@@ -0,0 +1,576 @@
import librosa
import torch
from torch import nn
from functools import partial
from math import prod
from typing import Callable, Tuple, List
import numpy as np
import torch.nn.functional as F
from torch.nn import Conv1d
from torch.nn.utils import weight_norm
from torch.nn.utils.parametrize import remove_parametrizations as remove_weight_norm
from diffusers.models.modeling_utils import ModelMixin
from diffusers.loaders import FromOriginalModelMixin
from diffusers.configuration_utils import ConfigMixin, register_to_config
try:
from music_log_mel import LogMelSpectrogram
except ImportError:
from .music_log_mel import LogMelSpectrogram
def drop_path(
x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
):
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
'survival rate' as the argument.
""" # noqa: E501
if drop_prob == 0.0 or not training:
return x
keep_prob = 1 - drop_prob
shape = (x.shape[0],) + (1,) * (
x.ndim - 1
) # work with diff dim tensors, not just 2D ConvNets
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
if keep_prob > 0.0 and scale_by_keep:
random_tensor.div_(keep_prob)
return x * random_tensor
class DropPath(nn.Module):
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" # noqa: E501
def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
super(DropPath, self).__init__()
self.drop_prob = drop_prob
self.scale_by_keep = scale_by_keep
def forward(self, x):
return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
def extra_repr(self):
return f"drop_prob={round(self.drop_prob,3):0.3f}"
class LayerNorm(nn.Module):
r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
shape (batch_size, height, width, channels) while channels_first corresponds to inputs
with shape (batch_size, channels, height, width).
""" # noqa: E501
def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
super().__init__()
self.weight = nn.Parameter(torch.ones(normalized_shape))
self.bias = nn.Parameter(torch.zeros(normalized_shape))
self.eps = eps
self.data_format = data_format
if self.data_format not in ["channels_last", "channels_first"]:
raise NotImplementedError
self.normalized_shape = (normalized_shape,)
def forward(self, x):
if self.data_format == "channels_last":
return F.layer_norm(
x, self.normalized_shape, self.weight, self.bias, self.eps
)
elif self.data_format == "channels_first":
u = x.mean(1, keepdim=True)
s = (x - u).pow(2).mean(1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.eps)
x = self.weight[:, None] * x + self.bias[:, None]
return x
class ConvNeXtBlock(nn.Module):
r"""ConvNeXt Block. There are two equivalent implementations:
(1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
(2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
We use (2) as we find it slightly faster in PyTorch
Args:
dim (int): Number of input channels.
drop_path (float): Stochastic depth rate. Default: 0.0
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0.
kernel_size (int): Kernel size for depthwise conv. Default: 7.
dilation (int): Dilation for depthwise conv. Default: 1.
""" # noqa: E501
def __init__(
self,
dim: int,
drop_path: float = 0.0,
layer_scale_init_value: float = 1e-6,
mlp_ratio: float = 4.0,
kernel_size: int = 7,
dilation: int = 1,
):
super().__init__()
self.dwconv = nn.Conv1d(
dim,
dim,
kernel_size=kernel_size,
padding=int(dilation * (kernel_size - 1) / 2),
groups=dim,
) # depthwise conv
self.norm = LayerNorm(dim, eps=1e-6)
self.pwconv1 = nn.Linear(
dim, int(mlp_ratio * dim)
) # pointwise/1x1 convs, implemented with linear layers
self.act = nn.GELU()
self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)
self.gamma = (
nn.Parameter(layer_scale_init_value *
torch.ones((dim)), requires_grad=True)
if layer_scale_init_value > 0
else None
)
self.drop_path = DropPath(
drop_path) if drop_path > 0.0 else nn.Identity()
def forward(self, x, apply_residual: bool = True):
input = x
x = self.dwconv(x)
x = x.permute(0, 2, 1) # (N, C, L) -> (N, L, C)
x = self.norm(x)
x = self.pwconv1(x)
x = self.act(x)
x = self.pwconv2(x)
if self.gamma is not None:
x = self.gamma * x
x = x.permute(0, 2, 1) # (N, L, C) -> (N, C, L)
x = self.drop_path(x)
if apply_residual:
x = input + x
return x
class ParallelConvNeXtBlock(nn.Module):
def __init__(self, kernel_sizes: List[int], *args, **kwargs):
super().__init__()
self.blocks = nn.ModuleList(
[
ConvNeXtBlock(kernel_size=kernel_size, *args, **kwargs)
for kernel_size in kernel_sizes
]
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.stack(
[block(x, apply_residual=False) for block in self.blocks] + [x],
dim=1,
).sum(dim=1)
class ConvNeXtEncoder(nn.Module):
def __init__(
self,
input_channels=3,
depths=[3, 3, 9, 3],
dims=[96, 192, 384, 768],
drop_path_rate=0.0,
layer_scale_init_value=1e-6,
kernel_sizes: Tuple[int] = (7,),
):
super().__init__()
assert len(depths) == len(dims)
self.channel_layers = nn.ModuleList()
stem = nn.Sequential(
nn.Conv1d(
input_channels,
dims[0],
kernel_size=7,
padding=3,
padding_mode="replicate",
),
LayerNorm(dims[0], eps=1e-6, data_format="channels_first"),
)
self.channel_layers.append(stem)
for i in range(len(depths) - 1):
mid_layer = nn.Sequential(
LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
nn.Conv1d(dims[i], dims[i + 1], kernel_size=1),
)
self.channel_layers.append(mid_layer)
block_fn = (
partial(ConvNeXtBlock, kernel_size=kernel_sizes[0])
if len(kernel_sizes) == 1
else partial(ParallelConvNeXtBlock, kernel_sizes=kernel_sizes)
)
self.stages = nn.ModuleList()
drop_path_rates = [
x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
]
cur = 0
for i in range(len(depths)):
stage = nn.Sequential(
*[
block_fn(
dim=dims[i],
drop_path=drop_path_rates[cur + j],
layer_scale_init_value=layer_scale_init_value,
)
for j in range(depths[i])
]
)
self.stages.append(stage)
cur += depths[i]
self.norm = LayerNorm(dims[-1], eps=1e-6, data_format="channels_first")
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, (nn.Conv1d, nn.Linear)):
nn.init.trunc_normal_(m.weight, std=0.02)
nn.init.constant_(m.bias, 0)
def forward(
self,
x: torch.Tensor,
) -> torch.Tensor:
for channel_layer, stage in zip(self.channel_layers, self.stages):
x = channel_layer(x)
x = stage(x)
return self.norm(x)
def init_weights(m, mean=0.0, std=0.01):
classname = m.__class__.__name__
if classname.find("Conv") != -1:
m.weight.data.normal_(mean, std)
def get_padding(kernel_size, dilation=1):
return (kernel_size * dilation - dilation) // 2
class ResBlock1(torch.nn.Module):
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
super().__init__()
self.convs1 = nn.ModuleList(
[
weight_norm(
Conv1d(
channels,
channels,
kernel_size,
1,
dilation=dilation[0],
padding=get_padding(kernel_size, dilation[0]),
)
),
weight_norm(
Conv1d(
channels,
channels,
kernel_size,
1,
dilation=dilation[1],
padding=get_padding(kernel_size, dilation[1]),
)
),
weight_norm(
Conv1d(
channels,
channels,
kernel_size,
1,
dilation=dilation[2],
padding=get_padding(kernel_size, dilation[2]),
)
),
]
)
self.convs1.apply(init_weights)
self.convs2 = nn.ModuleList(
[
weight_norm(
Conv1d(
channels,
channels,
kernel_size,
1,
dilation=1,
padding=get_padding(kernel_size, 1),
)
),
weight_norm(
Conv1d(
channels,
channels,
kernel_size,
1,
dilation=1,
padding=get_padding(kernel_size, 1),
)
),
weight_norm(
Conv1d(
channels,
channels,
kernel_size,
1,
dilation=1,
padding=get_padding(kernel_size, 1),
)
),
]
)
self.convs2.apply(init_weights)
def forward(self, x):
for c1, c2 in zip(self.convs1, self.convs2):
xt = F.silu(x)
xt = c1(xt)
xt = F.silu(xt)
xt = c2(xt)
x = xt + x
return x
def remove_weight_norm(self):
for conv in self.convs1:
remove_weight_norm(conv)
for conv in self.convs2:
remove_weight_norm(conv)
class HiFiGANGenerator(nn.Module):
def __init__(
self,
*,
hop_length: int = 512,
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]] = (
(1, 3, 5), (1, 3, 5), (1, 3, 5)),
num_mels: int = 128,
upsample_initial_channel: int = 512,
use_template: bool = True,
pre_conv_kernel_size: int = 7,
post_conv_kernel_size: int = 7,
post_activation: Callable = partial(nn.SiLU, inplace=True),
):
super().__init__()
assert (
prod(upsample_rates) == hop_length
), f"hop_length must be {prod(upsample_rates)}"
self.conv_pre = weight_norm(
nn.Conv1d(
num_mels,
upsample_initial_channel,
pre_conv_kernel_size,
1,
padding=get_padding(pre_conv_kernel_size),
)
)
self.num_upsamples = len(upsample_rates)
self.num_kernels = len(resblock_kernel_sizes)
self.noise_convs = nn.ModuleList()
self.use_template = use_template
self.ups = nn.ModuleList()
for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
c_cur = upsample_initial_channel // (2 ** (i + 1))
self.ups.append(
weight_norm(
nn.ConvTranspose1d(
upsample_initial_channel // (2**i),
upsample_initial_channel // (2 ** (i + 1)),
k,
u,
padding=(k - u) // 2,
)
)
)
if not use_template:
continue
if i + 1 < len(upsample_rates):
stride_f0 = np.prod(upsample_rates[i + 1:])
self.noise_convs.append(
Conv1d(
1,
c_cur,
kernel_size=stride_f0 * 2,
stride=stride_f0,
padding=stride_f0 // 2,
)
)
else:
self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
self.resblocks = nn.ModuleList()
for i in range(len(self.ups)):
ch = upsample_initial_channel // (2 ** (i + 1))
for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes):
self.resblocks.append(ResBlock1(ch, k, d))
self.activation_post = post_activation()
self.conv_post = weight_norm(
nn.Conv1d(
ch,
1,
post_conv_kernel_size,
1,
padding=get_padding(post_conv_kernel_size),
)
)
self.ups.apply(init_weights)
self.conv_post.apply(init_weights)
def forward(self, x, template=None):
x = self.conv_pre(x)
for i in range(self.num_upsamples):
x = F.silu(x, inplace=True)
x = self.ups[i](x)
if self.use_template:
x = x + self.noise_convs[i](template)
xs = None
for j in range(self.num_kernels):
if xs is None:
xs = self.resblocks[i * self.num_kernels + j](x)
else:
xs += self.resblocks[i * self.num_kernels + j](x)
x = xs / self.num_kernels
x = self.activation_post(x)
x = self.conv_post(x)
x = torch.tanh(x)
return x
def remove_weight_norm(self):
for up in self.ups:
remove_weight_norm(up)
for block in self.resblocks:
block.remove_weight_norm()
remove_weight_norm(self.conv_pre)
remove_weight_norm(self.conv_post)
class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
@register_to_config
def __init__(
self,
input_channels: int = 128,
depths: List[int] = [3, 3, 9, 3],
dims: List[int] = [128, 256, 384, 512],
drop_path_rate: float = 0.0,
kernel_sizes: Tuple[int] = (7,),
upsample_rates: Tuple[int] = (4, 4, 2, 2, 2, 2, 2),
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)),
num_mels: int = 512,
upsample_initial_channel: int = 1024,
use_template: bool = False,
pre_conv_kernel_size: int = 13,
post_conv_kernel_size: int = 13,
sampling_rate: int = 44100,
n_fft: int = 2048,
win_length: int = 2048,
hop_length: int = 512,
f_min: int = 40,
f_max: int = 16000,
n_mels: int = 128,
):
super().__init__()
self.backbone = ConvNeXtEncoder(
input_channels=input_channels,
depths=depths,
dims=dims,
drop_path_rate=drop_path_rate,
kernel_sizes=kernel_sizes,
)
self.head = HiFiGANGenerator(
hop_length=hop_length,
upsample_rates=upsample_rates,
upsample_kernel_sizes=upsample_kernel_sizes,
resblock_kernel_sizes=resblock_kernel_sizes,
resblock_dilation_sizes=resblock_dilation_sizes,
num_mels=num_mels,
upsample_initial_channel=upsample_initial_channel,
use_template=use_template,
pre_conv_kernel_size=pre_conv_kernel_size,
post_conv_kernel_size=post_conv_kernel_size,
)
self.sampling_rate = sampling_rate
self.mel_transform = LogMelSpectrogram(
sample_rate=sampling_rate,
n_fft=n_fft,
win_length=win_length,
hop_length=hop_length,
f_min=f_min,
f_max=f_max,
n_mels=n_mels,
)
self.eval()
@torch.no_grad()
def decode(self, mel):
y = self.backbone(mel)
y = self.head(y)
return y
@torch.no_grad()
def encode(self, x):
return self.mel_transform(x)
def forward(self, mel):
y = self.backbone(mel)
y = self.head(y)
return y
if __name__ == "__main__":
import soundfile as sf
x = "test_audio.flac"
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]
mel = model.encode(wav)
wav = model.decode(mel)[0].mT
sf.write("test_audio_vocoder_rec.flac", wav.cpu().numpy(), 44100)
pipeline_ace_step.py
+1129
View File
File diff suppressed because it is too large Load Diff
pyproject.toml
+15
View File
@@ -0,0 +1,15 @@
[project]
name = "ace-step"
description = "ACE-Step: A Step Towards Music Generation Foundation Model"
version = "1.0.0"
license = {file = "LICENSE"}
dependencies = ["diffusers", "librosa", "pypinyin", "pytorch_lightning", "soundfile", "hangul-romanize", "num2words", "spacy"]
[project.urls]
Repository = "https://github.com/billwuhao/ComfyUI_ACE-Step"
# Used by Comfy Registry https://comfyregistry.org
[tool.comfy]
PublisherId = "mw"
DisplayName = "ComfyUI_ACE-Step"
Icon = ""
requirements.txt
+8
View File
@@ -0,0 +1,8 @@
diffusers
librosa
pypinyin
pytorch_lightning
soundfile
hangul-romanize
num2words
spacy
schedulers/scheduling_flow_match_euler_discrete.py
+394
View File
@@ -0,0 +1,394 @@
# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union
import numpy as np
import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.utils import BaseOutput, logging
from diffusers.schedulers.scheduling_utils import SchedulerMixin
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
@dataclass
class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
"""
Output class for the scheduler's `step` function output.
Args:
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
denoising loop.
"""
prev_sample: torch.FloatTensor
class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
"""
Euler scheduler.
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
methods the library implements for all schedulers such as loading and saving.
Args:
num_train_timesteps (`int`, defaults to 1000):
The number of diffusion steps to train the model.
timestep_spacing (`str`, defaults to `"linspace"`):
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
shift (`float`, defaults to 1.0):
The shift value for the timestep schedule.
"""
_compatibles = []
order = 1
@register_to_config
def __init__(
self,
num_train_timesteps: int = 1000,
shift: float = 1.0,
use_dynamic_shifting=False,
base_shift: Optional[float] = 0.5,
max_shift: Optional[float] = 1.15,
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 = torch.from_numpy(timesteps).to(dtype=torch.float32)
sigmas = timesteps / num_train_timesteps
if not use_dynamic_shifting:
# when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
self.timesteps = sigmas * num_train_timesteps
self._step_index = None
self._begin_index = None
self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication
self.sigma_min = self.sigmas[-1].item()
self.sigma_max = self.sigmas[0].item()
@property
def step_index(self):
"""
The index counter for current timestep. It will increase 1 after each scheduler step.
"""
return self._step_index
@property
def begin_index(self):
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
"""
return self._begin_index
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
def set_begin_index(self, begin_index: int = 0):
"""
Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
Args:
begin_index (`int`):
The begin index for the scheduler.
"""
self._begin_index = begin_index
def scale_noise(
self,
sample: torch.FloatTensor,
timestep: Union[float, torch.FloatTensor],
noise: Optional[torch.FloatTensor] = None,
) -> torch.FloatTensor:
"""
Forward process in flow-matching
Args:
sample (`torch.FloatTensor`):
The input sample.
timestep (`int`, *optional*):
The current timestep in the diffusion chain.
Returns:
`torch.FloatTensor`:
A scaled input sample.
"""
# Make sure sigmas and timesteps have the same device and dtype as original_samples
sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
if sample.device.type == "mps" and torch.is_floating_point(timestep):
# mps does not support float64
schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
timestep = timestep.to(sample.device, dtype=torch.float32)
else:
schedule_timesteps = self.timesteps.to(sample.device)
timestep = timestep.to(sample.device)
# 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]
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]
else:
# add noise is called before first denoising step to create initial latent(img2img)
step_indices = [self.begin_index] * timestep.shape[0]
sigma = sigmas[step_indices].flatten()
while len(sigma.shape) < len(sample.shape):
sigma = sigma.unsqueeze(-1)
sample = sigma * noise + (1.0 - sigma) * sample
return sample
def _sigma_to_t(self, sigma):
return sigma * self.config.num_train_timesteps
def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
def set_timesteps(
self,
num_inference_steps: int = None,
device: Union[str, torch.device] = None,
sigmas: Optional[List[float]] = None,
mu: Optional[float] = None,
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
Args:
num_inference_steps (`int`):
The number of diffusion steps used when generating samples with a pre-trained model.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
"""
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`")
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
)
sigmas = timesteps / self.config.num_train_timesteps
if self.config.use_dynamic_shifting:
sigmas = self.time_shift(mu, 1.0, sigmas)
else:
sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
timesteps = sigmas * self.config.num_train_timesteps
self.timesteps = timesteps.to(device=device)
self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
self._step_index = None
self._begin_index = None
def index_for_timestep(self, timestep, schedule_timesteps=None):
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
indices = (schedule_timesteps == timestep).nonzero()
# The sigma index that is taken for the **very** first `step`
# is always the second index (or the last index if there is only 1)
# This way we can ensure we don't accidentally skip a sigma in
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
pos = 1 if len(indices) > 1 else 0
return indices[pos].item()
def _init_step_index(self, timestep):
if self.begin_index is None:
if isinstance(timestep, torch.Tensor):
timestep = timestep.to(self.timesteps.device)
self._step_index = self.index_for_timestep(timestep)
else:
self._step_index = self._begin_index
def step(
self,
model_output: torch.FloatTensor,
timestep: Union[float, torch.FloatTensor],
sample: torch.FloatTensor,
s_churn: float = 0.0,
s_tmin: float = 0.0,
s_tmax: float = float("inf"),
s_noise: float = 1.0,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
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
process from the learned model outputs (most often the predicted noise).
Args:
model_output (`torch.FloatTensor`):
The direct output from learned diffusion model.
timestep (`float`):
The current discrete timestep in the diffusion chain.
sample (`torch.FloatTensor`):
A current instance of a sample created by the diffusion process.
s_churn (`float`):
s_tmin (`float`):
s_tmax (`float`):
s_noise (`float`, defaults to 1.0):
Scaling factor for noise added to the sample.
generator (`torch.Generator`, *optional*):
A random number generator.
return_dict (`bool`):
Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
tuple.
Returns:
[`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
returned, otherwise a tuple is returned where the first element is the sample tensor.
"""
def logistic_function(x, L=0.9, U=1.1, x_0=0.0, k=1):
# L = Lower bound
# U = Upper bound
# x_0 = Midpoint (x corresponding to y = 1.0)
# k = Steepness, can adjust based on preference
if isinstance(x, torch.Tensor):
device_ = x.device
x = x.to(torch.float).cpu().numpy()
new_x = L + (U - L) / (1 + np.exp(-k * (x - x_0)))
if isinstance(new_x, np.ndarray):
new_x = torch.from_numpy(new_x).to(device_)
return new_x
self.omega_bef_rescale = omega
omega = logistic_function(omega, k=0.1)
self.omega_aft_rescale = omega
if (
isinstance(timestep, int)
or isinstance(timestep, torch.IntTensor)
or isinstance(timestep, torch.LongTensor)
):
raise ValueError(
(
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
" `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
" one of the `scheduler.timesteps` as a timestep."
),
)
if self.step_index is None:
self._init_step_index(timestep)
# Upcast to avoid precision issues when computing prev_sample
sample = sample.to(torch.float32)
sigma = self.sigmas[self.step_index]
sigma_next = self.sigmas[self.step_index + 1]
## --
## mean shift 1
dx = (sigma_next - sigma) * model_output
m = dx.mean()
# print(dx.shape) # torch.Size([1, 16, 128, 128])
# print(f'm: {m}') # m: -0.0014209747314453125
# raise NotImplementedError
dx_ = (dx - m) * omega + m
prev_sample = sample + dx_
# ## --
# ## mean shift 2
# m = model_output.mean()
# model_output_ = (model_output - m) * omega + m
# prev_sample = sample + (sigma_next - sigma) * model_output_
# ## --
# ## original
# prev_sample = sample + (sigma_next - sigma) * model_output * omega
# ## --
# ## spatial mean 1
# dx = (sigma_next - sigma) * model_output
# m = dx.mean(dim=(0, 1), keepdim=True)
# # print(dx.shape) # torch.Size([1, 16, 128, 128])
# # print(m.shape) # torch.Size([1, 1, 128, 128])
# # raise NotImplementedError
# dx_ = (dx - m) * omega + m
# prev_sample = sample + dx_
# ## --
# ## spatial mean 2
# m = model_output.mean(dim=(0, 1), keepdim=True)
# model_output_ = (model_output - m) * omega + m
# prev_sample = sample + (sigma_next - sigma) * model_output_
# ## --
# ## channel mean 1
# m = model_output.mean(dim=(2, 3), keepdim=True)
# # print(m.shape) # torch.Size([1, 16, 1, 1])
# model_output_ = (model_output - m) * omega + m
# prev_sample = sample + (sigma_next - sigma) * model_output_
# ## --
# ## channel mean 2
# dx = (sigma_next - sigma) * model_output
# m = dx.mean(dim=(2, 3), keepdim=True)
# # print(m.shape) # torch.Size([1, 16, 1, 1])
# dx_ = (dx - m) * omega + m
# prev_sample = sample + dx_
# ## --
# ## keep sample mean
# m_tgt = sample.mean()
# prev_sample_ = sample + (sigma_next - sigma) * model_output * omega
# m_src = prev_sample_.mean()
# prev_sample = prev_sample_ - m_src + m_tgt
# ## --
# ## test
# # print(sample.mean())
# prev_sample = sample + (sigma_next - sigma) * model_output * omega
# # raise NotImplementedError
# Cast sample back to model compatible dtype
prev_sample = prev_sample.to(model_output.dtype)
# upon completion increase step index by one
self._step_index += 1
if not return_dict:
return (prev_sample,)
return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
def __len__(self):
return self.config.num_train_timesteps
schedulers/scheduling_flow_match_heun_discrete.py
+348
View File
@@ -0,0 +1,348 @@
# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
from typing import Optional, Tuple, Union
import numpy as np
import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.utils import BaseOutput, logging
from diffusers.utils.torch_utils import randn_tensor
from diffusers.schedulers.scheduling_utils import SchedulerMixin
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
@dataclass
class FlowMatchHeunDiscreteSchedulerOutput(BaseOutput):
"""
Output class for the scheduler's `step` function output.
Args:
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
denoising loop.
"""
prev_sample: torch.FloatTensor
class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
"""
Heun scheduler.
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
methods the library implements for all schedulers such as loading and saving.
Args:
num_train_timesteps (`int`, defaults to 1000):
The number of diffusion steps to train the model.
timestep_spacing (`str`, defaults to `"linspace"`):
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
shift (`float`, defaults to 1.0):
The shift value for the timestep schedule.
"""
_compatibles = []
order = 2
@register_to_config
def __init__(
self,
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 = torch.from_numpy(timesteps).to(dtype=torch.float32)
sigmas = timesteps / num_train_timesteps
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
self.timesteps = sigmas * num_train_timesteps
self._step_index = None
self._begin_index = None
self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication
self.sigma_min = self.sigmas[-1].item()
self.sigma_max = self.sigmas[0].item()
@property
def step_index(self):
"""
The index counter for current timestep. It will increase 1 after each scheduler step.
"""
return self._step_index
@property
def begin_index(self):
"""
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
"""
return self._begin_index
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
def set_begin_index(self, begin_index: int = 0):
"""
Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
Args:
begin_index (`int`):
The begin index for the scheduler.
"""
self._begin_index = begin_index
def scale_noise(
self,
sample: torch.FloatTensor,
timestep: Union[float, torch.FloatTensor],
noise: Optional[torch.FloatTensor] = None,
) -> torch.FloatTensor:
"""
Forward process in flow-matching
Args:
sample (`torch.FloatTensor`):
The input sample.
timestep (`int`, *optional*):
The current timestep in the diffusion chain.
Returns:
`torch.FloatTensor`:
A scaled input sample.
"""
if self.step_index is None:
self._init_step_index(timestep)
sigma = self.sigmas[self.step_index]
sample = sigma * noise + (1.0 - sigma) * sample
return sample
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):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
Args:
num_inference_steps (`int`):
The number of diffusion steps used when generating samples with a pre-trained model.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
"""
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
)
sigmas = timesteps / self.config.num_train_timesteps
sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
timesteps = sigmas * self.config.num_train_timesteps
timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2)])
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:]])
# empty dt and derivative
self.prev_derivative = None
self.dt = None
self._step_index = None
self._begin_index = None
def index_for_timestep(self, timestep, schedule_timesteps=None):
if schedule_timesteps is None:
schedule_timesteps = self.timesteps
indices = (schedule_timesteps == timestep).nonzero()
# The sigma index that is taken for the **very** first `step`
# is always the second index (or the last index if there is only 1)
# This way we can ensure we don't accidentally skip a sigma in
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
pos = 1 if len(indices) > 1 else 0
return indices[pos].item()
def _init_step_index(self, timestep):
if self.begin_index is None:
if isinstance(timestep, torch.Tensor):
timestep = timestep.to(self.timesteps.device)
self._step_index = self.index_for_timestep(timestep)
else:
self._step_index = self._begin_index
@property
def state_in_first_order(self):
return self.dt is None
def step(
self,
model_output: torch.FloatTensor,
timestep: Union[float, torch.FloatTensor],
sample: torch.FloatTensor,
s_churn: float = 0.0,
s_tmin: float = 0.0,
s_tmax: float = float("inf"),
s_noise: float = 1.0,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
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
process from the learned model outputs (most often the predicted noise).
Args:
model_output (`torch.FloatTensor`):
The direct output from learned diffusion model.
timestep (`float`):
The current discrete timestep in the diffusion chain.
sample (`torch.FloatTensor`):
A current instance of a sample created by the diffusion process.
s_churn (`float`):
s_tmin (`float`):
s_tmax (`float`):
s_noise (`float`, defaults to 1.0):
Scaling factor for noise added to the sample.
generator (`torch.Generator`, *optional*):
A random number generator.
return_dict (`bool`):
Whether or not to return a [`~schedulers.scheduling_Heun_discrete.HeunDiscreteSchedulerOutput`] or
tuple.
Returns:
[`~schedulers.scheduling_Heun_discrete.HeunDiscreteSchedulerOutput`] or `tuple`:
If return_dict is `True`, [`~schedulers.scheduling_Heun_discrete.HeunDiscreteSchedulerOutput`] is
returned, otherwise a tuple is returned where the first element is the sample tensor.
"""
def logistic_function(x, L=0.9, U=1.1, x_0=0.0, k=1):
# L = Lower bound
# U = Upper bound
# x_0 = Midpoint (x corresponding to y = 1.0)
# k = Steepness, can adjust based on preference
if isinstance(x, torch.Tensor):
device_ = x.device
x = x.to(torch.float).cpu().numpy()
new_x = L + (U - L) / (1 + np.exp(-k * (x - x_0)))
if isinstance(new_x, np.ndarray):
new_x = torch.from_numpy(new_x).to(device_)
return new_x
self.omega_bef_rescale = omega
omega = logistic_function(omega, k=0.1)
self.omega_aft_rescale = omega
if (
isinstance(timestep, int)
or isinstance(timestep, torch.IntTensor)
or isinstance(timestep, torch.LongTensor)
):
raise ValueError(
(
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
" `HeunDiscreteScheduler.step()` is not supported. Make sure to pass"
" one of the `scheduler.timesteps` as a timestep."
),
)
if self.step_index is None:
self._init_step_index(timestep)
# Upcast to avoid precision issues when computing prev_sample
sample = sample.to(torch.float32)
if self.state_in_first_order:
sigma = self.sigmas[self.step_index]
sigma_next = self.sigmas[self.step_index + 1]
else:
# 2nd order / Heun's method
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
sigma_hat = sigma * (gamma + 1)
if gamma > 0:
noise = randn_tensor(
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
if self.state_in_first_order:
# 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
denoised = sample - model_output * sigma
# 2. convert to an ODE derivative for 1st order
derivative = (sample - denoised) / sigma_hat
# 3. Delta timestep
dt = sigma_next - sigma_hat
# store for 2nd order step
self.prev_derivative = derivative
self.dt = dt
self.sample = sample
else:
# 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
denoised = sample - model_output * sigma_next
# 2. 2nd order / Heun's method
derivative = (sample - denoised) / sigma_next
derivative = 0.5 * (self.prev_derivative + derivative)
# 3. take prev timestep & sample
dt = self.dt
sample = self.sample
# free dt and derivative
# Note, this puts the scheduler in "first order mode"
self.prev_derivative = None
self.dt = None
self.sample = None
# original sample way
# prev_sample = sample + derivative * dt
dx = derivative * dt
m = dx.mean()
dx_ = (dx - m) * omega + m
prev_sample = sample + dx_
# Cast sample back to model compatible dtype
prev_sample = prev_sample.to(model_output.dtype)
# upon completion increase step index by one
self._step_index += 1
if not return_dict:
return (prev_sample,)
return FlowMatchHeunDiscreteSchedulerOutput(prev_sample=prev_sample)
def __len__(self):
return self.config.num_train_timesteps
workflow-examples/ace-step-gen.json
+1
View File
File diff suppressed because one or more lines are too long
workflow-examples/edit.json
+1
View File
File diff suppressed because one or more lines are too long
workflow-examples/extend.json
+1
View File
File diff suppressed because one or more lines are too long
workflow-examples/rapainting.json
+1
View File
File diff suppressed because one or more lines are too long