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 !requirements.txt
*.log *.log
*.flac *.flac
*.wav
minio_config.yaml minio_config.yaml
.history/* .history/*
__pycache__/* __pycache__/*
@@ -200,3 +201,4 @@ app_demo.py
ui/components_demo.py ui/components_demo.py
data_sampler_demo.py data_sampler_demo.py
pipeline_ace_step_demo.py pipeline_ace_step_demo.py
*.wav
Dockerfile
+2 -2
View File
@@ -34,7 +34,7 @@ WORKDIR /app
# Clone the repository # Clone the repository
RUN git clone https://github.com/ace-step/ACE-Step.git . 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 && \ RUN pip3 install --no-cache-dir --upgrade pip && \
pip3 install --no-cache-dir hf_transfer peft && \ 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 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 CMD curl -f http://localhost:7865/ || exit 1
# Command to run the application with GPU support # 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
+103 -119
View File
@@ -1,22 +1,21 @@
<h1 align="center">ACE-Step</h1> # ACE-Step
<h1 align="center">A Step Towards Music Generation Foundation Model</h1>
<p align="center"> A Step Towards Music Generation Foundation Model
<a href="https://ace-step.github.io/">Project</a> |
<a href="https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B">Hugging Face</a> | [Project Page](https://ace-step.github.io/) |
<a href="https://modelscope.cn/models/ACE-Step/ACE-Step-v1-3.5B">ModelScope</a> | [Checkpoints](https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B) |
<a href="https://huggingface.co/spaces/ACE-Step/ACE-Step">Space Demo</a> | [Space Demo](https://huggingface.co/spaces/ACE-Step/ACE-Step) |
<a href="https://discord.gg/rjAZz2xBdG">Discord</a> [Discord](https://discord.gg/rjAZz2xBdG)
</p>
--- ---
<p align="center"> <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> </p>
## Table of Contents ## Table of Contents
- [Features](#-features)
- [Installation](#-installation) - [Installation](#-installation)
- [Features](#-features)
- [Usage](#-usage) - [Usage](#-usage)
- [Train](#-train) - [Train](#-train)
@@ -24,21 +23,43 @@
- 🚀 2025.05.06: Open source demo code and model - 🚀 2025.05.06: Open source demo code and model
## TODOs📋 ## 📋 Roadmap
- [x] Release training code 🔥 - [x] Release training code 🔥
- [x] Release LoRA training code 🔥 - [x] Release LoRA training code 🔥
- [ ] Release RapMachine lora 🎤 - [ ] Release RapMachine LoRA 🎤
- [ ] Release ControlNet training code 🔥 - [ ] Release ControlNet training code 🔥
- [ ] Release Singing2Accompaniment controlnet 🎮 - [ ] Release Singing2Accompaniment ControlNet 🎮
- [ ] Release evaluation performance and technical report 📄 - [ ] 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 ## 🏗️ Architecture
<p align="center"> <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> </p>
## 📝 Abstract ## 📝 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. 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 ## ✨ Features
<p align="center"> <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> </p>
### 🎯 Baseline Quality ### 🎯 Baseline Quality
#### 🌈 Diverse Styles & Genres #### 🌈 Diverse Styles & Genres
- 🎸 Supports all mainstream music styles with various description formats including short tags, descriptive text, or use-case scenarios - 🎸 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 - 🎷 Capable of generating music across different genres with appropriate instrumentation and style
#### 🌍 Multiple Languages #### 🌍 Multiple Languages
- 🗣️ Supports 19 languages with top 10 well-performing languages including: - 🗣️ Supports 19 languages with top 10 well-performing languages including:
- 🇺🇸 English, 🇨🇳 Chinese, 🇷🇺 Russian, 🇪🇸 Spanish, 🇯🇵 Japanese, 🇩🇪 German, 🇫🇷 French, 🇵🇹 Portuguese, 🇮🇹 Italian, 🇰🇷 Korean - 🇺🇸 English, 🇨🇳 Chinese, 🇷🇺 Russian, 🇪🇸 Spanish, 🇯🇵 Japanese, 🇩🇪 German, 🇫🇷 French, 🇵🇹 Portuguese, 🇮🇹 Italian, 🇰🇷 Korean
- ⚠️ Due to data imbalance, less common languages may underperform - ⚠️ Due to data imbalance, less common languages may underperform
#### 🎻 Instrumental Styles #### 🎻 Instrumental Styles
- 🎹 Supports various instrumental music generation across different genres and 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 - 🎺 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 - 🎼 Can generate complex arrangements with multiple instruments while maintaining musical coherence
#### 🎤 Vocal Techniques #### 🎤 Vocal Techniques
- 🎙️ Capable of rendering various vocal styles and techniques with good quality - 🎙️ Capable of rendering various vocal styles and techniques with good quality
- 🗣️ Supports different vocal expressions including various singing techniques and styles - 🗣️ Supports different vocal expressions including various singing techniques and styles
### 🎛️ Controllability ### 🎛️ Controllability
#### 🔄 Variations Generation #### 🔄 Variations Generation
- ⚙️ Implemented using training-free, inference-time optimization techniques - ⚙️ 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 - 🌊 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 - 🎚️ Adjustable mixing ratio between original initial noise and new Gaussian noise to control variation degree
#### 🎨 Repainting #### 🎨 Repainting
- 🖌️ Implemented by adding noise to the target audio input and applying mask constraints during the ODE process - 🖌️ 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 - 🔍 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 - 🔀 Can be combined with Variations Generation techniques to create localized variations in style, lyrics, or vocals
#### ✏️ Lyric Editing #### ✏️ Lyric Editing
- 💡 Innovatively applies flow-edit technology to enable localized lyric modifications while preserving melody, vocals, and accompaniment - 💡 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 - 🔄 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 - ℹ️ 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 ### 🚀 Applications
#### 🎤 Lyric2Vocal (LoRA) #### 🎤 Lyric2Vocal (LoRA)
- 🔊 Based on a LoRA fine-tuned on pure vocal data, allowing direct generation of vocal samples from lyrics - 🔊 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 - 🛠️ 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 - ⏱️ Provides a quick way to test how lyrics might sound when sung, helping songwriters iterate faster
#### 📝 Text2Samples (LoRA) #### 📝 Text2Samples (LoRA)
- 🎛️ Similar to Lyric2Vocal, but fine-tuned on pure instrumental and sample data - 🎛️ Similar to Lyric2Vocal, but fine-tuned on pure instrumental and sample data
- 🎵 Capable of generating conceptual music production samples from text descriptions - 🎵 Capable of generating conceptual music production samples from text descriptions
- 🧰 Useful for quickly creating instrument loops, sound effects, and musical elements for production - 🧰 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 ### 🔮 Coming Soon
#### 🎤 RapMachine #### 🎤 RapMachine
- 🔥 Fine-tuned on pure rap data to create an AI system specialized in rap generation - 🔥 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 - 🏆 Expected capabilities include AI rap battles and narrative expression through rap
- 📚 Rap has exceptional storytelling and expressive capabilities, offering extraordinary application potential - 📚 Rap has exceptional storytelling and expressive capabilities, offering extraordinary application potential
#### 🎛️ StemGen #### 🎛️ StemGen
- 🎚️ A controlnet-lora trained on multi-track data to generate individual instrument stems - 🎚️ 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 - 🎯 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 - 🎹 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 #### 🎤 Singing2Accompaniment
- 🔄 The reverse process of StemGen, generating a mixed master track from a single vocal track - 🔄 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 - 🎵 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 - 🎸 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. 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 ## 🚀 Usage
![Demo Interface](fig/demo_interface.png) ![Demo Interface](assets/demo_interface.png)
### 🔍 Basic Usage ### 🔍 Basic Usage
```bash ```bash
python app.py acestep
``` ```
### ⚙️ Advanced Usage ### ⚙️ Advanced Usage
```bash ```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 #### 🛠️ 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 - 📐 Specify left and right extension lengths
- 🔍 Choose the source audio to extend - 🔍 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. The `examples/input_params` directory contains sample input parameters that can be used as references for generating music.
## Train ## 🔨 Train
### Prerequisites ### Prerequisites
1. Prepare the environment as described in the installation section. 1. Prepare the environment as described in the installation section.
@@ -293,17 +277,17 @@ The `examples/input_params` directory contains sample input parameters that can
Example dataset entry: Example dataset entry:
```json ```json
{ {
"keys": "1ce52937-cd1d-456f-967d-0f1072fcbb58", "keys": "1ce52937-cd1d-456f-967d-0f1072fcbb58",
"tags": ["pop", "acoustic", "ballad", "romantic", "emotional"], "tags": ["pop", "acoustic", "ballad", "romantic", "emotional"],
"speaker_emb_path": "", "speaker_emb_path": "",
"norm_lyrics": "I love you, I love you, I love you", "norm_lyrics": "I love you, I love you, I love you",
"recaption": { "recaption": {
"simplified": "pop", "simplified": "pop",
"expanded": "pop, acoustic, ballad, romantic, emotional", "expanded": "pop, acoustic, ballad, romantic, emotional",
"descriptive": "The sound is soft and gentle, like a tender breeze on a quiet evening. It's soothing and full of longing.", "descriptive": "The sound is soft and gentle, like a tender breeze on a quiet evening. It's soothing and full of longing.",
"use_cases": "Suitable for background music in romantic films or during intimate moments.", "use_cases": "Suitable for background music in romantic films or during intimate moments.",
"analysis": "pop, ballad, piano, guitar, slow tempo, romantic, emotional" "analysis": "pop, ballad, piano, guitar, slow tempo, romantic, emotional"
} }
} }
``` ```
@@ -339,18 +323,18 @@ python trainer.py --dataset_path "path/to/your/dataset" --checkpoint_dir "path/t
Example LoRA configuration file (lora_config.json): Example LoRA configuration file (lora_config.json):
```json ```json
{ {
"r": 16, "r": 16,
"lora_alpha": 32, "lora_alpha": 32,
"target_modules": [ "target_modules": [
"speaker_embedder", "speaker_embedder",
"linear_q", "linear_q",
"linear_k", "linear_k",
"linear_v", "linear_v",
"to_q", "to_q",
"to_k", "to_k",
"to_v", "to_v",
"to_out.0" "to_out.0"
] ]
} }
``` ```
@@ -361,7 +345,7 @@ Example LoRA configuration file (lora_config.json):
- `--reload_dataloaders_every_n_epochs`: Frequency to reload dataloaders (default: 1) - `--reload_dataloaders_every_n_epochs`: Frequency to reload dataloaders (default: 1)
- `--val_check_interval`: Validation check interval (default: None) - `--val_check_interval`: Validation check interval (default: None)
## 📜 License&Disclaimer ## 📜 License & Disclaimer
This project is licensed under [Apache License 2.0](./LICENSE) This project is licensed under [Apache License 2.0](./LICENSE)
@@ -378,10 +362,10 @@ If you find this project useful for your research, please consider citing:
```BibTeX ```BibTeX
@misc{gong2025acestep, @misc{gong2025acestep,
title={ACE-Step: A Step Towards Music Generation Foundation Model}, title={ACE-Step: A Step Towards Music Generation Foundation Model},
author={Junmin Gong, Wenxiao Zhao, Sen Wang, Shengyuan Xu, Jing Guo}, author={Junmin Gong, Wenxiao Zhao, Sen Wang, Shengyuan Xu, Jing Guo},
howpublished={\url{https://github.com/ace-step/ACE-Step}}, howpublished={\url{https://github.com/ace-step/ACE-Step}},
year={2025}, year={2025},
note={GitHub repository} note={GitHub repository}
} }
``` ```
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
+21 -6
View File
@@ -25,7 +25,9 @@ def project(
v1 = torch.nn.functional.normalize(v1, dim=dims) v1 = torch.nn.functional.normalize(v1, dim=dims)
v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1 v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
v0_orthogonal = v0 - v0_parallel 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( def apg_forward(
@@ -65,7 +67,11 @@ def cfg_double_condition_forward(
guidance_scale_text, guidance_scale_text,
guidance_scale_lyric, 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): def optimized_scale(positive_flat, negative_flat):
@@ -74,7 +80,7 @@ def optimized_scale(positive_flat, negative_flat):
dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
# Squared norm of uncondition # Squared norm of uncondition
squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8 squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
# st_star = v_cond^T * v_uncond / ||v_uncond||^2 # st_star = v_cond^T * v_uncond / ||v_uncond||^2
st_star = dot_product / squared_norm st_star = dot_product / squared_norm
@@ -82,14 +88,23 @@ def optimized_scale(positive_flat, negative_flat):
return st_star 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] bsz = noise_pred_with_cond.shape[0]
positive_flat = noise_pred_with_cond.view(bsz, -1) positive_flat = noise_pred_with_cond.view(bsz, -1)
negative_flat = noise_pred_uncond.view(bsz, -1) negative_flat = noise_pred_uncond.view(bsz, -1)
alpha = optimized_scale(positive_flat, negative_flat) alpha = optimized_scale(positive_flat, negative_flat)
alpha = alpha.view(bsz, 1, 1, 1) alpha = alpha.view(bsz, 1, 1, 1)
if (i <= zero_steps) and use_zero_init: if (i <= zero_steps) and use_zero_init:
noise_pred = noise_pred_with_cond * 0. noise_pred = noise_pred_with_cond * 0.0
else: 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 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
+465 -259
View File
File diff suppressed because it is too large Load Diff
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
+63 -57
View File
@@ -21,19 +21,21 @@ import re
from collections import OrderedDict from collections import OrderedDict
from typing import List from typing import List
DIGITS = {str(i): tran for i, tran in enumerate('零一二三四五六七八九')} DIGITS = {str(i): tran for i, tran in enumerate("零一二三四五六七八九")}
UNITS = OrderedDict({ UNITS = OrderedDict(
1: '十', {
2: '百', 1: "十",
3: '千', 2: "百",
4: '万', 3: "千",
8: '亿', 4: "万",
}) 8: "亿",
}
)
COM_QUANTIFIERS = '(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)' COM_QUANTIFIERS = "(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)"
# 分数表达式 # 分数表达式
RE_FRAC = re.compile(r'(-?)(\d+)/(\d+)') RE_FRAC = re.compile(r"(-?)(\d+)/(\d+)")
def replace_frac(match) -> str: 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: def replace_percentage(match) -> str:
@@ -74,7 +76,7 @@ def replace_percentage(match) -> str:
# 整数表达式 # 整数表达式
# 带负号的整数 -10 # 带负号的整数 -10
RE_INTEGER = re.compile(r'(-)' r'(\d+)') RE_INTEGER = re.compile(r"(-)" r"(\d+)")
def replace_negative_num(match) -> str: def replace_negative_num(match) -> str:
@@ -94,7 +96,7 @@ def replace_negative_num(match) -> str:
# 编号-无符号整形 # 编号-无符号整形
# 00078 # 00078
RE_DEFAULT_NUM = re.compile(r'\d{3}\d*') RE_DEFAULT_NUM = re.compile(r"\d{3}\d*")
def replace_default_num(match): def replace_default_num(match):
@@ -112,15 +114,11 @@ def replace_default_num(match):
# RE_ASMD = re.compile( # RE_ASMD = re.compile(
# r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))([\+\-\×÷=])((-?)((\d+)(\.\d+)?)|(\.(\d+)))') # r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))([\+\-\×÷=])((-?)((\d+)(\.\d+)?)|(\.(\d+)))')
RE_ASMD = re.compile( 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: 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 = { power_map = {
'⁰': '0', "⁰": "0",
'¹': '1', "¹": "1",
'²': '2', "²": "2",
'³': '3', "³": "3",
'⁴': '4', "⁴": "4",
'⁵': '5', "⁵": "5",
'⁶': '6', "⁶": "6",
'⁷': '7', "⁷": "7",
'⁸': '8', "⁸": "8",
'⁹': '9', "⁹": "9",
'ˣ': 'x', "ˣ": "x",
'ʸ': 'y', "ʸ": "y",
'ⁿ': 'n' "ⁿ": "n",
} }
def replace_power(match) -> str: def replace_power(match) -> str:
""" """
Args: 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_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: def replace_positive_quantifier(match) -> str:
@@ -221,7 +220,9 @@ RE_RANGE = re.compile(
[-~] # 匹配范围分隔符 [-~] # 匹配范围分隔符
((-?)((\d+)(\.\d+)?)) # 匹配范围结束的负数或正数（整数或小数） ((-?)((\d+)(\.\d+)?)) # 匹配范围结束的负数或正数（整数或小数）
(?![\d\+\-\×÷=]) # 使用正向前瞻以确保数字范围之后没有其他数字和操作符 (?![\d\+\-\×÷=]) # 使用正向前瞻以确保数字范围之后没有其他数字和操作符
""", re.VERBOSE) """,
re.VERBOSE,
)
def replace_range(match) -> str: def replace_range(match) -> str:
@@ -240,7 +241,9 @@ def replace_range(match) -> str:
# ~至表达式 # ~至表达式
RE_TO_RANGE = re.compile( 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: def replace_to_range(match) -> str:
""" """
@@ -249,58 +252,61 @@ def replace_to_range(match) -> str:
Returns: Returns:
str str
""" """
result = match.group(0).replace('~', '至') result = match.group(0).replace("~", "至")
return result return result
def _get_value(value_string: str, use_zero: bool=True) -> List[str]: 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: if len(stripped) == 0:
return [] return []
elif len(stripped) == 1: elif len(stripped) == 1:
if use_zero and len(stripped) < len(value_string): if use_zero and len(stripped) < len(value_string):
return [DIGITS['0'], DIGITS[stripped]] return [DIGITS["0"], DIGITS[stripped]]
else: else:
return [DIGITS[stripped]] return [DIGITS[stripped]]
else: else:
largest_unit = next( 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] first_part = value_string[:-largest_unit]
second_part = value_string[-largest_unit:] second_part = value_string[-largest_unit:]
return _get_value(first_part) + [UNITS[largest_unit]] + _get_value( return _get_value(first_part) + [UNITS[largest_unit]] + _get_value(second_part)
second_part)
def verbalize_cardinal(value_string: str) -> str: def verbalize_cardinal(value_string: str) -> str:
if not value_string: if not value_string:
return '' return ""
# 000 -> '零' , 0 -> '零' # 000 -> '零' , 0 -> '零'
value_string = value_string.lstrip('0') value_string = value_string.lstrip("0")
if len(value_string) == 0: if len(value_string) == 0:
return DIGITS['0'] return DIGITS["0"]
result_symbols = _get_value(value_string) result_symbols = _get_value(value_string)
# verbalized number starting with '一十*' is abbreviated as `十*` # verbalized number starting with '一十*' is abbreviated as `十*`
if len(result_symbols) >= 2 and result_symbols[0] == DIGITS[ if (
'1'] and result_symbols[1] == UNITS[1]: len(result_symbols) >= 2
and result_symbols[0] == DIGITS["1"]
and result_symbols[1] == UNITS[1]
):
result_symbols = result_symbols[1:] result_symbols = result_symbols[1:]
return ''.join(result_symbols) return "".join(result_symbols)
def verbalize_digit(value_string: str, alt_one=False) -> str: def verbalize_digit(value_string: str, alt_one=False) -> str:
result_symbols = [DIGITS[digit] for digit in value_string] result_symbols = [DIGITS[digit] for digit in value_string]
result = ''.join(result_symbols) result = "".join(result_symbols)
if alt_one: if alt_one:
result = result.replace("一", "幺") result = result.replace("一", "幺")
return result return result
def num2str(value_string: str) -> str: def num2str(value_string: str) -> str:
integer_decimal = value_string.split('.') integer_decimal = value_string.split(".")
if len(integer_decimal) == 1: if len(integer_decimal) == 1:
integer = integer_decimal[0] integer = integer_decimal[0]
decimal = '' decimal = ""
elif len(integer_decimal) == 2: elif len(integer_decimal) == 2:
integer, decimal = integer_decimal integer, decimal = integer_decimal
else: else:
@@ -310,12 +316,12 @@ def num2str(value_string: str) -> str:
result = verbalize_cardinal(integer) result = verbalize_cardinal(integer)
decimal = decimal.rstrip('0') decimal = decimal.rstrip("0")
if decimal: if decimal:
# '.22' is verbalized as '零点二二' # '.22' is verbalized as '零点二二'
# '3.20' is verbalized as '三点二 # '3.20' is verbalized as '三点二
result = result if result else "零" result = result if result else "零"
result += '点' + verbalize_digit(decimal) result += "点" + verbalize_digit(decimal)
return result return result
language_segmentation/utils/__init__.py → acestep/models/__init__.py
View File
models/ace_step_transformer.py → acestep/models/ace_step_transformer.py
+154 -45
View File
@@ -31,9 +31,15 @@ from .lyrics_utils.lyric_encoder import ConformerEncoder as LyricEncoder
def cross_norm(hidden_states, controlnet_input): def cross_norm(hidden_states, controlnet_input):
# input N x T x c # input N x T x c
mean_hidden_states, std_hidden_states = hidden_states.mean(dim=(1,2), keepdim=True), hidden_states.std(dim=(1,2), keepdim=True) mean_hidden_states, std_hidden_states = hidden_states.mean(
mean_controlnet_input, std_controlnet_input = controlnet_input.mean(dim=(1,2), keepdim=True), controlnet_input.std(dim=(1,2), keepdim=True) dim=(1, 2), keepdim=True
controlnet_input = (controlnet_input - mean_controlnet_input) * (std_hidden_states / (std_controlnet_input + 1e-12)) + mean_hidden_states ), hidden_states.std(dim=(1, 2), keepdim=True)
mean_controlnet_input, std_controlnet_input = controlnet_input.mean(
dim=(1, 2), keepdim=True
), controlnet_input.std(dim=(1, 2), keepdim=True)
controlnet_input = (controlnet_input - mean_controlnet_input) * (
std_hidden_states / (std_controlnet_input + 1e-12)
) + mean_hidden_states
return controlnet_input return controlnet_input
@@ -45,17 +51,27 @@ class Qwen2RotaryEmbedding(nn.Module):
self.dim = dim self.dim = dim
self.max_position_embeddings = max_position_embeddings self.max_position_embeddings = max_position_embeddings
self.base = base 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) self.register_buffer("inv_freq", inv_freq, persistent=False)
# Build here to make `torch.jit.trace` work. # Build here to make `torch.jit.trace` work.
self._set_cos_sin_cache( 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): def _set_cos_sin_cache(self, seq_len, device, dtype):
self.max_seq_len_cached = seq_len 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) freqs = torch.outer(t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation # 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): def __init__(self, hidden_size, patch_size=[16, 1], out_channels=256):
super().__init__() super().__init__()
self.norm_final = nn.RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.norm_final = nn.RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.linear = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True) self.linear = nn.Linear(
self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5) hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True
)
self.scale_shift_table = nn.Parameter(
torch.randn(2, hidden_size) / hidden_size**0.5
)
self.out_channels = out_channels self.out_channels = out_channels
self.patch_size = patch_size self.patch_size = patch_size
@@ -95,14 +115,28 @@ class T2IFinalLayer(nn.Module):
# 4 unpatchify # 4 unpatchify
new_height, new_width = 1, hidden_states.size(1) new_height, new_width = 1, hidden_states.size(1)
hidden_states = hidden_states.reshape( 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() ).contiguous()
hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
output = hidden_states.reshape( 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() ).contiguous()
if width > new_width: 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: elif width < new_width:
output = output[:, :, :, :width] output = output[:, :, :, :width]
return output return output
@@ -131,9 +165,25 @@ class PatchEmbed(nn.Module):
super().__init__() super().__init__()
patch_size_h, patch_size_w = patch_size patch_size_h, patch_size_w = patch_size
self.early_conv_layers = nn.Sequential( self.early_conv_layers = nn.Sequential(
nn.Conv2d(in_channels, in_channels*256, kernel_size=patch_size, stride=patch_size, padding=0, bias=bias), nn.Conv2d(
torch.nn.GroupNorm(num_groups=32, num_channels=in_channels*256, eps=1e-6, affine=True), in_channels,
nn.Conv2d(in_channels*256, embed_dim, kernel_size=1, stride=1, padding=0, bias=bias) in_channels * 256,
kernel_size=patch_size,
stride=patch_size,
padding=0,
bias=bias,
),
torch.nn.GroupNorm(
num_groups=32, num_channels=in_channels * 256, eps=1e-6, affine=True
),
nn.Conv2d(
in_channels * 256,
embed_dim,
kernel_size=1,
stride=1,
padding=0,
bias=bias,
),
) )
self.patch_size = patch_size self.patch_size = patch_size
self.height, self.width = height // patch_size_h, width // patch_size_w 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 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 _supports_gradient_checkpointing = True
@register_to_config @register_to_config
@@ -217,9 +269,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
) )
self.num_layers = num_layers self.num_layers = num_layers
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) self.time_proj = Timesteps(
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim) num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0
self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True)) )
self.timestep_embedder = TimestepEmbedding(
in_channels=256, time_embed_dim=self.inner_dim
)
self.t_block = nn.Sequential(
nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True)
)
# speaker # speaker
self.speaker_embedder = nn.Linear(speaker_embedding_dim, self.inner_dim) self.speaker_embedder = nn.Linear(speaker_embedding_dim, self.inner_dim)
@@ -229,25 +287,30 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# lyric # lyric
self.lyric_embs = nn.Embedding(lyric_encoder_vocab_size, lyric_hidden_size) 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) self.lyric_proj = nn.Linear(lyric_hidden_size, self.inner_dim)
projector_dim = 2 * 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.Sequential(
nn.SiLU(), nn.Linear(self.inner_dim, projector_dim),
nn.Linear(projector_dim, projector_dim), nn.SiLU(),
nn.SiLU(), nn.Linear(projector_dim, projector_dim),
nn.Linear(projector_dim, ssl_dim), nn.SiLU(),
) for ssl_dim in ssl_latent_dims nn.Linear(projector_dim, ssl_dim),
]) )
for ssl_dim in ssl_latent_dims
]
)
self.ssl_latent_dims = ssl_latent_dims self.ssl_latent_dims = ssl_latent_dims
self.ssl_encoder_depths = ssl_encoder_depths 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.ssl_names = ssl_names
self.proj_in = PatchEmbed( self.proj_in = PatchEmbed(
@@ -258,11 +321,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
bias=True, 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 self.gradient_checkpointing = False
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking # 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 Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). 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 # By default chunk size is 1
chunk_size = chunk_size or 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"): if hasattr(module, "set_chunk_feed_forward"):
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) 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 # N x T x D
lyric_embs = self.lyric_embs(lyric_token_idx) 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) prompt_prenet_out = self.lyric_proj(prompt_prenet_out)
return prompt_prenet_out return prompt_prenet_out
@@ -331,8 +402,17 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
lyric_mask=lyric_mask, lyric_mask=lyric_mask,
) )
encoder_hidden_states = torch.cat([encoder_spk_hidden_states, encoder_text_hidden_states, encoder_lyric_hidden_states], dim=1) encoder_hidden_states = torch.cat(
encoder_hidden_mask = torch.cat([speaker_mask, text_attention_mask, lyric_mask], dim=1) [
encoder_spk_hidden_states,
encoder_text_hidden_states,
encoder_lyric_hidden_states,
],
dim=1,
)
encoder_hidden_mask = torch.cat(
[speaker_mask, text_attention_mask, lyric_mask], dim=1
)
return encoder_hidden_states, encoder_hidden_mask return encoder_hidden_states, encoder_hidden_mask
def decode( def decode(
@@ -344,12 +424,16 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
timestep: Optional[torch.Tensor], timestep: Optional[torch.Tensor],
ssl_hidden_states: Optional[List[torch.Tensor]] = None, ssl_hidden_states: Optional[List[torch.Tensor]] = None,
output_length: int = 0, 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, controlnet_scale: Union[float, torch.Tensor] = 1.0,
return_dict: bool = True, 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) temb = self.t_block(embedded_timestep)
hidden_states = self.proj_in(hidden_states) hidden_states = self.proj_in(hidden_states)
@@ -361,8 +445,12 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
inner_hidden_states = [] inner_hidden_states = []
rotary_freqs_cis = self.rotary_emb(hidden_states, seq_len=hidden_states.shape[1]) rotary_freqs_cis = self.rotary_emb(
encoder_rotary_freqs_cis = self.rotary_emb(encoder_hidden_states, seq_len=encoder_hidden_states.shape[1]) hidden_states, seq_len=hidden_states.shape[1]
)
encoder_rotary_freqs_cis = self.rotary_emb(
encoder_hidden_states, seq_len=encoder_hidden_states.shape[1]
)
for index_block, block in enumerate(self.transformer_blocks): for index_block, block in enumerate(self.transformer_blocks):
@@ -377,7 +465,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
return custom_forward 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( hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block), create_custom_forward(block),
hidden_states=hidden_states, hidden_states=hidden_states,
@@ -406,9 +496,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
inner_hidden_states.append(hidden_states) inner_hidden_states.append(hidden_states)
proj_losses = [] 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: if ssl_hidden_state is None:
continue continue
# 1. N x T x D1 -> N x D x D2 # 1. N x T x D1 -> N x D x D2
@@ -416,9 +512,20 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# 3. projection loss # 3. projection loss
bs = inner_hidden_state.shape[0] bs = inner_hidden_state.shape[0]
proj_loss = 0.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 # 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_tilde = torch.nn.functional.normalize(z_tilde, dim=-1)
z = torch.nn.functional.normalize(z, 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, lyric_mask: Optional[torch.LongTensor] = None,
timestep: Optional[torch.Tensor] = None, timestep: Optional[torch.Tensor] = None,
ssl_hidden_states: Optional[List[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, controlnet_scale: Union[float, torch.Tensor] = 1.0,
return_dict: bool = True, 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: try:
# from .dcformer import DCMHAttention # 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: except ImportError:
# from dcformer import DCMHAttention # 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__) logger = logging.get_logger(__name__)
@@ -55,13 +63,16 @@ def t2i_modulate(x, shift, scale):
return x * (1 + scale) + shift 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): if isinstance(kernel_size, tuple):
return tuple([get_same_padding(ks) for ks in kernel_size]) return tuple([get_same_padding(ks) for ks in kernel_size])
else: else:
assert kernel_size % 2 > 0, f"kernel size {kernel_size} should be odd number" assert kernel_size % 2 > 0, f"kernel size {kernel_size} should be odd number"
return kernel_size // 2 return kernel_size // 2
class ConvLayer(nn.Module): class ConvLayer(nn.Module):
def __init__( def __init__(
self, self,
@@ -187,6 +198,7 @@ class LinearTransformerBlock(nn.Module):
""" """
A Sana block with global shared adaptive layer norm (adaLN-single) conditioning. A Sana block with global shared adaptive layer norm (adaLN-single) conditioning.
""" """
def __init__( def __init__(
self, self,
dim, dim,
models/config.json → acestep/models/config.json
View File
models/customer_attention_processor.py → acestep/models/customer_attention_processor.py
+105 -30
View File
@@ -78,12 +78,16 @@ class CustomLiteLAProcessor2_0:
input_ndim = hidden_states.ndim input_ndim = hidden_states.ndim
if input_ndim == 4: if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape 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: if encoder_hidden_states is not None:
context_input_ndim = encoder_hidden_states.ndim context_input_ndim = encoder_hidden_states.ndim
if context_input_ndim == 4: if context_input_ndim == 4:
batch_size, channel, height, width = encoder_hidden_states.shape 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] batch_size = hidden_states.shape[0]
@@ -94,7 +98,11 @@ class CustomLiteLAProcessor2_0:
value = attn.to_v(hidden_states) value = attn.to_v(hidden_states)
# `context` projections. # `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: 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_query_proj = attn.add_q_proj(encoder_hidden_states)
encoder_hidden_states_key_proj = attn.add_k_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 head_dim = inner_dim // attn.heads
query = query.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1) 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) value = value.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
# RoPE需要 [B, H, S, D] 输入 # RoPE需要 [B, H, S, D] 输入
@@ -140,17 +152,33 @@ class CustomLiteLAProcessor2_0:
if attention_mask is not None: if attention_mask is not None:
# attention_mask: [B, S] -> [B, 1, S, 1] # attention_mask: [B, S] -> [B, 1, S, 1]
attention_mask = attention_mask[:, None, :, None].to(key.dtype) # [B, 1, S, 1] attention_mask = attention_mask[:, None, :, None].to(
query = query * attention_mask.permute(0, 1, 3, 2) # [B, H, S, D] * [B, 1, S, 1] key.dtype
) # [B, 1, S, 1]
query = query * attention_mask.permute(
0, 1, 3, 2
) # [B, H, S, D] * [B, 1, S, 1]
if not attn.is_cross_attention: if not attn.is_cross_attention:
key = key * attention_mask # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘 key = (
value = value * attention_mask.permute(0, 1, 3, 2) # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度 key * attention_mask
) # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘
value = value * attention_mask.permute(
0, 1, 3, 2
) # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度
if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj: if (
encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(key.dtype) # [B, 1, S_enc, 1] attn.is_cross_attention
and encoder_attention_mask is not None
and has_encoder_hidden_state_proj
):
encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(
key.dtype
) # [B, 1, S_enc, 1]
# 此时 key: [B, h, S_enc, D], value: [B, h, D, S_enc] # 此时 key: [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] 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) query = self.kernel_func(query)
key = self.kernel_func(key) key = self.kernel_func(key)
@@ -168,16 +196,22 @@ class CustomLiteLAProcessor2_0:
hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps) 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) hidden_states = hidden_states.to(dtype)
if encoder_hidden_states is not None: if encoder_hidden_states is not None:
encoder_hidden_states = encoder_hidden_states.to(dtype) encoder_hidden_states = encoder_hidden_states.to(dtype)
# Split the attention outputs. # 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, encoder_hidden_states = (
hidden_states[:, : hidden_states_len], hidden_states[:, :hidden_states_len],
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) hidden_states = attn.to_out[0](hidden_states)
# dropout # dropout
hidden_states = attn.to_out[1](hidden_states) 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) encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
if input_ndim == 4: 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: 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: if torch.get_autocast_gpu_dtype() == torch.float16:
hidden_states = hidden_states.clip(-65504, 65504) hidden_states = hidden_states.clip(-65504, 65504)
@@ -208,7 +251,9 @@ class CustomerAttnProcessor2_0:
def __init__(self): def __init__(self):
if not hasattr(F, "scaled_dot_product_attention"): 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( def apply_rotary_emb(
self, self,
@@ -258,23 +303,35 @@ class CustomerAttnProcessor2_0:
if input_ndim == 4: if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape 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, _ = ( 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: 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) query = attn.to_q(hidden_states)
if encoder_hidden_states is None: if encoder_hidden_states is None:
encoder_hidden_states = hidden_states encoder_hidden_states = hidden_states
elif attn.norm_cross: 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) key = attn.to_k(encoder_hidden_states)
value = attn.to_v(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: elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:
key = self.apply_rotary_emb(key, rotary_freqs_cis_cross) 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 # attention_mask: N x S1
# encoder_attention_mask: N x S2 # encoder_attention_mask: N x S2
# cross attention 整合attention_mask和encoder_attention_mask # 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 = 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: 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 # scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length) # (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) # the output of sdp = (batch, num_heads, seq_len, head_dim)
# TODO: add support for attn.scale when we move to Torch 2.1 # 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 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) hidden_states = hidden_states.to(query.dtype)
# linear proj # linear proj
@@ -329,7 +402,9 @@ class CustomerAttnProcessor2_0:
hidden_states = attn.to_out[1](hidden_states) hidden_states = attn.to_out[1](hidden_states)
if input_ndim == 4: 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: if attn.residual_connection:
hidden_states = hidden_states + residual 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
+161 -141
View File
@@ -3,16 +3,19 @@ import math
import torch import torch
from torch import nn from torch import nn
class ConvolutionModule(nn.Module): class ConvolutionModule(nn.Module):
"""ConvolutionModule in Conformer model.""" """ConvolutionModule in Conformer model."""
def __init__(self, def __init__(
channels: int, self,
kernel_size: int = 15, channels: int,
activation: nn.Module = nn.ReLU(), kernel_size: int = 15,
norm: str = "batch_norm", activation: nn.Module = nn.ReLU(),
causal: bool = False, norm: str = "batch_norm",
bias: bool = True): causal: bool = False,
bias: bool = True,
):
"""Construct an ConvolutionModule object. """Construct an ConvolutionModule object.
Args: Args:
channels (int): The number of channels of conv layers. channels (int): The number of channels of conv layers.
@@ -51,7 +54,7 @@ class ConvolutionModule(nn.Module):
bias=bias, bias=bias,
) )
assert norm in ['batch_norm', 'layer_norm'] assert norm in ["batch_norm", "layer_norm"]
if norm == "batch_norm": if norm == "batch_norm":
self.use_layer_norm = False self.use_layer_norm = False
self.norm = nn.BatchNorm1d(channels) self.norm = nn.BatchNorm1d(channels)
@@ -95,13 +98,13 @@ class ConvolutionModule(nn.Module):
if self.lorder > 0: if self.lorder > 0:
if cache.size(2) == 0: # cache_t == 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: else:
assert cache.size(0) == x.size(0) # equal batch assert cache.size(0) == x.size(0) # equal batch
assert cache.size(1) == x.size(1) # equal channel assert cache.size(1) == x.size(1) # equal channel
x = torch.cat((cache, x), dim=2) x = torch.cat((cache, x), dim=2)
assert (x.size(2) > self.lorder) assert x.size(2) > self.lorder
new_cache = x[:, :, -self.lorder:] new_cache = x[:, :, -self.lorder :]
else: else:
# It's better we just return None if no cache is required, # It's better we just return None if no cache is required,
# However, for JIT export, here we just fake one tensor instead of # However, for JIT export, here we just fake one tensor instead of
@@ -126,6 +129,7 @@ class ConvolutionModule(nn.Module):
return x.transpose(1, 2), new_cache return x.transpose(1, 2), new_cache
class PositionwiseFeedForward(torch.nn.Module): class PositionwiseFeedForward(torch.nn.Module):
"""Positionwise feed forward layer. """Positionwise feed forward layer.
@@ -140,11 +144,11 @@ class PositionwiseFeedForward(torch.nn.Module):
""" """
def __init__( def __init__(
self, self,
idim: int, idim: int,
hidden_units: int, hidden_units: int,
dropout_rate: float, dropout_rate: float,
activation: torch.nn.Module = torch.nn.ReLU(), activation: torch.nn.Module = torch.nn.ReLU(),
): ):
"""Construct a PositionwiseFeedForward object.""" """Construct a PositionwiseFeedForward object."""
super(PositionwiseFeedForward, self).__init__() super(PositionwiseFeedForward, self).__init__()
@@ -163,6 +167,7 @@ class PositionwiseFeedForward(torch.nn.Module):
""" """
return self.w_2(self.dropout(self.activation(self.w_1(xs)))) return self.w_2(self.dropout(self.activation(self.w_1(xs))))
class Swish(torch.nn.Module): class Swish(torch.nn.Module):
"""Construct an Swish object.""" """Construct an Swish object."""
@@ -170,6 +175,7 @@ class Swish(torch.nn.Module):
"""Return Swish activation function.""" """Return Swish activation function."""
return x * torch.sigmoid(x) return x * torch.sigmoid(x)
class MultiHeadedAttention(nn.Module): class MultiHeadedAttention(nn.Module):
"""Multi-Head Attention layer. """Multi-Head Attention layer.
@@ -180,11 +186,9 @@ class MultiHeadedAttention(nn.Module):
""" """
def __init__(self, def __init__(
n_head: int, self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True
n_feat: int, ):
dropout_rate: float,
key_bias: bool = True):
"""Construct an MultiHeadedAttention object.""" """Construct an MultiHeadedAttention object."""
super().__init__() super().__init__()
assert n_feat % n_head == 0 assert n_feat % n_head == 0
@@ -229,7 +233,7 @@ class MultiHeadedAttention(nn.Module):
self, self,
value: torch.Tensor, value: torch.Tensor,
scores: 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: ) -> torch.Tensor:
"""Compute attention context vector. """Compute attention context vector.
@@ -251,19 +255,20 @@ class MultiHeadedAttention(nn.Module):
if mask.size(2) > 0: # time2 > 0 if mask.size(2) > 0: # time2 > 0
mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2)
# For last chunk, time2 might be larger than scores.size(-1) # For last chunk, time2 might be larger than scores.size(-1)
mask = mask[:, :, :, :scores.size(-1)] # (batch, 1, *, time2) 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( attn = torch.softmax(scores, dim=-1).masked_fill(
mask, 0.0) # (batch, head, time1, time2) mask, 0.0
) # (batch, head, time1, time2)
else: else:
attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2)
p_attn = self.dropout(attn) p_attn = self.dropout(attn)
x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) x = torch.matmul(p_attn, value) # (batch, head, time1, d_k)
x = (x.transpose(1, 2).contiguous().view(n_batch, -1, x = (
self.h * self.d_k) x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
) # (batch, time1, d_model) ) # (batch, time1, d_model)
return self.linear_out(x) # (batch, time1, d_model) return self.linear_out(x) # (batch, time1, d_model)
@@ -274,7 +279,7 @@ class MultiHeadedAttention(nn.Module):
value: torch.Tensor, value: torch.Tensor,
mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
pos_emb: torch.Tensor = torch.empty(0), 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]: ) -> Tuple[torch.Tensor, torch.Tensor]:
"""Compute scaled dot product attention. """Compute scaled dot product attention.
@@ -308,9 +313,7 @@ class MultiHeadedAttention(nn.Module):
""" """
q, k, v = self.forward_qkv(query, key, value) q, k, v = self.forward_qkv(query, key, value)
if cache.size(0) > 0: if cache.size(0) > 0:
key_cache, value_cache = torch.split(cache, key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
cache.size(-1) // 2,
dim=-1)
k = torch.cat([key_cache, k], dim=2) k = torch.cat([key_cache, k], dim=2)
v = torch.cat([value_cache, v], dim=2) v = torch.cat([value_cache, v], dim=2)
new_cache = torch.cat((k, v), dim=-1) new_cache = torch.cat((k, v), dim=-1)
@@ -328,11 +331,9 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
dropout_rate (float): Dropout rate. dropout_rate (float): Dropout rate.
""" """
def __init__(self, def __init__(
n_head: int, self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True
n_feat: int, ):
dropout_rate: float,
key_bias: bool = True):
"""Construct an RelPositionMultiHeadedAttention object.""" """Construct an RelPositionMultiHeadedAttention object."""
super().__init__(n_head, n_feat, dropout_rate, key_bias) super().__init__(n_head, n_feat, dropout_rate, key_bias)
# linear transformation for positional encoding # linear transformation for positional encoding
@@ -355,14 +356,12 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
torch.Tensor: Output tensor. torch.Tensor: Output tensor.
""" """
zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1), zero_pad = torch.zeros(
device=x.device, (x.size()[0], x.size()[1], x.size()[2], 1), device=x.device, dtype=x.dtype
dtype=x.dtype) )
x_padded = torch.cat([zero_pad, x], dim=-1) x_padded = torch.cat([zero_pad, x], dim=-1)
x_padded = x_padded.view(x.size()[0], x_padded = x_padded.view(x.size()[0], x.size()[1], x.size(3) + 1, x.size(2))
x.size()[1],
x.size(3) + 1, x.size(2))
x = x_padded[:, :, 1:].view_as(x)[ x = x_padded[:, :, 1:].view_as(x)[
:, :, :, : x.size(-1) // 2 + 1 :, :, :, : x.size(-1) // 2 + 1
] # only keep the positions from 0 to time2 ] # only keep the positions from 0 to time2
@@ -375,7 +374,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
value: torch.Tensor, value: torch.Tensor,
mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
pos_emb: torch.Tensor = torch.empty(0), 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]: ) -> Tuple[torch.Tensor, torch.Tensor]:
"""Compute 'Scaled Dot Product Attention' with rel. positional encoding. """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
Args: Args:
@@ -399,9 +398,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
q = q.transpose(1, 2) # (batch, time1, head, d_k) q = q.transpose(1, 2) # (batch, time1, head, d_k)
if cache.size(0) > 0: if cache.size(0) > 0:
key_cache, value_cache = torch.split(cache, key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
cache.size(-1) // 2,
dim=-1)
k = torch.cat([key_cache, k], dim=2) k = torch.cat([key_cache, k], dim=2)
v = torch.cat([value_cache, v], dim=2) v = torch.cat([value_cache, v], dim=2)
# NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
@@ -431,15 +428,15 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
matrix_bd = self.rel_shift(matrix_bd) matrix_bd = self.rel_shift(matrix_bd)
scores = (matrix_ac + matrix_bd) / math.sqrt( 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 return self.forward_attention(v, scores, mask), new_cache
def subsequent_mask( def subsequent_mask(
size: int, size: int,
device: torch.device = torch.device("cpu"), device: torch.device = torch.device("cpu"),
) -> torch.Tensor: ) -> torch.Tensor:
"""Create mask for subsequent steps (size, size). """Create mask for subsequent steps (size, size).
@@ -474,11 +471,11 @@ def subsequent_mask(
def subsequent_chunk_mask( def subsequent_chunk_mask(
size: int, size: int,
chunk_size: int, chunk_size: int,
num_left_chunks: int = -1, num_left_chunks: int = -1,
device: torch.device = torch.device("cpu"), device: torch.device = torch.device("cpu"),
) -> torch.Tensor: ) -> torch.Tensor:
"""Create mask for subsequent steps (size, size) with chunk size, """Create mask for subsequent steps (size, size) with chunk size,
this is for streaming encoder this is for streaming encoder
@@ -510,15 +507,18 @@ def subsequent_chunk_mask(
ret[i, start:ending] = True ret[i, start:ending] = True
return ret return ret
def add_optional_chunk_mask(xs: torch.Tensor,
masks: torch.Tensor, def add_optional_chunk_mask(
use_dynamic_chunk: bool, xs: torch.Tensor,
use_dynamic_left_chunk: bool, masks: torch.Tensor,
decoding_chunk_size: int, use_dynamic_chunk: bool,
static_chunk_size: int, use_dynamic_left_chunk: bool,
num_decoding_left_chunks: int, decoding_chunk_size: int,
enable_full_context: bool = True): static_chunk_size: int,
""" Apply optional mask for encoder. num_decoding_left_chunks: int,
enable_full_context: bool = True,
):
"""Apply optional mask for encoder.
Args: Args:
xs (torch.Tensor): padded input, (B, L, D), L for max length xs (torch.Tensor): padded input, (B, L, D), L for max length
@@ -557,7 +557,7 @@ def add_optional_chunk_mask(xs: torch.Tensor,
# chunk size is either [1, 25] or full context(max_len). # chunk size is either [1, 25] or full context(max_len).
# Since we use 4 times subsampling and allow up to 1s(100 frames) # Since we use 4 times subsampling and allow up to 1s(100 frames)
# delay, the maximum frame is 100 / 4 = 25. # delay, the maximum frame is 100 / 4 = 25.
chunk_size = torch.randint(1, max_len, (1, )).item() chunk_size = torch.randint(1, max_len, (1,)).item()
num_left_chunks = -1 num_left_chunks = -1
if chunk_size > max_len // 2 and enable_full_context: if chunk_size > max_len // 2 and enable_full_context:
chunk_size = max_len chunk_size = max_len
@@ -565,18 +565,17 @@ def add_optional_chunk_mask(xs: torch.Tensor,
chunk_size = chunk_size % 25 + 1 chunk_size = chunk_size % 25 + 1
if use_dynamic_left_chunk: if use_dynamic_left_chunk:
max_left_chunks = (max_len - 1) // chunk_size max_left_chunks = (max_len - 1) // chunk_size
num_left_chunks = torch.randint(0, max_left_chunks, num_left_chunks = torch.randint(0, max_left_chunks, (1,)).item()
(1, )).item() chunk_masks = subsequent_chunk_mask(
chunk_masks = subsequent_chunk_mask(xs.size(1), chunk_size, xs.size(1), chunk_size, num_left_chunks, xs.device
num_left_chunks, ) # (L, L)
xs.device) # (L, L)
chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L) chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L)
chunk_masks = masks & chunk_masks # (B, L, L) chunk_masks = masks & chunk_masks # (B, L, L)
elif static_chunk_size > 0: elif static_chunk_size > 0:
num_left_chunks = num_decoding_left_chunks num_left_chunks = num_decoding_left_chunks
chunk_masks = subsequent_chunk_mask(xs.size(1), static_chunk_size, chunk_masks = subsequent_chunk_mask(
num_left_chunks, xs.size(1), static_chunk_size, num_left_chunks, xs.device
xs.device) # (L, L) ) # (L, L)
chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L) chunk_masks = chunk_masks.unsqueeze(0) # (1, L, L)
chunk_masks = masks & chunk_masks # (B, L, L) chunk_masks = masks & chunk_masks # (B, L, L)
else: else:
@@ -630,7 +629,8 @@ class ConformerEncoderLayer(nn.Module):
if self.conv_module is not None: if self.conv_module is not None:
self.norm_conv = nn.LayerNorm(size, eps=1e-5) # for the CNN module self.norm_conv = nn.LayerNorm(size, eps=1e-5) # for the CNN module
self.norm_final = nn.LayerNorm( 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.dropout = nn.Dropout(dropout_rate)
self.size = size self.size = size
self.normalize_before = normalize_before self.normalize_before = normalize_before
@@ -671,8 +671,7 @@ class ConformerEncoderLayer(nn.Module):
residual = x residual = x
if self.normalize_before: if self.normalize_before:
x = self.norm_ff_macaron(x) x = self.norm_ff_macaron(x)
x = residual + self.ff_scale * self.dropout( x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x))
self.feed_forward_macaron(x))
if not self.normalize_before: if not self.normalize_before:
x = self.norm_ff_macaron(x) x = self.norm_ff_macaron(x)
@@ -680,8 +679,7 @@ class ConformerEncoderLayer(nn.Module):
residual = x residual = x
if self.normalize_before: if self.normalize_before:
x = self.norm_mha(x) x = self.norm_mha(x)
x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, att_cache)
att_cache)
x = residual + self.dropout(x_att) x = residual + self.dropout(x_att)
if not self.normalize_before: if not self.normalize_before:
x = self.norm_mha(x) x = self.norm_mha(x)
@@ -714,7 +712,6 @@ class ConformerEncoderLayer(nn.Module):
return x, mask, new_att_cache, new_cnn_cache return x, mask, new_att_cache, new_cnn_cache
class EspnetRelPositionalEncoding(torch.nn.Module): class EspnetRelPositionalEncoding(torch.nn.Module):
"""Relative positional encoding module (new implementation). """Relative positional encoding module (new implementation).
@@ -770,8 +767,9 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
pe = torch.cat([pe_positive, pe_negative], dim=1) pe = torch.cat([pe_positive, pe_negative], dim=1)
self.pe = pe.to(device=x.device, dtype=x.dtype) self.pe = pe.to(device=x.device, dtype=x.dtype)
def forward(self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0) \ def forward(
-> Tuple[torch.Tensor, torch.Tensor]: self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Add positional encoding. """Add positional encoding.
Args: Args:
@@ -786,10 +784,10 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
pos_emb = self.position_encoding(size=x.size(1), offset=offset) pos_emb = self.position_encoding(size=x.size(1), offset=offset)
return self.dropout(x), self.dropout(pos_emb) return self.dropout(x), self.dropout(pos_emb)
def position_encoding(self, def position_encoding(
offset: Union[int, torch.Tensor], self, offset: Union[int, torch.Tensor], size: int
size: int) -> torch.Tensor: ) -> torch.Tensor:
""" For getting encoding in a streaming fashion """For getting encoding in a streaming fashion
Attention!!!!! Attention!!!!!
we apply dropout only once at the whole utterance level in a none we apply dropout only once at the whole utterance level in a none
@@ -806,12 +804,11 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
""" """
pos_emb = self.pe[ pos_emb = self.pe[
:, :,
self.pe.size(1) // 2 - size + 1: self.pe.size(1) // 2 + size, self.pe.size(1) // 2 - size + 1 : self.pe.size(1) // 2 + size,
] ]
return pos_emb return pos_emb
class LinearEmbed(torch.nn.Module): class LinearEmbed(torch.nn.Module):
"""Linear transform the input without subsampling """Linear transform the input without subsampling
@@ -822,8 +819,9 @@ class LinearEmbed(torch.nn.Module):
""" """
def __init__(self, idim: int, odim: int, dropout_rate: float, def __init__(
pos_enc_class: torch.nn.Module): self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module
):
"""Construct an linear object.""" """Construct an linear object."""
super().__init__() super().__init__()
self.out = torch.nn.Sequential( self.out = torch.nn.Sequential(
@@ -831,16 +829,15 @@ class LinearEmbed(torch.nn.Module):
torch.nn.LayerNorm(odim, eps=1e-5), torch.nn.LayerNorm(odim, eps=1e-5),
torch.nn.Dropout(dropout_rate), torch.nn.Dropout(dropout_rate),
) )
self.pos_enc = pos_enc_class #rel_pos_espnet self.pos_enc = pos_enc_class # rel_pos_espnet
def position_encoding(self, offset: Union[int, torch.Tensor], def position_encoding(
size: int) -> torch.Tensor: self, offset: Union[int, torch.Tensor], size: int
) -> torch.Tensor:
return self.pos_enc.position_encoding(offset, size) return self.pos_enc.position_encoding(offset, size)
def forward( def forward(
self, self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0
x: torch.Tensor,
offset: Union[int, torch.Tensor] = 0
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Input x. """Input x.
@@ -894,16 +891,14 @@ def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
""" """
batch_size = lengths.size(0) batch_size = lengths.size(0)
max_len = max_len if max_len > 0 else lengths.max().item() max_len = max_len if max_len > 0 else lengths.max().item()
seq_range = torch.arange(0, seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device)
max_len,
dtype=torch.int64,
device=lengths.device)
seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
seq_length_expand = lengths.unsqueeze(-1) seq_length_expand = lengths.unsqueeze(-1)
mask = seq_range_expand >= seq_length_expand mask = seq_range_expand >= seq_length_expand
return mask return mask
#https://github.com/FunAudioLLM/CosyVoice/blob/main/examples/magicdata-read/cosyvoice/conf/cosyvoice.yaml
# https://github.com/FunAudioLLM/CosyVoice/blob/main/examples/magicdata-read/cosyvoice/conf/cosyvoice.yaml
class ConformerEncoder(torch.nn.Module): class ConformerEncoder(torch.nn.Module):
"""Conformer encoder module.""" """Conformer encoder module."""
@@ -917,14 +912,14 @@ class ConformerEncoder(torch.nn.Module):
dropout_rate: float = 0.1, dropout_rate: float = 0.1,
positional_dropout_rate: float = 0.1, positional_dropout_rate: float = 0.1,
attention_dropout_rate: float = 0.0, attention_dropout_rate: float = 0.0,
input_layer: str = 'linear', input_layer: str = "linear",
pos_enc_layer_type: str = 'rel_pos_espnet', pos_enc_layer_type: str = "rel_pos_espnet",
normalize_before: bool = True, normalize_before: bool = True,
static_chunk_size: int = 1, # 1: causal_mask; 0: full_mask static_chunk_size: int = 1, # 1: causal_mask; 0: full_mask
use_dynamic_chunk: bool = False, use_dynamic_chunk: bool = False,
use_dynamic_left_chunk: bool = False, use_dynamic_left_chunk: bool = False,
positionwise_conv_kernel_size: int = 1, positionwise_conv_kernel_size: int = 1,
macaron_style: bool =False, macaron_style: bool = False,
selfattention_layer_type: str = "rel_selfattn", selfattention_layer_type: str = "rel_selfattn",
activation_type: str = "swish", activation_type: str = "swish",
use_cnn_module: bool = False, use_cnn_module: bool = False,
@@ -953,8 +948,12 @@ class ConformerEncoder(torch.nn.Module):
""" """
super().__init__() super().__init__()
self.output_size = output_size self.output_size = output_size
self.embed = LinearEmbed(input_size, output_size, dropout_rate, self.embed = LinearEmbed(
EspnetRelPositionalEncoding(output_size, positional_dropout_rate)) input_size,
output_size,
dropout_rate,
EspnetRelPositionalEncoding(output_size, positional_dropout_rate),
)
self.normalize_before = normalize_before self.normalize_before = normalize_before
self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5) self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5)
self.gradient_checkpointing = gradient_checkpointing self.gradient_checkpointing = gradient_checkpointing
@@ -980,40 +979,60 @@ class ConformerEncoder(torch.nn.Module):
activation, activation,
) )
# convolution module definition # convolution module definition
convolution_layer_args = (output_size, cnn_module_kernel, activation, convolution_layer_args = (
cnn_module_norm, causal) output_size,
cnn_module_kernel,
activation,
cnn_module_norm,
causal,
)
self.encoders = torch.nn.ModuleList([ self.encoders = torch.nn.ModuleList(
ConformerEncoderLayer( [
output_size, ConformerEncoderLayer(
RelPositionMultiHeadedAttention( output_size,
*encoder_selfattn_layer_args), RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args),
PositionwiseFeedForward(*positionwise_layer_args), PositionwiseFeedForward(*positionwise_layer_args),
PositionwiseFeedForward( (
*positionwise_layer_args) if macaron_style else None, PositionwiseFeedForward(*positionwise_layer_args)
ConvolutionModule( if macaron_style
*convolution_layer_args) if use_cnn_module else None, else None
dropout_rate, ),
normalize_before, (
) for _ in range(num_blocks) ConvolutionModule(*convolution_layer_args)
]) if use_cnn_module
else None
),
dropout_rate,
normalize_before,
)
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, pos_emb: torch.Tensor,
mask_pad: torch.Tensor) -> torch.Tensor: mask_pad: torch.Tensor,
) -> torch.Tensor:
for layer in self.encoders: for layer in self.encoders:
xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad) xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad)
return xs return xs
@torch.jit.unused @torch.jit.unused
def forward_layers_checkpointed(self, xs: torch.Tensor, def forward_layers_checkpointed(
chunk_masks: torch.Tensor, self,
pos_emb: torch.Tensor, xs: torch.Tensor,
mask_pad: torch.Tensor) -> torch.Tensor: chunk_masks: torch.Tensor,
pos_emb: torch.Tensor,
mask_pad: torch.Tensor,
) -> torch.Tensor:
for layer in self.encoders: for layer in self.encoders:
xs, chunk_masks, _, _ = ckpt.checkpoint(layer.__call__, xs, xs, chunk_masks, _, _ = ckpt.checkpoint(
chunk_masks, pos_emb, layer.__call__, xs, chunk_masks, pos_emb, mask_pad
mask_pad) )
return xs return xs
def forward( def forward(
@@ -1050,15 +1069,17 @@ class ConformerEncoder(torch.nn.Module):
masks = pad_mask.to(torch.bool).unsqueeze(1) # (B, 1, T) masks = pad_mask.to(torch.bool).unsqueeze(1) # (B, 1, T)
xs, pos_emb = self.embed(xs) xs, pos_emb = self.embed(xs)
mask_pad = masks # (B, 1, T/subsample_rate) mask_pad = masks # (B, 1, T/subsample_rate)
chunk_masks = add_optional_chunk_mask(xs, masks, chunk_masks = add_optional_chunk_mask(
self.use_dynamic_chunk, xs,
self.use_dynamic_left_chunk, masks,
decoding_chunk_size, self.use_dynamic_chunk,
self.static_chunk_size, self.use_dynamic_left_chunk,
num_decoding_left_chunks) decoding_chunk_size,
self.static_chunk_size,
num_decoding_left_chunks,
)
if self.gradient_checkpointing and self.training: if self.gradient_checkpointing and self.training:
xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb, xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb, mask_pad)
mask_pad)
else: else:
xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad) xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad)
if self.normalize_before: 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 # return the masks before encoder layers, and the masks will be used
# for cross attention with decoder later # for cross attention with decoder later
return xs, masks 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 from opencc import OpenCC
t2s_converter = OpenCC('t2s') t2s_converter = OpenCC("t2s")
s2t_converter = OpenCC('s2t') s2t_converter = OpenCC("s2t")
EMOJI_PATTERN = re.compile( EMOJI_PATTERN = re.compile(
"[" "["
"\U0001F600-\U0001F64F" # Emoticons "\U0001f600-\U0001f64f" # Emoticons
"]+", flags=re.UNICODE "]+",
flags=re.UNICODE,
) )
# 创建一个翻译表，用于替换和移除字符 # 创建一个翻译表，用于替换和移除字符
TRANSLATION_TABLE = str.maketrans({ TRANSLATION_TABLE = str.maketrans(
'-': ' ', # 将 '-' 替换为空格 {
',': None, "-": " ", # 将 '-' 替换为空格
'.': None, ",": None,
'，': None, ".": None,
'。': None, "，": None,
'!': None, "。": None,
'！': None, "!": None,
'?': None, "！": None,
'？': None, "?": None,
'…': None, "？": None,
';': None, "…": None,
'；': None, ";": None,
':': None, "；": None,
'：': None, ":": None,
'\u3000': ' ', # 将全角空格替换为空格 "：": None,
}) "\u3000": " ", # 将全角空格替换为空格
}
)
# 替换括号中的内容，包括中括号和小括号 # 替换括号中的内容，包括中括号和小括号
BACKSLASH_PATTERN = re.compile(r'\(.*?\)|\[.*?\]') BACKSLASH_PATTERN = re.compile(r"\(.*?\)|\[.*?\]")
SPACE_PATTERN = re.compile('(?<!^)\s+(?!$)') SPACE_PATTERN = re.compile("(?<!^)\s+(?!$)")
def normalize_text(text, language, strip=True): def normalize_text(text, language, strip=True):
@@ -45,10 +48,10 @@ def normalize_text(text, language, strip=True):
text = text.translate(TRANSLATION_TABLE) text = text.translate(TRANSLATION_TABLE)
# Step 2: 移除表情符号 # Step 2: 移除表情符号
text = EMOJI_PATTERN.sub('', text) text = EMOJI_PATTERN.sub("", text)
# Step 3: 连续空白字符替换为单个空格，首位除外 # Step 3: 连续空白字符替换为单个空格，首位除外
text = SPACE_PATTERN.sub(' ', text) text = SPACE_PATTERN.sub(" ", text)
# Step 4: 去除首尾空白字符（如果需要） # Step 4: 去除首尾空白字符（如果需要）
if strip: if strip:
models/lyrics_utils/lyric_tokenizer.py → acestep/models/lyrics_utils/lyric_tokenizer.py
+139 -35
View File
@@ -19,7 +19,7 @@ from .zh_num2words import TextNorm as zh_num2words
from typing import Dict, List, Optional, Set, Union from typing import Dict, List, Optional, Set, Union
#copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/tokenizer.py # copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/tokenizer.py
def get_spacy_lang(lang): def get_spacy_lang(lang):
if lang == "zh": if lang == "zh":
return Chinese() return Chinese()
@@ -446,7 +446,9 @@ _ordinal_re = {
"it": re.compile(r"([0-9]+)(º|°|ª|o|a|i|e)"), "it": re.compile(r"([0-9]+)(º|°|ª|o|a|i|e)"),
"pl": re.compile(r"([0-9]+)(º|ª|st|nd|rd|th)"), "pl": re.compile(r"([0-9]+)(º|ª|st|nd|rd|th)"),
"ar": re.compile(r"([0-9]+)(ون|ين|ث|ر|ى)"), "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]+)(-й|-я|-е|-ое|-ье|-го)"), "ru": re.compile(r"([0-9]+)(-й|-я|-е|-ое|-ье|-го)"),
"nl": re.compile(r"([0-9]+)(de|ste|e)"), "nl": re.compile(r"([0-9]+)(de|ste|e)"),
"tr": re.compile(r"([0-9]+)(\.|inci|nci|uncu|üncü|\.)"), "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"): def _expand_currency(m, lang="en", currency="USD"):
amount = float((re.sub(r"[^\d.]", "", m.group(0).replace(",", ".")))) 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 = { and_equivalents = {
"en": ", ", "en": ", ",
@@ -530,13 +534,21 @@ def expand_numbers_multilingual(text, lang="en"):
else: else:
text = re.sub(_dot_number_re, _remove_dots, text) text = re.sub(_dot_number_re, _remove_dots, text)
try: try:
text = re.sub(_currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text) text = re.sub(
text = re.sub(_currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text) _currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text
text = re.sub(_currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text) )
text = re.sub(
_currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text
)
text = re.sub(
_currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text
)
except: except:
pass pass
if lang != "tr": 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(_ordinal_re[lang], lambda m: _expand_ordinal(m, lang), text)
text = re.sub(_number_re, lambda m: _expand_number(m, lang), text) text = re.sub(_number_re, lambda m: _expand_number(m, lang), text)
return text return text
@@ -582,7 +594,15 @@ def basic_cleaners(text):
def chinese_transliterate(text): def chinese_transliterate(text):
return "".join( 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) 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: class VoiceBpeTokenizer:
@@ -639,7 +661,23 @@ class VoiceBpeTokenizer:
# ) # )
def preprocess_text(self, txt, lang): 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) txt = multilingual_cleaners(txt, lang)
if lang == "zh": if lang == "zh":
txt = chinese_transliterate(txt) txt = chinese_transliterate(txt)
@@ -672,11 +710,12 @@ class VoiceBpeTokenizer:
# txt = txt.replace("[UNK]", "") # txt = txt.replace("[UNK]", "")
return txt return txt
# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3936
#copy from https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3936
def batch_decode( def batch_decode(
self, 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, skip_special_tokens: bool = False,
) -> List[str]: ) -> List[str]:
""" """
@@ -693,12 +732,9 @@ class VoiceBpeTokenizer:
Returns: Returns:
`List[str]`: The list of decoded sentences. `List[str]`: The list of decoded sentences.
""" """
return [ return [self.decode(seq) for seq in sequences]
self.decode(seq)
for seq in sequences
]
#https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/layers/xtts/trainer/dataset.py#L202 # https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/layers/xtts/trainer/dataset.py#L202
# def pad(self): # def pad(self):
def __len__(self): def __len__(self):
@@ -716,15 +752,27 @@ def test_expand_numbers_multilingual():
("This is a 1st test", "This is a first test", "en"), ("This is a 1st test", "This is a first test", "en"),
("That will be $20 sir.", "That will be twenty dollars sir.", "en"), ("That will be $20 sir.", "That will be twenty dollars sir.", "en"),
("That will be 20€ sir.", "That will be twenty euro 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"), ("That's 100,000.5.", "That's one hundred thousand point five.", "en"),
# French # French
("En 12,5 secondes.", "En douze virgule cinq secondes.", "fr"), ("En 12,5 secondes.", "En douze virgule cinq secondes.", "fr"),
("Il y avait 50 soldats.", "Il y avait cinquante soldats.", "fr"), ("Il y avait 50 soldats.", "Il y avait cinquante soldats.", "fr"),
("Ceci est un 1er test", "Ceci est un premier test", "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€ 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"), ("Ce sera 100.000,5.", "Ce sera cent mille virgule cinq.", "fr"),
# German # German
("In 12,5 Sekunden.", "In zwölf Komma fünf Sekunden.", "de"), ("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 ("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 Dollar Herr.", "de"),
("Das macht 20€ Herr.", "Das macht zwanzig Euro 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 # Spanish
("En 12,5 segundos.", "En doce punto cinco segundos.", "es"), ("En 12,5 segundos.", "En doce punto cinco segundos.", "es"),
("Había 50 soldados.", "Había cincuenta soldados.", "es"), ("Había 50 soldados.", "Había cincuenta soldados.", "es"),
("Este es un 1er test", "Este es un primero test", "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 dólares señor.", "es"),
("Eso le costará 20€ señor.", "Eso le costará veinte euros 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 # Italian
("In 12,5 secondi.", "In dodici virgola cinque secondi.", "it"), ("In 12,5 secondi.", "In dodici virgola cinque secondi.", "it"),
("C'erano 50 soldati.", "C'erano cinquanta soldati.", "it"), ("C'erano 50 soldati.", "C'erano cinquanta soldati.", "it"),
("Questo è un 1° test", "Questo è un primo test", "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 dollari signore.", "it"),
("Ti costerà 20€ signore.", "Ti costerà venti euro 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 # Portuguese
("Em 12,5 segundos.", "Em doze vírgula cinco segundos.", "pt"), ("Em 12,5 segundos.", "Em doze vírgula cinco segundos.", "pt"),
("Havia 50 soldados.", "Havia cinquenta soldados.", "pt"), ("Havia 50 soldados.", "Havia cinquenta soldados.", "pt"),
@@ -761,8 +821,16 @@ def test_expand_numbers_multilingual():
# Polish # Polish
("W 12,5 sekundy.", "W dwanaście przecinek pięć sekundy.", "pl"), ("W 12,5 sekundy.", "W dwanaście przecinek pięć sekundy.", "pl"),
("Było 50 żołnierzy.", "Było pięćdziesiąt żołnierzy.", "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 # Arabic
("في الـ 12,5 ثانية.", "في الـ اثنا عشر , خمسون ثانية.", "ar"), ("في الـ 12,5 ثانية.", "في الـ اثنا عشر , خمسون ثانية.", "ar"),
("كان هناك 50 جنديًا.", "كان هناك خمسون جنديًا.", "ar"), ("كان هناك 50 جنديًا.", "كان هناك خمسون جنديًا.", "ar"),
@@ -776,8 +844,16 @@ def test_expand_numbers_multilingual():
# Russian # Russian
("Через 12.5 секунды.", "Через двенадцать запятая пять секунды.", "ru"), ("Через 12.5 секунды.", "Через двенадцать запятая пять секунды.", "ru"),
("Там было 50 солдат.", "Там было пятьдесят солдат.", "ru"), ("Там было 50 солдат.", "Там было пятьдесят солдат.", "ru"),
("Это будет 20.15€ сэр.", "Это будет двадцать евро, пятнадцать центов сэр.", "ru"), (
("Это будет стоить 20€ господин.", "Это будет стоить двадцать евро господин.", "ru"), "Это будет 20.15€ сэр.",
"Это будет двадцать евро, пятнадцать центов сэр.",
"ru",
),
(
"Это будет стоить 20€ господин.",
"Это будет стоить двадцать евро господин.",
"ru",
),
# Dutch # Dutch
("In 12,5 seconden.", "In twaalf komma vijf seconden.", "nl"), ("In 12,5 seconden.", "In twaalf komma vijf seconden.", "nl"),
("Er waren 50 soldaten.", "Er waren vijftig soldaten.", "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"), ("La Dra. Martinez es muy buena.", "La doctora Martinez es muy buena.", "es"),
# French # French
("Bonjour Mr. Dupond.", "Bonjour monsieur Dupond.", "fr"), ("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 # German
("Frau Dr. Müller ist sehr klug.", "Frau doktor Müller ist sehr klug.", "de"), ("Frau Dr. Müller ist sehr klug.", "Frau doktor Müller ist sehr klug.", "de"),
# Portuguese # Portuguese
("Olá Sr. Silva.", "Olá senhor Silva.", "pt"), ("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 # Italian
("Buongiorno, Sig. Rossi.", "Buongiorno, signore Rossi.", "it"), ("Buongiorno, Sig. Rossi.", "Buongiorno, signore Rossi.", "it"),
# ("Sig.ra Bianchi, posso aiutarti?", 'signora Bianchi, posso aiutarti?', 'it'), # Issue with matching that pattern # ("Sig.ra Bianchi, posso aiutarti?", 'signora Bianchi, posso aiutarti?', 'it'), # Issue with matching that pattern
# Polish # Polish
("Dzień dobry, P. Kowalski.", "Dzień dobry, pani Kowalski.", "pl"), ("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 # Czech
("P. Novák", "pan Novák", "cs"), ("P. Novák", "pan Novák", "cs"),
("Dr. Vojtěch", "doktor Vojtěch", "cs"), ("Dr. Vojtěch", "doktor Vojtěch", "cs"),
@@ -837,7 +925,11 @@ def test_abbreviations_multilingual():
("Mevr. de Vries", "mevrouw de Vries", "nl"), ("Mevr. de Vries", "mevrouw de Vries", "nl"),
# Russian # Russian
("Здравствуйте Г-н Иванов.", "Здравствуйте господин Иванов.", "ru"), ("Здравствуйте Г-н Иванов.", "Здравствуйте господин Иванов.", "ru"),
("Д-р Смирнов здесь, чтобы увидеть вас.", "доктор Смирнов здесь, чтобы увидеть вас.", "ru"), (
"Д-р Смирнов здесь, чтобы увидеть вас.",
"доктор Смирнов здесь, чтобы увидеть вас.",
"ru",
),
# Turkish # Turkish
("Merhaba B. Yılmaz.", "Merhaba bay Yılmaz.", "tr"), ("Merhaba B. Yılmaz.", "Merhaba bay Yılmaz.", "tr"),
("Dr. Ayşe burada.", "doktor Ayşe burada.", "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"), ("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"), ("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"), ("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"), ("Moja temperatura to 36.6°", "Moja temperatura to 36.6 stopnie", "pl"),
("Mám 14% baterie", "Mám 14 procento baterie", "cs"), ("Mám 14% baterie", "Mám 14 procento baterie", "cs"),
("Těším se na tebe @ party", "Těším se na tebe na party", "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 في المئة في البطارية", "ar"),
("我的电量为 14%", "我的电量为 14 百分之", "zh"), ("我的电量为 14%", "我的电量为 14 百分之", "zh"),
("Pilim %14 dolu.", "Pilim yüzde 14 dolu.", "tr"), ("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"), ("배터리 잔량이 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
+1 -1
View File
@@ -1,7 +1,7 @@
# Authors: # Authors:
# 2019.5 Zhiyang Zhou (https://github.com/Joee1995/chn_text_norm.git) # 2019.5 Zhiyang Zhou (https://github.com/Joee1995/chn_text_norm.git)
# 2019.9 - 2022 Jiayu DU # 2019.9 - 2022 Jiayu DU
#copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/zh_num2words.py # copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/zh_num2words.py
import argparse import argparse
import csv import csv
import os import os
acestep/music_dcae/__init__.py
View File
music_dcae/music_dcae_pipeline.py → acestep/music_dcae/music_dcae_pipeline.py
+38 -11
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 os
import torch import torch
from diffusers import AutoencoderDC from diffusers import AutoencoderDC
@@ -21,7 +29,12 @@ VOCODER_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_vocoder")
class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin): class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
@register_to_config @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__() super(MusicDCAE, self).__init__()
self.dcae = AutoencoderDC.from_pretrained(dcae_checkpoint_path) 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.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
self.transform = transforms.Compose([ self.transform = transforms.Compose(
transforms.Normalize(0.5, 0.5), [
]) transforms.Normalize(0.5, 0.5),
]
)
self.min_mel_value = -11.0 self.min_mel_value = -11.0
self.max_mel_value = 3.0 self.max_mel_value = 3.0
self.audio_chunk_size = int(round((1024 * 512 / 44100 * 48000))) 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] max_audio_len = audio.shape[-1]
if max_audio_len % (8 * 512) != 0: 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 = self.forward_mel(audio)
mels = (mels - self.min_mel_value) / (self.max_mel_value - self.min_mel_value) 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)) latent = self.dcae.encoder(mel.unsqueeze(0))
latents.append(latent) latents.append(latent)
latents = torch.cat(latents, dim=0) 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 latents = (latents - self.shift_factor) * self.scale_factor
return latents, latent_lengths return latents, latent_lengths
@@ -103,18 +122,26 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
wav = self.vocoder.decode(mels[0]).squeeze(1) wav = self.vocoder.decode(mels[0]).squeeze(1)
if sr is not None: 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) wav = resampler(wav)
else: else:
sr = 44100 sr = 44100
pred_wavs.append(wav) pred_wavs.append(wav)
if audio_lengths is not None: 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 return sr, pred_wavs
def forward(self, audios, audio_lengths=None, sr=None): 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) sr, pred_wavs = self.decode(latents=latents, audio_lengths=audio_lengths, sr=sr)
return sr, pred_wavs, latents, latent_lengths return sr, pred_wavs, latents, latent_lengths
@@ -137,5 +164,5 @@ if __name__ == "__main__":
print("latents shape: ", latents.shape) print("latents shape: ", latents.shape)
print("latent_lengths: ", latent_lengths) print("latent_lengths: ", latent_lengths)
print("sr: ", sr) print("sr: ", sr)
torchaudio.save("test_reconstructed.flac", pred_wavs[0], sr) torchaudio.save("test_reconstructed.wav", pred_wavs[0], sr)
print("test_reconstructed.flac") 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
import torch.nn as nn import torch.nn as nn
from torch import Tensor from torch import Tensor
music_dcae/music_vocoder.py → acestep/music_dcae/music_vocoder.py
+22 -11
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 librosa
import torch import torch
from torch import nn from torch import nn
@@ -132,13 +140,11 @@ class ConvNeXtBlock(nn.Module):
self.act = nn.GELU() self.act = nn.GELU()
self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim) self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)
self.gamma = ( self.gamma = (
nn.Parameter(layer_scale_init_value * nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
torch.ones((dim)), requires_grad=True)
if layer_scale_init_value > 0 if layer_scale_init_value > 0
else None else None
) )
self.drop_path = DropPath( self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
drop_path) if drop_path > 0.0 else nn.Identity()
def forward(self, x, apply_residual: bool = True): def forward(self, x, apply_residual: bool = True):
input = x input = x
@@ -367,8 +373,7 @@ class HiFiGANGenerator(nn.Module):
upsample_rates: Tuple[int] = (8, 8, 2, 2, 2), upsample_rates: Tuple[int] = (8, 8, 2, 2, 2),
upsample_kernel_sizes: Tuple[int] = (16, 16, 8, 2, 2), upsample_kernel_sizes: Tuple[int] = (16, 16, 8, 2, 2),
resblock_kernel_sizes: Tuple[int] = (3, 7, 11), resblock_kernel_sizes: Tuple[int] = (3, 7, 11),
resblock_dilation_sizes: Tuple[Tuple[int]] = ( resblock_dilation_sizes: Tuple[Tuple[int]] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)),
(1, 3, 5), (1, 3, 5), (1, 3, 5)),
num_mels: int = 128, num_mels: int = 128,
upsample_initial_channel: int = 512, upsample_initial_channel: int = 512,
use_template: bool = True, use_template: bool = True,
@@ -417,7 +422,7 @@ class HiFiGANGenerator(nn.Module):
continue continue
if i + 1 < len(upsample_rates): if i + 1 < len(upsample_rates):
stride_f0 = np.prod(upsample_rates[i + 1:]) stride_f0 = np.prod(upsample_rates[i + 1 :])
self.noise_convs.append( self.noise_convs.append(
Conv1d( Conv1d(
1, 1,
@@ -498,7 +503,11 @@ class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
upsample_kernel_sizes: Tuple[int] = (8, 8, 4, 4, 4, 4, 4), upsample_kernel_sizes: Tuple[int] = (8, 8, 4, 4, 4, 4, 4),
resblock_kernel_sizes: Tuple[int] = (3, 7, 11, 13), resblock_kernel_sizes: Tuple[int] = (3, 7, 11, 13),
resblock_dilation_sizes: Tuple[Tuple[int]] = ( resblock_dilation_sizes: Tuple[Tuple[int]] = (
(1, 3, 5), (1, 3, 5), (1, 3, 5), (1, 3, 5)), (1, 3, 5),
(1, 3, 5),
(1, 3, 5),
(1, 3, 5),
),
num_mels: int = 512, num_mels: int = 512,
upsample_initial_channel: int = 1024, upsample_initial_channel: int = 1024,
use_template: bool = False, use_template: bool = False,
@@ -565,12 +574,14 @@ class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
if __name__ == "__main__": if __name__ == "__main__":
import soundfile as sf import soundfile as sf
x = "test_audio.flac" x = "test_audio.wav"
model = ADaMoSHiFiGANV1.from_pretrained("./checkpoints/music_vocoder", local_files_only=True) model = ADaMoSHiFiGANV1.from_pretrained(
"./checkpoints/music_vocoder", local_files_only=True
)
wav, sr = librosa.load(x, sr=44100, mono=True) wav, sr = librosa.load(x, sr=44100, mono=True)
wav = torch.from_numpy(wav).float()[None] wav = torch.from_numpy(wav).float()[None]
mel = model.encode(wav) mel = model.encode(wav)
wav = model.decode(mel)[0].mT 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
+633 -164
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, base_image_seq_len: Optional[int] = 256,
max_image_seq_len: Optional[int] = 4096, 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) timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
sigmas = timesteps / num_train_timesteps 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 # 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: 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: elif self.step_index is not None:
# add_noise is called after first denoising step (for inpainting) # add_noise is called after first denoising step (for inpainting)
step_indices = [self.step_index] * timestep.shape[0] 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: 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: if sigmas is None:
self.num_inference_steps = num_inference_steps self.num_inference_steps = num_inference_steps
timesteps = np.linspace( 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 sigmas = timesteps / self.config.num_train_timesteps
@@ -243,7 +251,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
s_noise: float = 1.0, s_noise: float = 1.0,
generator: Optional[torch.Generator] = None, generator: Optional[torch.Generator] = None,
return_dict: bool = True, return_dict: bool = True,
omega: Union[float, np.array] = 0.0 omega: Union[float, np.array] = 0.0,
) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]: ) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:
""" """
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion 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, num_train_timesteps: int = 1000,
shift: float = 1.0, 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) timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
sigmas = timesteps / num_train_timesteps sigmas = timesteps / num_train_timesteps
@@ -137,7 +139,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
def _sigma_to_t(self, sigma): def _sigma_to_t(self, sigma):
return sigma * self.config.num_train_timesteps 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). 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 self.num_inference_steps = num_inference_steps
timesteps = np.linspace( 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 sigmas = timesteps / self.config.num_train_timesteps
@@ -162,7 +168,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
self.timesteps = timesteps.to(device=device) self.timesteps = timesteps.to(device=device)
sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.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 # empty dt and derivative
self.prev_derivative = None self.prev_derivative = None
@@ -208,7 +216,7 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
s_noise: float = 1.0, s_noise: float = 1.0,
generator: Optional[torch.Generator] = None, generator: Optional[torch.Generator] = None,
return_dict: bool = True, return_dict: bool = True,
omega: Union[float, np.array] = 0.0 omega: Union[float, np.array] = 0.0,
) -> Union[FlowMatchHeunDiscreteSchedulerOutput, Tuple]: ) -> Union[FlowMatchHeunDiscreteSchedulerOutput, Tuple]:
""" """
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion 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 = self.sigmas[self.step_index - 1]
sigma_next = self.sigmas[self.step_index] 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) sigma_hat = sigma * (gamma + 1)
if gamma > 0: if gamma > 0:
noise = randn_tensor( 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 eps = noise * s_noise
sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5 sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
text2music_dataset.py → acestep/text2music_dataset.py
+173 -33
View File
@@ -17,6 +17,7 @@ warnings.simplefilter("ignore", category=FutureWarning)
DEFAULT_TRAIN_PATH = "./data/example_dataset" DEFAULT_TRAIN_PATH = "./data/example_dataset"
def is_silent_audio(audio_tensor, silence_threshold=0.95): def is_silent_audio(audio_tensor, silence_threshold=0.95):
""" """
Determine if an audio is silent and should be discarded 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 # Supported languages for tokenization
SUPPORT_LANGUAGES = { SUPPORT_LANGUAGES = {
"en": 259, "de": 260, "fr": 262, "es": 284, "it": 285, "en": 259,
"pt": 286, "pl": 294, "tr": 295, "ru": 267, "cs": 293, "de": 260,
"nl": 297, "ar": 5022, "zh": 5023, "ja": 5412, "hu": 5753, "fr": 262,
"ko": 6152, "hi": 6680 "es": 284,
"it": 285,
"pt": 286,
"pl": 294,
"tr": 295,
"ru": 267,
"cs": 293,
"nl": 297,
"ar": 5022,
"zh": 5023,
"ja": 5412,
"hu": 5753,
"ko": 6152,
"hi": 6680,
} }
# Regex pattern for structure markers like [Verse], [Chorus], etc. # 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 Dataset for text-to-music generation that processes lyrics and audio files
""" """
def __init__(self, train=True, train_dataset_path=DEFAULT_TRAIN_PATH, def __init__(
max_duration=240.0, sample_size=None, shuffle=True, self,
minibatch_size=1): train=True,
train_dataset_path=DEFAULT_TRAIN_PATH,
max_duration=240.0,
sample_size=None,
shuffle=True,
minibatch_size=1,
):
""" """
Initialize the Text2Music dataset Initialize the Text2Music dataset
@@ -75,14 +95,107 @@ class Text2MusicDataset(Dataset):
# Initialize language segmentation # Initialize language segmentation
self.lang_segment = LangSegment() self.lang_segment = LangSegment()
self.lang_segment.setfilters([ self.lang_segment.setfilters(
'af', 'am', 'an', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'br', 'bs', 'ca', 'cs', 'cy', 'da', 'de', 'dz', 'el', [
'en', 'eo', 'es', 'et', 'eu', 'fa', 'fi', 'fo', 'fr', 'ga', 'gl', 'gu', 'he', 'hi', 'hr', 'ht', 'hu', 'hy', "af",
'id', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'ku', 'ky', 'la', 'lb', 'lo', 'lt', 'lv', 'mg', "am",
'mk', 'ml', 'mn', 'mr', 'ms', 'mt', 'nb', 'ne', 'nl', 'nn', 'no', 'oc', 'or', 'pa', 'pl', 'ps', 'pt', 'qu', "an",
'ro', 'ru', 'rw', 'se', 'si', 'sk', 'sl', 'sq', 'sr', 'sv', 'sw', 'ta', 'te', 'th', 'tl', 'tr', 'ug', 'uk', "ar",
'ur', 'vi', 'vo', 'wa', 'xh', 'zh', 'zu' "as",
]) "az",
"be",
"bg",
"bn",
"br",
"bs",
"ca",
"cs",
"cy",
"da",
"de",
"dz",
"el",
"en",
"eo",
"es",
"et",
"eu",
"fa",
"fi",
"fo",
"fr",
"ga",
"gl",
"gu",
"he",
"hi",
"hr",
"ht",
"hu",
"hy",
"id",
"is",
"it",
"ja",
"jv",
"ka",
"kk",
"km",
"kn",
"ko",
"ku",
"ky",
"la",
"lb",
"lo",
"lt",
"lv",
"mg",
"mk",
"ml",
"mn",
"mr",
"ms",
"mt",
"nb",
"ne",
"nl",
"nn",
"no",
"oc",
"or",
"pa",
"pl",
"ps",
"pt",
"qu",
"ro",
"ru",
"rw",
"se",
"si",
"sk",
"sl",
"sq",
"sr",
"sv",
"sw",
"ta",
"te",
"th",
"tl",
"tr",
"ug",
"uk",
"ur",
"vi",
"vo",
"wa",
"xh",
"zh",
"zu",
]
)
# Initialize lyric tokenizer # Initialize lyric tokenizer
self.lyric_tokenizer = VoiceBpeTokenizer() self.lyric_tokenizer = VoiceBpeTokenizer()
@@ -196,22 +309,30 @@ class Text2MusicDataset(Dataset):
# If debug mode, show tokenization results # If debug mode, show tokenization results
if debug: if debug:
toks = self.lyric_tokenizer.batch_decode([[tok_id] for tok_id in token_idx]) toks = self.lyric_tokenizer.batch_decode(
logger.info(f"debug using most_common_lang {line} --> {most_common_lang} --> {toks}") [[tok_id] for tok_id in token_idx]
)
logger.info(
f"debug using most_common_lang {line} --> {most_common_lang} --> {toks}"
)
# If tokenization contains unknown token (1), try with segment language # If tokenization contains unknown token (1), try with segment language
if 1 in token_idx: if 1 in token_idx:
token_idx = self.lyric_tokenizer.encode(line, lang) token_idx = self.lyric_tokenizer.encode(line, lang)
if debug: 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}") logger.info(f"debug {line} --> {lang} --> {toks}")
# Add tokens and line break # Add tokens and line break
lyric_token_idx = lyric_token_idx + token_idx + [2] lyric_token_idx = lyric_token_idx + token_idx + [2]
except Exception as e: 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 return lyric_token_idx
@@ -302,7 +423,9 @@ class Text2MusicDataset(Dataset):
# Pad to minimum 3 seconds if needed # Pad to minimum 3 seconds if needed
if audio.shape[-1] < 48000 * 3: 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 # Check if audio is silent
if is_silent_audio(audio): if is_silent_audio(audio):
@@ -368,9 +491,11 @@ class Text2MusicDataset(Dataset):
lyrics = item["norm_lyrics"] lyrics = item["norm_lyrics"]
lyrics_lines = lyrics.split("\n") lyrics_lines = lyrics.split("\n")
for lyric_line in lyrics_lines: for lyric_line in lyrics_lines:
candidate_lyric_chunk.append({ candidate_lyric_chunk.append(
"lyric": lyric_line, {
}) "lyric": lyric_line,
}
)
# Limit audio length # Limit audio length
longest_length = 24 * 10 * 48000 # 240 seconds longest_length = 24 * 10 * 48000 # 240 seconds
@@ -482,28 +607,43 @@ class Text2MusicDataset(Dataset):
elif k in ["src_wavs", "target_wavs", "vocal_wavs"]: elif k in ["src_wavs", "target_wavs", "vocal_wavs"]:
# Pad audio to max length # Pad audio to max length
max_length = max(seq.shape[1] for seq in v) max_length = max(seq.shape[1] for seq in v)
padded_input_list = torch.stack([ padded_input_list = torch.stack(
torch.nn.functional.pad(seq, (0, max_length - seq.shape[1]), 'constant', 0) [
for seq in v torch.nn.functional.pad(
]) seq, (0, max_length - seq.shape[1]), "constant", 0
)
for seq in v
]
)
elif k in ["clap_conditions"]: elif k in ["clap_conditions"]:
# Pad time dimension of embeddings # Pad time dimension of embeddings
max_length = max(seq.shape[0] for seq in v) max_length = max(seq.shape[0] for seq in v)
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 for seq in v
] ]
padded_input_list = torch.stack(v) padded_input_list = torch.stack(v)
elif k == "speaker_embs": elif k == "speaker_embs":
# Stack speaker embeddings # Stack speaker embeddings
padded_input_list = torch.stack(v) 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 # Pad sequence tensors
max_length = max(len(seq) for seq in v) max_length = max(len(seq) for seq in v)
padded_input_list = torch.stack([ padded_input_list = torch.stack(
torch.nn.functional.pad(seq, (0, max_length - len(seq)), 'constant', 0) [
for seq in v torch.nn.functional.pad(
]) seq, (0, max_length - len(seq)), "constant", 0
)
for seq in v
]
)
output[k] = padded_input_list 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 gradio as gr
import librosa import librosa
@@ -41,8 +49,6 @@ In this moment we take flight
""" """
def create_output_ui(task_name="Text2Music"): def create_output_ui(task_name="Text2Music"):
# For many consumer-grade GPU devices, only one batch can be run # 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") 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.Column():
with gr.Row(equal_height=True): with gr.Row(equal_height=True):
# add markdown, tags and lyrics examples are from ai music generation community # 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) sample_bnt = gr.Button("Sample", variant="primary", scale=1)
prompt = gr.Textbox(lines=2, label="Tags", max_lines=4, value=TAG_DEFAULT, info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community") prompt = gr.Textbox(
lyrics = gr.Textbox(lines=9, label="Lyrics", max_lines=13, value=LYRIC_DEFAULT, info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics") lines=2,
label="Tags",
max_lines=4,
value=TAG_DEFAULT,
info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community",
)
lyrics = gr.Textbox(
lines=9,
label="Lyrics",
max_lines=13,
value=LYRIC_DEFAULT,
info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics",
)
with gr.Accordion("Basic Settings", open=False): with gr.Accordion("Basic Settings", open=False):
infer_step = gr.Slider(minimum=1, maximum=1000, step=1, value=27, label="Infer Steps", interactive=True) infer_step = gr.Slider(
guidance_scale = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=15.0, label="Guidance Scale", interactive=True, info="When guidance_scale_lyric > 1 and guidance_scale_text > 1, the guidance scale will not be applied.") minimum=1,
guidance_scale_text = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Text", interactive=True, info="Guidance scale for text condition. It can only apply to cfg. set guidance_scale_text=5.0, guidance_scale_lyric=1.5 for start") maximum=1000,
guidance_scale_lyric = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Lyric", interactive=True) step=1,
value=27,
label="Infer Steps",
interactive=True,
)
guidance_scale = gr.Slider(
minimum=0.0,
maximum=200.0,
step=0.1,
value=15.0,
label="Guidance Scale",
interactive=True,
info="When guidance_scale_lyric > 1 and guidance_scale_text > 1, the guidance scale will not be applied.",
)
guidance_scale_text = gr.Slider(
minimum=0.0,
maximum=10.0,
step=0.1,
value=0.0,
label="Guidance Scale Text",
interactive=True,
info="Guidance scale for text condition. It can only apply to cfg. set guidance_scale_text=5.0, guidance_scale_lyric=1.5 for start",
)
guidance_scale_lyric = gr.Slider(
minimum=0.0,
maximum=10.0,
step=0.1,
value=0.0,
label="Guidance Scale Lyric",
interactive=True,
)
manual_seeds = gr.Textbox(label="manual seeds (default None)", placeholder="1,2,3,4", value=None, info="Seed for the generation") manual_seeds = gr.Textbox(
label="manual seeds (default None)",
placeholder="1,2,3,4",
value=None,
info="Seed for the generation",
)
with gr.Accordion("Advanced Settings", open=False): with gr.Accordion("Advanced Settings", open=False):
scheduler_type = gr.Radio(["euler", "heun"], value="euler", label="Scheduler Type", elem_id="scheduler_type", info="Scheduler type for the generation. euler is recommended. heun will take more time.") scheduler_type = gr.Radio(
cfg_type = gr.Radio(["cfg", "apg", "cfg_star"], value="apg", label="CFG Type", elem_id="cfg_type", info="CFG type for the generation. apg is recommended. cfg and cfg_star are almost the same.") ["euler", "heun"],
use_erg_tag = gr.Checkbox(label="use ERG for tag", value=True, info="Use Entropy Rectifying Guidance for tag. It will multiple a temperature to the attention to make a weaker tag condition and make better diversity.") value="euler",
use_erg_lyric = gr.Checkbox(label="use ERG for lyric", value=True, info="The same but apply to lyric encoder's attention.") label="Scheduler Type",
use_erg_diffusion = gr.Checkbox(label="use ERG for diffusion", value=True, info="The same but apply to diffusion model's attention.") elem_id="scheduler_type",
info="Scheduler type for the generation. euler is recommended. heun will take more time.",
)
cfg_type = gr.Radio(
["cfg", "apg", "cfg_star"],
value="apg",
label="CFG Type",
elem_id="cfg_type",
info="CFG type for the generation. apg is recommended. cfg and cfg_star are almost the same.",
)
use_erg_tag = gr.Checkbox(
label="use ERG for tag",
value=True,
info="Use Entropy Rectifying Guidance for tag. It will multiple a temperature to the attention to make a weaker tag condition and make better diversity.",
)
use_erg_lyric = gr.Checkbox(
label="use ERG for lyric",
value=True,
info="The same but apply to lyric encoder's attention.",
)
use_erg_diffusion = gr.Checkbox(
label="use ERG for diffusion",
value=True,
info="The same but apply to diffusion model's attention.",
)
omega_scale = gr.Slider(minimum=-100.0, maximum=100.0, step=0.1, value=10.0, label="Granularity Scale", interactive=True, info="Granularity scale for the generation. Higher values can reduce artifacts") omega_scale = gr.Slider(
minimum=-100.0,
maximum=100.0,
step=0.1,
value=10.0,
label="Granularity Scale",
interactive=True,
info="Granularity scale for the generation. Higher values can reduce artifacts",
)
guidance_interval = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Guidance Interval", interactive=True, info="Guidance interval for the generation. 0.5 means only apply guidance in the middle steps (0.25 * infer_steps to 0.75 * infer_steps)") guidance_interval = gr.Slider(
guidance_interval_decay = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.0, label="Guidance Interval Decay", interactive=True, info="Guidance interval decay for the generation. Guidance scale will decay from guidance_scale to min_guidance_scale in the interval. 0.0 means no decay.") minimum=0.0,
min_guidance_scale = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=3.0, label="Min Guidance Scale", interactive=True, info="Min guidance scale for guidance interval decay's end scale") maximum=1.0,
oss_steps = gr.Textbox(label="OSS Steps", placeholder="16, 29, 52, 96, 129, 158, 172, 183, 189, 200", value=None, info="Optimal Steps for the generation. But not test well") step=0.01,
value=0.5,
label="Guidance Interval",
interactive=True,
info="Guidance interval for the generation. 0.5 means only apply guidance in the middle steps (0.25 * infer_steps to 0.75 * infer_steps)",
)
guidance_interval_decay = gr.Slider(
minimum=0.0,
maximum=1.0,
step=0.01,
value=0.0,
label="Guidance Interval Decay",
interactive=True,
info="Guidance interval decay for the generation. Guidance scale will decay from guidance_scale to min_guidance_scale in the interval. 0.0 means no decay.",
)
min_guidance_scale = gr.Slider(
minimum=0.0,
maximum=200.0,
step=0.1,
value=3.0,
label="Min Guidance Scale",
interactive=True,
info="Min guidance scale for guidance interval decay's end scale",
)
oss_steps = gr.Textbox(
label="OSS Steps",
placeholder="16, 29, 52, 96, 129, 158, 172, 183, 189, 200",
value=None,
info="Optimal Steps for the generation. But not test well",
)
text2music_bnt = gr.Button("Generate", variant="primary") text2music_bnt = gr.Button("Generate", variant="primary")
with gr.Column(): with gr.Column():
outputs, input_params_json = create_output_ui() outputs, input_params_json = create_output_ui()
with gr.Tab("retake"): with gr.Tab("retake"):
retake_variance = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance") retake_variance = gr.Slider(
retake_seeds = gr.Textbox(label="retake seeds (default None)", placeholder="", value=None) minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance"
)
retake_seeds = gr.Textbox(
label="retake seeds (default None)", placeholder="", value=None
)
retake_bnt = gr.Button("Retake", variant="primary") retake_bnt = gr.Button("Retake", variant="primary")
retake_outputs, retake_input_params_json = create_output_ui("Retake") 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_lyric"],
json_data["use_erg_diffusion"], json_data["use_erg_diffusion"],
", ".join(map(str, json_data["oss_steps"])), ", ".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_seeds=retake_seeds,
retake_variance=retake_variance, retake_variance=retake_variance,
task="retake", task="retake",
@@ -141,15 +276,45 @@ def create_text2music_ui(
outputs=retake_outputs + [retake_input_params_json], outputs=retake_outputs + [retake_input_params_json],
) )
with gr.Tab("repainting"): with gr.Tab("repainting"):
retake_variance = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance") retake_variance = gr.Slider(
retake_seeds = gr.Textbox(label="repaint seeds (default None)", placeholder="", value=None) minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance"
repaint_start = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=0.0, label="Repaint Start Time", interactive=True) )
repaint_end = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=30.0, label="Repaint End Time", interactive=True) retake_seeds = gr.Textbox(
repaint_source = gr.Radio(["text2music", "last_repaint", "upload"], value="text2music", label="Repaint Source", elem_id="repaint_source") label="repaint seeds (default None)", placeholder="", value=None
)
repaint_start = gr.Slider(
minimum=0.0,
maximum=240.0,
step=0.01,
value=0.0,
label="Repaint Start Time",
interactive=True,
)
repaint_end = gr.Slider(
minimum=0.0,
maximum=240.0,
step=0.01,
value=30.0,
label="Repaint End Time",
interactive=True,
)
repaint_source = gr.Radio(
["text2music", "last_repaint", "upload"],
value="text2music",
label="Repaint Source",
elem_id="repaint_source",
)
repaint_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="repaint_source_audio_upload") repaint_source_audio_upload = gr.Audio(
label="Upload Audio",
type="filepath",
visible=False,
elem_id="repaint_source_audio_upload",
)
repaint_source.change( 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], inputs=[repaint_source],
outputs=[repaint_source_audio_upload], outputs=[repaint_source_audio_upload],
) )
@@ -187,9 +352,7 @@ def create_text2music_ui(
if repaint_source == "upload": if repaint_source == "upload":
src_audio_path = repaint_source_audio_upload src_audio_path = repaint_source_audio_upload
audio_duration = librosa.get_duration(filename=src_audio_path) audio_duration = librosa.get_duration(filename=src_audio_path)
json_data = { json_data = {"audio_duration": audio_duration}
"audio_duration": audio_duration
}
elif repaint_source == "text2music": elif repaint_source == "text2music":
json_data = text2music_json_data json_data = text2music_json_data
src_audio_path = json_data["audio_path"] src_audio_path = json_data["audio_path"]
@@ -258,11 +421,33 @@ def create_text2music_ui(
with gr.Tab("edit"): with gr.Tab("edit"):
edit_prompt = gr.Textbox(lines=2, label="Edit Tags", max_lines=4) edit_prompt = gr.Textbox(lines=2, label="Edit Tags", max_lines=4)
edit_lyrics = gr.Textbox(lines=9, label="Edit Lyrics", max_lines=13) edit_lyrics = gr.Textbox(lines=9, label="Edit Lyrics", max_lines=13)
retake_seeds = gr.Textbox(label="edit seeds (default None)", placeholder="", value=None) retake_seeds = gr.Textbox(
label="edit seeds (default None)", placeholder="", value=None
)
edit_type = gr.Radio(["only_lyrics", "remix"], value="only_lyrics", label="Edit Type", elem_id="edit_type", info="`only_lyrics` will keep the whole song the same except lyrics difference. Make your diffrence smaller, e.g. one lyrc line change.\nremix can change the song melody and genre") edit_type = gr.Radio(
edit_n_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.6, label="edit_n_min", interactive=True) ["only_lyrics", "remix"],
edit_n_max = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="edit_n_max", interactive=True) value="only_lyrics",
label="Edit Type",
elem_id="edit_type",
info="`only_lyrics` will keep the whole song the same except lyrics difference. Make your diffrence smaller, e.g. one lyrc line change.\nremix can change the song melody and genre",
)
edit_n_min = gr.Slider(
minimum=0.0,
maximum=1.0,
step=0.01,
value=0.6,
label="edit_n_min",
interactive=True,
)
edit_n_max = gr.Slider(
minimum=0.0,
maximum=1.0,
step=0.01,
value=1.0,
label="edit_n_max",
interactive=True,
)
def edit_type_change_func(edit_type): def edit_type_change_func(edit_type):
if edit_type == "only_lyrics": if edit_type == "only_lyrics":
@@ -276,13 +461,25 @@ def create_text2music_ui(
edit_type.change( edit_type.change(
edit_type_change_func, edit_type_change_func,
inputs=[edit_type], inputs=[edit_type],
outputs=[edit_n_min, edit_n_max] outputs=[edit_n_min, edit_n_max],
) )
edit_source = gr.Radio(["text2music", "last_edit", "upload"], value="text2music", label="Edit Source", elem_id="edit_source") edit_source = gr.Radio(
edit_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="edit_source_audio_upload") ["text2music", "last_edit", "upload"],
value="text2music",
label="Edit Source",
elem_id="edit_source",
)
edit_source_audio_upload = gr.Audio(
label="Upload Audio",
type="filepath",
visible=False,
elem_id="edit_source_audio_upload",
)
edit_source.change( 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], inputs=[edit_source],
outputs=[edit_source_audio_upload], outputs=[edit_source_audio_upload],
) )
@@ -321,9 +518,7 @@ def create_text2music_ui(
if edit_source == "upload": if edit_source == "upload":
src_audio_path = edit_source_audio_upload src_audio_path = edit_source_audio_upload
audio_duration = librosa.get_duration(filename=src_audio_path) audio_duration = librosa.get_duration(filename=src_audio_path)
json_data = { json_data = {"audio_duration": audio_duration}
"audio_duration": audio_duration
}
elif edit_source == "text2music": elif edit_source == "text2music":
json_data = text2music_json_data json_data = text2music_json_data
src_audio_path = json_data["audio_path"] src_audio_path = json_data["audio_path"]
@@ -397,14 +592,42 @@ def create_text2music_ui(
outputs=edit_outputs + [edit_input_params_json], outputs=edit_outputs + [edit_input_params_json],
) )
with gr.Tab("extend"): with gr.Tab("extend"):
extend_seeds = gr.Textbox(label="extend seeds (default None)", placeholder="", value=None) extend_seeds = gr.Textbox(
left_extend_length = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=0.0, label="Left Extend Length", interactive=True) label="extend seeds (default None)", placeholder="", value=None
right_extend_length = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=30.0, label="Right Extend Length", interactive=True) )
extend_source = gr.Radio(["text2music", "last_extend", "upload"], value="text2music", label="Extend Source", elem_id="extend_source") left_extend_length = gr.Slider(
minimum=0.0,
maximum=240.0,
step=0.01,
value=0.0,
label="Left Extend Length",
interactive=True,
)
right_extend_length = gr.Slider(
minimum=0.0,
maximum=240.0,
step=0.01,
value=30.0,
label="Right Extend Length",
interactive=True,
)
extend_source = gr.Radio(
["text2music", "last_extend", "upload"],
value="text2music",
label="Extend Source",
elem_id="extend_source",
)
extend_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="extend_source_audio_upload") extend_source_audio_upload = gr.Audio(
label="Upload Audio",
type="filepath",
visible=False,
elem_id="extend_source_audio_upload",
)
extend_source.change( 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], inputs=[extend_source],
outputs=[extend_source_audio_upload], outputs=[extend_source_audio_upload],
) )
@@ -442,9 +665,7 @@ def create_text2music_ui(
src_audio_path = extend_source_audio_upload src_audio_path = extend_source_audio_upload
# get audio duration # get audio duration
audio_duration = librosa.get_duration(filename=src_audio_path) audio_duration = librosa.get_duration(filename=src_audio_path)
json_data = { json_data = {"audio_duration": audio_duration}
"audio_duration": audio_duration
}
elif extend_source == "text2music": elif extend_source == "text2music":
json_data = text2music_json_data json_data = text2music_json_data
src_audio_path = json_data["audio_path"] src_audio_path = json_data["audio_path"]
@@ -531,8 +752,16 @@ def create_text2music_ui(
json_data["use_erg_lyric"], json_data["use_erg_lyric"],
json_data["use_erg_diffusion"], json_data["use_erg_diffusion"],
", ".join(map(str, json_data["oss_steps"])), ", ".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( sample_bnt.click(
@@ -580,7 +809,8 @@ def create_text2music_ui(
oss_steps, oss_steps,
guidance_scale_text, guidance_scale_text,
guidance_scale_lyric, guidance_scale_lyric,
], outputs=outputs + [input_params_json] ],
outputs=outputs + [input_params_json],
) )
@@ -594,7 +824,8 @@ def create_main_demo_ui(
gr.Markdown( gr.Markdown(
""" """
<h1 style="text-align: center;">ACE-Step: A Step Towards Music Generation Foundation Model</h1> <h1 style="text-align: center;">ACE-Step: A Step Towards Music Generation Foundation Model</h1>
""") """
)
with gr.Tab("text2music"): with gr.Tab("text2music"):
create_text2music_ui( 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
+84 -75
View File
@@ -1,48 +1,54 @@
import argparse import click
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 os import os
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device_id) from acestep.pipeline_ace_step import ACEStepPipeline
from acestep.data_sampler import DataSampler
from pipeline_ace_step import ACEStepPipeline
from data_sampler import DataSampler
def sample_data(json_data): def sample_data(json_data):
return ( return (
json_data["audio_duration"], json_data["audio_duration"],
json_data["prompt"], json_data["prompt"],
json_data["lyrics"], json_data["lyrics"],
json_data["infer_step"], json_data["infer_step"],
json_data["guidance_scale"], json_data["guidance_scale"],
json_data["scheduler_type"], json_data["scheduler_type"],
json_data["cfg_type"], json_data["cfg_type"],
json_data["omega_scale"], json_data["omega_scale"],
", ".join(map(str, json_data["actual_seeds"])), ", ".join(map(str, json_data["actual_seeds"])),
json_data["guidance_interval"], json_data["guidance_interval"],
json_data["guidance_interval_decay"], json_data["guidance_interval_decay"],
json_data["min_guidance_scale"], json_data["min_guidance_scale"],
json_data["use_erg_tag"], json_data["use_erg_tag"],
json_data["use_erg_lyric"], json_data["use_erg_lyric"],
json_data["use_erg_diffusion"], json_data["use_erg_diffusion"],
", ".join(map(str, json_data["oss_steps"])), ", ".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_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( model_demo = ACEStepPipeline(
checkpoint_dir=args.checkpoint_path, checkpoint_dir=checkpoint_path,
dtype="bfloat16" if args.bf16 else "float32", dtype="bfloat16" if bf16 else "float32",
torch_compile=args.torch_compile torch_compile=torch_compile,
) )
print(model_demo) print(model_demo)
@@ -52,46 +58,49 @@ def main(args):
json_data = sample_data(json_data) json_data = sample_data(json_data)
print(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
audio_duration,\ model_demo(
prompt, \ audio_duration,
lyrics,\ prompt,
infer_step, \ lyrics,
guidance_scale,\ infer_step,
scheduler_type, \ guidance_scale,
cfg_type, \ scheduler_type,
omega_scale, \ cfg_type,
manual_seeds, \ omega_scale,
guidance_interval, \ manual_seeds,
guidance_interval_decay, \ guidance_interval,
min_guidance_scale, \ guidance_interval_decay,
use_erg_tag, \ min_guidance_scale,
use_erg_lyric, \ use_erg_tag,
use_erg_diffusion, \ use_erg_lyric,
oss_steps, \ use_erg_diffusion,
guidance_scale_text, \ oss_steps,
guidance_scale_lyric = json_data guidance_scale_text,
guidance_scale_lyric,
save_path=output_path,
)
model_demo(audio_duration, \
prompt, \
lyrics,\
infer_step, \
guidance_scale,\
scheduler_type, \
cfg_type, \
omega_scale, \
manual_seeds, \
guidance_interval, \
guidance_interval_decay, \
min_guidance_scale, \
use_erg_tag, \
use_erg_lyric, \
use_erg_diffusion, \
oss_steps, \
guidance_scale_text, \
guidance_scale_lyric,
save_path=args.output_path)
if __name__ == "__main__": 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 accelerate==1.6.0
cutlet cutlet
fugashi[unidic-lite] 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 import torch.utils.data
from pytorch_lightning.core import LightningModule from pytorch_lightning.core import LightningModule
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
from schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from acestep.schedulers.scheduling_flow_match_euler_discrete import (
from text2music_dataset import Text2MusicDataset FlowMatchEulerDiscreteScheduler,
)
from acestep.text2music_dataset import Text2MusicDataset
from loguru import logger from loguru import logger
from transformers import AutoModel, Wav2Vec2FeatureExtractor from transformers import AutoModel, Wav2Vec2FeatureExtractor
import torchaudio 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 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 from tqdm import tqdm
import random import random
import os import os
from pipeline_ace_step import ACEStepPipeline from acestep.pipeline_ace_step import ACEStepPipeline
matplotlib.use("Agg") matplotlib.use("Agg")
torch.backends.cudnn.benchmark = False torch.backends.cudnn.benchmark = False
torch.set_float32_matmul_precision('high') torch.set_float32_matmul_precision("high")
class Pipeline(LightningModule): class Pipeline(LightningModule):
@@ -47,7 +51,7 @@ class Pipeline(LightningModule):
max_steps: int = 200000, max_steps: int = 200000,
warmup_steps: int = 4000, warmup_steps: int = 4000,
dataset_path: str = "./data/your_dataset_path", dataset_path: str = "./data/your_dataset_path",
lora_config_path: str = None lora_config_path: str = None,
): ):
super().__init__() super().__init__()
@@ -91,15 +95,31 @@ class Pipeline(LightningModule):
if self.is_train: if self.is_train:
self.transformers.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.mert_model.requires_grad_(False)
self.resampler_mert = torchaudio.transforms.Resample(orig_freq=48000, new_freq=24000) self.resampler_mert = torchaudio.transforms.Resample(
self.processor_mert = Wav2Vec2FeatureExtractor.from_pretrained("m-a-p/MERT-v1-330M", trust_remote_code=True) orig_freq=48000, new_freq=24000
)
self.processor_mert = Wav2Vec2FeatureExtractor.from_pretrained(
"m-a-p/MERT-v1-330M", trust_remote_code=True
)
self.hubert_model = AutoModel.from_pretrained("utter-project/mHuBERT-147", local_files_only=True, cache_dir=checkpoint_dir).eval() self.hubert_model = AutoModel.from_pretrained(
"utter-project/mHuBERT-147",
local_files_only=True,
cache_dir=checkpoint_dir,
).eval()
self.hubert_model.requires_grad_(False) self.hubert_model.requires_grad_(False)
self.resampler_mhubert = torchaudio.transforms.Resample(orig_freq=48000, new_freq=16000) self.resampler_mhubert = torchaudio.transforms.Resample(
self.processor_mhubert = Wav2Vec2FeatureExtractor.from_pretrained("utter-project/mHuBERT-147", local_files_only=True, cache_dir=checkpoint_dir) orig_freq=48000, new_freq=16000
)
self.processor_mhubert = Wav2Vec2FeatureExtractor.from_pretrained(
"utter-project/mHuBERT-147",
local_files_only=True,
cache_dir=checkpoint_dir,
)
self.ssl_coeff = ssl_coeff self.ssl_coeff = ssl_coeff
@@ -110,9 +130,21 @@ class Pipeline(LightningModule):
actual_lengths_24k = wav_lengths // 2 # 48kHz -> 24kHz actual_lengths_24k = wav_lengths // 2 # 48kHz -> 24kHz
# Normalize the actual audio part # Normalize the actual audio part
means = torch.stack([mert_input_wavs_mono_24k[i, :actual_lengths_24k[i]].mean() for i in range(bsz)]) means = torch.stack(
vars = torch.stack([mert_input_wavs_mono_24k[i, :actual_lengths_24k[i]].var() for i in range(bsz)]) [
mert_input_wavs_mono_24k = (mert_input_wavs_mono_24k - means.view(-1, 1)) / torch.sqrt(vars.view(-1, 1) + 1e-7) mert_input_wavs_mono_24k[i, : actual_lengths_24k[i]].mean()
for i in range(bsz)
]
)
vars = torch.stack(
[
mert_input_wavs_mono_24k[i, : actual_lengths_24k[i]].var()
for i in range(bsz)
]
)
mert_input_wavs_mono_24k = (
mert_input_wavs_mono_24k - means.view(-1, 1)
) / torch.sqrt(vars.view(-1, 1) + 1e-7)
# MERT SSL constraint # MERT SSL constraint
# Define the length of each chunk (5 seconds of samples) # Define the length of each chunk (5 seconds of samples)
@@ -131,7 +163,9 @@ class Pipeline(LightningModule):
end = min(start + chunk_size, actual_length) end = min(start + chunk_size, actual_length)
chunk = audio[start:end] chunk = audio[start:end]
if len(chunk) < chunk_size: 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) all_chunks.append(chunk)
chunk_actual_lengths.append(end - start) 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] chunk_num_features = [(length + 319) // 320 for length in chunk_actual_lengths]
# Trim the hidden states of each chunk # 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 # Organize hidden states by audio
mert_ssl_hidden_states_list = [] mert_ssl_hidden_states_list = []
chunk_idx = 0 chunk_idx = 0
for i in range(bsz): for i in range(bsz):
audio_chunks = chunk_hidden_states[chunk_idx:chunk_idx + num_chunks_per_audio[i]] audio_chunks = chunk_hidden_states[
audio_hidden = torch.cat(audio_chunks, dim=0) # Concatenate chunks of the same audio chunk_idx : chunk_idx + num_chunks_per_audio[i]
]
audio_hidden = torch.cat(
audio_chunks, dim=0
) # Concatenate chunks of the same audio
mert_ssl_hidden_states_list.append(audio_hidden) mert_ssl_hidden_states_list.append(audio_hidden)
chunk_idx += num_chunks_per_audio[i] 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 actual_lengths_16k = wav_lengths // 3 # Convert lengths from 48kHz to 16kHz
# Step 2: Zero-mean unit-variance normalization (only on actual audio) # Step 2: Zero-mean unit-variance normalization (only on actual audio)
means = torch.stack([mhubert_input_wavs_mono_16k[i, :actual_lengths_16k[i]].mean() means = torch.stack(
for i in range(bsz)]) [
vars = torch.stack([mhubert_input_wavs_mono_16k[i, :actual_lengths_16k[i]].var() mhubert_input_wavs_mono_16k[i, : actual_lengths_16k[i]].mean()
for i in range(bsz)]) for i in range(bsz)
mhubert_input_wavs_mono_16k = (mhubert_input_wavs_mono_16k - means.view(-1, 1)) / \ ]
torch.sqrt(vars.view(-1, 1) + 1e-7) )
vars = torch.stack(
[
mhubert_input_wavs_mono_16k[i, : actual_lengths_16k[i]].var()
for i in range(bsz)
]
)
mhubert_input_wavs_mono_16k = (
mhubert_input_wavs_mono_16k - means.view(-1, 1)
) / torch.sqrt(vars.view(-1, 1) + 1e-7)
# Step 3: Define chunk size for MHubert (30 seconds at 16kHz) # Step 3: Define chunk size for MHubert (30 seconds at 16kHz)
chunk_size = 16000 * 30 # 30 seconds = 480,000 samples chunk_size = 16000 * 30 # 30 seconds = 480,000 samples
# Step 4: Split audio into chunks # 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 = [] all_chunks = []
chunk_actual_lengths = [] chunk_actual_lengths = []
@@ -206,20 +258,33 @@ class Pipeline(LightningModule):
chunk_num_features = [(length + 319) // 320 for length in chunk_actual_lengths] chunk_num_features = [(length + 319) // 320 for length in chunk_actual_lengths]
# Step 8: Trim hidden states to remove padding effects # 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 # Step 9: Reorganize hidden states by original audio
mhubert_ssl_hidden_states_list = [] mhubert_ssl_hidden_states_list = []
chunk_idx = 0 chunk_idx = 0
for i in range(bsz): for i in range(bsz):
audio_chunks = chunk_hidden_states[chunk_idx:chunk_idx + num_chunks_per_audio[i]] audio_chunks = chunk_hidden_states[
audio_hidden = torch.cat(audio_chunks, dim=0) # Concatenate chunks for this audio chunk_idx : chunk_idx + num_chunks_per_audio[i]
]
audio_hidden = torch.cat(
audio_chunks, dim=0
) # Concatenate chunks for this audio
mhubert_ssl_hidden_states_list.append(audio_hidden) mhubert_ssl_hidden_states_list.append(audio_hidden)
chunk_idx += num_chunks_per_audio[i] chunk_idx += num_chunks_per_audio[i]
return mhubert_ssl_hidden_states_list return mhubert_ssl_hidden_states_list
def get_text_embeddings(self, texts, device, text_max_length=256): 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()} inputs = {key: value.to(device) for key, value in inputs.items()}
if self.text_encoder_model.device != device: if self.text_encoder_model.device != device:
self.text_encoder_model.to(device) self.text_encoder_model.to(device)
@@ -243,15 +308,21 @@ class Pipeline(LightningModule):
if train: if train:
with torch.amp.autocast(device_type="cuda", dtype=dtype): with torch.amp.autocast(device_type="cuda", dtype=dtype):
mert_ssl_hidden_states = self.infer_mert_ssl(target_wavs, wav_lengths) 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 # 1: text embedding
texts = batch["prompts"] 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) encoder_text_hidden_states = encoder_text_hidden_states.to(dtype)
target_latents, _ = self.dcae.encode(target_wavs, wav_lengths) 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) speaker_embds = batch["speaker_embs"].to(dtype)
keys = batch["keys"] keys = batch["keys"]
@@ -263,27 +334,43 @@ class Pipeline(LightningModule):
full_cfg_condition_mask = torch.where( full_cfg_condition_mask = torch.where(
(torch.rand(size=(bs,), device=device) < 0.15), (torch.rand(size=(bs,), device=device) < 0.15),
torch.zeros(size=(bs,), device=device), torch.zeros(size=(bs,), device=device),
torch.ones(size=(bs,), device=device) torch.ones(size=(bs,), device=device),
).long() ).long()
# N x T x 768 # 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( full_cfg_condition_mask = torch.where(
(torch.rand(size=(bs,), device=device) < 0.50), (torch.rand(size=(bs,), device=device) < 0.50),
torch.zeros(size=(bs,), device=device), torch.zeros(size=(bs,), device=device),
torch.ones(size=(bs,), device=device) torch.ones(size=(bs,), device=device),
).long() ).long()
# N x 512 # 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 # Lyrics
full_cfg_condition_mask = torch.where( full_cfg_condition_mask = torch.where(
(torch.rand(size=(bs,), device=device) < 0.15), (torch.rand(size=(bs,), device=device) < 0.15),
torch.zeros(size=(bs,), device=device), torch.zeros(size=(bs,), device=device),
torch.ones(size=(bs,), device=device) torch.ones(size=(bs,), device=device),
).long() ).long()
lyric_token_ids = torch.where(full_cfg_condition_mask.unsqueeze(1).bool(), lyric_token_ids, torch.zeros_like(lyric_token_ids)) lyric_token_ids = torch.where(
lyric_mask = torch.where(full_cfg_condition_mask.unsqueeze(1).bool(), lyric_mask, torch.zeros_like(lyric_mask)) full_cfg_condition_mask.unsqueeze(1).bool(),
lyric_token_ids,
torch.zeros_like(lyric_token_ids),
)
lyric_mask = torch.where(
full_cfg_condition_mask.unsqueeze(1).bool(),
lyric_mask,
torch.zeros_like(lyric_mask),
)
return ( return (
keys, keys,
@@ -305,10 +392,12 @@ class Pipeline(LightningModule):
) )
def configure_optimizers(self): 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( optimizer = torch.optim.AdamW(
params=[ params=[
{'params': trainable_params}, {"params": trainable_params},
], ],
lr=self.hparams.learning_rate, lr=self.hparams.learning_rate,
weight_decay=self.hparams.weight_decay, weight_decay=self.hparams.weight_decay,
@@ -324,13 +413,13 @@ class Pipeline(LightningModule):
return float(current_step) / float(max(1, warmup_steps)) return float(current_step) / float(max(1, warmup_steps))
else: else:
# Linear decay from learning_rate to 0 # 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) return max(0.0, 1.0 - progress)
lr_scheduler = torch.optim.lr_scheduler.LambdaLR( lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer, optimizer, lr_lambda, last_epoch=-1
lr_lambda,
last_epoch=-1
) )
return [optimizer], lr_scheduler return [optimizer], lr_scheduler
@@ -362,10 +451,17 @@ class Pipeline(LightningModule):
# See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$). # 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) # 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 # 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) u = torch.nn.functional.sigmoid(u)
indices = (u * self.scheduler.config.num_train_timesteps).long() 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) timesteps = self.scheduler.timesteps[indices].to(device)
return timesteps return timesteps
@@ -394,7 +490,9 @@ class Pipeline(LightningModule):
timesteps = self.get_timestep(bsz, device) timesteps = self.get_timestep(bsz, device)
# Add noise according to flow matching. # 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 noisy_image = sigmas * noise + (1.0 - sigmas) * target_image
# This is the flow-matching target for vanilla SD3. # 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 # Compute loss. Only calculate loss where chunk_mask is 1 and there is no padding
# N x T x 64 # N x T x 64
# N x T -> N x c x W x T # 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_model_pred = (model_pred * mask).reshape(bsz, -1).contiguous()
selected_target = (target * mask).reshape(bsz, -1).contiguous() selected_target = (target * mask).reshape(bsz, -1).contiguous()
@@ -443,11 +545,19 @@ class Pipeline(LightningModule):
prefix = "train" 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 total_proj_loss = 0.0
for k, v in proj_losses: 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 total_proj_loss += v
if len(proj_losses) > 0: if len(proj_losses) > 0:
@@ -459,7 +569,13 @@ class Pipeline(LightningModule):
# Log learning rate if scheduler exists # Log learning rate if scheduler exists
if self.lr_schedulers() is not None: if self.lr_schedulers() is not None:
learning_rate = self.lr_schedulers().get_last_lr()[0] 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(): # with torch.autograd.detect_anomaly():
# self.manual_backward(loss) # self.manual_backward(loss)
return loss return loss
@@ -496,18 +612,35 @@ class Pipeline(LightningModule):
) )
frame_length = int(duration * 44100 / 512 / 8) 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) attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
if do_classifier_free_guidance: if do_classifier_free_guidance:
attention_mask = torch.cat([attention_mask] * 2, dim=0) 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) 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) lyric_mask = torch.cat([lyric_mask, torch.zeros_like(lyric_mask)], 0)
momentum_buffer = MomentumBuffer() momentum_buffer = MomentumBuffer()
@@ -515,7 +648,9 @@ class Pipeline(LightningModule):
for i, t in tqdm(enumerate(timesteps), total=num_inference_steps): for i, t in tqdm(enumerate(timesteps), total=num_inference_steps):
# expand the latents if we are doing classifier free guidance # expand the latents if we are doing classifier free guidance
latents = target_latents 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]) timestep = t.expand(latent_model_input.shape[0])
noise_pred = self.transformers( noise_pred = self.transformers(
hidden_states=latent_model_input, hidden_states=latent_model_input,
@@ -537,7 +672,13 @@ class Pipeline(LightningModule):
momentum_buffer=momentum_buffer, 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 return target_latents
@@ -582,7 +723,9 @@ class Pipeline(LightningModule):
) )
audio_lengths = batch["wav_lengths"] 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 { return {
"target_wavs": batch["target_wavs"], "target_wavs": batch["target_wavs"],
"pred_wavs": pred_wavs, "pred_wavs": pred_wavs,
@@ -603,7 +746,12 @@ class Pipeline(LightningModule):
def plot_step(self, batch, batch_idx): def plot_step(self, batch, batch_idx):
global_step = self.global_step 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 return
results = self.predict_step(batch) results = self.predict_step(batch)
@@ -615,7 +763,9 @@ class Pipeline(LightningModule):
sr = results["sr"] sr = results["sr"]
seeds = results["seeds"] seeds = results["seeds"]
i = 0 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 key = key
prompt = prompt prompt = prompt
lyric = self.construct_lyrics(candidate_lyric_chunk) 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}" save_dir = f"{log_dir}/eval_results/step_{self.global_step}"
if not os.path.exists(save_dir): if not os.path.exists(save_dir):
os.makedirs(save_dir, exist_ok=True) os.makedirs(save_dir, exist_ok=True)
torchaudio.save(f"{save_dir}/target_wav_{key}_{i}.flac", target_wav.float().cpu(), sr) torchaudio.save(
torchaudio.save(f"{save_dir}/pred_wav_{key}_{i}.flac", pred_wav.float().cpu(), sr) f"{save_dir}/target_wav_{key}_{i}.wav", target_wav.float().cpu(), sr
with open(f"{save_dir}/key_prompt_lyric_{key}_{i}.txt", "w", encoding="utf-8") as f: )
torchaudio.save(
f"{save_dir}/pred_wav_{key}_{i}.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) f.write(key_prompt_lyric)
i += 1 i += 1
@@ -642,11 +798,14 @@ def main(args):
checkpoint_dir=args.checkpoint_dir, checkpoint_dir=args.checkpoint_dir,
) )
checkpoint_callback = ModelCheckpoint( 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 # add datetime str to version
logger_callback = TensorBoardLogger( 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( trainer = Trainer(
accelerator="gpu", accelerator="gpu",