mrleo1nid/ACE-Step
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #17 from fakerybakery/main

Browse Source 
Make ACE-Step `pip`-installable package
This commit is contained in:
Sayo
2025-05-07 23:34:34 +08:00
committed by GitHub
parent 9bf891fb28 df2b9d5f77
commit be5c795257
44 changed files with 3078 additions and 1405 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -176,6 +176,7 @@ cython_debug/
!requirements.txt
*.log
*.flac
*.wav
minio_config.yaml
.history/*
__pycache__/*
@@ -200,3 +201,4 @@ app_demo.py
ui/components_demo.py
data_sampler_demo.py
pipeline_ace_step_demo.py
*.wav
Dockerfile
+2 -2
View File
@@ -34,7 +34,7 @@ WORKDIR /app
# Clone the repository
RUN git clone https://github.com/ace-step/ACE-Step.git .
# Install specific PyTorch version compatible with CUDA 12.2
# Install specific PyTorch version compatible with CUDA 12.6
RUN pip3 install --no-cache-dir --upgrade pip && \
pip3 install --no-cache-dir hf_transfer peft && \
pip3 install --no-cache-dir -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu126
@@ -55,4 +55,4 @@ HEALTHCHECK --interval=60s --timeout=10s --start-period=5s --retries=5 \
CMD curl -f http://localhost:7865/ || exit 1
# Command to run the application with GPU support
CMD ["python3", "app.py", "--server_name", "0.0.0.0", "--bf16", "true"]
CMD ["python3", "acestep/gui.py", "--server_name", "0.0.0.0", "--bf16", "true"]
README.md
+73 -89
View File
@@ -1,22 +1,21 @@
<h1 align="center">ACE-Step</h1>
<h1 align="center">A Step Towards Music Generation Foundation Model</h1>
<p align="center">
<a href="https://ace-step.github.io/">Project</a> |
<a href="https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B">Hugging Face</a> |
<a href="https://modelscope.cn/models/ACE-Step/ACE-Step-v1-3.5B">ModelScope</a> |
<a href="https://huggingface.co/spaces/ACE-Step/ACE-Step">Space Demo</a> |
<a href="https://discord.gg/rjAZz2xBdG">Discord</a>
</p>
# ACE-Step
A Step Towards Music Generation Foundation Model
[Project Page](https://ace-step.github.io/) |
[Checkpoints](https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B) |
[Space Demo](https://huggingface.co/spaces/ACE-Step/ACE-Step) |
[Discord](https://discord.gg/rjAZz2xBdG)
---
<p align="center">
<img src="./fig/orgnization_logos.png" width="100%" alt="StepFun Logo">
<img src="./assets/orgnization_logos.png" width="100%" alt="StepFun Logo">
</p>
## Table of Contents
- [Features](#-features)
- [Installation](#-installation)
- [Features](#-features)
- [Usage](#-usage)
- [Train](#-train)
@@ -24,21 +23,43 @@
- 🚀 2025.05.06: Open source demo code and model
## TODOs📋
## 📋 Roadmap
- [x] Release training code 🔥
- [x] Release LoRA training code 🔥
- [ ] Release RapMachine lora 🎤
- [ ] Release RapMachine LoRA 🎤
- [ ] Release ControlNet training code 🔥
- [ ] Release Singing2Accompaniment controlnet 🎮
- [ ] Release Singing2Accompaniment ControlNet 🎮
- [ ] Release evaluation performance and technical report 📄
## 📦 Installation
### Quick Start
```bash
pip install git+https://github.com/ace-step/ACE-Step
```
On Windows, you may need to run the following before running the command above:
```bash
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
```
Then, launch the GUI:
```bash
acestep # Launch the GUI
```
The GUI works on Windows, macOS, and Linux.
## 🏗️ Architecture
<p align="center">
<img src="./fig/ACE-Step_framework.png" width="100%" alt="ACE-Step Framework">
<img src="./assets/ACE-Step_framework.png" width="100%" alt="ACE-Step Framework">
</p>
## 📝 Abstract
We introduce ACE-Step, a novel open-source foundation model for music generation that overcomes key limitations of existing approaches and achieves state-of-the-art performance through a holistic architectural design. Current methods face inherent trade-offs between generation speed, musical coherence, and controllability. For instance, LLM-based models (e.g., Yue, SongGen) excel at lyric alignment but suffer from slow inference and structural artifacts. Diffusion models (e.g., DiffRhythm), on the other hand, enable faster synthesis but often lack long-range structural coherence.
@@ -50,42 +71,49 @@ Rather than building yet another end-to-end text-to-music pipeline, our vision i
## ✨ Features
<p align="center">
<img src="./fig/application_map.png" width="100%" alt="ACE-Step Framework">
<img src="./assets/application_map.png" width="100%" alt="ACE-Step Framework">
</p>
### 🎯 Baseline Quality
#### 🌈 Diverse Styles & Genres
- 🎸 Supports all mainstream music styles with various description formats including short tags, descriptive text, or use-case scenarios
- 🎷 Capable of generating music across different genres with appropriate instrumentation and style
#### 🌍 Multiple Languages
- 🗣️ Supports 19 languages with top 10 well-performing languages including:
- 🇺🇸 English, 🇨🇳 Chinese, 🇷🇺 Russian, 🇪🇸 Spanish, 🇯🇵 Japanese, 🇩🇪 German, 🇫🇷 French, 🇵🇹 Portuguese, 🇮🇹 Italian, 🇰🇷 Korean
- ⚠️ Due to data imbalance, less common languages may underperform
#### 🎻 Instrumental Styles
- 🎹 Supports various instrumental music generation across different genres and styles
- 🎺 Capable of producing realistic instrumental tracks with appropriate timbre and expression for each instrument
- 🎼 Can generate complex arrangements with multiple instruments while maintaining musical coherence
#### 🎤 Vocal Techniques
- 🎙️ Capable of rendering various vocal styles and techniques with good quality
- 🗣️ Supports different vocal expressions including various singing techniques and styles
### 🎛️ Controllability
#### 🔄 Variations Generation
- ⚙️ Implemented using training-free, inference-time optimization techniques
- 🌊 Flow-matching model generates initial noise, then uses trigFlow's noise formula to add additional Gaussian noise
- 🎚️ Adjustable mixing ratio between original initial noise and new Gaussian noise to control variation degree
#### 🎨 Repainting
- 🖌️ Implemented by adding noise to the target audio input and applying mask constraints during the ODE process
- 🔍 When input conditions change from the original generation, only specific aspects can be modified while preserving the rest
- 🔀 Can be combined with Variations Generation techniques to create localized variations in style, lyrics, or vocals
#### ✏️ Lyric Editing
- 💡 Innovatively applies flow-edit technology to enable localized lyric modifications while preserving melody, vocals, and accompaniment
- 🔄 Works with both generated content and uploaded audio, greatly enhancing creative possibilities
- ℹ️ Current limitation: can only modify small segments of lyrics at once to avoid distortion, but multiple edits can be applied sequentially
@@ -93,11 +121,13 @@ Rather than building yet another end-to-end text-to-music pipeline, our vision i
### 🚀 Applications
#### 🎤 Lyric2Vocal (LoRA)
- 🔊 Based on a LoRA fine-tuned on pure vocal data, allowing direct generation of vocal samples from lyrics
- 🛠️ Offers numerous practical applications such as vocal demos, guide tracks, songwriting assistance, and vocal arrangement experimentation
- ⏱️ Provides a quick way to test how lyrics might sound when sung, helping songwriters iterate faster
#### 📝 Text2Samples (LoRA)
- 🎛️ Similar to Lyric2Vocal, but fine-tuned on pure instrumental and sample data
- 🎵 Capable of generating conceptual music production samples from text descriptions
- 🧰 Useful for quickly creating instrument loops, sound effects, and musical elements for production
@@ -105,16 +135,19 @@ Rather than building yet another end-to-end text-to-music pipeline, our vision i
### 🔮 Coming Soon
#### 🎤 RapMachine
- 🔥 Fine-tuned on pure rap data to create an AI system specialized in rap generation
- 🏆 Expected capabilities include AI rap battles and narrative expression through rap
- 📚 Rap has exceptional storytelling and expressive capabilities, offering extraordinary application potential
#### 🎛️ StemGen
- 🎚️ A controlnet-lora trained on multi-track data to generate individual instrument stems
- 🎯 Takes a reference track and specified instrument (or instrument reference audio) as input
- 🎹 Outputs an instrument stem that complements the reference track, such as creating a piano accompaniment for a flute melody or adding jazz drums to a lead guitar
#### 🎤 Singing2Accompaniment
- 🔄 The reverse process of StemGen, generating a mixed master track from a single vocal track
- 🎵 Takes a vocal track and specified style as input to produce a complete vocal accompaniment
- 🎸 Creates full instrumental backing that complements the input vocals, making it easy to add professional-sounding accompaniment to any vocal recording
@@ -134,89 +167,40 @@ We have evaluated ACE-Step across different hardware setups, yielding the follow
We use RTF (Real-Time Factor) to measure the performance of ACE-Step. Higher values indicate faster generation speed. 27.27x means to generate 1 minute of music, it takes 2.2 seconds (60/27.27). The performance is measured on a single GPU with batch size 1 and 27 steps.
## 💻 Installation
### Prerequisites
* Make sure you have Python installed. You can download it from [python.org](https://www.python.org/).
* You will also need either `Conda` (recommended) or `venv`.
### Environment Setup
It is highly recommended to use a virtual environment to manage project dependencies and avoid conflicts. Choose **one** of the following methods (Conda or venv):
#### Option 1: Using Conda
1. **Create the environment** named `ace_step` with Python 3.10:
```bash
conda create -n ace_step python=3.10 -y
```
2. **Activate the environment:**
```bash
conda activate ace_step
```
#### Option 2: Using venv
1. **Ensure you are using the correct Python version.**
2. **Create the virtual environment** (commonly named `venv`):
```bash
python -m venv venv
```
3. **Activate the environment:**
* **On Windows (cmd.exe):**
```bash
venv\Scripts\activate.bat
```
* **On Windows (PowerShell):**
```powershell
.\venv\Scripts\Activate.ps1
```
*(If you encounter execution policy errors, you might need to run `Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope Process` first)*
* **On Linux / macOS (bash/zsh):**
```bash
source venv/bin/activate
```
3. **Install dependencies** from the `requirements.txt` file:
for macOS/Linux users:
```bash
pip install -r requirements.txt
```
for Windows users:
```bash
# Install PyTorch, TorchAudio, and TorchVision for Windows
# replace cu126 with your CUDA version
# replace torchvision and torchaudio with your version
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
# then install other dependencies
pip install -r requirements.txt
```
## 🚀 Usage
![Demo Interface](fig/demo_interface.png)
![Demo Interface](assets/demo_interface.png)
### 🔍 Basic Usage
```bash
python app.py
acestep
```
### ⚙️ Advanced Usage
```bash
python app.py --checkpoint_path /path/to/checkpoint --port 7865 --device_id 0 --share true --bf16 true
acestep --checkpoint_path /path/to/checkpoint --port 7865 --device_id 0 --share true --bf16 true
```
If you are using MacOS, please use `--bf16 false` to avoid errors.
If you are using macOS, please use `--bf16 false` to avoid errors.
#### 🔍 API Usage
```python
from acestep.api import ACEStep
model = ACEStep()
audio_out = model.infer(
prompt="upbeat pop, catchy melody, female singer",
lyrics="[verse]\nSun is shining bright today\nFeeling happy, come what may",
audio_duration=5.0, # 5 seconds
infer_step=20 # Fewer steps for speed
)
audio_out.save_wav("output.wav")
```
#### 🛠️ Command Line Arguments
@@ -268,11 +252,11 @@ The ACE-Step interface provides several tabs for different music generation and
- 📐 Specify left and right extension lengths
- 🔍 Choose the source audio to extend
## Examples
## 📂 Examples
The `examples/input_params` directory contains sample input parameters that can be used as references for generating music.
## Train
## 🔨 Train
### Prerequisites
1. Prepare the environment as described in the installation section.
acestep/__init__.py
+7
View File
@@ -0,0 +1,7 @@
"""
ACE-Step: A Step Towards Music Generation Foundation Model
https://github.com/ace-step/ACE-Step
Apache 2.0 License
"""
apg_guidance.py → acestep/apg_guidance.py
+20 -5
View File
@@ -25,7 +25,9 @@ def project(
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).to(device_type), v0_orthogonal.to(dtype).to(device_type)
return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(
device_type
)
def apg_forward(
@@ -65,7 +67,11 @@ def cfg_double_condition_forward(
guidance_scale_text,
guidance_scale_lyric,
):
return (1 - guidance_scale_text) * uncond_output + (guidance_scale_text - guidance_scale_lyric) * only_text_cond_output + guidance_scale_lyric * cond_output
return (
(1 - guidance_scale_text) * uncond_output
+ (guidance_scale_text - guidance_scale_lyric) * only_text_cond_output
+ guidance_scale_lyric * cond_output
)
def optimized_scale(positive_flat, negative_flat):
@@ -82,14 +88,23 @@ def optimized_scale(positive_flat, negative_flat):
return st_star
def cfg_zero_star(noise_pred_with_cond, noise_pred_uncond, guidance_scale, i, zero_steps=1, use_zero_init=True):
def cfg_zero_star(
noise_pred_with_cond,
noise_pred_uncond,
guidance_scale,
i,
zero_steps=1,
use_zero_init=True,
):
bsz = noise_pred_with_cond.shape[0]
positive_flat = noise_pred_with_cond.view(bsz, -1)
negative_flat = noise_pred_uncond.view(bsz, -1)
alpha = optimized_scale(positive_flat, negative_flat)
alpha = alpha.view(bsz, 1, 1, 1)
if (i <= zero_steps) and use_zero_init:
noise_pred = noise_pred_with_cond * 0.
noise_pred = noise_pred_with_cond * 0.0
else:
noise_pred = noise_pred_uncond * alpha + guidance_scale * (noise_pred_with_cond - noise_pred_uncond * alpha)
noise_pred = noise_pred_uncond * alpha + guidance_scale * (
noise_pred_with_cond - noise_pred_uncond * alpha
)
return noise_pred
data_sampler.py → acestep/data_sampler.py
View File
acestep/gui.py
+71
View File
@@ -0,0 +1,71 @@
"""
ACE-Step: A Step Towards Music Generation Foundation Model
https://github.com/ace-step/ACE-Step
Apache 2.0 License
"""
import os
import click
from acestep.ui.components import create_main_demo_ui
from acestep.pipeline_ace_step import ACEStepPipeline
from acestep.data_sampler import DataSampler
@click.command()
@click.option(
"--checkpoint_path",
type=str,
default="",
help="Path to the checkpoint directory. Downloads automatically if empty.",
)
@click.option(
"--server_name",
type=str,
default="127.0.0.1",
help="The server name to use for the Gradio app.",
)
@click.option(
"--port", type=int, default=7865, help="The port to use for the Gradio app."
)
@click.option("--device_id", type=int, default=0, help="The CUDA device ID to use.")
@click.option(
"--share",
is_flag=True,
default=False,
help="Whether to create a public, shareable link for the Gradio app.",
)
@click.option(
"--bf16",
is_flag=True,
default=True,
help="Whether to use bfloat16 precision. Turn off if using MPS.",
)
@click.option(
"--torch_compile", is_flag=True, default=False, help="Whether to use torch.compile."
)
def main(checkpoint_path, server_name, port, device_id, share, bf16, torch_compile):
"""
Main function to launch the ACE Step pipeline demo.
"""
os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
model_demo = ACEStepPipeline(
checkpoint_dir=checkpoint_path,
dtype="bfloat16" if bf16 else "float32",
torch_compile=torch_compile,
)
data_sampler = DataSampler()
demo = create_main_demo_ui(
text2music_process_func=model_demo.__call__,
sample_data_func=data_sampler.sample,
)
demo.launch(server_name=server_name, server_port=port, share=share)
if __name__ == "__main__":
main()
language_segmentation/LangSegment.py → acestep/language_segmentation/LangSegment.py
+349 -143
View File
@@ -11,6 +11,41 @@ https://github.com/adbar/py3langid
Projects:
https://github.com/juntaosun/LangSegment
LICENSE:
py3langid - Language Identifier
BSD 3-Clause License
Modifications (fork): Copyright (c) 2021, Adrien Barbaresi.
Original code: Copyright (c) 2011 Marco Lui <saffsd@gmail.com>.
Based on research by Marco Lui and Tim Baldwin.
All rights reserved.
Redistribution and use in source and binary forms, with or without modification, are
permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.
3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER ``AS IS'' AND ANY EXPRESS OR IMPLIED
WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
import os
@@ -29,12 +64,7 @@ from collections import defaultdict
# For probability normalization in library use, the user must instantiate their own . An example of such usage is as follows:
from py3langid.langid import LanguageIdentifier, MODEL_FILE
# Digital processing
try:from .utils.num import num2str
except ImportError:
try:from utils.num import num2str
except ImportError as e:
raise e
from acestep.language_segmentation.utils.num import num2str
# -----------------------------------
# 更新日志：新版本分词更加精准。
@@ -107,16 +137,16 @@ class LangSSML:
def __init__(self):
# 纯数字
self._zh_numerals_number = {
'0': '零',
'1': '一',
'2': '二',
'3': '三',
'4': '四',
'5': '五',
'6': '六',
'7': '七',
'8': '八',
'9': '九'
"0": "零",
"1": "一",
"2": "二",
"3": "三",
"4": "四",
"5": "五",
"6": "六",
"7": "七",
"8": "八",
"9": "九",
}
# 将2024/8/24, 2024-08, 08-24, 24 标准化“年月日”
@@ -124,36 +154,46 @@ class LangSSML:
def _format_chinese_data(self, date_str: str):
# 处理日期格式
input_date = date_str
if date_str is None or date_str.strip() == "":return ""
if date_str is None or date_str.strip() == "":
return ""
date_str = re.sub(r"[\/\._|年|月]", "-", date_str)
date_str = re.sub(r"日", r"", date_str)
date_arrs = date_str.split(' ')
date_arrs = date_str.split(" ")
if len(date_arrs) == 1 and ":" in date_arrs[0]:
time_str = date_arrs[0]
date_arrs = []
else:
time_str = date_arrs[1] if len(date_arrs) >= 2 else ""
def nonZero(num, cn, func=None):
if func is not None:num=func(num)
if func is not None:
num = func(num)
return f"{num}{cn}" if num is not None and num != "" and num != "0" else ""
f_number = self.to_chinese_number
f_currency = self.to_chinese_currency
# year, month, day
year_month_day = ""
if len(date_arrs) > 0:
year, month, day = "", "", ""
parts = date_arrs[0].split('-')
parts = date_arrs[0].split("-")
if len(parts) == 3: # 格式为 YYYY-MM-DD
year, month, day = parts
elif len(parts) == 2: # 格式为 MM-DD 或 YYYY-MM
if len(parts[0]) == 4: # 年-月
year, month = parts
else:month, day = parts # 月-日
else:
month, day = parts # 月-日
elif len(parts[0]) > 0: # 仅有月-日或年
if len(parts[0]) == 4:
year = parts[0]
else:day = parts[0]
year,month,day = nonZero(year,"年",f_number),nonZero(month,"月",f_currency),nonZero(day,"日",f_currency)
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)
@@ -163,24 +203,35 @@ class LangSSML:
hours, minutes, seconds = time_arrs
elif len(time_arrs) == 2: # H/M
hours, minutes = time_arrs
elif len(time_arrs[0]) > 0:hours = f'{time_arrs[0]}点' # H
elif len(time_arrs[0]) > 0:
hours = f"{time_arrs[0]}点" # H
if len(time_arrs) > 1:
hours, minutes, seconds = nonZero(hours,"点",f_currency),nonZero(minutes,"分",f_currency),nonZero(seconds,"秒",f_currency)
hours_minutes_seconds = re.sub(r"([点|分|秒])+",r"\1",f"{hours}{minutes}{seconds}")
hours, minutes, seconds = (
nonZero(hours, "点", f_currency),
nonZero(minutes, "分", f_currency),
nonZero(seconds, "秒", f_currency),
)
hours_minutes_seconds = re.sub(
r"([点|分|秒])+", r"\1", f"{hours}{minutes}{seconds}"
)
output_date = f"{year_month_day}{hours_minutes_seconds}"
return output_date
# 【SSML】number=中文大写数字读法（单字）
# Chinese Numbers(single word)
def to_chinese_number(self, num: str):
pattern = r'(\d+)'
pattern = r"(\d+)"
zh_numerals = self._zh_numerals_number
arrs = re.split(pattern, num)
output = ""
for item in arrs:
if re.match(pattern, item):
output += ''.join(zh_numerals[digit] if digit in zh_numerals else "" for digit in str(item))
else:output += 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
@@ -194,13 +245,14 @@ class LangSSML:
# 【SSML】currency=按金额发音。
# Digital processing from GPT_SoVITS num.py （thanks）
def to_chinese_currency(self, num: str):
pattern = r'(\d+)'
pattern = r"(\d+)"
arrs = re.split(pattern, num)
output = ""
for item in arrs:
if re.match(pattern, item):
output += num2str(item)
else:output += item
else:
output += item
output = output.replace(".", "点")
return output
@@ -225,7 +277,7 @@ class LangSegment:
# 可自定义语言匹配标签：カスタマイズ可能な言語対応タグ:사용자 지정 가능한 언어 일치 태그:
# Customizable language matching tags: These are supported，이 표현들은 모두 지지합니다
# <zh>你好<zh> , <ja>佐々木</ja> , <en>OK<en> , <ko>오빠</ko> 这些写法均支持
self.SYMBOLS_PATTERN = r'(<([a-zA-Z|-]*)>(.*?)<\/*[a-zA-Z|-]*>)'
self.SYMBOLS_PATTERN = r"(<([a-zA-Z|-]*)>(.*?)<\/*[a-zA-Z|-]*>)"
# 语言过滤组功能, 可以指定保留语言。不在过滤组中的语言将被清除。您可随心搭配TTS语音合成所支持的语言。
# 언어 필터 그룹 기능을 사용하면 예약된 언어를 지정할 수 있습니다. 필터 그룹에 없는 언어는 지워집니다. TTS 텍스트에서 지원하는 언어를 원하는 대로 일치시킬 수 있습니다.
@@ -287,7 +339,7 @@ class LangSegment:
self.keepPinyin = False
# DEFINITION
self.PARSE_TAG = re.compile(r'(⑥\$*\d+[\d]{6,}⑥)')
self.PARSE_TAG = re.compile(r"(⑥\$*\d+[\d]{6,}⑥)")
self.LangSSML = LangSSML()
@@ -300,26 +352,26 @@ class LangSegment:
self._lang_eos = None
def _is_english_word(self, word):
return bool(re.match(r'^[a-zA-Z]+$', word))
return bool(re.match(r"^[a-zA-Z]+$", word))
def _is_chinese(self, word):
for char in word:
if '\u4e00' <= char <= '\u9fff':
if "\u4e00" <= char <= "\u9fff":
return True
return False
def _is_japanese_kana(self, word):
pattern = re.compile(r'[\u3040-\u309F\u30A0-\u30FF]+')
pattern = re.compile(r"[\u3040-\u309F\u30A0-\u30FF]+")
matches = pattern.findall(word)
return len(matches) > 0
def _insert_english_uppercase(self, word):
modified_text = re.sub(r'(?<!\b)([A-Z])', r' \1', word)
modified_text = modified_text.strip('-')
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)
return re.sub(r"(?<!^)(?=[A-Z])", " ", word)
def _statistics(self, language, text):
# Language word statistics:
@@ -328,13 +380,16 @@ class LangSegment:
self._lang_count = defaultdict(int)
lang_count = self._lang_count
if not "|" in language:
lang_count[language] += int(len(text)*2) if language == "zh" else len(text)
lang_count[language] += (
int(len(text) * 2) if language == "zh" else len(text)
)
self._lang_count = lang_count
def _clear_text_number(self, text):
if text == "\n":return text,False # Keep Line Breaks
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
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):
@@ -342,12 +397,15 @@ class LangSegment:
clear_text, is_number = self._clear_text_number(text)
# Merge the same language and save the results
preData = words[-1] if len(words) > 0 else None
if symbol is not None:pass
if symbol is not None:
pass
elif preData is not None and preData["symbol"] is None:
if len(clear_text) == 0:language = preData["lang"]
elif is_number == True:language = preData["lang"]
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):
if preData["lang"] == language:
self._statistics(preData["lang"], text)
text = preData["text"] + text
preData["text"] = text
@@ -357,70 +415,108 @@ class LangSegment:
words.pop()
elif is_number == True:
priority_language = self._get_filters_string()[:2]
if priority_language in "ja-zh-en-ko-fr-vi":language = priority_language
if priority_language in "ja-zh-en-ko-fr-vi":
language = priority_language
data = {"lang": language, "text": text, "score": score, "symbol": symbol}
filters = self.Langfilters
if filters is None or len(filters) == 0 or "?" in language or \
language in filters or language in filters[0] or \
filters[0] == "*" or filters[0] in "alls-mixs-autos":
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 = ""
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)
if language == "en":
text = self._insert_english_uppercase(text)
# text = re.sub(r'[(（）)]', ',' , text) # Keep it.
text_waits = self._text_waits
ispre_waits = len(text_waits) > 0
preResult = text_waits.pop() if ispre_waits else None
if preResult is None:preResult = words[-1] if len(words) > 0 else None
if preResult is None:
preResult = words[-1] if len(words) > 0 else None
if preResult and ("|" in preResult["lang"]):
pre_lang = preResult["lang"]
if language in pre_lang:preResult["lang"] = language = language.split("|")[0]
else:preResult["lang"]=pre_lang.split("|")[0]
if ispre_waits:preResult = self._saveData(words,preResult["lang"],preResult["text"],preResult["score"],preResult["symbol"])
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)
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"])
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
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)
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
if text is None or (text.strip()) == "":
return None, 0.0
arrs = self._split_camel_case(text).split(" ")
langs = []
for t in arrs:
if len(t.strip()) <= 3:continue
if len(t.strip()) <= 3:
continue
language, score = self.langid.classify(t)
langs.append({"lang": language})
if len(langs) == 0:return None , 0.0
return Counter([item['lang'] for item in langs]).most_common(1)[0][0],1.0
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"):
if (
score is not None
and isinstance(score, np.generic)
and hasattr(score, "item")
):
score = score.item()
score = round(score, 3)
return language, score
@@ -432,49 +528,80 @@ class LangSegment:
def _parse_language(self, words, segment):
LANG_JA = "ja"
LANG_ZH = "zh"
LANG_ZH_JA = f'{LANG_ZH}|{LANG_JA}'
LANG_JA_ZH = f'{LANG_JA}|{LANG_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]+)')
regex_pattern = re.compile(r"([^\w\s]+)")
lines = regex_pattern.split(segment)
lines_max = len(lines)
LANG_EOS = self._lang_eos
for index, text in enumerate(lines):
if len(text) == 0:continue
if len(text) == 0:
continue
EOS = index >= (lines_max - 1)
nextId = index + 1
nextText = lines[nextId] if not EOS else ""
nextPunc = len(re.sub(regex_pattern,'',re.sub(r'\n+','',nextText)).strip()) == 0
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}'
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)
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}'
lines[nextId] = f"{text}{nextText}"
continue
language, score = self._lang_classify(cleans_text)
prev_language, prev_text = self._get_prev_data(words)
if language != LANG_ZH and all('\u4e00' <= c <= '\u9fff' for c in re.sub(r'\s','',cleans_text)):language,score = LANG_ZH,1
if language != LANG_ZH and all(
"\u4e00" <= c <= "\u9fff" for c in re.sub(r"\s", "", cleans_text)
):
language, score = LANG_ZH, 1
if len(cleans_text) <= 5 and self._is_chinese(cleans_text):
filters_string = self._get_filters_string()
if score < self.LangPriorityThreshold and len(filters_string) > 0:
index_ja , index_zh = filters_string.find(LANG_JA) , filters_string.find(LANG_ZH)
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
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
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}|…"
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)
@@ -518,7 +645,8 @@ class LangSegment:
return value
def _pattern_symbols(self, item, text):
if text is None:return text
if text is None:
return text
tag, pattern, process = item
matches = pattern.findall(text)
if len(matches) == 1 and "".join(matches[0]) == text:
@@ -549,19 +677,25 @@ class LangSegment:
enablePreview = self.EnablePreview
if enablePreview == True:
# Experimental: Other language support
regex_pattern = re.compile(r'(.*?[。.?？!！]+[\n]{,1})')
regex_pattern = re.compile(r"(.*?[。.?？!！]+[\n]{,1})")
lines = regex_pattern.split(text)
for index, text in enumerate(lines):
if len(text.strip()) == 0:continue
if len(text.strip()) == 0:
continue
cleans_text = self._cleans_text(text)
language, score = self._lang_classify(cleans_text)
if language not in filters:
language, score = self._mean_processing(cleans_text)
if language is None or score <= 0.0:continue
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"
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
@@ -627,10 +761,12 @@ class LangSegment:
segments = re.split(self.PARSE_TAG, text)
segments_len = len(segments) - 1
for index, text in enumerate(segments):
if root_tag:self._lang_eos = index >= segments_len
if root_tag:
self._lang_eos = index >= segments_len
if self.PARSE_TAG.match(text):
process, data = text_cache[text]
if process:process(words , data)
if process:
process(words, data)
else:
self._parse_language(words, text)
return words
@@ -638,19 +774,31 @@ class LangSegment:
def _merge_results(self, words):
new_word = []
for index, cur_data in enumerate(words):
if "symbol" in cur_data:del cur_data["symbol"]
if index == 0:new_word.append(cur_data)
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)
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]*)+)')
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(rf'(([【《（(“‘"\']*[LANGUAGE]+[\W\s]*)+)')
# Get custom language filter
filters = self.Langfilters
filters = filters if filters is not None else ""
@@ -663,51 +811,107 @@ class LangSegment:
# -------------------------------------------------------------------------------------------------------
# Preview feature, other language support
enablePreview = self.EnablePreview
if "fr" in filters or \
"vi" in filters:enablePreview = True
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 "đơưăáàảãạắằẳẵặấầẩẫậéèẻẽẹếềểễệíìỉĩịóòỏõọốồổỗộớờởỡợúùủũụứừửữựôâêơưỷỹ"
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.
(
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.
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))
lines = re.findall(r".*\n*", re.sub(self.PARSE_TAG, "", text))
for index, text in enumerate(lines):
if len(text.strip()) == 0:continue
if len(text.strip()) == 0:
continue
self._lang_eos = False
self._text_cache = {}
for item in process_list:
text = self._pattern_symbols(item, text)
cur_word = self._process_tags([], text, True)
if len(cur_word) == 0:continue
if len(cur_word) == 0:
continue
cur_data = cur_word[0] if len(cur_word) > 0 else None
pre_data = words[-1] if len(words) > 0 else None
if cur_data and pre_data and cur_data["lang"] == pre_data["lang"] \
and cur_data["symbol"] == False and pre_data["symbol"] :
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)
if self.isLangMerge == True:
words = self._merge_results(words)
lang_count = self._lang_count
if lang_count and len(lang_count) > 0:
lang_count = dict(sorted(lang_count.items(), key=lambda x: x[1], reverse=True))
lang_count = dict(
sorted(lang_count.items(), key=lambda x: x[1], reverse=True)
)
lang_count = list(lang_count.items())
self._lang_count = lang_count
return words
@@ -732,12 +936,19 @@ class LangSegment:
def getCounts(self):
lang_count = self._lang_count
if lang_count is not None:return lang_count
if lang_count is not None:
return lang_count
text_langs = self._text_langs
if text_langs is None or len(text_langs) == 0:return [("zh",0)]
if text_langs is None or len(text_langs) == 0:
return [("zh", 0)]
lang_counts = defaultdict(int)
for d in text_langs:lang_counts[d['lang']] += int(len(d['text'])*2) if d['lang'] == "zh" else len(d['text'])
lang_counts = dict(sorted(lang_counts.items(), key=lambda x: x[1], reverse=True))
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
@@ -748,7 +959,8 @@ class LangSegment:
return []
# lasts
text_langs = self._text_langs
if self._text_lasts == text and text_langs is not None:return text_langs
if self._text_lasts == text and text_langs is not None:
return text_langs
# parse
self._text_waits = []
self._lang_count = None
@@ -760,6 +972,7 @@ class LangSegment:
def classify(self, text: str):
return self.getTexts(text)
def printList(langlist):
"""
功能：打印数组结果
@@ -776,7 +989,6 @@ def printList(langlist):
pass
def main():
# -----------------------------------
@@ -795,11 +1007,9 @@ def main():
# 输入示例3：（包含日文，中文）Input Example 1: (including Japanese, Chinese)
# text = "明日、私たちは海辺にバカンスに行きます。你会说日语吗：“中国語、話せますか” 你的日语真好啊！"
# 输入示例4：（包含日文，中文，韩语，英文）Input Example 4: (including Japanese, Chinese, Korean, English)
# text = "你的名字叫<ja>佐々木？<ja>吗？韩语中的안녕 오빠读什么呢？あなたの体育の先生は誰ですか? 此次发布会带来了四款iPhone 15系列机型和三款Apple Watch等一系列新品，这次的iPad Air采用了LCD屏幕"
# 试验性支持："fr"法语 , "vi"越南语 , "ru"俄语 , "th"泰语。Experimental: Other language support.
langsegment = LangSegment()
langsegment.setfilters(["fr", "vi", "ja", "zh", "ko", "en", "ru", "th"])
@@ -814,13 +1024,10 @@ Tôi thích nghe nhạc vào những ngày mưa.
ฉันชอบฟังเพลงในวันที่ฝนตก
"""
# 进行分词：（接入TTS项目仅需一行代码调用）Segmentation: (Only one line of code is required to access the TTS project)
langlist = langsegment.getTexts(text)
printList(langlist)
# 语种统计:Language statistics:
print("\n===================【语种统计】===================")
# 获取所有语种数组结果，根据内容字数降序排列
@@ -834,7 +1041,6 @@ Tôi thích nghe nhạc vào những ngày mưa.
print(f"输入内容的主要语言为 = {lang} ，字数 = {count}")
print("==================================================\n")
# 分词输出：lang=语言，text=内容。Word output: lang = language, text = content
# ===================【打印结果】===================
# {'lang': 'zh', 'text': '你的名字叫'}
acestep/language_segmentation/__init__.py
+9
View File
@@ -0,0 +1,9 @@
from acestep.language_segmentation.LangSegment import LangSegment
# release
__version__ = "0.3.5"
# develop
__develop__ = "dev-0.0.1"
acestep/language_segmentation/utils/__init__.py
+1
View File
@@ -0,0 +1 @@
# This file intentionally left blank for Python to recognize the directory as a package.
language_segmentation/utils/num.py → acestep/language_segmentation/utils/num.py
+62 -56
View File
@@ -21,19 +21,21 @@ 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: '亿',
})
DIGITS = {str(i): tran for i, tran in enumerate("零一二三四五六七八九")}
UNITS = OrderedDict(
{
1: "十",
2: "百",
3: "千",
4: "万",
8: "亿",
}
)
COM_QUANTIFIERS = '(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)'
COM_QUANTIFIERS = "(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)"
# 分数表达式
RE_FRAC = re.compile(r'(-?)(\d+)/(\d+)')
RE_FRAC = re.compile(r"(-?)(\d+)/(\d+)")
def replace_frac(match) -> str:
@@ -54,7 +56,7 @@ def replace_frac(match) -> str:
# 百分数表达式
RE_PERCENTAGE = re.compile(r'(-?)(\d+(\.\d+)?)%')
RE_PERCENTAGE = re.compile(r"(-?)(\d+(\.\d+)?)%")
def replace_percentage(match) -> str:
@@ -74,7 +76,7 @@ def replace_percentage(match) -> str:
# 整数表达式
# 带负号的整数 -10
RE_INTEGER = re.compile(r'(-)' r'(\d+)')
RE_INTEGER = re.compile(r"(-)" r"(\d+)")
def replace_negative_num(match) -> str:
@@ -94,7 +96,7 @@ def replace_negative_num(match) -> str:
# 编号-无符号整形
# 00078
RE_DEFAULT_NUM = re.compile(r'\d{3}\d*')
RE_DEFAULT_NUM = re.compile(r"\d{3}\d*")
def replace_default_num(match):
@@ -112,15 +114,11 @@ def replace_default_num(match):
# RE_ASMD = re.compile(
# r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))([\+\-\×÷=])((-?)((\d+)(\.\d+)?)|(\.(\d+)))')
RE_ASMD = re.compile(
r'((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))([\+\-\×÷=])((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))')
r"((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))([\+\-\×÷=])((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))"
)
asmd_map = {"+": "加", "-": "减", "×": "乘", "÷": "除", "=": "等于"}
asmd_map = {
'+': '加',
'-': '减',
'×': '乘',
'÷': '除',
'=': '等于'
}
def replace_asmd(match) -> str:
"""
@@ -134,24 +132,25 @@ def replace_asmd(match) -> str:
# 次方专项
RE_POWER = re.compile(r'[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]+')
RE_POWER = re.compile(r"[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]+")
power_map = {
'⁰': '0',
'¹': '1',
'²': '2',
'³': '3',
'⁴': '4',
'⁵': '5',
'⁶': '6',
'⁷': '7',
'⁸': '8',
'⁹': '9',
'ˣ': 'x',
'ʸ': 'y',
'ⁿ': 'n'
"⁰": "0",
"¹": "1",
"²": "2",
"³": "3",
"⁴": "4",
"⁵": "5",
"⁶": "6",
"⁷": "7",
"⁸": "8",
"⁹": "9",
"ˣ": "x",
"ʸ": "y",
"ⁿ": "n",
}
def replace_power(match) -> str:
"""
Args:
@@ -168,10 +167,10 @@ def replace_power(match) -> str:
# 数字表达式
# 纯小数
RE_DECIMAL_NUM = re.compile(r'(-?)((\d+)(\.\d+))' r'|(\.(\d+))')
RE_DECIMAL_NUM = re.compile(r"(-?)((\d+)(\.\d+))" r"|(\.(\d+))")
# 正整数 + 量词
RE_POSITIVE_QUANTIFIERS = re.compile(r"(\d+)([多余几\+])?" + COM_QUANTIFIERS)
RE_NUMBER = re.compile(r'(-?)((\d+)(\.\d+)?)' r'|(\.(\d+))')
RE_NUMBER = re.compile(r"(-?)((\d+)(\.\d+)?)" r"|(\.(\d+))")
def replace_positive_quantifier(match) -> str:
@@ -221,7 +220,9 @@ RE_RANGE = re.compile(
[-~] # 匹配范围分隔符
((-?)((\d+)(\.\d+)?)) # 匹配范围结束的负数或正数（整数或小数）
(?![\d\+\-\×÷=]) # 使用正向前瞻以确保数字范围之后没有其他数字和操作符
""", re.VERBOSE)
""",
re.VERBOSE,
)
def replace_range(match) -> str:
@@ -240,7 +241,9 @@ def replace_range(match) -> str:
# ~至表达式
RE_TO_RANGE = re.compile(
r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)[~]((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)')
r"((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)[~]((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)"
)
def replace_to_range(match) -> str:
"""
@@ -249,58 +252,61 @@ def replace_to_range(match) -> str:
Returns:
str
"""
result = match.group(0).replace('~', '至')
result = match.group(0).replace("~", "至")
return result
def _get_value(value_string: str, use_zero: bool = True) -> List[str]:
stripped = value_string.lstrip('0')
stripped = value_string.lstrip("0")
if len(stripped) == 0:
return []
elif len(stripped) == 1:
if use_zero and len(stripped) < len(value_string):
return [DIGITS['0'], DIGITS[stripped]]
return [DIGITS["0"], DIGITS[stripped]]
else:
return [DIGITS[stripped]]
else:
largest_unit = next(
power for power in reversed(UNITS.keys()) if power < len(stripped))
power for power in reversed(UNITS.keys()) if power < len(stripped)
)
first_part = value_string[:-largest_unit]
second_part = value_string[-largest_unit:]
return _get_value(first_part) + [UNITS[largest_unit]] + _get_value(
second_part)
return _get_value(first_part) + [UNITS[largest_unit]] + _get_value(second_part)
def verbalize_cardinal(value_string: str) -> str:
if not value_string:
return ''
return ""
# 000 -> '零' , 0 -> '零'
value_string = value_string.lstrip('0')
value_string = value_string.lstrip("0")
if len(value_string) == 0:
return DIGITS['0']
return DIGITS["0"]
result_symbols = _get_value(value_string)
# verbalized number starting with '一十*' is abbreviated as `十*`
if len(result_symbols) >= 2 and result_symbols[0] == DIGITS[
'1'] and result_symbols[1] == UNITS[1]:
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)
return "".join(result_symbols)
def verbalize_digit(value_string: str, alt_one=False) -> str:
result_symbols = [DIGITS[digit] for digit in value_string]
result = ''.join(result_symbols)
result = "".join(result_symbols)
if alt_one:
result = result.replace("一", "幺")
return result
def num2str(value_string: str) -> str:
integer_decimal = value_string.split('.')
integer_decimal = value_string.split(".")
if len(integer_decimal) == 1:
integer = integer_decimal[0]
decimal = ''
decimal = ""
elif len(integer_decimal) == 2:
integer, decimal = integer_decimal
else:
@@ -310,12 +316,12 @@ def num2str(value_string: str) -> str:
result = verbalize_cardinal(integer)
decimal = decimal.rstrip('0')
decimal = decimal.rstrip("0")
if decimal:
# '.22' is verbalized as '零点二二'
# '3.20' is verbalized as '三点二
result = result if result else "零"
result += '点' + verbalize_digit(decimal)
result += "点" + verbalize_digit(decimal)
return result
language_segmentation/utils/__init__.py → acestep/models/__init__.py
View File
models/ace_step_transformer.py → acestep/models/ace_step_transformer.py
+148 -39
View File
@@ -31,9 +31,15 @@ from .lyrics_utils.lyric_encoder import ConformerEncoder as LyricEncoder
def cross_norm(hidden_states, controlnet_input):
# input N x T x c
mean_hidden_states, std_hidden_states = hidden_states.mean(dim=(1,2), keepdim=True), hidden_states.std(dim=(1,2), keepdim=True)
mean_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
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
@@ -45,17 +51,27 @@ class Qwen2RotaryEmbedding(nn.Module):
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.base = base
inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
inv_freq = 1.0 / (
self.base
** (
torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device)
/ self.dim
)
)
self.register_buffer("inv_freq", inv_freq, persistent=False)
# Build here to make `torch.jit.trace` work.
self._set_cos_sin_cache(
seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
seq_len=max_position_embeddings,
device=self.inv_freq.device,
dtype=torch.get_default_dtype(),
)
def _set_cos_sin_cache(self, seq_len, device, dtype):
self.max_seq_len_cached = seq_len
t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
t = torch.arange(
self.max_seq_len_cached, device=device, dtype=torch.int64
).type_as(self.inv_freq)
freqs = torch.outer(t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
@@ -82,8 +98,12 @@ class T2IFinalLayer(nn.Module):
def __init__(self, hidden_size, patch_size=[16, 1], out_channels=256):
super().__init__()
self.norm_final = nn.RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.linear = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True)
self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5)
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
@@ -95,14 +115,28 @@ class T2IFinalLayer(nn.Module):
# 4 unpatchify
new_height, new_width = 1, hidden_states.size(1)
hidden_states = hidden_states.reshape(
shape=(hidden_states.shape[0], new_height, new_width, self.patch_size[0], self.patch_size[1], self.out_channels)
shape=(
hidden_states.shape[0],
new_height,
new_width,
self.patch_size[0],
self.patch_size[1],
self.out_channels,
)
).contiguous()
hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
output = hidden_states.reshape(
shape=(hidden_states.shape[0], self.out_channels, new_height * self.patch_size[0], new_width * self.patch_size[1])
shape=(
hidden_states.shape[0],
self.out_channels,
new_height * self.patch_size[0],
new_width * self.patch_size[1],
)
).contiguous()
if width > new_width:
output = torch.nn.functional.pad(output, (0, width - new_width, 0, 0), 'constant', 0)
output = torch.nn.functional.pad(
output, (0, width - new_width, 0, 0), "constant", 0
)
elif width < new_width:
output = output[:, :, :, :width]
return output
@@ -131,9 +165,25 @@ class PatchEmbed(nn.Module):
super().__init__()
patch_size_h, patch_size_w = patch_size
self.early_conv_layers = nn.Sequential(
nn.Conv2d(in_channels, in_channels*256, kernel_size=patch_size, stride=patch_size, padding=0, bias=bias),
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)
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
@@ -153,7 +203,9 @@ class Transformer2DModelOutput(BaseOutput):
proj_losses: Optional[Tuple[Tuple[str, torch.Tensor]]] = None
class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
class ACEStepTransformer2DModel(
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
):
_supports_gradient_checkpointing = True
@register_to_config
@@ -217,9 +269,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
)
self.num_layers = num_layers
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
self.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))
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)
@@ -229,25 +287,30 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# lyric
self.lyric_embs = nn.Embedding(lyric_encoder_vocab_size, lyric_hidden_size)
self.lyric_encoder = LyricEncoder(input_size=lyric_hidden_size, static_chunk_size=0)
self.lyric_encoder = LyricEncoder(
input_size=lyric_hidden_size, static_chunk_size=0
)
self.lyric_proj = nn.Linear(lyric_hidden_size, self.inner_dim)
projector_dim = 2 * self.inner_dim
self.projectors = nn.ModuleList([
self.projectors = nn.ModuleList(
[
nn.Sequential(
nn.Linear(self.inner_dim, projector_dim),
nn.SiLU(),
nn.Linear(projector_dim, projector_dim),
nn.SiLU(),
nn.Linear(projector_dim, ssl_dim),
) for ssl_dim in ssl_latent_dims
])
)
for ssl_dim in ssl_latent_dims
]
)
self.ssl_latent_dims = ssl_latent_dims
self.ssl_encoder_depths = ssl_encoder_depths
self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction='mean')
self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction="mean")
self.ssl_names = ssl_names
self.proj_in = PatchEmbed(
@@ -258,11 +321,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
bias=True,
)
self.final_layer = T2IFinalLayer(self.inner_dim, patch_size=patch_size, out_channels=out_channels)
self.final_layer = T2IFinalLayer(
self.inner_dim, patch_size=patch_size, out_channels=out_channels
)
self.gradient_checkpointing = False
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
def enable_forward_chunking(
self, chunk_size: Optional[int] = None, dim: int = 0
) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -281,7 +348,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# By default chunk size is 1
chunk_size = chunk_size or 1
def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
def fn_recursive_feed_forward(
module: torch.nn.Module, chunk_size: int, dim: int
):
if hasattr(module, "set_chunk_feed_forward"):
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
@@ -302,7 +371,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
):
# N x T x D
lyric_embs = self.lyric_embs(lyric_token_idx)
prompt_prenet_out, _mask = self.lyric_encoder(lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1)
prompt_prenet_out, _mask = self.lyric_encoder(
lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1
)
prompt_prenet_out = self.lyric_proj(prompt_prenet_out)
return prompt_prenet_out
@@ -331,8 +402,17 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
lyric_mask=lyric_mask,
)
encoder_hidden_states = torch.cat([encoder_spk_hidden_states, encoder_text_hidden_states, encoder_lyric_hidden_states], dim=1)
encoder_hidden_mask = torch.cat([speaker_mask, text_attention_mask, lyric_mask], dim=1)
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(
@@ -344,12 +424,16 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
timestep: Optional[torch.Tensor],
ssl_hidden_states: Optional[List[torch.Tensor]] = None,
output_length: int = 0,
block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
block_controlnet_hidden_states: Optional[
Union[List[torch.Tensor], torch.Tensor]
] = None,
controlnet_scale: Union[float, torch.Tensor] = 1.0,
return_dict: bool = True,
):
embedded_timestep = self.timestep_embedder(self.time_proj(timestep).to(dtype=hidden_states.dtype))
embedded_timestep = self.timestep_embedder(
self.time_proj(timestep).to(dtype=hidden_states.dtype)
)
temb = self.t_block(embedded_timestep)
hidden_states = self.proj_in(hidden_states)
@@ -361,8 +445,12 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
inner_hidden_states = []
rotary_freqs_cis = self.rotary_emb(hidden_states, seq_len=hidden_states.shape[1])
encoder_rotary_freqs_cis = self.rotary_emb(encoder_hidden_states, seq_len=encoder_hidden_states.shape[1])
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):
@@ -377,7 +465,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
return custom_forward
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
ckpt_kwargs: Dict[str, Any] = (
{"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
)
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states=hidden_states,
@@ -406,9 +496,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
inner_hidden_states.append(hidden_states)
proj_losses = []
if len(inner_hidden_states) > 0 and ssl_hidden_states is not None and len(ssl_hidden_states) > 0:
if (
len(inner_hidden_states) > 0
and ssl_hidden_states is not None
and len(ssl_hidden_states) > 0
):
for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names):
for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(
inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names
):
if ssl_hidden_state is None:
continue
# 1. N x T x D1 -> N x D x D2
@@ -416,9 +512,20 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# 3. projection loss
bs = inner_hidden_state.shape[0]
proj_loss = 0.0
for i, (z, z_tilde) in enumerate(zip(ssl_hidden_state, est_ssl_hidden_state)):
for i, (z, z_tilde) in enumerate(
zip(ssl_hidden_state, est_ssl_hidden_state)
):
# 2. interpolate
z_tilde = F.interpolate(z_tilde.unsqueeze(0).transpose(1, 2), size=len(z), mode='linear', align_corners=False).transpose(1, 2).squeeze(0)
z_tilde = (
F.interpolate(
z_tilde.unsqueeze(0).transpose(1, 2),
size=len(z),
mode="linear",
align_corners=False,
)
.transpose(1, 2)
.squeeze(0)
)
z_tilde = torch.nn.functional.normalize(z_tilde, dim=-1)
z = torch.nn.functional.normalize(z, dim=-1)
@@ -445,7 +552,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
lyric_mask: Optional[torch.LongTensor] = None,
timestep: Optional[torch.Tensor] = None,
ssl_hidden_states: Optional[List[torch.Tensor]] = None,
block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
block_controlnet_hidden_states: Optional[
Union[List[torch.Tensor], torch.Tensor]
] = None,
controlnet_scale: Union[float, torch.Tensor] = 1.0,
return_dict: bool = True,
):
models/attention.py → acestep/models/attention.py
+15 -3
View File
@@ -23,10 +23,18 @@ from diffusers.models.normalization import RMSNorm
try:
# from .dcformer import DCMHAttention
from .customer_attention_processor import Attention, CustomLiteLAProcessor2_0, CustomerAttnProcessor2_0
from .customer_attention_processor import (
Attention,
CustomLiteLAProcessor2_0,
CustomerAttnProcessor2_0,
)
except ImportError:
# from dcformer import DCMHAttention
from customer_attention_processor import Attention, CustomLiteLAProcessor2_0, CustomerAttnProcessor2_0
from customer_attention_processor import (
Attention,
CustomLiteLAProcessor2_0,
CustomerAttnProcessor2_0,
)
logger = logging.get_logger(__name__)
@@ -55,13 +63,16 @@ def t2i_modulate(x, shift, scale):
return x * (1 + scale) + shift
def get_same_padding(kernel_size: Union[int, Tuple[int, ...]]) -> Union[int, Tuple[int, ...]]:
def get_same_padding(
kernel_size: Union[int, Tuple[int, ...]],
) -> Union[int, Tuple[int, ...]]:
if isinstance(kernel_size, tuple):
return tuple([get_same_padding(ks) for ks in kernel_size])
else:
assert kernel_size % 2 > 0, f"kernel size {kernel_size} should be odd number"
return kernel_size // 2
class ConvLayer(nn.Module):
def __init__(
self,
@@ -187,6 +198,7 @@ class LinearTransformerBlock(nn.Module):
"""
A Sana block with global shared adaptive layer norm (adaLN-single) conditioning.
"""
def __init__(
self,
dim,
models/config.json → acestep/models/config.json
View File
models/customer_attention_processor.py → acestep/models/customer_attention_processor.py
+104 -29
View File
@@ -78,12 +78,16 @@ class CustomLiteLAProcessor2_0:
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
hidden_states = hidden_states.view(
batch_size, channel, height * width
).transpose(1, 2)
if encoder_hidden_states is not None:
context_input_ndim = encoder_hidden_states.ndim
if context_input_ndim == 4:
batch_size, channel, height, width = encoder_hidden_states.shape
encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
encoder_hidden_states = encoder_hidden_states.view(
batch_size, channel, height * width
).transpose(1, 2)
batch_size = hidden_states.shape[0]
@@ -94,7 +98,11 @@ class CustomLiteLAProcessor2_0:
value = attn.to_v(hidden_states)
# `context` projections.
has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
has_encoder_hidden_state_proj = (
hasattr(attn, "add_q_proj")
and hasattr(attn, "add_k_proj")
and hasattr(attn, "add_v_proj")
)
if encoder_hidden_states is not None and has_encoder_hidden_state_proj:
encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
@@ -114,7 +122,11 @@ class CustomLiteLAProcessor2_0:
head_dim = inner_dim // attn.heads
query = query.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
key = key.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1).transpose(-1, -2)
key = (
key.transpose(-1, -2)
.reshape(batch_size, attn.heads, head_dim, -1)
.transpose(-1, -2)
)
value = value.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
# RoPE需要 [B, H, S, D] 输入
@@ -140,17 +152,33 @@ class CustomLiteLAProcessor2_0:
if attention_mask is not None:
# attention_mask: [B, S] -> [B, 1, S, 1]
attention_mask = attention_mask[:, None, :, None].to(key.dtype) # [B, 1, S, 1]
query = query * attention_mask.permute(0, 1, 3, 2) # [B, H, S, D] * [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维度
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]
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]
value = value * encoder_attention_mask.permute(
0, 1, 3, 2
) # [B, h, D, S_enc] * [B, 1, 1, S_enc]
query = self.kernel_func(query)
key = self.kernel_func(key)
@@ -168,14 +196,20 @@ class CustomLiteLAProcessor2_0:
hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps)
hidden_states = hidden_states.view(batch_size, attn.heads * head_dim, -1).permute(0, 2, 1)
hidden_states = hidden_states.view(
batch_size, attn.heads * head_dim, -1
).permute(0, 2, 1)
hidden_states = hidden_states.to(dtype)
if encoder_hidden_states is not None:
encoder_hidden_states = encoder_hidden_states.to(dtype)
# Split the attention outputs.
if encoder_hidden_states is not None and not attn.is_cross_attention and has_encoder_hidden_state_proj:
if (
encoder_hidden_states is not None
and not attn.is_cross_attention
and has_encoder_hidden_state_proj
):
hidden_states, encoder_hidden_states = (
hidden_states[:, :hidden_states_len],
hidden_states[:, hidden_states_len:],
@@ -185,13 +219,22 @@ class CustomLiteLAProcessor2_0:
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if encoder_hidden_states is not None and not attn.context_pre_only and not attn.is_cross_attention and hasattr(attn, "to_add_out"):
if (
encoder_hidden_states is not None
and not attn.context_pre_only
and not attn.is_cross_attention
and hasattr(attn, "to_add_out")
):
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
hidden_states = hidden_states.transpose(-1, -2).reshape(
batch_size, channel, height, width
)
if encoder_hidden_states is not None and context_input_ndim == 4:
encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(
batch_size, channel, height, width
)
if torch.get_autocast_gpu_dtype() == torch.float16:
hidden_states = hidden_states.clip(-65504, 65504)
@@ -208,7 +251,9 @@ class CustomerAttnProcessor2_0:
def __init__(self):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
raise ImportError(
"AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
)
def apply_rotary_emb(
self,
@@ -258,23 +303,35 @@ class CustomerAttnProcessor2_0:
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
hidden_states = hidden_states.view(
batch_size, channel, height * width
).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
hidden_states.shape
if encoder_hidden_states is None
else encoder_hidden_states.shape
)
has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
has_encoder_hidden_state_proj = (
hasattr(attn, "add_q_proj")
and hasattr(attn, "add_k_proj")
and hasattr(attn, "add_v_proj")
)
if attn.group_norm is not None:
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
1, 2
)
query = attn.to_q(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
elif attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
encoder_hidden_states = attn.norm_encoder_hidden_states(
encoder_hidden_states
)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
@@ -300,19 +357,33 @@ class CustomerAttnProcessor2_0:
elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:
key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)
if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
if (
attn.is_cross_attention
and encoder_attention_mask is not None
and has_encoder_hidden_state_proj
):
# attention_mask: N x S1
# encoder_attention_mask: N x S2
# cross attention 整合attention_mask和encoder_attention_mask
combined_mask = attention_mask[:, :, None] * encoder_attention_mask[:, None, :]
combined_mask = (
attention_mask[:, :, None] * encoder_attention_mask[:, None, :]
)
attention_mask = torch.where(combined_mask == 1, 0.0, -torch.inf)
attention_mask = attention_mask[:, None, :, :].expand(-1, attn.heads, -1, -1).to(query.dtype)
attention_mask = (
attention_mask[:, None, :, :]
.expand(-1, attn.heads, -1, -1)
.to(query.dtype)
)
elif not attn.is_cross_attention and attention_mask is not None:
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
attention_mask = attn.prepare_attention_mask(
attention_mask, sequence_length, batch_size
)
# scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
attention_mask = attention_mask.view(
batch_size, attn.heads, -1, attention_mask.shape[-1]
)
# the output of sdp = (batch, num_heads, seq_len, head_dim)
# TODO: add support for attn.scale when we move to Torch 2.1
@@ -320,7 +391,9 @@ class CustomerAttnProcessor2_0:
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.transpose(1, 2).reshape(
batch_size, -1, attn.heads * head_dim
)
hidden_states = hidden_states.to(query.dtype)
# linear proj
@@ -329,7 +402,9 @@ class CustomerAttnProcessor2_0:
hidden_states = attn.to_out[1](hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
hidden_states = hidden_states.transpose(-1, -2).reshape(
batch_size, channel, height, width
)
if attn.residual_connection:
hidden_states = hidden_states + residual
music_dcae/__init__.py → acestep/models/lyrics_utils/__init__.py
View File
models/lyrics_utils/lyric_encoder.py → acestep/models/lyrics_utils/lyric_encoder.py
+116 -96
View File
@@ -3,16 +3,19 @@ import math
import torch
from torch import nn
class ConvolutionModule(nn.Module):
"""ConvolutionModule in Conformer model."""
def __init__(self,
def __init__(
self,
channels: int,
kernel_size: int = 15,
activation: nn.Module = nn.ReLU(),
norm: str = "batch_norm",
causal: bool = False,
bias: bool = True):
bias: bool = True,
):
"""Construct an ConvolutionModule object.
Args:
channels (int): The number of channels of conv layers.
@@ -51,7 +54,7 @@ class ConvolutionModule(nn.Module):
bias=bias,
)
assert norm in ['batch_norm', 'layer_norm']
assert norm in ["batch_norm", "layer_norm"]
if norm == "batch_norm":
self.use_layer_norm = False
self.norm = nn.BatchNorm1d(channels)
@@ -95,12 +98,12 @@ class ConvolutionModule(nn.Module):
if self.lorder > 0:
if cache.size(2) == 0: # cache_t == 0
x = nn.functional.pad(x, (self.lorder, 0), 'constant', 0.0)
x = nn.functional.pad(x, (self.lorder, 0), "constant", 0.0)
else:
assert cache.size(0) == x.size(0) # equal batch
assert cache.size(1) == x.size(1) # equal channel
x = torch.cat((cache, x), dim=2)
assert (x.size(2) > self.lorder)
assert x.size(2) > self.lorder
new_cache = x[:, :, -self.lorder :]
else:
# It's better we just return None if no cache is required,
@@ -126,6 +129,7 @@ class ConvolutionModule(nn.Module):
return x.transpose(1, 2), new_cache
class PositionwiseFeedForward(torch.nn.Module):
"""Positionwise feed forward layer.
@@ -163,6 +167,7 @@ class PositionwiseFeedForward(torch.nn.Module):
"""
return self.w_2(self.dropout(self.activation(self.w_1(xs))))
class Swish(torch.nn.Module):
"""Construct an Swish object."""
@@ -170,6 +175,7 @@ class Swish(torch.nn.Module):
"""Return Swish activation function."""
return x * torch.sigmoid(x)
class MultiHeadedAttention(nn.Module):
"""Multi-Head Attention layer.
@@ -180,11 +186,9 @@ class MultiHeadedAttention(nn.Module):
"""
def __init__(self,
n_head: int,
n_feat: int,
dropout_rate: float,
key_bias: bool = True):
def __init__(
self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True
):
"""Construct an MultiHeadedAttention object."""
super().__init__()
assert n_feat % n_head == 0
@@ -229,7 +233,7 @@ class MultiHeadedAttention(nn.Module):
self,
value: torch.Tensor,
scores: torch.Tensor,
mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
) -> torch.Tensor:
"""Compute attention context vector.
@@ -252,17 +256,18 @@ class MultiHeadedAttention(nn.Module):
mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2)
# For last chunk, time2 might be larger than scores.size(-1)
mask = mask[:, :, :, : scores.size(-1)] # (batch, 1, *, time2)
scores = scores.masked_fill(mask, -float('inf'))
scores = scores.masked_fill(mask, -float("inf"))
attn = torch.softmax(scores, dim=-1).masked_fill(
mask, 0.0) # (batch, head, time1, time2)
mask, 0.0
) # (batch, head, time1, time2)
else:
attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2)
p_attn = self.dropout(attn)
x = torch.matmul(p_attn, value) # (batch, head, time1, d_k)
x = (x.transpose(1, 2).contiguous().view(n_batch, -1,
self.h * self.d_k)
x = (
x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
) # (batch, time1, d_model)
return self.linear_out(x) # (batch, time1, d_model)
@@ -274,7 +279,7 @@ class MultiHeadedAttention(nn.Module):
value: torch.Tensor,
mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
pos_emb: torch.Tensor = torch.empty(0),
cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Compute scaled dot product attention.
@@ -308,9 +313,7 @@ class MultiHeadedAttention(nn.Module):
"""
q, k, v = self.forward_qkv(query, key, value)
if cache.size(0) > 0:
key_cache, value_cache = torch.split(cache,
cache.size(-1) // 2,
dim=-1)
key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
k = torch.cat([key_cache, k], dim=2)
v = torch.cat([value_cache, v], dim=2)
new_cache = torch.cat((k, v), dim=-1)
@@ -328,11 +331,9 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
dropout_rate (float): Dropout rate.
"""
def __init__(self,
n_head: int,
n_feat: int,
dropout_rate: float,
key_bias: bool = True):
def __init__(
self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True
):
"""Construct an RelPositionMultiHeadedAttention object."""
super().__init__(n_head, n_feat, dropout_rate, key_bias)
# linear transformation for positional encoding
@@ -355,14 +356,12 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
torch.Tensor: Output tensor.
"""
zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1),
device=x.device,
dtype=x.dtype)
zero_pad = torch.zeros(
(x.size()[0], x.size()[1], x.size()[2], 1), device=x.device, dtype=x.dtype
)
x_padded = torch.cat([zero_pad, x], dim=-1)
x_padded = x_padded.view(x.size()[0],
x.size()[1],
x.size(3) + 1, x.size(2))
x_padded = x_padded.view(x.size()[0], x.size()[1], x.size(3) + 1, x.size(2))
x = x_padded[:, :, 1:].view_as(x)[
:, :, :, : x.size(-1) // 2 + 1
] # only keep the positions from 0 to time2
@@ -375,7 +374,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
value: torch.Tensor,
mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
pos_emb: torch.Tensor = torch.empty(0),
cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Compute 'Scaled Dot Product Attention' with rel. positional encoding.
Args:
@@ -399,9 +398,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
q = q.transpose(1, 2) # (batch, time1, head, d_k)
if cache.size(0) > 0:
key_cache, value_cache = torch.split(cache,
cache.size(-1) // 2,
dim=-1)
key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
k = torch.cat([key_cache, k], dim=2)
v = torch.cat([value_cache, v], dim=2)
# NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
@@ -431,12 +428,12 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
matrix_bd = self.rel_shift(matrix_bd)
scores = (matrix_ac + matrix_bd) / math.sqrt(
self.d_k) # (batch, head, time1, time2)
self.d_k
) # (batch, head, time1, time2)
return self.forward_attention(v, scores, mask), new_cache
def subsequent_mask(
size: int,
device: torch.device = torch.device("cpu"),
@@ -510,14 +507,17 @@ def subsequent_chunk_mask(
ret[i, start:ending] = True
return ret
def add_optional_chunk_mask(xs: torch.Tensor,
def add_optional_chunk_mask(
xs: torch.Tensor,
masks: torch.Tensor,
use_dynamic_chunk: bool,
use_dynamic_left_chunk: bool,
decoding_chunk_size: int,
static_chunk_size: int,
num_decoding_left_chunks: int,
enable_full_context: bool = True):
enable_full_context: bool = True,
):
"""Apply optional mask for encoder.
Args:
@@ -565,18 +565,17 @@ def add_optional_chunk_mask(xs: torch.Tensor,
chunk_size = chunk_size % 25 + 1
if use_dynamic_left_chunk:
max_left_chunks = (max_len - 1) // chunk_size
num_left_chunks = torch.randint(0, max_left_chunks,
(1, )).item()
chunk_masks = subsequent_chunk_mask(xs.size(1), chunk_size,
num_left_chunks,
xs.device) # (L, L)
num_left_chunks = torch.randint(0, max_left_chunks, (1,)).item()
chunk_masks = subsequent_chunk_mask(
xs.size(1), chunk_size, num_left_chunks, xs.device
) # (L, L)
chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L)
chunk_masks = masks & chunk_masks # (B, L, L)
elif static_chunk_size > 0:
num_left_chunks = num_decoding_left_chunks
chunk_masks = subsequent_chunk_mask(xs.size(1), static_chunk_size,
num_left_chunks,
xs.device) # (L, L)
chunk_masks = subsequent_chunk_mask(
xs.size(1), static_chunk_size, num_left_chunks, xs.device
) # (L, L)
chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L)
chunk_masks = masks & chunk_masks # (B, L, L)
else:
@@ -630,7 +629,8 @@ class ConformerEncoderLayer(nn.Module):
if self.conv_module is not None:
self.norm_conv = nn.LayerNorm(size, eps=1e-5) # for the CNN module
self.norm_final = nn.LayerNorm(
size, eps=1e-5) # for the final output of the block
size, eps=1e-5
) # for the final output of the block
self.dropout = nn.Dropout(dropout_rate)
self.size = size
self.normalize_before = normalize_before
@@ -671,8 +671,7 @@ class ConformerEncoderLayer(nn.Module):
residual = x
if self.normalize_before:
x = self.norm_ff_macaron(x)
x = residual + self.ff_scale * self.dropout(
self.feed_forward_macaron(x))
x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x))
if not self.normalize_before:
x = self.norm_ff_macaron(x)
@@ -680,8 +679,7 @@ class ConformerEncoderLayer(nn.Module):
residual = x
if self.normalize_before:
x = self.norm_mha(x)
x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb,
att_cache)
x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, att_cache)
x = residual + self.dropout(x_att)
if not self.normalize_before:
x = self.norm_mha(x)
@@ -714,7 +712,6 @@ class ConformerEncoderLayer(nn.Module):
return x, mask, new_att_cache, new_cnn_cache
class EspnetRelPositionalEncoding(torch.nn.Module):
"""Relative positional encoding module (new implementation).
@@ -770,8 +767,9 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
pe = torch.cat([pe_positive, pe_negative], dim=1)
self.pe = pe.to(device=x.device, dtype=x.dtype)
def forward(self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0) \
-> Tuple[torch.Tensor, torch.Tensor]:
def forward(
self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Add positional encoding.
Args:
@@ -786,9 +784,9 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
pos_emb = self.position_encoding(size=x.size(1), offset=offset)
return self.dropout(x), self.dropout(pos_emb)
def position_encoding(self,
offset: Union[int, torch.Tensor],
size: int) -> torch.Tensor:
def position_encoding(
self, offset: Union[int, torch.Tensor], size: int
) -> torch.Tensor:
"""For getting encoding in a streaming fashion
Attention!!!!!
@@ -811,7 +809,6 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
return pos_emb
class LinearEmbed(torch.nn.Module):
"""Linear transform the input without subsampling
@@ -822,8 +819,9 @@ class LinearEmbed(torch.nn.Module):
"""
def __init__(self, idim: int, odim: int, dropout_rate: float,
pos_enc_class: torch.nn.Module):
def __init__(
self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module
):
"""Construct an linear object."""
super().__init__()
self.out = torch.nn.Sequential(
@@ -833,14 +831,13 @@ class LinearEmbed(torch.nn.Module):
)
self.pos_enc = pos_enc_class # rel_pos_espnet
def position_encoding(self, offset: Union[int, torch.Tensor],
size: int) -> torch.Tensor:
def position_encoding(
self, offset: Union[int, torch.Tensor], size: int
) -> torch.Tensor:
return self.pos_enc.position_encoding(offset, size)
def forward(
self,
x: torch.Tensor,
offset: Union[int, torch.Tensor] = 0
self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Input x.
@@ -894,15 +891,13 @@ def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
"""
batch_size = lengths.size(0)
max_len = max_len if max_len > 0 else lengths.max().item()
seq_range = torch.arange(0,
max_len,
dtype=torch.int64,
device=lengths.device)
seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device)
seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
seq_length_expand = lengths.unsqueeze(-1)
mask = seq_range_expand >= seq_length_expand
return mask
# https://github.com/FunAudioLLM/CosyVoice/blob/main/examples/magicdata-read/cosyvoice/conf/cosyvoice.yaml
class ConformerEncoder(torch.nn.Module):
"""Conformer encoder module."""
@@ -917,8 +912,8 @@ class ConformerEncoder(torch.nn.Module):
dropout_rate: float = 0.1,
positional_dropout_rate: float = 0.1,
attention_dropout_rate: float = 0.0,
input_layer: str = 'linear',
pos_enc_layer_type: str = 'rel_pos_espnet',
input_layer: str = "linear",
pos_enc_layer_type: str = "rel_pos_espnet",
normalize_before: bool = True,
static_chunk_size: int = 1, # 1: causal_mask; 0: full_mask
use_dynamic_chunk: bool = False,
@@ -953,8 +948,12 @@ class ConformerEncoder(torch.nn.Module):
"""
super().__init__()
self.output_size = output_size
self.embed = LinearEmbed(input_size, output_size, dropout_rate,
EspnetRelPositionalEncoding(output_size, positional_dropout_rate))
self.embed = LinearEmbed(
input_size,
output_size,
dropout_rate,
EspnetRelPositionalEncoding(output_size, positional_dropout_rate),
)
self.normalize_before = normalize_before
self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5)
self.gradient_checkpointing = gradient_checkpointing
@@ -980,40 +979,60 @@ class ConformerEncoder(torch.nn.Module):
activation,
)
# convolution module definition
convolution_layer_args = (output_size, cnn_module_kernel, activation,
cnn_module_norm, causal)
convolution_layer_args = (
output_size,
cnn_module_kernel,
activation,
cnn_module_norm,
causal,
)
self.encoders = torch.nn.ModuleList([
self.encoders = torch.nn.ModuleList(
[
ConformerEncoderLayer(
output_size,
RelPositionMultiHeadedAttention(
*encoder_selfattn_layer_args),
RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args),
PositionwiseFeedForward(*positionwise_layer_args),
PositionwiseFeedForward(
*positionwise_layer_args) if macaron_style else None,
ConvolutionModule(
*convolution_layer_args) if use_cnn_module else None,
(
PositionwiseFeedForward(*positionwise_layer_args)
if macaron_style
else None
),
(
ConvolutionModule(*convolution_layer_args)
if use_cnn_module
else None
),
dropout_rate,
normalize_before,
) for _ in range(num_blocks)
])
)
for _ in range(num_blocks)
]
)
def forward_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor,
def forward_layers(
self,
xs: torch.Tensor,
chunk_masks: torch.Tensor,
pos_emb: torch.Tensor,
mask_pad: torch.Tensor) -> torch.Tensor:
mask_pad: torch.Tensor,
) -> torch.Tensor:
for layer in self.encoders:
xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad)
return xs
@torch.jit.unused
def forward_layers_checkpointed(self, xs: torch.Tensor,
def forward_layers_checkpointed(
self,
xs: torch.Tensor,
chunk_masks: torch.Tensor,
pos_emb: torch.Tensor,
mask_pad: torch.Tensor) -> torch.Tensor:
mask_pad: torch.Tensor,
) -> torch.Tensor:
for layer in self.encoders:
xs, chunk_masks, _, _ = ckpt.checkpoint(layer.__call__, xs,
chunk_masks, pos_emb,
mask_pad)
xs, chunk_masks, _, _ = ckpt.checkpoint(
layer.__call__, xs, chunk_masks, pos_emb, mask_pad
)
return xs
def forward(
@@ -1050,15 +1069,17 @@ class ConformerEncoder(torch.nn.Module):
masks = pad_mask.to(torch.bool).unsqueeze(1) # (B, 1, T)
xs, pos_emb = self.embed(xs)
mask_pad = masks # (B, 1, T/subsample_rate)
chunk_masks = add_optional_chunk_mask(xs, masks,
chunk_masks = add_optional_chunk_mask(
xs,
masks,
self.use_dynamic_chunk,
self.use_dynamic_left_chunk,
decoding_chunk_size,
self.static_chunk_size,
num_decoding_left_chunks)
num_decoding_left_chunks,
)
if self.gradient_checkpointing and self.training:
xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb,
mask_pad)
xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb, mask_pad)
else:
xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad)
if self.normalize_before:
@@ -1067,4 +1088,3 @@ class ConformerEncoder(torch.nn.Module):
# return the masks before encoder layers, and the masks will be used
# for cross attention with decoder later
return xs, masks
models/lyrics_utils/lyric_normalizer.py → acestep/models/lyrics_utils/lyric_normalizer.py
+28 -25
View File
@@ -2,39 +2,42 @@ import re
from opencc import OpenCC
t2s_converter = OpenCC('t2s')
s2t_converter = OpenCC('s2t')
t2s_converter = OpenCC("t2s")
s2t_converter = OpenCC("s2t")
EMOJI_PATTERN = re.compile(
"["
"\U0001F600-\U0001F64F" # Emoticons
"]+", flags=re.UNICODE
"\U0001f600-\U0001f64f" # Emoticons
"]+",
flags=re.UNICODE,
)
# 创建一个翻译表，用于替换和移除字符
TRANSLATION_TABLE = str.maketrans({
'-': ' ', # 将 '-' 替换为空格
',': None,
'.': None,
'，': None,
'。': None,
'!': None,
'！': None,
'?': None,
'？': None,
'…': None,
';': None,
'；': None,
':': None,
'：': None,
'\u3000': ' ', # 将全角空格替换为空格
})
TRANSLATION_TABLE = str.maketrans(
{
"-": " ", # 将 '-' 替换为空格
",": None,
".": None,
"，": None,
"。": None,
"!": None,
"！": None,
"?": None,
"？": None,
"…": None,
";": None,
"；": None,
":": None,
"：": None,
"\u3000": " ", # 将全角空格替换为空格
}
)
# 替换括号中的内容，包括中括号和小括号
BACKSLASH_PATTERN = re.compile(r'\(.*?\)|\[.*?\]')
BACKSLASH_PATTERN = re.compile(r"\(.*?\)|\[.*?\]")
SPACE_PATTERN = re.compile('(?<!^)\s+(?!$)')
SPACE_PATTERN = re.compile("(?<!^)\s+(?!$)")
def normalize_text(text, language, strip=True):
@@ -45,10 +48,10 @@ def normalize_text(text, language, strip=True):
text = text.translate(TRANSLATION_TABLE)
# Step 2: 移除表情符号
text = EMOJI_PATTERN.sub('', text)
text = EMOJI_PATTERN.sub("", text)
# Step 3: 连续空白字符替换为单个空格，首位除外
text = SPACE_PATTERN.sub(' ', text)
text = SPACE_PATTERN.sub(" ", text)
# Step 4: 去除首尾空白字符（如果需要）
if strip:
models/lyrics_utils/lyric_tokenizer.py → acestep/models/lyrics_utils/lyric_tokenizer.py
+136 -32
View File
@@ -446,7 +446,9 @@ _ordinal_re = {
"it": re.compile(r"([0-9]+)(º|°|ª|o|a|i|e)"),
"pl": re.compile(r"([0-9]+)(º|ª|st|nd|rd|th)"),
"ar": re.compile(r"([0-9]+)(ون|ين|ث|ر|ى)"),
"cs": re.compile(r"([0-9]+)\.(?=\s|$)"), # In Czech, a dot is often used after the number to indicate ordinals.
"cs": re.compile(
r"([0-9]+)\.(?=\s|$)"
), # In Czech, a dot is often used after the number to indicate ordinals.
"ru": re.compile(r"([0-9]+)(-й|-я|-е|-ое|-ье|-го)"),
"nl": re.compile(r"([0-9]+)(de|ste|e)"),
"tr": re.compile(r"([0-9]+)(\.|inci|nci|uncu|üncü|\.)"),
@@ -486,7 +488,9 @@ def _expand_decimal_point(m, lang="en"):
def _expand_currency(m, lang="en", currency="USD"):
amount = float((re.sub(r"[^\d.]", "", m.group(0).replace(",", "."))))
full_amount = num2words(amount, to="currency", currency=currency, lang=lang if lang != "cs" else "cz")
full_amount = num2words(
amount, to="currency", currency=currency, lang=lang if lang != "cs" else "cz"
)
and_equivalents = {
"en": ", ",
@@ -530,13 +534,21 @@ def expand_numbers_multilingual(text, lang="en"):
else:
text = re.sub(_dot_number_re, _remove_dots, text)
try:
text = re.sub(_currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text)
text = re.sub(_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)
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(
_decimal_number_re, lambda m: _expand_decimal_point(m, lang), text
)
text = re.sub(_ordinal_re[lang], lambda m: _expand_ordinal(m, lang), text)
text = re.sub(_number_re, lambda m: _expand_number(m, lang), text)
return text
@@ -582,7 +594,15 @@ def basic_cleaners(text):
def chinese_transliterate(text):
return "".join(
[p[0] for p in pypinyin.pinyin(text, style=pypinyin.Style.TONE3, heteronym=False, neutral_tone_with_five=True)]
[
p[0]
for p in pypinyin.pinyin(
text,
style=pypinyin.Style.TONE3,
heteronym=False,
neutral_tone_with_five=True,
)
]
)
@@ -597,7 +617,9 @@ def korean_transliterate(text):
return r.translit(text)
DEFAULT_VOCAB_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "vocab.json")
DEFAULT_VOCAB_FILE = os.path.join(
os.path.dirname(os.path.realpath(__file__)), "vocab.json"
)
class VoiceBpeTokenizer:
@@ -639,7 +661,23 @@ class VoiceBpeTokenizer:
# )
def preprocess_text(self, txt, lang):
if lang in {"ar", "cs", "de", "en", "es", "fr", "hu", "it", "nl", "pl", "pt", "ru", "tr", "zh", "ko"}:
if lang in {
"ar",
"cs",
"de",
"en",
"es",
"fr",
"hu",
"it",
"nl",
"pl",
"pt",
"ru",
"tr",
"zh",
"ko",
}:
txt = multilingual_cleaners(txt, lang)
if lang == "zh":
txt = chinese_transliterate(txt)
@@ -672,11 +710,12 @@ class VoiceBpeTokenizer:
# txt = txt.replace("[UNK]", "")
return txt
# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3936
def batch_decode(
self,
sequences: Union[List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"],
sequences: Union[
List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"
],
skip_special_tokens: bool = False,
) -> List[str]:
"""
@@ -693,10 +732,7 @@ class VoiceBpeTokenizer:
Returns:
`List[str]`: The list of decoded sentences.
"""
return [
self.decode(seq)
for seq in sequences
]
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):
@@ -716,15 +752,27 @@ def test_expand_numbers_multilingual():
("This is a 1st test", "This is a first test", "en"),
("That will be $20 sir.", "That will be twenty dollars sir.", "en"),
("That will be 20€ sir.", "That will be twenty euro sir.", "en"),
("That will be 20.15€ sir.", "That will be twenty euro, fifteen cents sir.", "en"),
(
"That will be 20.15€ sir.",
"That will be twenty euro, fifteen cents sir.",
"en",
),
("That's 100,000.5.", "That's one hundred thousand point five.", "en"),
# French
("En 12,5 secondes.", "En douze virgule cinq secondes.", "fr"),
("Il y avait 50 soldats.", "Il y avait cinquante soldats.", "fr"),
("Ceci est un 1er test", "Ceci est un premier test", "fr"),
("Cela vous fera $20 monsieur.", "Cela vous fera vingt dollars monsieur.", "fr"),
(
"Cela vous fera $20 monsieur.",
"Cela vous fera vingt dollars monsieur.",
"fr",
),
("Cela vous fera 20€ monsieur.", "Cela vous fera vingt euros monsieur.", "fr"),
("Cela vous fera 20,15€ monsieur.", "Cela vous fera vingt euros et quinze centimes monsieur.", "fr"),
(
"Cela vous fera 20,15€ monsieur.",
"Cela vous fera vingt euros et quinze centimes monsieur.",
"fr",
),
("Ce sera 100.000,5.", "Ce sera cent mille virgule cinq.", "fr"),
# German
("In 12,5 Sekunden.", "In zwölf Komma fünf Sekunden.", "de"),
@@ -732,21 +780,33 @@ def test_expand_numbers_multilingual():
("Dies ist ein 1. Test", "Dies ist ein erste Test", "de"), # Issue with gender
("Das macht $20 Herr.", "Das macht zwanzig Dollar Herr.", "de"),
("Das macht 20€ Herr.", "Das macht zwanzig Euro Herr.", "de"),
("Das macht 20,15€ Herr.", "Das macht zwanzig Euro und fünfzehn Cent Herr.", "de"),
(
"Das macht 20,15€ Herr.",
"Das macht zwanzig Euro und fünfzehn Cent Herr.",
"de",
),
# Spanish
("En 12,5 segundos.", "En doce punto cinco segundos.", "es"),
("Había 50 soldados.", "Había cincuenta soldados.", "es"),
("Este es un 1er test", "Este es un primero test", "es"),
("Eso le costará $20 señor.", "Eso le costará veinte dólares señor.", "es"),
("Eso le costará 20€ señor.", "Eso le costará veinte euros señor.", "es"),
("Eso le costará 20,15€ señor.", "Eso le costará veinte euros con quince céntimos señor.", "es"),
(
"Eso le costará 20,15€ señor.",
"Eso le costará veinte euros con quince céntimos señor.",
"es",
),
# Italian
("In 12,5 secondi.", "In dodici virgola cinque secondi.", "it"),
("C'erano 50 soldati.", "C'erano cinquanta soldati.", "it"),
("Questo è un 1° test", "Questo è un primo test", "it"),
("Ti costerà $20 signore.", "Ti costerà venti dollari signore.", "it"),
("Ti costerà 20€ signore.", "Ti costerà venti euro signore.", "it"),
("Ti costerà 20,15€ signore.", "Ti costerà venti euro e quindici centesimi signore.", "it"),
(
"Ti costerà 20,15€ signore.",
"Ti costerà venti euro e quindici centesimi signore.",
"it",
),
# Portuguese
("Em 12,5 segundos.", "Em doze vírgula cinco segundos.", "pt"),
("Havia 50 soldados.", "Havia cinquenta soldados.", "pt"),
@@ -761,8 +821,16 @@ def test_expand_numbers_multilingual():
# Polish
("W 12,5 sekundy.", "W dwanaście przecinek pięć sekundy.", "pl"),
("Było 50 żołnierzy.", "Było pięćdziesiąt żołnierzy.", "pl"),
("To będzie kosztować 20€ panie.", "To będzie kosztować dwadzieścia euro panie.", "pl"),
("To będzie kosztować 20,15€ panie.", "To będzie kosztować dwadzieścia euro, piętnaście centów panie.", "pl"),
(
"To będzie kosztować 20€ panie.",
"To będzie kosztować dwadzieścia euro panie.",
"pl",
),
(
"To będzie kosztować 20,15€ panie.",
"To będzie kosztować dwadzieścia euro, piętnaście centów panie.",
"pl",
),
# Arabic
("في الـ 12,5 ثانية.", "في الـ اثنا عشر , خمسون ثانية.", "ar"),
("كان هناك 50 جنديًا.", "كان هناك خمسون جنديًا.", "ar"),
@@ -776,8 +844,16 @@ def test_expand_numbers_multilingual():
# Russian
("Через 12.5 секунды.", "Через двенадцать запятая пять секунды.", "ru"),
("Там было 50 солдат.", "Там было пятьдесят солдат.", "ru"),
("Это будет 20.15€ сэр.", "Это будет двадцать евро, пятнадцать центов сэр.", "ru"),
("Это будет стоить 20€ господин.", "Это будет стоить двадцать евро господин.", "ru"),
(
"Это будет 20.15€ сэр.",
"Это будет двадцать евро, пятнадцать центов сэр.",
"ru",
),
(
"Это будет стоить 20€ господин.",
"Это будет стоить двадцать евро господин.",
"ru",
),
# Dutch
("In 12,5 seconden.", "In twaalf komma vijf seconden.", "nl"),
("Er waren 50 soldaten.", "Er waren vijftig soldaten.", "nl"),
@@ -817,18 +893,30 @@ def test_abbreviations_multilingual():
("La Dra. Martinez es muy buena.", "La doctora Martinez es muy buena.", "es"),
# French
("Bonjour Mr. Dupond.", "Bonjour monsieur Dupond.", "fr"),
("Mme. Moreau est absente aujourd'hui.", "madame Moreau est absente aujourd'hui.", "fr"),
(
"Mme. Moreau est absente aujourd'hui.",
"madame Moreau est absente aujourd'hui.",
"fr",
),
# German
("Frau Dr. Müller ist sehr klug.", "Frau doktor Müller ist sehr klug.", "de"),
# Portuguese
("Olá Sr. Silva.", "Olá senhor Silva.", "pt"),
("Dra. Costa, você está disponível?", "doutora Costa, você está disponível?", "pt"),
(
"Dra. Costa, você está disponível?",
"doutora Costa, você está disponível?",
"pt",
),
# Italian
("Buongiorno, Sig. Rossi.", "Buongiorno, signore Rossi.", "it"),
# ("Sig.ra Bianchi, posso aiutarti?", 'signora Bianchi, posso aiutarti?', 'it'), # Issue with matching that pattern
# Polish
("Dzień dobry, P. Kowalski.", "Dzień dobry, pani Kowalski.", "pl"),
("M. Nowak, czy mogę zadać pytanie?", "pan Nowak, czy mogę zadać pytanie?", "pl"),
(
"M. Nowak, czy mogę zadać pytanie?",
"pan Nowak, czy mogę zadać pytanie?",
"pl",
),
# Czech
("P. Novák", "pan Novák", "cs"),
("Dr. Vojtěch", "doktor Vojtěch", "cs"),
@@ -837,7 +925,11 @@ def test_abbreviations_multilingual():
("Mevr. de Vries", "mevrouw de Vries", "nl"),
# Russian
("Здравствуйте Г-н Иванов.", "Здравствуйте господин Иванов.", "ru"),
("Д-р Смирнов здесь, чтобы увидеть вас.", "доктор Смирнов здесь, чтобы увидеть вас.", "ru"),
(
"Д-р Смирнов здесь, чтобы увидеть вас.",
"доктор Смирнов здесь, чтобы увидеть вас.",
"ru",
),
# Turkish
("Merhaba B. Yılmaz.", "Merhaba bay Yılmaz.", "tr"),
("Dr. Ayşe burada.", "doktor Ayşe burada.", "tr"),
@@ -856,8 +948,16 @@ def test_symbols_multilingual():
("Te veo @ la fiesta", "Te veo arroba la fiesta", "es"),
("J'ai 14° de fièvre", "J'ai 14 degrés de fièvre", "fr"),
("Die Rechnung beträgt £ 20", "Die Rechnung beträgt pfund 20", "de"),
("O meu email é ana&joao@gmail.com", "O meu email é ana e joao arroba gmail.com", "pt"),
("linguaggio di programmazione C#", "linguaggio di programmazione C cancelletto", "it"),
(
"O meu email é ana&joao@gmail.com",
"O meu email é ana e joao arroba gmail.com",
"pt",
),
(
"linguaggio di programmazione C#",
"linguaggio di programmazione C cancelletto",
"it",
),
("Moja temperatura to 36.6°", "Moja temperatura to 36.6 stopnie", "pl"),
("Mám 14% baterie", "Mám 14 procento baterie", "cs"),
("Těším se na tebe @ party", "Těším se na tebe na party", "cs"),
@@ -868,7 +968,11 @@ def test_symbols_multilingual():
("لدي 14% في البطارية", "لدي 14 في المئة في البطارية", "ar"),
("我的电量为 14%", "我的电量为 14 百分之", "zh"),
("Pilim %14 dolu.", "Pilim yüzde 14 dolu.", "tr"),
("Az akkumulátorom töltöttsége 14%", "Az akkumulátorom töltöttsége 14 százalék", "hu"),
(
"Az akkumulátorom töltöttsége 14%",
"Az akkumulátorom töltöttsége 14 százalék",
"hu",
),
("배터리 잔량이 14%입니다.", "배터리 잔량이 14 퍼센트입니다.", "ko"),
]
models/lyrics_utils/vocab.json → acestep/models/lyrics_utils/vocab.json
View File
models/lyrics_utils/zh_num2words.py → acestep/models/lyrics_utils/zh_num2words.py
View File
acestep/music_dcae/__init__.py
View File
music_dcae/music_dcae_pipeline.py → acestep/music_dcae/music_dcae_pipeline.py
+37 -10
View File
@@ -1,3 +1,11 @@
"""
ACE-Step: A Step Towards Music Generation Foundation Model
https://github.com/ace-step/ACE-Step
Apache 2.0 License
"""
import os
import torch
from diffusers import AutoencoderDC
@@ -21,7 +29,12 @@ VOCODER_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_vocoder")
class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
@register_to_config
def __init__(self, source_sample_rate=None, dcae_checkpoint_path=DEFAULT_PRETRAINED_PATH, vocoder_checkpoint_path=VOCODER_PRETRAINED_PATH):
def __init__(
self,
source_sample_rate=None,
dcae_checkpoint_path=DEFAULT_PRETRAINED_PATH,
vocoder_checkpoint_path=VOCODER_PRETRAINED_PATH,
):
super(MusicDCAE, self).__init__()
self.dcae = AutoencoderDC.from_pretrained(dcae_checkpoint_path)
@@ -32,9 +45,11 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
self.transform = transforms.Compose([
self.transform = transforms.Compose(
[
transforms.Normalize(0.5, 0.5),
])
]
)
self.min_mel_value = -11.0
self.max_mel_value = 3.0
self.audio_chunk_size = int(round((1024 * 512 / 44100 * 48000)))
@@ -76,7 +91,9 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
max_audio_len = audio.shape[-1]
if max_audio_len % (8 * 512) != 0:
audio = torch.nn.functional.pad(audio, (0, 8 * 512 - max_audio_len % (8 * 512)))
audio = torch.nn.functional.pad(
audio, (0, 8 * 512 - max_audio_len % (8 * 512))
)
mels = self.forward_mel(audio)
mels = (mels - self.min_mel_value) / (self.max_mel_value - self.min_mel_value)
@@ -86,7 +103,9 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
latent = self.dcae.encoder(mel.unsqueeze(0))
latents.append(latent)
latents = torch.cat(latents, dim=0)
latent_lengths = (audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple).long()
latent_lengths = (
audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple
).long()
latents = (latents - self.shift_factor) * self.scale_factor
return latents, latent_lengths
@@ -103,18 +122,26 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
wav = self.vocoder.decode(mels[0]).squeeze(1)
if sr is not None:
resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
resampler = (
torchaudio.transforms.Resample(44100, sr)
.to(latents.device)
.to(latents.dtype)
)
wav = resampler(wav)
else:
sr = 44100
pred_wavs.append(wav)
if audio_lengths is not None:
pred_wavs = [wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)]
pred_wavs = [
wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)
]
return sr, pred_wavs
def forward(self, audios, audio_lengths=None, sr=None):
latents, latent_lengths = self.encode(audios=audios, audio_lengths=audio_lengths, sr=sr)
latents, latent_lengths = self.encode(
audios=audios, audio_lengths=audio_lengths, sr=sr
)
sr, pred_wavs = self.decode(latents=latents, audio_lengths=audio_lengths, sr=sr)
return sr, pred_wavs, latents, latent_lengths
@@ -137,5 +164,5 @@ if __name__ == "__main__":
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")
torchaudio.save("test_reconstructed.wav", pred_wavs[0], sr)
print("test_reconstructed.wav")
music_dcae/music_log_mel.py → acestep/music_dcae/music_log_mel.py
+8
View File
@@ -1,3 +1,11 @@
"""
ACE-Step: A Step Towards Music Generation Foundation Model
https://github.com/ace-step/ACE-Step
Apache 2.0 License
"""
import torch
import torch.nn as nn
from torch import Tensor
music_dcae/music_vocoder.py → acestep/music_dcae/music_vocoder.py
+21 -10
View File
@@ -1,3 +1,11 @@
"""
ACE-Step: A Step Towards Music Generation Foundation Model
https://github.com/ace-step/ACE-Step
Apache 2.0 License
"""
import librosa
import torch
from torch import nn
@@ -132,13 +140,11 @@ class ConvNeXtBlock(nn.Module):
self.act = nn.GELU()
self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)
self.gamma = (
nn.Parameter(layer_scale_init_value *
torch.ones((dim)), requires_grad=True)
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()
self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
def forward(self, x, apply_residual: bool = True):
input = x
@@ -367,8 +373,7 @@ class HiFiGANGenerator(nn.Module):
upsample_rates: Tuple[int] = (8, 8, 2, 2, 2),
upsample_kernel_sizes: Tuple[int] = (16, 16, 8, 2, 2),
resblock_kernel_sizes: Tuple[int] = (3, 7, 11),
resblock_dilation_sizes: Tuple[Tuple[int]] = (
(1, 3, 5), (1, 3, 5), (1, 3, 5)),
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,
@@ -498,7 +503,11 @@ class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
upsample_kernel_sizes: Tuple[int] = (8, 8, 4, 4, 4, 4, 4),
resblock_kernel_sizes: Tuple[int] = (3, 7, 11, 13),
resblock_dilation_sizes: Tuple[Tuple[int]] = (
(1, 3, 5), (1, 3, 5), (1, 3, 5), (1, 3, 5)),
(1, 3, 5),
(1, 3, 5),
(1, 3, 5),
(1, 3, 5),
),
num_mels: int = 512,
upsample_initial_channel: int = 1024,
use_template: bool = False,
@@ -565,12 +574,14 @@ class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
if __name__ == "__main__":
import soundfile as sf
x = "test_audio.flac"
model = ADaMoSHiFiGANV1.from_pretrained("./checkpoints/music_vocoder", local_files_only=True)
x = "test_audio.wav"
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)
sf.write("test_audio_vocoder_rec.wav", wav.cpu().numpy(), 44100)
pipeline_ace_step.py → acestep/pipeline_ace_step.py
+621 -152
View File
File diff suppressed because it is too large Load Diff
acestep/schedulers/__init__.py
View File
schedulers/scheduling_flow_match_euler_discrete.py → acestep/schedulers/scheduling_flow_match_euler_discrete.py
+13 -5
View File
@@ -72,7 +72,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
base_image_seq_len: Optional[int] = 256,
max_image_seq_len: Optional[int] = 4096,
):
timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
timesteps = np.linspace(
1, num_train_timesteps, num_train_timesteps, dtype=np.float32
)[::-1].copy()
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
sigmas = timesteps / num_train_timesteps
@@ -146,7 +148,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
# self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
if self.begin_index is None:
step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
step_indices = [
self.index_for_timestep(t, schedule_timesteps) for t in timestep
]
elif self.step_index is not None:
# add_noise is called after first denoising step (for inpainting)
step_indices = [self.step_index] * timestep.shape[0]
@@ -186,12 +190,16 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
"""
if self.config.use_dynamic_shifting and mu is None:
raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`")
raise ValueError(
" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
)
if sigmas is None:
self.num_inference_steps = num_inference_steps
timesteps = np.linspace(
self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
self._sigma_to_t(self.sigma_max),
self._sigma_to_t(self.sigma_min),
num_inference_steps,
)
sigmas = timesteps / self.config.num_train_timesteps
@@ -243,7 +251,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
s_noise: float = 1.0,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
omega: Union[float, np.array] = 0.0
omega: Union[float, np.array] = 0.0,
) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
schedulers/scheduling_flow_match_heun_discrete.py → acestep/schedulers/scheduling_flow_match_heun_discrete.py
+22 -7
View File
@@ -67,7 +67,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
num_train_timesteps: int = 1000,
shift: float = 1.0,
):
timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
timesteps = np.linspace(
1, num_train_timesteps, num_train_timesteps, dtype=np.float32
)[::-1].copy()
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
sigmas = timesteps / num_train_timesteps
@@ -137,7 +139,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
def _sigma_to_t(self, sigma):
return sigma * self.config.num_train_timesteps
def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
def set_timesteps(
self, num_inference_steps: int, device: Union[str, torch.device] = None
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -150,7 +154,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
self.num_inference_steps = num_inference_steps
timesteps = np.linspace(
self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
self._sigma_to_t(self.sigma_max),
self._sigma_to_t(self.sigma_min),
num_inference_steps,
)
sigmas = timesteps / self.config.num_train_timesteps
@@ -162,7 +168,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
self.timesteps = timesteps.to(device=device)
sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
self.sigmas = torch.cat(
[sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]]
)
# empty dt and derivative
self.prev_derivative = None
@@ -208,7 +216,7 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
s_noise: float = 1.0,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
omega: Union[float, np.array] = 0.0
omega: Union[float, np.array] = 0.0,
) -> Union[FlowMatchHeunDiscreteSchedulerOutput, Tuple]:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
@@ -285,13 +293,20 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
sigma = self.sigmas[self.step_index - 1]
sigma_next = self.sigmas[self.step_index]
gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0
gamma = (
min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1)
if s_tmin <= sigma <= s_tmax
else 0.0
)
sigma_hat = sigma * (gamma + 1)
if gamma > 0:
noise = randn_tensor(
model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
model_output.shape,
dtype=model_output.dtype,
device=model_output.device,
generator=generator,
)
eps = noise * s_noise
sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
text2music_dataset.py → acestep/text2music_dataset.py
+170 -30
View File
@@ -17,6 +17,7 @@ warnings.simplefilter("ignore", category=FutureWarning)
DEFAULT_TRAIN_PATH = "./data/example_dataset"
def is_silent_audio(audio_tensor, silence_threshold=0.95):
"""
Determine if an audio is silent and should be discarded
@@ -39,10 +40,23 @@ def is_silent_audio(audio_tensor, silence_threshold=0.95):
# Supported languages for tokenization
SUPPORT_LANGUAGES = {
"en": 259, "de": 260, "fr": 262, "es": 284, "it": 285,
"pt": 286, "pl": 294, "tr": 295, "ru": 267, "cs": 293,
"nl": 297, "ar": 5022, "zh": 5023, "ja": 5412, "hu": 5753,
"ko": 6152, "hi": 6680
"en": 259,
"de": 260,
"fr": 262,
"es": 284,
"it": 285,
"pt": 286,
"pl": 294,
"tr": 295,
"ru": 267,
"cs": 293,
"nl": 297,
"ar": 5022,
"zh": 5023,
"ja": 5412,
"hu": 5753,
"ko": 6152,
"hi": 6680,
}
# Regex pattern for structure markers like [Verse], [Chorus], etc.
@@ -54,9 +68,15 @@ class Text2MusicDataset(Dataset):
Dataset for text-to-music generation that processes lyrics and audio files
"""
def __init__(self, train=True, train_dataset_path=DEFAULT_TRAIN_PATH,
max_duration=240.0, sample_size=None, shuffle=True,
minibatch_size=1):
def __init__(
self,
train=True,
train_dataset_path=DEFAULT_TRAIN_PATH,
max_duration=240.0,
sample_size=None,
shuffle=True,
minibatch_size=1,
):
"""
Initialize the Text2Music dataset
@@ -75,14 +95,107 @@ class Text2MusicDataset(Dataset):
# Initialize language segmentation
self.lang_segment = LangSegment()
self.lang_segment.setfilters([
'af', 'am', 'an', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'br', 'bs', 'ca', 'cs', 'cy', 'da', 'de', 'dz', 'el',
'en', 'eo', 'es', 'et', 'eu', 'fa', 'fi', 'fo', 'fr', 'ga', 'gl', 'gu', 'he', 'hi', 'hr', 'ht', 'hu', 'hy',
'id', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'ku', 'ky', 'la', 'lb', 'lo', 'lt', 'lv', 'mg',
'mk', 'ml', 'mn', 'mr', 'ms', 'mt', 'nb', 'ne', 'nl', 'nn', 'no', 'oc', 'or', 'pa', 'pl', 'ps', 'pt', 'qu',
'ro', 'ru', 'rw', 'se', 'si', 'sk', 'sl', 'sq', 'sr', 'sv', 'sw', 'ta', 'te', 'th', 'tl', 'tr', 'ug', 'uk',
'ur', 'vi', 'vo', 'wa', 'xh', 'zh', 'zu'
])
self.lang_segment.setfilters(
[
"af",
"am",
"an",
"ar",
"as",
"az",
"be",
"bg",
"bn",
"br",
"bs",
"ca",
"cs",
"cy",
"da",
"de",
"dz",
"el",
"en",
"eo",
"es",
"et",
"eu",
"fa",
"fi",
"fo",
"fr",
"ga",
"gl",
"gu",
"he",
"hi",
"hr",
"ht",
"hu",
"hy",
"id",
"is",
"it",
"ja",
"jv",
"ka",
"kk",
"km",
"kn",
"ko",
"ku",
"ky",
"la",
"lb",
"lo",
"lt",
"lv",
"mg",
"mk",
"ml",
"mn",
"mr",
"ms",
"mt",
"nb",
"ne",
"nl",
"nn",
"no",
"oc",
"or",
"pa",
"pl",
"ps",
"pt",
"qu",
"ro",
"ru",
"rw",
"se",
"si",
"sk",
"sl",
"sq",
"sr",
"sv",
"sw",
"ta",
"te",
"th",
"tl",
"tr",
"ug",
"uk",
"ur",
"vi",
"vo",
"wa",
"xh",
"zh",
"zu",
]
)
# Initialize lyric tokenizer
self.lyric_tokenizer = VoiceBpeTokenizer()
@@ -196,22 +309,30 @@ class Text2MusicDataset(Dataset):
# If debug mode, show tokenization results
if debug:
toks = self.lyric_tokenizer.batch_decode([[tok_id] for tok_id in token_idx])
logger.info(f"debug using most_common_lang {line} --> {most_common_lang} --> {toks}")
toks = self.lyric_tokenizer.batch_decode(
[[tok_id] for tok_id in token_idx]
)
logger.info(
f"debug using most_common_lang {line} --> {most_common_lang} --> {toks}"
)
# If tokenization contains unknown token (1), try with segment language
if 1 in token_idx:
token_idx = self.lyric_tokenizer.encode(line, lang)
if debug:
toks = self.lyric_tokenizer.batch_decode([[tok_id] for tok_id in token_idx])
toks = self.lyric_tokenizer.batch_decode(
[[tok_id] for tok_id in token_idx]
)
logger.info(f"debug {line} --> {lang} --> {toks}")
# Add tokens and line break
lyric_token_idx = lyric_token_idx + token_idx + [2]
except Exception as e:
logger.error(f"Tokenize error: {e} for line: {line}, major_language: {lang}")
logger.error(
f"Tokenize error: {e} for line: {line}, major_language: {lang}"
)
return lyric_token_idx
@@ -302,7 +423,9 @@ class Text2MusicDataset(Dataset):
# Pad to minimum 3 seconds if needed
if audio.shape[-1] < 48000 * 3:
audio = torch.nn.functional.pad(audio, (0, 48000 * 3 - audio.shape[-1]), 'constant', 0)
audio = torch.nn.functional.pad(
audio, (0, 48000 * 3 - audio.shape[-1]), "constant", 0
)
# Check if audio is silent
if is_silent_audio(audio):
@@ -368,9 +491,11 @@ class Text2MusicDataset(Dataset):
lyrics = item["norm_lyrics"]
lyrics_lines = lyrics.split("\n")
for lyric_line in lyrics_lines:
candidate_lyric_chunk.append({
candidate_lyric_chunk.append(
{
"lyric": lyric_line,
})
}
)
# Limit audio length
longest_length = 24 * 10 * 48000 # 240 seconds
@@ -482,28 +607,43 @@ class Text2MusicDataset(Dataset):
elif k in ["src_wavs", "target_wavs", "vocal_wavs"]:
# Pad audio to max length
max_length = max(seq.shape[1] for seq in v)
padded_input_list = torch.stack([
torch.nn.functional.pad(seq, (0, max_length - seq.shape[1]), 'constant', 0)
padded_input_list = torch.stack(
[
torch.nn.functional.pad(
seq, (0, max_length - seq.shape[1]), "constant", 0
)
for seq in v
])
]
)
elif k in ["clap_conditions"]:
# Pad time dimension of embeddings
max_length = max(seq.shape[0] for seq in v)
v = [
torch.nn.functional.pad(seq, (0, 0, 0, max_length - seq.shape[0]), 'constant', 0)
torch.nn.functional.pad(
seq, (0, 0, 0, max_length - seq.shape[0]), "constant", 0
)
for seq in v
]
padded_input_list = torch.stack(v)
elif k == "speaker_embs":
# Stack speaker embeddings
padded_input_list = torch.stack(v)
elif k in ["chunk_masks", "clap_attention_masks", "lyric_token_ids", "lyric_masks"]:
elif k in [
"chunk_masks",
"clap_attention_masks",
"lyric_token_ids",
"lyric_masks",
]:
# Pad sequence tensors
max_length = max(len(seq) for seq in v)
padded_input_list = torch.stack([
torch.nn.functional.pad(seq, (0, max_length - len(seq)), 'constant', 0)
padded_input_list = torch.stack(
[
torch.nn.functional.pad(
seq, (0, max_length - len(seq)), "constant", 0
)
for seq in v
])
]
)
output[k] = padded_input_list
ui/components.py → acestep/ui/components.py
+289 -58
View File
@@ -1,3 +1,11 @@
"""
ACE-Step: A Step Towards Music Generation Foundation Model
https://github.com/ace-step/ACE-Step
Apache 2.0 License
"""
import gradio as gr
import librosa
@@ -41,8 +49,6 @@ In this moment we take flight
"""
def create_output_ui(task_name="Text2Music"):
# For many consumer-grade GPU devices, only one batch can be run
output_audio1 = gr.Audio(type="filepath", label=f"{task_name} Generated Audio 1")
@@ -68,41 +74,162 @@ def create_text2music_ui(
with gr.Column():
with gr.Row(equal_height=True):
# add markdown, tags and lyrics examples are from ai music generation community
audio_duration = gr.Slider(-1, 240.0, step=0.00001, value=-1, label="Audio Duration", interactive=True, info="-1 means random duration (30 ~ 240).", scale=9)
audio_duration = gr.Slider(
-1,
240.0,
step=0.00001,
value=-1,
label="Audio Duration",
interactive=True,
info="-1 means random duration (30 ~ 240).",
scale=9,
)
sample_bnt = gr.Button("Sample", variant="primary", scale=1)
prompt = gr.Textbox(lines=2, label="Tags", max_lines=4, value=TAG_DEFAULT, info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community")
lyrics = gr.Textbox(lines=9, label="Lyrics", max_lines=13, value=LYRIC_DEFAULT, info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics")
prompt = gr.Textbox(
lines=2,
label="Tags",
max_lines=4,
value=TAG_DEFAULT,
info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community",
)
lyrics = gr.Textbox(
lines=9,
label="Lyrics",
max_lines=13,
value=LYRIC_DEFAULT,
info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics",
)
with gr.Accordion("Basic Settings", open=False):
infer_step = gr.Slider(minimum=1, maximum=1000, step=1, value=27, label="Infer Steps", interactive=True)
guidance_scale = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=15.0, label="Guidance Scale", interactive=True, info="When guidance_scale_lyric > 1 and guidance_scale_text > 1, the guidance scale will not be applied.")
guidance_scale_text = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Text", interactive=True, info="Guidance scale for text condition. It can only apply to cfg. set guidance_scale_text=5.0, guidance_scale_lyric=1.5 for start")
guidance_scale_lyric = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Lyric", interactive=True)
infer_step = gr.Slider(
minimum=1,
maximum=1000,
step=1,
value=27,
label="Infer Steps",
interactive=True,
)
guidance_scale = gr.Slider(
minimum=0.0,
maximum=200.0,
step=0.1,
value=15.0,
label="Guidance Scale",
interactive=True,
info="When guidance_scale_lyric > 1 and guidance_scale_text > 1, the guidance scale will not be applied.",
)
guidance_scale_text = gr.Slider(
minimum=0.0,
maximum=10.0,
step=0.1,
value=0.0,
label="Guidance Scale Text",
interactive=True,
info="Guidance scale for text condition. It can only apply to cfg. set guidance_scale_text=5.0, guidance_scale_lyric=1.5 for start",
)
guidance_scale_lyric = gr.Slider(
minimum=0.0,
maximum=10.0,
step=0.1,
value=0.0,
label="Guidance Scale Lyric",
interactive=True,
)
manual_seeds = gr.Textbox(label="manual seeds (default None)", placeholder="1,2,3,4", value=None, info="Seed for the generation")
manual_seeds = gr.Textbox(
label="manual seeds (default None)",
placeholder="1,2,3,4",
value=None,
info="Seed for the generation",
)
with gr.Accordion("Advanced Settings", open=False):
scheduler_type = gr.Radio(["euler", "heun"], value="euler", label="Scheduler Type", elem_id="scheduler_type", info="Scheduler type for the generation. euler is recommended. heun will take more time.")
cfg_type = gr.Radio(["cfg", "apg", "cfg_star"], value="apg", label="CFG Type", elem_id="cfg_type", info="CFG type for the generation. apg is recommended. cfg and cfg_star are almost the same.")
use_erg_tag = gr.Checkbox(label="use ERG for tag", value=True, info="Use Entropy Rectifying Guidance for tag. It will multiple a temperature to the attention to make a weaker tag condition and make better diversity.")
use_erg_lyric = gr.Checkbox(label="use ERG for lyric", value=True, info="The same but apply to lyric encoder's attention.")
use_erg_diffusion = gr.Checkbox(label="use ERG for diffusion", value=True, info="The same but apply to diffusion model's attention.")
scheduler_type = gr.Radio(
["euler", "heun"],
value="euler",
label="Scheduler Type",
elem_id="scheduler_type",
info="Scheduler type for the generation. euler is recommended. heun will take more time.",
)
cfg_type = gr.Radio(
["cfg", "apg", "cfg_star"],
value="apg",
label="CFG Type",
elem_id="cfg_type",
info="CFG type for the generation. apg is recommended. cfg and cfg_star are almost the same.",
)
use_erg_tag = gr.Checkbox(
label="use ERG for tag",
value=True,
info="Use Entropy Rectifying Guidance for tag. It will multiple a temperature to the attention to make a weaker tag condition and make better diversity.",
)
use_erg_lyric = gr.Checkbox(
label="use ERG for lyric",
value=True,
info="The same but apply to lyric encoder's attention.",
)
use_erg_diffusion = gr.Checkbox(
label="use ERG for diffusion",
value=True,
info="The same but apply to diffusion model's attention.",
)
omega_scale = gr.Slider(minimum=-100.0, maximum=100.0, step=0.1, value=10.0, label="Granularity Scale", interactive=True, info="Granularity scale for the generation. Higher values can reduce artifacts")
omega_scale = gr.Slider(
minimum=-100.0,
maximum=100.0,
step=0.1,
value=10.0,
label="Granularity Scale",
interactive=True,
info="Granularity scale for the generation. Higher values can reduce artifacts",
)
guidance_interval = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Guidance Interval", interactive=True, info="Guidance interval for the generation. 0.5 means only apply guidance in the middle steps (0.25 * infer_steps to 0.75 * infer_steps)")
guidance_interval_decay = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.0, label="Guidance Interval Decay", interactive=True, info="Guidance interval decay for the generation. Guidance scale will decay from guidance_scale to min_guidance_scale in the interval. 0.0 means no decay.")
min_guidance_scale = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=3.0, label="Min Guidance Scale", interactive=True, info="Min guidance scale for guidance interval decay's end scale")
oss_steps = gr.Textbox(label="OSS Steps", placeholder="16, 29, 52, 96, 129, 158, 172, 183, 189, 200", value=None, info="Optimal Steps for the generation. But not test well")
guidance_interval = gr.Slider(
minimum=0.0,
maximum=1.0,
step=0.01,
value=0.5,
label="Guidance Interval",
interactive=True,
info="Guidance interval for the generation. 0.5 means only apply guidance in the middle steps (0.25 * infer_steps to 0.75 * infer_steps)",
)
guidance_interval_decay = gr.Slider(
minimum=0.0,
maximum=1.0,
step=0.01,
value=0.0,
label="Guidance Interval Decay",
interactive=True,
info="Guidance interval decay for the generation. Guidance scale will decay from guidance_scale to min_guidance_scale in the interval. 0.0 means no decay.",
)
min_guidance_scale = gr.Slider(
minimum=0.0,
maximum=200.0,
step=0.1,
value=3.0,
label="Min Guidance Scale",
interactive=True,
info="Min guidance scale for guidance interval decay's end scale",
)
oss_steps = gr.Textbox(
label="OSS Steps",
placeholder="16, 29, 52, 96, 129, 158, 172, 183, 189, 200",
value=None,
info="Optimal Steps for the generation. But not test well",
)
text2music_bnt = gr.Button("Generate", variant="primary")
with gr.Column():
outputs, input_params_json = create_output_ui()
with gr.Tab("retake"):
retake_variance = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance")
retake_seeds = gr.Textbox(label="retake seeds (default None)", placeholder="", value=None)
retake_variance = gr.Slider(
minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance"
)
retake_seeds = gr.Textbox(
label="retake seeds (default None)", placeholder="", value=None
)
retake_bnt = gr.Button("Retake", variant="primary")
retake_outputs, retake_input_params_json = create_output_ui("Retake")
@@ -124,8 +251,16 @@ def create_text2music_ui(
json_data["use_erg_lyric"],
json_data["use_erg_diffusion"],
", ".join(map(str, json_data["oss_steps"])),
json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
(
json_data["guidance_scale_text"]
if "guidance_scale_text" in json_data
else 0.0
),
(
json_data["guidance_scale_lyric"]
if "guidance_scale_lyric" in json_data
else 0.0
),
retake_seeds=retake_seeds,
retake_variance=retake_variance,
task="retake",
@@ -141,15 +276,45 @@ def create_text2music_ui(
outputs=retake_outputs + [retake_input_params_json],
)
with gr.Tab("repainting"):
retake_variance = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance")
retake_seeds = gr.Textbox(label="repaint seeds (default None)", placeholder="", value=None)
repaint_start = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=0.0, label="Repaint Start Time", interactive=True)
repaint_end = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=30.0, label="Repaint End Time", interactive=True)
repaint_source = gr.Radio(["text2music", "last_repaint", "upload"], value="text2music", label="Repaint Source", elem_id="repaint_source")
retake_variance = gr.Slider(
minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance"
)
retake_seeds = gr.Textbox(
label="repaint seeds (default None)", placeholder="", value=None
)
repaint_start = gr.Slider(
minimum=0.0,
maximum=240.0,
step=0.01,
value=0.0,
label="Repaint Start Time",
interactive=True,
)
repaint_end = gr.Slider(
minimum=0.0,
maximum=240.0,
step=0.01,
value=30.0,
label="Repaint End Time",
interactive=True,
)
repaint_source = gr.Radio(
["text2music", "last_repaint", "upload"],
value="text2music",
label="Repaint Source",
elem_id="repaint_source",
)
repaint_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="repaint_source_audio_upload")
repaint_source_audio_upload = gr.Audio(
label="Upload Audio",
type="filepath",
visible=False,
elem_id="repaint_source_audio_upload",
)
repaint_source.change(
fn=lambda x: gr.update(visible=x == "upload", elem_id="repaint_source_audio_upload"),
fn=lambda x: gr.update(
visible=x == "upload", elem_id="repaint_source_audio_upload"
),
inputs=[repaint_source],
outputs=[repaint_source_audio_upload],
)
@@ -187,9 +352,7 @@ def create_text2music_ui(
if repaint_source == "upload":
src_audio_path = repaint_source_audio_upload
audio_duration = librosa.get_duration(filename=src_audio_path)
json_data = {
"audio_duration": audio_duration
}
json_data = {"audio_duration": audio_duration}
elif repaint_source == "text2music":
json_data = text2music_json_data
src_audio_path = json_data["audio_path"]
@@ -258,11 +421,33 @@ def create_text2music_ui(
with gr.Tab("edit"):
edit_prompt = gr.Textbox(lines=2, label="Edit Tags", max_lines=4)
edit_lyrics = gr.Textbox(lines=9, label="Edit Lyrics", max_lines=13)
retake_seeds = gr.Textbox(label="edit seeds (default None)", placeholder="", value=None)
retake_seeds = gr.Textbox(
label="edit seeds (default None)", placeholder="", value=None
)
edit_type = gr.Radio(["only_lyrics", "remix"], value="only_lyrics", label="Edit Type", elem_id="edit_type", info="`only_lyrics` will keep the whole song the same except lyrics difference. Make your diffrence smaller, e.g. one lyrc line change.\nremix can change the song melody and genre")
edit_n_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.6, label="edit_n_min", interactive=True)
edit_n_max = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="edit_n_max", interactive=True)
edit_type = gr.Radio(
["only_lyrics", "remix"],
value="only_lyrics",
label="Edit Type",
elem_id="edit_type",
info="`only_lyrics` will keep the whole song the same except lyrics difference. Make your diffrence smaller, e.g. one lyrc line change.\nremix can change the song melody and genre",
)
edit_n_min = gr.Slider(
minimum=0.0,
maximum=1.0,
step=0.01,
value=0.6,
label="edit_n_min",
interactive=True,
)
edit_n_max = gr.Slider(
minimum=0.0,
maximum=1.0,
step=0.01,
value=1.0,
label="edit_n_max",
interactive=True,
)
def edit_type_change_func(edit_type):
if edit_type == "only_lyrics":
@@ -276,13 +461,25 @@ def create_text2music_ui(
edit_type.change(
edit_type_change_func,
inputs=[edit_type],
outputs=[edit_n_min, edit_n_max]
outputs=[edit_n_min, edit_n_max],
)
edit_source = gr.Radio(["text2music", "last_edit", "upload"], value="text2music", label="Edit Source", elem_id="edit_source")
edit_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="edit_source_audio_upload")
edit_source = gr.Radio(
["text2music", "last_edit", "upload"],
value="text2music",
label="Edit Source",
elem_id="edit_source",
)
edit_source_audio_upload = gr.Audio(
label="Upload Audio",
type="filepath",
visible=False,
elem_id="edit_source_audio_upload",
)
edit_source.change(
fn=lambda x: gr.update(visible=x == "upload", elem_id="edit_source_audio_upload"),
fn=lambda x: gr.update(
visible=x == "upload", elem_id="edit_source_audio_upload"
),
inputs=[edit_source],
outputs=[edit_source_audio_upload],
)
@@ -321,9 +518,7 @@ def create_text2music_ui(
if edit_source == "upload":
src_audio_path = edit_source_audio_upload
audio_duration = librosa.get_duration(filename=src_audio_path)
json_data = {
"audio_duration": audio_duration
}
json_data = {"audio_duration": audio_duration}
elif edit_source == "text2music":
json_data = text2music_json_data
src_audio_path = json_data["audio_path"]
@@ -397,14 +592,42 @@ def create_text2music_ui(
outputs=edit_outputs + [edit_input_params_json],
)
with gr.Tab("extend"):
extend_seeds = gr.Textbox(label="extend seeds (default None)", placeholder="", value=None)
left_extend_length = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=0.0, label="Left Extend Length", interactive=True)
right_extend_length = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=30.0, label="Right Extend Length", interactive=True)
extend_source = gr.Radio(["text2music", "last_extend", "upload"], value="text2music", label="Extend Source", elem_id="extend_source")
extend_seeds = gr.Textbox(
label="extend seeds (default None)", placeholder="", value=None
)
left_extend_length = gr.Slider(
minimum=0.0,
maximum=240.0,
step=0.01,
value=0.0,
label="Left Extend Length",
interactive=True,
)
right_extend_length = gr.Slider(
minimum=0.0,
maximum=240.0,
step=0.01,
value=30.0,
label="Right Extend Length",
interactive=True,
)
extend_source = gr.Radio(
["text2music", "last_extend", "upload"],
value="text2music",
label="Extend Source",
elem_id="extend_source",
)
extend_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="extend_source_audio_upload")
extend_source_audio_upload = gr.Audio(
label="Upload Audio",
type="filepath",
visible=False,
elem_id="extend_source_audio_upload",
)
extend_source.change(
fn=lambda x: gr.update(visible=x == "upload", elem_id="extend_source_audio_upload"),
fn=lambda x: gr.update(
visible=x == "upload", elem_id="extend_source_audio_upload"
),
inputs=[extend_source],
outputs=[extend_source_audio_upload],
)
@@ -442,9 +665,7 @@ def create_text2music_ui(
src_audio_path = extend_source_audio_upload
# get audio duration
audio_duration = librosa.get_duration(filename=src_audio_path)
json_data = {
"audio_duration": audio_duration
}
json_data = {"audio_duration": audio_duration}
elif extend_source == "text2music":
json_data = text2music_json_data
src_audio_path = json_data["audio_path"]
@@ -531,8 +752,16 @@ def create_text2music_ui(
json_data["use_erg_lyric"],
json_data["use_erg_diffusion"],
", ".join(map(str, json_data["oss_steps"])),
json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
(
json_data["guidance_scale_text"]
if "guidance_scale_text" in json_data
else 0.0
),
(
json_data["guidance_scale_lyric"]
if "guidance_scale_lyric" in json_data
else 0.0
),
)
sample_bnt.click(
@@ -580,7 +809,8 @@ def create_text2music_ui(
oss_steps,
guidance_scale_text,
guidance_scale_lyric,
], outputs=outputs + [input_params_json]
],
outputs=outputs + [input_params_json],
)
@@ -594,7 +824,8 @@ def create_main_demo_ui(
gr.Markdown(
"""
<h1 style="text-align: center;">ACE-Step: A Step Towards Music Generation Foundation Model</h1>
""")
"""
)
with gr.Tab("text2music"):
create_text2music_ui(
app.py
-43
View File
@@ -1,43 +0,0 @@
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint_path", type=str, default="")
parser.add_argument("--server_name", type=str, default="127.0.0.1")
parser.add_argument("--port", type=int, default=7865)
parser.add_argument("--device_id", type=int, default=0)
parser.add_argument("--share", type=bool, default=False)
parser.add_argument("--bf16", type=bool, default=True)
parser.add_argument("--torch_compile", type=bool, default=False)
args = parser.parse_args()
import os
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device_id)
from ui.components import create_main_demo_ui
from pipeline_ace_step import ACEStepPipeline
from data_sampler import DataSampler
def main(args):
model_demo = ACEStepPipeline(
checkpoint_dir=args.checkpoint_path,
dtype="bfloat16" if args.bf16 else "float32",
torch_compile=args.torch_compile
)
data_sampler = DataSampler()
demo = create_main_demo_ui(
text2music_process_func=model_demo.__call__,
sample_data_func=data_sampler.sample,
)
demo.launch(
server_name=args.server_name,
server_port=args.port,
share=args.share
)
if __name__ == "__main__":
main(args)
fig/ACE-Step_framework.png → assets/ACE-Step_framework.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 1.1 MiB

After

Width:  |  Height:  |  Size: 1.1 MiB

fig/Logo_StepFun.png → assets/Logo_StepFun.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 29 KiB

After

Width:  |  Height:  |  Size: 29 KiB

fig/acestudio_logo.png → assets/acestudio_logo.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 126 KiB

After

Width:  |  Height:  |  Size: 126 KiB

fig/application_map.png → assets/application_map.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 253 KiB

After

Width:  |  Height:  |  Size: 253 KiB

fig/demo_interface.png → assets/demo_interface.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 621 KiB

After

Width:  |  Height:  |  Size: 621 KiB

fig/orgnization_logos.png → assets/orgnization_logos.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 302 KiB

After

Width:  |  Height:  |  Size: 302 KiB

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